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()); + } +}