248
/// context window, so the cap can afford to be generous
249
const MAX_TOOL_ROUNDS: usize = 24;
250
251
+/// Separator between two stretches of assistant prose that a tool round split
252
+/// apart. ACP clients concatenate consecutive agent-message chunks into a
253
+/// single block, so a plain newline (or nothing at all) leaves the rounds
254
+/// reading as one run-on paragraph.
255
+const PARAGRAPH_BREAK: &str = "\n\n";
256
+
257
+/// The assistant reply being streamed over the course of one prompt, tracked
258
+/// across every tool round so `<think>` stripping stays correct and no text is
259
+/// sent twice.
260
+#[derive(Default)]
261
+struct StreamedReply {
262
+ /// Raw text as it arrives, reasoning tags included: think-block stripping
263
+ /// needs the whole string, since a tag can span chunk boundaries.
264
+ assembled: String,
265
+ /// The visible text already on the wire.
266
+ sent: String,
267
+ /// Whether any chunk has been sent, i.e. the client already has a reply.
268
+ streamed_any: bool,
269
+ /// Whether the next visible text should open a new paragraph.
270
+ paragraph_pending: bool,
271
+}
272
+
273
+impl StreamedReply {
274
+ /// Record that a tool round has interrupted the prose, so the next chunk
275
+ /// starts a paragraph. A round that opened with tool calls and no text has
276
+ /// no sentence to break away from, hence the `streamed_any` guard.
277
+ fn interrupt(&mut self) {
278
+ self.paragraph_pending = self.streamed_any;
279
+ }
280
+
281
+ /// Fold one streamed fragment into the reply and return the visible text it
282
+ /// newly reveals, or `None` when it reveals nothing to show (it landed
283
+ /// inside a `<think>` block, say).
284
+ fn push(&mut self, piece: &str) -> Option<String> {
285
+ if self.paragraph_pending {
286
+ // A tool round interrupted the prose. Clients concatenate
287
+ // consecutive agent-message chunks into one block, so the new round
288
+ // has to open a paragraph or it runs onto the end of the previous
289
+ // sentence ("…parsed with unwrap_or_else.Let me broaden the
290
+ // search:"). Whatever whitespace the model puts at the seam gives
291
+ // way to the break.
292
+ let piece = piece.trim_start();
293
+ if piece.is_empty() {
294
+ return None;
295
+ }
296
+ self.assembled.push_str(PARAGRAPH_BREAK);
297
+ self.assembled.push_str(piece);
298
+ self.paragraph_pending = false;
299
+ } else {
300
+ self.assembled.push_str(piece);
301
+ }
302
+
303
+ let (_think, visible) = chat::strip_think_blocks(&self.assembled);
304
+ match visible.strip_prefix(self.sent.as_str()) {
305
+ Some(extra) if !extra.is_empty() => {
306
+ let extra = extra.to_string();
307
+ self.sent = visible;
308
+ self.streamed_any = true;
309
+ Some(extra)
310
+ }
311
+ // No new visible text, or the visible prefix changed retroactively
312
+ // (rare, e.g. a late-closing think tag): just resync without
313
+ // resending what's already on the wire.
314
+ _ => {
315
+ self.sent = visible;
316
+ None
317
+ }
318
+ }
319
+ }
320
+}
321
+
322
/// Outcome of asking the client for permission to run one tool call.
323
enum PermissionVerdict {
324
/// Run the tool.
799
/// tokens to the editor. The sink receiver is drained as the future runs, so
800
/// chunks reach the client live rather than all at once when it resolves.
801
///
731
- /// `assembled`/`sent`/`streamed_any` persist across the turns of a single
732
- /// prompt so reasoning is stripped consistently and we never re-send text.
733
- #[allow(clippy::too_many_arguments)]
802
+ /// `reply` persists across the turns of a single prompt so reasoning is
803
+ /// stripped consistently and we never re-send text.
804
async fn drain_turn<F>(
805
&self,
806
cx: &ConnectionTo<Client>,
807
session_id: &SessionId,
808
fut: F,
809
sink_rx: &mut tokio::sync::mpsc::UnboundedReceiver<String>,
740
- assembled: &mut String,
741
- sent: &mut String,
742
- streamed_any: &mut bool,
810
+ reply: &mut StreamedReply,
811
) -> Result<TurnResult, backend::BackendError>
812
where
813
F: std::future::Future<Output = Result<TurnResult, backend::BackendError>>,
817
tokio::select! {
818
done = &mut fut => break done,
819
Some(piece) = sink_rx.recv() => {
752
- self.emit_visible_chunk(cx, session_id, &piece, assembled, sent, streamed_any);
820
+ self.emit_visible_chunk(cx, session_id, &piece, reply);
821
}
822
}
823
};
824
// Flush tokens that landed between the last poll and the future resolving.
825
while let Ok(piece) = sink_rx.try_recv() {
758
- self.emit_visible_chunk(cx, session_id, &piece, assembled, sent, streamed_any);
826
+ self.emit_visible_chunk(cx, session_id, &piece, reply);
827
}
828
result
829
}
830
763
- /// Append a streamed fragment, strip `<think>` reasoning from the running
764
- /// text, and send only the newly revealed visible suffix as a chunk. Tracking
765
- /// the assembled text (not just deltas) keeps think-block stripping correct
766
- /// even when a tag spans chunk boundaries.
831
+ /// Fold a streamed fragment into `reply` and send whatever it newly reveals
832
+ /// as an agent-message chunk.
833
fn emit_visible_chunk(
834
&self,
835
cx: &ConnectionTo<Client>,
836
session_id: &SessionId,
837
piece: &str,
772
- assembled: &mut String,
773
- sent: &mut String,
774
- streamed_any: &mut bool,
838
+ reply: &mut StreamedReply,
839
) {
776
- assembled.push_str(piece);
777
- let (_think, visible) = chat::strip_think_blocks(assembled);
778
- match visible.strip_prefix(sent.as_str()) {
779
- Some(extra) if !extra.is_empty() => {
780
- let extra = extra.to_string();
781
- *sent = visible;
782
- *streamed_any = true;
783
- self.send_assistant_message(cx, session_id.clone(), extra)
784
- .ok();
785
- }
786
- // No new visible text, or the visible prefix changed retroactively
787
- // (rare, e.g. a late-closing think tag): just resync without
788
- // resending what's already on the wire.
789
- _ => *sent = visible,
840
+ if let Some(extra) = reply.push(piece) {
841
+ self.send_assistant_message(cx, session_id.clone(), extra)
842
+ .ok();
843
}
844
}
845
1522
// alive for the whole prompt so `recv()` only ends when a turn future
1523
// resolves, never because every sender was dropped.
1524
let (sink, mut sink_rx) = tokio::sync::mpsc::unbounded_channel::<String>();
1472
- let mut assembled = String::new();
1473
- let mut sent = String::new();
1474
- let mut streamed_any = false;
1525
+ let mut reply = StreamedReply::default();
1526
let mut repeated_tool_calls = std::collections::HashMap::<String, usize>::new();
1527
1528
let mut result = self
1531
&session_id,
1532
backend.send_message_with_tools(&user_text, &tools, Some(&sink)),
1533
&mut sink_rx,
1483
- &mut assembled,
1484
- &mut sent,
1485
- &mut streamed_any,
1534
+ &mut reply,
1535
)
1536
.await
1537
.map_err(|error| {
1721
None // last round: force text
1722
};
1723
1724
+ // Whatever this round says starts a new paragraph rather than
1725
+ // continuing the sentence the tool calls interrupted.
1726
+ reply.interrupt();
1727
+
1728
result = self
1729
.drain_turn(
1730
cx,
1731
&session_id,
1732
backend.send_tool_results(tool_results, next_tools, Some(&sink)),
1733
&mut sink_rx,
1681
- &mut assembled,
1682
- &mut sent,
1683
- &mut streamed_any,
1734
+ &mut reply,
1735
)
1736
.await
1737
.map_err(|e| agent_client_protocol::Error::new(-32603, e.to_string()))?;
1741
// If anything streamed, the visible reply is already on the wire; only
1742
// send a trailing block for the non-streamed path (e.g. on-device direct
1743
// answers, which onde can't stream while tools are on offer).
1693
- if !streamed_any {
1744
+ if !reply.streamed_any {
1745
let reply_text = result.text.trim().to_string();
1746
let final_text = if reply_text.is_empty() {
1747
if round > 0 {
3669
mod tests {
3670
use super::*;
3671
3672
+ #[test]
3673
+ fn streamed_reply_sends_each_fragment_as_it_is_revealed() {
3674
+ let mut reply = StreamedReply::default();
3675
+ assert_eq!(reply.push("Hello"), Some("Hello".to_string()));
3676
+ assert_eq!(reply.push(", world"), Some(", world".to_string()));
3677
+ assert!(reply.streamed_any);
3678
+ }
3679
+
3680
+ #[test]
3681
+ fn streamed_reply_opens_a_paragraph_after_a_tool_round() {
3682
+ let mut reply = StreamedReply::default();
3683
+ reply.push("I'll search for this pattern.");
3684
+
3685
+ // …tool calls run, then the next round starts talking again.
3686
+ reply.interrupt();
3687
+
3688
+ assert_eq!(
3689
+ reply.push("Let me broaden the search:"),
3690
+ Some("\n\nLet me broaden the search:".to_string()),
3691
+ "the new round must not run onto the end of the last sentence"
3692
+ );
3693
+ // Only the seam gets a break; the rest of the round streams as usual.
3694
+ assert_eq!(reply.push(" Found them."), Some(" Found them.".to_string()));
3695
+ }
3696
+
3697
+ #[test]
3698
+ fn streamed_reply_drops_the_models_own_whitespace_at_the_seam() {
3699
+ let mut reply = StreamedReply::default();
3700
+ reply.push("Checking that.");
3701
+ reply.interrupt();
3702
+
3703
+ // A fragment of pure whitespace reveals nothing and must not consume
3704
+ // the pending break.
3705
+ assert_eq!(reply.push("\n"), None);
3706
+ assert_eq!(
3707
+ reply.push(" Done."),
3708
+ Some("\n\nDone.".to_string()),
3709
+ "the break replaces the model's leading whitespace"
3710
+ );
3711
+ }
3712
+
3713
+ #[test]
3714
+ fn streamed_reply_does_not_open_with_a_blank_paragraph() {
3715
+ // First round was tool calls only, so nothing is on the wire yet and
3716
+ // there is no sentence to break away from.
3717
+ let mut reply = StreamedReply::default();
3718
+ reply.interrupt();
3719
+
3720
+ assert_eq!(reply.push("Found them."), Some("Found them.".to_string()));
3721
+ }
3722
+
3723
+ #[test]
3724
+ fn streamed_reply_hides_reasoning_and_still_breaks_the_paragraph() {
3725
+ let mut reply = StreamedReply::default();
3726
+ reply.push("Looking now.");
3727
+ reply.interrupt();
3728
+
3729
+ // A round that opens with reasoning reveals nothing until the visible
3730
+ // text arrives — and that text still starts the new paragraph.
3731
+ assert_eq!(reply.push("<think>weigh"), None);
3732
+ assert_eq!(reply.push(" options</think>"), None);
3733
+ assert_eq!(
3734
+ reply.push("Here it is."),
3735
+ Some("\n\nHere it is.".to_string()),
3736
+ "the break waits for the reasoning to end and rides out with the text"
3737
+ );
3738
+ }
3739
+
3740
#[test]
3741
fn explicit_acp_flag_is_accepted_without_other_arguments() {
3742
assert_eq!(parse_explicit_acp(&["--acp".to_string()]), Ok(true));