From d92b4d512ba9d2fee0ddb20727b5171aa6fba64f Mon Sep 17 00:00:00 2001 From: salimlaimeche Date: Sun, 30 Aug 2026 21:58:38 +0200 Subject: [PATCH 1/2] fix(stream): bound and coalesce live tool output previews --- crates/aionui-ai-agent/src/session_agent.rs | 102 +++- .../src/session_agent/tool_output.rs | 174 ++++++ .../tests/tool_output_preview.rs | 560 ++++++++++++++++++ 3 files changed, 809 insertions(+), 27 deletions(-) create mode 100644 crates/aionui-ai-agent/src/session_agent/tool_output.rs create mode 100644 crates/aionui-conversation/tests/tool_output_preview.rs diff --git a/crates/aionui-ai-agent/src/session_agent.rs b/crates/aionui-ai-agent/src/session_agent.rs index 9c62c546d..28ec75dde 100644 --- a/crates/aionui-ai-agent/src/session_agent.rs +++ b/crates/aionui-ai-agent/src/session_agent.rs @@ -39,6 +39,9 @@ use aionui_realtime::EventBroadcaster; const EVENT_CHANNEL_CAPACITY: usize = 512; +mod tool_output; +use tool_output::{PREVIEW_PERIOD, PREVIEW_QUEUE_HIGH_WATER, ToolOutputPreviews}; + // Option ids for the generic tool-approval card. `confirm()` maps the incoming // `data` string against these to pick the PermissionDecision; anything else is // treated as an AskUserQuestion answer label (Approved + `selected`). @@ -2707,14 +2710,13 @@ fn spawn_event_pump( // removing it from the manager map) drops that Arc → backend `Drop` → reader // abort + `kill_on_drop` → `event_tx` drops → this stream Closes → the loop ends. tokio::spawn(async move { - // Per-tool accumulated live output for codex `ToolOutputDelta` (streamed - // command stdout). The frontend merges `tool_call` frames by call_id with a - // shallow REPLACE of `output` (hooks.ts: `{...existing, ...new}`), so we must - // send the CUMULATIVE text each time, not the delta — otherwise each chunk - // overwrites the last and only the final chunk shows. Keyed by item_id (== - // the ToolCall tool_use_id). The authoritative full output still arrives on - // the completed ToolResult, which harmlessly replaces this live view. - let mut tool_output: std::collections::HashMap = std::collections::HashMap::new(); + // Replace-style live output: coalesce deltas into bounded UTF-8 tails, + // NOT an ever-growing string cloned into the 512-slot broadcast queue on + // every delta (#946). Small outputs remain cumulative. Large previews + // explicitly mark truncation; ToolResult uses the existing full-result + // path, bypassing this buffer. The timer admits at most 16 queued previews + // (1 MiB of output payload); other event types retain their existing policy. + let mut tool_output = ToolOutputPreviews::default(); // In-flight workflow/subagent refs, mirroring `state::background_active` // (any non-terminal roster entry ⇒ in-flight). claude's non-blocking // Workflow turn emits MULTIPLE `result` frames: the LAUNCH result arrives @@ -2812,12 +2814,35 @@ fn spawn_event_pump( // select arm is disabled entirely while no card is open. let mut progress_tick = tokio::time::interval(std::time::Duration::from_secs(1)); progress_tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + let mut preview_tick = tokio::time::interval_at(tokio::time::Instant::now() + PREVIEW_PERIOD, PREVIEW_PERIOD); + preview_tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); loop { let env = tokio::select! { maybe_env = events.next() => match maybe_env { Some(env) => env, None => break, }, + _ = preview_tick.tick(), if tool_output.has_pending() => { + // Reset relative to NOW: even after a stall, consecutive + // previews stay >=100ms apart (no catch-up burst). + preview_tick.reset(); + if runtime.tx.len() < PREVIEW_QUEUE_HIGH_WATER + && let Some((call_id, output)) = tool_output.take_next() + { + let name = tool_name.get(&call_id).cloned().unwrap_or_default(); + let _ = runtime.tx.send(AgentStreamEvent::ToolCall(ToolCallEventData { + call_id, + name, + args: serde_json::Value::Null, + status: ToolCallStatus::Running, + input: None, + output: Some(output), + description: None, + parent_call_id: None, + })); + } + continue; + } _ = progress_tick.tick(), if !workflow_cards.is_empty() => { let now = aionui_common::now_ms(); for card in workflow_cards.values_mut() { @@ -2850,6 +2875,9 @@ fn spawn_event_pump( "session-pump: turn gen advanced; per-turn suppression state reset" ); } + if last_seen_turn_gen > 0 { + tool_output.retain(|id| open_tools.contains_key(id) && is_detached_exec_call(tool_args.get(id))); + } last_seen_turn_gen = env.turn_gen; runtime.set_status(ConversationStatus::Running); terminal_result_seen = false; @@ -2919,20 +2947,14 @@ fn spawn_event_pump( if let SessionEvent::ToolOutputDelta { item_id, text } = &env.event { // Streamed tool stdout is user-visible output — this turn is not blank. saw_visible_output = true; - let acc = tool_output.entry(item_id.clone()).or_default(); - acc.push_str(text); - let _ = runtime.tx.send(AgentStreamEvent::ToolCall(ToolCallEventData { - call_id: item_id.clone(), - // The wire delta carries no name; use the remembered one so this - // live-output frame doesn't overwrite the persisted row's name to "". - name: tool_name.get(item_id).cloned().unwrap_or_default(), - args: serde_json::Value::Null, - status: ToolCallStatus::Running, - input: None, - output: Some(acc.clone()), - description: None, - parent_call_id: None, - })); + if tool_output.append(item_id, text) { + tracing::info!( + conv_id = %conversation_id, + call_id = %item_id, + preview_limit_bytes = tool_output::PREVIEW_BYTES, + "session-pump: live tool preview truncated; final result is unaffected" + ); + } continue; } @@ -3132,13 +3154,14 @@ fn spawn_event_pump( // close every tool call left open as Canceled BEFORE the // Finish (the relay stops forwarding the turn at Finish). for (call_id, name) in open_tools.drain() { + let output = tool_output.remove(&call_id); let _ = runtime.tx.send(AgentStreamEvent::ToolCall(ToolCallEventData { call_id, name, args: serde_json::Value::Null, status: ToolCallStatus::Canceled, input: None, - output: None, + output, description: None, parent_call_id: None, })); @@ -3202,6 +3225,10 @@ fn spawn_event_pump( tool_args.insert(tool_use_id.clone(), input.clone()); } SessionEvent::ToolResult { tool_use_id, .. } => { + // A full terminal result supersedes any pending live preview. + // Remove it BEFORE forwarding the result; no later timer may + // regress the card to Running or overwrite its full output. + tool_output.discard(tool_use_id); open_tools.remove(tool_use_id); } SessionEvent::TurnResult { .. } | SessionEvent::Detached { .. } if !suppress_intermediate_finish => { @@ -3265,13 +3292,14 @@ fn spawn_event_pump( tool = %name, "session-pump: closing tool call left open at turn end as canceled" ); + let output = tool_output.remove(&call_id); let _ = runtime.tx.send(AgentStreamEvent::ToolCall(ToolCallEventData { call_id, name, args: serde_json::Value::Null, status: ToolCallStatus::Canceled, input: None, - output: None, + output, description: None, parent_call_id: None, })); @@ -3308,7 +3336,7 @@ fn spawn_event_pump( // still open past this turn end (detached exec): their terminal // arrives minutes later and `stamp_tool_name` must still find the // name, or the card re-renders nameless. - tool_output.retain(|call_id, _| open_tools.contains_key(call_id)); + tool_output.retain(|call_id| open_tools.contains_key(call_id)); tool_name.retain(|call_id, _| open_tools.contains_key(call_id)); // Reset the per-turn visibility flag for the next turn. saw_visible_output = false; @@ -3421,6 +3449,25 @@ fn spawn_event_pump( } } + // Preserve the latest bounded output even if the stream ends before a + // ToolResult (including delta-only streams). Never flush it as Running. + for (call_id, output) in tool_output.drain() { + let name = open_tools + .remove(&call_id) + .or_else(|| tool_name.remove(&call_id)) + .unwrap_or_default(); + let _ = runtime.tx.send(AgentStreamEvent::ToolCall(ToolCallEventData { + call_id, + name, + args: serde_json::Value::Null, + status: ToolCallStatus::Canceled, + input: None, + output: Some(output), + description: None, + parent_call_id: None, + })); + } + // The event stream ended: the backend (and its process group) is being // torn down. Settle every card still open and push the frames NOW — // the out-of-turn watcher is still subscribed at this instant and drains @@ -9768,8 +9815,9 @@ mod pump_tests { _ => None, }) .collect(); - // Two frames: first the 1st chunk, then the cumulative 1st+2nd (not just "line-2"). - assert_eq!(outputs, vec!["line-1\n".to_string(), "line-1\nline-2\n".to_string()]); + // Adjacent chunks are coalesced; stream-end preserves the latest cumulative + // preview. Timed delivery/backpressure is covered by tool_output_preview.rs. + assert_eq!(outputs, vec!["line-1\nline-2\n".to_string()]); } // ── Defect 1: process-reap on task drop ─────────────────────────────── diff --git a/crates/aionui-ai-agent/src/session_agent/tool_output.rs b/crates/aionui-ai-agent/src/session_agent/tool_output.rs new file mode 100644 index 000000000..e2f1e95ad --- /dev/null +++ b/crates/aionui-ai-agent/src/session_agent/tool_output.rs @@ -0,0 +1,174 @@ +//! Bounded, replace-style live previews. Authoritative ToolResult output never +//! passes through this buffer. Limits apply to previews, not all session memory. + +use std::collections::{HashMap, VecDeque}; +use std::time::Duration; + +pub(super) const PREVIEW_BYTES: usize = 64 * 1024; +pub(super) const PREVIEW_QUEUE_HIGH_WATER: usize = 16; +pub(super) const PREVIEW_PERIOD: Duration = Duration::from_millis(100); +const TRUNCATED: &str = "[Live preview truncated; showing latest output]\n"; + +struct Preview { + tail: String, + truncated: bool, + dirty: bool, +} + +impl Default for Preview { + fn default() -> Self { + Self { + // Fixed capacity avoids briefly allocating a huge delta, or String's + // geometric growth allocating more than the per-tool budget. + tail: String::with_capacity(PREVIEW_BYTES), + truncated: false, + dirty: false, + } + } +} + +impl Preview { + /// Returns true only when this tool first crosses the preview limit. + fn append(&mut self, text: &str) -> bool { + let was_truncated = self.truncated; + self.truncated |= text.len() > PREVIEW_BYTES - self.tail.len(); + let limit = PREVIEW_BYTES - if self.truncated { TRUNCATED.len() } else { 0 }; + if text.len() >= limit { + self.tail.clear(); + self.tail.push_str(suffix(text, limit)); + } else { + let keep = limit - text.len(); + if self.tail.len() > keep { + let start = self.tail.ceil_char_boundary(self.tail.len() - keep); + self.tail.drain(..start); + } + self.tail.push_str(text); + } + !was_truncated && self.truncated + } + + fn snapshot(&self) -> String { + if self.truncated { + let mut output = String::with_capacity(TRUNCATED.len() + self.tail.len()); + output.push_str(TRUNCATED); + output.push_str(&self.tail); + output + } else { + self.tail.clone() + } + } +} + +fn suffix(text: &str, max_bytes: usize) -> &str { + &text[text.ceil_char_boundary(text.len().saturating_sub(max_bytes))..] +} + +#[derive(Default)] +pub(super) struct ToolOutputPreviews { + tools: HashMap, + // Each dirty call appears once. A hot tool re-enters at the back after its + // snapshot is taken, so it cannot starve another tool's pending update. + pending: VecDeque, +} + +impl ToolOutputPreviews { + pub(super) fn append(&mut self, call_id: &str, text: &str) -> bool { + if text.is_empty() { + return false; + } + let preview = self.tools.entry(call_id.to_owned()).or_default(); + if !preview.dirty { + self.pending.push_back(call_id.to_owned()); + preview.dirty = true; + } + preview.append(text) + } + + pub(super) fn has_pending(&self) -> bool { + !self.pending.is_empty() + } + + pub(super) fn take_next(&mut self) -> Option<(String, String)> { + let call_id = self.pending.pop_front()?; + let preview = self.tools.get_mut(&call_id).expect("pending preview exists"); + preview.dirty = false; + Some((call_id, preview.snapshot())) + } + + pub(super) fn remove(&mut self, call_id: &str) -> Option { + self.pending.retain(|id| id != call_id); + self.tools.remove(call_id).map(|p| p.snapshot()) + } + + pub(super) fn discard(&mut self, call_id: &str) { + self.pending.retain(|id| id != call_id); + self.tools.remove(call_id); + } + + pub(super) fn retain(&mut self, mut keep: impl FnMut(&str) -> bool) { + self.tools.retain(|id, _| keep(id)); + self.pending.retain(|id| self.tools.contains_key(id)); + } + + pub(super) fn clear(&mut self) { + self.tools.clear(); + self.pending.clear(); + } + + pub(super) fn drain(&mut self) -> impl Iterator + '_ { + self.pending.clear(); + self.tools.drain().map(|(id, p)| (id, p.snapshot())) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn huge_delta_never_grows_backing_buffer_and_keeps_utf8_tail() { + let mut previews = ToolOutputPreviews::default(); + let text = format!("{}THE-END", "é🙂".repeat(1024 * 1024)); + assert!(previews.append("a", &text)); + assert!(previews.tools["a"].tail.capacity() <= PREVIEW_BYTES); + let (_, output) = previews.take_next().unwrap(); + assert!(output.len() <= PREVIEW_BYTES); + assert!(output.starts_with(TRUNCATED)); + assert!(output.ends_with("THE-END")); + assert!(text.ends_with(output.strip_prefix(TRUNCATED).unwrap())); + assert!(!previews.append("a", "é🙂"), "log truncation only once per tool"); + } + + #[test] + fn exact_boundary_and_fragmented_unicode() { + let mut previews = ToolOutputPreviews::default(); + let text = "x".repeat(PREVIEW_BYTES); + assert!(!previews.append("a", &text)); + assert_eq!(previews.take_next().unwrap().1, text); + assert!(previews.append("a", "🙂")); + for _ in 0..10_000 { + previews.append("a", "é🙂"); + } + let output = previews.take_next().unwrap().1; + assert!(output.starts_with(TRUNCATED)); + assert!(output.len() <= PREVIEW_BYTES); + assert!(output.ends_with("é🙂")); + assert_eq!(previews.tools["a"].tail.capacity(), PREVIEW_BYTES); + } + + #[test] + fn coalescing_is_fair_and_removal_clears_pending_entries() { + let mut previews = ToolOutputPreviews::default(); + previews.append("a", "one"); + previews.append("b", "two"); + previews.append("a", "three"); + assert_eq!(previews.take_next(), Some(("a".into(), "onethree".into()))); + previews.append("a", "four"); + assert_eq!(previews.take_next(), Some(("b".into(), "two".into()))); + previews.append("b", "five"); + previews.retain(|id| id == "b"); + assert_eq!(previews.remove("b"), Some("twofive".into())); + assert!(!previews.has_pending()); + assert!(previews.tools.is_empty()); + } +} diff --git a/crates/aionui-conversation/tests/tool_output_preview.rs b/crates/aionui-conversation/tests/tool_output_preview.rs new file mode 100644 index 000000000..04e419c00 --- /dev/null +++ b/crates/aionui-conversation/tests/tool_output_preview.rs @@ -0,0 +1,560 @@ +//! Regression for #946 through the real SessionAgentTask event pump. +//! Synthetic backend events are not claims about any CLI's wire protocol. + +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use aionui_ai_agent::{ + AgentStreamEvent, IAgentTask, + protocol::events::{ToolCallEventData, ToolCallStatus}, + session_agent::SessionAgentTask, +}; +use aionui_common::AgentType; +use aionui_session::{ + Admission, BackendError, Capabilities, Command, CommandReceipt, SessionBackend, SessionEnvelope, SessionEvent, + ToolResultContent, TurnOutcome, +}; +use futures_util::{StreamExt, stream::BoxStream}; +use tokio::sync::{broadcast, mpsc, oneshot}; + +type Input = (SessionEnvelope, oneshot::Sender<()>); + +struct ControlledBackend(Mutex>>); + +#[async_trait::async_trait] +impl SessionBackend for ControlledBackend { + async fn dispatch(&self, _: Command) -> Result { + Ok(CommandReceipt { + accepted: true, + admission: Admission::NoTurn, + turn_gen: 1, + }) + } + + fn events(&self) -> BoxStream<'static, SessionEnvelope> { + let rx = self.0.lock().unwrap().take().unwrap(); + // Acknowledge on the NEXT poll: the real pump has handled the preceding + // event by then. Tests never need sleeps or assumptions about scheduling. + futures_util::stream::unfold((rx, None::>), |(mut rx, ack)| async move { + if let Some(ack) = ack { + let _ = ack.send(()); + } + let (event, ack) = rx.recv().await?; + Some((event, (rx, Some(ack)))) + }) + .boxed() + } + + fn capabilities(&self) -> Capabilities { + Capabilities::default() + } +} + +struct Session { + task: Arc, + tx: mpsc::UnboundedSender, + rx: broadcast::Receiver, +} + +impl Session { + fn new() -> Self { + let (tx, rx) = mpsc::unbounded_channel(); + let task = SessionAgentTask::new( + AgentType::Acp, + "conv-preview".into(), + "user-preview".into(), + "/synthetic-workspace".into(), + Arc::new(ControlledBackend(Mutex::new(Some(rx)))), + None, + ); + let rx = task.subscribe(); + Self { task, tx, rx } + } + + async fn send(&self, event: SessionEvent) { + self.send_gen(1, event).await; + } + + async fn send_gen(&self, turn_gen: u64, event: SessionEvent) { + let (ack, processed) = oneshot::channel(); + self.tx + .send(( + SessionEnvelope { + session_id: "conv-preview".into(), + turn_gen, + event, + }, + ack, + )) + .unwrap(); + processed.await.unwrap(); + } + + async fn delta(&self, text: &str) { + self.delta_for("call-preview", text).await; + } + + async fn delta_for(&self, id: &str, text: &str) { + self.send(SessionEvent::ToolOutputDelta { + item_id: id.into(), + text: text.into(), + }) + .await; + } + + async fn call(&self, input: serde_json::Value) { + self.send(SessionEvent::ToolCall { + tool_use_id: "call-preview".into(), + name: "synthetic-tool".into(), + subagent: Default::default(), + input, + parent_tool_use_id: Some("parent-tool".into()), + }) + .await; + } + + async fn result(&self, output: &str, is_error: bool) { + self.send(SessionEvent::ToolResult { + tool_use_id: "call-preview".into(), + is_error, + content: vec![ToolResultContent::Text(output.into())], + parent_tool_use_id: Some("parent-tool".into()), + }) + .await; + } +} + +fn end(outcome: TurnOutcome, is_error: bool) -> SessionEvent { + SessionEvent::TurnResult { + is_error, + api_error_status: None, + result_text: if is_error { + "synthetic failure".into() + } else { + String::new() + }, + epoch: 0, + outcome, + } +} + +fn drain(rx: &mut broadcast::Receiver) -> Vec { + let mut frames = Vec::new(); + loop { + match rx.try_recv() { + Ok(frame) => frames.push(frame), + Err(broadcast::error::TryRecvError::Empty | broadcast::error::TryRecvError::Closed) => return frames, + Err(err) => panic!("events must not be lost: {err}"), + } + } +} + +fn outputs(frames: &[AgentStreamEvent]) -> Vec<&ToolCallEventData> { + frames + .iter() + .filter_map(|frame| match frame { + AgentStreamEvent::ToolCall(data) if data.output.is_some() => Some(data), + _ => None, + }) + .collect() +} + +async fn tick() { + tokio::time::advance(Duration::from_millis(100)).await; + tokio::task::yield_now().await; +} + +#[tokio::test(start_paused = true)] +async fn stalled_reader_retains_at_most_one_mib_of_previews() { + let mut session = Session::new(); + for _ in 0..1024 { + session.delta(&"x".repeat(256)).await; + tick().await; + } + let mut bytes = 0; + let mut frames = 0; + let mut lagged = 0; + loop { + match session.rx.try_recv() { + Ok(AgentStreamEvent::ToolCall(data)) => { + bytes += data.output.as_deref().map_or(0, str::len); + frames += 1; + } + Ok(_) => {} + Err(broadcast::error::TryRecvError::Lagged(n)) => lagged += n, + Err(broadcast::error::TryRecvError::Empty) => break, + Err(err) => panic!("unexpected closed session: {err}"), + } + } + eprintln!("256 KiB output: retained_preview_bytes={bytes}, frames={frames}, lagged={lagged}"); + assert!(bytes <= 1024 * 1024, "unbounded cumulative copies: {bytes} bytes"); + assert_eq!(lagged, 0, "output previews must not overrun a stalled reader"); + assert!(frames > 0, "must exercise actual preview delivery"); + session.delta("LATEST-AFTER-BACKPRESSURE").await; + tick().await; + let resumed = drain(&mut session.rx); + let resumed = outputs(&resumed); + assert_eq!(resumed.len(), 1); + assert!( + resumed[0] + .output + .as_ref() + .unwrap() + .ends_with("LATEST-AFTER-BACKPRESSURE") + ); + assert!(resumed[0].output.as_ref().unwrap().len() <= 64 * 1024); + drop(session.task); +} + +#[tokio::test(start_paused = true)] +async fn small_outputs_are_cumulative_but_coalesced_and_do_not_catch_up() { + let mut session = Session::new(); + session.delta("first\n").await; + session.delta("second\n").await; + assert!(drain(&mut session.rx).is_empty()); + tick().await; + let frames = drain(&mut session.rx); + assert_eq!(outputs(&frames)[0].output.as_deref(), Some("first\nsecond\n")); + + session.delta("third\n").await; + tokio::time::advance(Duration::from_millis(99)).await; + tokio::task::yield_now().await; + assert!(drain(&mut session.rx).is_empty()); + tokio::time::advance(Duration::from_secs(5)).await; + tokio::task::yield_now().await; + let frames = drain(&mut session.rx); + assert_eq!(outputs(&frames).len(), 1); + assert_eq!(outputs(&frames)[0].output.as_deref(), Some("first\nsecond\nthird\n")); + session.delta("fourth\n").await; + tokio::time::advance(Duration::from_millis(99)).await; + tokio::task::yield_now().await; + assert!(drain(&mut session.rx).is_empty(), "no catch-up burst after stall"); + tokio::time::advance(Duration::from_millis(1)).await; + tokio::task::yield_now().await; + let frames = drain(&mut session.rx); + assert_eq!( + outputs(&frames)[0].output.as_deref(), + Some("first\nsecond\nthird\nfourth\n") + ); +} + +#[tokio::test(start_paused = true)] +async fn hot_tool_does_not_starve_other_pending_tools() { + let mut session = Session::new(); + for id in ["a", "b", "c"] { + session.delta_for(id, id).await; + } + for expected in ["a", "b", "c"] { + tick().await; + let frames = drain(&mut session.rx); + let previews = outputs(&frames); + assert_eq!(previews.len(), 1, "only one preview per session per tick"); + assert_eq!(previews[0].call_id, expected); + for _ in 0..100 { + session.delta_for("a", "hot").await; + } + } +} + +#[tokio::test(start_paused = true)] +async fn huge_single_line_and_unicode_are_bounded_in_real_pump() { + let mut session = Session::new(); + let text = format!("{}END", "é🙂".repeat(1024 * 1024)); + session.delta(&text).await; + tick().await; + let frames = drain(&mut session.rx); + let previews = outputs(&frames); + assert_eq!(previews.len(), 1); + let output = previews[0].output.as_ref().unwrap(); + assert!(output.len() <= 64 * 1024); + assert!(output.starts_with("[Live preview truncated;")); + assert!(output.ends_with("END")); + assert!(text.ends_with(output.split_once('\n').unwrap().1)); +} + +#[tokio::test(start_paused = true)] +async fn final_result_bypasses_preview_pressure_and_is_never_overwritten() { + for is_error in [false, true] { + let mut session = Session::new(); + session.call(serde_json::json!({"synthetic": true})).await; + for _ in 0..100 { + session.delta(&"x".repeat(8192)).await; + tick().await; + } + let full = format!("{}FINAL", "f".repeat(256 * 1024)); + session.result(&full, is_error).await; + session.send(end(TurnOutcome::EndTurn, false)).await; + let frames = drain(&mut session.rx); + let result = outputs(&frames).last().copied().unwrap(); + assert_eq!(result.output.as_deref(), Some(full.as_str())); + assert_eq!( + result.status, + if is_error { + ToolCallStatus::Error + } else { + ToolCallStatus::Completed + } + ); + assert_eq!(result.name, "synthetic-tool"); + assert_eq!(result.parent_call_id.as_deref(), Some("parent-tool")); + assert!(matches!(frames.last(), Some(AgentStreamEvent::Finish(_)))); + tick().await; + assert!(drain(&mut session.rx).is_empty(), "pending preview removed at result"); + } +} + +#[tokio::test(start_paused = true)] +async fn immediate_result_has_no_obsolete_running_preview() { + let mut session = Session::new(); + session.delta("intermediate").await; + session.result("authoritative", false).await; + tick().await; + let frames = drain(&mut session.rx); + let results = outputs(&frames); + assert_eq!(results.len(), 1); + assert_eq!(results[0].status, ToolCallStatus::Completed); + assert_eq!(results[0].output.as_deref(), Some("authoritative")); +} + +#[tokio::test(start_paused = true)] +async fn cancellation_and_error_preserve_last_preview_before_turn_terminal() { + for (outcome, is_error) in [ + ( + TurnOutcome::Cancelled { + reason: aionui_session::CancelReason::UserCancel, + }, + false, + ), + (TurnOutcome::EndTurn, true), + ] { + let mut session = Session::new(); + session.call(serde_json::Value::Null).await; + session.delta(&format!("{}LATEST", "x".repeat(1024 * 1024))).await; + session.send(end(outcome, is_error)).await; + let frames = drain(&mut session.rx); + let tool = outputs(&frames)[0]; + assert_eq!(tool.status, ToolCallStatus::Canceled); + assert!(tool.output.as_ref().unwrap().ends_with("LATEST")); + assert!(tool.output.as_ref().unwrap().len() <= 64 * 1024); + assert!(matches!( + frames.last(), + Some(AgentStreamEvent::Finish(_) | AgentStreamEvent::Error(_)) + )); + tick().await; + assert!(drain(&mut session.rx).is_empty()); + } +} + +#[tokio::test(start_paused = true)] +async fn stream_end_preserves_unsent_preview_as_canceled() { + let session = Session::new(); + session.call(serde_json::Value::Null).await; + session.delta("last output before stream closed").await; + let Session { task, tx, mut rx } = session; + drop(tx); + tokio::task::yield_now().await; + let frames = drain(&mut rx); + let tool = outputs(&frames)[0]; + assert_eq!(tool.status, ToolCallStatus::Canceled); + assert_eq!(tool.output.as_deref(), Some("last output before stream closed")); + tick().await; + assert!(drain(&mut rx).is_empty()); + drop(task); +} + +#[tokio::test(start_paused = true)] +async fn detached_output_survives_clean_end_and_next_turn_until_its_result() { + let mut session = Session::new(); + session.call(serde_json::json!({"source": "unifiedExecStartup"})).await; + session.delta("before\n").await; + session.send(end(TurnOutcome::EndTurn, false)).await; + let frames = drain(&mut session.rx); + assert!(outputs(&frames).is_empty(), "clean end must not cancel detached tool"); + session.send_gen(2, SessionEvent::TurnStarted { epoch: 0 }).await; + session + .send_gen( + 2, + SessionEvent::ToolOutputDelta { + item_id: "call-preview".into(), + text: "after\n".into(), + }, + ) + .await; + tick().await; + let frames = drain(&mut session.rx); + assert_eq!(outputs(&frames)[0].output.as_deref(), Some("before\nafter\n")); + session.result("FULL-DETACHED-RESULT", false).await; + tick().await; + let frames = drain(&mut session.rx); + assert_eq!(outputs(&frames).len(), 1); + assert_eq!(outputs(&frames)[0].status, ToolCallStatus::Completed); +} + +#[tokio::test(start_paused = true)] +async fn new_generation_drops_non_detached_pending_preview() { + let mut session = Session::new(); + session.delta("old turn").await; + session.send_gen(2, SessionEvent::TurnStarted { epoch: 0 }).await; + tick().await; + assert!(outputs(&drain(&mut session.rx)).is_empty()); + session + .send_gen( + 2, + SessionEvent::ToolOutputDelta { + item_id: "call-preview".into(), + text: "new turn".into(), + }, + ) + .await; + tick().await; + let frames = drain(&mut session.rx); + assert_eq!(outputs(&frames)[0].output.as_deref(), Some("new turn")); +} + +#[tokio::test(start_paused = true)] +async fn eight_sessions_five_mib_each_with_slow_readers() { + let mut sessions: Vec<_> = (0..8).map(|_| Session::new()).collect(); + let chunk = "s".repeat(5120); + let mut delivered_bytes = 0; + let mut peak_queued_bytes = 0; + for i in 0..1024 { + for session in &sessions { + session.delta(&chunk).await; + } + tick().await; + // Each reader runs only every two seconds (20 producer ticks). + if i % 20 == 19 || i == 1023 { + let mut queued_bytes = 0; + for session in &mut sessions { + let frames = drain(&mut session.rx); + let previews = outputs(&frames); + let bytes: usize = previews.iter().map(|d| d.output.as_ref().unwrap().len()).sum(); + assert!(bytes <= 1024 * 1024); + assert!(previews.iter().all(|d| d.output.as_ref().unwrap().len() <= 64 * 1024)); + queued_bytes += bytes; + } + peak_queued_bytes = peak_queued_bytes.max(queued_bytes); + delivered_bytes += queued_bytes; + } + } + assert!(delivered_bytes > 0); + assert!(peak_queued_bytes <= 8 * 1024 * 1024); + eprintln!( + "8 sessions x 5 MiB: peak_retained_preview_bytes={peak_queued_bytes}, delivered_preview_bytes={delivered_bytes}" + ); +} + +#[tokio::test] +async fn full_final_output_reaches_real_relay_websocket_event_and_sqlite() { + use aionui_common::now_ms; + use aionui_conversation::stream_relay::StreamRelay; + use aionui_db::{ + IConversationRepository, MessagePageDirection, MessagePageParams, SqliteConversationRepository, + init_database_memory, models::ConversationRow, + }; + use aionui_realtime::BroadcastEventBus; + + let db = init_database_memory().await.unwrap(); + let repo = Arc::new(SqliteConversationRepository::new(db.pool().clone())); + let now = now_ms(); + repo.create(&ConversationRow { + id: "conv-preview".into(), + user_id: "system_default_user".into(), + name: "Synthetic preview regression".into(), + r#type: "aionrs".into(), + extra: "{}".into(), + model: None, + status: Some("running".into()), + source: Some("aionui".into()), + channel_chat_id: None, + pinned: false, + pinned_at: None, + created_at: now, + updated_at: now, + project_id: None, + folder_id: None, + name_source: None, + }) + .await + .unwrap(); + let bus = Arc::new(BroadcastEventBus::new(64)); + let mut ws_rx = bus.subscribe(); + let mut session = Session::new(); + // Subscribe before any event; hold this reader while output is produced. + let relay_rx = session.task.subscribe(); + session.call(serde_json::json!({"synthetic": true})).await; + let full = format!("{}FINAL", "é🙂".repeat(100_000)); + session.delta(&full).await; + // Here use real time: SQLite workers need to run independently of Tokio's + // paused clock. Wait for an actual preview before adding the final result. + loop { + let frame = tokio::time::timeout(Duration::from_secs(5), session.rx.recv()) + .await + .unwrap() + .unwrap(); + if let AgentStreamEvent::ToolCall(data) = frame + && data.output.is_some() + { + assert!(data.output.as_ref().unwrap().len() <= 64 * 1024); + break; + } + } + session.result(&full, false).await; + session.send(end(TurnOutcome::EndTurn, false)).await; + let relay = StreamRelay::new( + "conv-preview".into(), + "assistant-preview".into(), + "turn-preview".into(), + "system_default_user".into(), + repo.clone(), + bus, + ); + tokio::time::timeout(Duration::from_secs(10), relay.consume(relay_rx)) + .await + .unwrap(); + + let messages = repo + .list_messages_page( + "system_default_user", + "conv-preview", + &MessagePageParams { + limit: 100, + direction: MessagePageDirection::InitialLatest, + }, + ) + .await + .unwrap(); + let row = messages.items.iter().find(|row| row.id == "call-preview").unwrap(); + assert_eq!(row.status.as_deref(), Some("finish")); + let content: serde_json::Value = serde_json::from_str(&row.content).unwrap(); + assert_eq!(content["output"].as_str(), Some(full.as_str())); + assert_eq!(content["status"], "completed"); + assert_eq!(content["name"], "synthetic-tool"); + assert_eq!(content["args"]["synthetic"], true); + assert_eq!(content["parent_call_id"], "parent-tool"); + + let mut saw_preview = false; + let mut saw_final = false; + while let Ok(frame) = ws_rx.try_recv() { + // Existing response envelope: body carries the serialized ToolCall. + let data = &frame.data["data"]; + if frame.data["type"] != "tool_call" { + continue; + } + if data["status"] == "running" + && let Some(text) = data["output"].as_str() + { + assert!(text.len() <= 64 * 1024); + assert!(!saw_final, "no preview after authoritative final"); + saw_preview = true; + } + if data["status"] == "completed" { + assert_eq!(data["output"].as_str(), Some(full.as_str())); + saw_final = true; + } + } + assert!( + saw_preview && saw_final, + "real relay must forward preview AND complete final" + ); +} From 1af08fd3e087a1ef1337f0e17943b46a84234d08 Mon Sep 17 00:00:00 2001 From: salimlaimeche Date: Mon, 31 Aug 2026 01:11:26 +0200 Subject: [PATCH 2/2] test(stream): validate bounded previews with real Codex Verified: crates/aionui-ai-agent/tests/live_tool_output_preview.rs with Codex CLI 0.147.0 and a 1 MiB synthetic command. The old main fails the 1 MiB retained-preview bound; the rebuilt patched pump passes and preserves the final result. The live test remains opt-in and uses a disposable CLI home. --- .../tests/live_tool_output_preview.rs | 326 ++++++++++++++++++ 1 file changed, 326 insertions(+) create mode 100644 crates/aionui-ai-agent/tests/live_tool_output_preview.rs diff --git a/crates/aionui-ai-agent/tests/live_tool_output_preview.rs b/crates/aionui-ai-agent/tests/live_tool_output_preview.rs new file mode 100644 index 000000000..efd8cbc01 --- /dev/null +++ b/crates/aionui-ai-agent/tests/live_tool_output_preview.rs @@ -0,0 +1,326 @@ +//! Opt-in real-agent validation for #946, complementing the deterministic +//! regressions in aionui-conversation/tests/tool_output_preview.rs. +//! +//! Needs a logged-in Codex CLI and python3; spends model tokens. Credentials are +//! copied into a disposable CLI home, never logged. User config, skills, MCP +//! settings and project data are not copied. No production service is contacted. +//! Only aggregate evidence prints; this is not a sandbox/privacy audit of Codex. +//! +//! Run with AIONUI_LIVE_CODEX_BIN and AIONUI_LIVE_CODEX_AUTH_FILE explicitly set: +//! cargo test -p aionui-ai-agent --test live_tool_output_preview -- --ignored --nocapture +//! Optional AIONUI_LIVE_CODEX_MODEL selects a model via the CLI's documented -c flag. +//! Uses the production RealSpawner -> CodexConnection -> SessionAgentTask path; +//! ObservedBackend below only observes, never synthesizes or delays input events. + +use std::collections::{HashMap, HashSet}; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use aionui_ai_agent::{ + AgentStreamEvent, IAgentTask, + protocol::events::{ToolCallEventData, ToolCallStatus}, + session_agent::SessionAgentTask, +}; +use aionui_common::{AgentType, EnvVar}; +use aionui_session::{ + BackendConnection, BackendError, Capabilities, CodexConnection, Command, CommandMeta, CommandReceipt, ContentBlock, + SessionBackend, SessionConfig, SessionEnvelope, SessionEvent, SessionSpec, ToolResultContent, +}; +use futures_util::{StreamExt, stream::BoxStream}; +use sha2::{Digest, Sha256}; +use tokio::sync::broadcast; + +const HALF: usize = 512 * 1024; +const MIB: usize = 1024 * 1024; +const GENERATOR: &str = r#"import pathlib, sys, time +chunk = ('é🙂' + 'x' * 8186).encode('utf-8') +assert len(chunk) == 8192 +for phase in range(2): + for _ in range(64): + sys.stdout.buffer.write(chunk) + sys.stdout.buffer.flush() + time.sleep(0.04) + if phase == 0: + deadline = time.monotonic() + 30 + while not pathlib.Path('continue').exists(): + if time.monotonic() > deadline: + raise RuntimeError('test reader did not release the checkpoint') + time.sleep(0.02) +"#; + +#[derive(Default)] +struct Observed { + delta_bytes: usize, + deltas: usize, + finals: HashMap, bool)>, + terminal: Option, + notices: usize, +} + +struct ObservedBackend { + real: Arc, + observed: Arc>, +} + +#[async_trait::async_trait] +impl SessionBackend for ObservedBackend { + async fn dispatch(&self, command: Command) -> Result { + self.real.dispatch(command).await + } + + fn capabilities(&self) -> Capabilities { + self.real.capabilities() + } + + fn events(&self) -> BoxStream<'static, SessionEnvelope> { + let observed = self.observed.clone(); + self.real + .events() + .inspect(move |envelope| { + let mut seen = observed.lock().unwrap(); + match &envelope.event { + SessionEvent::ToolOutputDelta { text, .. } => { + seen.delta_bytes += text.len(); + seen.deltas += 1; + } + SessionEvent::ToolResult { + tool_use_id, + content, + is_error, + .. + } => { + let texts: Vec<_> = content + .iter() + .filter_map(|c| match c { + ToolResultContent::Text(text) => Some(text.as_str()), + _ => None, + }) + .collect(); + // Match the existing translator's text-block join, not + // the original stdout (the CLI may truncate its final). + let text = texts.join("\n"); + seen.finals.insert( + tool_use_id.clone(), + (text.len(), Sha256::digest(text.as_bytes()).to_vec(), *is_error), + ); + } + SessionEvent::TurnResult { is_error, .. } => seen.terminal = Some(*is_error), + SessionEvent::Notice { .. } => seen.notices += 1, + _ => {} + } + }) + .boxed() + } +} + +fn take_pending(rx: &mut broadcast::Receiver) -> Vec { + let mut events = Vec::new(); + loop { + match rx.try_recv() { + Ok(event) => events.push(event), + Err(broadcast::error::TryRecvError::Empty) => return events, + Err(error) => panic!("live subscriber lost events: {error}"), + } + } +} + +fn inspect_frame(data: &ToolCallEventData, terminal_ids: &mut HashSet) -> usize { + if data.status == ToolCallStatus::Running { + assert!(!terminal_ids.contains(&data.call_id), "preview after final"); + let bytes = data.output.as_deref().map_or(0, str::len); + assert!(bytes <= 64 * 1024, "oversized preview: {bytes}"); + bytes + } else { + terminal_ids.insert(data.call_id.clone()); + 0 + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[ignore = "LIVE: real Codex login required; executes synthetic output and spends tokens"] +async fn real_codex_with_stalled_then_slow_reader() { + let binary = std::env::var_os("AIONUI_LIVE_CODEX_BIN").expect("set AIONUI_LIVE_CODEX_BIN"); + let auth = std::env::var_os("AIONUI_LIVE_CODEX_AUTH_FILE").expect("set AIONUI_LIVE_CODEX_AUTH_FILE"); + let mut version_probe = aionui_runtime::Builder::clean_cli(&binary); + version_probe.arg("--version"); + let version = version_probe.output().await.expect("read Codex version"); + assert!(version.status.success()); + eprintln!("live CLI: {}", String::from_utf8_lossy(&version.stdout).trim()); + let tmp = tempfile::tempdir().unwrap(); + let workspace = tmp.path().join("workspace"); + let cli_home = tmp.path().join("codex-home"); + std::fs::create_dir(&workspace).unwrap(); + std::fs::create_dir(&cli_home).unwrap(); + std::fs::copy(auth, cli_home.join("auth.json")).unwrap(); + std::fs::write(workspace.join("emit.py"), GENERATOR).unwrap(); + // Only this disposable synthetic workspace is trusted. This never edits + // the user's global config or weakens the workspace-write sandbox. + std::fs::write( + cli_home.join("config.toml"), + format!( + "[projects.{}]\ntrust_level = \"trusted\"\n", + serde_json::to_string(&workspace.to_string_lossy()).unwrap() + ), + ) + .unwrap(); + let mut extra_args = Vec::new(); + if let Ok(model) = std::env::var("AIONUI_LIVE_CODEX_MODEL") { + extra_args.extend(["-c".into(), format!("model={}", serde_json::to_string(&model).unwrap())]); + } + let spawner = Arc::new(aionui_process::RealSpawner::new( + Arc::new(aionui_process::FileRegistryStore::new(tmp.path())), + uuid::Uuid::now_v7(), + "live-preview-test", + )); + let real = tokio::time::timeout( + Duration::from_secs(30), + CodexConnection::new(spawner).open_session( + SessionSpec::Fresh { + session_id: "live-preview".into(), + }, + SessionConfig { + cwd: Some(workspace.to_string_lossy().into_owned()), + cli_program: Some(binary.into()), + approval_policy: Some("never".into()), + spawn_env: vec![EnvVar { + name: "CODEX_HOME".into(), + value: cli_home.to_string_lossy().into_owned(), + }], + extra_args, + ..Default::default() + }, + ), + ) + .await + .expect("real Codex handshake timeout") + .expect("real Codex session opens"); + let observed = Arc::new(Mutex::new(Observed::default())); + let task = SessionAgentTask::new( + AgentType::Acp, + "live-preview".into(), + "test-user".into(), + workspace.to_string_lossy().into_owned(), + Arc::new(ObservedBackend { + real: real.clone(), + observed: observed.clone(), + }), + None, + ); + let mut rx = task.subscribe(); + tokio::time::timeout( + Duration::from_secs(30), + real.dispatch(Command::Send { + content: vec![ContentBlock::Text( + "Run exactly python3 emit.py in the current directory. This is an output-streaming test: \ + do not read or modify any files, do not redirect, filter or capture stdout. \ + The test harness releases a checkpoint file after reading the first half of stdout. \ + Allow up to 60 seconds, wait for the command to finish, then reply only DONE." + .into(), + )], + metadata: CommandMeta::default(), + }), + ) + .await + .expect("real dispatch timeout") + .expect("real prompt accepted"); + + // The producer pauses after 512 KiB until we release its checkpoint. The + // production backend/pump keep draining; only the outer subscriber stalls. + tokio::time::timeout(Duration::from_secs(90), async { + loop { + { + let seen = observed.lock().unwrap(); + assert!(seen.terminal.is_none(), "turn ended before the streaming checkpoint"); + if seen.delta_bytes >= HALF { + break; + } + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + }) + .await + .unwrap_or_else(|_| { + let seen = observed.lock().unwrap(); + panic!( + "no checkpoint: delta_bytes={}, deltas={}, notices={}", + seen.delta_bytes, seen.deltas, seen.notices + ) + }); + tokio::time::sleep(Duration::from_millis(200)).await; + let pending = take_pending(&mut rx); + let preview_bytes: usize = pending + .iter() + .filter_map(|event| match event { + AgentStreamEvent::ToolCall(data) if data.status == ToolCallStatus::Running => { + Some(data.output.as_deref().map_or(0, str::len)) + } + _ => None, + }) + .sum(); + eprintln!( + "live checkpoint: retained_preview_bytes={preview_bytes}, queued_events={}", + pending.len() + ); + assert!(preview_bytes > 0, "must exercise actual output preview delivery"); + assert!(preview_bytes <= MIB, "unbounded live preview copies: {preview_bytes}"); + let mut terminal_ids = HashSet::new(); + for event in pending { + if let AgentStreamEvent::ToolCall(data) = event { + inspect_frame(&data, &mut terminal_ids); + } + } + std::fs::write(workspace.join("continue"), b"release").unwrap(); + let mut received_finals = HashMap::new(); + tokio::time::timeout(Duration::from_secs(90), async { + loop { + let event = rx.recv().await.expect("slow reader must not lag"); + if let AgentStreamEvent::ToolCall(data) = &event { + inspect_frame(data, &mut terminal_ids); + if data.status != ToolCallStatus::Running { + assert_eq!( + data.status, + ToolCallStatus::Completed, + "synthetic command must complete" + ); + let text = data.output.as_deref().unwrap_or_default(); + received_finals.insert( + data.call_id.clone(), + (text.len(), Sha256::digest(text.as_bytes()).to_vec()), + ); + } + } + if matches!(event, AgentStreamEvent::Finish(_)) { + break; + } + tokio::time::sleep(Duration::from_millis(200)).await; + } + }) + .await + .expect("real turn finishes with slow subscriber"); + tokio::time::sleep(Duration::from_millis(300)).await; + for event in take_pending(&mut rx) { + if let AgentStreamEvent::ToolCall(data) = event { + inspect_frame(&data, &mut terminal_ids); + } + } + let seen = observed.lock().unwrap(); + assert!(seen.delta_bytes >= MIB, "real CLI must stream the full synthetic MiB"); + assert!(seen.deltas >= 16, "exercise multiple real CLI deltas"); + assert_eq!(seen.terminal, Some(false)); + assert!(!seen.finals.is_empty()); + for (id, (len, hash, error)) in &seen.finals { + assert!(!error, "synthetic command failed"); + assert_eq!( + received_finals.get(id), + Some(&(*len, hash.clone())), + "final changed in pump" + ); + } + eprintln!( + "live final: delta_bytes={}, deltas={}, final_bytes={}, finals={}, terminal_error=false; no late preview", + seen.delta_bytes, + seen.deltas, + seen.finals.values().map(|(len, _, _)| len).sum::(), + seen.finals.len(), + ); +}