From e2c96e9cb35fcb1fc864dd4a09d79206b8ac3ab0 Mon Sep 17 00:00:00 2001 From: Noah Shipley Date: Mon, 27 Jul 2026 23:13:08 -0500 Subject: [PATCH] feat(acp): stream an agent's reply into the channel as it is written MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ACP delivers an answer as a run of `agent_message_chunk` updates. The harness logged each one to `acp::stream` and dropped it, so a channel saw nothing until the agent posted a finished message. Every token was already arriving — it just never left the log. Publish it instead: post once on the first meaningful text, then edit that same message as more arrives. A stream is one message that grows, not a wall of fragments. This works cleanly because edits (kind 40003) are already excluded from CHANNEL_MESSAGE_EVENT_KINDS, so a streamed answer does not turn into hundreds of unread pings. Relaying every chunk would be worse than not streaming, so `stream_flush` holds a coalescing policy: - the first publish ignores the throttle — delaying the first sign of life is exactly backwards, that is when a reader most needs to know the agent is working; - mid-turn edits need both a minimum delta and a throttle interval; - end-of-turn always flushes, so the tail is never dropped; - an empty turn publishes nothing and never blanks an existing message; - bodies are clamped under build_edit's 64 KiB limit, truncating on a char boundary — a rejected edit would strand the message mid-answer, and a naive byte slice would panic on a multi-byte codepoint. The sender is scope-local in the turn driver on purpose: dropping it IS the end-of-turn signal, so success, cancel, idle timeout and agent exit all flush without any path having to remember to. The driver awaits the streamer so the complete reply is on the relay before the turn reports its outcome. Off by default (BUZZ_ACP_STREAM_REPLIES) — this changes what every channel sees. BUZZ_ACP_STREAM_THROTTLE_MS and BUZZ_ACP_STREAM_MIN_DELTA tune it. 18 tests: every policy branch, plus the post-once-then-edit contract driven through an injected publisher, covering the empty turn, the dropped tail, and a failed post — which must not leave the policy believing there is a message to edit, or the reply would vanish into orphan edits. Signed-off-by: Noah Shipley --- crates/buzz-acp/src/acp.rs | 20 + crates/buzz-acp/src/config.rs | 23 ++ crates/buzz-acp/src/lib.rs | 4 + crates/buzz-acp/src/pool.rs | 114 ++++++ crates/buzz-acp/src/stream_flush.rs | 570 ++++++++++++++++++++++++++++ 5 files changed, 731 insertions(+) create mode 100644 crates/buzz-acp/src/stream_flush.rs diff --git a/crates/buzz-acp/src/acp.rs b/crates/buzz-acp/src/acp.rs index c0147baf1b..29fa128899 100644 --- a/crates/buzz-acp/src/acp.rs +++ b/crates/buzz-acp/src/acp.rs @@ -172,6 +172,10 @@ pub struct AcpClient { observer_agent_index: Option, /// Best-effort context attached to raw ACP wire events. observer_context: ObserverContext, + /// Per-turn sink for `agent_message_chunk` text. Set by the turn driver when + /// streaming is enabled; `None` means chunks are logged and dropped, which + /// is the behaviour every deployment had before streaming existed. + stream_sink: Option>, /// Most recently observed `_meta.goose.activeRunId` from a /// `session/update` notification of kind `session_info_update`. /// @@ -493,6 +497,7 @@ impl AcpClient { observer: None, observer_agent_index: None, observer_context: ObserverContext::default(), + stream_sink: None, active_run_id: None, steer_rx: None, goose_usage: UsageTracker::default(), @@ -500,6 +505,15 @@ impl AcpClient { } /// Attach a local observer feed to this ACP client. + /// Attach (or clear) the per-turn streaming sink. Cleared between turns so a + /// stale sender can never receive a later turn's text. + pub(crate) fn set_stream_sink( + &mut self, + sink: Option>, + ) { + self.stream_sink = sink; + } + pub fn set_observer(&mut self, observer: Option, agent_index: usize) { self.observer = observer; self.observer_agent_index = Some(agent_index); @@ -1543,6 +1557,12 @@ impl AcpClient { "agent_message_chunk" => { if let Some(text) = update["content"]["text"].as_str() { tracing::info!(target: "acp::stream", "{text}"); + // Forward to the turn's streamer when one is attached. A + // closed receiver is not an error worth failing a turn over + // — the answer still lands via the normal path. + if let Some(sink) = &self.stream_sink { + let _ = sink.send(text.to_string()); + } } false } diff --git a/crates/buzz-acp/src/config.rs b/crates/buzz-acp/src/config.rs index 9a1b74c276..59c1631d38 100644 --- a/crates/buzz-acp/src/config.rs +++ b/crates/buzz-acp/src/config.rs @@ -261,6 +261,22 @@ pub struct CliArgs { #[arg(long, env = "BUZZ_ACP_MCP_COMMAND", default_value = "")] pub mcp_command: String, + /// Stream the agent's reply into the channel as it is generated: post once, + /// then edit that message as more text arrives. Off by default — it changes + /// what every channel sees. + #[arg(long, env = "BUZZ_ACP_STREAM_REPLIES")] + pub stream_replies: bool, + + /// Minimum milliseconds between streamed edits. The first publish of a turn + /// ignores this, so the channel sees the agent start work immediately. + #[arg(long, env = "BUZZ_ACP_STREAM_THROTTLE_MS", default_value_t = 1_000)] + pub stream_throttle_ms: u64, + + /// Minimum newly-buffered characters before a mid-turn edit is worth a round + /// trip. Guards against an edit per token. + #[arg(long, env = "BUZZ_ACP_STREAM_MIN_DELTA", default_value_t = 24)] + pub stream_min_delta: usize, + /// Idle timeout: max seconds of silence before killing a turn. /// Resets on any agent stdout activity. #[arg(long, env = "BUZZ_ACP_IDLE_TIMEOUT")] @@ -495,6 +511,7 @@ pub struct Config { pub agent_command: String, pub agent_args: Vec, pub mcp_command: String, + pub stream_flush: crate::stream_flush::StreamFlushConfig, pub idle_timeout_secs: u64, pub max_turn_duration_secs: u64, pub agents: u32, @@ -1035,6 +1052,11 @@ impl Config { agent_command, agent_args, mcp_command: args.mcp_command, + stream_flush: crate::stream_flush::StreamFlushConfig { + enabled: args.stream_replies, + throttle_ms: args.stream_throttle_ms, + min_delta_chars: args.stream_min_delta, + }, idle_timeout_secs, max_turn_duration_secs, agents: args.agents, @@ -1413,6 +1435,7 @@ mod tests { agent_command: "goose".into(), agent_args: vec!["acp".into()], mcp_command: "".into(), + stream_flush: crate::stream_flush::StreamFlushConfig::default(), idle_timeout_secs: DEFAULT_IDLE_TIMEOUT_SECS, max_turn_duration_secs: DEFAULT_MAX_TURN_DURATION_SECS, agents: 1, diff --git a/crates/buzz-acp/src/lib.rs b/crates/buzz-acp/src/lib.rs index b11d96d8f7..7214a36e0e 100644 --- a/crates/buzz-acp/src/lib.rs +++ b/crates/buzz-acp/src/lib.rs @@ -10,6 +10,7 @@ mod pool_lifecycle; mod queue; mod relay; mod setup_mode; +pub(crate) mod stream_flush; mod usage; pub use usage::TurnUsage; @@ -1528,6 +1529,7 @@ async fn tokio_main() -> Result<()> { let base_prompt_content = config.base_prompt_content.take(); let ctx = Arc::new(PromptContext { + stream_flush: config.stream_flush, mcp_servers: build_mcp_servers(&config), initial_message: config.initial_message.clone(), idle_timeout: Duration::from_secs(config.idle_timeout_secs), @@ -4963,6 +4965,7 @@ mod build_mcp_servers_tests { agent_command: "goose".into(), agent_args: vec!["acp".into()], mcp_command: "test-mcp-server".into(), + stream_flush: stream_flush::StreamFlushConfig::default(), idle_timeout_secs: config::DEFAULT_IDLE_TIMEOUT_SECS, max_turn_duration_secs: config::DEFAULT_MAX_TURN_DURATION_SECS, agents: 1, @@ -5184,6 +5187,7 @@ mod error_outcome_emission_tests { agent_command: "true".into(), agent_args: vec![], mcp_command: "test-mcp-server".into(), + stream_flush: stream_flush::StreamFlushConfig::default(), idle_timeout_secs: config::DEFAULT_IDLE_TIMEOUT_SECS, max_turn_duration_secs: config::DEFAULT_MAX_TURN_DURATION_SECS, agents: 1, diff --git a/crates/buzz-acp/src/pool.rs b/crates/buzz-acp/src/pool.rs index 0c51fe954f..545e7c0568 100644 --- a/crates/buzz-acp/src/pool.rs +++ b/crates/buzz-acp/src/pool.rs @@ -480,6 +480,8 @@ impl ChannelInfoResolver { } pub struct PromptContext { + /// Streaming policy for replies. Disabled by default. + pub stream_flush: crate::stream_flush::StreamFlushConfig, pub mcp_servers: Vec, pub initial_message: Option, pub idle_timeout: Duration, @@ -1879,6 +1881,41 @@ pub async fn run_prompt_task( None => prompt_sections.iter().map(String::as_str).collect(), }; + // Streamed replies: attach a chunk sink for the duration of the prompt so + // the channel sees the answer being written instead of arriving whole. + // + // `stream_tx` is scope-local ON PURPOSE. Dropping it is the end-of-turn + // signal, so every exit below — success, cancel, idle timeout, agent exit, + // early return — flushes the tail without needing to remember to. Only + // channel turns stream; a heartbeat has no channel to write into. + let stream_handle = match (ctx.stream_flush.enabled, &source) { + (true, PromptSource::Channel(channel_id)) => { + let (tx, rx) = tokio::sync::mpsc::unbounded_channel::(); + let publisher = RelayStreamPublisher { + rest: ctx.rest_client.clone(), + channel_id: *channel_id, + thread_tags: batch + .as_ref() + .and_then(|b| b.events.last()) + .map(|be| crate::queue::parse_thread_tags(&be.event)) + .unwrap_or_default(), + }; + let cfg = ctx.stream_flush; + agent.acp.set_stream_sink(Some(tx.clone())); + let join = tokio::spawn(async move { + crate::stream_flush::run_stream(rx, publisher, cfg, || { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_millis() as u64) + .unwrap_or(0) + }) + .await + }); + Some((tx, join)) + } + _ => None, + }; + // When control_rx is Some (channel tasks), wrap the prompt in select! so // the main loop can cancel, interrupt, or rotate it. Heartbeats // (control_rx=None) take the simple await path — they are not controllable. @@ -2046,6 +2083,19 @@ pub async fn run_prompt_task( } }; + // Turn is over: detach the sink and drop the sender so the streamer sees the + // channel close, performs its final flush, and exits. Awaiting the join + // means the complete reply is on the relay before the outcome is reported — + // otherwise a caller could observe "turn complete" against a message still + // missing its last edit. + if let Some((tx, join)) = stream_handle { + agent.acp.set_stream_sink(None); + drop(tx); + if let Err(e) = join.await { + tracing::debug!("stream task join failed: {e}"); + } + } + match prompt_result { Ok(stop_reason) => { log_stop_reason(&source, &stop_reason); @@ -3706,6 +3756,69 @@ async fn clear_reactions(rest: crate::relay::RestClient, event_ids: Vec) } } +/// Publishes a streamed reply into a channel: one kind:9 post, then kind:40003 +/// edits against it. +/// +/// Edits are deliberately unthreaded — they carry only the `e` tag identifying +/// the message being replaced, so a stream never fans out into the thread. +struct RelayStreamPublisher { + rest: crate::relay::RestClient, + channel_id: Uuid, + thread_tags: ThreadTags, +} + +impl crate::stream_flush::StreamPublisher for RelayStreamPublisher { + async fn post(&mut self, body: String) -> Option { + let thread_ref = self.thread_tags.root_event_id.as_deref().and_then(|root| { + let root_id = nostr::EventId::from_hex(root).ok()?; + let parent_id = self + .thread_tags + .parent_event_id + .as_deref() + .and_then(|p| nostr::EventId::from_hex(p).ok()) + .unwrap_or(root_id); + Some(buzz_sdk::ThreadRef { + root_event_id: root_id, + parent_event_id: parent_id, + }) + }); + let builder = + buzz_sdk::build_message(self.channel_id, &body, thread_ref.as_ref(), &[], false, &[]) + .ok()?; + let event = builder.sign_with_keys(&self.rest.keys).ok()?; + match tokio::time::timeout(Duration::from_secs(5), self.rest.submit_event(&event)).await { + Ok(Ok(_)) => Some(event.id.to_hex()), + Ok(Err(e)) => { + tracing::debug!(channel = %self.channel_id, "stream post failed: {e}"); + None + } + Err(_) => { + tracing::debug!(channel = %self.channel_id, "stream post timed out"); + None + } + } + } + + async fn edit(&mut self, event_id: &str, body: String) { + let Ok(target) = nostr::EventId::from_hex(event_id) else { + return; + }; + let Ok(builder) = buzz_sdk::build_edit(self.channel_id, target, &body) else { + return; + }; + let Ok(event) = builder.sign_with_keys(&self.rest.keys) else { + return; + }; + // Best effort: a dropped edit only costs freshness — the next edit, or + // the end-of-turn flush, carries the full text anyway. + match tokio::time::timeout(Duration::from_secs(5), self.rest.submit_event(&event)).await { + Ok(Ok(_)) => {} + Ok(Err(e)) => tracing::debug!(channel = %self.channel_id, "stream edit failed: {e}"), + Err(_) => tracing::debug!(channel = %self.channel_id, "stream edit timed out"), + } + } +} + #[cfg(test)] mod tests { use super::*; @@ -5335,6 +5448,7 @@ mod tests { ) -> PromptContext { use crate::relay::RestClient; PromptContext { + stream_flush: crate::stream_flush::StreamFlushConfig::default(), mcp_servers: vec![], initial_message: None, idle_timeout: Duration::from_secs(60), diff --git a/crates/buzz-acp/src/stream_flush.rs b/crates/buzz-acp/src/stream_flush.rs new file mode 100644 index 0000000000..de5c06ddbb --- /dev/null +++ b/crates/buzz-acp/src/stream_flush.rs @@ -0,0 +1,570 @@ +//! Coalescing policy for streaming an agent's reply into a channel. +//! +//! ACP delivers an answer as many `agent_message_chunk` updates. Today those +//! are logged and dropped, so a channel sees nothing until the agent posts a +//! finished message — the reply appears to arrive all at once, long after the +//! agent started talking. +//! +//! Relaying every chunk is not the fix: a token-rate edit stream would spend +//! the connection's event budget on typography. This module decides *when* an +//! accumulated buffer is worth publishing. It is pure — no clock, no I/O; the +//! caller supplies `now_ms` and performs the post/edit. +//! +//! Three properties the caller depends on: +//! +//! - **The first publish is a post, every later one is an edit.** Edits target +//! the posted event id, so a stream is one message that grows, not a wall of +//! fragments. Edits are excluded from unread triggers upstream, so this does +//! not turn one answer into hundreds of notifications. +//! - **A finished turn always flushes**, even if the throttle would say wait — +//! otherwise the last few tokens of every answer are silently dropped. +//! - **Nothing is published for an empty turn.** A turn that produced no text +//! must not leave an empty message behind. + +/// Hard ceiling on a single message body. `build_edit` rejects content over +/// 64 KiB, and a rejected edit would strand the message mid-answer, so the +/// policy truncates before the builder can refuse. +pub const MAX_STREAM_BYTES: usize = 64 * 1024; + +/// Appended when the buffer is truncated, so a reader can tell a clipped +/// answer from a complete one rather than silently seeing less. +pub const TRUNCATION_MARKER: &str = "\n\n… (streamed reply truncated)"; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct StreamFlushConfig { + /// Minimum gap between publishes. Zero disables throttling. + pub throttle_ms: u64, + /// Minimum newly-buffered characters before a mid-turn publish is worth it. + pub min_delta_chars: usize, + /// When false the policy never publishes — the pre-existing behaviour. + pub enabled: bool, +} + +impl Default for StreamFlushConfig { + fn default() -> Self { + // Off by default: streaming changes what every channel sees, so it is + // opted into rather than inherited by existing deployments. + Self { + throttle_ms: 1_000, + min_delta_chars: 24, + enabled: false, + } + } +} + +/// Publication state for one turn. `posted` flips once the first message exists. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct StreamFlushState { + /// Everything received this turn. + pub buffer: String, + /// Whether a message has been posted for this turn yet. + pub posted: bool, + /// Byte length of the buffer at the last publish. + pub published_len: usize, + /// `now_ms` of the last publish. + pub last_publish_ms: u64, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum StreamFlushDecision { + /// Accumulate; nothing to publish yet. + Buffer, + /// Create the turn's message with this body. + Post { body: String }, + /// Replace the turn's message body. + Edit { body: String }, +} + +/// Whether this call is a mid-turn chunk or the end of the turn. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum StreamEvent<'a> { + Chunk(&'a str), + /// The turn ended — flush whatever is buffered. + TurnEnd, +} + +fn clamp_body(buffer: &str) -> String { + if buffer.len() <= MAX_STREAM_BYTES { + return buffer.to_string(); + } + // Truncate on a char boundary: slicing mid-codepoint would panic. + let budget = MAX_STREAM_BYTES - TRUNCATION_MARKER.len(); + let mut end = budget.min(buffer.len()); + while end > 0 && !buffer.is_char_boundary(end) { + end -= 1; + } + format!("{}{}", &buffer[..end], TRUNCATION_MARKER) +} + +/// Fold one stream event into the state and decide whether to publish. +/// +/// The state is advanced in place so the caller only records the outcome of a +/// publish it actually performed — a failed post must not leave the policy +/// believing a message exists to edit. +pub fn decide_stream_flush( + state: &mut StreamFlushState, + event: StreamEvent<'_>, + now_ms: u64, + config: StreamFlushConfig, +) -> StreamFlushDecision { + if let StreamEvent::Chunk(text) = event { + state.buffer.push_str(text); + } + + if !config.enabled { + return StreamFlushDecision::Buffer; + } + + let trimmed_is_empty = state.buffer.trim().is_empty(); + let is_end = matches!(event, StreamEvent::TurnEnd); + + // An empty turn publishes nothing — and if it somehow posted earlier, an + // end-of-turn edit to empty would blank a real message, so hold. + if trimmed_is_empty { + return StreamFlushDecision::Buffer; + } + + let delta = state.buffer.len().saturating_sub(state.published_len); + if !is_end { + if delta < config.min_delta_chars { + return StreamFlushDecision::Buffer; + } + let elapsed = now_ms.saturating_sub(state.last_publish_ms); + // `posted == false` is the first publish: never delay the first sign of + // life behind the throttle, or the channel stays silent exactly when a + // reader most needs to know the agent is working. + if state.posted && elapsed < config.throttle_ms { + return StreamFlushDecision::Buffer; + } + } else if delta == 0 && state.posted { + // Nothing new since the last publish and the message already reflects + // it — no redundant final edit. + return StreamFlushDecision::Buffer; + } + + let body = clamp_body(&state.buffer); + state.published_len = state.buffer.len(); + state.last_publish_ms = now_ms; + + if state.posted { + StreamFlushDecision::Edit { body } + } else { + state.posted = true; + StreamFlushDecision::Post { body } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn on() -> StreamFlushConfig { + StreamFlushConfig { + throttle_ms: 1_000, + min_delta_chars: 10, + enabled: true, + } + } + + #[test] + fn disabled_never_publishes_but_still_accumulates() { + let cfg = StreamFlushConfig { + enabled: false, + ..on() + }; + let mut st = StreamFlushState::default(); + assert_eq!( + decide_stream_flush(&mut st, StreamEvent::Chunk("hello world"), 0, cfg), + StreamFlushDecision::Buffer + ); + assert_eq!( + decide_stream_flush(&mut st, StreamEvent::TurnEnd, 5_000, cfg), + StreamFlushDecision::Buffer + ); + assert_eq!(st.buffer, "hello world"); + assert!(!st.posted); + } + + #[test] + fn first_publish_is_a_post_and_ignores_the_throttle() { + let mut st = StreamFlushState::default(); + // now_ms = 0 and last_publish_ms = 0: a throttle check would block this. + let d = decide_stream_flush(&mut st, StreamEvent::Chunk("a first sentence"), 0, on()); + assert!(matches!(d, StreamFlushDecision::Post { .. })); + assert!(st.posted); + } + + #[test] + fn later_publishes_are_edits() { + let mut st = StreamFlushState::default(); + decide_stream_flush(&mut st, StreamEvent::Chunk("a first sentence"), 0, on()); + let d = decide_stream_flush( + &mut st, + StreamEvent::Chunk(" and a second one"), + 2_000, + on(), + ); + assert!(matches!(d, StreamFlushDecision::Edit { .. })); + } + + #[test] + fn small_deltas_buffer_until_they_are_worth_a_round_trip() { + let mut st = StreamFlushState::default(); + decide_stream_flush(&mut st, StreamEvent::Chunk("a first sentence"), 0, on()); + assert_eq!( + decide_stream_flush(&mut st, StreamEvent::Chunk("hi"), 9_000, on()), + StreamFlushDecision::Buffer + ); + } + + #[test] + fn throttle_holds_edits_that_are_too_close_together() { + let mut st = StreamFlushState::default(); + decide_stream_flush(&mut st, StreamEvent::Chunk("a first sentence"), 0, on()); + // Big enough delta, but only 100ms later. + assert_eq!( + decide_stream_flush( + &mut st, + StreamEvent::Chunk("plenty of new text here"), + 100, + on() + ), + StreamFlushDecision::Buffer + ); + } + + #[test] + fn turn_end_flushes_through_the_throttle() { + let mut st = StreamFlushState::default(); + decide_stream_flush(&mut st, StreamEvent::Chunk("a first sentence"), 0, on()); + decide_stream_flush(&mut st, StreamEvent::Chunk(" tail"), 10, on()); // buffered + let d = decide_stream_flush(&mut st, StreamEvent::TurnEnd, 20, on()); + match d { + StreamFlushDecision::Edit { body } => assert!(body.ends_with(" tail")), + other => panic!("expected a final edit, got {other:?}"), + } + } + + // The whole point of streaming is that the reader sees the finished answer. + // Dropping the tail because the throttle was mid-window would be worse than + // not streaming at all. + #[test] + fn no_tokens_are_lost_between_the_last_edit_and_turn_end() { + let mut st = StreamFlushState::default(); + decide_stream_flush(&mut st, StreamEvent::Chunk("opening statement"), 0, on()); + decide_stream_flush( + &mut st, + StreamEvent::Chunk(" middle part here"), + 5_000, + on(), + ); + decide_stream_flush(&mut st, StreamEvent::Chunk(" final words"), 5_050, on()); + let d = decide_stream_flush(&mut st, StreamEvent::TurnEnd, 5_060, on()); + match d { + StreamFlushDecision::Edit { body } => { + assert_eq!(body, "opening statement middle part here final words") + } + other => panic!("expected a final edit, got {other:?}"), + } + } + + #[test] + fn an_empty_turn_publishes_nothing() { + let mut st = StreamFlushState::default(); + assert_eq!( + decide_stream_flush(&mut st, StreamEvent::Chunk(" \n"), 0, on()), + StreamFlushDecision::Buffer + ); + assert_eq!( + decide_stream_flush(&mut st, StreamEvent::TurnEnd, 1_000, on()), + StreamFlushDecision::Buffer + ); + assert!(!st.posted); + } + + #[test] + fn turn_end_does_not_re_edit_when_nothing_changed() { + let mut st = StreamFlushState::default(); + decide_stream_flush(&mut st, StreamEvent::Chunk("a complete answer"), 0, on()); + assert_eq!( + decide_stream_flush(&mut st, StreamEvent::TurnEnd, 9_000, on()), + StreamFlushDecision::Buffer + ); + } + + #[test] + fn a_turn_that_only_ends_still_posts_what_it_buffered() { + // Chunks arrived while disabled-by-throttle; the turn ends without any + // publish having happened yet. It must still post once. + let mut st = StreamFlushState::default(); + st.buffer.push_str("buffered without publishing"); + let d = decide_stream_flush(&mut st, StreamEvent::TurnEnd, 1, on()); + assert!(matches!(d, StreamFlushDecision::Post { .. })); + } + + #[test] + fn oversize_buffers_are_truncated_below_the_builder_limit() { + let mut st = StreamFlushState::default(); + let huge = "x".repeat(MAX_STREAM_BYTES * 2); + let d = decide_stream_flush(&mut st, StreamEvent::Chunk(&huge), 0, on()); + match d { + StreamFlushDecision::Post { body } => { + assert!( + body.len() <= MAX_STREAM_BYTES, + "body {} exceeds cap", + body.len() + ); + assert!(body.ends_with(TRUNCATION_MARKER)); + } + other => panic!("expected a post, got {other:?}"), + } + } + + // Truncation slices bytes; a multi-byte codepoint straddling the cap would + // panic on a naive slice. + #[test] + fn truncation_never_splits_a_multibyte_character() { + let mut st = StreamFlushState::default(); + let huge = "é".repeat(MAX_STREAM_BYTES); // 2 bytes each + let d = decide_stream_flush(&mut st, StreamEvent::Chunk(&huge), 0, on()); + match d { + StreamFlushDecision::Post { body } => assert!(body.len() <= MAX_STREAM_BYTES), + other => panic!("expected a post, got {other:?}"), + } + } + + #[test] + fn zero_throttle_publishes_on_every_sufficient_delta() { + let cfg = StreamFlushConfig { + throttle_ms: 0, + min_delta_chars: 1, + enabled: true, + }; + let mut st = StreamFlushState::default(); + assert!(matches!( + decide_stream_flush(&mut st, StreamEvent::Chunk("one"), 0, cfg), + StreamFlushDecision::Post { .. } + )); + assert!(matches!( + decide_stream_flush(&mut st, StreamEvent::Chunk("two"), 0, cfg), + StreamFlushDecision::Edit { .. } + )); + } +} + +/// Drive one turn's streamed reply. +/// +/// Owns the buffer and publishes through `sink`, which is the seam a test +/// injects a fake at — the policy above is pure, but "post once then edit that +/// same event" is a stateful contract worth exercising end to end. +pub trait StreamPublisher { + /// Create the turn's message; returns the event id later edits target. + fn post(&mut self, body: String) -> impl std::future::Future> + Send; + /// Replace the body of a previously posted message. + fn edit( + &mut self, + event_id: &str, + body: String, + ) -> impl std::future::Future + Send; +} + +/// Consume chunks until the sender is dropped, publishing per the policy, then +/// flush whatever remains. +/// +/// Dropping the sender IS the end-of-turn signal: every turn exit path in the +/// driver — success, idle timeout, agent exit, cancellation — drops it, so no +/// path can silently skip the final flush and strand a half-written answer. +pub async fn run_stream( + mut rx: tokio::sync::mpsc::UnboundedReceiver, + mut publisher: P, + config: StreamFlushConfig, + now_ms: impl Fn() -> u64 + Send, +) -> StreamFlushState { + let mut state = StreamFlushState::default(); + let mut event_id: Option = None; + + while let Some(chunk) = rx.recv().await { + let decision = + decide_stream_flush(&mut state, StreamEvent::Chunk(&chunk), now_ms(), config); + apply(&mut publisher, &mut state, &mut event_id, decision).await; + } + + let decision = decide_stream_flush(&mut state, StreamEvent::TurnEnd, now_ms(), config); + apply(&mut publisher, &mut state, &mut event_id, decision).await; + state +} + +async fn apply( + publisher: &mut P, + state: &mut StreamFlushState, + event_id: &mut Option, + decision: StreamFlushDecision, +) { + match decision { + StreamFlushDecision::Buffer => {} + StreamFlushDecision::Post { body } => match publisher.post(body).await { + Some(id) => *event_id = Some(id), + None => { + // The post failed, so there is nothing to edit. Roll the state + // back so the next decision retries a post rather than issuing + // edits against an event that was never created. + state.posted = false; + state.published_len = 0; + } + }, + StreamFlushDecision::Edit { body } => { + if let Some(id) = event_id.as_deref() { + publisher.edit(id, body).await; + } + } + } +} + +#[cfg(test)] +mod run_stream_tests { + use super::*; + use std::sync::{Arc, Mutex}; + + #[derive(Default, Clone)] + struct FakePublisher { + posts: Arc>>, + edits: Arc>>, + fail_post: bool, + } + + impl StreamPublisher for FakePublisher { + async fn post(&mut self, body: String) -> Option { + if self.fail_post { + return None; + } + self.posts.lock().unwrap().push(body); + Some("evt-1".to_string()) + } + async fn edit(&mut self, event_id: &str, body: String) { + self.edits + .lock() + .unwrap() + .push((event_id.to_string(), body)); + } + } + + fn cfg() -> StreamFlushConfig { + StreamFlushConfig { + throttle_ms: 0, + min_delta_chars: 1, + enabled: true, + } + } + + #[tokio::test] + async fn posts_once_then_edits_the_same_event() { + let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); + let pubr = FakePublisher::default(); + let (posts, edits) = (pubr.posts.clone(), pubr.edits.clone()); + + tx.send("first".to_string()).unwrap(); + tx.send(" second".to_string()).unwrap(); + drop(tx); // end of turn + + run_stream(rx, pubr, cfg(), || 0).await; + + assert_eq!(posts.lock().unwrap().len(), 1, "exactly one post per turn"); + let edits = edits.lock().unwrap(); + assert!(!edits.is_empty(), "later chunks must edit"); + assert!( + edits.iter().all(|(id, _)| id == "evt-1"), + "every edit targets the posted event" + ); + assert_eq!( + edits.last().unwrap().1, + "first second", + "final edit carries the whole reply" + ); + } + + // Dropping the sender is the only end-of-turn signal, so a turn that ends + // before the throttle would have fired must still publish. + #[tokio::test] + async fn a_dropped_sender_flushes_the_tail() { + let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); + let pubr = FakePublisher::default(); + let posts = pubr.posts.clone(); + + tx.send("only chunk".to_string()).unwrap(); + drop(tx); + + let state = run_stream( + rx, + pubr, + StreamFlushConfig { + throttle_ms: 60_000, + ..cfg() + }, + || 0, + ) + .await; + assert_eq!( + posts.lock().unwrap().as_slice(), + &["only chunk".to_string()] + ); + assert!(state.posted); + } + + #[tokio::test] + async fn an_empty_turn_publishes_nothing() { + let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); + let pubr = FakePublisher::default(); + let (posts, edits) = (pubr.posts.clone(), pubr.edits.clone()); + drop(tx); + + run_stream(rx, pubr, cfg(), || 0).await; + assert!(posts.lock().unwrap().is_empty()); + assert!(edits.lock().unwrap().is_empty()); + } + + // A failed post must not leave the policy believing a message exists — the + // edits that followed would target nothing and the reply would vanish. + #[tokio::test] + async fn a_failed_post_never_produces_orphan_edits() { + let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); + let pubr = FakePublisher { + fail_post: true, + ..Default::default() + }; + let edits = pubr.edits.clone(); + + tx.send("first".to_string()).unwrap(); + tx.send(" second".to_string()).unwrap(); + drop(tx); + + run_stream(rx, pubr, cfg(), || 0).await; + assert!( + edits.lock().unwrap().is_empty(), + "no edits without a posted event" + ); + } + + #[tokio::test] + async fn disabled_streaming_publishes_nothing_at_all() { + let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); + let pubr = FakePublisher::default(); + let (posts, edits) = (pubr.posts.clone(), pubr.edits.clone()); + + tx.send("text that would otherwise publish".to_string()) + .unwrap(); + drop(tx); + + run_stream( + rx, + pubr, + StreamFlushConfig { + enabled: false, + ..cfg() + }, + || 0, + ) + .await; + assert!(posts.lock().unwrap().is_empty()); + assert!(edits.lock().unwrap().is_empty()); + } +}