@sigitcodereview / sigit / commits / 9441ee0

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>

Seto Elkahfi committed Sep 2, 2026 at 22:29 UTC 9441ee0287b8188648d17ac2d6314778a8dca1ac
4 files changed +334 -76
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
@@ -1454,9 +1507,7 @@ impl SiGitAgent {
1507 // alive for the whole prompt so `recv()` only ends when a turn future
1508 // resolves, never because every sender was dropped.
1509 let (sink, mut sink_rx) = tokio::sync::mpsc::unbounded_channel::<String>();
1457 - let mut assembled = String::new();
1458 - let mut sent = String::new();
1459 - let mut streamed_any = false;
1510 + let mut reply = StreamedReply::default();
1511 let mut repeated_tool_calls = std::collections::HashMap::<String, usize>::new();
1512
1513 let mut result = self
@@ -1465,9 +1516,7 @@ impl SiGitAgent {
1516 &session_id,
1517 backend.send_message_with_tools(&user_text, &tools, Some(&sink)),
1518 &mut sink_rx,
1468 - &mut assembled,
1469 - &mut sent,
1470 - &mut streamed_any,
1519 + &mut reply,
1520 )
1521 .await
1522 .map_err(|error| {
@@ -1657,15 +1706,17 @@ impl SiGitAgent {
1706 None // last round: force text
1707 };
1708
1709 + // Whatever this round says starts a new paragraph rather than
1710 + // continuing the sentence the tool calls interrupted.
1711 + reply.interrupt();
1712 +
1713 result = self
1714 .drain_turn(
1715 cx,
1716 &session_id,
1717 backend.send_tool_results(tool_results, next_tools, Some(&sink)),
1718 &mut sink_rx,
1666 - &mut assembled,
1667 - &mut sent,
1668 - &mut streamed_any,
1719 + &mut reply,
1720 )
1721 .await
1722 .map_err(|e| agent_client_protocol::Error::new(-32603, e.to_string()))?;
@@ -1675,7 +1726,7 @@ impl SiGitAgent {
1726 // If anything streamed, the visible reply is already on the wire; only
1727 // send a trailing block for the non-streamed path (e.g. on-device direct
1728 // answers, which onde can't stream while tools are on offer).
1678 - if !streamed_any {
1729 + if !reply.streamed_any {
1730 let reply_text = result.text.trim().to_string();
1731 let final_text = if reply_text.is_empty() {
1732 if round > 0 {
@@ -3586,6 +3637,74 @@ async fn main() -> anyhow::Result<()> {
3637 mod tests {
3638 use super::*;
3639
3640 + #[test]
3641 + fn streamed_reply_sends_each_fragment_as_it_is_revealed() {
3642 + let mut reply = StreamedReply::default();
3643 + assert_eq!(reply.push("Hello"), Some("Hello".to_string()));
3644 + assert_eq!(reply.push(", world"), Some(", world".to_string()));
3645 + assert!(reply.streamed_any);
3646 + }
3647 +
3648 + #[test]
3649 + fn streamed_reply_opens_a_paragraph_after_a_tool_round() {
3650 + let mut reply = StreamedReply::default();
3651 + reply.push("I'll search for this pattern.");
3652 +
3653 + // …tool calls run, then the next round starts talking again.
3654 + reply.interrupt();
3655 +
3656 + assert_eq!(
3657 + reply.push("Let me broaden the search:"),
3658 + Some("\n\nLet me broaden the search:".to_string()),
3659 + "the new round must not run onto the end of the last sentence"
3660 + );
3661 + // Only the seam gets a break; the rest of the round streams as usual.
3662 + assert_eq!(reply.push(" Found them."), Some(" Found them.".to_string()));
3663 + }
3664 +
3665 + #[test]
3666 + fn streamed_reply_drops_the_models_own_whitespace_at_the_seam() {
3667 + let mut reply = StreamedReply::default();
3668 + reply.push("Checking that.");
3669 + reply.interrupt();
3670 +
3671 + // A fragment of pure whitespace reveals nothing and must not consume
3672 + // the pending break.
3673 + assert_eq!(reply.push("\n"), None);
3674 + assert_eq!(
3675 + reply.push(" Done."),
3676 + Some("\n\nDone.".to_string()),
3677 + "the break replaces the model's leading whitespace"
3678 + );
3679 + }
3680 +
3681 + #[test]
3682 + fn streamed_reply_does_not_open_with_a_blank_paragraph() {
3683 + // First round was tool calls only, so nothing is on the wire yet and
3684 + // there is no sentence to break away from.
3685 + let mut reply = StreamedReply::default();
3686 + reply.interrupt();
3687 +
3688 + assert_eq!(reply.push("Found them."), Some("Found them.".to_string()));
3689 + }
3690 +
3691 + #[test]
3692 + fn streamed_reply_hides_reasoning_and_still_breaks_the_paragraph() {
3693 + let mut reply = StreamedReply::default();
3694 + reply.push("Looking now.");
3695 + reply.interrupt();
3696 +
3697 + // A round that opens with reasoning reveals nothing until the visible
3698 + // text arrives — and that text still starts the new paragraph.
3699 + assert_eq!(reply.push("<think>weigh"), None);
3700 + assert_eq!(reply.push(" options</think>"), None);
3701 + assert_eq!(
3702 + reply.push("Here it is."),
3703 + Some("\n\nHere it is.".to_string()),
3704 + "the break waits for the reasoning to end and rides out with the text"
3705 + );
3706 + }
3707 +
3708 #[test]
3709 fn explicit_acp_flag_is_accepted_without_other_arguments() {
3710 assert_eq!(parse_explicit_acp(&["--acp".to_string()]), Ok(true));
tests/acp_permissions.rs
+100
@@ -13,6 +13,10 @@
13 //! would reject the whole session.
14 //! 2. `selected: allow_once` — the tool must actually execute and its output
15 //! travel back to the endpoint as a tool result.
16 +//!
17 +//! A second test covers what the client renders: text from two tool rounds must
18 +//! reach it as separate paragraphs, since ACP clients concatenate consecutive
19 +//! agent-message chunks into one block.
20
21 use std::collections::VecDeque;
22 use std::io::{BufRead, BufReader, Read, Write};
@@ -53,6 +57,21 @@ fn sse_text(text: &str) -> String {
57 sse_body(&[json!({"choices": [{"delta": {"content": text}}]})])
58 }
59
60 +/// One round that says something and then asks for a tool — the shape that used
61 +/// to leave the next round's text glued to this sentence.
62 +fn sse_text_then_tool_call(text: &str, id: &str, name: &str, arguments: &str) -> String {
63 + sse_body(&[
64 + json!({"choices": [{"delta": {"content": text}}]}),
65 + json!({
66 + "choices": [{"delta": {"tool_calls": [{
67 + "index": 0,
68 + "id": id,
69 + "function": {"name": name, "arguments": arguments},
70 + }]}}]
71 + }),
72 + ])
73 +}
74 +
75 /// Serves one scripted SSE response per request and records each request body.
76 struct FakeEndpoint {
77 port: u16,
@@ -208,6 +227,30 @@ impl AgentUnderTest {
227 response
228 }
229
230 + /// The response to `session/prompt`, paired with every `agent_message_chunk`
231 + /// seen on the way, concatenated exactly as a client renders them.
232 + fn wait_for_prompt(&mut self, id: u64) -> (Value, String) {
233 + let mut rendered = String::new();
234 + let deadline = Instant::now() + TIMEOUT;
235 + loop {
236 + let remaining = deadline.saturating_duration_since(Instant::now());
237 + let Ok(message) = self.incoming.recv_timeout(remaining) else {
238 + panic!("timed out waiting for the prompt response");
239 + };
240 + let update = &message["params"]["update"];
241 + if message["method"] == "session/update"
242 + && update["sessionUpdate"] == "agent_message_chunk"
243 + && let Some(chunk) = update["content"]["text"].as_str()
244 + {
245 + rendered.push_str(chunk);
246 + }
247 + if message["id"] == id && message.get("method").is_none() {
248 + assert!(message.get("error").is_none(), "prompt failed: {message}");
249 + return (message, rendered);
250 + }
251 + }
252 + }
253 +
254 /// A request *from* the agent (has a `method` and its own id).
255 fn wait_for_agent_request(&mut self, method: &str) -> Value {
256 self.wait_for(&format!("agent request {method}"), |message| {
@@ -357,3 +400,60 @@ fn permission_round_trip_cancel_then_allow() {
400 drop(agent);
401 let _ = std::fs::remove_dir_all(&scratch);
402 }
403 +
404 +#[test]
405 +fn text_from_two_tool_rounds_reaches_the_client_as_separate_paragraphs() {
406 + let endpoint = start_fake_endpoint(vec![
407 + // `list_directory` is read-only, so it runs without a permission gate
408 + // and the two rounds follow each other straight away.
409 + sse_text_then_tool_call(
410 + "I'll search for this pattern.",
411 + "call_1",
412 + "list_directory",
413 + r#"{"path":"."}"#,
414 + ),
415 + sse_text("Let me broaden the search:"),
416 + ]);
417 +
418 + let scratch = std::env::temp_dir().join(format!("sigit_acp_spacing_{}", std::process::id()));
419 + let config_dir = scratch.join("config");
420 + let cwd = scratch.join("cwd");
421 + std::fs::create_dir_all(&config_dir).unwrap();
422 + std::fs::create_dir_all(&cwd).unwrap();
423 +
424 + let mut agent = spawn_agent(endpoint.port, &config_dir);
425 +
426 + let id = agent.request(
427 + "initialize",
428 + json!({"protocolVersion": 1, "clientCapabilities": {}}),
429 + );
430 + agent.wait_for_response(id);
431 +
432 + let id = agent.request("session/new", json!({"cwd": cwd, "mcpServers": []}));
433 + let session_id = agent.wait_for_response(id)["result"]["sessionId"]
434 + .as_str()
435 + .expect("session id")
436 + .to_string();
437 +
438 + let prompt_id = agent.request(
439 + "session/prompt",
440 + json!({
441 + "sessionId": session_id,
442 + "prompt": [{"type": "text", "text": "find it"}],
443 + }),
444 + );
445 + let (response, rendered) = agent.wait_for_prompt(prompt_id);
446 + assert_eq!(response["result"]["stopReason"], "end_turn");
447 +
448 + assert!(
449 + rendered.contains("pattern.\n\nLet me broaden"),
450 + "the second round must open a new paragraph, got: {rendered:?}"
451 + );
452 + assert!(
453 + !rendered.contains("pattern.Let me"),
454 + "rounds must not run together into one sentence, got: {rendered:?}"
455 + );
456 +
457 + drop(agent);
458 + let _ = std::fs::remove_dir_all(&scratch);
459 +}
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![]);