Break the paragraph between tool rounds

Text from consecutive tool rounds arrived in Zed as one run-on block: "…parsed with unwrap_or_else.Let me broaden the search:Found them." Each round streams its own agent-message chunks, and ACP clients concatenate consecutive chunks, so round two picked up exactly where round one's last sentence ended. The streaming bookkeeping (assembled text, what's already on the wire, whether anything streamed) moves into a StreamedReply that both the ACP and headless loops share, and the loops mark the seam when a tool round interrupts the prose. The next visible fragment then opens with a blank line, taking the model's own leading whitespace with it. A round that produced only tool calls has nothing to break away from, so a reply never starts with a stray blank line. Covered end to end on both surfaces: a scripted two-round turn asserts the break reaches the client, and reverting the seam marker fails both tests with the run-on from the issue. Fixes #55 Co-Authored-By: siGit Code <noreply@sigit.si> (cherry picked from commit 9441ee0287b8188648d17ac2d6314778a8dca1ac)

Seto Elkahfi committed Sep 2, 2026 at 22:29 UTC e4eec6ac8011e689508a0854cbe516ebc7878b1a
4 files changed +334 -77
src/headless.rs
+17 -37
@@ -228,16 +228,12 @@ async fn run_prompt(
228 // sink is passed and only the final message is printed.
229 let (sink, mut sink_rx) = tokio::sync::mpsc::unbounded_channel::<String>();
230 let sink_opt = if config.quiet { None } else { Some(&sink) };
231 - let mut assembled = String::new();
232 - let mut sent = String::new();
233 - let mut streamed_any = false;
231 + let mut reply = crate::StreamedReply::default();
232
233 let mut result = drain_to_stdout(
234 backend.send_message_with_tools(&config.prompt, &tools, sink_opt),
235 &mut sink_rx,
238 - &mut assembled,
239 - &mut sent,
240 - &mut streamed_any,
236 + &mut reply,
237 )
238 .await?;
239
@@ -300,12 +296,14 @@ async fn run_prompt(
296 None // last round: force text
297 };
298
299 + // Whatever this round says starts a new paragraph rather than
300 + // continuing the sentence the tool calls interrupted.
301 + reply.interrupt();
302 +
303 result = drain_to_stdout(
304 backend.send_tool_results(tool_results, next_tools, sink_opt),
305 &mut sink_rx,
306 - &mut assembled,
307 - &mut sent,
308 - &mut streamed_any,
306 + &mut reply,
307 )
308 .await?;
309 }
@@ -318,7 +316,7 @@ async fn run_prompt(
316 if !final_visible.is_empty() {
317 let _ = writeln!(stdout, "{final_visible}");
318 }
321 - } else if streamed_any {
319 + } else if reply.streamed_any {
320 // The reply is already on stdout; end the line for the shell.
321 let _ = writeln!(stdout);
322 } else if !final_visible.is_empty() {
@@ -335,9 +333,7 @@ async fn run_prompt(
333 async fn drain_to_stdout<F>(
334 fut: F,
335 sink_rx: &mut tokio::sync::mpsc::UnboundedReceiver<String>,
338 - assembled: &mut String,
339 - sent: &mut String,
340 - streamed_any: &mut bool,
336 + reply: &mut crate::StreamedReply,
337 ) -> Result<TurnResult, backend::BackendError>
338 where
339 F: std::future::Future<Output = Result<TurnResult, backend::BackendError>>,
@@ -347,39 +343,23 @@ where
343 tokio::select! {
344 done = &mut fut => break done,
345 Some(piece) = sink_rx.recv() => {
350 - emit_visible_chunk(&piece, assembled, sent, streamed_any);
346 + emit_visible_chunk(&piece, reply);
347 }
348 }
349 };
350 // Flush tokens that landed between the last poll and the future resolving.
351 while let Ok(piece) = sink_rx.try_recv() {
356 - emit_visible_chunk(&piece, assembled, sent, streamed_any);
352 + emit_visible_chunk(&piece, reply);
353 }
354 result
355 }
356
361 -/// Append a streamed fragment, strip `<think>` reasoning from the running
362 -/// text, and print only the newly revealed visible suffix. Tracking the
363 -/// assembled text (not just deltas) keeps think-block stripping correct even
364 -/// when a tag spans chunk boundaries — same approach as the ACP path.
365 -fn emit_visible_chunk(
366 - piece: &str,
367 - assembled: &mut String,
368 - sent: &mut String,
369 - streamed_any: &mut bool,
370 -) {
371 - assembled.push_str(piece);
372 - let (_think, visible) = crate::chat::strip_think_blocks(assembled);
373 - match visible.strip_prefix(sent.as_str()) {
374 - Some(extra) if !extra.is_empty() => {
375 - print!("{extra}");
376 - let _ = std::io::stdout().flush();
377 - *sent = visible;
378 - *streamed_any = true;
379 - }
380 - // No new visible text, or the visible prefix changed retroactively
381 - // (rare, e.g. a late-closing think tag): resync without reprinting.
382 - _ => *sent = visible,
357 +/// Fold a streamed fragment into `reply` and print whatever it newly reveals —
358 +/// the stdio counterpart of `SiGitAgent::emit_visible_chunk`.
359 +fn emit_visible_chunk(piece: &str, reply: &mut crate::StreamedReply) {
360 + if let Some(extra) = reply.push(piece) {
361 + print!("{extra}");
362 + let _ = std::io::stdout().flush();
363 }
364 }
365
src/main.rs
+158 -39
@@ -248,6 +248,77 @@ pub(crate) fn system_prompt_for_model(tool_calling: bool) -> &'static str {
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.
@@ -728,18 +799,15 @@ impl SiGitAgent {
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>>,
@@ -749,44 +817,29 @@ impl SiGitAgent {
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
@@ -1469,9 +1522,7 @@ impl SiGitAgent {
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
@@ -1480,9 +1531,7 @@ impl SiGitAgent {
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| {
@@ -1672,15 +1721,17 @@ impl SiGitAgent {
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()))?;
@@ -1690,7 +1741,7 @@ impl SiGitAgent {
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 {
@@ -3618,6 +3669,74 @@ async fn main() -> anyhow::Result<()> {
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));
tests/acp_permissions.rs
+100 -1
@@ -16,7 +16,10 @@
16 //!
17 //! A second test covers live progress: a slow command's `in_progress` tool call
18 //! must reach the client while the command is still running, which is what the
19 -//! editor's spinner is driven from.
19 +//! editor's spinner is driven from. A third test covers what the client
20 +//! renders: text from two tool rounds must reach it as separate paragraphs,
21 +//! since ACP clients concatenate consecutive agent-message chunks into one
22 +//! block.
23
24 use std::collections::VecDeque;
25 use std::io::{BufRead, BufReader, Read, Write};
@@ -57,6 +60,21 @@ fn sse_text(text: &str) -> String {
60 sse_body(&[json!({"choices": [{"delta": {"content": text}}]})])
61 }
62
63 +/// One round that says something and then asks for a tool — the shape that used
64 +/// to leave the next round's text glued to this sentence.
65 +fn sse_text_then_tool_call(text: &str, id: &str, name: &str, arguments: &str) -> String {
66 + sse_body(&[
67 + json!({"choices": [{"delta": {"content": text}}]}),
68 + json!({
69 + "choices": [{"delta": {"tool_calls": [{
70 + "index": 0,
71 + "id": id,
72 + "function": {"name": name, "arguments": arguments},
73 + }]}}]
74 + }),
75 + ])
76 +}
77 +
78 /// Serves one scripted SSE response per request and records each request body.
79 struct FakeEndpoint {
80 port: u16,
@@ -241,6 +259,30 @@ impl AgentUnderTest {
259 message["params"]["update"].clone()
260 }
261
262 + /// The response to `session/prompt`, paired with every `agent_message_chunk`
263 + /// seen on the way, concatenated exactly as a client renders them.
264 + fn wait_for_prompt(&mut self, id: u64) -> (Value, String) {
265 + let mut rendered = String::new();
266 + let deadline = Instant::now() + TIMEOUT;
267 + loop {
268 + let remaining = deadline.saturating_duration_since(Instant::now());
269 + let Ok(message) = self.incoming.recv_timeout(remaining) else {
270 + panic!("timed out waiting for the prompt response");
271 + };
272 + let update = &message["params"]["update"];
273 + if message["method"] == "session/update"
274 + && update["sessionUpdate"] == "agent_message_chunk"
275 + && let Some(chunk) = update["content"]["text"].as_str()
276 + {
277 + rendered.push_str(chunk);
278 + }
279 + if message["id"] == id && message.get("method").is_none() {
280 + assert!(message.get("error").is_none(), "prompt failed: {message}");
281 + return (message, rendered);
282 + }
283 + }
284 + }
285 +
286 /// A request *from* the agent (has a `method` and its own id).
287 fn wait_for_agent_request(&mut self, method: &str) -> Value {
288 self.wait_for(&format!("agent request {method}"), |message| {
@@ -482,3 +524,60 @@ fn a_slow_commands_progress_reaches_the_client_while_it_runs() {
524 drop(agent);
525 let _ = std::fs::remove_dir_all(&scratch);
526 }
527 +
528 +#[test]
529 +fn text_from_two_tool_rounds_reaches_the_client_as_separate_paragraphs() {
530 + let endpoint = start_fake_endpoint(vec![
531 + // `list_directory` is read-only, so it runs without a permission gate
532 + // and the two rounds follow each other straight away.
533 + sse_text_then_tool_call(
534 + "I'll search for this pattern.",
535 + "call_1",
536 + "list_directory",
537 + r#"{"path":"."}"#,
538 + ),
539 + sse_text("Let me broaden the search:"),
540 + ]);
541 +
542 + let scratch = std::env::temp_dir().join(format!("sigit_acp_spacing_{}", std::process::id()));
543 + let config_dir = scratch.join("config");
544 + let cwd = scratch.join("cwd");
545 + std::fs::create_dir_all(&config_dir).unwrap();
546 + std::fs::create_dir_all(&cwd).unwrap();
547 +
548 + let mut agent = spawn_agent(endpoint.port, &config_dir);
549 +
550 + let id = agent.request(
551 + "initialize",
552 + json!({"protocolVersion": 1, "clientCapabilities": {}}),
553 + );
554 + agent.wait_for_response(id);
555 +
556 + let id = agent.request("session/new", json!({"cwd": cwd, "mcpServers": []}));
557 + let session_id = agent.wait_for_response(id)["result"]["sessionId"]
558 + .as_str()
559 + .expect("session id")
560 + .to_string();
561 +
562 + let prompt_id = agent.request(
563 + "session/prompt",
564 + json!({
565 + "sessionId": session_id,
566 + "prompt": [{"type": "text", "text": "find it"}],
567 + }),
568 + );
569 + let (response, rendered) = agent.wait_for_prompt(prompt_id);
570 + assert_eq!(response["result"]["stopReason"], "end_turn");
571 +
572 + assert!(
573 + rendered.contains("pattern.\n\nLet me broaden"),
574 + "the second round must open a new paragraph, got: {rendered:?}"
575 + );
576 + assert!(
577 + !rendered.contains("pattern.Let me"),
578 + "rounds must not run together into one sentence, got: {rendered:?}"
579 + );
580 +
581 + drop(agent);
582 + let _ = std::fs::remove_dir_all(&scratch);
583 +}
tests/headless_mode.rs
+59
@@ -46,6 +46,21 @@ fn sse_text(text: &str) -> String {
46 sse_body(&[json!({"choices": [{"delta": {"content": text}}]})])
47 }
48
49 +/// One round that says something and then asks for a tool — the shape that used
50 +/// to leave the next round's text glued to this sentence.
51 +fn sse_text_then_tool_call(text: &str, id: &str, name: &str, arguments: &str) -> String {
52 + sse_body(&[
53 + json!({"choices": [{"delta": {"content": text}}]}),
54 + json!({
55 + "choices": [{"delta": {"tool_calls": [{
56 + "index": 0,
57 + "id": id,
58 + "function": {"name": name, "arguments": arguments},
59 + }]}}]
60 + }),
61 + ])
62 +}
63 +
64 /// Serves one scripted SSE response per request and records each request body.
65 struct FakeEndpoint {
66 port: u16,
@@ -292,6 +307,50 @@ fn ask_level_tool_without_allow_flag_is_denied_but_run_completes() {
307 );
308 }
309
310 +#[test]
311 +fn text_from_two_tool_rounds_is_not_run_together() {
312 + let endpoint = start_fake_endpoint(vec![
313 + sse_text_then_tool_call(
314 + "I'll search for this pattern.",
315 + "call_1",
316 + "run_command",
317 + r#"{"command":"echo first"}"#,
318 + ),
319 + sse_text("Let me broaden the search:"),
320 + ]);
321 + let scratch = scratch("spacing");
322 + let cwd = scratch.cwd.to_str().unwrap().to_string();
323 +
324 + let output = run_headless(
325 + endpoint.port,
326 + &scratch,
327 + &[
328 + "-p",
329 + "find it",
330 + "--allow-tool",
331 + "run_command",
332 + "--cwd",
333 + &cwd,
334 + ],
335 + );
336 +
337 + assert_eq!(
338 + output.status.code(),
339 + Some(0),
340 + "stderr: {}",
341 + String::from_utf8_lossy(&output.stderr)
342 + );
343 + let stdout = stdout_of(&output);
344 + assert!(
345 + stdout.contains("pattern.\n\nLet me broaden"),
346 + "the second round must open a new paragraph, got: {stdout:?}"
347 + );
348 + assert!(
349 + !stdout.contains("pattern.Let me"),
350 + "rounds must not run together into one sentence, got: {stdout:?}"
351 + );
352 +}
353 +
354 #[test]
355 fn unknown_flag_exits_2_with_usage() {
356 let endpoint = start_fake_endpoint(vec![]);