94
use onde::inference::{ChatEngine, GgufModelConfig};
95
96
use crate::backend::{
97
- ImageInput, InferenceBackend, LocalBackend, OpenAiBackend, ToolResult as BackendToolResult,
98
- ToolSpec, TurnResult,
97
+ ImageInput, InferenceBackend, LocalBackend, OpenAiBackend, TokenChunk,
98
+ ToolResult as BackendToolResult, ToolSpec, TurnResult,
99
};
100
use std::path::PathBuf;
101
use std::sync::atomic::{AtomicBool, Ordering};
355
assembled: String,
356
/// The visible text already on the wire.
357
sent: String,
358
+ /// The inline-`<think>` reasoning text already on the wire.
359
+ sent_think: String,
360
/// Whether any chunk has been sent, i.e. the client already has a reply.
361
streamed_any: bool,
362
/// Whether the next visible text should open a new paragraph.
363
paragraph_pending: bool,
364
}
365
366
+/// What a streamed fragment newly revealed once `<think>` blocks were split:
367
+/// visible reply text and/or inline reasoning text, either of which may be
368
+/// absent for a given chunk.
369
+#[derive(Default)]
370
+struct RevealedText {
371
+ visible: Option<String>,
372
+ reasoning: Option<String>,
373
+}
374
+
375
impl StreamedReply {
376
/// Record that a tool round has interrupted the prose, so the next chunk
377
/// starts a paragraph. A round that opened with tool calls and no text has
380
self.paragraph_pending = self.streamed_any;
381
}
382
372
- /// Fold one streamed fragment into the reply and return the visible text it
373
- /// newly reveals, or `None` when it reveals nothing to show (it landed
374
- /// inside a `<think>` block, say).
375
- fn push(&mut self, piece: &str) -> Option<String> {
383
+ /// Fold one streamed fragment into the reply and return the text it newly
384
+ /// reveals: visible reply text (sent as an agent message) and/or inline
385
+ /// `<think>` reasoning (sent as an agent thought). Either may be absent when
386
+ /// the fragment reveals nothing of that kind.
387
+ fn push(&mut self, piece: &str) -> RevealedText {
388
if self.paragraph_pending {
389
// A tool round interrupted the prose. Clients concatenate
390
// consecutive agent-message chunks into one block, so the new round
394
// way to the break.
395
let piece = piece.trim_start();
396
if piece.is_empty() {
385
- return None;
397
+ return RevealedText::default();
398
}
399
self.assembled.push_str(PARAGRAPH_BREAK);
400
self.assembled.push_str(piece);
403
self.assembled.push_str(piece);
404
}
405
394
- let (_think, visible) = chat::strip_think_blocks(&self.assembled);
406
+ let (think, visible) = chat::strip_think_blocks(&self.assembled);
407
+ let mut revealed = RevealedText::default();
408
+
409
match visible.strip_prefix(self.sent.as_str()) {
410
Some(extra) if !extra.is_empty() => {
397
- let extra = extra.to_string();
411
+ revealed.visible = Some(extra.to_string());
412
self.sent = visible;
413
self.streamed_any = true;
400
- Some(extra)
414
}
415
// No new visible text, or the visible prefix changed retroactively
416
// (rare, e.g. a late-closing think tag): just resync without
417
// resending what's already on the wire.
418
_ => {
419
self.sent = visible;
407
- None
420
}
421
}
422
+
423
+ match think.strip_prefix(self.sent_think.as_str()) {
424
+ Some(extra) if !extra.is_empty() => {
425
+ revealed.reasoning = Some(extra.to_string());
426
+ self.sent_think = think;
427
+ }
428
+ // No new reasoning, or it changed retroactively: resync quietly.
429
+ _ => {
430
+ self.sent_think = think;
431
+ }
432
+ }
433
+
434
+ revealed
435
}
436
}
437
1425
))
1426
}
1427
1428
+ /// Stream a fragment of the model's reasoning as a thought chunk. Clients
1429
+ /// render these in a collapsed "thinking" section, separate from the reply.
1430
+ fn send_agent_thought(
1431
+ &self,
1432
+ cx: &ConnectionTo<Client>,
1433
+ session_id: SessionId,
1434
+ text: impl Into<String>,
1435
+ ) -> agent_client_protocol::Result<()> {
1436
+ cx.send_notification(SessionNotification::new(
1437
+ session_id,
1438
+ SessionUpdate::AgentThoughtChunk(ContentChunk::new(ContentBlock::from(text.into()))),
1439
+ ))
1440
+ }
1441
+
1442
fn send_system_status(
1443
&self,
1444
cx: &ConnectionTo<Client>,
1469
cx: &ConnectionTo<Client>,
1470
session_id: &SessionId,
1471
fut: F,
1433
- sink_rx: &mut tokio::sync::mpsc::UnboundedReceiver<String>,
1472
+ sink_rx: &mut tokio::sync::mpsc::UnboundedReceiver<TokenChunk>,
1473
reply: &mut StreamedReply,
1474
cancellation: &PromptCancellation,
1475
) -> Result<TurnResult, DrainTurnError>
1489
break Err(DrainTurnError::Cancelled);
1490
}
1491
}
1453
- Some(piece) = sink_rx.recv() => {
1454
- self.emit_visible_chunk(cx, session_id, &piece, reply);
1492
+ Some(chunk) = sink_rx.recv() => {
1493
+ self.emit_token_chunk(cx, session_id, chunk, reply);
1494
}
1495
}
1496
};
1497
// Flush tokens that landed between the last poll and the future resolving.
1459
- while let Ok(piece) = sink_rx.try_recv() {
1460
- self.emit_visible_chunk(cx, session_id, &piece, reply);
1498
+ while let Ok(chunk) = sink_rx.try_recv() {
1499
+ self.emit_token_chunk(cx, session_id, chunk, reply);
1500
}
1501
result
1502
}
1503
1465
- /// Fold a streamed fragment into `reply` and send whatever it newly reveals
1466
- /// as an agent-message chunk.
1467
- fn emit_visible_chunk(
1504
+ /// Route one streamed token chunk. A `Reasoning` chunk comes from a
1505
+ /// dedicated reasoning field and is unambiguous, so it goes straight to the
1506
+ /// client as a thought. A `Visible` chunk is folded into `reply`, which may
1507
+ /// reveal visible reply text (sent as an agent message) and/or inline
1508
+ /// `<think>` reasoning (sent as a thought).
1509
+ fn emit_token_chunk(
1510
&self,
1511
cx: &ConnectionTo<Client>,
1512
session_id: &SessionId,
1471
- piece: &str,
1513
+ chunk: TokenChunk,
1514
reply: &mut StreamedReply,
1515
) {
1474
- if let Some(extra) = reply.push(piece) {
1475
- self.send_assistant_message(cx, session_id.clone(), extra)
1476
- .ok();
1516
+ match chunk {
1517
+ TokenChunk::Reasoning(text) => {
1518
+ self.send_agent_thought(cx, session_id.clone(), text).ok();
1519
+ }
1520
+ TokenChunk::Visible(piece) => {
1521
+ let revealed = reply.push(&piece);
1522
+ if let Some(reasoning) = revealed.reasoning {
1523
+ self.send_agent_thought(cx, session_id.clone(), reasoning)
1524
+ .ok();
1525
+ }
1526
+ if let Some(visible) = revealed.visible {
1527
+ self.send_assistant_message(cx, session_id.clone(), visible)
1528
+ .ok();
1529
+ }
1530
+ }
1531
}
1532
}
1533
3057
3058
// Token sink: backends stream assistant text through this while a turn
3059
// runs. We forward the visible portion to the editor as agent-message
3006
- // chunks live (see `drain_turn` / `emit_visible_chunk`). The sink stays
3007
- // alive for the whole prompt so `recv()` only ends when a turn future
3008
- // resolves, never because every sender was dropped.
3009
- let (sink, mut sink_rx) = tokio::sync::mpsc::unbounded_channel::<String>();
3060
+ // chunks live, and the reasoning portion as thought chunks (see
3061
+ // `drain_turn` / `emit_token_chunk`). The sink stays alive for the whole
3062
+ // prompt so `recv()` only ends when a turn future resolves, never
3063
+ // because every sender was dropped.
3064
+ let (sink, mut sink_rx) = tokio::sync::mpsc::unbounded_channel::<TokenChunk>();
3065
let mut reply = StreamedReply::default();
3066
let mut repeated_tool_calls = std::collections::HashMap::<String, usize>::new();
3067
6812
#[test]
6813
fn streamed_reply_sends_each_fragment_as_it_is_revealed() {
6814
let mut reply = StreamedReply::default();
6760
- assert_eq!(reply.push("Hello"), Some("Hello".to_string()));
6761
- assert_eq!(reply.push(", world"), Some(", world".to_string()));
6815
+ assert_eq!(reply.push("Hello").visible, Some("Hello".to_string()));
6816
+ assert_eq!(reply.push(", world").visible, Some(", world".to_string()));
6817
assert!(reply.streamed_any);
6818
}
6819
6826
reply.interrupt();
6827
6828
assert_eq!(
6774
- reply.push("Let me broaden the search:"),
6829
+ reply.push("Let me broaden the search:").visible,
6830
Some("\n\nLet me broaden the search:".to_string()),
6831
"the new round must not run onto the end of the last sentence"
6832
);
6833
// Only the seam gets a break; the rest of the round streams as usual.
6779
- assert_eq!(reply.push(" Found them."), Some(" Found them.".to_string()));
6834
+ assert_eq!(
6835
+ reply.push(" Found them.").visible,
6836
+ Some(" Found them.".to_string())
6837
+ );
6838
}
6839
6840
#[test]
6845
6846
// A fragment of pure whitespace reveals nothing and must not consume
6847
// the pending break.
6790
- assert_eq!(reply.push("\n"), None);
6848
+ assert_eq!(reply.push("\n").visible, None);
6849
assert_eq!(
6792
- reply.push(" Done."),
6850
+ reply.push(" Done.").visible,
6851
Some("\n\nDone.".to_string()),
6852
"the break replaces the model's leading whitespace"
6853
);
6860
let mut reply = StreamedReply::default();
6861
reply.interrupt();
6862
6805
- assert_eq!(reply.push("Found them."), Some("Found them.".to_string()));
6863
+ assert_eq!(
6864
+ reply.push("Found them.").visible,
6865
+ Some("Found them.".to_string())
6866
+ );
6867
}
6868
6869
#[test]
6809
- fn streamed_reply_hides_reasoning_and_still_breaks_the_paragraph() {
6870
+ fn streamed_reply_keeps_reasoning_out_of_the_reply_and_still_breaks_the_paragraph() {
6871
let mut reply = StreamedReply::default();
6872
reply.push("Looking now.");
6873
reply.interrupt();
6874
6814
- // A round that opens with reasoning reveals nothing until the visible
6815
- // text arrives — and that text still starts the new paragraph.
6816
- assert_eq!(reply.push("<think>weigh"), None);
6817
- assert_eq!(reply.push(" options</think>"), None);
6875
+ // A round that opens with reasoning reveals no *visible* text until the
6876
+ // reply arrives — and that text still starts the new paragraph.
6877
+ assert_eq!(reply.push("<think>weigh").visible, None);
6878
+ assert_eq!(reply.push(" options</think>").visible, None);
6879
assert_eq!(
6819
- reply.push("Here it is."),
6880
+ reply.push("Here it is.").visible,
6881
Some("\n\nHere it is.".to_string()),
6882
"the break waits for the reasoning to end and rides out with the text"
6883
);
6884
}
6885
6886
+ #[test]
6887
+ fn streamed_reply_surfaces_inline_think_reasoning_incrementally() {
6888
+ let mut reply = StreamedReply::default();
6889
+
6890
+ // Inline `<think>` content is routed as reasoning, not reply text, and
6891
+ // only the newly-revealed suffix is emitted each time.
6892
+ let r = reply.push("<think>weigh");
6893
+ assert_eq!(r.visible, None);
6894
+ assert_eq!(r.reasoning, Some("weigh".to_string()));
6895
+
6896
+ let r = reply.push(" options</think>");
6897
+ assert_eq!(r.visible, None);
6898
+ assert_eq!(r.reasoning, Some(" options".to_string()));
6899
+
6900
+ // The visible answer carries no reasoning.
6901
+ let r = reply.push("Here it is.");
6902
+ assert_eq!(r.visible, Some("Here it is.".to_string()));
6903
+ assert_eq!(r.reasoning, None);
6904
+ }
6905
+
6906
#[test]
6907
fn explicit_acp_flag_is_accepted_without_other_arguments() {
6908
assert_eq!(parse_explicit_acp(&["--acp".to_string()]), Ok(true));