From b8e91bd712a63dbde3521502c2a4b327a96cb314 Mon Sep 17 00:00:00 2001 From: daniel Date: Sun, 13 Sep 2026 16:37:03 +0100 Subject: [PATCH] feat(acp): stream shell output as v2 terminals --- src/protocols/acp.rs | 14 +- src/protocols/acp/tool_projection.rs | 64 +++++- src/protocols/acp/tool_projection/terminal.rs | 117 ++++++++++ src/protocols/acp/tool_projection/tests.rs | 215 ++++++++++++++++++ src/tools/shell.rs | 23 +- 5 files changed, 410 insertions(+), 23 deletions(-) create mode 100644 src/protocols/acp/tool_projection/terminal.rs diff --git a/src/protocols/acp.rs b/src/protocols/acp.rs index 0119592c..ead61abd 100644 --- a/src/protocols/acp.rs +++ b/src/protocols/acp.rs @@ -1092,11 +1092,15 @@ impl LoopObserver for ResponseInterruptionNoticeObserver { let client = self.client.clone(); let session_id = self.session_id.clone(); tool_projection::Subscription::start(event.session_id.0.clone(), move |update| { - update.v1().is_ok_and(|update| { - client - .notify_session(SessionNotification::new(session_id.clone(), update)) - .is_ok() - }) + update.v2_only() + || update.v1().is_ok_and(|update| { + client + .notify_session(SessionNotification::new( + session_id.clone(), + update, + )) + .is_ok() + }) }) }); } diff --git a/src/protocols/acp/tool_projection.rs b/src/protocols/acp/tool_projection.rs index 998b09c6..0dbcc5e2 100644 --- a/src/protocols/acp/tool_projection.rs +++ b/src/protocols/acp/tool_projection.rs @@ -1,7 +1,8 @@ //! Best-effort projection at the evaluated hidden-tool invocation boundary. -//! No source/input/output retention and no execution waits. The compose call +//! No execution waits or accumulated source/input/output retention. The compose //! budget bounds cards per run; the bus and receiver bound queued/active cards. -use std::{collections::HashSet, path::Path, sync::OnceLock}; +//! Rich-content patches carry bounded payloads only for the lifetime of delivery. +use std::{collections::HashMap, path::Path, sync::OnceLock}; use agentkit_tools_core::ToolRequest; use serde_json::{Map, Value}; @@ -16,6 +17,8 @@ use tokio::sync::broadcast; )] mod tests; +pub(crate) mod terminal; + const CAPACITY: usize = crate::runlet_progress::MAX_NODES; const MAX_ID: usize = 256; const MAX_PATH: usize = 4096; @@ -167,6 +170,12 @@ impl Update { }) } + pub(crate) fn v2_only(&self) -> bool { + self.patch + .as_ref() + .is_some_and(|patch| patch.get("sessionUpdate").is_some()) + } + pub(crate) fn v1(&self) -> Result { if self.start.is_some() { serde_json::from_value(self.value()).map(agentkit_acp::SessionUpdate::ToolCall) @@ -176,8 +185,12 @@ impl Update { } pub(crate) fn v2(&self) -> Result { - serde_json::from_value(self.value()) - .map(agentkit_acp::v2::wire::SessionUpdate::ToolCallUpdate) + if self.v2_only() { + serde_json::from_value(self.value()) + } else { + serde_json::from_value(self.value()) + .map(agentkit_acp::v2::wire::SessionUpdate::ToolCallUpdate) + } } } @@ -229,7 +242,8 @@ async fn forward( mut drains: tokio::sync::mpsc::Receiver, ) { use futures_util::future::{Either, select}; - let mut active = HashSet::new(); + // The bool tracks whether this card has a live v2 terminal. + let mut active = HashMap::new(); let mut drains_open = true; loop { let next = if drains_open { @@ -289,20 +303,24 @@ fn forward_event( event: Result, receiver: &mut broadcast::Receiver, session: &str, - active: &mut HashSet, + active: &mut HashMap, send: &impl Fn(Update) -> bool, ) -> Result<(), agentkit_acp::AcpRuntimeError> { match event { Ok(update) if update.session == session => { if update.start.is_some() { - if active.len() >= CAPACITY || !active.insert(update.call.clone()) { + if active.len() >= CAPACITY || active.contains_key(&update.call) { return Ok(()); } - } else if update.patch.is_some() { - if !active.contains(&update.call) { + active.insert(update.call.clone(), false); + } else if let Some(patch) = &update.patch { + let Some(terminal_running) = active.get_mut(&update.call) else { return Ok(()); + }; + if patch.get("sessionUpdate").and_then(Value::as_str) == Some("terminal_update") { + *terminal_running = patch.get("exitStatus").is_none(); } - } else if !active.remove(&update.call) { + } else if active.remove(&update.call).is_none() { return Ok(()); } if !send(update) { @@ -311,7 +329,31 @@ fn forward_event( } Ok(_) => {} Err(error) => { - for call in active.drain() { + for (call, terminal_running) in active.drain() { + // Loss invalidates the stream as well as its card. Do not leave + // an editor waiting for a terminal exit frame that was dropped. + if terminal_running + && !send(Update { + session: session.into(), + call: call.clone(), + start: None, + patch: Some(Value::Object(Map::from_iter([ + ("sessionUpdate".into(), Value::from("terminal_update")), + ("terminalId".into(), Value::from(call.clone())), + ("exitStatus".into(), Value::Object(Map::new())), + ( + "_meta".into(), + Value::Object(Map::from_iter([( + "kit/outputIncomplete".into(), + Value::from(true), + )])), + ), + ]))), + ok: false, + }) + { + return Err(delivery_error()); + } if !send(Update { session: session.into(), call, diff --git a/src/protocols/acp/tool_projection/terminal.rs b/src/protocols/acp/tool_projection/terminal.rs new file mode 100644 index 00000000..d6688769 --- /dev/null +++ b/src/protocols/acp/tool_projection/terminal.rs @@ -0,0 +1,117 @@ +//! Agent-owned terminals are v2 only. Frames share the bounded invocation bus. +use super::{MAX_ID, Update, bus}; +use agentkit_tools_core::ToolRequest; +use base64::{Engine as _, engine::general_purpose::STANDARD}; +use serde_json::{Map, Value}; +use std::path::Path; + +/// A reader retains only routing IDs, never command text or accumulated output. +#[derive(Clone)] +pub(crate) struct Output { + session: String, + call: String, +} + +impl Output { + fn publish(&self, patch: Value) { + let _ = bus().send(Update { + session: self.session.clone(), + call: self.call.clone(), + start: None, + patch: Some(patch), + ok: false, + }); + } + + pub(crate) fn chunk(&self, bytes: &[u8]) { + // The real shell reader uses this same bounded buffer size. Keep the + // boundary safe for other callers too, including non-UTF-8 output. + for bytes in bytes.chunks(8192) { + self.publish(Value::Object(Map::from_iter([ + ("sessionUpdate".into(), Value::from("terminal_output_chunk")), + ("terminalId".into(), Value::from(self.call.clone())), + ("data".into(), Value::from(STANDARD.encode(bytes))), + ]))); + } + } + + fn exited(&self, code: Option) { + self.publish(Value::Object(Map::from_iter([ + ("sessionUpdate".into(), Value::from("terminal_update")), + ("terminalId".into(), Value::from(self.call.clone())), + ( + "exitStatus".into(), + Value::Object(Map::from_iter([( + "exitCode".into(), + code.and_then(|code| u32::try_from(code).ok()) + .map_or(Value::Null, Value::from), + )])), + ), + ]))); + } +} + +/// Dropped before Observed terminalizes the call, also on errors/cancellation. +pub(crate) struct Terminal(Option); + +impl Terminal { + pub(crate) fn start(request: &ToolRequest, command: &str, cwd: &Path) -> Self { + if bus().receiver_count() == 0 + || request.session_id.0.len() > MAX_ID + || request.call_id.0.len() > MAX_ID + || !request.call_id.0.contains(":compose:") + { + return Self(None); + } + let output = Output { + session: request.session_id.0.clone(), + call: request.call_id.0.clone(), + }; + let mut patch = Value::Object(Map::from_iter([ + ("sessionUpdate".into(), Value::from("terminal_update")), + ("terminalId".into(), Value::from(output.call.clone())), + ])); + // Omit oversized metadata rather than retaining or truncating commands. + if command.len() <= crate::runlet_progress::MAX_SOURCE { + patch["command"] = Value::from(command); + } + if cwd.is_absolute() + && let Some(cwd) = cwd.to_str().filter(|cwd| cwd.len() <= super::MAX_PATH) + { + patch["cwd"] = Value::from(cwd); + } + output.publish(patch); + output.publish(Value::Object(Map::from_iter([ + ("sessionUpdate".into(), Value::from("tool_call_update")), + ("toolCallId".into(), Value::from(output.call.clone())), + ( + "content".into(), + Value::Array(vec![Value::Object(Map::from_iter([ + ("type".into(), Value::from("terminal")), + ("terminalId".into(), Value::from(output.call.clone())), + ]))]), + ), + ]))); + Self(Some(output)) + } + + pub(crate) fn output(&self) -> Option { + self.0.clone() + } + + pub(crate) fn finish(mut self, code: Option) { + if let Some(output) = self.0.take() { + output.exited(code); + } + } +} + +impl Drop for Terminal { + fn drop(&mut self) { + if let Some(output) = self.0.take() { + // An empty exit object marks an exited terminal without inventing + // an exit code for cancellation, timeout, I/O errors or unwind. + output.exited(None); + } + } +} diff --git a/src/protocols/acp/tool_projection/tests.rs b/src/protocols/acp/tool_projection/tests.rs index 8f1f3fc7..d0071a6d 100644 --- a/src/protocols/acp/tool_projection/tests.rs +++ b/src/protocols/acp/tool_projection/tests.rs @@ -236,3 +236,218 @@ async fn lag_invalidates_active_cards_and_recovers_for_fresh_calls() { task.await.unwrap(); assert!(messages.try_recv().is_err()); } + +#[cfg(unix)] +#[tokio::test] +async fn real_shell_streams_bytes_before_exit_and_drains_terminal_before_call() { + use agentkit_tools_core::{AllowAllPermissions, OwnedToolContext, Tool}; + use base64::{Engine as _, engine::general_purpose::STANDARD}; + use std::sync::Arc; + let root = tempfile::tempdir().unwrap(); + let tool = crate::tools::Observed::new(crate::tools::ShellTool::new(root.path().into())); + let context = OwnedToolContext { + session_id: SessionId::new("real-terminal"), + turn_id: TurnId::new("turn"), + metadata: MetadataMap::new(), + permissions: Arc::new(AllowAllPermissions), + resources: Arc::new(()), + cancellation: None, + execution_scope: None, + approved_request: None, + }; + let (send, mut receive) = tokio::sync::mpsc::unbounded_channel(); + let subscription = Subscription::start("real-terminal".into(), move |update| { + send.send(update).is_ok() + }); + let call = request( + "real-terminal", + "shell", + json!({ + "command": "printf '\\377A'; while [ ! -f release ]; do sleep 0.01; done; printf 'err' >&2; exit 7", + "timeout_seconds": 5, + }), + ); + let execution = tokio::spawn(async move { tool.invoke(call, &mut context.borrowed()).await }); + let mut bytes = Vec::new(); + let mut updates = Vec::new(); + // Release the actual process only after receiving a streamed byte, rather + // than asserting a fragile wall-clock latency or accepting buffered output. + while bytes.is_empty() { + let update = tokio::time::timeout(std::time::Duration::from_secs(10), receive.recv()) + .await + .unwrap() + .unwrap(); + let value = serde_json::to_value(update.v2().unwrap()).unwrap(); + if value["sessionUpdate"] == "terminal_output_chunk" { + bytes.extend(STANDARD.decode(value["data"].as_str().unwrap()).unwrap()); + } + updates.push(update); + } + assert!(!execution.is_finished()); + std::fs::write(root.path().join("release"), "").unwrap(); + let result = execution.await.unwrap().unwrap(); + let agentkit_core::ToolOutput::Structured(result) = result.result.output else { + panic!() + }; + assert_eq!(result["exit_code"], 7); + assert_eq!(result["stderr"], "err"); + subscription.drain().await.unwrap(); + while let Ok(update) = receive.try_recv() { + let value = serde_json::to_value(update.v2().unwrap()).unwrap(); + if value["sessionUpdate"] == "terminal_output_chunk" { + bytes.extend(STANDARD.decode(value["data"].as_str().unwrap()).unwrap()); + } + updates.push(update); + } + assert_eq!(bytes, b"\xffAerr"); + let values: Vec<_> = updates + .iter() + .map(|u| serde_json::to_value(u.v2().unwrap()).unwrap()) + .collect(); + assert_eq!(values[1]["sessionUpdate"], "terminal_update"); + assert_eq!(values[2]["content"][0]["type"], "terminal"); + assert_eq!(values[2]["content"][0]["terminalId"], "parent:compose:node"); + assert_eq!(values[values.len() - 2]["exitStatus"]["exitCode"], 7); + assert_eq!(values.last().unwrap()["status"], "completed"); + // v1 observes only the existing invocation lifecycle, no agent-owned terminal. + assert_eq!(updates.iter().filter(|u| !u.v2_only()).count(), 2); +} + +#[test] +fn terminal_drop_marks_exit_and_bounds_binary_chunks_and_metadata() { + let mut receiver = bus().subscribe(); + let request = request("terminal-bounds", "shell", json!({})); + let invocation = Invocation::start(&request, None).unwrap(); + let terminal = terminal::Terminal::start( + &request, + &"x".repeat(crate::runlet_progress::MAX_SOURCE + 1), + Path::new("/tmp"), + ); + terminal.output().unwrap().chunk(&vec![255; 20_000]); + drop(terminal); + invocation.finish(false); + let updates: Vec<_> = std::iter::from_fn(|| receiver.try_recv().ok()) + .filter(|update| update.session == "terminal-bounds") + .collect(); + let values: Vec<_> = updates + .iter() + .map(|u| serde_json::to_value(u.v2().unwrap()).unwrap()) + .collect(); + assert!(values[1].get("command").is_none()); + for value in &values { + if value["sessionUpdate"] == "terminal_output_chunk" { + assert!(value["data"].as_str().unwrap().len() <= 10_924); + } + } + assert!(values[values.len() - 2]["exitStatus"].is_object()); + assert_eq!(values.last().unwrap()["status"], "failed"); +} + +#[tokio::test] +async fn lag_exits_live_terminal_before_invalidating_its_card() { + let (sender, receiver) = broadcast::channel(2); + let (_drain, commands) = tokio::sync::mpsc::channel(1); + let (send, mut receive) = tokio::sync::mpsc::unbounded_channel(); + let task = tokio::spawn(forward( + receiver, + "terminal-lag".into(), + move |u| send.send(u).is_ok(), + commands, + )); + let start = Update { + session: "terminal-lag".into(), + call: "parent:compose:node".into(), + start: Some(json!({"toolCallId": "parent:compose:node", "status": "in_progress"})), + patch: None, + ok: false, + }; + sender.send(start.clone()).unwrap(); + receive.recv().await.unwrap(); + let mut terminal = start.clone(); + terminal.start = None; + terminal.patch = Some(json!({"sessionUpdate": "terminal_update", "terminalId": terminal.call})); + sender.send(terminal.clone()).unwrap(); + receive.recv().await.unwrap(); + // No await: force the bounded receiver to observe lag deterministically. + for _ in 0..3 { + sender.send(terminal.clone()).unwrap(); + } + let exit = receive.recv().await.unwrap(); + let value = serde_json::to_value(exit.v2().unwrap()).unwrap(); + assert_eq!(value["sessionUpdate"], "terminal_update"); + assert!(value["exitStatus"].is_object()); + assert_eq!(value["_meta"]["kit/outputIncomplete"], true); + let end = receive.recv().await.unwrap(); + assert_eq!(end.value()["status"], "failed"); + drop(sender); + task.await.unwrap(); + assert!(receive.try_recv().is_err()); +} + +#[cfg(unix)] +#[tokio::test] +async fn real_shell_cancellation_and_timeout_exit_terminal_before_failed_card() { + use agentkit_tools_core::{AllowAllPermissions, OwnedToolContext, Tool}; + use std::sync::Arc; + for cancel in [true, false] { + let session = if cancel { + "terminal-cancel" + } else { + "terminal-timeout" + }; + let root = tempfile::tempdir().unwrap(); + let tool = crate::tools::Observed::new(crate::tools::ShellTool::new(root.path().into())); + let controller = agentkit_core::CancellationController::new(); + let context = OwnedToolContext { + session_id: SessionId::new(session), + turn_id: TurnId::new("turn"), + metadata: MetadataMap::new(), + permissions: Arc::new(AllowAllPermissions), + resources: Arc::new(()), + cancellation: Some(controller.handle().checkpoint()), + execution_scope: None, + approved_request: None, + }; + let (send, mut receive) = tokio::sync::mpsc::unbounded_channel(); + let subscription = Subscription::start(session.into(), move |u| send.send(u).is_ok()); + let call = request( + session, + "shell", + json!({"command": "printf ready; sleep 30", "timeout_seconds": 1}), + ); + let execution = + tokio::spawn(async move { tool.invoke(call, &mut context.borrowed()).await }); + loop { + let update = tokio::time::timeout(std::time::Duration::from_secs(5), receive.recv()) + .await + .unwrap() + .unwrap(); + if update.value()["sessionUpdate"] == "terminal_output_chunk" { + break; + } + } + if cancel { + controller.interrupt(); + } + let result = tokio::time::timeout(std::time::Duration::from_secs(5), execution) + .await + .unwrap() + .unwrap(); + if cancel { + assert!(matches!( + result, + Err(agentkit_tools_core::ToolError::Cancelled) + )); + } else { + assert!( + matches!(result, Err(agentkit_tools_core::ToolError::ExecutionFailed(message)) if message.contains("timed out")) + ); + } + subscription.drain().await.unwrap(); + let values: Vec<_> = std::iter::from_fn(|| receive.try_recv().ok()) + .map(|u| serde_json::to_value(u.v2().unwrap()).unwrap()) + .collect(); + assert!(values[values.len() - 2]["exitStatus"].is_object()); + assert_eq!(values.last().unwrap()["status"], "failed"); + } +} diff --git a/src/tools/shell.rs b/src/tools/shell.rs index bf6566a0..a844a66e 100644 --- a/src/tools/shell.rs +++ b/src/tools/shell.rs @@ -11,6 +11,7 @@ use serde_json::{Map, Value}; use tokio::{io::AsyncReadExt, process::Command}; use crate::process_tree::{isolate_tokio_process_tree, terminate_tokio_process_tree}; +use crate::protocols::acp::tool_projection::terminal::{Output, Terminal}; const MAX_INTERNAL_OUTPUT_BYTES: usize = 64 * 1024 * 1024; @@ -113,11 +114,11 @@ impl Tool for ShellTool { async fn invoke( &self, - request: ToolRequest, + mut request: ToolRequest, context: &mut ToolContext<'_>, ) -> Result { let cancellation = context.cancellation.clone(); - let input: ShellInput = serde_json::from_value(request.input) + let input: ShellInput = serde_json::from_value(std::mem::take(&mut request.input)) .map_err(|error| ToolError::InvalidInput(error.to_string()))?; if input.command.is_empty() || !(1..=3600).contains(&input.timeout_seconds) { return Err(ToolError::InvalidInput( @@ -146,8 +147,9 @@ impl Tool for ShellTool { .stderr .take() .ok_or_else(|| ToolError::Internal("shell stderr was not piped".into()))?; - let mut stdout_task = tokio::spawn(read_output(stdout)); - let mut stderr_task = tokio::spawn(read_output(stderr)); + let terminal = Terminal::start(&request, &input.command, &self.root); + let mut stdout_task = tokio::spawn(read_output(stdout, terminal.output())); + let mut stderr_task = tokio::spawn(read_output(stderr, terminal.output())); let mut stdout_finished = false; let mut stderr_finished = false; let mut status = None; @@ -231,6 +233,7 @@ impl Tool for ShellTool { stdout.ok_or_else(|| ToolError::Internal("shell stdout was not collected".into()))?; let stderr = stderr.ok_or_else(|| ToolError::Internal("shell stderr was not collected".into()))?; + terminal.finish(status.code()); let output = Value::Object(Map::from_iter([ ( "exit_code".into(), @@ -285,7 +288,10 @@ async fn abort_output_task(task: &mut OutputTask, finished: bool) { } } -async fn read_output(mut reader: impl tokio::io::AsyncRead + Unpin) -> std::io::Result { +async fn read_output( + mut reader: impl tokio::io::AsyncRead + Unpin, + output: Option, +) -> std::io::Result { let mut content = Vec::new(); let mut buffer = [0_u8; 8192]; loop { @@ -298,6 +304,9 @@ async fn read_output(mut reader: impl tokio::io::AsyncRead + Unpin) -> std::io:: "shell output exceeds {MAX_INTERNAL_OUTPUT_BYTES} bytes" ))); } + if let Some(output) = &output { + output.chunk(&buffer[..read]); + } content.extend_from_slice(&buffer[..read]); } Ok(String::from_utf8(content) @@ -435,7 +444,7 @@ mod tests { writer.write_all(&data).await.unwrap(); writer.shutdown().await.unwrap(); - let captured = read_output(reader).await.unwrap(); + let captured = read_output(reader, None).await.unwrap(); assert_eq!(captured.as_bytes(), data); } @@ -446,7 +455,7 @@ mod tests { writer.write_all(b"small output").await.unwrap(); writer.shutdown().await.unwrap(); - let captured = read_output(reader).await.unwrap(); + let captured = read_output(reader, None).await.unwrap(); assert_eq!(captured, "small output"); }