diff --git a/docs/user/security-limits-and-troubleshooting.md b/docs/user/security-limits-and-troubleshooting.md index f87b1ebe..7d74efd8 100644 --- a/docs/user/security-limits-and-troubleshooting.md +++ b/docs/user/security-limits-and-troubleshooting.md @@ -88,6 +88,7 @@ Use `kit tui --help` for credential-store options. If nested Kit agents need the The following are fixed runtime limits, not configurable policy controls: - Shell timeout: 120 seconds by default; accepted values are 1 through 3600 seconds. Timeout reports `shell command timed out`. Shell stdout and stderr remain complete inside compose and fail if either stream exceeds the 64 MiB internal safety limit. +- ACP v2 shell terminals stream stdout and stderr as binary-safe output chunks. Per session projection subscription, Kit admits at most 128 chunks and 1 MiB total for encoded output and optional command/cwd metadata (including worst-case JSON escaping). These cumulative limits do not reset between shell calls or turns. Further output is omitted and the terminal carries `_meta["kit/outputIncomplete"] = true`; terminal exits and tool completion still report normally. This preview limit does not change the complete structured shell result or enable client-owned terminals in ACP v1. - Shell and Git timeout or output-limit cleanup targets the spawned process tree. On Unix, Kit starts the direct child in a separate process group and terminates that group; a descendant that deliberately creates a new session or process group can escape this cleanup. On Windows, Kit makes a best-effort `taskkill /PID /T /F` request, which is not a guarantee that every descendant stops. On other platforms, only direct-child termination is available. Always inspect for partial side effects after interruption or failure. - Git plugin source commands have a fixed 120-second per-command timeout and hard-bounded stdout and stderr pipes. Fetches use backoff-based live object-store checks and a final 256 MiB validation. Archive output streams directly into extraction; selected content also uses the archive entry, per-file, and expanded-size limits. Final compose results from 8 KiB through the 64 MiB result limit spill at the model-context boundary, which receives a bounded head-and-tail preview and artifact path. - Subagents: nesting depth is two and at most 120 live subagent sessions are retained per main session. Errors include `subagent depth limit (2) reached` and `live subagent session limit (120) reached`. Reuse completed sessions or release unneeded ones with `close` instead of creating unbounded children. diff --git a/src/protocols/acp/tool_projection.rs b/src/protocols/acp/tool_projection.rs index d4fb3735..7bb2bfd8 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 raw 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; + #[cfg(test)] #[allow( clippy::unwrap_used, @@ -37,9 +40,46 @@ pub(crate) struct Update { ok: bool, } +fn subscribers() -> usize { + bus().receiver_count() + v2_bus().receiver_count() +} + +// Separate ingress queues: v1 must never lag because of v2-only traffic, +// including when clients of both protocol versions coexist. +struct Buses { + v1: broadcast::Sender, + v2: broadcast::Sender, +} + +impl Buses { + fn new() -> Self { + Self { + v1: broadcast::channel(CAPACITY).0, + v2: broadcast::channel(CAPACITY).0, + } + } + + fn publish(&self, update: Update) { + if !update.v2_only() { + let _ = self.v1.send(update.clone()); + } + let _ = self.v2.send(update); + } +} + +fn buses() -> &'static Buses { + static BUSES: OnceLock = OnceLock::new(); + BUSES.get_or_init(Buses::new) +} + +fn v2_bus() -> &'static broadcast::Sender { + &buses().v2 +} fn bus() -> &'static broadcast::Sender { - static BUS: OnceLock> = OnceLock::new(); - BUS.get_or_init(|| broadcast::channel(CAPACITY).0) + &buses().v1 +} +fn publish(update: Update) { + buses().publish(update); } /// An invocation owns its terminal update, including cancellation/unwind. @@ -47,7 +87,7 @@ pub(crate) struct Invocation(Update); impl Invocation { pub(crate) fn start(request: &ToolRequest, root: Option<&Path>) -> Option { - if bus().receiver_count() == 0 + if subscribers() == 0 || request.session_id.0.len() > MAX_ID || request.call_id.0.len() > MAX_ID { @@ -106,7 +146,7 @@ impl Invocation { patch: None, ok: false, }; - let _ = bus().send(update.clone()); + publish(update.clone()); Some(Self(Update { start: None, ..update @@ -120,13 +160,13 @@ impl Invocation { impl Drop for Invocation { fn drop(&mut self) { - let _ = bus().send(self.0.clone()); + publish(self.0.clone()); } } /// Publish a location established by the tool itself, not a guessed source line. pub(crate) fn location(request: &ToolRequest, path: &Path, line: u32) { - if bus().receiver_count() == 0 + if subscribers() == 0 || request.session_id.0.len() > MAX_ID || request.call_id.0.len() > MAX_ID || !request.call_id.0.contains(":compose:") @@ -136,7 +176,7 @@ pub(crate) fn location(request: &ToolRequest, path: &Path, line: u32) { let Some(locations) = location_value(path, Some(line)) else { return; }; - let _ = bus().send(Update { + publish(Update { session: request.session_id.0.clone(), call: request.call_id.0.clone(), start: None, @@ -157,7 +197,7 @@ const MAX_DIFF_TEXT: usize = 16 * 1024; /// otherwise valid delete. Reads stop at the bound even if the file grows. pub(crate) fn deletion_text(path: &Path) -> Option { use std::io::Read; - if bus().receiver_count() == 0 { + if subscribers() == 0 { return None; } let metadata = std::fs::symlink_metadata(path).ok()?; @@ -177,7 +217,7 @@ pub(crate) fn deletion_text(path: &Path) -> Option { /// text. The patch contains both wire shapes; each protocol's typed decoder /// retains only its own fields. No renderable v2 git patch is synthesized. pub(crate) fn diff(request: &ToolRequest, path: &Path, old: Option<&str>, new: Option<&str>) { - if bus().receiver_count() == 0 + if subscribers() == 0 || request.session_id.0.len() > MAX_ID || request.call_id.0.len() > MAX_ID || !request.call_id.0.contains(":compose:") @@ -195,7 +235,7 @@ pub(crate) fn diff(request: &ToolRequest, path: &Path, old: Option<&str>, new: O (Some(_), Some(_)) => "modify", (None, None) => return, }; - let _ = bus().send(Update { + publish(Update { session: request.session_id.0.clone(), call: request.call_id.0.clone(), start: None, @@ -250,6 +290,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) @@ -259,8 +305,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) + } } } @@ -275,7 +325,21 @@ type Drain = tokio::sync::oneshot::Sender bool + Send + 'static) -> Self { - let receiver = bus().subscribe(); + Self::with_receiver(bus().subscribe(), session, send) + } + + pub(super) fn start_v2( + session: String, + send: impl Fn(Update) -> bool + Send + 'static, + ) -> Self { + Self::with_receiver(v2_bus().subscribe(), session, send) + } + + fn with_receiver( + receiver: broadcast::Receiver, + session: String, + send: impl Fn(Update) -> bool + Send + 'static, + ) -> Self { let (drains, commands) = tokio::sync::mpsc::channel(1); Self { task: tokio::spawn(forward(receiver, session, send, commands)), @@ -312,7 +376,8 @@ async fn forward( mut drains: tokio::sync::mpsc::Receiver, ) { use futures_util::future::{Either, select}; - let mut active = HashSet::new(); + let mut active = HashMap::new(); + let mut budget = terminal::Budget::default(); let mut drains_open = true; loop { let next = if drains_open { @@ -348,7 +413,14 @@ async fn forward( Err(broadcast::error::RecvError::Closed) } }; - result = forward_event(event, &mut receiver, &session, &mut active, &send); + result = forward_event( + event, + &mut receiver, + &session, + &mut active, + &mut budget, + &send, + ); if result.is_err() { break; } @@ -360,7 +432,16 @@ async fn forward( } } Either::Right(event) => { - if forward_event(event, &mut receiver, &session, &mut active, &send).is_err() { + if forward_event( + event, + &mut receiver, + &session, + &mut active, + &mut budget, + &send, + ) + .is_err() + { return; } } @@ -372,20 +453,28 @@ fn forward_event( event: Result, receiver: &mut broadcast::Receiver, session: &str, - active: &mut HashSet, + active: &mut HashMap, + budget: &mut terminal::Budget, send: &impl Fn(Update) -> bool, ) -> Result<(), agentkit_acp::AcpRuntimeError> { match event { - Ok(update) if update.session == session => { + Ok(mut 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(), terminal::State::default()); + } else if let Some(patch) = &update.patch { + let Some(state) = active.get_mut(&update.call) else { return Ok(()); + }; + if patch.get("sessionUpdate").and_then(Value::as_str) == Some("terminal_update") { + state.running = patch.get("exitStatus").is_none(); } - } else if !active.remove(&update.call) { + if !budget.admit(&mut update, state) { + return Ok(()); + } + } else if active.remove(&update.call).is_none() { return Ok(()); } if !send(update) { @@ -394,7 +483,31 @@ fn forward_event( } Ok(_) => {} Err(error) => { - for call in active.drain() { + for (call, state) 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 state.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..5cf1953f --- /dev/null +++ b/src/protocols/acp/tool_projection/terminal.rs @@ -0,0 +1,120 @@ +//! Agent-owned terminals are v2 only. Frames share the bounded invocation bus. +use super::{MAX_ID, Update, publish, v2_bus}; +use agentkit_tools_core::ToolRequest; +use base64::{Engine as _, engine::general_purpose::STANDARD}; +use serde_json::{Map, Value}; +use std::path::Path; + +mod budget; +pub(super) use budget::{Budget, State}; + +/// 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) { + publish(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 v2_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/terminal/budget.rs b/src/protocols/acp/tool_projection/terminal/budget.rs new file mode 100644 index 00000000..fca80f09 --- /dev/null +++ b/src/protocols/acp/tool_projection/terminal/budget.rs @@ -0,0 +1,84 @@ +//! Cumulative admission, not a per-call rate limit: a stalled SDK transport can +//! retain only this much extra stream payload for the subscription's lifetime. +//! Invocation/terminal lifecycle delivery is never charged against this budget. +use super::super::Update; +use serde_json::{Map, Value}; + +pub(super) const MAX_BYTES: usize = 1024 * 1024; +pub(super) const MAX_CHUNKS: usize = 128; + +#[derive(Default)] +pub(in super::super) struct State { + pub(in super::super) running: bool, + incomplete: bool, +} + +pub(in super::super) struct Budget { + bytes: usize, + chunks: usize, +} + +impl Default for Budget { + fn default() -> Self { + Self { + bytes: MAX_BYTES, + chunks: MAX_CHUNKS, + } + } +} + +impl Budget { + pub(in super::super) fn admit(&mut self, update: &mut Update, state: &mut State) -> bool { + let Some(patch) = update.patch.as_mut().and_then(Value::as_object_mut) else { + return true; + }; + match patch.get("sessionUpdate").and_then(Value::as_str) { + Some("terminal_output_chunk") => { + if state.incomplete { + return false; + } + let bytes = patch + .get("data") + .and_then(Value::as_str) + .map_or(0, str::len); + if self.chunks > 0 && bytes <= self.bytes { + self.bytes -= bytes; + self.chunks -= 1; + } else { + state.incomplete = true; + // Replace the first omitted chunk with one small, explicit + // notice; suppress the rest without dropping the exit. + *patch = Map::from_iter([ + ("sessionUpdate".into(), Value::from("terminal_update")), + ("terminalId".into(), Value::from(update.call.clone())), + ( + "_meta".into(), + Value::Object(Map::from_iter([( + "kit/outputIncomplete".into(), + Value::from(true), + )])), + ), + ]); + } + } + Some("terminal_update") => { + // Command/cwd are optional, variable-size metadata. Charge the + // worst-case JSON escape expansion too, across all shell calls. + for key in ["command", "cwd"] { + let bytes = patch + .get(key) + .and_then(Value::as_str) + .map_or(0, str::len) + .saturating_mul(6); + if bytes <= self.bytes { + self.bytes -= bytes; + } else { + patch.remove(key); + } + } + } + _ => {} + } + true + } +} diff --git a/src/protocols/acp/tool_projection/tests.rs b/src/protocols/acp/tool_projection/tests.rs index 67e12a49..30e86d14 100644 --- a/src/protocols/acp/tool_projection/tests.rs +++ b/src/protocols/acp/tool_projection/tests.rs @@ -273,3 +273,314 @@ 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_v2("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 = v2_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_v2(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"); + } +} + +#[tokio::test] +async fn terminal_bursts_never_enter_the_v1_queue_with_mixed_clients() { + let buses = Buses::new(); + let (send, mut receive) = tokio::sync::mpsc::unbounded_channel(); + let (drains, commands) = tokio::sync::mpsc::channel(1); + let subscription = Subscription { + task: tokio::spawn(forward( + buses.v1.subscribe(), + "v1-isolated".into(), + move |u| send.send(u).is_ok(), + commands, + )), + drains, + }; + let mut v2 = buses.v2.subscribe(); + let start = Update { + session: "v1-isolated".into(), + call: "parent:compose:node".into(), + start: Some(json!({"toolCallId": "parent:compose:node", "status": "in_progress"})), + patch: None, + ok: false, + }; + buses.publish(start.clone()); + for _ in 0..CAPACITY + 1 { + buses.publish(Update { + session: "v2-burst".into(), call: "parent:compose:other".into(), + start: None, patch: Some(json!({"sessionUpdate": "terminal_output_chunk", "terminalId": "parent:compose:other", "data": "YQ=="})), ok: false, + }); + } + buses.publish(Update { + start: None, + ok: true, + ..start + }); + // The v2 receiver really did lag, while v1 had only lifecycle traffic. + assert!(matches!( + v2.try_recv(), + Err(broadcast::error::TryRecvError::Lagged(_)) + )); + subscription.drain().await.unwrap(); + let start = receive.try_recv().unwrap(); + let end = receive.try_recv().unwrap(); + assert!(start.start.is_some()); + assert_eq!(end.value()["status"], "completed"); + assert!(receive.try_recv().is_err()); +} + +#[tokio::test] +async fn cumulative_budget_bounds_a_stalled_transport_across_calls() { + for size in [1, 8192] { + // This external sink deliberately accepts without consumption, like the + // SDK's unbounded queue. The budget must hold even without bus lag. + let (send, mut receive) = tokio::sync::mpsc::unbounded_channel(); + let subscription = + Subscription::start_v2("slow-terminal".into(), move |u| send.send(u).is_ok()); + for call in 0..4 { + let mut request = request("slow-terminal", "shell", json!({})); + request.call_id = ToolCallId::new(format!("parent:compose:{call}")); + let invocation = Invocation::start(&request, None).unwrap(); + let terminal = terminal::Terminal::start(&request, "printf lots", Path::new("/tmp")); + subscription.drain().await.unwrap(); + for _ in 0..64 { + terminal.output().unwrap().chunk(&vec![255; size]); + subscription.drain().await.unwrap(); + } + terminal.finish(Some(0)); + invocation.finish(true); + 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(); + let chunks: Vec<_> = values.iter().filter_map(|v| v["data"].as_str()).collect(); + assert!(!chunks.is_empty()); + assert!(chunks.len() <= 128); + assert!(chunks.iter().map(|s| s.len()).sum::() <= 1024 * 1024); + assert!( + values + .iter() + .any(|v| v["_meta"]["kit/outputIncomplete"] == true) + ); + assert_eq!( + values.iter().filter(|v| v["status"] == "completed").count(), + 4 + ); + assert_eq!( + values + .iter() + .filter(|v| v["exitStatus"]["exitCode"] == 0) + .count(), + 4 + ); + assert!(!values.iter().any(|v| v["status"] == "failed")); + } +} diff --git a/src/protocols/acp/v2.rs b/src/protocols/acp/v2.rs index f272113b..8adaa750 100644 --- a/src/protocols/acp/v2.rs +++ b/src/protocols/acp/v2.rs @@ -429,7 +429,7 @@ where self.activity.tool_projection.get_or_init(|| { let sink = self.sink.clone(); let session_id = self.session_id.clone(); - super::tool_projection::Subscription::start( + super::tool_projection::Subscription::start_v2( event.session_id.0.clone(), move |update| { update.v2().is_ok_and(|update| { 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"); }