From 66305be4e17dcffd7df5287af205aa8e4bc59b62 Mon Sep 17 00:00:00 2001 From: daniel Date: Sun, 30 Aug 2026 12:21:46 +0100 Subject: [PATCH 1/7] feat(acp): add unstable v2 session injection --- Cargo.toml | 4 + book/src/acp.md | 38 +- crates/agentkit-acp/Cargo.toml | 5 + crates/agentkit-acp/README.md | 49 +- crates/agentkit-acp/src/v2.rs | 3890 +++++++++++++++++++++++++++++--- docs/acp.md | 31 +- 6 files changed, 3653 insertions(+), 364 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index c64d394..1206993 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -60,6 +60,10 @@ repository = "https://github.com/danielkov/agentkit" rust-version = "1.92" version = "0.10.5" +# Intentionally unpublishable until the upstream SDK releases session injection. +[patch.crates-io] +agent-client-protocol = { git = "https://github.com/danielkov/rust-sdk", rev = "2f039993d1d6ed8da35b38c31f54a7cbb7338c70" } + [workspace.dependencies] async-trait = "0.1.89" dotenvy = "0.15.7" diff --git a/book/src/acp.md b/book/src/acp.md index 02d0769..48ad854 100644 --- a/book/src/acp.md +++ b/book/src/acp.md @@ -2,27 +2,28 @@ The [Agent Client Protocol (ACP)](https://agentclientprotocol.com) standardizes communication between clients (code editors, IDEs, desktop apps) and coding agents. Where MCP connects an agent to external tools, ACP connects a client to the agent itself: session lifecycle, prompt turns, streamed updates, tool call reporting, and permission prompts all travel over JSON-RPC — usually with the agent running as an editor child process on stdio. This chapter covers [`agentkit-acp`](https://github.com/danielkov/agentkit/tree/main/crates/agentkit-acp): how an agentkit host becomes ACP-addressable, and how a standalone agent serves ACP directly. -## Built on the official SDK +## Built on the ACP Rust SDK -Like `agentkit-mcp`, this crate does not define a parallel protocol vocabulary. It builds on the official Rust SDK, [`agent-client-protocol`](https://crates.io/crates/agent-client-protocol), and re-exports the stable v1 wire types (`SessionId`, `ContentBlock`, `SessionUpdate`, `ToolCallUpdate`, `StopReason`, …) at the crate root and under `agentkit_acp::wire`. The full upstream SDK is available as `agentkit_acp::sdk`. Agentkit owns only the host-facing glue: session binding, observer routing, prompt conversion, cancellation handles, and approval resolution. +Like `agentkit-mcp`, this crate does not define a parallel protocol vocabulary. It currently builds on a commit-pinned [`agent-client-protocol` fork](https://github.com/danielkov/rust-sdk), needed for experimental session injection support, and re-exports the stable v1 wire types (`SessionId`, `ContentBlock`, `SessionUpdate`, `ToolCallUpdate`, `StopReason`, …) at the crate root and under `agentkit_acp::wire`. The full SDK is available as `agentkit_acp::sdk`. Agentkit owns only the host-facing glue: session binding, observer routing, prompt conversion, cancellation handles, and approval resolution. - **Protocol docs:** [agentclientprotocol.com](https://agentclientprotocol.com/protocol/v1/overview) -- **Rust SDK:** [`agent-client-protocol` on crates.io](https://crates.io/crates/agent-client-protocol) +- **Rust SDK fork:** [`danielkov/rust-sdk`](https://github.com/danielkov/rust-sdk), pinned in `Cargo.toml` and `Cargo.lock` ## Opt-in ACP v2 runtime ACP v2 support is additive and disabled by default. Enable it explicitly: ```toml -agentkit-acp = { version = "0.10.8", features = ["protocol-v2"] } +agentkit-acp = { version = "0.10.11", features = ["protocol-v2"] } ``` -`protocol-v2` enables the official upstream +`protocol-v2` enables the fork's `agent-client-protocol/unstable_protocol_v2` feature. The root API and `agentkit_acp::wire` continue to expose stable v1 behavior. Experimental v2 -runtime APIs and official v2 wire types are isolated under +runtime APIs and v2 wire types are isolated under `agentkit_acp::v2` and `agentkit_acp::v2::wire`; v1 wire types are not part of -that namespace. +that namespace. The additive `unstable-inject` feature implies `protocol-v2` +and enables the unstable session-injection methods. Build a v2 server with `agentkit_acp::v2::AcpHeadlessRuntime`. Its factory is called once for each `session/new` and receives a v2 session ID, an agentkit @@ -47,11 +48,22 @@ progress concurrently, while a second prompt for a running session is rejected. work and drops the session worker. `session/list` and `session/resume` cover active in-memory sessions; replay is not supported. +Session injection is steer-only and finishes an in-flight model stream rather +than interrupting it. Acceptance follows response-frame enqueue, including for +JSON-RPC batches; the full `ContentBlock` list is emitted unchanged at the next +safe model/tool boundary. Request or session cancellation cannot remove a +response-committed steer, so it carries to the next valid boundary when the +current turn is cancelled. Close may discard it. Revoke is serialized with +acknowledged user-message forwarding. Queueing and replacement are not +supported. Delivered IDs retain `already_delivered` classification for the +session lifetime; a 4,096-accept lifetime cap bounds that history. + The initial v2 foundation routes text, reasoning, and tool lifecycle updates. -ACP v2 permission callbacks are intentionally deferred; an unsupported approval -interrupt retains the transcript and ends the prompt with the custom `_error` -stop reason rather than `refusal`. Upstream labels the v2 protocol unstable, so -opt-in callers should expect the `v2` namespace to track official SDK changes. +ACP v2 permission callbacks are intentionally deferred; unsupported approval +requests are denied while accepted steers remain pending for the next safe +boundary. The SDK labels v2 unstable, so opt-in callers should expect the `v2` +namespace to track the pinned fork. The workspace patch is intentionally +unpublishable until the required APIs are released upstream. ## Two integration shapes @@ -214,7 +226,7 @@ On the umbrella crate, ACP is behind the `acp` feature (implies `loop`): ```toml [dependencies] -agentkit = { version = "0.10.8", features = ["acp"] } +agentkit = { version = "0.10.11", features = ["acp"] } ``` The `agentkit-acp` crate itself has a default `stdio` feature that gates `serve_stdio`, which uses the upstream SDK's built-in `Stdio` transport. Its `unstable-acp` feature forwards to the upstream SDK's unstable protocol surface. @@ -228,4 +240,4 @@ The `agentkit-acp` crate itself has a default `stdio` feature that gates `serve_ > **Example:** [`openrouter-acp-trio`](https://github.com/danielkov/agentkit/tree/main/examples/openrouter-acp-trio) runs three OpenRouter-backed agents (orchestrator, worker, reviewer) that call each other over in-memory ACP endpoints while a REPL drives the orchestrator through a persistent ACP session — session binding, streamed updates, tool call reporting, and agent-to-agent handoffs in one program. > -> **Crate:** [`agentkit-acp`](https://github.com/danielkov/agentkit/tree/main/crates/agentkit-acp) — depends on [`agentkit-core`](https://github.com/danielkov/agentkit/tree/main/crates/agentkit-core), [`agentkit-loop`](https://github.com/danielkov/agentkit/tree/main/crates/agentkit-loop), [`agentkit-tools-core`](https://github.com/danielkov/agentkit/tree/main/crates/agentkit-tools-core), and [`agent-client-protocol`](https://crates.io/crates/agent-client-protocol). Design notes: [`docs/acp.md`](https://github.com/danielkov/agentkit/blob/main/docs/acp.md). +> **Crate:** [`agentkit-acp`](https://github.com/danielkov/agentkit/tree/main/crates/agentkit-acp) — depends on [`agentkit-core`](https://github.com/danielkov/agentkit/tree/main/crates/agentkit-core), [`agentkit-loop`](https://github.com/danielkov/agentkit/tree/main/crates/agentkit-loop), [`agentkit-tools-core`](https://github.com/danielkov/agentkit/tree/main/crates/agentkit-tools-core), and a commit-pinned [`agent-client-protocol` fork](https://github.com/danielkov/rust-sdk). Design notes: [`docs/acp.md`](https://github.com/danielkov/agentkit/blob/main/docs/acp.md). diff --git a/crates/agentkit-acp/Cargo.toml b/crates/agentkit-acp/Cargo.toml index 993b3f6..fe158be 100644 --- a/crates/agentkit-acp/Cargo.toml +++ b/crates/agentkit-acp/Cargo.toml @@ -14,6 +14,10 @@ default = ["stdio"] stdio = [] unstable-acp = ["agent-client-protocol/unstable"] protocol-v2 = ["agent-client-protocol/unstable_protocol_v2"] +unstable-inject = [ + "protocol-v2", + "agent-client-protocol/unstable_session_inject", +] [dependencies] agent-client-protocol = "=2.0.0" @@ -28,6 +32,7 @@ tokio = { workspace = true, features = ["sync", "time"] } tracing.workspace = true [dev-dependencies] +futures-util.workspace = true agentkit-integration-tests = { path = "../agentkit-integration-tests" } agentkit-task-manager = { path = "../agentkit-task-manager" } tokio = { workspace = true, features = ["macros", "rt-multi-thread", "time"] } diff --git a/crates/agentkit-acp/README.md b/crates/agentkit-acp/README.md index 07a1e58..4ac26ae 100644 --- a/crates/agentkit-acp/README.md +++ b/crates/agentkit-acp/README.md @@ -2,7 +2,7 @@ Agent Client Protocol integration for agentkit hosts. -This crate re-exports upstream ACP wire types from `agent-client-protocol` and +This crate re-exports ACP wire types from a pinned `agent-client-protocol` fork and adds only agentkit-specific glue: session binding, observer routing, prompt conversion, cancellation handles, and approval resolver abstractions. @@ -21,28 +21,51 @@ The crate root, default features, and `wire` module remain ACP v1. To use the experimental upstream ACP v2 protocol, enable the additive feature: ```toml -agentkit-acp = { version = "0.10.8", features = ["protocol-v2"] } +agentkit-acp = { version = "0.10.11", features = ["protocol-v2"] } ``` -The feature maps directly to the official -`agent-client-protocol/unstable_protocol_v2` feature. V2 APIs and official v2 -wire types live only under `agentkit_acp::v2` (and -`agentkit_acp::v2::wire`). +The feature maps directly to the pinned fork's +`agent-client-protocol/unstable_protocol_v2` feature. V2 APIs and v2 wire +types live only under `agentkit_acp::v2` (and +`agentkit_acp::v2::wire`). Enable `unstable-inject` instead to add the unstable +ACP v2 session-injection surface; it implies `protocol-v2`. `v2::AcpHeadlessRuntime` supports ACP v2 initialize, new/list/resume session, prompt, cancel, session updates, and close. Listing and resume cover the runtime's active in-memory sessions; replay is not supported. Each session owns a worker and loop driver, so work in one session does not block request -handling for another. A prompt is accepted before model work completes, then -the runtime emits ordered `UserMessage`, `Running`, streamed output, and `Idle` -updates. User, visible-agent, and thought message IDs are distinct and stable -for the lifetime of a prompt. +handling for another. A prompt is acknowledged before model work completes. `Running` and `Idle` are +derived from authoritative loop `TurnStarted` and `TurnFinished` events, not +from calls that merely poll the driver. The session worker also wakes for +autonomous task-manager updates, while prompt ownership, cancellation, and +logical-turn cleanup are serialized so a prompt finishes only once. Session +updates retain structured background content and acknowledged updates do not +overtake earlier notifications. User, visible-agent, and thought message IDs +are distinct and stable for the lifetime of a prompt. + +With `unstable-inject`, the runtime advertises only `steer` delivery with +finish-current-stream behavior. `session/inject` returns an agent-owned message +ID after its response frame is enqueued, preserves every `ContentBlock`, and +delivers at the next safe model/tool boundary. Single and batch requests use the +same receipt-backed acceptance path. Once response commitment succeeds, request +or session cancellation cannot remove the reservation; a committed steer that +misses the cancelled turn carries into the next prompt. Closing the session may +discard it. `session/revoke_inject` is mandatory and is serialized with +acknowledged `UserMessage` forwarding. Queueing, stream interruption, and +replacement are not supported. Delivered IDs retain their `already_delivered` +classification for the session lifetime. To bound that history safely, each +session accepts at most 4,096 injections and rejects later accepts with a +lifetime-limit error. This first v2 foundation streams text, reasoning, and tool lifecycle updates. The v1 permission bridge is not exposed through v2 wire types. Unsupported -approval interrupts retain the transcript and therefore end with the custom -`_error` stop reason rather than `Refusal`. Because upstream marks protocol v2 -unstable, all APIs in the `v2` namespace can evolve with the official SDK. +approval requests are resolved as denials while retaining a provider-valid +transcript, so the prompt ends with the custom `_error` stop reason rather than +`Refusal`. Already accepted steers remain pending and are delivered at the next +safe boundary. Because the SDK marks +protocol v2 unstable, all APIs in the `v2` namespace can evolve with the pinned +SDK fork. The workspace patch is intentionally unpublishable until these SDK +APIs are available in an upstream release. Run the stable v1 in-memory end-to-end example with: diff --git a/crates/agentkit-acp/src/v2.rs b/crates/agentkit-acp/src/v2.rs index ba5b8b3..fb7c138 100644 --- a/crates/agentkit-acp/src/v2.rs +++ b/crates/agentkit-acp/src/v2.rs @@ -1,15 +1,17 @@ //! Opt-in runtime foundation for the experimental ACP protocol v2. //! //! Enable the `protocol-v2` crate feature to use this module. The feature maps -//! directly to the official `agent-client-protocol/unstable_protocol_v2` -//! feature. Root-level APIs remain the stable ACP v1 integration. +//! directly to the pinned SDK fork's +//! `agent-client-protocol/unstable_protocol_v2` feature. Root-level APIs remain the stable ACP v1 integration. use std::collections::HashMap; +#[cfg(feature = "unstable-inject")] +use std::collections::VecDeque; use std::path::PathBuf; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::{Arc, Mutex, RwLock}; -use agent_client_protocol::{Client, ConnectionTo, Handled}; +use agent_client_protocol::{Client, Handled, V2ConnectionTo}; use agentkit_core::{ CancellationController, CancellationHandle, DataRef, Delta, FilePart, FinishReason, Item, ItemKind, MediaPart, MetadataMap, Modality, Part, PartId, PartKind, @@ -22,26 +24,67 @@ use agentkit_loop::{ }; use async_trait::async_trait; use serde_json::json; +#[cfg(feature = "unstable-inject")] +use tokio::sync::Notify; use tokio::sync::{Mutex as AsyncMutex, mpsc, oneshot}; use crate::AcpRuntimeError; -/// Official upstream ACP v2 wire types. +#[cfg(feature = "unstable-inject")] +fn sdk_v2_error(error: wire::Error) -> agent_client_protocol::Error { + agent_client_protocol::Error::new(i32::from(error.code), error.message).data(error.data) +} + +#[cfg(feature = "unstable-inject")] +fn session_not_found_error(session_id: &wire::SessionId) -> agent_client_protocol::Error { + sdk_v2_error(wire::Error::resource_not_found(None)).data(serde_json::json!({ + "sessionId": session_id, + })) +} + +/// ACP v2 wire types from the pinned SDK fork. /// -/// These are gated by upstream's `unstable_protocol_v2` feature and can change +/// These are gated by the SDK's `unstable_protocol_v2` feature and can change /// while ACP v2 is under development. No stable v1 wire type is re-exported /// from this namespace. pub use agent_client_protocol::schema::ProtocolVersion; pub use agent_client_protocol::schema::v2::*; -/// Explicit namespace for the official upstream ACP v2 wire types. +/// Explicit namespace for ACP v2 wire types from the pinned SDK fork. pub mod wire { pub use agent_client_protocol::schema::ProtocolVersion; pub use agent_client_protocol::schema::v2::*; } +/// Host-provided destination for ACP v2 session updates. +/// +/// This unstable v2 API deliberately hides the runtime's internal channels. A +/// host can forward updates through its own ACP connection, event loop, or test +/// sink. Calls for one binding arrive in protocol order; implementations must +/// preserve that order when they enqueue them. Acknowledged updates must +/// complete only after the notification has been accepted by the destination +/// and must not overtake earlier updates. +#[async_trait] +pub trait AcpSessionUpdateSink: Send + Sync + 'static { + /// Forwards a session update without waiting for delivery acknowledgement. + fn update(&self, notification: wire::UpdateSessionNotification) -> Result<(), AcpRuntimeError>; + + /// Forwards a session update and waits until the destination accepts it. + async fn update_acknowledged( + &self, + notification: wire::UpdateSessionNotification, + ) -> Result<(), AcpRuntimeError>; + + /// Waits until all earlier updates have been accepted by the destination. + async fn flush(&self) -> Result<(), AcpRuntimeError>; +} + enum ClientMessage { Update(Box), + AcknowledgedUpdate { + notification: Box, + acknowledged: oneshot::Sender>, + }, Flush(oneshot::Sender<()>), } @@ -74,7 +117,7 @@ impl ClientHandle { (client, rx) } - fn update( + fn update_for( &self, session_id: wire::SessionId, update: wire::SessionUpdate, @@ -83,13 +126,32 @@ impl ClientHandle { if let Some(hook) = &self.before_update { hook(&update); } + self.update(wire::UpdateSessionNotification::new(session_id, update)) + } +} + +#[async_trait] +impl AcpSessionUpdateSink for ClientHandle { + fn update(&self, notification: wire::UpdateSessionNotification) -> Result<(), AcpRuntimeError> { self.tx - .send(ClientMessage::Update(Box::new( - wire::UpdateSessionNotification::new(session_id, update), - ))) + .send(ClientMessage::Update(Box::new(notification))) .map_err(|_| AcpRuntimeError::ClientClosed) } + async fn update_acknowledged( + &self, + notification: wire::UpdateSessionNotification, + ) -> Result<(), AcpRuntimeError> { + let (tx, rx) = oneshot::channel(); + self.tx + .send(ClientMessage::AcknowledgedUpdate { + notification: Box::new(notification), + acknowledged: tx, + }) + .map_err(|_| AcpRuntimeError::ClientClosed)?; + rx.await.map_err(|_| AcpRuntimeError::ClientClosed)? + } + async fn flush(&self) -> Result<(), AcpRuntimeError> { let (tx, rx) = oneshot::channel(); self.tx @@ -101,7 +163,7 @@ impl ClientHandle { async fn drain_client_messages( mut rx: mpsc::UnboundedReceiver, - cx: ConnectionTo, + cx: V2ConnectionTo, ) { while let Some(message) = rx.recv().await { match message { @@ -111,6 +173,19 @@ async fn drain_client_messages( break; } } + ClientMessage::AcknowledgedUpdate { + notification, + acknowledged, + } => { + let result = cx + .send_notification(*notification) + .map_err(|_| AcpRuntimeError::ClientClosed); + let failed = result.is_err(); + let _ = acknowledged.send(result); + if failed { + break; + } + } ClientMessage::Flush(response) => { let _ = response.send(()); } @@ -126,7 +201,13 @@ struct CurrentMessageIds { struct IntegrationSession { acp_session_id: wire::SessionId, - client: ClientHandle, + agentkit_session_id: AgentkitSessionId, + sink: Arc, + cancellation: CancellationController, + closed: AtomicBool, + lifecycle: Arc>, + #[cfg(feature = "unstable-inject")] + injection: Arc, next_message: AtomicU64, current_messages: Mutex>, part_kinds: Mutex>, @@ -134,6 +215,17 @@ struct IntegrationSession { prompt_state: Mutex>, } +impl IntegrationSession { + fn update_for( + &self, + session_id: wire::SessionId, + update: wire::SessionUpdate, + ) -> Result<(), AcpRuntimeError> { + self.sink + .update(wire::UpdateSessionNotification::new(session_id, update)) + } +} + struct PromptState { active_prompt: Arc, lifecycle: Arc>, @@ -147,28 +239,161 @@ struct PromptOwner { cancellation: TurnCancellation, } +/// Unstable host-owned binding for one ACP v2 session. +/// +/// This API follows the experimental ACP v2 schema and may change with the +/// pinned SDK. The update sink lets hosts keep transport ownership. +pub struct AcpSessionBinding { + acp_session_id: wire::SessionId, + agentkit_session_id: AgentkitSessionId, + sink: Arc, + cancellation: Option, +} + +impl AcpSessionBinding { + /// Creates an unstable ACP v2 session binding. + #[must_use] + pub fn new( + acp_session_id: wire::SessionId, + agentkit_session_id: AgentkitSessionId, + sink: impl AcpSessionUpdateSink, + ) -> Self { + Self { + acp_session_id, + agentkit_session_id, + sink: Arc::new(sink), + cancellation: None, + } + } + + /// Uses a host-owned cancellation controller for this unstable v2 session. + #[must_use] + pub fn cancellation(mut self, cancellation: CancellationController) -> Self { + self.cancellation = Some(cancellation); + self + } +} + +/// Unstable host handle for one bound ACP v2 session. +#[derive(Clone)] +pub struct AcpSessionHandle { + session: Arc, +} + +impl AcpSessionHandle { + /// Returns the client-visible ACP v2 session ID. + #[must_use] + pub fn acp_session_id(&self) -> &wire::SessionId { + &self.session.acp_session_id + } + + /// Returns the agentkit loop session ID. + #[must_use] + pub fn agentkit_session_id(&self) -> &AgentkitSessionId { + &self.session.agentkit_session_id + } + + /// Returns the cancellation handle for this unstable v2 session. + #[must_use] + pub fn cancellation_handle(&self) -> CancellationHandle { + self.session.cancellation.handle() + } + + /// Interrupts the active turn, preserving response-committed steers. + pub fn interrupt(&self) { + let _lifecycle = self + .session + .lifecycle + .lock() + .unwrap_or_else(|error| error.into_inner()); + if !self.session.closed.load(Ordering::Acquire) { + #[cfg(feature = "unstable-inject")] + self.session.injection.cancel_turn(); + self.session.cancellation.interrupt(); + } + } + + /// Closes this unstable v2 session and discards deliverable steers. + pub fn close(&self) { + let _lifecycle = self + .session + .lifecycle + .lock() + .unwrap_or_else(|error| error.into_inner()); + self.session.closed.store(true, Ordering::Release); + #[cfg(feature = "unstable-inject")] + self.session.injection.close_session(); + self.session.cancellation.interrupt(); + } + + fn is_closed(&self) -> bool { + self.session.closed.load(Ordering::Acquire) + } + + /// Prepares a prompt admission for unstable ACP v2 injection. + #[cfg(feature = "unstable-inject")] + pub fn prepare_injection_turn(&self) { + self.session.injection.reset_cancellation(); + } + + /// Makes the admitted turn available to unstable ACP v2 injection. + #[cfg(feature = "unstable-inject")] + pub fn start_injection_turn(&self) { + self.session.injection.start_turn(); + } + + /// Stops accepting unstable ACP v2 injection for the current turn. + #[cfg(feature = "unstable-inject")] + pub fn stop_injection_turn(&self) { + self.session.injection.stop_turn(); + } +} + #[derive(Default)] struct IntegrationInner { by_acp: HashMap>, by_agentkit: HashMap, } -/// Routes agentkit loop output to ACP v2 `session/update` notifications. +#[cfg(feature = "unstable-inject")] +#[derive(Default)] +struct InjectRequestOrder { + next: AtomicU64, + current: AtomicU64, + changed: Notify, +} + +/// Routes agentkit loop output and coordinates unstable ACP v2 host sessions. /// -/// Agent factories should install this value as their loop observer. The -/// headless runtime binds and unbinds sessions automatically. -#[derive(Clone, Default)] +/// Agent factories should install this value as their loop observer. Hosts can +/// bind sessions directly; [`AcpHeadlessRuntime`] delegates to the same public +/// coordinator. This entire API follows the experimental ACP v2 schema. +#[derive(Clone)] pub struct AcpIntegration { inner: Arc>, + #[cfg(feature = "unstable-inject")] + inject_requests: Arc, +} + +impl Default for AcpIntegration { + fn default() -> Self { + Self { + inner: Arc::new(RwLock::new(IntegrationInner::default())), + #[cfg(feature = "unstable-inject")] + inject_requests: Arc::new(InjectRequestOrder::default()), + } + } } impl AcpIntegration { - fn bind( + /// Binds an unstable ACP v2 session to the public coordinator. + pub fn bind_session( &self, - acp_session_id: wire::SessionId, - agentkit_session_id: AgentkitSessionId, - client: ClientHandle, - ) -> Result<(), AcpRuntimeError> { + binding: AcpSessionBinding, + ) -> Result { + let cancellation = binding.cancellation.unwrap_or_default(); + let acp_session_id = binding.acp_session_id; + let agentkit_session_id = binding.agentkit_session_id; let mut inner = self .inner .write() @@ -185,23 +410,43 @@ impl AcpIntegration { } inner .by_agentkit - .insert(agentkit_session_id, acp_session_id.clone()); - inner.by_acp.insert( - acp_session_id.clone(), - Arc::new(IntegrationSession { - acp_session_id, - client, - next_message: AtomicU64::new(1), - current_messages: Mutex::new(None), - part_kinds: Mutex::new(HashMap::new()), - unsupported_approval: Mutex::new(None), - prompt_state: Mutex::new(None), - }), - ); - Ok(()) + .insert(agentkit_session_id.clone(), acp_session_id.clone()); + let session = Arc::new(IntegrationSession { + acp_session_id: acp_session_id.clone(), + agentkit_session_id, + sink: binding.sink, + cancellation, + closed: AtomicBool::new(false), + lifecycle: Arc::new(Mutex::new(())), + #[cfg(feature = "unstable-inject")] + injection: Arc::new(InjectionController::default()), + next_message: AtomicU64::new(1), + current_messages: Mutex::new(None), + part_kinds: Mutex::new(HashMap::new()), + unsupported_approval: Mutex::new(None), + prompt_state: Mutex::new(None), + }); + inner.by_acp.insert(acp_session_id, Arc::clone(&session)); + Ok(AcpSessionHandle { session }) + } + + #[cfg(test)] + fn bind( + &self, + acp_session_id: wire::SessionId, + agentkit_session_id: AgentkitSessionId, + client: ClientHandle, + ) -> Result<(), AcpRuntimeError> { + self.bind_session(AcpSessionBinding::new( + acp_session_id, + agentkit_session_id, + client, + )) + .map(|_| ()) } - fn unbind(&self, session_id: &wire::SessionId) -> Result<(), AcpRuntimeError> { + /// Unbinds and closes an unstable ACP v2 session. + pub fn unbind_session(&self, session_id: &wire::SessionId) -> Result<(), AcpRuntimeError> { let mut inner = self .inner .write() @@ -213,6 +458,8 @@ impl AcpIntegration { inner .by_agentkit .retain(|_, mapped| mapped != &session.acp_session_id); + drop(inner); + AcpSessionHandle { session }.close(); Ok(()) } @@ -229,6 +476,52 @@ impl AcpIntegration { .ok_or_else(|| AcpRuntimeError::SessionNotFound(session_id.to_string())) } + fn next_user_message_id( + &self, + session_id: &wire::SessionId, + ) -> Result { + let session = self.session(session_id)?; + let sequence = session.next_message.fetch_add(1, Ordering::Relaxed); + Ok(wire::MessageId::new(format!( + "{session_id}-user-{sequence}" + ))) + } + + /// Converts an unstable ACP v2 prompt into agentkit input items. + pub fn prompt_to_items( + &self, + request: &wire::PromptRequest, + ) -> Result, AcpRuntimeError> { + self.session(&request.session_id)?; + prompt_to_items(request) + } + + /// Starts unstable ACP v2 prompt message routing. + pub fn begin_prompt( + &self, + session_id: &wire::SessionId, + ) -> Result { + let message_id = self.next_user_message_id(session_id)?; + let session = self.session(session_id)?; + finish_model_message(&session); + Ok(message_id) + } + + /// Finishes unstable ACP v2 prompt message routing. + pub fn finish_prompt(&self, session_id: &wire::SessionId) { + if let Ok(session) = self.session(session_id) { + finish_model_message(&session); + } + } + + /// Flushes updates already submitted for an unstable ACP v2 session. + pub async fn flush_session_updates( + &self, + session_id: &wire::SessionId, + ) -> Result<(), AcpRuntimeError> { + self.session(session_id)?.sink.flush().await + } + fn install_prompt_state( &self, session_id: &wire::SessionId, @@ -248,7 +541,7 @@ impl AcpIntegration { Ok(()) } - fn begin_prompt( + fn begin_prompt_owner( &self, session_id: &wire::SessionId, owner: u64, @@ -269,13 +562,12 @@ impl AcpIntegration { }); } let sequence = session.next_message.fetch_add(1, Ordering::Relaxed); - finish_model_message(&session); Ok(wire::MessageId::new(format!( "{session_id}-user-{sequence}" ))) } - fn finish_prompt(&self, session_id: &wire::SessionId, owner: u64) { + fn finish_prompt_owner(&self, session_id: &wire::SessionId, owner: u64) { if let Ok(session) = self.session(session_id) { if let Some(prompt_state) = session .prompt_state @@ -355,7 +647,7 @@ impl AcpIntegration { } }; release_prompt(&active_prompt, owner.id); - if let Err(error) = session.client.update( + if let Err(error) = session.update_for( session.acp_session_id.clone(), wire::SessionUpdate::StateUpdate(wire::StateUpdate::Idle( wire::IdleStateUpdate::new().stop_reason(stop_reason), @@ -392,7 +684,7 @@ impl AcpIntegration { .is_some(); if prompt_owned { start_model_message(&session); - if let Err(error) = session.client.update( + if let Err(error) = session.update_for( session.acp_session_id.clone(), wire::SessionUpdate::StateUpdate(wire::StateUpdate::Running( wire::RunningStateUpdate::new(), @@ -455,10 +747,10 @@ impl AcpIntegration { let Some(update) = event_to_update(&event, message_ids.as_ref(), &mut part_kinds) else { return; }; - if let Err(error) = session - .client - .update(session.acp_session_id.clone(), update) - { + if let Err(error) = session.sink.update(wire::UpdateSessionNotification::new( + session.acp_session_id.clone(), + update, + )) { tracing::debug!(%error, "failed to queue ACP v2 session update"); } } @@ -556,6 +848,7 @@ where M: ModelAdapter, { factory: Option>>, + integration: AcpIntegration, name: String, version: String, } @@ -567,6 +860,7 @@ where fn default() -> Self { Self { factory: None, + integration: AcpIntegration::default(), name: "agentkit".into(), version: env!("CARGO_PKG_VERSION").into(), } @@ -598,6 +892,16 @@ where self } + /// Uses a host-visible unstable ACP v2 session coordinator. + /// + /// The headless runtime delegates binding, updates, cancellation, and + /// feature-gated injection handling to this same public value. + #[must_use] + pub fn integration(mut self, integration: AcpIntegration) -> Self { + self.integration = integration; + self + } + /// Sets the implementation name reported by `initialize`. #[must_use] pub fn name(mut self, name: impl Into) -> Self { @@ -618,7 +922,7 @@ where self.serve(agent_client_protocol::Stdio::new()).await } - /// Serves ACP v2 over a custom upstream SDK transport. + /// Serves ACP v2 over a custom SDK transport. pub async fn serve( self, transport: impl agent_client_protocol::ConnectTo + 'static, @@ -626,9 +930,14 @@ where let factory = self .factory .ok_or(AcpRuntimeError::MissingField("agent_factory"))?; - let state = Arc::new(RuntimeState::new(factory, self.name, self.version)); + let state = Arc::new(RuntimeState::new( + factory, + Arc::new(self.integration), + self.name, + self.version, + )); let (shutdown, mut shutdown_rx) = oneshot::channel(); - let connection = agent_client_protocol::Agent + let agent = agent_client_protocol::Agent .v2() .name(state.name.as_str()) .on_receive_request( @@ -707,7 +1016,32 @@ where } }, agent_client_protocol::on_receive_request!(), + ); + #[cfg(feature = "unstable-inject")] + let agent = agent + .on_receive_request( + { + let integration = Arc::clone(&state.integration); + async move |request: wire::InjectSessionRequest, responder, cx| { + let integration = Arc::clone(&integration); + cx.spawn(async move { + integration.handle_inject_request(request, responder).await + })?; + Ok(()) + } + }, + agent_client_protocol::on_receive_request!(), ) + .on_receive_request( + { + let integration = Arc::clone(&state.integration); + async move |request: wire::RevokeInjectSessionRequest, responder, _cx| { + responder.respond_with_result(integration.revoke_inject(request).await) + } + }, + agent_client_protocol::on_receive_request!(), + ); + let connection = agent .on_receive_notification( { let state = Arc::clone(&state); @@ -753,84 +1087,733 @@ where } } -struct SessionEntry { - commands: mpsc::UnboundedSender, - cancellation: CancellationController, - info: wire::SessionInfo, - active_prompt: Arc, - driving_prompt: Arc, - cancelled_prompt: Arc, - next_prompt_owner: AtomicU64, - closed: AtomicBool, - lifecycle: Arc>, - task: Mutex>>, - drain_task: Mutex>>, +/// Maximum number of pending unstable ACP v2 injections per session. +#[cfg(feature = "unstable-inject")] +pub const MAX_PENDING_INJECTIONS: usize = 64; +/// Maximum serialized content bytes retained by pending unstable injections. +#[cfg(feature = "unstable-inject")] +pub const MAX_PENDING_INJECTION_BYTES: usize = 256 * 1024; +/// Maximum accepted unstable injections tracked during one session lifetime. +/// +/// This caps every accepted injection ID retained for `already_delivered` +/// classification. At 4,096 compact IDs, lifetime tracking remains bounded +/// without evicting classifications before the session closes. +#[cfg(feature = "unstable-inject")] +pub const MAX_ACCEPTED_INJECTIONS: usize = 4_096; + +#[cfg(feature = "unstable-inject")] +struct PendingInject { + message_id: wire::MessageId, + content: Vec, + items: Vec, + bytes: usize, + commitment: InjectCommitment, } -enum SessionCommand { - Prompt { - request: wire::PromptRequest, - items: Vec, - prompt_cancellation: TurnCancellation, - cancellation_generation: u64, - owner: u64, - start: oneshot::Receiver<()>, - }, - Shutdown, +#[cfg(feature = "unstable-inject")] +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum InjectCommitment { + Reserved, + Committed, + Ready, } -struct RuntimeState -where - M: ModelAdapter, -{ - factory: Arc>, - integration: Arc, - sessions: AsyncMutex>>, - next_session: AtomicU64, - name: String, - version: String, +#[cfg(feature = "unstable-inject")] +#[derive(Default)] +struct InjectionState { + running: bool, + cancelled: bool, + at_boundary: bool, + pending: VecDeque, + delivering: Option, + pending_bytes: usize, + accepted_count: usize, + delivered: VecDeque, } -impl RuntimeState -where - M: ModelAdapter + Send + Sync + 'static, - M::Session: Send + 'static, -{ - fn new(factory: Arc>, name: String, version: String) -> Self { - Self { - factory, - integration: Arc::new(AcpIntegration::default()), - sessions: AsyncMutex::new(HashMap::new()), - next_session: AtomicU64::new(1), - name, - version, - } +#[cfg(feature = "unstable-inject")] +#[derive(Default)] +struct InjectionController { + state: Mutex, + changed: Notify, +} + +#[cfg(feature = "unstable-inject")] +enum BoundaryAction { + Wait, + Deliver(PendingInject), + Complete(AcpInjectionBoundary), +} + +#[cfg(feature = "unstable-inject")] +enum RevokeTransition { + Revoked, + WaitForDelivery, + AlreadyDelivered, + Unknown, +} + +#[cfg(feature = "unstable-inject")] +impl InjectionController { + fn start_turn(&self) { + let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + state.running = !state.cancelled; + state.at_boundary = false; } - fn initialize( - &self, - request: wire::InitializeRequest, - ) -> Result { - if request.protocol_version != wire::ProtocolVersion::V2 { - return Err(AcpRuntimeError::Unsupported( - "ACP v2 runtime requires protocol version 2".into(), + fn reset_cancellation(&self) { + self.state + .lock() + .unwrap_or_else(|error| error.into_inner()) + .cancelled = false; + } + + fn reserve(&self, pending: PendingInject) -> Result<(), agent_client_protocol::Error> { + let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + if !state.running { + return Err(sdk_v2_error(wire::Error::inject_no_running_turn())); + } + let reserved_count = state + .pending + .iter() + .filter(|pending| pending.commitment == InjectCommitment::Reserved) + .count(); + if state.accepted_count.saturating_add(reserved_count) >= MAX_ACCEPTED_INJECTIONS { + return Err(agent_client_protocol::Error::new( + -32000, + "session injection lifetime limit exceeded", + ) + .data(serde_json::json!({ + "reason": "lifetime_limit_exceeded", + "limit": MAX_ACCEPTED_INJECTIONS, + }))); + } + if state.pending.len() + usize::from(state.delivering.is_some()) >= MAX_PENDING_INJECTIONS + || state.pending_bytes.saturating_add(pending.bytes) > MAX_PENDING_INJECTION_BYTES + { + return Err(agent_client_protocol::Error::new( + -32602, + "pending session injection budget exceeded", )); } - Ok(wire::InitializeResponse::new( - wire::ProtocolVersion::V2, - wire::Implementation::new(self.name.clone(), self.version.clone()), - ) - .capabilities(headless_capabilities())) + state.pending_bytes += pending.bytes; + state.pending.push_back(pending); + Ok(()) } - async fn new_session( - self: &Arc, - request: wire::NewSessionRequest, - cx: ConnectionTo, - ) -> Result { - let sequence = self.next_session.fetch_add(1, Ordering::Relaxed); - let acp_session_id = wire::SessionId::new(format!("session-{sequence}")); - let agentkit_session_id = AgentkitSessionId::new(acp_session_id.to_string()); + fn commit(&self, message_id: &wire::MessageId) { + let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + let pending = state + .pending + .iter_mut() + .find(|pending| &pending.message_id == message_id) + .expect("reserved injection remains until response commitment"); + debug_assert_eq!(pending.commitment, InjectCommitment::Reserved); + pending.commitment = InjectCommitment::Committed; + state.accepted_count += 1; + } + + fn activate(&self, message_id: &wire::MessageId) { + let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + if let Some(pending) = state + .pending + .iter_mut() + .find(|pending| &pending.message_id == message_id) + { + debug_assert_eq!(pending.commitment, InjectCommitment::Committed); + pending.commitment = InjectCommitment::Ready; + } + drop(state); + self.changed.notify_waiters(); + } + + fn discard(&self, message_id: &wire::MessageId) { + let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + Self::remove_pending(&mut state, message_id); + drop(state); + self.changed.notify_waiters(); + } + + fn remove_pending(state: &mut InjectionState, message_id: &wire::MessageId) -> bool { + let Some(index) = state + .pending + .iter() + .position(|pending| &pending.message_id == message_id) + else { + return false; + }; + if let Some(pending) = state.pending.remove(index) { + state.pending_bytes = state.pending_bytes.saturating_sub(pending.bytes); + } + true + } + + fn boundary_action(&self, terminal: bool, delivered_any: bool) -> BoundaryAction { + let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + state.at_boundary = true; + if !state.running { + state.at_boundary = false; + return BoundaryAction::Complete(AcpInjectionBoundary::Stopped); + } + if state + .pending + .front() + .is_some_and(|pending| pending.commitment != InjectCommitment::Ready) + { + return BoundaryAction::Wait; + } + if let Some(pending) = state.pending.pop_front() { + state.delivering = Some(pending.message_id.clone()); + return BoundaryAction::Deliver(pending); + } + state.at_boundary = false; + if delivered_any { + BoundaryAction::Complete(AcpInjectionBoundary::Delivered) + } else if terminal { + state.running = false; + BoundaryAction::Complete(AcpInjectionBoundary::Finished) + } else { + BoundaryAction::Complete(AcpInjectionBoundary::Continue) + } + } + + fn finish_delivery(&self, pending: &PendingInject, delivered: bool) { + let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + debug_assert_eq!(state.delivering.as_ref(), Some(&pending.message_id)); + state.delivering = None; + state.pending_bytes = state.pending_bytes.saturating_sub(pending.bytes); + if delivered { + state.delivered.push_back(pending.message_id.clone()); + } + drop(state); + self.changed.notify_waiters(); + } + + fn revoke_transition(&self, message_id: &wire::MessageId) -> RevokeTransition { + let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + let ready = state.pending.iter().any(|pending| { + &pending.message_id == message_id && pending.commitment == InjectCommitment::Ready + }); + if ready && Self::remove_pending(&mut state, message_id) { + drop(state); + self.changed.notify_waiters(); + RevokeTransition::Revoked + } else if state.delivering.as_ref() == Some(message_id) + || state + .pending + .iter() + .any(|pending| &pending.message_id == message_id) + { + RevokeTransition::WaitForDelivery + } else if state.delivered.contains(message_id) { + RevokeTransition::AlreadyDelivered + } else { + RevokeTransition::Unknown + } + } + + fn cancel_turn(&self) { + let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + state.running = false; + state.cancelled = true; + state.at_boundary = false; + // Accepted steers survive cancellation and become deliverable at the + // next valid boundary, potentially in the next prompt. + drop(state); + self.changed.notify_waiters(); + } + + fn stop_turn(&self) { + let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + state.running = false; + state.at_boundary = false; + drop(state); + self.changed.notify_waiters(); + } + + fn close_session(&self) { + let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + state.running = false; + state.cancelled = true; + state.at_boundary = false; + let discarded_bytes = state + .pending + .iter() + .filter(|pending| pending.commitment == InjectCommitment::Ready) + .map(|pending| pending.bytes) + .sum::(); + state + .pending + .retain(|pending| pending.commitment != InjectCommitment::Ready); + state.pending_bytes = state.pending_bytes.saturating_sub(discarded_bytes); + drop(state); + self.changed.notify_waiters(); + } +} + +#[cfg(feature = "unstable-inject")] +struct ReservedInject { + session: Arc, + message_id: Option, +} + +#[cfg(feature = "unstable-inject")] +impl ReservedInject { + fn response(&self) -> wire::InjectSessionResponse { + wire::InjectSessionResponse::new( + self.message_id + .as_ref() + .expect("reserved injection has a message id") + .clone(), + ) + } + + fn message_id(&self) -> &wire::MessageId { + self.message_id + .as_ref() + .expect("reserved injection has a message id") + } + + fn commit(&mut self) { + let message_id = self.message_id(); + self.session.injection.commit(message_id); + } + + fn activate(mut self) { + let message_id = self + .message_id + .take() + .expect("reserved injection has a message id"); + let lifecycle = self + .session + .lifecycle + .lock() + .unwrap_or_else(|error| error.into_inner()); + if self.session.closed.load(Ordering::Acquire) { + self.session.injection.discard(&message_id); + } else { + self.session.injection.activate(&message_id); + } + drop(lifecycle); + } + + fn discard(mut self) { + let message_id = self + .message_id + .take() + .expect("reserved injection has a message id"); + self.session.injection.discard(&message_id); + } +} + +#[cfg(feature = "unstable-inject")] +impl Drop for ReservedInject { + fn drop(&mut self) { + if let Some(message_id) = self.message_id.take() { + self.session.injection.discard(&message_id); + } + } +} + +#[cfg(feature = "unstable-inject")] +struct InjectRequestPermit { + integration: AcpIntegration, + sequence: Option, +} + +#[cfg(feature = "unstable-inject")] +impl Drop for InjectRequestPermit { + fn drop(&mut self) { + if let Some(sequence) = self.sequence.take() { + self.integration.finish_inject_request(sequence); + } + } +} + +/// Opaque reservation for one globally ordered unstable ACP v2 inject request. +/// +/// The handle owns the SDK responder so the response cannot be committed +/// outside the session lifecycle lock or against a different cancellation +/// token. Dropping it discards the reservation and advances request ordering. +#[cfg(feature = "unstable-inject")] +#[must_use = "the reserved inject request must be responded to or dropped"] +pub struct AcpInjectRequest { + reserved: Option, + responder: Option>, + cancellation: agent_client_protocol::RequestCancellation, + _permit: InjectRequestPermit, +} + +#[cfg(feature = "unstable-inject")] +impl AcpInjectRequest { + /// Returns the response that will be committed by [`respond_tracked`](Self::respond_tracked). + #[must_use] + pub fn response(&self) -> wire::InjectSessionResponse { + self.reserved + .as_ref() + .expect("inject request has a reservation") + .response() + } + + /// Commits the response under the session lifecycle lock. + /// + /// This checks close and request cancellation immediately before calling the + /// SDK's tracked responder. `Ok(None)` means the coordinator sent a close or + /// cancellation error instead. `Ok(Some(_))` returns an acceptance handle + /// bound to the resulting [`agent_client_protocol::ResponseReceipt`]. + pub fn respond_tracked( + mut self, + ) -> Result, agent_client_protocol::Error> { + let session = Arc::clone( + &self + .reserved + .as_ref() + .expect("inject request has a reservation") + .session, + ); + let lifecycle = session + .lifecycle + .lock() + .unwrap_or_else(|error| error.into_inner()); + let rejection = if session.closed.load(Ordering::Acquire) { + Some(session_not_found_error(&session.acp_session_id)) + } else if self.cancellation.is_cancelled() { + Some(agent_client_protocol::Error::request_cancelled()) + } else { + None + }; + if let Some(error) = rejection { + drop(lifecycle); + self.reserved.take().expect("inject reservation").discard(); + self.responder + .take() + .expect("inject request has a responder") + .respond_with_result(Err(error))?; + return Ok(None); + } + + let response = self.response(); + let receipt = match self + .responder + .take() + .expect("inject request has a responder") + .respond_tracked(response) + { + Ok(receipt) => receipt, + Err(error) => { + drop(lifecycle); + self.reserved.take().expect("inject reservation").discard(); + return Err(error); + } + }; + let mut reserved = self.reserved.take().expect("inject reservation"); + reserved.commit(); + drop(lifecycle); + Ok(Some(AcpInjectAcceptance { + reserved: Some(reserved), + receipt: Some(receipt), + })) + } +} + +/// Opaque response-committed unstable ACP v2 injection acceptance. +/// +/// This handle owns both the reservation and the SDK response receipt. Dropping +/// it before activation discards the reservation. +#[cfg(feature = "unstable-inject")] +#[must_use = "the response receipt must be activated or the injection is discarded"] +pub struct AcpInjectAcceptance { + reserved: Option, + receipt: Option, +} + +#[cfg(feature = "unstable-inject")] +impl AcpInjectAcceptance { + /// Returns the agent-owned message ID reserved for this acceptance. + #[must_use] + pub fn message_id(&self) -> &wire::MessageId { + self.reserved + .as_ref() + .expect("inject acceptance has a reservation") + .message_id() + } + + /// Waits for the SDK response receipt, then makes the steer deliverable. + /// + /// Batch safety requires calling this only after the receive callback has + /// returned, normally in a spawned task. Receipt failure or future + /// cancellation drops and discards the reservation. + pub async fn activate_after_response(mut self) -> Result<(), agent_client_protocol::Error> { + self.receipt + .take() + .expect("inject acceptance has a response receipt") + .await?; + self.reserved + .take() + .expect("inject acceptance has a reservation") + .activate(); + Ok(()) + } +} + +#[cfg(feature = "unstable-inject")] +impl AcpIntegration { + async fn wait_for_inject_request(&self, sequence: u64) { + loop { + let changed = self.inject_requests.changed.notified(); + tokio::pin!(changed); + changed.as_mut().enable(); + if self.inject_requests.current.load(Ordering::Acquire) == sequence { + return; + } + changed.await; + } + } + + fn finish_inject_request(&self, sequence: u64) { + let advanced = self.inject_requests.current.compare_exchange( + sequence, + sequence + 1, + Ordering::AcqRel, + Ordering::Acquire, + ); + debug_assert_eq!(advanced, Ok(sequence)); + self.inject_requests.changed.notify_waiters(); + } + + fn reserve_inject( + &self, + request: wire::InjectSessionRequest, + ) -> Result { + if !matches!(request.mode, wire::SessionInjectMode::Steer) { + return Err(agent_client_protocol::Error::new( + -32602, + "unsupported session injection mode", + )); + } + let items = content_blocks_to_items(&request.content).map_err(crate::sdk_error)?; + let bytes = serde_json::to_vec(&request.content) + .map_err(|error| agent_client_protocol::Error::new(-32603, error.to_string()))? + .len(); + let session = self + .session(&request.session_id) + .map_err(|_| session_not_found_error(&request.session_id))?; + let _lifecycle = session + .lifecycle + .lock() + .unwrap_or_else(|error| error.into_inner()); + if session.closed.load(Ordering::Acquire) { + return Err(session_not_found_error(&request.session_id)); + } + let message_id = self + .next_user_message_id(&request.session_id) + .map_err(crate::sdk_error)?; + session.injection.reserve(PendingInject { + message_id: message_id.clone(), + content: request.content, + items, + bytes, + commitment: InjectCommitment::Reserved, + })?; + drop(_lifecycle); + Ok(ReservedInject { + session, + message_id: Some(message_id), + }) + } + + /// Reserves one globally ordered unstable ACP v2 inject request. + /// + /// The SDK responder is intentionally consumed here: its cancellation token + /// and tracked response must remain paired with this reservation. This method + /// responds to pre-reservation cancellation or validation errors itself and + /// returns `Ok(None)`. A returned handle owns FIFO ordering and Drop cleanup. + pub async fn reserve_inject_request( + &self, + request: wire::InjectSessionRequest, + responder: agent_client_protocol::Responder, + ) -> Result, agent_client_protocol::Error> { + let sequence = self.inject_requests.next.fetch_add(1, Ordering::Relaxed); + self.wait_for_inject_request(sequence).await; + let permit = InjectRequestPermit { + integration: self.clone(), + sequence: Some(sequence), + }; + let cancellation = responder.cancellation(); + // Return the receive callback first so a queued cancellation can commit + // before reservation or response commitment. + tokio::task::yield_now().await; + if cancellation.is_cancelled() { + responder + .respond_with_result(Err(agent_client_protocol::Error::request_cancelled()))?; + drop(permit); + return Ok(None); + } + let reserved = match self.reserve_inject(request) { + Ok(reserved) => reserved, + Err(error) => { + responder.respond_with_result(Err(error))?; + drop(permit); + return Ok(None); + } + }; + Ok(Some(AcpInjectRequest { + reserved: Some(reserved), + responder: Some(responder), + cancellation, + _permit: permit, + })) + } + + /// Handles one unstable ACP v2 `session/inject` request end to end. + /// + /// This convenience path uses [`reserve_inject_request`](Self::reserve_inject_request), + /// [`AcpInjectRequest::respond_tracked`], and + /// [`AcpInjectAcceptance::activate_after_response`]. Hosts that need staged + /// control can call those same methods without reimplementing the race. + pub async fn handle_inject_request( + &self, + request: wire::InjectSessionRequest, + responder: agent_client_protocol::Responder, + ) -> Result<(), agent_client_protocol::Error> { + let Some(request) = self.reserve_inject_request(request, responder).await? else { + return Ok(()); + }; + let Some(acceptance) = request.respond_tracked()? else { + return Ok(()); + }; + let _receipt_task = tokio::spawn(async move { + if let Err(error) = acceptance.activate_after_response().await { + tracing::debug!(%error, "ACP v2 inject response was not accepted"); + } + }); + Ok(()) + } + + /// Handles the state transition for unstable ACP v2 `session/revoke_inject`. + pub async fn revoke_inject( + &self, + request: wire::RevokeInjectSessionRequest, + ) -> Result { + let session = self + .session(&request.session_id) + .map_err(|_| session_not_found_error(&request.session_id))?; + loop { + let changed = session.injection.changed.notified(); + tokio::pin!(changed); + changed.as_mut().enable(); + let _lifecycle = session + .lifecycle + .lock() + .unwrap_or_else(|error| error.into_inner()); + if session.closed.load(Ordering::Acquire) { + return Err(session_not_found_error(&request.session_id)); + } + match session.injection.revoke_transition(&request.message_id) { + RevokeTransition::Revoked => { + return Ok(wire::RevokeInjectSessionResponse::new()); + } + RevokeTransition::WaitForDelivery => drop(_lifecycle), + RevokeTransition::AlreadyDelivered => { + return Err(sdk_v2_error(wire::Error::inject_already_delivered( + request.message_id, + ))); + } + RevokeTransition::Unknown => { + return Err(sdk_v2_error(wire::Error::inject_unknown_message_id( + request.message_id, + ))); + } + } + changed.await; + } + } +} + +struct SessionEntry { + commands: mpsc::UnboundedSender, + session: AcpSessionHandle, + cancellation: CancellationController, + info: wire::SessionInfo, + active_prompt: Arc, + driving_prompt: Arc, + cancelled_prompt: Arc, + next_prompt_owner: AtomicU64, + closed: AtomicBool, + lifecycle: Arc>, + task: Mutex>>, + drain_task: Mutex>>, +} + +enum SessionCommand { + Prompt { + request: wire::PromptRequest, + items: Vec, + prompt_cancellation: TurnCancellation, + cancellation_generation: u64, + owner: u64, + start: oneshot::Receiver<()>, + }, + Shutdown, +} + +struct RuntimeState +where + M: ModelAdapter, +{ + factory: Arc>, + integration: Arc, + sessions: AsyncMutex>>, + next_session: AtomicU64, + name: String, + version: String, +} + +impl RuntimeState +where + M: ModelAdapter + Send + Sync + 'static, + M::Session: Send + 'static, +{ + fn new( + factory: Arc>, + integration: Arc, + name: String, + version: String, + ) -> Self { + Self { + factory, + integration, + sessions: AsyncMutex::new(HashMap::new()), + next_session: AtomicU64::new(1), + name, + version, + } + } + + fn initialize( + &self, + request: wire::InitializeRequest, + ) -> Result { + if request.protocol_version != wire::ProtocolVersion::V2 { + return Err(AcpRuntimeError::Unsupported( + "ACP v2 runtime requires protocol version 2".into(), + )); + } + Ok(wire::InitializeResponse::new( + wire::ProtocolVersion::V2, + wire::Implementation::new(self.name.clone(), self.version.clone()), + ) + .capabilities(agent_capabilities())) + } + + async fn new_session( + self: &Arc, + request: wire::NewSessionRequest, + cx: V2ConnectionTo, + ) -> Result { + let sequence = self.next_session.fetch_add(1, Ordering::Relaxed); + let acp_session_id = wire::SessionId::new(format!("session-{sequence}")); + let agentkit_session_id = AgentkitSessionId::new(acp_session_id.to_string()); let cancellation = CancellationController::new(); let (client, client_messages) = ClientHandle::channel(); let info = wire::SessionInfo::new(acp_session_id.clone(), request.cwd.clone()) @@ -847,24 +1830,27 @@ where let active_prompt = Arc::new(AtomicU64::new(0)); let driving_prompt = Arc::new(AtomicU64::new(0)); let cancelled_prompt = Arc::new(AtomicU64::new(0)); - let lifecycle = Arc::new(Mutex::new(())); - self.integration.bind( - acp_session_id.clone(), - agentkit_session_id.clone(), - client.clone(), + let session = self.integration.bind_session( + AcpSessionBinding::new( + acp_session_id.clone(), + agentkit_session_id.clone(), + client.clone(), + ) + .cancellation(cancellation.clone()), )?; + let lifecycle = Arc::clone(&session.session.lifecycle); if let Err(error) = self.integration.install_prompt_state( &acp_session_id, Arc::clone(&active_prompt), Arc::clone(&lifecycle), ) { - let _ = self.integration.unbind(&acp_session_id); + let _ = self.integration.unbind_session(&acp_session_id); return Err(error); } let drain_task = tokio::spawn(drain_client_messages(client_messages, cx)); let ctx = AcpAgentFactoryContext { acp_session_id: acp_session_id.clone(), - agentkit_session_id, + agentkit_session_id: agentkit_session_id.clone(), cwd: request.cwd.into_inner(), additional_directories: request .additional_directories @@ -878,7 +1864,7 @@ where let driver = match self.factory.start(ctx).await { Ok(driver) => driver, Err(error) => { - let _ = self.integration.unbind(&acp_session_id); + let _ = self.integration.unbind_session(&acp_session_id); drain_task.abort(); let _ = drain_task.await; return Err(error); @@ -891,11 +1877,11 @@ where let worker_cancelled_prompt = Arc::clone(&cancelled_prompt); let worker_lifecycle = Arc::clone(&lifecycle); let integration = Arc::clone(&self.integration); - let worker_session_id = acp_session_id.clone(); + let worker_session = session.clone(); let worker_cancellation = cancellation.handle(); let task = tokio::spawn(async move { session_worker( - worker_session_id, + worker_session, driver, client, integration, @@ -910,6 +1896,7 @@ where }); let entry = Arc::new(SessionEntry { commands, + session, cancellation, info, active_prompt, @@ -941,7 +1928,7 @@ where let mut infos = sessions .values() .filter(|entry| { - !entry.closed.load(Ordering::Acquire) + !entry.session.is_closed() && request .cwd .as_ref() @@ -973,7 +1960,7 @@ where .get(&request.session_id) .cloned() .ok_or_else(|| AcpRuntimeError::SessionNotFound(request.session_id.to_string()))?; - if entry.closed.load(Ordering::Acquire) || entry.info.cwd != request.cwd { + if entry.session.is_closed() || entry.info.cwd != request.cwd { return Err(AcpRuntimeError::SessionNotFound( request.session_id.to_string(), )); @@ -990,7 +1977,7 @@ where &self, request: wire::PromptRequest, ) -> Result, AcpRuntimeError> { - let items = prompt_to_items(&request)?; + let items = self.integration.prompt_to_items(&request)?; let entry = self .sessions .lock() @@ -1001,10 +1988,12 @@ where let (start_tx, start_rx) = oneshot::channel(); { let _lifecycle = entry + .session + .session .lifecycle .lock() .unwrap_or_else(|error| error.into_inner()); - if entry.closed.load(Ordering::Acquire) { + if entry.session.is_closed() { return Err(AcpRuntimeError::SessionNotFound( request.session_id.to_string(), )); @@ -1023,6 +2012,8 @@ where )); } let prompt_cancellation = entry.cancellation.handle().checkpoint(); + #[cfg(feature = "unstable-inject")] + entry.session.prepare_injection_turn(); let cancellation_generation = prompt_cancellation.generation(); if entry .commands @@ -1061,12 +2052,16 @@ where if !entry.closed.load(Ordering::Acquire) { let driving = entry.driving_prompt.load(Ordering::Acquire); if driving != 0 { + #[cfg(feature = "unstable-inject")] + entry.session.session.injection.cancel_turn(); entry.cancellation.interrupt(); } else { let queued = entry.active_prompt.load(Ordering::Acquire); if queued != 0 { entry.cancelled_prompt.store(queued, Ordering::Release); } else { + #[cfg(feature = "unstable-inject")] + entry.session.session.injection.cancel_turn(); entry.cancellation.interrupt(); } } @@ -1086,7 +2081,7 @@ where .remove(&request.session_id) .ok_or_else(|| AcpRuntimeError::SessionNotFound(request.session_id.to_string()))?; stop_session(Arc::clone(&entry)).await; - self.integration.unbind(&request.session_id)?; + self.integration.unbind_session(&request.session_id)?; stop_client(entry).await; Ok(wire::CloseSessionResponse::new()) } @@ -1103,7 +2098,7 @@ where if let Some(task) = take_task(&entry.task) { session_tasks.push(task); } - let _ = self.integration.unbind(session_id); + let _ = self.integration.unbind_session(session_id); } join_tasks_until(deadline, session_tasks).await; @@ -1117,12 +2112,8 @@ where } fn signal_session_stop(entry: &Arc) { - let _lifecycle = entry - .lifecycle - .lock() - .unwrap_or_else(|error| error.into_inner()); entry.closed.store(true, Ordering::Release); - entry.cancellation.interrupt(); + entry.session.close(); let _ = entry.commands.send(SessionCommand::Shutdown); } @@ -1164,7 +2155,9 @@ async fn stop_session(entry: Arc) { } async fn stop_client(entry: Arc) { - if let Some(task) = take_task(&entry.drain_task) { + let task = take_task(&entry.drain_task); + drop(entry); + if let Some(task) = task { let _ = task.await; } } @@ -1198,10 +2191,10 @@ fn fail_accepted_prompt( ) { let _lifecycle = lifecycle.lock().unwrap_or_else(|error| error.into_inner()); if prompt_began { - integration.finish_prompt(session_id, owner); + integration.finish_prompt_owner(session_id, owner); } clear_prompt_tracking(active_prompt, driving_prompt, cancelled_prompt, owner); - if let Err(error) = client.update( + if let Err(error) = client.update_for( session_id.clone(), wire::SessionUpdate::StateUpdate(wire::StateUpdate::Idle( wire::IdleStateUpdate::new().stop_reason(error_stop_reason()), @@ -1212,7 +2205,7 @@ fn fail_accepted_prompt( } async fn session_worker( - session_id: wire::SessionId, + session: AcpSessionHandle, mut driver: agentkit_loop::LoopDriver, client: ClientHandle, integration: Arc, @@ -1225,6 +2218,7 @@ async fn session_worker( ) where S: ModelSession + Send + 'static, { + let session_id = session.acp_session_id().clone(); enum SessionAction { Command(Option), LoopUpdate(Result<(), LoopError>), @@ -1262,6 +2256,8 @@ async fn session_worker( &session_id, &cancellation, cancellation.generation(), + #[cfg(feature = "unstable-inject")] + &session, ) .await; if let Err(error) = client.flush().await { @@ -1299,11 +2295,17 @@ async fn session_worker( } }; if !should_start { + #[cfg(feature = "unstable-inject")] + session.stop_injection_turn(); continue; } + #[cfg(feature = "unstable-inject")] + session.start_injection_turn(); if let Err(error) = driver.submit_input(items) { tracing::debug!(%error, owner, "failed to submit accepted ACP v2 prompt"); + #[cfg(feature = "unstable-inject")] + session.stop_injection_turn(); fail_accepted_prompt( &client, &integration, @@ -1318,10 +2320,12 @@ async fn session_worker( continue; } let user_message_id = - match integration.begin_prompt(&session_id, owner, prompt_cancellation) { + match integration.begin_prompt_owner(&session_id, owner, prompt_cancellation) { Ok(message_id) => message_id, Err(error) => { tracing::debug!(%error, owner, "failed to begin accepted ACP v2 prompt"); + #[cfg(feature = "unstable-inject")] + session.stop_injection_turn(); fail_accepted_prompt( &client, &integration, @@ -1337,13 +2341,15 @@ async fn session_worker( } }; - if let Err(error) = client.update( + if let Err(error) = client.update_for( session_id.clone(), wire::SessionUpdate::UserMessage( wire::UserMessage::new(user_message_id).content(request.prompt), ), ) { tracing::debug!(%error, owner, "failed to publish accepted ACP v2 prompt"); + #[cfg(feature = "unstable-inject")] + session.stop_injection_turn(); fail_accepted_prompt( &client, &integration, @@ -1364,13 +2370,17 @@ async fn session_worker( &session_id, &cancellation, cancellation_generation, + #[cfg(feature = "unstable-inject")] + &session, ) .await; { let _lifecycle = lifecycle.lock().unwrap_or_else(|error| error.into_inner()); - integration.finish_prompt(&session_id, owner); + integration.finish_prompt_owner(&session_id, owner); clear_prompt_tracking(&active_prompt, &driving_prompt, &cancelled_prompt, owner); } + #[cfg(feature = "unstable-inject")] + session.stop_injection_turn(); if let Err(error) = client.flush().await { tracing::debug!(%error, "failed to flush ACP v2 output"); } @@ -1378,49 +2388,185 @@ async fn session_worker( } } +/// Outcome of an unstable ACP v2 injection boundary. +#[cfg(feature = "unstable-inject")] +#[derive(Debug, Eq, PartialEq)] +pub enum AcpInjectionBoundary { + /// No steer was delivered; continue the current non-terminal loop. + Continue, + /// At least one steer was delivered and the loop must continue. + Delivered, + /// The terminal boundary had no steer and the turn can finish. + Finished, + /// Cancellation or close stopped delivery. + Stopped, +} + +#[cfg(feature = "unstable-inject")] +impl AcpSessionHandle { + /// Delivers ready unstable ACP v2 steers at one safe loop boundary. + /// + /// `terminal` must be true for terminal model/input boundaries and false + /// after tool results. The returned outcome tells the host whether to + /// continue, finish, or report cancellation. + pub async fn handle_injection_boundary( + &self, + driver: &mut agentkit_loop::LoopDriver, + terminal: bool, + ) -> Result + where + S: ModelSession + Send + 'static, + { + let injection = &self.session.injection; + let mut delivered_any = false; + loop { + let changed = injection.changed.notified(); + tokio::pin!(changed); + changed.as_mut().enable(); + match injection.boundary_action(terminal, delivered_any) { + BoundaryAction::Wait => changed.await, + BoundaryAction::Complete(outcome) => return Ok(outcome), + BoundaryAction::Deliver(pending) => { + if let Err(error) = driver + .submit_input(pending.items.clone()) + .map_err(|error| AcpRuntimeError::Loop(error.to_string())) + { + injection.finish_delivery(&pending, false); + return Err(error); + } + let result = self + .session + .sink + .update_acknowledged(wire::UpdateSessionNotification::new( + self.session.acp_session_id.clone(), + wire::SessionUpdate::UserMessage( + wire::UserMessage::new(pending.message_id.clone()) + .content(pending.content.clone()), + ), + )) + .await; + injection.finish_delivery(&pending, result.is_ok()); + result?; + delivered_any = true; + } + } + } + } +} + async fn drive_prompt( driver: &mut agentkit_loop::LoopDriver, integration: &AcpIntegration, session_id: &wire::SessionId, cancellation: &CancellationHandle, generation: u64, -) where + #[cfg(feature = "unstable-inject")] session: &AcpSessionHandle, +) -> wire::StopReason +where S: ModelSession + Send + 'static, { loop { - // Cancellation is installed on the driver and its model/tool work. Keep - // polling the driver so it can close interrupted tool calls and leave a - // resumable transcript before the session becomes idle. - match driver.next().await { - Ok(LoopStep::Finished(result)) => { - if result.finish_reason == FinishReason::ToolCall { - continue; - } - return; + let step = match driver.next().await { + Ok(step) => step, + Err(error) => { + tracing::debug!(%error, "ACP v2 agent loop failed"); + #[cfg(feature = "unstable-inject")] + session.stop_injection_turn(); + return if cancellation.is_cancelled_since(generation) { + wire::StopReason::Cancelled + } else { + error_stop_reason() + }; + } + }; + if cancellation.is_cancelled_since(generation) { + #[cfg(feature = "unstable-inject")] + session.stop_injection_turn(); + return wire::StopReason::Cancelled; + } + match step { + LoopStep::Finished(result) => { + if result.finish_reason == FinishReason::ToolCall { + continue; + } + #[cfg(feature = "unstable-inject")] + match session.handle_injection_boundary(driver, true).await { + Ok(AcpInjectionBoundary::Delivered | AcpInjectionBoundary::Continue) => { + continue; + } + Ok(AcpInjectionBoundary::Stopped) => return wire::StopReason::Cancelled, + Ok(AcpInjectionBoundary::Finished) => { + return finish_reason_to_stop_reason(&result.finish_reason); + } + Err(error) => { + tracing::debug!(%error, "failed to deliver ACP v2 injected message"); + session.stop_injection_turn(); + return error_stop_reason(); + } + } + #[cfg(not(feature = "unstable-inject"))] + return finish_reason_to_stop_reason(&result.finish_reason); + } + LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_)) => { + #[cfg(feature = "unstable-inject")] + match session.handle_injection_boundary(driver, true).await { + Ok(AcpInjectionBoundary::Delivered | AcpInjectionBoundary::Continue) => { + continue; + } + Ok(AcpInjectionBoundary::Stopped) => return wire::StopReason::Cancelled, + Ok(AcpInjectionBoundary::Finished) => return wire::StopReason::EndTurn, + Err(error) => { + tracing::debug!(%error, "failed to deliver ACP v2 injected message"); + session.stop_injection_turn(); + return error_stop_reason(); + } + } + #[cfg(not(feature = "unstable-inject"))] + return wire::StopReason::EndTurn; + } + LoopStep::Interrupt(LoopInterrupt::AfterToolResult(_)) => { + #[cfg(feature = "unstable-inject")] + match session.handle_injection_boundary(driver, false).await { + Ok(AcpInjectionBoundary::Stopped) => return wire::StopReason::Cancelled, + Err(error) => { + tracing::debug!(%error, "failed to deliver ACP v2 injected message"); + session.stop_injection_turn(); + return error_stop_reason(); + } + _ => {} + } } - Ok(LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_))) => return, - Ok(LoopStep::Interrupt(LoopInterrupt::AfterToolResult(_))) => continue, - Ok(LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(_))) => { + LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(_)) => { integration.mark_unsupported_approval(session_id, cancellation.clone(), generation); if let Err(error) = driver.cancel_pending_approvals().await { integration.clear_unsupported_approval(session_id); tracing::debug!(%error, "failed to cancel unsupported ACP v2 approval"); + #[cfg(feature = "unstable-inject")] + session.stop_injection_turn(); + return error_stop_reason(); } - return; - } - Err(error) => { - tracing::debug!(%error, "ACP v2 agent loop failed"); - return; + #[cfg(not(feature = "unstable-inject"))] + return error_stop_reason(); } } } } fn prompt_to_items(request: &wire::PromptRequest) -> Result, AcpRuntimeError> { + content_blocks_to_items(&request.prompt) +} + +/// Converts unstable ACP v2 content blocks into agentkit input items. +/// +/// The conversion preserves text, images, audio, resource links, and embedded +/// resources using the same path as prompts and session injection. +pub fn content_blocks_to_items( + content: &[wire::ContentBlock], +) -> Result, AcpRuntimeError> { let mut user_parts = Vec::new(); let mut context_items = Vec::new(); - for block in &request.prompt { + for block in content { match block { wire::ContentBlock::Text(text) => { user_parts.push(Part::Text(TextPart::new(text.text.clone()))); @@ -1757,15 +2903,31 @@ fn finish_reason_to_stop_reason(reason: &FinishReason) -> wire::StopReason { } } -fn headless_capabilities() -> wire::AgentCapabilities { - wire::AgentCapabilities::new().session( - wire::SessionCapabilities::new().prompt( - wire::PromptCapabilities::new() - .image(wire::PromptImageCapabilities::new()) - .audio(wire::PromptAudioCapabilities::new()) - .embedded_context(wire::PromptEmbeddedContextCapabilities::new()), - ), - ) +/// Constructs the exact unstable ACP v2 injection capability implemented here. +/// +/// This function is available only with `unstable-inject`, matching the request +/// handlers and turn-boundary APIs it advertises. +#[cfg(feature = "unstable-inject")] +pub fn session_inject_capabilities() -> wire::SessionInjectCapabilities { + wire::SessionInjectCapabilities::new(vec![wire::SessionInjectMode::Steer]) + .steer_in_stream(vec![wire::SessionInjectSteerInStream::Finish]) +} + +/// Constructs the honest ACP v2 capabilities supported by this build. +/// +/// Injection is advertised only when `unstable-inject` is enabled, and uses +/// [`session_inject_capabilities`] so hosts and [`AcpHeadlessRuntime`] report +/// the same behavior. +pub fn agent_capabilities() -> wire::AgentCapabilities { + let session = wire::SessionCapabilities::new().prompt( + wire::PromptCapabilities::new() + .image(wire::PromptImageCapabilities::new()) + .audio(wire::PromptAudioCapabilities::new()) + .embedded_context(wire::PromptEmbeddedContextCapabilities::new()), + ); + #[cfg(feature = "unstable-inject")] + let session = session.inject(session_inject_capabilities()); + wire::AgentCapabilities::new().session(session) } #[cfg(test)] @@ -1776,9 +2938,9 @@ mod tests { use std::time::Duration; use agent_client_protocol::Channel; - use agentkit_core::{ - ItemKind, ToolCallId, ToolOutput, ToolResultPart, TurnCancellation, TurnId, - }; + #[cfg(feature = "unstable-inject")] + use agent_client_protocol::{RawJsonRpcMessage, TransportBatch, TransportFrame}; + use agentkit_core::{ItemKind, ToolCallId, ToolOutput, ToolResultPart, TurnCancellation}; use agentkit_integration_tests::mock_model::{MockAdapter, TurnScript}; use agentkit_integration_tests::mock_tool::BlockingTool; use agentkit_loop::{ @@ -1786,11 +2948,49 @@ mod tests { TurnRequest, TurnResult, }; use agentkit_task_manager::{AsyncTaskManager, RoutingDecision}; + #[cfg(feature = "unstable-inject")] + use agentkit_tools_core::{ + ApprovalReason, ApprovalRequest, PermissionChecker, PermissionDecision, PermissionRequest, + }; use agentkit_tools_core::{ Tool, ToolContext, ToolError, ToolRegistry, ToolRequest, ToolResult, ToolSpec, }; + #[cfg(feature = "unstable-inject")] + use futures_util::StreamExt as _; use tokio::sync::Notify; + #[derive(Clone, Default)] + struct RecordingSink { + updates: Arc>>, + acknowledged: Arc, + flushes: Arc, + } + + #[async_trait] + impl AcpSessionUpdateSink for RecordingSink { + fn update( + &self, + notification: wire::UpdateSessionNotification, + ) -> Result<(), AcpRuntimeError> { + self.updates.lock().unwrap().push(notification); + Ok(()) + } + + async fn update_acknowledged( + &self, + notification: wire::UpdateSessionNotification, + ) -> Result<(), AcpRuntimeError> { + self.updates.lock().unwrap().push(notification); + self.acknowledged.fetch_add(1, Ordering::Release); + Ok(()) + } + + async fn flush(&self) -> Result<(), AcpRuntimeError> { + self.flushes.fetch_add(1, Ordering::Release); + Ok(()) + } + } + #[derive(Clone)] struct TestFactory { adapter: A, @@ -1908,6 +3108,298 @@ mod tests { } } + #[cfg(feature = "unstable-inject")] + struct ApprovalPermissionRequest { + metadata: MetadataMap, + } + + #[cfg(feature = "unstable-inject")] + impl PermissionRequest for ApprovalPermissionRequest { + fn kind(&self) -> &'static str { + "custom.approval-test" + } + + fn summary(&self) -> String { + "approve test tool".into() + } + + fn metadata(&self) -> &MetadataMap { + &self.metadata + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + } + + #[cfg(feature = "unstable-inject")] + #[derive(Clone)] + struct ApprovalTool { + spec: ToolSpec, + } + + #[cfg(feature = "unstable-inject")] + impl ApprovalTool { + fn new() -> Self { + Self { + spec: ToolSpec::new( + "approval_tool", + "requires approval", + json!({ "type": "object" }), + ), + } + } + } + + #[cfg(feature = "unstable-inject")] + #[async_trait] + impl Tool for ApprovalTool { + fn spec(&self) -> &ToolSpec { + &self.spec + } + + fn proposed_requests( + &self, + _request: &ToolRequest, + ) -> Result>, ToolError> { + Ok(vec![Box::new(ApprovalPermissionRequest { + metadata: MetadataMap::new(), + })]) + } + + async fn invoke( + &self, + request: ToolRequest, + _ctx: &mut ToolContext<'_>, + ) -> Result { + Ok(ToolResult::new(ToolResultPart::success( + request.call_id, + ToolOutput::text("approved"), + ))) + } + } + + #[cfg(feature = "unstable-inject")] + #[derive(Clone, Copy)] + struct RequireApproval; + + #[cfg(feature = "unstable-inject")] + impl PermissionChecker for RequireApproval { + fn evaluate(&self, request: &dyn PermissionRequest) -> PermissionDecision { + PermissionDecision::RequireApproval(ApprovalRequest::new( + "approval-test", + request.kind(), + ApprovalReason::PolicyRequiresConfirmation, + request.summary(), + )) + } + } + + #[cfg(feature = "unstable-inject")] + #[derive(Clone)] + struct ApprovalAdapter { + permits: Arc, + next_turn: Arc, + } + + #[cfg(feature = "unstable-inject")] + impl ApprovalAdapter { + fn new() -> Self { + Self { + permits: Arc::new(tokio::sync::Semaphore::new(0)), + next_turn: Arc::new(AtomicUsize::new(0)), + } + } + + fn release(&self) { + self.permits.add_permits(1); + } + } + + #[cfg(feature = "unstable-inject")] + struct ApprovalSession { + permits: Arc, + next_turn: Arc, + } + + #[cfg(feature = "unstable-inject")] + #[async_trait] + impl ModelAdapter for ApprovalAdapter { + type Session = ApprovalSession; + + async fn start_session(&self, _config: SessionConfig) -> Result { + Ok(ApprovalSession { + permits: Arc::clone(&self.permits), + next_turn: Arc::clone(&self.next_turn), + }) + } + } + + #[cfg(feature = "unstable-inject")] + #[async_trait] + impl ModelSession for ApprovalSession { + type Turn = GatedTurn; + + async fn begin_turn( + &mut self, + _request: TurnRequest, + _cancellation: Option, + ) -> Result { + let turn = self.next_turn.fetch_add(1, Ordering::AcqRel); + let events = if turn == 0 { + let call = + ToolCallPart::new(ToolCallId::new("approval-call"), "approval_tool", json!({})); + VecDeque::from([ + ModelTurnEvent::ToolCall(call.clone()), + ModelTurnEvent::Finished(ModelTurnResult { + model: None, + response_id: None, + finish_reason: FinishReason::ToolCall, + output_items: vec![Item::new( + ItemKind::Assistant, + vec![Part::ToolCall(call)], + )], + usage: None, + metadata: MetadataMap::new(), + }), + ]) + } else { + VecDeque::from([ModelTurnEvent::Finished(ModelTurnResult { + model: None, + response_id: None, + finish_reason: FinishReason::Completed, + output_items: vec![Item::text(ItemKind::Assistant, "after approval")], + usage: None, + metadata: MetadataMap::new(), + })]) + }; + Ok(GatedTurn { + permits: Arc::clone(&self.permits), + events, + started: false, + }) + } + } + + #[cfg(feature = "unstable-inject")] + #[derive(Clone)] + struct ApprovalFactory { + adapter: ApprovalAdapter, + } + + #[cfg(feature = "unstable-inject")] + #[async_trait] + impl AcpAgentFactory for ApprovalFactory { + async fn start( + &self, + ctx: AcpAgentFactoryContext, + ) -> Result, AcpRuntimeError> { + Agent::builder() + .model(self.adapter.clone()) + .add_tool_source(ToolRegistry::new().with(ApprovalTool::new())) + .permissions(RequireApproval) + .observer(ctx.integration.as_ref().clone()) + .cancellation(ctx.cancellation) + .build() + .map_err(|error| AcpRuntimeError::Loop(error.to_string()))? + .start(SessionConfig::new(ctx.agentkit_session_id).with_metadata(ctx.metadata)) + .await + .map_err(|error| AcpRuntimeError::Loop(error.to_string())) + } + } + + #[cfg(feature = "unstable-inject")] + #[derive(Clone)] + struct CompletionGatedTool { + spec: ToolSpec, + entered: Arc, + release: Arc, + } + + #[cfg(feature = "unstable-inject")] + impl CompletionGatedTool { + fn new() -> Self { + Self { + spec: ToolSpec::new( + "blocking_tool", + "waits for deterministic release", + json!({ "type": "object" }), + ), + entered: Arc::new(AtomicUsize::new(0)), + release: Arc::new(tokio::sync::Semaphore::new(0)), + } + } + + async fn wait_for_entered(&self) { + tokio::time::timeout(Duration::from_secs(2), async { + while self.entered.load(Ordering::Acquire) == 0 { + tokio::task::yield_now().await; + } + }) + .await + .expect("tool did not start"); + } + + fn release(&self) { + self.release.add_permits(1); + } + } + + #[cfg(feature = "unstable-inject")] + #[async_trait] + impl Tool for CompletionGatedTool { + fn spec(&self) -> &ToolSpec { + &self.spec + } + + async fn invoke( + &self, + request: ToolRequest, + _ctx: &mut ToolContext<'_>, + ) -> Result { + self.entered.fetch_add(1, Ordering::AcqRel); + Arc::clone(&self.release) + .acquire_owned() + .await + .expect("tool gate stays open") + .forget(); + Ok(ToolResult::new(ToolResultPart::success( + request.call_id, + ToolOutput::text("done"), + ))) + } + } + + #[cfg(feature = "unstable-inject")] + #[derive(Clone)] + struct InjectToolFactory { + adapter: MockAdapter, + tool: CompletionGatedTool, + } + + #[cfg(feature = "unstable-inject")] + #[async_trait] + impl AcpAgentFactory for InjectToolFactory { + async fn start( + &self, + ctx: AcpAgentFactoryContext, + ) -> Result< + agentkit_loop::LoopDriver<::Session>, + AcpRuntimeError, + > { + Agent::builder() + .model(self.adapter.clone()) + .add_tool_source(ToolRegistry::new().with(self.tool.clone())) + .observer(ctx.integration.as_ref().clone()) + .cancellation(ctx.cancellation) + .build() + .map_err(|error| AcpRuntimeError::Loop(error.to_string()))? + .start(SessionConfig::new(ctx.agentkit_session_id).with_metadata(ctx.metadata)) + .await + .map_err(|error| AcpRuntimeError::Loop(error.to_string())) + } + } + #[async_trait] impl AcpAgentFactory for BackgroundToolTestFactory where @@ -2906,33 +4398,218 @@ mod tests { let _ = client.await; } + #[cfg(not(feature = "unstable-inject"))] #[test] - fn tool_execution_boundary_rotates_message_ids_without_a_terminal_result() { + fn public_capabilities_omit_inject_without_feature() { + let capabilities = serde_json::to_value(agent_capabilities()).unwrap(); + assert!(capabilities.pointer("/session/inject").is_none()); + } + + #[cfg(feature = "unstable-inject")] + #[test] + fn public_host_helpers_expose_conversion_capabilities_and_limits() { + assert_eq!(MAX_PENDING_INJECTIONS, 64); + assert_eq!(MAX_PENDING_INJECTION_BYTES, 256 * 1024); + assert_eq!(MAX_ACCEPTED_INJECTIONS, 4_096); + + let inject = session_inject_capabilities(); + assert_eq!(inject.modes, vec![wire::SessionInjectMode::Steer]); + assert_eq!( + inject.steer_in_stream, + Some(vec![wire::SessionInjectSteerInStream::Finish]) + ); + let advertised = agent_capabilities() + .session + .and_then(|session| session.inject) + .expect("inject capability must match compiled handlers"); + assert_eq!(advertised, inject); + + let items = content_blocks_to_items(&[wire::ContentBlock::Text(wire::TextContent::new( + "converted", + ))]) + .expect("public content conversion"); + assert_eq!(items.len(), 1); + assert!(matches!( + &items[0].parts[0], + Part::Text(text) if text.text == "converted" + )); + } + + #[cfg(feature = "unstable-inject")] + #[tokio::test] + async fn public_host_staged_inject_api_commits_and_activates_receipt() { let integration = AcpIntegration::default(); - let (client, mut messages) = ClientHandle::channel(); - let acp_id = wire::SessionId::new("acp-session"); - let agentkit_id = AgentkitSessionId::new("agentkit-session"); - integration - .bind(acp_id.clone(), agentkit_id.clone(), client) - .expect("bind session"); - integration - .install_prompt_state( - &acp_id, - Arc::new(AtomicU64::new(1)), - Arc::new(Mutex::new(())), - ) - .expect("install prompt state"); - let cancellation = CancellationController::new(); - integration - .begin_prompt(&acp_id, 1, cancellation.handle().checkpoint()) - .expect("begin prompt"); + let session_id = wire::SessionId::new("external-host"); + let session = integration + .bind_session(AcpSessionBinding::new( + session_id.clone(), + AgentkitSessionId::new("external-host-loop"), + RecordingSink::default(), + )) + .expect("bind external host session"); + session.prepare_injection_turn(); + session.start_injection_turn(); - for event in [ - AgentEvent::ContentDelta(Delta::BeginPart { - part_id: PartId::new("before-tool"), - kind: PartKind::Text, - }), - AgentEvent::ContentDelta(Delta::AppendText { + let stages = Arc::new(Mutex::new(Vec::new())); + let (client_transport, agent_transport) = Channel::duplex(); + let server = tokio::spawn({ + let integration = integration.clone(); + let stages = Arc::clone(&stages); + async move { + agent_client_protocol::Agent + .v2() + .on_receive_request( + async move |request: wire::InitializeRequest, responder, _cx| { + responder.respond( + wire::InitializeResponse::new( + request.protocol_version, + wire::Implementation::new("external-host", "1"), + ) + .capabilities(agent_capabilities()), + ) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + { + let integration = integration.clone(); + let stages = Arc::clone(&stages); + async move |request: wire::InjectSessionRequest, responder, cx| { + let integration = integration.clone(); + let stages = Arc::clone(&stages); + cx.spawn(async move { + let reserved = integration + .reserve_inject_request(request, responder) + .await? + .expect("inject reservation"); + let response_id = reserved.response().message_id; + stages.lock().unwrap().push("reserved"); + let acceptance = reserved + .respond_tracked()? + .expect("tracked response acceptance"); + assert_eq!(acceptance.message_id(), &response_id); + stages.lock().unwrap().push("responded"); + acceptance.activate_after_response().await?; + stages.lock().unwrap().push("activated"); + Ok(()) + })?; + Ok(()) + } + }, + agent_client_protocol::on_receive_request!(), + ) + .connect_to(agent_transport) + .await + } + }); + + let client = agent_client_protocol::Client + .v2() + .connect_with(client_transport, { + let session_id = session_id.clone(); + async move |cx| { + cx.send_request(wire::InitializeRequest::new( + wire::ProtocolVersion::V2, + wire::Implementation::new("external-client", "1"), + )) + .block_task() + .await?; + let response = cx + .send_request(wire::InjectSessionRequest::new( + session_id, + wire::SessionInjectMode::Steer, + vec![wire::ContentBlock::Text(wire::TextContent::new( + "external steer", + ))], + )) + .block_task() + .await?; + assert!(response.message_id.to_string().contains("-user-")); + Ok(()) + } + }); + tokio::time::timeout(Duration::from_secs(2), client) + .await + .expect("external host client timed out") + .expect("external host client failed"); + tokio::time::timeout(Duration::from_secs(2), async { + while stages.lock().unwrap().len() != 3 { + tokio::task::yield_now().await; + } + }) + .await + .expect("response receipt was not activated"); + assert_eq!( + stages.lock().unwrap().as_slice(), + ["reserved", "responded", "activated"] + ); + + session.stop_injection_turn(); + integration.unbind_session(&session_id).unwrap(); + server.abort(); + let _ = server.await; + } + + #[tokio::test] + async fn public_host_binding_routes_updates_through_sink() { + let integration = AcpIntegration::default(); + let sink = RecordingSink::default(); + let acp_id = wire::SessionId::new("public-host"); + let agentkit_id = AgentkitSessionId::new("public-host-loop"); + let session = integration + .bind_session(AcpSessionBinding::new( + acp_id.clone(), + agentkit_id.clone(), + sink.clone(), + )) + .expect("bind public host session"); + + assert_eq!(session.acp_session_id(), &acp_id); + assert_eq!(session.agentkit_session_id(), &agentkit_id); + let prompt = wire::PromptRequest::new( + acp_id.clone(), + vec![wire::ContentBlock::Text(wire::TextContent::new("hello"))], + ); + assert_eq!(integration.prompt_to_items(&prompt).unwrap().len(), 1); + integration.begin_prompt(&acp_id).unwrap(); + integration.route_event( + &agentkit_id, + AgentEvent::ContentDelta(Delta::AppendText { + part_id: PartId::new("part"), + chunk: "response".into(), + }), + ); + integration.flush_session_updates(&acp_id).await.unwrap(); + + let updates = sink.updates.lock().unwrap(); + assert_eq!(updates.len(), 1); + assert_eq!(updates[0].session_id, acp_id); + assert_eq!(sink.flushes.load(Ordering::Acquire), 1); + drop(updates); + integration.unbind_session(&acp_id).unwrap(); + } + + #[test] + fn public_host_binding_rotates_message_ids_at_tool_boundary() { + let integration = AcpIntegration::default(); + let (client, mut messages) = ClientHandle::channel(); + let acp_id = wire::SessionId::new("acp-session"); + let agentkit_id = AgentkitSessionId::new("agentkit-session"); + integration + .bind_session(AcpSessionBinding::new( + acp_id.clone(), + agentkit_id.clone(), + client, + )) + .expect("bind session"); + integration.begin_prompt(&acp_id).expect("begin prompt"); + + for event in [ + AgentEvent::ContentDelta(Delta::BeginPart { + part_id: PartId::new("before-tool"), + kind: PartKind::Text, + }), + AgentEvent::ContentDelta(Delta::AppendText { part_id: PartId::new("before-tool"), chunk: "before".into(), }), @@ -2959,7 +4636,7 @@ mod tests { wire::SessionUpdate::AgentMessageChunk(chunk) => Some(chunk.message_id), _ => None, }, - ClientMessage::Flush(_) => None, + ClientMessage::AcknowledgedUpdate { .. } | ClientMessage::Flush(_) => None, }) .collect::>(); assert_eq!( @@ -2972,7 +4649,7 @@ mod tests { } #[test] - fn integration_rejects_duplicate_agentkit_session_ids() { + fn public_host_binding_rejects_duplicate_agentkit_session_ids() { let integration = AcpIntegration::default(); let (first_client, _first_rx) = ClientHandle::channel(); let (second_client, _second_rx) = ClientHandle::channel(); @@ -2980,17 +4657,28 @@ mod tests { let first_acp = wire::SessionId::new("first-acp-session"); let second_acp = wire::SessionId::new("second-acp-session"); integration - .bind(first_acp.clone(), agentkit_id.clone(), first_client) + .bind_session(AcpSessionBinding::new( + first_acp.clone(), + agentkit_id.clone(), + first_client, + )) .expect("first binding"); let error = integration - .bind(second_acp.clone(), agentkit_id, second_client) - .expect_err("duplicate AgentKit session id must be rejected"); + .bind_session(AcpSessionBinding::new( + second_acp.clone(), + agentkit_id, + second_client, + )) + .err() + .expect("duplicate AgentKit session id must be rejected"); assert!(matches!(error, AcpRuntimeError::SessionAlreadyBound(_))); assert!(matches!( integration.session(&second_acp), Err(AcpRuntimeError::SessionNotFound(_)) )); - integration.unbind(&first_acp).expect("unbind first"); + integration + .unbind_session(&first_acp) + .expect("unbind first"); } #[test] @@ -3022,7 +4710,7 @@ mod tests { ) .expect("install prompt state"); integration - .begin_prompt(&acp_id, 1, cancellation.handle().checkpoint()) + .begin_prompt_owner(&acp_id, 1, cancellation.handle().checkpoint()) .expect("begin prompt"); integration.route_event( @@ -3086,7 +4774,7 @@ mod tests { .install_prompt_state(&acp_id, active_prompt, Arc::new(Mutex::new(()))) .expect("install prompt state"); integration - .begin_prompt(&acp_id, index, cancellation.handle().checkpoint()) + .begin_prompt_owner(&acp_id, index, cancellation.handle().checkpoint()) .expect("begin prompt"); let turn_id = TurnId::new(format!("turn-{index}")); integration.route_event( @@ -3163,7 +4851,7 @@ mod tests { .expect("install prompt state"); let cancellation = CancellationController::new(); integration - .begin_prompt(&acp_id, owner, cancellation.handle().checkpoint()) + .begin_prompt_owner(&acp_id, owner, cancellation.handle().checkpoint()) .expect("begin prompt"); fail_accepted_prompt( @@ -3234,7 +4922,7 @@ mod tests { .expect("install prompt state"); let cancellation = CancellationController::new(); integration - .begin_prompt(&acp_id, 7, cancellation.handle().checkpoint()) + .begin_prompt_owner(&acp_id, 7, cancellation.handle().checkpoint()) .expect("begin prompt"); let turn_id = TurnId::new("turn-1"); integration.route_event( @@ -3296,7 +4984,7 @@ mod tests { .expect("install prompt state"); let cancellation = CancellationController::new(); integration - .begin_prompt(&acp_id, 9, cancellation.handle().checkpoint()) + .begin_prompt_owner(&acp_id, 9, cancellation.handle().checkpoint()) .expect("begin prompt"); let turn_id = TurnId::new("turn-race"); integration.route_event( @@ -3580,134 +5268,1676 @@ mod tests { } } - #[tokio::test] - async fn same_session_prompt_race_admits_one_and_other_sessions_stay_independent() { - let updates = Arc::new(Mutex::new(Vec::new())); - let (client_transport, agent_transport) = Channel::duplex(); - let server = tokio::spawn(async move { - AcpHeadlessRuntime::::builder() - .agent_factory(TestFactory { - adapter: BlockingAdapter, - }) - .serve(agent_transport) - .await - }); + #[cfg(feature = "unstable-inject")] + #[derive(Clone)] + struct GatedAdapter { + permits: Arc, + next_turn: Arc, + requests: Arc>>, + } - let client = - agent_client_protocol::Client - .v2() - .on_receive_notification( - { - let updates = Arc::clone(&updates); - async move |notification: wire::UpdateSessionNotification, _cx| { - updates - .lock() - .unwrap() - .push((notification.session_id, notification.update)); - Ok(()) - } - }, - agent_client_protocol::on_receive_notification!(), - ) - .connect_with(client_transport, { - let updates = Arc::clone(&updates); - async move |cx| { - cx.send_request(wire::InitializeRequest::new( - wire::ProtocolVersion::V2, - wire::Implementation::new("test-client", "1"), - )) - .block_task() - .await?; - let cwd = std::env::current_dir() - .map_err(agent_client_protocol::Error::into_internal_error)?; - let first = cx - .send_request(wire::NewSessionRequest::new(cwd.clone())) - .block_task() - .await?; - let second = cx - .send_request(wire::NewSessionRequest::new(cwd)) - .block_task() - .await?; + #[cfg(feature = "unstable-inject")] + impl GatedAdapter { + fn new() -> Self { + Self { + permits: Arc::new(tokio::sync::Semaphore::new(0)), + next_turn: Arc::new(AtomicUsize::new(0)), + requests: Arc::new(Mutex::new(Vec::new())), + } + } - let first_prompt = cx - .send_request(wire::PromptRequest::new( - first.session_id.clone(), - vec![wire::ContentBlock::Text(wire::TextContent::new("first"))], - )) - .block_task(); - let competing_prompt = cx - .send_request(wire::PromptRequest::new( - first.session_id.clone(), - vec![wire::ContentBlock::Text(wire::TextContent::new( - "competing", - ))], - )) - .block_task(); - let (first_result, competing_result) = - tokio::time::timeout(Duration::from_millis(250), async { - tokio::join!(first_prompt, competing_prompt) - }) - .await - .expect("same-session prompt admission timed out"); - assert_ne!( - first_result.is_ok(), - competing_result.is_ok(), - "exactly one same-session prompt must be admitted" - ); + fn release(&self, count: usize) { + self.permits.add_permits(count); + } - tokio::time::timeout( - Duration::from_millis(250), - cx.send_request(wire::PromptRequest::new( - second.session_id.clone(), - vec![wire::ContentBlock::Text(wire::TextContent::new("wait"))], - )) - .block_task(), - ) - .await - .expect("prompt acceptance was blocked by another session")?; + fn requests(&self) -> Vec { + self.requests + .lock() + .unwrap_or_else(|error| error.into_inner()) + .clone() + } + } - cx.send_notification(wire::CancelSessionNotification::new( - first.session_id.clone(), - ))?; - wait_for_idle(&updates, &first.session_id).await; - assert!(!updates.lock().unwrap().iter().any(|(id, update)| { - id == &second.session_id - && matches!( - update, - wire::SessionUpdate::StateUpdate(wire::StateUpdate::Idle(_)) - ) - })); + #[cfg(feature = "unstable-inject")] + #[derive(Clone)] + struct SecondStartGatedFactory { + adapter: GatedAdapter, + starts: Arc, + second_started: Arc, + second_release: Arc, + } - cx.send_notification(wire::CancelSessionNotification::new( - second.session_id.clone(), - ))?; - wait_for_idle(&updates, &second.session_id).await; - for session_id in [&first.session_id, &second.session_id] { - let stop_reason = updates.lock().unwrap().iter().find_map( - |(id, update)| match update { - wire::SessionUpdate::StateUpdate(wire::StateUpdate::Idle( - idle, - )) if id == session_id => idle.stop_reason.clone(), - _ => None, - }, - ); - assert_eq!(stop_reason, Some(wire::StopReason::Cancelled)); - } - cx.send_request(wire::CloseSessionRequest::new(first.session_id)) - .block_task() - .await?; - cx.send_request(wire::CloseSessionRequest::new(second.session_id)) - .block_task() - .await?; - Ok(()) + #[cfg(feature = "unstable-inject")] + impl SecondStartGatedFactory { + fn new(adapter: GatedAdapter) -> Self { + Self { + adapter, + starts: Arc::new(AtomicUsize::new(0)), + second_started: Arc::new(Notify::new()), + second_release: Arc::new(tokio::sync::Semaphore::new(0)), + } + } + + async fn wait_for_second_start(&self) { + if self.starts.load(Ordering::Acquire) < 2 { + tokio::time::timeout(Duration::from_secs(2), self.second_started.notified()) + .await + .expect("batched sibling did not start"); + } + } + + fn release_second_start(&self) { + self.second_release.add_permits(1); + } + } + + #[cfg(feature = "unstable-inject")] + #[async_trait] + impl AcpAgentFactory for SecondStartGatedFactory { + async fn start( + &self, + ctx: AcpAgentFactoryContext, + ) -> Result, AcpRuntimeError> { + if self.starts.fetch_add(1, Ordering::AcqRel) == 1 { + self.second_started.notify_one(); + Arc::clone(&self.second_release) + .acquire_owned() + .await + .expect("second-start gate stays open") + .forget(); + } + Agent::builder() + .model(self.adapter.clone()) + .observer(ctx.integration.as_ref().clone()) + .cancellation(ctx.cancellation) + .build() + .map_err(|error| AcpRuntimeError::Loop(error.to_string()))? + .start(SessionConfig::new(ctx.agentkit_session_id).with_metadata(ctx.metadata)) + .await + .map_err(|error| AcpRuntimeError::Loop(error.to_string())) + } + } + + #[cfg(feature = "unstable-inject")] + async fn raw_request( + channel: &mut Channel, + id: i64, + method: &str, + params: serde_json::Value, + ) -> serde_json::Value { + channel + .tx + .unbounded_send(TransportFrame::Single( + RawJsonRpcMessage::request(method.to_string(), params, id.into()) + .expect("valid raw request"), + )) + .expect("server channel stays open"); + raw_response(channel, id).await + } + + #[cfg(feature = "unstable-inject")] + async fn raw_response(channel: &mut Channel, id: i64) -> serde_json::Value { + tokio::time::timeout(Duration::from_secs(2), async { + loop { + let frame = channel.rx.next().await.expect("server channel closed"); + if let TransportFrame::Single(message @ RawJsonRpcMessage::Response(_)) = frame { + let response = serde_json::to_value(message).expect("response serializes"); + if response.get("id") == Some(&json!(id)) { + return response; } - }); + } + } + }) + .await + .expect("raw request timed out") + } - tokio::time::timeout(Duration::from_secs(5), client) - .await - .expect("client timed out") - .expect("client run"); - server.abort(); - let _ = server.await; + #[cfg(feature = "unstable-inject")] + struct GatedSession { + permits: Arc, + next_turn: Arc, + requests: Arc>>, + } + + #[cfg(feature = "unstable-inject")] + struct GatedTurn { + permits: Arc, + events: VecDeque, + started: bool, + } + + #[cfg(feature = "unstable-inject")] + #[async_trait] + impl ModelAdapter for GatedAdapter { + type Session = GatedSession; + + async fn start_session(&self, _config: SessionConfig) -> Result { + Ok(GatedSession { + permits: Arc::clone(&self.permits), + next_turn: Arc::clone(&self.next_turn), + requests: Arc::clone(&self.requests), + }) + } + } + + #[cfg(feature = "unstable-inject")] + #[async_trait] + impl ModelSession for GatedSession { + type Turn = GatedTurn; + + async fn begin_turn( + &mut self, + request: TurnRequest, + _cancellation: Option, + ) -> Result { + self.requests + .lock() + .unwrap_or_else(|error| error.into_inner()) + .push(request); + let turn = self.next_turn.fetch_add(1, Ordering::AcqRel); + let text = format!("turn-{turn}"); + Ok(GatedTurn { + permits: Arc::clone(&self.permits), + events: VecDeque::from([ + ModelTurnEvent::Delta(Delta::BeginPart { + part_id: PartId::new(format!("part-{turn}")), + kind: PartKind::Text, + }), + ModelTurnEvent::Delta(Delta::AppendText { + part_id: PartId::new(format!("part-{turn}")), + chunk: text.clone(), + }), + ModelTurnEvent::Finished(ModelTurnResult { + model: None, + response_id: None, + finish_reason: FinishReason::Completed, + output_items: vec![Item::text(ItemKind::Assistant, text)], + usage: None, + metadata: MetadataMap::new(), + }), + ]), + started: false, + }) + } + } + + #[cfg(feature = "unstable-inject")] + #[async_trait] + impl ModelTurn for GatedTurn { + async fn next_event( + &mut self, + cancellation: Option, + ) -> Result, LoopError> { + if !self.started { + self.started = true; + if let Some(cancellation) = cancellation { + tokio::select! { + permit = Arc::clone(&self.permits).acquire_owned() => { + permit.expect("gate stays open").forget(); + } + _ = cancellation.cancelled() => { + return Ok(Some(ModelTurnEvent::Finished(ModelTurnResult { + model: None, + response_id: None, + finish_reason: FinishReason::Cancelled, + output_items: Vec::new(), + usage: None, + metadata: MetadataMap::new(), + }))); + } + } + } + } + Ok(self.events.pop_front()) + } + } + + #[cfg(feature = "unstable-inject")] + #[tokio::test] + async fn batched_session_inject_activates_after_aggregate_response_enqueue() { + let adapter = GatedAdapter::new(); + let factory = SecondStartGatedFactory::new(adapter.clone()); + let coordinator = AcpIntegration::default(); + let (mut client, agent_transport) = Channel::duplex(); + let server = tokio::spawn({ + let factory = factory.clone(); + let coordinator = coordinator.clone(); + async move { + AcpHeadlessRuntime::::builder() + .agent_factory(factory) + .integration(coordinator) + .serve(agent_transport) + .await + } + }); + + let initialize = raw_request( + &mut client, + 1, + "initialize", + serde_json::to_value(wire::InitializeRequest::new( + wire::ProtocolVersion::V2, + wire::Implementation::new("batch-test", "1"), + )) + .unwrap(), + ) + .await; + assert!(initialize.get("result").is_some()); + let cwd = std::env::current_dir().unwrap(); + let new_session = raw_request( + &mut client, + 2, + "session/new", + serde_json::to_value(wire::NewSessionRequest::new(cwd.clone())).unwrap(), + ) + .await; + let session_id = new_session["result"]["sessionId"] + .as_str() + .unwrap() + .to_string(); + let prompt = raw_request( + &mut client, + 3, + "session/prompt", + serde_json::to_value(wire::PromptRequest::new( + wire::SessionId::new(session_id.clone()), + vec![wire::ContentBlock::Text(wire::TextContent::new("start"))], + )) + .unwrap(), + ) + .await; + assert!(prompt.get("result").is_some()); + + let inject = RawJsonRpcMessage::request( + "session/inject".to_string(), + serde_json::to_value(wire::InjectSessionRequest::new( + wire::SessionId::new(session_id.clone()), + wire::SessionInjectMode::Steer, + vec![wire::ContentBlock::Text(wire::TextContent::new( + "batched steer", + ))], + )) + .unwrap(), + 10.into(), + ) + .unwrap(); + let second_inject = RawJsonRpcMessage::request( + "session/inject".to_string(), + serde_json::to_value(wire::InjectSessionRequest::new( + wire::SessionId::new(session_id.clone()), + wire::SessionInjectMode::Steer, + vec![wire::ContentBlock::Text(wire::TextContent::new( + "second batched steer", + ))], + )) + .unwrap(), + 12.into(), + ) + .unwrap(); + let slow_sibling = RawJsonRpcMessage::request( + "session/new".to_string(), + serde_json::to_value(wire::NewSessionRequest::new(cwd)).unwrap(), + 11.into(), + ) + .unwrap(); + client + .tx + .unbounded_send(TransportFrame::Batch( + TransportBatch::from_messages([inject, second_inject, slow_sibling]).unwrap(), + )) + .unwrap(); + + factory.wait_for_second_start().await; + adapter.release(1); + tokio::time::sleep(Duration::from_millis(30)).await; + assert_eq!( + adapter.requests().len(), + 1, + "batch injections became deliverable before their aggregate response receipt" + ); + factory.release_second_start(); + + let response = tokio::time::timeout(Duration::from_secs(2), async { + loop { + if let Some(TransportFrame::Batch(batch)) = client.rx.next().await { + break serde_json::to_value(batch).unwrap(); + } + } + }) + .await + .expect("batched response timed out"); + let entries = response.as_array().expect("aggregate batch response"); + let inject_response = entries + .iter() + .find(|entry| entry.get("id") == Some(&json!(10))) + .expect("inject response slot"); + assert!(inject_response.get("result").is_some()); + assert!(inject_response.get("error").is_none()); + assert_eq!( + coordinator.inject_requests.next.load(Ordering::Acquire), + 2, + "headless inject handlers did not use the supplied public coordinator" + ); + assert!( + entries + .iter() + .find(|entry| entry.get("id") == Some(&json!(12))) + .is_some_and(|entry| entry.get("result").is_some()) + ); + + adapter.release(1); + let delivered = tokio::time::timeout(Duration::from_secs(2), async { + loop { + let frame = client.rx.next().await.expect("server channel stays open"); + let value = match frame { + TransportFrame::Single(message) => serde_json::to_value(message).unwrap(), + TransportFrame::Batch(batch) => serde_json::to_value(batch).unwrap(), + TransportFrame::Malformed { raw, .. } => serde_json::Value::String(raw), + }; + if serde_json::to_string(&value) + .unwrap() + .contains("batched steer") + { + break true; + } + } + }) + .await + .expect("batched injection was not delivered"); + assert!(delivered); + let requests = adapter.requests(); + assert!(requests.len() >= 2); + let steers = requests[1] + .transcript + .iter() + .filter(|item| item.kind == ItemKind::User) + .flat_map(|item| &item.parts) + .filter_map(|part| match part { + Part::Text(text) if text.text.contains("batched steer") => Some(text.text.as_str()), + _ => None, + }) + .collect::>(); + assert_eq!(steers, ["batched steer", "second batched steer"]); + + drop(client); + server.abort(); + let _ = server.await; + } + + #[cfg(feature = "unstable-inject")] + #[tokio::test] + async fn request_cancellation_respects_response_commit() { + let adapter = GatedAdapter::new(); + let (mut client, agent_transport) = Channel::duplex(); + let server = tokio::spawn({ + let adapter = adapter.clone(); + async move { + AcpHeadlessRuntime::::builder() + .agent_factory(TestFactory { adapter }) + .serve(agent_transport) + .await + } + }); + raw_request( + &mut client, + 1, + "initialize", + serde_json::to_value(wire::InitializeRequest::new( + wire::ProtocolVersion::V2, + wire::Implementation::new("cancel-test", "1"), + )) + .unwrap(), + ) + .await; + let created = raw_request( + &mut client, + 2, + "session/new", + serde_json::to_value(wire::NewSessionRequest::new( + std::env::current_dir().unwrap(), + )) + .unwrap(), + ) + .await; + let session_id = + wire::SessionId::new(created["result"]["sessionId"].as_str().unwrap().to_string()); + raw_request( + &mut client, + 3, + "session/prompt", + serde_json::to_value(wire::PromptRequest::new( + session_id.clone(), + vec![wire::ContentBlock::Text(wire::TextContent::new("start"))], + )) + .unwrap(), + ) + .await; + + let large_steer = "x".repeat(200_000); + client + .tx + .unbounded_send(TransportFrame::Single( + RawJsonRpcMessage::request( + "session/inject".to_string(), + serde_json::to_value(wire::InjectSessionRequest::new( + session_id.clone(), + wire::SessionInjectMode::Steer, + vec![wire::ContentBlock::Text(wire::TextContent::new( + large_steer, + ))], + )) + .unwrap(), + 10.into(), + ) + .unwrap(), + )) + .unwrap(); + client + .tx + .unbounded_send(TransportFrame::Single( + RawJsonRpcMessage::notification( + "$/cancel_request".to_string(), + json!({ "requestId": 10 }), + ) + .unwrap(), + )) + .unwrap(); + let cancelled_response = raw_response(&mut client, 10).await; + assert!(cancelled_response.get("result").is_none()); + assert_eq!(cancelled_response["error"]["code"], json!(-32800)); + + let committed = raw_request( + &mut client, + 11, + "session/inject", + serde_json::to_value(wire::InjectSessionRequest::new( + session_id, + wire::SessionInjectMode::Steer, + vec![wire::ContentBlock::Text(wire::TextContent::new( + "committed", + ))], + )) + .unwrap(), + ) + .await; + assert!(committed.get("result").is_some()); + client + .tx + .unbounded_send(TransportFrame::Single( + RawJsonRpcMessage::notification( + "$/cancel_request".to_string(), + json!({ "requestId": 11 }), + ) + .unwrap(), + )) + .unwrap(); + adapter.release(2); + + tokio::time::timeout(Duration::from_secs(2), async { + while adapter.requests().len() < 2 { + tokio::task::yield_now().await; + } + }) + .await + .expect("committed injection was removed by late request cancellation"); + assert!(adapter.requests()[1].transcript.iter().any(|item| { + item.parts + .iter() + .any(|part| matches!(part, Part::Text(text) if text.text == "committed")) + })); + drop(client); + server.abort(); + let _ = server.await; + } + + #[cfg(feature = "unstable-inject")] + #[tokio::test] + async fn closing_session_discards_pending_inject_and_returns_not_found() { + let adapter = GatedAdapter::new(); + let factory = SecondStartGatedFactory::new(adapter); + let (mut client, agent_transport) = Channel::duplex(); + let server = tokio::spawn({ + let factory = factory.clone(); + async move { + AcpHeadlessRuntime::::builder() + .agent_factory(factory) + .serve(agent_transport) + .await + } + }); + raw_request( + &mut client, + 1, + "initialize", + serde_json::to_value(wire::InitializeRequest::new( + wire::ProtocolVersion::V2, + wire::Implementation::new("close-test", "1"), + )) + .unwrap(), + ) + .await; + let cwd = std::env::current_dir().unwrap(); + let created = raw_request( + &mut client, + 2, + "session/new", + serde_json::to_value(wire::NewSessionRequest::new(cwd.clone())).unwrap(), + ) + .await; + let session_id = + wire::SessionId::new(created["result"]["sessionId"].as_str().unwrap().to_string()); + raw_request( + &mut client, + 3, + "session/prompt", + serde_json::to_value(wire::PromptRequest::new( + session_id.clone(), + vec![wire::ContentBlock::Text(wire::TextContent::new("start"))], + )) + .unwrap(), + ) + .await; + let inject = RawJsonRpcMessage::request( + "session/inject".to_string(), + serde_json::to_value(wire::InjectSessionRequest::new( + session_id.clone(), + wire::SessionInjectMode::Steer, + vec![wire::ContentBlock::Text(wire::TextContent::new("pending"))], + )) + .unwrap(), + 10.into(), + ) + .unwrap(); + let sibling = RawJsonRpcMessage::request( + "session/new".to_string(), + serde_json::to_value(wire::NewSessionRequest::new(cwd)).unwrap(), + 11.into(), + ) + .unwrap(); + client + .tx + .unbounded_send(TransportFrame::Batch( + TransportBatch::from_messages([inject, sibling]).unwrap(), + )) + .unwrap(); + factory.wait_for_second_start().await; + let close = raw_request( + &mut client, + 20, + "session/close", + serde_json::to_value(wire::CloseSessionRequest::new(session_id.clone())).unwrap(), + ) + .await; + assert!(close.get("result").is_some()); + factory.release_second_start(); + tokio::time::sleep(Duration::from_millis(30)).await; + let missing = raw_request( + &mut client, + 21, + "session/inject", + serde_json::to_value(wire::InjectSessionRequest::new( + session_id, + wire::SessionInjectMode::Steer, + vec![wire::ContentBlock::Text(wire::TextContent::new("late"))], + )) + .unwrap(), + ) + .await; + assert_eq!(missing["error"]["code"], json!(-32002)); + + drop(client); + server.abort(); + let _ = server.await; + } + + #[cfg(feature = "unstable-inject")] + #[tokio::test] + async fn same_batch_close_prevents_inject_acceptance() { + let adapter = GatedAdapter::new(); + let (mut client, agent_transport) = Channel::duplex(); + let server = tokio::spawn({ + let adapter = adapter.clone(); + async move { + AcpHeadlessRuntime::::builder() + .agent_factory(TestFactory { adapter }) + .serve(agent_transport) + .await + } + }); + raw_request( + &mut client, + 1, + "initialize", + serde_json::to_value(wire::InitializeRequest::new( + wire::ProtocolVersion::V2, + wire::Implementation::new("same-batch-close-test", "1"), + )) + .unwrap(), + ) + .await; + let created = raw_request( + &mut client, + 2, + "session/new", + serde_json::to_value(wire::NewSessionRequest::new( + std::env::current_dir().unwrap(), + )) + .unwrap(), + ) + .await; + let session_id = + wire::SessionId::new(created["result"]["sessionId"].as_str().unwrap().to_string()); + raw_request( + &mut client, + 3, + "session/prompt", + serde_json::to_value(wire::PromptRequest::new( + session_id.clone(), + vec![wire::ContentBlock::Text(wire::TextContent::new("start"))], + )) + .unwrap(), + ) + .await; + + let inject = RawJsonRpcMessage::request( + "session/inject".to_string(), + serde_json::to_value(wire::InjectSessionRequest::new( + session_id.clone(), + wire::SessionInjectMode::Steer, + vec![wire::ContentBlock::Text(wire::TextContent::new("discard"))], + )) + .unwrap(), + 10.into(), + ) + .unwrap(); + let close = RawJsonRpcMessage::request( + "session/close".to_string(), + serde_json::to_value(wire::CloseSessionRequest::new(session_id)).unwrap(), + 11.into(), + ) + .unwrap(); + client + .tx + .unbounded_send(TransportFrame::Batch( + TransportBatch::from_messages([inject, close]).unwrap(), + )) + .unwrap(); + + let response = tokio::time::timeout(Duration::from_secs(2), async { + loop { + if let Some(TransportFrame::Batch(batch)) = client.rx.next().await { + break serde_json::to_value(batch).unwrap(); + } + } + }) + .await + .expect("same-batch close response timed out"); + let entries = response.as_array().expect("aggregate batch response"); + let inject_response = entries + .iter() + .find(|entry| entry.get("id") == Some(&json!(10))) + .expect("inject response slot"); + assert!(inject_response.get("result").is_none()); + assert_eq!(inject_response["error"]["code"], json!(-32002)); + let close_response = entries + .iter() + .find(|entry| entry.get("id") == Some(&json!(11))) + .expect("close response slot"); + assert!(close_response.get("result").is_some()); + assert!(close_response.get("error").is_none()); + assert_eq!(adapter.requests().len(), 1); + + drop(client); + server.abort(); + let _ = server.await; + } + + #[cfg(feature = "unstable-inject")] + #[test] + fn accepted_injection_lifetime_cap_is_enforced() { + let injection = InjectionController::default(); + injection.start_turn(); + for index in 0..MAX_ACCEPTED_INJECTIONS { + let message_id = wire::MessageId::new(format!("accepted-{index}")); + injection + .reserve(PendingInject { + message_id: message_id.clone(), + content: Vec::new(), + items: Vec::new(), + bytes: 0, + commitment: InjectCommitment::Reserved, + }) + .expect("injection below lifetime cap"); + injection.commit(&message_id); + injection.activate(&message_id); + assert!(matches!( + injection.revoke_transition(&message_id), + RevokeTransition::Revoked + )); + } + + let error = injection + .reserve(PendingInject { + message_id: wire::MessageId::new("over-limit"), + content: Vec::new(), + items: Vec::new(), + bytes: 0, + commitment: InjectCommitment::Reserved, + }) + .expect_err("accepted injection lifetime cap must reject new reservations"); + assert_eq!(i32::from(error.code), -32000); + assert_eq!( + error.data, + Some(json!({ + "reason": "lifetime_limit_exceeded", + "limit": MAX_ACCEPTED_INJECTIONS, + })) + ); + } + + #[cfg(feature = "unstable-inject")] + #[tokio::test] + async fn accepted_steer_survives_approval_resolution() { + let adapter = ApprovalAdapter::new(); + let updates = Arc::new(Mutex::new(Vec::new())); + let (client_transport, agent_transport) = Channel::duplex(); + let server = tokio::spawn({ + let adapter = adapter.clone(); + async move { + AcpHeadlessRuntime::::builder() + .agent_factory(ApprovalFactory { adapter }) + .serve(agent_transport) + .await + } + }); + let client = agent_client_protocol::Client + .v2() + .on_receive_notification( + { + let updates = Arc::clone(&updates); + async move |notification: wire::UpdateSessionNotification, _cx| { + updates + .lock() + .unwrap() + .push((notification.session_id, notification.update)); + Ok(()) + } + }, + agent_client_protocol::on_receive_notification!(), + ) + .connect_with(client_transport, { + let updates = Arc::clone(&updates); + async move |cx| { + cx.send_request(wire::InitializeRequest::new( + wire::ProtocolVersion::V2, + wire::Implementation::new("approval-test", "1"), + )) + .block_task() + .await?; + let session_id = cx + .send_request(wire::NewSessionRequest::new( + std::env::current_dir() + .map_err(agent_client_protocol::Error::into_internal_error)?, + )) + .block_task() + .await? + .session_id; + cx.send_request(wire::PromptRequest::new( + session_id.clone(), + vec![wire::ContentBlock::Text(wire::TextContent::new("start"))], + )) + .block_task() + .await?; + cx.send_request(wire::InjectSessionRequest::new( + session_id.clone(), + wire::SessionInjectMode::Steer, + vec![wire::ContentBlock::Text(wire::TextContent::new( + "approval steer", + ))], + )) + .block_task() + .await?; + adapter.release(); + adapter.release(); + wait_for_idle(&updates, &session_id).await; + assert!(updates.lock().unwrap().iter().any(|(id, update)| { + id == &session_id + && matches!( + update, + wire::SessionUpdate::UserMessage(message) + if matches!( + &message.content, + agent_client_protocol::schema::MaybeUndefined::Value(content) + if content.iter().any(|block| { + matches!(block, wire::ContentBlock::Text(text) + if text.text == "approval steer") + }) + ) + ) + })); + Ok(()) + } + }); + tokio::time::timeout(Duration::from_secs(4), client) + .await + .expect("approval client timed out") + .expect("approval client failed"); + server.abort(); + let _ = server.await; + } + + #[cfg(feature = "unstable-inject")] + #[tokio::test] + async fn session_inject_accepted_during_tool_runs_after_tool_boundary() { + let adapter = MockAdapter::new(); + adapter.enqueue_many([ + tool_turn("post-tool"), + TurnScript::new([ModelTurnEvent::Finished(ModelTurnResult { + model: None, + response_id: None, + finish_reason: FinishReason::Completed, + output_items: Vec::new(), + usage: None, + metadata: MetadataMap::new(), + })]), + ]); + let tool = CompletionGatedTool::new(); + let (client_transport, agent_transport) = Channel::duplex(); + let server = tokio::spawn({ + let factory = InjectToolFactory { + adapter: adapter.clone(), + tool: tool.clone(), + }; + async move { + AcpHeadlessRuntime::::builder() + .agent_factory(factory) + .serve(agent_transport) + .await + } + }); + + let client = + agent_client_protocol::Client + .v2() + .connect_with(client_transport, async move |cx| { + cx.send_request(wire::InitializeRequest::new( + wire::ProtocolVersion::V2, + wire::Implementation::new("test-client", "1"), + )) + .block_task() + .await?; + let cwd = std::env::current_dir() + .map_err(agent_client_protocol::Error::into_internal_error)?; + let session_id = cx + .send_request(wire::NewSessionRequest::new(cwd)) + .block_task() + .await? + .session_id; + cx.send_request(wire::PromptRequest::new( + session_id.clone(), + vec![wire::ContentBlock::Text(wire::TextContent::new("start"))], + )) + .block_task() + .await?; + tool.wait_for_entered().await; + cx.send_request(wire::InjectSessionRequest::new( + session_id.clone(), + wire::SessionInjectMode::Steer, + vec![wire::ContentBlock::Text(wire::TextContent::new( + "after tool", + ))], + )) + .block_task() + .await?; + tool.release(); + tokio::time::timeout(Duration::from_secs(2), async { + while adapter.observed().len() < 2 { + tokio::task::yield_now().await; + } + }) + .await + .expect("continuation did not start"); + let observed = adapter.observed(); + assert_eq!(observed.len(), 2); + let contains_injection = |item: &Item| { + item.kind == ItemKind::User + && item.parts.iter().any(|part| { + matches!(part, Part::Text(text) if text.text == "after tool") + }) + }; + assert!(!observed[0].transcript.iter().any(contains_injection)); + assert!(observed[1].transcript.iter().any(contains_injection)); + cx.send_request(wire::CloseSessionRequest::new(session_id)) + .block_task() + .await?; + Ok(()) + }); + + tokio::time::timeout(Duration::from_secs(5), client) + .await + .expect("client timed out") + .expect("client run"); + server.abort(); + let _ = server.await; + } + + #[cfg(feature = "unstable-inject")] + #[tokio::test] + async fn delivery_and_revoke_race_has_one_linearized_outcome() { + let adapter = GatedAdapter::new(); + let updates = Arc::new(Mutex::new(Vec::new())); + let (client_transport, agent_transport) = Channel::duplex(); + let server = tokio::spawn({ + let adapter = adapter.clone(); + async move { + AcpHeadlessRuntime::::builder() + .agent_factory(TestFactory { adapter }) + .serve(agent_transport) + .await + } + }); + let client = agent_client_protocol::Client + .v2() + .on_receive_notification( + { + let updates = Arc::clone(&updates); + async move |notification: wire::UpdateSessionNotification, _cx| { + updates + .lock() + .unwrap() + .push((notification.session_id, notification.update)); + Ok(()) + } + }, + agent_client_protocol::on_receive_notification!(), + ) + .connect_with(client_transport, { + let updates = Arc::clone(&updates); + async move |cx| { + cx.send_request(wire::InitializeRequest::new( + wire::ProtocolVersion::V2, + wire::Implementation::new("race-test", "1"), + )) + .block_task() + .await?; + let session_id = cx + .send_request(wire::NewSessionRequest::new( + std::env::current_dir() + .map_err(agent_client_protocol::Error::into_internal_error)?, + )) + .block_task() + .await? + .session_id; + cx.send_request(wire::PromptRequest::new( + session_id.clone(), + vec![wire::ContentBlock::Text(wire::TextContent::new("race"))], + )) + .block_task() + .await?; + let accepted = cx + .send_request(wire::InjectSessionRequest::new( + session_id.clone(), + wire::SessionInjectMode::Steer, + vec![wire::ContentBlock::Text(wire::TextContent::new( + "racing steer", + ))], + )) + .block_task() + .await?; + let revoke = cx + .send_request(wire::RevokeInjectSessionRequest::new( + session_id.clone(), + accepted.message_id.clone(), + )) + .block_task(); + adapter.release(2); + let revoked = match revoke.await { + Ok(_) => true, + Err(error) => { + assert_eq!(i32::from(error.code), -32010); + false + } + }; + wait_for_idle(&updates, &session_id).await; + let was_delivered = updates.lock().unwrap().iter().any(|(id, update)| { + id == &session_id + && matches!(update, wire::SessionUpdate::UserMessage(message) + if message.message_id == accepted.message_id) + }); + assert_ne!(revoked, was_delivered); + Ok(()) + } + }); + tokio::time::timeout(Duration::from_secs(4), client) + .await + .expect("race client timed out") + .expect("race client failed"); + server.abort(); + let _ = server.await; + } + + #[cfg(feature = "unstable-inject")] + #[tokio::test] + async fn session_inject_steers_preserves_content_and_serializes_revoke_cancel() { + let adapter = GatedAdapter::new(); + let updates = Arc::new(Mutex::new(Vec::new())); + let (client_transport, agent_transport) = Channel::duplex(); + let server = tokio::spawn({ + let adapter = adapter.clone(); + async move { + AcpHeadlessRuntime::::builder() + .agent_factory(TestFactory { adapter }) + .serve(agent_transport) + .await + } + }); + + let client = agent_client_protocol::Client + .v2() + .on_receive_notification( + { + let updates = Arc::clone(&updates); + async move |notification: wire::UpdateSessionNotification, _cx| { + updates + .lock() + .unwrap() + .push((notification.session_id, notification.update)); + Ok(()) + } + }, + agent_client_protocol::on_receive_notification!(), + ) + .connect_with(client_transport, { + let adapter = adapter.clone(); + let updates = Arc::clone(&updates); + async move |cx| { + let initialize = cx + .send_request(wire::InitializeRequest::new( + wire::ProtocolVersion::V2, + wire::Implementation::new("test-client", "1"), + )) + .block_task() + .await?; + let inject = initialize + .capabilities + .session + .and_then(|session| session.inject) + .expect("inject capability advertised"); + assert_eq!(inject.modes, vec![wire::SessionInjectMode::Steer]); + assert_eq!( + inject.steer_in_stream, + Some(vec![wire::SessionInjectSteerInStream::Finish]) + ); + assert!(inject.pending.is_none()); + + let cwd = std::env::current_dir() + .map_err(agent_client_protocol::Error::into_internal_error)?; + let session_id = cx + .send_request(wire::NewSessionRequest::new(cwd)) + .block_task() + .await? + .session_id; + let idle_error = cx + .send_request(wire::InjectSessionRequest::new( + session_id.clone(), + wire::SessionInjectMode::Steer, + vec![wire::ContentBlock::Text(wire::TextContent::new("idle"))], + )) + .block_task() + .await + .expect_err("idle injection must fail"); + assert_eq!(i32::from(idle_error.code), -32010); + assert_eq!( + idle_error.data, + Some(json!({ "reason": "no_running_turn" })) + ); + let replace_error = cx + .send_request(wire::ReplaceInjectSessionRequest::new( + session_id.clone(), + "not-pending", + vec![wire::ContentBlock::Text(wire::TextContent::new("replace"))], + )) + .block_task() + .await + .expect_err("replace handler must not be registered"); + assert_eq!(i32::from(replace_error.code), -32601); + + cx.send_request(wire::PromptRequest::new( + session_id.clone(), + vec![wire::ContentBlock::Text(wire::TextContent::new("start"))], + )) + .block_task() + .await?; + let content = vec![ + wire::ContentBlock::Text(wire::TextContent::new("steer")), + wire::ContentBlock::ResourceLink(wire::ResourceLink::new( + "notes", + "file:///tmp/notes.txt", + )), + ]; + let accepted = tokio::time::timeout( + Duration::from_millis(250), + cx.send_request(wire::InjectSessionRequest::new( + session_id.clone(), + wire::SessionInjectMode::Steer, + content.clone(), + )) + .block_task(), + ) + .await + .expect("inject acceptance waited for the model")?; + assert!( + !updates.lock().unwrap().iter().any(|(id, update)| { + id == &session_id + && matches!(update, wire::SessionUpdate::UserMessage(message) + if message.message_id == accepted.message_id) + }), + "UserMessage was emitted before the inject response" + ); + adapter.release(2); + wait_for_idle_count(&updates, &session_id, 1).await; + let delivered = + updates + .lock() + .unwrap() + .iter() + .find_map(|(id, update)| match update { + wire::SessionUpdate::UserMessage(message) + if id == &session_id + && message.message_id == accepted.message_id => + { + Some(message.content.clone()) + } + _ => None, + }); + assert_eq!( + delivered, + Some(agent_client_protocol::schema::MaybeUndefined::Value( + content + )) + ); + { + let ordered = updates.lock().unwrap(); + let first_agent = ordered + .iter() + .position(|(id, update)| { + id == &session_id + && matches!(update, wire::SessionUpdate::AgentMessageChunk(_)) + }) + .expect("first model stream"); + let injected = ordered + .iter() + .position(|(id, update)| { + id == &session_id + && matches!(update, wire::SessionUpdate::UserMessage(message) + if message.message_id == accepted.message_id) + }) + .expect("injected user message"); + let second_agent = ordered + .iter() + .rposition(|(id, update)| { + id == &session_id + && matches!(update, wire::SessionUpdate::AgentMessageChunk(_)) + }) + .expect("continuation model stream"); + assert!(first_agent < injected && injected < second_agent); + assert_eq!( + ordered + .iter() + .filter(|(id, update)| id == &session_id + && matches!( + update, + wire::SessionUpdate::StateUpdate( + wire::StateUpdate::Running(_) + ) + )) + .count(), + 1 + ); + } + cx.send_request(wire::PromptRequest::new( + session_id.clone(), + vec![wire::ContentBlock::Text(wire::TextContent::new("revoke"))], + )) + .block_task() + .await?; + let queue_error = cx + .send_request(wire::InjectSessionRequest::new( + session_id.clone(), + wire::SessionInjectMode::Queue, + vec![wire::ContentBlock::Text(wire::TextContent::new("queue"))], + )) + .block_task() + .await + .expect_err("queue mode must not be accepted"); + assert_eq!(i32::from(queue_error.code), -32602); + let revoked = cx + .send_request(wire::InjectSessionRequest::new( + session_id.clone(), + wire::SessionInjectMode::Steer, + vec![wire::ContentBlock::Text(wire::TextContent::new("remove"))], + )) + .block_task() + .await?; + cx.send_request(wire::RevokeInjectSessionRequest::new( + session_id.clone(), + revoked.message_id.clone(), + )) + .block_task() + .await?; + adapter.release(1); + wait_for_idle_count(&updates, &session_id, 2).await; + assert!(!updates.lock().unwrap().iter().any(|(id, update)| { + id == &session_id + && matches!(update, wire::SessionUpdate::UserMessage(message) + if message.message_id == revoked.message_id) + })); + let delivered_error = cx + .send_request(wire::RevokeInjectSessionRequest::new( + session_id.clone(), + accepted.message_id.clone(), + )) + .block_task() + .await + .expect_err("old delivered injection cannot be revoked in a later turn"); + assert_eq!(i32::from(delivered_error.code), -32010); + assert_eq!( + delivered_error.data, + Some(json!({ + "reason": "already_delivered", + "messageId": accepted.message_id, + })) + ); + + cx.send_request(wire::PromptRequest::new( + session_id.clone(), + vec![wire::ContentBlock::Text(wire::TextContent::new("cancel"))], + )) + .block_task() + .await?; + let cancelled = cx + .send_request(wire::InjectSessionRequest::new( + session_id.clone(), + wire::SessionInjectMode::Steer, + vec![wire::ContentBlock::Text(wire::TextContent::new("drop"))], + )) + .block_task() + .await?; + cx.send_notification(wire::CancelSessionNotification::new(session_id.clone()))?; + wait_for_idle_count(&updates, &session_id, 3).await; + assert!(!updates.lock().unwrap().iter().any(|(id, update)| { + id == &session_id + && matches!(update, wire::SessionUpdate::UserMessage(message) + if message.message_id == cancelled.message_id) + })); + let stop_reason = updates.lock().unwrap().iter().rev().find_map( + |(id, update)| match update { + wire::SessionUpdate::StateUpdate(wire::StateUpdate::Idle(idle)) + if id == &session_id => + { + idle.stop_reason.clone() + } + _ => None, + }, + ); + assert_eq!(stop_reason, Some(wire::StopReason::Cancelled)); + + cx.send_request(wire::PromptRequest::new( + session_id.clone(), + vec![wire::ContentBlock::Text(wire::TextContent::new( + "after cancel", + ))], + )) + .block_task() + .await?; + adapter.release(2); + wait_for_idle_count(&updates, &session_id, 4).await; + assert!( + updates.lock().unwrap().iter().any(|(id, update)| { + id == &session_id + && matches!(update, wire::SessionUpdate::UserMessage(message) + if message.message_id == cancelled.message_id) + }), + "committed injection did not survive session cancellation" + ); + cx.send_request(wire::CloseSessionRequest::new(session_id)) + .block_task() + .await?; + Ok(()) + } + }); + + tokio::time::timeout(Duration::from_secs(8), client) + .await + .expect("client timed out") + .expect("client run"); + server.abort(); + let _ = server.await; + } + + #[tokio::test] + async fn same_session_prompt_race_admits_one_and_other_sessions_stay_independent() { + let updates = Arc::new(Mutex::new(Vec::new())); + let (client_transport, agent_transport) = Channel::duplex(); + let server = tokio::spawn(async move { + AcpHeadlessRuntime::::builder() + .agent_factory(TestFactory { + adapter: BlockingAdapter, + }) + .serve(agent_transport) + .await + }); + + let client = + agent_client_protocol::Client + .v2() + .on_receive_notification( + { + let updates = Arc::clone(&updates); + async move |notification: wire::UpdateSessionNotification, _cx| { + updates + .lock() + .unwrap() + .push((notification.session_id, notification.update)); + Ok(()) + } + }, + agent_client_protocol::on_receive_notification!(), + ) + .connect_with(client_transport, { + let updates = Arc::clone(&updates); + async move |cx| { + cx.send_request(wire::InitializeRequest::new( + wire::ProtocolVersion::V2, + wire::Implementation::new("test-client", "1"), + )) + .block_task() + .await?; + let cwd = std::env::current_dir() + .map_err(agent_client_protocol::Error::into_internal_error)?; + let first = cx + .send_request(wire::NewSessionRequest::new(cwd.clone())) + .block_task() + .await?; + let second = cx + .send_request(wire::NewSessionRequest::new(cwd)) + .block_task() + .await?; + + let first_prompt = cx + .send_request(wire::PromptRequest::new( + first.session_id.clone(), + vec![wire::ContentBlock::Text(wire::TextContent::new("first"))], + )) + .block_task(); + let competing_prompt = cx + .send_request(wire::PromptRequest::new( + first.session_id.clone(), + vec![wire::ContentBlock::Text(wire::TextContent::new( + "competing", + ))], + )) + .block_task(); + let (first_result, competing_result) = + tokio::time::timeout(Duration::from_millis(250), async { + tokio::join!(first_prompt, competing_prompt) + }) + .await + .expect("same-session prompt admission timed out"); + assert_ne!( + first_result.is_ok(), + competing_result.is_ok(), + "exactly one same-session prompt must be admitted" + ); + + tokio::time::timeout( + Duration::from_millis(250), + cx.send_request(wire::PromptRequest::new( + second.session_id.clone(), + vec![wire::ContentBlock::Text(wire::TextContent::new("wait"))], + )) + .block_task(), + ) + .await + .expect("prompt acceptance was blocked by another session")?; + + cx.send_notification(wire::CancelSessionNotification::new( + first.session_id.clone(), + ))?; + wait_for_idle(&updates, &first.session_id).await; + assert!(!updates.lock().unwrap().iter().any(|(id, update)| { + id == &second.session_id + && matches!( + update, + wire::SessionUpdate::StateUpdate(wire::StateUpdate::Idle(_)) + ) + })); + + cx.send_notification(wire::CancelSessionNotification::new( + second.session_id.clone(), + ))?; + wait_for_idle(&updates, &second.session_id).await; + for session_id in [&first.session_id, &second.session_id] { + let stop_reason = updates.lock().unwrap().iter().find_map( + |(id, update)| match update { + wire::SessionUpdate::StateUpdate(wire::StateUpdate::Idle( + idle, + )) if id == session_id => idle.stop_reason.clone(), + _ => None, + }, + ); + assert_eq!(stop_reason, Some(wire::StopReason::Cancelled)); + } + cx.send_request(wire::CloseSessionRequest::new(first.session_id)) + .block_task() + .await?; + cx.send_request(wire::CloseSessionRequest::new(second.session_id)) + .block_task() + .await?; + Ok(()) + } + }); + + tokio::time::timeout(Duration::from_secs(5), client) + .await + .expect("client timed out") + .expect("client run"); + server.abort(); + let _ = server.await; + } + #[tokio::test] + async fn independent_prompts_are_accepted_immediately_and_cancel_separately() { + let updates = Arc::new(Mutex::new(Vec::new())); + let (client_transport, agent_transport) = Channel::duplex(); + let server = tokio::spawn(async move { + AcpHeadlessRuntime::::builder() + .agent_factory(TestFactory { + adapter: BlockingAdapter, + }) + .serve(agent_transport) + .await + }); + + let client = + agent_client_protocol::Client + .v2() + .on_receive_notification( + { + let updates = Arc::clone(&updates); + async move |notification: wire::UpdateSessionNotification, _cx| { + updates + .lock() + .unwrap() + .push((notification.session_id, notification.update)); + Ok(()) + } + }, + agent_client_protocol::on_receive_notification!(), + ) + .connect_with(client_transport, { + let updates = Arc::clone(&updates); + async move |cx| { + cx.send_request(wire::InitializeRequest::new( + wire::ProtocolVersion::V2, + wire::Implementation::new("test-client", "1"), + )) + .block_task() + .await?; + let cwd = std::env::current_dir() + .map_err(agent_client_protocol::Error::into_internal_error)?; + let first = cx + .send_request(wire::NewSessionRequest::new(cwd.clone())) + .block_task() + .await?; + let second = cx + .send_request(wire::NewSessionRequest::new(cwd)) + .block_task() + .await?; + + for session_id in [&first.session_id, &second.session_id] { + tokio::time::timeout( + Duration::from_millis(250), + cx.send_request(wire::PromptRequest::new( + session_id.clone(), + vec![wire::ContentBlock::Text(wire::TextContent::new("wait"))], + )) + .block_task(), + ) + .await + .expect("prompt acceptance was blocked by another session")?; + } + + cx.send_notification(wire::CancelSessionNotification::new( + first.session_id.clone(), + ))?; + wait_for_idle(&updates, &first.session_id).await; + assert!(!updates.lock().unwrap().iter().any(|(id, update)| { + id == &second.session_id + && matches!( + update, + wire::SessionUpdate::StateUpdate(wire::StateUpdate::Idle(_)) + ) + })); + + cx.send_notification(wire::CancelSessionNotification::new( + second.session_id.clone(), + ))?; + wait_for_idle(&updates, &second.session_id).await; + for session_id in [&first.session_id, &second.session_id] { + let stop_reason = updates.lock().unwrap().iter().find_map( + |(id, update)| match update { + wire::SessionUpdate::StateUpdate(wire::StateUpdate::Idle( + idle, + )) if id == session_id => idle.stop_reason.clone(), + _ => None, + }, + ); + assert_eq!(stop_reason, Some(wire::StopReason::Cancelled)); + } + cx.send_request(wire::CloseSessionRequest::new(first.session_id)) + .block_task() + .await?; + cx.send_request(wire::CloseSessionRequest::new(second.session_id)) + .block_task() + .await?; + Ok(()) + } + }); + + tokio::time::timeout(Duration::from_secs(5), client) + .await + .expect("client timed out") + .expect("client run"); + server.abort(); + let _ = server.await; + } + + #[test] + fn tool_execution_boundary_rotates_message_ids_without_a_terminal_result() { + let integration = AcpIntegration::default(); + let (client, mut messages) = ClientHandle::channel(); + let acp_id = wire::SessionId::new("acp-session"); + let agentkit_id = AgentkitSessionId::new("agentkit-session"); + integration + .bind(acp_id.clone(), agentkit_id.clone(), client) + .expect("bind session"); + integration + .install_prompt_state( + &acp_id, + Arc::new(AtomicU64::new(1)), + Arc::new(Mutex::new(())), + ) + .expect("install prompt state"); + let cancellation = CancellationController::new(); + integration + .begin_prompt_owner(&acp_id, 1, cancellation.handle().checkpoint()) + .expect("begin prompt"); + + for event in [ + AgentEvent::ContentDelta(Delta::BeginPart { + part_id: PartId::new("before-tool"), + kind: PartKind::Text, + }), + AgentEvent::ContentDelta(Delta::AppendText { + part_id: PartId::new("before-tool"), + chunk: "before".into(), + }), + AgentEvent::ToolExecutionStarted(ToolCallPart::new( + ToolCallId::new("call"), + "background_tool", + json!({}), + )), + AgentEvent::ContentDelta(Delta::BeginPart { + part_id: PartId::new("after-tool"), + kind: PartKind::Text, + }), + AgentEvent::ContentDelta(Delta::AppendText { + part_id: PartId::new("after-tool"), + chunk: "after".into(), + }), + ] { + integration.route_event(&agentkit_id, event); + } + + let message_ids = std::iter::from_fn(|| messages.try_recv().ok()) + .filter_map(|message| match message { + ClientMessage::Update(notification) => match notification.update { + wire::SessionUpdate::AgentMessageChunk(chunk) => Some(chunk.message_id), + _ => None, + }, + ClientMessage::AcknowledgedUpdate { .. } | ClientMessage::Flush(_) => None, + }) + .collect::>(); + assert_eq!( + message_ids, + [ + wire::MessageId::new("acp-session-agent-2"), + wire::MessageId::new("acp-session-agent-3"), + ] + ); + } + + #[test] + fn integration_rejects_duplicate_agentkit_session_ids() { + let integration = AcpIntegration::default(); + let (first_client, _first_rx) = ClientHandle::channel(); + let (second_client, _second_rx) = ClientHandle::channel(); + let agentkit_id = AgentkitSessionId::new("shared-agentkit-session"); + let first_acp = wire::SessionId::new("first-acp-session"); + let second_acp = wire::SessionId::new("second-acp-session"); + integration + .bind(first_acp.clone(), agentkit_id.clone(), first_client) + .expect("first binding"); + let error = integration + .bind(second_acp.clone(), agentkit_id, second_client) + .expect_err("duplicate AgentKit session id must be rejected"); + assert!(matches!(error, AcpRuntimeError::SessionAlreadyBound(_))); + assert!(matches!( + integration.session(&second_acp), + Err(AcpRuntimeError::SessionNotFound(_)) + )); + integration + .unbind_session(&first_acp) + .expect("unbind first"); } } diff --git a/docs/acp.md b/docs/acp.md index bfa4c3e..ab8e62a 100644 --- a/docs/acp.md +++ b/docs/acp.md @@ -1,10 +1,16 @@ # agentkit-acp design > **Implementation status:** Root APIs and default features implement stable ACP -> v1. Version 0.10.8 also provides an opt-in ACP v2 runtime foundation under -> `agentkit_acp::v2`; enable it with `protocol-v2`. It uses only the official -> `agent-client-protocol` 2.0.0 `unstable_protocol_v2` feature. See the crate -> README or book chapter for the supported v2 lifecycle and current limits. +> v1. Version 0.10.11 also provides an opt-in ACP v2 runtime foundation under +> `agentkit_acp::v2`; enable it with `protocol-v2`. It currently uses a workspace- +> patched `agent-client-protocol` fork with `unstable_protocol_v2`. The additive +> `unstable-inject` feature enables steer-only injection and mandatory revoke. +> Receipt-committed injections survive cancellation for the next valid boundary; +> close may discard them. Delivered-ID history is lifetime-capped at 4,096 +> accepts per session. V2 running/idle state follows loop turn lifecycle, and +> autonomous task wakeups, cancellation, and logical prompt cleanup use that +> same authoritative state. See the crate README or book chapter for the +> supported v2 lifecycle and current limits. ## Purpose @@ -37,7 +43,7 @@ There is already an official Rust ACP SDK: - `agent-client-protocol-schema` - `agent-client-protocol-rmcp` -So agentkit should not build a parallel protocol crate by default. A separate `racp` crate only makes sense if the official SDK becomes unsuitable. The first implementation should mirror `agentkit-mcp`: re-export upstream wire types and make agentkit own only the lifecycle and conversion glue. +So agentkit should not build a parallel protocol crate by default. The current implementation uses a commit-pinned source fork of that SDK for experimental session injection; it still re-exports the SDK wire types and owns only the lifecycle and conversion glue. A separate `racp` crate would make sense only if the SDK API became unsuitable. ## Non-goals @@ -52,7 +58,7 @@ So agentkit should not build a parallel protocol crate by default. A separate `r - long-term session persistence - policy storage for remembered approval choices -The crate integrates agentkit into ACP. It should not fork ACP or turn agentkit into a single opinionated coding-agent product. +The crate integrates agentkit into ACP. Its pinned SDK source fork must not become a divergent ACP protocol, and the crate must not turn agentkit into a single opinionated coding-agent product. ## Dependencies @@ -60,7 +66,8 @@ Recommended initial crate: ```toml [dependencies] -agent-client-protocol = "2.0.0" +# Registry-shaped so every workspace consumer resolves one patched type identity. +agent-client-protocol = "=2.0.0" agentkit-core = { path = "../agentkit-core", version = "0.10.5" } agentkit-loop = { path = "../agentkit-loop", version = "0.10.7" } agentkit-tools-core = { path = "../agentkit-tools-core", version = "0.10.5" } @@ -71,11 +78,19 @@ thiserror = { workspace = true } tokio = { workspace = true, features = ["sync"] } tracing = { workspace = true } +[patch.crates-io] +# Intentionally unpublishable until these SDK APIs are released upstream. +agent-client-protocol = { git = "https://github.com/danielkov/rust-sdk", rev = "2f039993d1d6ed8da35b38c31f54a7cbb7338c70" } + [features] default = ["stdio"] stdio = [] unstable-acp = ["agent-client-protocol/unstable"] protocol-v2 = ["agent-client-protocol/unstable_protocol_v2"] +unstable-inject = [ + "protocol-v2", + "agent-client-protocol/unstable_session_inject", +] ``` The umbrella crate should later add: @@ -727,6 +742,6 @@ The key should be explicit and policy-owned. A good default key should include: - ACP introduction: - ACP Rust SDK: - ACP protocol docs: -- Rust crate docs: +- ACP Rust SDK fork used by this crate: - Existing MCP integration design: [`mcp.md`](./mcp.md) - Agentkit permission design: [`permissions.md`](./permissions.md) From cf2ef410cdfa6323f3f9294e53db3492f662ce3b Mon Sep 17 00:00:00 2001 From: daniel Date: Wed, 2 Sep 2026 11:40:23 +0100 Subject: [PATCH 2/7] fix(loop): preserve interrupted streamed responses --- crates/agentkit-loop/src/lib.rs | 648 ++++++++++++++++++++++++++++++-- 1 file changed, 625 insertions(+), 23 deletions(-) diff --git a/crates/agentkit-loop/src/lib.rs b/crates/agentkit-loop/src/lib.rs index ec8aef7..10b26f4 100644 --- a/crates/agentkit-loop/src/lib.rs +++ b/crates/agentkit-loop/src/lib.rs @@ -66,8 +66,8 @@ use std::sync::Arc; use agentkit_core::{ CancellationHandle, DataRef, Delta, FinishReason, Item, ItemKind, MetadataMap, Modality, Part, - SessionId, TaskId, TextPart, Timestamp, ToolCallId, ToolCallPart, ToolOutput, ToolResultPart, - TurnCancellation, Usage, + PartId, PartKind, ReasoningPart, SessionId, StructuredPart, TaskId, TextPart, Timestamp, + ToolCallId, ToolCallPart, ToolOutput, ToolResultPart, TurnCancellation, Usage, }; use agentkit_task_manager::{ PendingLoopUpdates, SimpleTaskManager, TOOL_RESULT_NOT_STARTED_METADATA_KEY, TaskApproval, @@ -2398,6 +2398,9 @@ where let mut saw_tool_call = false; let mut finished_result = None; let mut latest_usage = None; + let mut streamed_content = cancellation + .is_some() + .then(StreamedAssistantContent::default); while let Some(event) = match turn .next_event(cancellation.clone()) @@ -2410,10 +2413,22 @@ where .on_turn_interrupted(&turn_id) .await .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?; - return self.finish_cancelled(turn_id, interrupted_assistant_items()); + return self.finish_cancelled( + turn_id, + interrupted_stream_items(streamed_content.as_ref()), + ); } Err(error) => return Err(error), } { + let attempt_superseded = matches!(event, ModelTurnEvent::ResponseAttemptSuperseded); + if attempt_superseded { + saw_tool_call = false; + latest_usage = None; + if let Some(content) = &mut streamed_content { + content.clear(); + } + self.emit(AgentEvent::ResponseAttemptSuperseded); + } if cancellation .as_ref() .is_some_and(TurnCancellation::is_cancelled) @@ -2423,22 +2438,34 @@ where .on_turn_interrupted(&turn_id) .await .map_err(|error| LoopError::Tool(ToolError::Internal(error.to_string())))?; - return self.finish_cancelled(turn_id, interrupted_assistant_items()); + return self.finish_cancelled( + turn_id, + interrupted_stream_items(streamed_content.as_ref()), + ); + } + if attempt_superseded { + continue; } match event { - ModelTurnEvent::Delta(delta) => self.emit(AgentEvent::ContentDelta(delta)), + ModelTurnEvent::Delta(delta) => { + if let Some(content) = &mut streamed_content { + content.apply_delta(&delta); + } + self.emit(AgentEvent::ContentDelta(delta)); + } ModelTurnEvent::Usage(usage) => { latest_usage = Some(usage.clone()); self.emit(AgentEvent::UsageUpdated(usage)); } ModelTurnEvent::ToolCall(call) => { saw_tool_call = true; - self.emit(AgentEvent::ToolCallRequested(call.clone())); + if let Some(content) = &mut streamed_content { + content.commit_tool_call(call.clone()); + } + self.emit(AgentEvent::ToolCallRequested(call)); } ModelTurnEvent::ResponseAttemptSuperseded => { - saw_tool_call = false; - latest_usage = None; - self.emit(AgentEvent::ResponseAttemptSuperseded); + unreachable!("response-attempt supersession is handled before cancellation") } ModelTurnEvent::Finished(result) => { finished_result = Some(result); @@ -2703,8 +2730,8 @@ where ) -> Result { let pending = self.drain_pending_approval_items(); self.reject_drained_approvals(pending); - self.close_interrupted_tool_calls(); self.extend_transcript(items.clone()); + self.close_interrupted_tool_calls(); let turn_result = TurnResult { turn_id, finish_reason: FinishReason::Cancelled, @@ -4141,20 +4168,253 @@ fn main() { "#; } +#[derive(Default)] +struct StreamedAssistantContent { + parts: Vec, +} + +enum StreamedPart { + Open { + id: PartId, + kind: PartKind, + text: String, + bytes: Vec, + structured: Option, + metadata: MetadataMap, + }, + Committed(Part), +} + +impl StreamedAssistantContent { + fn apply_delta(&mut self, delta: &Delta) { + match delta { + Delta::BeginPart { part_id, kind } => { + let open = StreamedPart::Open { + id: part_id.clone(), + kind: *kind, + text: String::new(), + bytes: Vec::new(), + structured: None, + metadata: MetadataMap::new(), + }; + if let Some(slot) = self + .parts + .iter_mut() + .find(|slot| matches!(slot, StreamedPart::Open { id, .. } if id == part_id)) + { + *slot = open; + } else { + self.parts.push(open); + } + } + Delta::AppendText { part_id, chunk } => { + if let Some(StreamedPart::Open { text, .. }) = self.open_part_mut(part_id) { + text.push_str(chunk); + } + } + Delta::AppendBytes { part_id, chunk } => { + if let Some(StreamedPart::Open { kind, bytes, .. }) = self.open_part_mut(part_id) + && !matches!(kind, PartKind::Media) + { + bytes.extend_from_slice(chunk); + } + } + Delta::ReplaceStructured { part_id, value } => { + if let Some(StreamedPart::Open { structured, .. }) = self.open_part_mut(part_id) { + *structured = Some(value.clone()); + } + } + Delta::SetMetadata { part_id, metadata } => { + if let Some(StreamedPart::Open { + metadata: target, .. + }) = self.open_part_mut(part_id) + { + *target = metadata.clone(); + } + } + Delta::CommitPart { part } => self.commit_part(part.clone()), + } + } + + fn open_part_mut(&mut self, id: &PartId) -> Option<&mut StreamedPart> { + self.parts + .iter_mut() + .find(|part| matches!(part, StreamedPart::Open { id: open_id, .. } if open_id == id)) + } + + fn commit_part(&mut self, part: Part) { + if let Part::ToolCall(call) = &part + && let Some(slot) = self.parts.iter_mut().find(|slot| { + matches!(slot, StreamedPart::Committed(Part::ToolCall(existing)) if existing.id == call.id) + }) + { + *slot = StreamedPart::Committed(part); + return; + } + + let kind = part_kind(&part); + let matching_open = self + .parts + .iter() + .position(|slot| slot.open_matches_part(&part)); + let mut open_with_kind = self.parts.iter().enumerate().filter_map(|(index, slot)| { + matches!(slot, StreamedPart::Open { kind: open_kind, .. } if *open_kind == kind) + .then_some(index) + }); + let only_open_with_kind = match (open_with_kind.next(), open_with_kind.next()) { + (Some(index), None) => Some(index), + _ => None, + }; + if let Some(index) = matching_open.or(only_open_with_kind) { + self.parts[index] = StreamedPart::Committed(part); + } else { + self.parts.push(StreamedPart::Committed(part)); + } + } + + fn commit_tool_call(&mut self, call: ToolCallPart) { + self.commit_part(Part::ToolCall(call)); + } + + fn clear(&mut self) { + self.parts.clear(); + } + + fn interrupted_items(&self) -> Vec { + let parts = self + .parts + .iter() + .filter_map(StreamedPart::preserved_part) + .collect::>(); + if parts.is_empty() { + return interrupted_assistant_items(); + } + vec![Item::new(ItemKind::Assistant, parts).with_metadata(interrupted_metadata("assistant"))] + } +} + +impl StreamedPart { + fn open_matches_part(&self, committed: &Part) -> bool { + match (self, committed) { + ( + Self::Open { + kind: PartKind::Text, + text, + .. + }, + Part::Text(part), + ) => text == &part.text, + ( + Self::Open { + kind: PartKind::Reasoning, + text, + bytes, + .. + }, + Part::Reasoning(part), + ) => { + part.summary.as_deref() == (!text.is_empty()).then_some(text.as_str()) + && match &part.data { + Some(DataRef::InlineBytes(data)) => data == bytes, + None => bytes.is_empty(), + _ => false, + } + } + ( + Self::Open { + kind: PartKind::Structured, + structured: Some(value), + .. + }, + Part::Structured(part), + ) => value == &part.value, + _ => false, + } + } + + fn preserved_part(&self) -> Option { + match self { + Self::Committed(part) if useful_streamed_part(part) => Some(part.clone()), + Self::Committed(_) => None, + Self::Open { + kind, + text, + bytes, + structured, + metadata, + .. + } => match kind { + PartKind::Text if !text.is_empty() => Some(Part::Text( + TextPart::new(text).with_metadata(metadata.clone()), + )), + PartKind::Reasoning if !text.is_empty() || !bytes.is_empty() => { + Some(Part::Reasoning(ReasoningPart { + summary: (!text.is_empty()).then(|| text.clone()), + data: (!bytes.is_empty()).then(|| DataRef::inline_bytes(bytes.clone())), + redacted: false, + metadata: metadata.clone(), + })) + } + PartKind::Structured => structured.clone().map(|value| { + Part::Structured(StructuredPart::new(value).with_metadata(metadata.clone())) + }), + // BeginPart does not carry media modality or MIME type, so an + // uncommitted media part cannot be reconstructed faithfully. + PartKind::Media => None, + _ => None, + }, + } + } +} + +fn part_kind(part: &Part) -> PartKind { + match part { + Part::Text(_) => PartKind::Text, + Part::Media(_) => PartKind::Media, + Part::File(_) => PartKind::File, + Part::Structured(_) => PartKind::Structured, + Part::Reasoning(_) => PartKind::Reasoning, + Part::ToolCall(_) => PartKind::ToolCall, + Part::ToolResult(_) => PartKind::ToolResult, + Part::Custom(_) => PartKind::Custom, + } +} + +fn useful_streamed_part(part: &Part) -> bool { + match part { + Part::Text(text) => !text.text.is_empty(), + Part::Reasoning(reasoning) => { + reasoning + .summary + .as_ref() + .is_some_and(|text| !text.is_empty()) + || reasoning.data.is_some() + || reasoning.redacted + } + Part::ToolResult(_) => false, + _ => true, + } +} + +fn interrupted_stream_items(content: Option<&StreamedAssistantContent>) -> Vec { + content.map_or_else( + interrupted_assistant_items, + StreamedAssistantContent::interrupted_items, + ) +} + fn interrupted_assistant_items() -> Vec { - vec![Item { - id: None, - kind: ItemKind::Assistant, - parts: vec![Part::Text(TextPart { - text: "Previous assistant response was interrupted by the user before completion." - .into(), - metadata: interrupted_metadata("assistant"), - })], - metadata: interrupted_metadata("assistant"), - usage: None, - finish_reason: None, - created_at: None, - }] + vec![ + Item::new( + ItemKind::Assistant, + vec![Part::Text(TextPart { + text: "Previous assistant response was interrupted by the user before completion." + .into(), + metadata: interrupted_metadata("assistant"), + })], + ) + .with_metadata(interrupted_metadata("assistant")), + ] } /// Tool calls in `transcript` that no `tool_result` answers, in call order. @@ -4397,6 +4657,10 @@ mod tests { struct FakeAdapter; struct SupersedingAdapter; + struct InterruptedStreamAdapter { + controller: CancellationController, + scenario: InterruptedStreamScenario, + } struct SlowAdapter; struct RecordingAdapter { seen_descriptions: StdArc>>>, @@ -4409,6 +4673,10 @@ mod tests { struct SupersedingSession { supersession_enabled: bool, } + struct InterruptedStreamSession { + controller: CancellationController, + scenario: InterruptedStreamScenario, + } struct SlowSession; struct RecordingSession { seen_descriptions: StdArc>>>, @@ -4425,6 +4693,18 @@ mod tests { events: VecDeque, } + #[derive(Clone, Copy)] + enum InterruptedStreamScenario { + PreserveContent, + CancelOnSupersession, + } + + struct InterruptedStreamTurn { + controller: CancellationController, + events: VecDeque, + cancel_on_supersession: bool, + } + struct SlowTurn { emitted: bool, } @@ -4669,6 +4949,18 @@ mod tests { } } + #[async_trait] + impl ModelAdapter for InterruptedStreamAdapter { + type Session = InterruptedStreamSession; + + async fn start_session(&self, _config: SessionConfig) -> Result { + Ok(InterruptedStreamSession { + controller: self.controller.clone(), + scenario: self.scenario, + }) + } + } + #[async_trait] impl ModelAdapter for SlowAdapter { type Session = SlowSession; @@ -4740,6 +5032,89 @@ mod tests { } } + #[async_trait] + impl ModelSession for InterruptedStreamSession { + type Turn = InterruptedStreamTurn; + + async fn begin_turn( + &mut self, + _request: TurnRequest, + _cancellation: Option, + ) -> Result { + let events = match self.scenario { + InterruptedStreamScenario::PreserveContent => { + let mut reasoning_metadata = MetadataMap::new(); + reasoning_metadata.insert("provider.detail".into(), true.into()); + VecDeque::from([ + ModelTurnEvent::Delta(Delta::BeginPart { + part_id: PartId::new("text"), + kind: PartKind::Text, + }), + ModelTurnEvent::Delta(Delta::AppendText { + part_id: PartId::new("text"), + chunk: "partial answer".into(), + }), + ModelTurnEvent::Delta(Delta::CommitPart { + part: Part::text("partial answer"), + }), + ModelTurnEvent::Delta(Delta::BeginPart { + part_id: PartId::new("reasoning"), + kind: PartKind::Reasoning, + }), + ModelTurnEvent::Delta(Delta::AppendText { + part_id: PartId::new("reasoning"), + chunk: "partial thought".into(), + }), + ModelTurnEvent::Delta(Delta::SetMetadata { + part_id: PartId::new("reasoning"), + metadata: reasoning_metadata, + }), + ModelTurnEvent::Delta(Delta::BeginPart { + part_id: PartId::new("structured"), + kind: PartKind::Structured, + }), + ModelTurnEvent::Delta(Delta::ReplaceStructured { + part_id: PartId::new("structured"), + value: json!({ "complete": false }), + }), + ModelTurnEvent::Delta(Delta::BeginPart { + part_id: PartId::new("media"), + kind: PartKind::Media, + }), + ModelTurnEvent::Delta(Delta::AppendBytes { + part_id: PartId::new("media"), + chunk: vec![1, 2, 3], + }), + ModelTurnEvent::ToolCall(ToolCallPart::new( + "partial-call", + "unfinished-tool", + json!({}), + )), + ]) + } + InterruptedStreamScenario::CancelOnSupersession => VecDeque::from([ + ModelTurnEvent::Delta(Delta::BeginPart { + part_id: PartId::new("stale"), + kind: PartKind::Text, + }), + ModelTurnEvent::Delta(Delta::AppendText { + part_id: PartId::new("stale"), + chunk: "discard me".into(), + }), + ModelTurnEvent::ResponseAttemptSuperseded, + ]), + }; + Ok(InterruptedStreamTurn { + controller: self.controller.clone(), + events, + cancel_on_supersession: matches!( + self.scenario, + InterruptedStreamScenario::CancelOnSupersession + ), + }) + } + } + #[async_trait] impl ModelSession for FakeSession { type Turn = FakeTurn; @@ -5059,6 +5434,25 @@ mod tests { } } + #[async_trait] + impl ModelTurn for InterruptedStreamTurn { + async fn next_event( + &mut self, + _cancellation: Option, + ) -> Result, LoopError> { + if let Some(event) = self.events.pop_front() { + if self.cancel_on_supersession + && matches!(event, ModelTurnEvent::ResponseAttemptSuperseded) + { + self.controller.interrupt(); + } + return Ok(Some(event)); + } + self.controller.interrupt(); + Err(LoopError::Cancelled) + } + } + #[async_trait] impl ModelTurn for SlowTurn { async fn next_event( @@ -5525,6 +5919,214 @@ mod tests { assert!(tool_call < superseded); } + #[tokio::test] + async fn cancellation_preserves_streamed_content_in_result_and_transcript() { + let controller = CancellationController::new(); + let agent = Agent::builder() + .model(InterruptedStreamAdapter { + controller: controller.clone(), + scenario: InterruptedStreamScenario::PreserveContent, + }) + .cancellation(controller.handle()) + .build() + .unwrap(); + let mut driver = agent + .start(SessionConfig::new("preserve-interrupted-stream")) + .await + .unwrap(); + driver + .submit_input(vec![Item::text(ItemKind::User, "start")]) + .unwrap(); + + let LoopStep::Finished(result) = run_until_finished(&mut driver).await else { + panic!("turn did not finish"); + }; + assert_eq!(result.finish_reason, FinishReason::Cancelled); + assert_eq!(result.items.len(), 1); + assert_eq!( + result.items[0].metadata.get(INTERRUPTED_METADATA_KEY), + Some(&Value::Bool(true)) + ); + assert_eq!(result.items[0].parts.len(), 4); + assert!(matches!( + &result.items[0].parts[0], + Part::Text(text) if text.text == "partial answer" + )); + assert!(matches!( + &result.items[0].parts[1], + Part::Reasoning(reasoning) + if reasoning.summary.as_deref() == Some("partial thought") + && reasoning.metadata.get("provider.detail") == Some(&Value::Bool(true)) + )); + assert!(matches!( + &result.items[0].parts[2], + Part::Structured(structured) if structured.value == json!({ "complete": false }) + )); + assert!(matches!( + &result.items[0].parts[3], + Part::ToolCall(call) if call.id == ToolCallId::new("partial-call") + )); + assert!( + !result.items[0] + .parts + .iter() + .any(|part| matches!(part, Part::Media(_))) + ); + assert!(!result.items[0].parts.iter().any(|part| { + matches!(part, Part::Text(text) if text.text.contains("response was interrupted")) + })); + + let transcript = driver.snapshot().transcript; + let preserved = transcript + .iter() + .find(|item| { + item.kind == ItemKind::Assistant + && item.metadata.get(INTERRUPTED_METADATA_KEY) == Some(&Value::Bool(true)) + }) + .expect("preserved assistant item in transcript"); + assert_eq!(preserved.parts, result.items[0].parts); + assert!(transcript.iter().any(|item| { + item.parts.iter().any(|part| { + matches!( + part, + Part::ToolResult(tool_result) + if tool_result.call_id == ToolCallId::new("partial-call") + && tool_result.is_error + ) + }) + })); + validate_transcript_invariants(&transcript).unwrap(); + } + + #[test] + fn committed_part_matches_the_correct_interleaved_open_part() { + let mut content = StreamedAssistantContent::default(); + for delta in [ + Delta::BeginPart { + part_id: PartId::new("first"), + kind: PartKind::Text, + }, + Delta::AppendText { + part_id: PartId::new("first"), + chunk: "one".into(), + }, + Delta::BeginPart { + part_id: PartId::new("second"), + kind: PartKind::Text, + }, + Delta::AppendText { + part_id: PartId::new("second"), + chunk: "two".into(), + }, + Delta::CommitPart { + part: Part::text("two"), + }, + ] { + content.apply_delta(&delta); + } + + let items = content.interrupted_items(); + assert!(matches!( + items[0].parts.as_slice(), + [Part::Text(first), Part::Text(second)] + if first.text == "one" && second.text == "two" + )); + } + + #[test] + fn ambiguous_committed_part_does_not_duplicate_identical_open_parts() { + let mut content = StreamedAssistantContent::default(); + for part_id in [PartId::new("first"), PartId::new("second")] { + content.apply_delta(&Delta::BeginPart { + part_id: part_id.clone(), + kind: PartKind::Text, + }); + content.apply_delta(&Delta::AppendText { + part_id, + chunk: "same".into(), + }); + } + content.apply_delta(&Delta::CommitPart { + part: Part::text("same"), + }); + + let items = content.interrupted_items(); + assert_eq!(items[0].parts.len(), 2); + assert!( + items[0] + .parts + .iter() + .all(|part| matches!(part, Part::Text(text) if text.text == "same")) + ); + } + + #[test] + fn deltas_without_begin_part_are_not_guessed() { + let mut content = StreamedAssistantContent::default(); + content.apply_delta(&Delta::AppendText { + part_id: PartId::new("missing"), + chunk: "orphan".into(), + }); + + let items = content.interrupted_items(); + assert!(matches!( + items[0].parts.as_slice(), + [Part::Text(text)] if text.text.contains("response was interrupted") + )); + } + + #[tokio::test] + async fn supersession_is_processed_before_concurrent_cancellation() { + let controller = CancellationController::new(); + let events = StdArc::new(StdMutex::new(Vec::new())); + let agent = Agent::builder() + .model(InterruptedStreamAdapter { + controller: controller.clone(), + scenario: InterruptedStreamScenario::CancelOnSupersession, + }) + .cancellation(controller.handle()) + .observer(RecordingObserver { + events: events.clone(), + }) + .build() + .unwrap(); + let mut driver = agent + .start(SessionConfig::new("supersede-then-cancel").with_response_attempt_supersession()) + .await + .unwrap(); + driver + .submit_input(vec![Item::text(ItemKind::User, "start")]) + .unwrap(); + + let LoopStep::Finished(result) = run_until_finished(&mut driver).await else { + panic!("turn did not finish"); + }; + assert_eq!(result.finish_reason, FinishReason::Cancelled); + assert_eq!(result.items.len(), 1); + assert!(matches!( + result.items[0].parts.as_slice(), + [Part::Text(text)] + if text.text.contains("response was interrupted") + && !text.text.contains("discard me") + )); + assert!(!driver.snapshot().transcript.iter().any(|item| { + item.parts + .iter() + .any(|part| matches!(part, Part::Text(text) if text.text.contains("discard me"))) + })); + + let events = events.lock().unwrap(); + let superseded = events + .iter() + .position(|event| matches!(event, AgentEvent::ResponseAttemptSuperseded)) + .expect("supersession event"); + let finished = events + .iter() + .position(|event| matches!(event, AgentEvent::TurnFinished(_))) + .expect("turn-finished event"); + assert!(superseded < finished); + } + fn turn_lifecycle_events( events: &[AgentEvent], ) -> Vec<(agentkit_core::TurnId, Option)> { From fbd0b81ffb8b03e7f3dd03a838b539750c79a33b Mon Sep 17 00:00:00 2001 From: daniel Date: Wed, 2 Sep 2026 14:37:47 +0100 Subject: [PATCH 3/7] fix(loop): bound interrupted response buffer --- crates/agentkit-loop/src/lib.rs | 488 +++++++++++++++++++++++++++++--- 1 file changed, 450 insertions(+), 38 deletions(-) diff --git a/crates/agentkit-loop/src/lib.rs b/crates/agentkit-loop/src/lib.rs index 10b26f4..4a2ac04 100644 --- a/crates/agentkit-loop/src/lib.rs +++ b/crates/agentkit-loop/src/lib.rs @@ -103,6 +103,8 @@ const USER_CANCELLED_REASON: &str = "user_cancelled"; const DETACHED_NOTIFICATION_TEXT_MAX_CHARS: usize = 512; const DETACHED_TEXT_PREVIEW_MAX_CHARS: usize = 160; const DETACHED_CALL_ID_MAX_CHARS: usize = 80; +const MAX_STREAMED_ASSISTANT_CONTENT_BYTES: usize = 1024 * 1024; +const MAX_STREAMED_ASSISTANT_CONTENT_PARTS: usize = 256; /// Metadata key used by adapters to retain provider-native finish reasons. pub const PROVIDER_FINISH_REASONS_METADATA_KEY: &str = "agentkit.provider_finish_reasons"; @@ -2425,7 +2427,7 @@ where saw_tool_call = false; latest_usage = None; if let Some(content) = &mut streamed_content { - content.clear(); + content.reset(); } self.emit(AgentEvent::ResponseAttemptSuperseded); } @@ -2460,7 +2462,7 @@ where ModelTurnEvent::ToolCall(call) => { saw_tool_call = true; if let Some(content) = &mut streamed_content { - content.commit_tool_call(call.clone()); + content.commit_tool_call(&call); } self.emit(AgentEvent::ToolCallRequested(call)); } @@ -4171,6 +4173,8 @@ fn main() { #[derive(Default)] struct StreamedAssistantContent { parts: Vec, + retained_bytes: usize, + overflowed: bool, } enum StreamedPart { @@ -4180,83 +4184,160 @@ enum StreamedPart { text: String, bytes: Vec, structured: Option, + structured_bytes: usize, metadata: MetadataMap, + metadata_bytes: usize, + retained_bytes: usize, + }, + Committed { + part: Part, + retained_bytes: usize, }, - Committed(Part), } impl StreamedAssistantContent { fn apply_delta(&mut self, delta: &Delta) { + if self.overflowed { + return; + } + match delta { Delta::BeginPart { part_id, kind } => { + let index = self.open_part_index(part_id); + let retained_bytes = part_id.0.len(); + if !self.reserve_slot(index, retained_bytes) { + return; + } let open = StreamedPart::Open { id: part_id.clone(), kind: *kind, text: String::new(), bytes: Vec::new(), structured: None, + structured_bytes: 0, metadata: MetadataMap::new(), + metadata_bytes: 0, + retained_bytes, }; - if let Some(slot) = self - .parts - .iter_mut() - .find(|slot| matches!(slot, StreamedPart::Open { id, .. } if id == part_id)) - { - *slot = open; + if let Some(index) = index { + self.parts[index] = open; } else { self.parts.push(open); } } Delta::AppendText { part_id, chunk } => { - if let Some(StreamedPart::Open { text, .. }) = self.open_part_mut(part_id) { + let Some(index) = self.open_part_index(part_id) else { + return; + }; + if self.grow_slot(index, chunk.len()) + && let StreamedPart::Open { text, .. } = &mut self.parts[index] + { text.push_str(chunk); } } Delta::AppendBytes { part_id, chunk } => { - if let Some(StreamedPart::Open { kind, bytes, .. }) = self.open_part_mut(part_id) - && !matches!(kind, PartKind::Media) + let Some(index) = self.open_part_index(part_id) else { + return; + }; + if matches!( + self.parts[index], + StreamedPart::Open { + kind: PartKind::Media, + .. + } + ) { + return; + } + if self.grow_slot(index, chunk.len()) + && let StreamedPart::Open { bytes, .. } = &mut self.parts[index] { bytes.extend_from_slice(chunk); } } Delta::ReplaceStructured { part_id, value } => { - if let Some(StreamedPart::Open { structured, .. }) = self.open_part_mut(part_id) { + let Some(index) = self.open_part_index(part_id) else { + return; + }; + let old_bytes = match &self.parts[index] { + StreamedPart::Open { + structured_bytes, .. + } => *structured_bytes, + StreamedPart::Committed { .. } => unreachable!(), + }; + let limit = self.available_after_replacing(old_bytes); + let Some(new_bytes) = serialized_size_with_limit(value, limit) else { + self.overflow(); + return; + }; + self.replace_slot_bytes(index, old_bytes, new_bytes); + if let StreamedPart::Open { + structured, + structured_bytes, + .. + } = &mut self.parts[index] + { *structured = Some(value.clone()); + *structured_bytes = new_bytes; } } Delta::SetMetadata { part_id, metadata } => { - if let Some(StreamedPart::Open { - metadata: target, .. - }) = self.open_part_mut(part_id) + let Some(index) = self.open_part_index(part_id) else { + return; + }; + let old_bytes = match &self.parts[index] { + StreamedPart::Open { metadata_bytes, .. } => *metadata_bytes, + StreamedPart::Committed { .. } => unreachable!(), + }; + let limit = self.available_after_replacing(old_bytes); + let Some(new_bytes) = serialized_size_with_limit(metadata, limit) else { + self.overflow(); + return; + }; + self.replace_slot_bytes(index, old_bytes, new_bytes); + if let StreamedPart::Open { + metadata: target, + metadata_bytes, + .. + } = &mut self.parts[index] { *target = metadata.clone(); + *metadata_bytes = new_bytes; } } - Delta::CommitPart { part } => self.commit_part(part.clone()), + Delta::CommitPart { part } => self.commit_part(part), } } - fn open_part_mut(&mut self, id: &PartId) -> Option<&mut StreamedPart> { - self.parts - .iter_mut() - .find(|part| matches!(part, StreamedPart::Open { id: open_id, .. } if open_id == id)) + fn open_part_index(&self, id: &PartId) -> Option { + self.parts.iter().position( + |part| matches!(part, StreamedPart::Open { id: open_id, .. } if open_id == id), + ) } - fn commit_part(&mut self, part: Part) { - if let Part::ToolCall(call) = &part - && let Some(slot) = self.parts.iter_mut().find(|slot| { - matches!(slot, StreamedPart::Committed(Part::ToolCall(existing)) if existing.id == call.id) - }) - { - *slot = StreamedPart::Committed(part); + fn commit_part(&mut self, part: &Part) { + if self.overflowed { return; } - let kind = part_kind(&part); + let duplicate_tool_call = if let Part::ToolCall(call) = part { + self.parts.iter().position(|slot| { + matches!( + slot, + StreamedPart::Committed { + part: Part::ToolCall(existing), + .. + } if existing.id == call.id + ) + }) + } else { + None + }; + + let kind = part_kind(part); let matching_open = self .parts .iter() - .position(|slot| slot.open_matches_part(&part)); + .position(|slot| slot.open_matches_part(part)); let mut open_with_kind = self.parts.iter().enumerate().filter_map(|(index, slot)| { matches!(slot, StreamedPart::Open { kind: open_kind, .. } if *open_kind == kind) .then_some(index) @@ -4265,22 +4346,139 @@ impl StreamedAssistantContent { (Some(index), None) => Some(index), _ => None, }; - if let Some(index) = matching_open.or(only_open_with_kind) { - self.parts[index] = StreamedPart::Committed(part); + let index = duplicate_tool_call + .or(matching_open) + .or(only_open_with_kind); + let limit = self.available_for_slot(index); + let Some(retained_bytes) = serialized_size_with_limit(part, limit) else { + self.overflow(); + return; + }; + if !self.reserve_slot(index, retained_bytes) { + return; + } + let committed = StreamedPart::Committed { + part: part.clone(), + retained_bytes, + }; + if let Some(index) = index { + self.parts[index] = committed; + } else { + self.parts.push(committed); + } + } + + fn commit_tool_call(&mut self, call: &ToolCallPart) { + if self.overflowed { + return; + } + let duplicate_tool_call = self.parts.iter().position(|slot| { + matches!( + slot, + StreamedPart::Committed { + part: Part::ToolCall(existing), + .. + } if existing.id == call.id + ) + }); + let mut open_tool_calls = self.parts.iter().enumerate().filter_map(|(index, slot)| { + matches!( + slot, + StreamedPart::Open { + kind: PartKind::ToolCall, + .. + } + ) + .then_some(index) + }); + let only_open_tool_call = match (open_tool_calls.next(), open_tool_calls.next()) { + (Some(index), None) => Some(index), + _ => None, + }; + let index = duplicate_tool_call.or(only_open_tool_call); + let limit = self.available_for_slot(index); + let borrowed_part = BorrowedPart::ToolCall(call); + let Some(retained_bytes) = serialized_size_with_limit(&borrowed_part, limit) else { + self.overflow(); + return; + }; + if !self.reserve_slot(index, retained_bytes) { + return; + } + let committed = StreamedPart::Committed { + part: Part::ToolCall(call.clone()), + retained_bytes, + }; + if let Some(index) = index { + self.parts[index] = committed; } else { - self.parts.push(StreamedPart::Committed(part)); + self.parts.push(committed); + } + } + + fn reserve_slot(&mut self, index: Option, retained_bytes: usize) -> bool { + if index.is_none() && self.parts.len() >= MAX_STREAMED_ASSISTANT_CONTENT_PARTS { + self.overflow(); + return false; } + let replaced_bytes = index.map_or(0, |index| self.parts[index].retained_bytes()); + let Some(total) = self + .retained_bytes + .checked_sub(replaced_bytes) + .and_then(|bytes| bytes.checked_add(retained_bytes)) + .filter(|bytes| *bytes <= MAX_STREAMED_ASSISTANT_CONTENT_BYTES) + else { + self.overflow(); + return false; + }; + self.retained_bytes = total; + true } - fn commit_tool_call(&mut self, call: ToolCallPart) { - self.commit_part(Part::ToolCall(call)); + fn grow_slot(&mut self, index: usize, additional_bytes: usize) -> bool { + let Some(total) = self + .retained_bytes + .checked_add(additional_bytes) + .filter(|bytes| *bytes <= MAX_STREAMED_ASSISTANT_CONTENT_BYTES) + else { + self.overflow(); + return false; + }; + self.retained_bytes = total; + *self.parts[index].retained_bytes_mut() += additional_bytes; + true + } + + fn available_after_replacing(&self, replaced_bytes: usize) -> usize { + MAX_STREAMED_ASSISTANT_CONTENT_BYTES - (self.retained_bytes - replaced_bytes) + } + + fn available_for_slot(&self, index: Option) -> usize { + self.available_after_replacing(index.map_or(0, |index| self.parts[index].retained_bytes())) } - fn clear(&mut self) { + fn replace_slot_bytes(&mut self, index: usize, old_bytes: usize, new_bytes: usize) { + self.retained_bytes = self.retained_bytes - old_bytes + new_bytes; + let retained_bytes = self.parts[index].retained_bytes_mut(); + *retained_bytes = *retained_bytes - old_bytes + new_bytes; + } + + fn overflow(&mut self) { + self.parts.clear(); + self.retained_bytes = 0; + self.overflowed = true; + } + + fn reset(&mut self) { self.parts.clear(); + self.retained_bytes = 0; + self.overflowed = false; } fn interrupted_items(&self) -> Vec { + if self.overflowed { + return interrupted_assistant_items(); + } let parts = self .parts .iter() @@ -4294,6 +4492,22 @@ impl StreamedAssistantContent { } impl StreamedPart { + fn retained_bytes(&self) -> usize { + match self { + Self::Open { retained_bytes, .. } | Self::Committed { retained_bytes, .. } => { + *retained_bytes + } + } + } + + fn retained_bytes_mut(&mut self) -> &mut usize { + match self { + Self::Open { retained_bytes, .. } | Self::Committed { retained_bytes, .. } => { + retained_bytes + } + } + } + fn open_matches_part(&self, committed: &Part) -> bool { match (self, committed) { ( @@ -4334,8 +4548,8 @@ impl StreamedPart { fn preserved_part(&self) -> Option { match self { - Self::Committed(part) if useful_streamed_part(part) => Some(part.clone()), - Self::Committed(_) => None, + Self::Committed { part, .. } if useful_streamed_part(part) => Some(part.clone()), + Self::Committed { .. } => None, Self::Open { kind, text, @@ -4367,6 +4581,42 @@ impl StreamedPart { } } +#[derive(Serialize)] +enum BorrowedPart<'a> { + ToolCall(&'a ToolCallPart), +} + +struct PayloadSizeWriter { + remaining: usize, + written: usize, +} + +impl std::io::Write for PayloadSizeWriter { + fn write(&mut self, bytes: &[u8]) -> std::io::Result { + if bytes.len() > self.remaining { + return Err(std::io::Error::other( + "streamed assistant payload limit exceeded", + )); + } + self.remaining -= bytes.len(); + self.written += bytes.len(); + Ok(bytes.len()) + } + + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } +} + +fn serialized_size_with_limit(value: &T, limit: usize) -> Option { + let mut writer = PayloadSizeWriter { + remaining: limit, + written: 0, + }; + serde_json::to_writer(&mut writer, value).ok()?; + Some(writer.written) +} + fn part_kind(part: &Part) -> PartKind { match part { Part::Text(_) => PartKind::Text, @@ -4697,6 +4947,7 @@ mod tests { enum InterruptedStreamScenario { PreserveContent, CancelOnSupersession, + OverflowThenSupersession, } struct InterruptedStreamTurn { @@ -5103,6 +5354,25 @@ mod tests { }), ModelTurnEvent::ResponseAttemptSuperseded, ]), + InterruptedStreamScenario::OverflowThenSupersession => VecDeque::from([ + ModelTurnEvent::Delta(Delta::BeginPart { + part_id: PartId::new("oversized"), + kind: PartKind::Text, + }), + ModelTurnEvent::Delta(Delta::AppendText { + part_id: PartId::new("oversized"), + chunk: "x".repeat(MAX_STREAMED_ASSISTANT_CONTENT_BYTES), + }), + ModelTurnEvent::ResponseAttemptSuperseded, + ModelTurnEvent::Delta(Delta::BeginPart { + part_id: PartId::new("fresh"), + kind: PartKind::Text, + }), + ModelTurnEvent::Delta(Delta::AppendText { + part_id: PartId::new("fresh"), + chunk: "preserve me".into(), + }), + ]), }; Ok(InterruptedStreamTurn { controller: self.controller.clone(), @@ -6075,6 +6345,117 @@ mod tests { )); } + #[test] + fn streamed_content_byte_budget_covers_every_retained_payload() { + fn open_content(kind: PartKind) -> StreamedAssistantContent { + let mut content = StreamedAssistantContent::default(); + content.apply_delta(&Delta::BeginPart { + part_id: PartId::new("part"), + kind, + }); + content + } + + let mut content = open_content(PartKind::Text); + content.apply_delta(&Delta::AppendText { + part_id: PartId::new("part"), + chunk: "x".repeat(MAX_STREAMED_ASSISTANT_CONTENT_BYTES), + }); + assert!(content.overflowed); + + let mut content = open_content(PartKind::Reasoning); + content.apply_delta(&Delta::AppendBytes { + part_id: PartId::new("part"), + chunk: vec![0; MAX_STREAMED_ASSISTANT_CONTENT_BYTES], + }); + assert!(content.overflowed); + + let mut content = open_content(PartKind::Structured); + content.apply_delta(&Delta::ReplaceStructured { + part_id: PartId::new("part"), + value: Value::String("x".repeat(MAX_STREAMED_ASSISTANT_CONTENT_BYTES)), + }); + assert!(content.overflowed); + + let mut metadata = MetadataMap::new(); + metadata.insert( + "large".into(), + Value::String("x".repeat(MAX_STREAMED_ASSISTANT_CONTENT_BYTES)), + ); + let mut content = open_content(PartKind::Text); + content.apply_delta(&Delta::SetMetadata { + part_id: PartId::new("part"), + metadata, + }); + assert!(content.overflowed); + + let mut content = StreamedAssistantContent::default(); + content.apply_delta(&Delta::CommitPart { + part: Part::text("x".repeat(MAX_STREAMED_ASSISTANT_CONTENT_BYTES)), + }); + assert!(content.overflowed); + + let mut content = StreamedAssistantContent::default(); + content.commit_tool_call(&ToolCallPart::new( + "call", + "tool", + Value::String("x".repeat(MAX_STREAMED_ASSISTANT_CONTENT_BYTES)), + )); + assert!(content.overflowed); + } + + #[test] + fn streamed_content_part_budget_releases_all_parts() { + let mut content = StreamedAssistantContent::default(); + for index in 0..MAX_STREAMED_ASSISTANT_CONTENT_PARTS { + content.apply_delta(&Delta::BeginPart { + part_id: PartId::new(format!("part-{index}")), + kind: PartKind::Text, + }); + } + assert_eq!(content.parts.len(), MAX_STREAMED_ASSISTANT_CONTENT_PARTS); + + content.apply_delta(&Delta::BeginPart { + part_id: PartId::new("overflow"), + kind: PartKind::Text, + }); + + assert!(content.overflowed); + assert!(content.parts.is_empty()); + assert_eq!(content.retained_bytes, 0); + } + + #[test] + fn streamed_content_overflow_uses_generic_interruption_fallback() { + let mut content = StreamedAssistantContent::default(); + content.apply_delta(&Delta::BeginPart { + part_id: PartId::new("text"), + kind: PartKind::Text, + }); + content.apply_delta(&Delta::AppendText { + part_id: PartId::new("text"), + chunk: "release me".into(), + }); + content.apply_delta(&Delta::AppendText { + part_id: PartId::new("text"), + chunk: "x".repeat(MAX_STREAMED_ASSISTANT_CONTENT_BYTES), + }); + content.apply_delta(&Delta::BeginPart { + part_id: PartId::new("ignored after overflow"), + kind: PartKind::Text, + }); + + assert!(content.parts.is_empty()); + assert_eq!(content.retained_bytes, 0); + let items = content.interrupted_items(); + assert!(matches!( + items[0].parts.as_slice(), + [Part::Text(text)] + if text.text.contains("response was interrupted") + && !text.text.contains("release me") + )); + } + #[tokio::test] async fn supersession_is_processed_before_concurrent_cancellation() { let controller = CancellationController::new(); @@ -6127,6 +6508,37 @@ mod tests { assert!(superseded < finished); } + #[tokio::test] + async fn supersession_resets_streamed_content_budget() { + let controller = CancellationController::new(); + let agent = Agent::builder() + .model(InterruptedStreamAdapter { + controller: controller.clone(), + scenario: InterruptedStreamScenario::OverflowThenSupersession, + }) + .cancellation(controller.handle()) + .build() + .unwrap(); + let mut driver = agent + .start( + SessionConfig::new("overflow-then-supersede").with_response_attempt_supersession(), + ) + .await + .unwrap(); + driver + .submit_input(vec![Item::text(ItemKind::User, "start")]) + .unwrap(); + + let LoopStep::Finished(result) = run_until_finished(&mut driver).await else { + panic!("turn did not finish"); + }; + assert_eq!(result.finish_reason, FinishReason::Cancelled); + assert!(matches!( + result.items[0].parts.as_slice(), + [Part::Text(text)] if text.text == "preserve me" + )); + } + fn turn_lifecycle_events( events: &[AgentEvent], ) -> Vec<(agentkit_core::TurnId, Option)> { From 79d779d2ae2c3c4976bf605d278472becaa14b5a Mon Sep 17 00:00:00 2001 From: daniel Date: Sat, 5 Sep 2026 12:42:28 +0100 Subject: [PATCH 4/7] fix(loop): retire interrupted turns without resuming execution --- crates/agentkit-loop/src/lib.rs | 289 +++++++++++++++++++++++++++++++- 1 file changed, 285 insertions(+), 4 deletions(-) diff --git a/crates/agentkit-loop/src/lib.rs b/crates/agentkit-loop/src/lib.rs index 4a2ac04..33f7c43 100644 --- a/crates/agentkit-loop/src/lib.rs +++ b/crates/agentkit-loop/src/lib.rs @@ -2944,6 +2944,42 @@ where self.finish_cancelled(turn_id, Vec::new()).map(Some) } + /// Retire the active logical turn without resuming model or tool execution. + /// + /// Use this after an interrupt (including [`LoopInterrupt::AfterToolResult`]) + /// when the host wants to abandon the continuation. Pending approvals are + /// denied, foreground task cleanup is requested, and unanswered tool calls + /// receive cancellation results. Background work and queued input are kept; + /// the next call to [`Self::next`] can start a fresh logical turn. + /// + /// Returns `None` when no logical turn is active, including on repeated calls. + /// Otherwise emits exactly one cancelled [`AgentEvent::TurnFinished`]. + /// + /// # Errors + /// + /// Returns the first task cleanup error, after completing local cleanup and + /// emitting the cancelled terminal event. The turn is retired even on error; + /// a failing task manager may leave external tasks running. + pub async fn retire_interrupted_turn(&mut self) -> Result, LoopError> { + let Some(turn_id) = self.lifecycle.active_turn.clone() else { + return Ok(None); + }; + // This is a logical presentation id, not a task-manager turn id. + // At AfterToolResult the foreground task round has already completed. + self.pending_round_resume = None; + let cleanup = self.cleanup_interrupted_turn().await; + let result = TurnResult { + turn_id, + finish_reason: FinishReason::Cancelled, + items: Vec::new(), + usage: None, + metadata: interrupted_metadata("turn"), + }; + self.finish_logical_turn(&result); + cleanup?; + Ok(Some(result)) + } + /// Take a read-only snapshot of the driver's current transcript and input queue. pub fn snapshot(&self) -> LoopSnapshot { LoopSnapshot { @@ -3019,7 +3055,9 @@ where Ok(LoopStep::Finished(turn)) => self.finish_logical_turn(turn), Err(_) => { if let Some(turn_id) = self.lifecycle.active_turn.clone() { - self.recover_from_next_error().await; + if let Err(error) = self.cleanup_interrupted_turn().await { + tracing::debug!(%error, "failed to clean up turn after loop error"); + } self.finish_logical_turn(&TurnResult { turn_id, finish_reason: FinishReason::Error, @@ -3034,7 +3072,7 @@ where result } - async fn recover_from_next_error(&mut self) { + async fn cleanup_interrupted_turn(&mut self) -> Result<(), LoopError> { let mut seen_turns = HashSet::new(); let mut interrupted_turns = Vec::new(); if let Some(active) = self.active_tool_round.take() @@ -3055,13 +3093,20 @@ where } let pending = self.drain_pending_approval_items(); + let mut cleanup_error = None; for turn_id in interrupted_turns { - if let Err(error) = self.task_manager.on_turn_interrupted(&turn_id).await { - tracing::debug!(%error, %turn_id, "failed to clean up turn after loop error"); + if let Err(error) = self.task_manager.on_turn_interrupted(&turn_id).await + && cleanup_error.is_none() + { + cleanup_error = Some(LoopError::Tool(ToolError::Internal(error.to_string()))); } } self.reject_drained_approvals(pending); self.close_interrupted_tool_calls(); + match cleanup_error { + Some(error) => Err(error), + None => Ok(()), + } } async fn next_inner(&mut self) -> Result { @@ -9607,6 +9652,242 @@ mod tests { assert_eq!(seen[1], Some(override_cache)); } + // Count actual model requests, including continuations that emit no deltas. + struct RetirementCountingAdapter(StdArc); + struct RetirementCountingSession(StdArc); + + #[async_trait] + impl ModelAdapter for RetirementCountingAdapter { + type Session = RetirementCountingSession; + + async fn start_session(&self, _config: SessionConfig) -> Result { + Ok(RetirementCountingSession(self.0.clone())) + } + } + + #[async_trait] + impl ModelSession for RetirementCountingSession { + type Turn = FakeTurn; + + async fn begin_turn( + &mut self, + request: TurnRequest, + cancellation: Option, + ) -> Result { + self.0.fetch_add(1, Ordering::SeqCst); + FakeSession.begin_turn(request, cancellation).await + } + } + + #[tokio::test] + async fn retire_interrupted_turn_after_tool_result_starts_fresh_without_continuing() { + let calls = StdArc::new(AtomicUsize::new(0)); + let events = StdArc::new(StdMutex::new(Vec::new())); + let agent = Agent::builder() + .model(RetirementCountingAdapter(calls.clone())) + .add_tool_source(ToolRegistry::new().with(EchoTool::default())) + .permissions(AllowAllPermissions) + .observer(RecordingObserver { + events: events.clone(), + }) + .build() + .unwrap(); + let mut driver = agent + .start(SessionConfig::new("retire-after-tool")) + .await + .unwrap(); + assert!(driver.retire_interrupted_turn().await.unwrap().is_none()); + driver + .submit_input(vec![Item::text(ItemKind::User, "ping")]) + .unwrap(); + let turn_id = match driver.next().await.unwrap() { + LoopStep::Interrupt(LoopInterrupt::AfterToolResult(info)) => info.turn_id, + other => panic!("expected AfterToolResult, got {other:?}"), + }; + let transcript = driver.snapshot().transcript; + let queued = vec![Item::text(ItemKind::User, "fresh request")]; + driver.submit_input(queued.clone()).unwrap(); + let retired = driver.retire_interrupted_turn().await.unwrap().unwrap(); + assert_eq!(retired.turn_id, turn_id); + assert_eq!(retired.finish_reason, FinishReason::Cancelled); + assert!(retired.items.is_empty()); + assert!(driver.retire_interrupted_turn().await.unwrap().is_none()); + assert_eq!(calls.load(Ordering::SeqCst), 1); + assert_eq!(driver.snapshot().transcript, transcript); + assert_eq!(driver.snapshot().pending_input, queued); + assert!(driver.pending_round_resume.is_none()); + assert!(driver.lifecycle.active_turn.is_none()); + let fresh = match driver.next().await.unwrap() { + LoopStep::Finished(turn) => turn, + other => panic!("expected fresh turn to finish, got {other:?}"), + }; + assert_ne!(fresh.turn_id, turn_id); + assert_eq!(fresh.finish_reason, FinishReason::Completed); + assert_eq!(calls.load(Ordering::SeqCst), 2); + assert!(driver.snapshot().pending_input.is_empty()); + assert!( + driver + .snapshot() + .transcript + .iter() + .any(|item| { item.kind == ItemKind::User && item.parts == queued[0].parts }) + ); + assert!(matches!( + driver.next().await.unwrap(), + LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_)) + )); + assert!(driver.retire_interrupted_turn().await.unwrap().is_none()); + assert_eq!(calls.load(Ordering::SeqCst), 2); + let lifecycle = turn_lifecycle_events(&events.lock().unwrap()); + assert_eq!( + lifecycle, + vec![ + (turn_id.clone(), None), + (turn_id, Some(FinishReason::Cancelled)), + (fresh.turn_id.clone(), None), + (fresh.turn_id, Some(FinishReason::Completed)), + ] + ); + } + + #[tokio::test] + async fn retire_interrupted_turn_denies_pending_approval_even_on_cleanup_error() { + for fail_cleanup in [false, true] { + let calls = StdArc::new(AtomicUsize::new(0)); + let events = StdArc::new(StdMutex::new(Vec::new())); + let manager = TestTaskManager::new(SimpleTaskManager::new()); + let manager = if fail_cleanup { + manager.fail_interrupt("retirement cleanup failed") + } else { + manager + }; + let agent = Agent::builder() + .model(RetirementCountingAdapter(calls.clone())) + .add_tool_source(ToolRegistry::new().with(EchoTool::default())) + .permissions(ApproveFsReads) + .task_manager(manager) + .observer(RecordingObserver { + events: events.clone(), + }) + .build() + .unwrap(); + let mut driver = agent + .start(SessionConfig::new("retire-approval")) + .await + .unwrap(); + driver + .submit_input(vec![Item::text(ItemKind::User, "ping")]) + .unwrap(); + assert!(matches!( + driver.next().await.unwrap(), + LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(_)) + )); + let turn_id = driver.lifecycle.active_turn.clone().unwrap(); + let result = driver.retire_interrupted_turn().await; + if fail_cleanup { + assert!( + result + .unwrap_err() + .to_string() + .contains("retirement cleanup failed") + ); + } else { + assert_eq!( + result.unwrap().unwrap().finish_reason, + FinishReason::Cancelled + ); + } + assert!(driver.pending_approvals.is_empty()); + assert!(driver.pending_approval_order.is_empty()); + assert!(driver.active_tool_round.is_none()); + assert!(driver.lifecycle.active_turn.is_none()); + assert!(driver.retire_interrupted_turn().await.unwrap().is_none()); + assert!(matches!( + driver.next().await.unwrap(), + LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_)) + )); + assert_eq!(calls.load(Ordering::SeqCst), 1); + validate_transcript_invariants(&driver.snapshot().transcript).unwrap(); + let events = events.lock().unwrap(); + assert_eq!( + turn_lifecycle_events(&events), + vec![ + (turn_id.clone(), None), + (turn_id, Some(FinishReason::Cancelled)) + ] + ); + assert_eq!(events.iter().filter(|event| matches!(event, + AgentEvent::ToolResultReceived(result) if result.call_id == ToolCallId::new("call-1") && result.is_error + )).count(), 1); + } + } + + #[tokio::test] + async fn retire_interrupted_turn_preserves_detached_background_task() { + let calls = StdArc::new(AtomicUsize::new(0)); + let entered = StdArc::new(AtomicBool::new(false)); + let release = StdArc::new(Notify::new()); + let task_manager = AsyncTaskManager::new().routing(NameRoutingPolicy::new([( + "detaching-wait", + RoutingDecision::ForegroundThenDetachAfter(Duration::from_millis(10)), + )])); + let handle = task_manager.handle(); + let agent = Agent::builder() + .model(RetirementCountingAdapter(calls.clone())) + .add_tool_source(ToolRegistry::new().with(BlockingTool::new( + "detaching-wait", + entered.clone(), + release.clone(), + "background-done", + ))) + .permissions(AllowAllPermissions) + .task_manager(task_manager) + .build() + .unwrap(); + let mut driver = agent + .start(SessionConfig::new("retire-background")) + .await + .unwrap(); + driver + .submit_input(vec![Item::text(ItemKind::User, "ping")]) + .unwrap(); + assert!(matches!( + driver.next().await.unwrap(), + LoopStep::Interrupt(LoopInterrupt::AfterToolResult(_)) + )); + assert!(matches!( + wait_for_task_event(&handle).await, + TaskEvent::Started(_) + )); + assert!(matches!( + wait_for_task_event(&handle).await, + TaskEvent::Detached(_) + )); + wait_until_entered(entered.as_ref()).await; + let transcript = driver.snapshot().transcript; + assert_eq!( + driver + .retire_interrupted_turn() + .await + .unwrap() + .unwrap() + .finish_reason, + FinishReason::Cancelled + ); + assert_eq!(calls.load(Ordering::SeqCst), 1); + assert_eq!(driver.snapshot().transcript, transcript); + let running = handle.list_running().await; + assert_eq!(running.len(), 1); + assert_eq!(running[0].tool_name, "detaching-wait"); + release.notify_one(); + match wait_for_task_event(&handle).await { + TaskEvent::Completed(_, result) => { + assert_eq!(result.output, ToolOutput::Text("background-done".into())) + } + other => panic!("retired turn must preserve background completion, got {other:?}"), + } + } + #[tokio::test] async fn loop_yields_after_tool_result_between_rounds() { let tools = ToolRegistry::new().with(EchoTool::default()); From bc5d6b131f380528a64417c32118c6ee8b6acd91 Mon Sep 17 00:00:00 2001 From: daniel Date: Sat, 5 Sep 2026 13:29:58 +0100 Subject: [PATCH 5/7] feat(acp): replace pending session injections in place --- book/src/acp.md | 11 +- crates/agentkit-acp/README.md | 11 +- crates/agentkit-acp/src/v2.rs | 180 +++++++- crates/agentkit-acp/src/v2/replace_tests.rs | 448 ++++++++++++++++++++ 4 files changed, 625 insertions(+), 25 deletions(-) create mode 100644 crates/agentkit-acp/src/v2/replace_tests.rs diff --git a/book/src/acp.md b/book/src/acp.md index 48ad854..cac2810 100644 --- a/book/src/acp.md +++ b/book/src/acp.md @@ -54,9 +54,14 @@ JSON-RPC batches; the full `ContentBlock` list is emitted unchanged at the next safe model/tool boundary. Request or session cancellation cannot remove a response-committed steer, so it carries to the next valid boundary when the current turn is cancelled. Close may discard it. Revoke is serialized with -acknowledged user-message forwarding. Queueing and replacement are not -supported. Delivered IDs retain `already_delivered` classification for the -session lifetime; a 4,096-accept lifetime cap bounds that history. +acknowledged user-message forwarding. With `pending.replace` advertised, +`session/replace_inject` atomically replaces a pending steer's full content while +preserving its ID, mode, and queue position. It applies the same validation and +pending-byte budget as injection. Replacement waits for in-flight delivery; +delivered IDs return `already_delivered` (`-32010`), and revoked or unknown IDs +return `unknown_message_id` (`-32002`). Queueing is not supported. Delivered IDs +retain their classification for the session lifetime; a 4,096-accept lifetime +cap bounds that history. The initial v2 foundation routes text, reasoning, and tool lifecycle updates. ACP v2 permission callbacks are intentionally deferred; unsupported approval diff --git a/crates/agentkit-acp/README.md b/crates/agentkit-acp/README.md index 4ac26ae..f5903e4 100644 --- a/crates/agentkit-acp/README.md +++ b/crates/agentkit-acp/README.md @@ -51,9 +51,14 @@ same receipt-backed acceptance path. Once response commitment succeeds, request or session cancellation cannot remove the reservation; a committed steer that misses the cancelled turn carries into the next prompt. Closing the session may discard it. `session/revoke_inject` is mandatory and is serialized with -acknowledged `UserMessage` forwarding. Queueing, stream interruption, and -replacement are not supported. Delivered IDs retain their `already_delivered` -classification for the session lifetime. To bound that history safely, each +acknowledged `UserMessage` forwarding. The runtime advertises `pending.replace`: +`session/replace_inject` atomically replaces the full content of a pending steer, +preserving its ID, mode, and queue position. Replacement uses injection content +validation and the same pending-byte budget; rejected replacements leave the +original content intact. Replacement waits for in-flight delivery to settle: +delivered IDs return `already_delivered` (`-32010`), while revoked or unknown IDs +return `unknown_message_id` (`-32002`). Queueing and stream interruption are not +supported. Delivered IDs retain their classification for the session lifetime. To bound that history safely, each session accepts at most 4,096 injections and rejects later accepts with a lifetime-limit error. diff --git a/crates/agentkit-acp/src/v2.rs b/crates/agentkit-acp/src/v2.rs index fb7c138..7b3a0f2 100644 --- a/crates/agentkit-acp/src/v2.rs +++ b/crates/agentkit-acp/src/v2.rs @@ -1040,6 +1040,19 @@ where } }, agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + { + let integration = Arc::clone(&state.integration); + async move |request: wire::ReplaceInjectSessionRequest, responder, cx| { + let integration = Arc::clone(&integration); + cx.spawn(async move { + responder.respond_with_result(integration.replace_inject(request).await) + })?; + Ok(()) + } + }, + agent_client_protocol::on_receive_request!(), ); let connection = agent .on_receive_notification( @@ -1101,6 +1114,17 @@ pub const MAX_PENDING_INJECTION_BYTES: usize = 256 * 1024; #[cfg(feature = "unstable-inject")] pub const MAX_ACCEPTED_INJECTIONS: usize = 4_096; +#[cfg(feature = "unstable-inject")] +fn validate_inject_content( + content: &[wire::ContentBlock], +) -> Result<(Vec, usize), agent_client_protocol::Error> { + let items = content_blocks_to_items(content).map_err(crate::sdk_error)?; + let bytes = serde_json::to_vec(content) + .map_err(|error| agent_client_protocol::Error::new(-32603, error.to_string()))? + .len(); + Ok((items, bytes)) +} + #[cfg(feature = "unstable-inject")] struct PendingInject { message_id: wire::MessageId, @@ -1146,8 +1170,8 @@ enum BoundaryAction { } #[cfg(feature = "unstable-inject")] -enum RevokeTransition { - Revoked, +enum PendingTransition { + Applied, WaitForDelivery, AlreadyDelivered, Unknown, @@ -1289,7 +1313,49 @@ impl InjectionController { self.changed.notify_waiters(); } - fn revoke_transition(&self, message_id: &wire::MessageId) -> RevokeTransition { + fn replace_transition( + &self, + message_id: &wire::MessageId, + content: &[wire::ContentBlock], + items: &[Item], + bytes: usize, + ) -> Result { + let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + if let Some(index) = state + .pending + .iter() + .position(|pending| &pending.message_id == message_id) + { + if state.pending[index].commitment != InjectCommitment::Ready { + return Ok(PendingTransition::WaitForDelivery); + } + // Include other pending entries and any in-flight delivery in the + // budget. Check before mutation so failure preserves the old content. + let pending_bytes = state.pending_bytes - state.pending[index].bytes; + if pending_bytes.saturating_add(bytes) > MAX_PENDING_INJECTION_BYTES { + return Err(agent_client_protocol::Error::new( + -32602, + "pending session injection budget exceeded", + )); + } + let pending = &mut state.pending[index]; + pending.content = content.to_vec(); + pending.items = items.to_vec(); + pending.bytes = bytes; + state.pending_bytes = pending_bytes + bytes; + // Do not remove/re-enqueue: identity, commitment, and FIFO position + // must survive replacement. The state lock also serializes delivery. + Ok(PendingTransition::Applied) + } else if state.delivering.as_ref() == Some(message_id) { + Ok(PendingTransition::WaitForDelivery) + } else if state.delivered.contains(message_id) { + Ok(PendingTransition::AlreadyDelivered) + } else { + Ok(PendingTransition::Unknown) + } + } + + fn revoke_transition(&self, message_id: &wire::MessageId) -> PendingTransition { let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); let ready = state.pending.iter().any(|pending| { &pending.message_id == message_id && pending.commitment == InjectCommitment::Ready @@ -1297,18 +1363,18 @@ impl InjectionController { if ready && Self::remove_pending(&mut state, message_id) { drop(state); self.changed.notify_waiters(); - RevokeTransition::Revoked + PendingTransition::Applied } else if state.delivering.as_ref() == Some(message_id) || state .pending .iter() .any(|pending| &pending.message_id == message_id) { - RevokeTransition::WaitForDelivery + PendingTransition::WaitForDelivery } else if state.delivered.contains(message_id) { - RevokeTransition::AlreadyDelivered + PendingTransition::AlreadyDelivered } else { - RevokeTransition::Unknown + PendingTransition::Unknown } } @@ -1591,10 +1657,7 @@ impl AcpIntegration { "unsupported session injection mode", )); } - let items = content_blocks_to_items(&request.content).map_err(crate::sdk_error)?; - let bytes = serde_json::to_vec(&request.content) - .map_err(|error| agent_client_protocol::Error::new(-32603, error.to_string()))? - .len(); + let (items, bytes) = validate_inject_content(&request.content)?; let session = self .session(&request.session_id) .map_err(|_| session_not_found_error(&request.session_id))?; @@ -1690,6 +1753,57 @@ impl AcpIntegration { Ok(()) } + /// Atomically replaces the complete content of a pending injection without + /// changing its message ID, steer mode, or queue position. + /// + /// Like revoke, this waits for acceptance or in-flight delivery to settle. + /// A delivered ID returns `already_delivered`; revoked or unknown IDs return + /// `unknown_message_id`. Cancellation does not prevent replacing a retained + /// pending steer, but closing the session does. + pub async fn replace_inject( + &self, + request: wire::ReplaceInjectSessionRequest, + ) -> Result { + let (items, bytes) = validate_inject_content(&request.content)?; + let session = self + .session(&request.session_id) + .map_err(|_| session_not_found_error(&request.session_id))?; + loop { + let changed = session.injection.changed.notified(); + tokio::pin!(changed); + changed.as_mut().enable(); + let lifecycle = session + .lifecycle + .lock() + .unwrap_or_else(|error| error.into_inner()); + if session.closed.load(Ordering::Acquire) { + return Err(session_not_found_error(&request.session_id)); + } + match session.injection.replace_transition( + &request.message_id, + &request.content, + &items, + bytes, + )? { + PendingTransition::Applied => { + return Ok(wire::ReplaceInjectSessionResponse::new(request.message_id)); + } + PendingTransition::WaitForDelivery => drop(lifecycle), + PendingTransition::AlreadyDelivered => { + return Err(sdk_v2_error(wire::Error::inject_already_delivered( + request.message_id, + ))); + } + PendingTransition::Unknown => { + return Err(sdk_v2_error(wire::Error::inject_unknown_message_id( + request.message_id, + ))); + } + } + changed.await; + } + } + /// Handles the state transition for unstable ACP v2 `session/revoke_inject`. pub async fn revoke_inject( &self, @@ -1710,16 +1824,16 @@ impl AcpIntegration { return Err(session_not_found_error(&request.session_id)); } match session.injection.revoke_transition(&request.message_id) { - RevokeTransition::Revoked => { + PendingTransition::Applied => { return Ok(wire::RevokeInjectSessionResponse::new()); } - RevokeTransition::WaitForDelivery => drop(_lifecycle), - RevokeTransition::AlreadyDelivered => { + PendingTransition::WaitForDelivery => drop(_lifecycle), + PendingTransition::AlreadyDelivered => { return Err(sdk_v2_error(wire::Error::inject_already_delivered( request.message_id, ))); } - RevokeTransition::Unknown => { + PendingTransition::Unknown => { return Err(sdk_v2_error(wire::Error::inject_unknown_message_id( request.message_id, ))); @@ -2911,6 +3025,7 @@ fn finish_reason_to_stop_reason(reason: &FinishReason) -> wire::StopReason { pub fn session_inject_capabilities() -> wire::SessionInjectCapabilities { wire::SessionInjectCapabilities::new(vec![wire::SessionInjectMode::Steer]) .steer_in_stream(vec![wire::SessionInjectSteerInStream::Finish]) + .pending(wire::SessionInjectPendingCapabilities::new().replace(true)) } /// Constructs the honest ACP v2 capabilities supported by this build. @@ -6012,7 +6127,7 @@ mod tests { injection.activate(&message_id); assert!(matches!( injection.revoke_transition(&message_id), - RevokeTransition::Revoked + PendingTransition::Applied )); } @@ -6368,7 +6483,10 @@ mod tests { inject.steer_in_stream, Some(vec![wire::SessionInjectSteerInStream::Finish]) ); - assert!(inject.pending.is_none()); + assert_eq!( + inject.pending.and_then(|pending| pending.replace), + Some(true) + ); let cwd = std::env::current_dir() .map_err(agent_client_protocol::Error::into_internal_error)?; @@ -6399,8 +6517,15 @@ mod tests { )) .block_task() .await - .expect_err("replace handler must not be registered"); - assert_eq!(i32::from(replace_error.code), -32601); + .expect_err("unknown pending message must fail"); + assert_eq!(i32::from(replace_error.code), -32002); + assert_eq!( + replace_error.data, + Some(json!({ + "reason": "unknown_message_id", + "messageId": "not-pending", + })) + ); cx.send_request(wire::PromptRequest::new( session_id.clone(), @@ -6434,6 +6559,19 @@ mod tests { }), "UserMessage was emitted before the inject response" ); + let content = vec![ + wire::ContentBlock::Text(wire::TextContent::new("replacement")), + content[1].clone(), + ]; + let replaced = cx + .send_request(wire::ReplaceInjectSessionRequest::new( + session_id.clone(), + accepted.message_id.clone(), + content.clone(), + )) + .block_task() + .await?; + assert_eq!(replaced.message_id, accepted.message_id); adapter.release(2); wait_for_idle_count(&updates, &session_id, 1).await; let delivered = @@ -6941,3 +7079,7 @@ mod tests { .expect("unbind first"); } } + +#[cfg(all(test, feature = "unstable-inject"))] +#[path = "v2/replace_tests.rs"] +mod replace_tests; diff --git a/crates/agentkit-acp/src/v2/replace_tests.rs b/crates/agentkit-acp/src/v2/replace_tests.rs new file mode 100644 index 0000000..f108f4b --- /dev/null +++ b/crates/agentkit-acp/src/v2/replace_tests.rs @@ -0,0 +1,448 @@ +use super::*; +use std::time::Duration; + +fn text_content(text: &str) -> Vec { + vec![wire::ContentBlock::Text(wire::TextContent::new(text))] +} + +fn fixture() -> (AcpIntegration, AcpSessionHandle) { + let integration = AcpIntegration::default(); + let (sink, _receiver) = ClientHandle::channel(); + let session = integration + .bind_session(AcpSessionBinding::new( + wire::SessionId::new("replace-session"), + AgentkitSessionId::new("replace-loop"), + sink, + )) + .unwrap(); + session.prepare_injection_turn(); + session.start_injection_turn(); + (integration, session) +} + +fn reserve(integration: &AcpIntegration, session: &AcpSessionHandle, text: &str) -> ReservedInject { + integration + .reserve_inject(wire::InjectSessionRequest::new( + session.acp_session_id().clone(), + wire::SessionInjectMode::Steer, + text_content(text), + )) + .unwrap() +} + +fn accept(integration: &AcpIntegration, session: &AcpSessionHandle, text: &str) -> wire::MessageId { + let mut reserved = reserve(integration, session, text); + let id = reserved.message_id().clone(); + reserved.commit(); + reserved.activate(); + id +} + +fn request( + session: &AcpSessionHandle, + id: &wire::MessageId, + text: &str, +) -> wire::ReplaceInjectSessionRequest { + wire::ReplaceInjectSessionRequest::new( + session.acp_session_id().clone(), + id.clone(), + text_content(text), + ) +} + +fn take_delivery(injection: &InjectionController) -> PendingInject { + match injection.boundary_action(false, false) { + BoundaryAction::Deliver(pending) => pending, + _ => panic!("expected a ready injection at the boundary"), + } +} + +fn assert_text(pending: &PendingInject, expected: &str) { + assert_eq!(pending.content, text_content(expected)); + assert_eq!(pending.items.len(), 1); + assert_eq!(pending.items[0].kind, ItemKind::User); + assert_eq!(pending.items[0].parts.len(), 1); + assert!(matches!(&pending.items[0].parts[0], Part::Text(text) if text.text == expected)); + assert_eq!( + pending.bytes, + serde_json::to_vec(&pending.content).unwrap().len() + ); + assert_eq!(pending.commitment, InjectCommitment::Ready); +} + +fn assert_inject_error(error: agent_client_protocol::Error, reason: &str, id: &wire::MessageId) { + assert_eq!( + i32::from(error.code), + if reason == "already_delivered" { + -32010 + } else { + -32002 + } + ); + assert_eq!( + error.data, + Some(json!({ "reason": reason, "messageId": id })) + ); +} + +#[tokio::test] +async fn replacement_preserves_fifo_identity_and_acceptance_count() { + let (integration, session) = fixture(); + let first = accept(&integration, &session, "first"); + let second = accept(&integration, &session, "old second"); + let third = accept(&integration, &session, "third"); + let response = integration + .replace_inject(request(&session, &second, "new second")) + .await + .unwrap(); + assert_eq!(response.message_id, second); + integration + .replace_inject(request(&session, &second, "final second")) + .await + .unwrap(); + + let injection = &session.session.injection; + { + let state = injection.state.lock().unwrap(); + assert_eq!(state.accepted_count, 3); + assert_eq!(state.pending.len(), 3); + assert_eq!( + state.pending_bytes, + state.pending.iter().map(|p| p.bytes).sum::() + ); + } + for (id, text) in [(first, "first"), (second, "final second"), (third, "third")] { + let pending = take_delivery(injection); + assert_eq!(pending.message_id, id); + assert_text(&pending, text); + injection.finish_delivery(&pending, true); + } + let state = injection.state.lock().unwrap(); + assert!(state.pending.is_empty()); + assert_eq!(state.pending_bytes, 0); + assert_eq!(state.accepted_count, 3); +} + +#[tokio::test] +async fn replacement_grows_shrinks_and_rejects_over_budget_atomically() { + let (integration, session) = fixture(); + let first = accept(&integration, &session, "first"); + let second = accept(&integration, &session, "second"); + let injection = &session.session.injection; + let in_flight = take_delivery(injection); + assert_eq!(in_flight.message_id, first); + let overhead = serde_json::to_vec(&text_content("")).unwrap().len(); + let full = "x".repeat(MAX_PENDING_INJECTION_BYTES - in_flight.bytes - overhead); + integration + .replace_inject(request(&session, &second, &full)) + .await + .unwrap(); + assert_eq!( + injection.state.lock().unwrap().pending_bytes, + MAX_PENDING_INJECTION_BYTES + ); + + let error = integration + .replace_inject(request(&session, &second, &(full.clone() + "x"))) + .await + .unwrap_err(); + assert_eq!(i32::from(error.code), -32602); + { + let state = injection.state.lock().unwrap(); + assert_eq!(state.pending_bytes, MAX_PENDING_INJECTION_BYTES); + assert_eq!(state.accepted_count, 2); + assert_eq!(state.delivering.as_ref(), Some(&first)); + assert_eq!(state.pending.len(), 1); + assert_eq!(state.pending[0].message_id, second); + assert_text(&state.pending[0], &full); + } + + integration + .replace_inject(request(&session, &second, "small")) + .await + .unwrap(); + let small_bytes = serde_json::to_vec(&text_content("small")).unwrap().len(); + assert_eq!( + injection.state.lock().unwrap().pending_bytes, + in_flight.bytes + small_bytes + ); + injection.finish_delivery(&in_flight, true); + assert_eq!(injection.state.lock().unwrap().pending_bytes, small_bytes); + let pending = take_delivery(injection); + assert_text(&pending, "small"); + injection.finish_delivery(&pending, true); + assert_eq!(injection.state.lock().unwrap().pending_bytes, 0); +} + +#[tokio::test] +async fn replacement_waits_until_response_is_committed_and_activated() { + let (integration, session) = fixture(); + let mut reserved = reserve(&integration, &session, "original"); + let id = reserved.message_id().clone(); + let replacement = integration.replace_inject(request(&session, &id, "replacement")); + tokio::pin!(replacement); + assert!(futures_util::poll!(replacement.as_mut()).is_pending()); + assert!(matches!( + session.session.injection.boundary_action(false, false), + BoundaryAction::Wait + )); + assert_eq!( + session.session.injection.state.lock().unwrap().pending[0].content, + text_content("original") + ); + reserved.commit(); + assert!(futures_util::poll!(replacement.as_mut()).is_pending()); + assert_eq!( + session.session.injection.state.lock().unwrap().pending[0].commitment, + InjectCommitment::Committed + ); + reserved.activate(); + tokio::time::timeout(Duration::from_secs(1), replacement) + .await + .unwrap() + .unwrap(); + let pending = take_delivery(&session.session.injection); + assert_eq!(pending.message_id, id); + assert_text(&pending, "replacement"); + session.session.injection.finish_delivery(&pending, true); +} + +#[tokio::test] +async fn replacement_of_abandoned_reservation_wakes_with_unknown_id() { + let (integration, session) = fixture(); + let reserved = reserve(&integration, &session, "original"); + let id = reserved.message_id().clone(); + let replacement = integration.replace_inject(request(&session, &id, "replacement")); + tokio::pin!(replacement); + assert!(futures_util::poll!(replacement.as_mut()).is_pending()); + drop(reserved); + let error = tokio::time::timeout(Duration::from_secs(1), replacement) + .await + .unwrap() + .unwrap_err(); + assert_inject_error(error, "unknown_message_id", &id); + let state = session.session.injection.state.lock().unwrap(); + assert!(state.pending.is_empty()); + assert_eq!(state.pending_bytes, 0); + assert_eq!(state.accepted_count, 0); +} + +#[tokio::test] +async fn replacement_reports_unknown_revoked_and_delivered_ids() { + let (integration, session) = fixture(); + let unknown = wire::MessageId::new("never-accepted"); + let error = integration + .replace_inject(request(&session, &unknown, "new")) + .await + .unwrap_err(); + assert_inject_error(error, "unknown_message_id", &unknown); + + let revoked = accept(&integration, &session, "revoked"); + integration + .revoke_inject(wire::RevokeInjectSessionRequest::new( + session.acp_session_id().clone(), + revoked.clone(), + )) + .await + .unwrap(); + let error = integration + .replace_inject(request(&session, &revoked, "new")) + .await + .unwrap_err(); + assert_inject_error(error, "unknown_message_id", &revoked); + + let delivered = accept(&integration, &session, "delivered"); + let pending = take_delivery(&session.session.injection); + session.session.injection.finish_delivery(&pending, true); + let error = integration + .replace_inject(request(&session, &delivered, "new")) + .await + .unwrap_err(); + assert_inject_error(error, "already_delivered", &delivered); +} + +#[tokio::test] +async fn delivery_winner_blocks_replacement_until_delivery_outcome() { + for delivered in [true, false] { + let (integration, session) = fixture(); + let id = accept(&integration, &session, "original"); + let pending = take_delivery(&session.session.injection); + let replacement = integration.replace_inject(request(&session, &id, "replacement")); + tokio::pin!(replacement); + assert!(futures_util::poll!(replacement.as_mut()).is_pending()); + assert_text(&pending, "original"); + session + .session + .injection + .finish_delivery(&pending, delivered); + let error = tokio::time::timeout(Duration::from_secs(1), replacement) + .await + .unwrap() + .unwrap_err(); + assert_inject_error( + error, + if delivered { + "already_delivered" + } else { + "unknown_message_id" + }, + &id, + ); + assert_eq!( + session + .session + .injection + .state + .lock() + .unwrap() + .pending_bytes, + 0 + ); + } +} + +#[test] +fn simultaneous_delivery_and_replacement_have_one_linearized_outcome() { + for _ in 0..32 { + let (integration, session) = fixture(); + let id = accept(&integration, &session, "original"); + let injection = &session.session.injection; + let content = text_content("replacement"); + let items = content_blocks_to_items(&content).unwrap(); + let bytes = serde_json::to_vec(&content).unwrap().len(); + let barrier = std::sync::Barrier::new(2); + let (transition, pending) = std::thread::scope(|scope| { + let replacing = scope.spawn(|| { + barrier.wait(); + injection + .replace_transition(&id, &content, &items, bytes) + .unwrap() + }); + barrier.wait(); + let pending = take_delivery(injection); + (replacing.join().unwrap(), pending) + }); + assert_eq!(pending.message_id, id); + match transition { + PendingTransition::Applied => assert_text(&pending, "replacement"), + PendingTransition::WaitForDelivery => assert_text(&pending, "original"), + _ => panic!("replacement either wins or waits for the in-flight delivery"), + } + injection.finish_delivery(&pending, true); + let state = injection.state.lock().unwrap(); + assert!(state.pending.is_empty()); + assert_eq!(state.pending_bytes, 0); + assert_eq!(state.accepted_count, 1); + assert_eq!(state.delivered.len(), 1); + } +} + +#[tokio::test] +async fn replacement_survives_cancellation_and_delivers_in_next_turn() { + let (integration, session) = fixture(); + let id = accept(&integration, &session, "original"); + session.interrupt(); + integration + .replace_inject(request(&session, &id, "after cancellation")) + .await + .unwrap(); + assert!(matches!( + session.session.injection.boundary_action(false, false), + BoundaryAction::Complete(AcpInjectionBoundary::Stopped) + )); + session.prepare_injection_turn(); + session.start_injection_turn(); + let pending = take_delivery(&session.session.injection); + assert_eq!(pending.message_id, id); + assert_text(&pending, "after cancellation"); + session.session.injection.finish_delivery(&pending, true); +} + +#[tokio::test] +async fn close_wakes_waiting_replacements_and_discards_pending_bytes() { + for committed in [false, true] { + let (integration, session) = fixture(); + let mut reserved = reserve(&integration, &session, "original"); + let id = reserved.message_id().clone(); + if committed { + reserved.commit(); + } + let replacement = integration.replace_inject(request(&session, &id, "replacement")); + tokio::pin!(replacement); + assert!(futures_util::poll!(replacement.as_mut()).is_pending()); + session.close(); + let error = tokio::time::timeout(Duration::from_secs(1), replacement) + .await + .unwrap() + .unwrap_err(); + assert_eq!(i32::from(error.code), -32002); + assert_eq!( + error.data, + Some(json!({ "sessionId": session.acp_session_id() })) + ); + if committed { + reserved.activate(); + } else { + drop(reserved); + } + let state = session.session.injection.state.lock().unwrap(); + assert!(state.pending.is_empty()); + assert_eq!(state.pending_bytes, 0); + } + + let (integration, session) = fixture(); + let id = accept(&integration, &session, "original"); + integration + .replace_inject(request(&session, &id, "replacement")) + .await + .unwrap(); + session.close(); + let error = integration + .replace_inject(request(&session, &id, "too late")) + .await + .unwrap_err(); + assert_eq!(i32::from(error.code), -32002); + assert_eq!( + session + .session + .injection + .state + .lock() + .unwrap() + .pending_bytes, + 0 + ); +} + +#[tokio::test] +async fn replacement_reuses_inject_content_validation_without_mutating_queue() { + let (integration, session) = fixture(); + let id = accept(&integration, &session, "original"); + let invalid: wire::ContentBlock = serde_json::from_value(json!({ + "type": "future-unsupported-block", "value": "unsupported" + })) + .unwrap(); + let inject_error = match integration.reserve_inject(wire::InjectSessionRequest::new( + session.acp_session_id().clone(), + wire::SessionInjectMode::Steer, + vec![invalid.clone()], + )) { + Ok(_) => panic!("unsupported inject content must fail"), + Err(error) => error, + }; + let replace_error = integration + .replace_inject(wire::ReplaceInjectSessionRequest::new( + session.acp_session_id().clone(), + id.clone(), + vec![invalid], + )) + .await + .unwrap_err(); + assert_eq!(i32::from(replace_error.code), i32::from(inject_error.code)); + assert_eq!(replace_error.message, inject_error.message); + assert_eq!(replace_error.data, inject_error.data); + let pending = take_delivery(&session.session.injection); + assert_eq!(pending.message_id, id); + assert_text(&pending, "original"); + session.session.injection.finish_delivery(&pending, true); +} From f892ad82dd8c91cfc2b18712f6ff8a8aadad703d Mon Sep 17 00:00:00 2001 From: David Danialy Date: Fri, 18 Sep 2026 19:24:09 -0700 Subject: [PATCH 6/7] fix(acp): budget inline media separately for session injections --- crates/agentkit-acp/src/v2.rs | 150 +++++- .../src/v2/injection_budget_tests.rs | 446 ++++++++++++++++++ crates/agentkit-acp/src/v2/replace_tests.rs | 54 ++- 3 files changed, 615 insertions(+), 35 deletions(-) create mode 100644 crates/agentkit-acp/src/v2/injection_budget_tests.rs diff --git a/crates/agentkit-acp/src/v2.rs b/crates/agentkit-acp/src/v2.rs index 7b3a0f2..a70aa40 100644 --- a/crates/agentkit-acp/src/v2.rs +++ b/crates/agentkit-acp/src/v2.rs @@ -1103,7 +1103,8 @@ where /// Maximum number of pending unstable ACP v2 injections per session. #[cfg(feature = "unstable-inject")] pub const MAX_PENDING_INJECTIONS: usize = 64; -/// Maximum serialized content bytes retained by pending unstable injections. +/// Maximum serialized non-media content bytes retained by pending unstable injections. +/// Includes text, metadata, URIs, and JSON structure, but excludes inline payloads. #[cfg(feature = "unstable-inject")] pub const MAX_PENDING_INJECTION_BYTES: usize = 256 * 1024; /// Maximum accepted unstable injections tracked during one session lifetime. @@ -1114,14 +1115,106 @@ pub const MAX_PENDING_INJECTION_BYTES: usize = 256 * 1024; #[cfg(feature = "unstable-inject")] pub const MAX_ACCEPTED_INJECTIONS: usize = 4_096; +/// Maximum serialized inline image, audio, and embedded blob payload bytes per +/// session, including in-flight delivery. 28 MiB accommodates 20 MiB of source +/// media after base64 expansion and padding. Both original content and converted +/// Items retain payloads, so retained media can occupy twice this limit, plus +/// bounded metadata/data-URL overhead. This is not a transport or transcript cap. +#[cfg(feature = "unstable-inject")] +pub const MAX_PENDING_INJECTION_MEDIA_BYTES: usize = 28 * 1024 * 1024; + +#[cfg(feature = "unstable-inject")] +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +struct InjectionBytes { + content: usize, + media: usize, +} + +#[cfg(feature = "unstable-inject")] +impl InjectionBytes { + fn saturating_add(self, other: Self) -> Self { + Self { + content: self.content.saturating_add(other.content), + media: self.media.saturating_add(other.media), + } + } + + fn saturating_sub(self, other: Self) -> Self { + Self { + content: self.content.saturating_sub(other.content), + media: self.media.saturating_sub(other.media), + } + } + + fn exceeds_budget(self) -> bool { + self.content > MAX_PENDING_INJECTION_BYTES || self.media > MAX_PENDING_INJECTION_MEDIA_BYTES + } +} + +// Count serialized bytes without allocating a second serialized media buffer. +// Stop once even the combined budgets cannot accommodate the input. +#[cfg(feature = "unstable-inject")] +#[derive(Default)] +struct InjectionByteCounter(usize); + +#[cfg(feature = "unstable-inject")] +impl std::io::Write for InjectionByteCounter { + fn write(&mut self, bytes: &[u8]) -> std::io::Result { + self.0 = self.0.saturating_add(bytes.len()); + if self.0 > MAX_PENDING_INJECTION_BYTES + MAX_PENDING_INJECTION_MEDIA_BYTES { + return Err(std::io::Error::other( + "pending session injection budget exceeded", + )); + } + Ok(bytes.len()) + } + + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } +} + +#[cfg(feature = "unstable-inject")] +fn injection_budget_error() -> agent_client_protocol::Error { + agent_client_protocol::Error::new(-32602, "pending session injection budget exceeded") +} + #[cfg(feature = "unstable-inject")] fn validate_inject_content( content: &[wire::ContentBlock], -) -> Result<(Vec, usize), agent_client_protocol::Error> { +) -> Result<(Vec, InjectionBytes), agent_client_protocol::Error> { + let mut total = InjectionByteCounter::default(); + serde_json::to_writer(&mut total, content).map_err(|_| injection_budget_error())?; + let mut media = 0usize; + for block in content { + let payload: Option<&str> = match block { + wire::ContentBlock::Image(image) => Some(image.data.as_ref()), + wire::ContentBlock::Audio(audio) => Some(audio.data.as_ref()), + wire::ContentBlock::Resource(resource) => match &resource.resource { + wire::EmbeddedResourceResource::BlobResourceContents(blob) => { + Some(blob.blob.as_ref()) + } + _ => None, + }, + _ => None, + }; + if let Some(payload) = payload { + let mut counter = InjectionByteCounter::default(); + serde_json::to_writer(&mut counter, payload).map_err(|_| injection_budget_error())?; + // Keep the string quotes in the non-media budget. Counting serialized + // strings also charges JSON escapes if the payload is malformed. + media = media.saturating_add(counter.0 - 2); + } + } + let bytes = InjectionBytes { + content: total.0 - media, + media, + }; + if bytes.exceeds_budget() { + return Err(injection_budget_error()); + } + // Validate bounds before cloning media into the converted Items. let items = content_blocks_to_items(content).map_err(crate::sdk_error)?; - let bytes = serde_json::to_vec(content) - .map_err(|error| agent_client_protocol::Error::new(-32603, error.to_string()))? - .len(); Ok((items, bytes)) } @@ -1130,7 +1223,7 @@ struct PendingInject { message_id: wire::MessageId, content: Vec, items: Vec, - bytes: usize, + bytes: InjectionBytes, commitment: InjectCommitment, } @@ -1150,7 +1243,7 @@ struct InjectionState { at_boundary: bool, pending: VecDeque, delivering: Option, - pending_bytes: usize, + pending_bytes: InjectionBytes, accepted_count: usize, delivered: VecDeque, } @@ -1158,6 +1251,11 @@ struct InjectionState { #[cfg(feature = "unstable-inject")] #[derive(Default)] struct InjectionController { + // Both counters cover the queue plus an in-flight delivery. Synchronous + // accounting transitions prepare allocations/clones before committing and + // invoke no external callbacks between mutation and counter updates. This + // preserves accounting on unwind before a commit and permits the existing + // poison recovery below; it does not make async delivery cancellation-safe. state: Mutex, changed: Notify, } @@ -1213,15 +1311,19 @@ impl InjectionController { }))); } if state.pending.len() + usize::from(state.delivering.is_some()) >= MAX_PENDING_INJECTIONS - || state.pending_bytes.saturating_add(pending.bytes) > MAX_PENDING_INJECTION_BYTES + || state + .pending_bytes + .saturating_add(pending.bytes) + .exceeds_budget() { return Err(agent_client_protocol::Error::new( -32602, "pending session injection budget exceeded", )); } - state.pending_bytes += pending.bytes; + let pending_bytes = state.pending_bytes.saturating_add(pending.bytes); state.pending.push_back(pending); + state.pending_bytes = pending_bytes; Ok(()) } @@ -1318,7 +1420,7 @@ impl InjectionController { message_id: &wire::MessageId, content: &[wire::ContentBlock], items: &[Item], - bytes: usize, + bytes: InjectionBytes, ) -> Result { let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); if let Some(index) = state @@ -1331,18 +1433,26 @@ impl InjectionController { } // Include other pending entries and any in-flight delivery in the // budget. Check before mutation so failure preserves the old content. - let pending_bytes = state.pending_bytes - state.pending[index].bytes; - if pending_bytes.saturating_add(bytes) > MAX_PENDING_INJECTION_BYTES { + let pending_bytes = state + .pending_bytes + .saturating_sub(state.pending[index].bytes); + if pending_bytes.saturating_add(bytes).exceeds_budget() { return Err(agent_client_protocol::Error::new( -32602, "pending session injection budget exceeded", )); } + // Prepare both clones before changing any guarded state. These are + // concrete data types; cloning invokes no external callbacks. + let content = content.to_vec(); + let items = items.to_vec(); let pending = &mut state.pending[index]; - pending.content = content.to_vec(); - pending.items = items.to_vec(); + let old_content = std::mem::replace(&mut pending.content, content); + let old_items = std::mem::replace(&mut pending.items, items); pending.bytes = bytes; - state.pending_bytes = pending_bytes + bytes; + state.pending_bytes = pending_bytes.saturating_add(bytes); + drop(state); + drop((old_content, old_items)); // Do not remove/re-enqueue: identity, commitment, and FIFO position // must survive replacement. The state lock also serializes delivery. Ok(PendingTransition::Applied) @@ -1407,7 +1517,7 @@ impl InjectionController { .iter() .filter(|pending| pending.commitment == InjectCommitment::Ready) .map(|pending| pending.bytes) - .sum::(); + .fold(InjectionBytes::default(), InjectionBytes::saturating_add); state .pending .retain(|pending| pending.commitment != InjectCommitment::Ready); @@ -6119,7 +6229,7 @@ mod tests { message_id: message_id.clone(), content: Vec::new(), items: Vec::new(), - bytes: 0, + bytes: InjectionBytes::default(), commitment: InjectCommitment::Reserved, }) .expect("injection below lifetime cap"); @@ -6136,7 +6246,7 @@ mod tests { message_id: wire::MessageId::new("over-limit"), content: Vec::new(), items: Vec::new(), - bytes: 0, + bytes: InjectionBytes::default(), commitment: InjectCommitment::Reserved, }) .expect_err("accepted injection lifetime cap must reject new reservations"); @@ -7083,3 +7193,7 @@ mod tests { #[cfg(all(test, feature = "unstable-inject"))] #[path = "v2/replace_tests.rs"] mod replace_tests; + +#[cfg(all(test, feature = "unstable-inject"))] +#[path = "v2/injection_budget_tests.rs"] +mod injection_budget_tests; diff --git a/crates/agentkit-acp/src/v2/injection_budget_tests.rs b/crates/agentkit-acp/src/v2/injection_budget_tests.rs new file mode 100644 index 0000000..63217a2 --- /dev/null +++ b/crates/agentkit-acp/src/v2/injection_budget_tests.rs @@ -0,0 +1,446 @@ +//! Exercise validated content and real reservation transitions, not synthetic byte charges. +use super::*; + +// Avoid multiplying the largest payload allocations under the parallel test runner. +static LARGE_MEDIA_TEST: Mutex<()> = Mutex::new(()); + +fn image(payload_bytes: usize) -> wire::ContentBlock { + wire::ContentBlock::Image(wire::ImageContent::new( + "A".repeat(payload_bytes), + "image/png", + )) +} + +fn text_at_size(serialized_bytes: usize) -> Vec { + let empty = vec![wire::ContentBlock::Text(wire::TextContent::new(""))]; + let overhead = serde_json::to_vec(&empty).unwrap().len(); + vec![wire::ContentBlock::Text(wire::TextContent::new( + "x".repeat(serialized_bytes - overhead), + ))] +} + +fn fixture() -> (AcpIntegration, wire::SessionId, Arc) { + let integration = AcpIntegration::default(); + let id = wire::SessionId::new("budget-session"); + let (client, _messages) = ClientHandle::channel(); + integration + .bind_session(AcpSessionBinding::new( + id.clone(), + AgentkitSessionId::new("budget-agentkit"), + client, + )) + .unwrap(); + let session = integration.session(&id).unwrap(); + session.injection.start_turn(); + (integration, id, session) +} + +fn reserve( + integration: &AcpIntegration, + id: &wire::SessionId, + content: Vec, +) -> Result { + integration.reserve_inject(wire::InjectSessionRequest::new( + id.clone(), + wire::SessionInjectMode::Steer, + content, + )) +} + +fn ready(mut reservation: ReservedInject) -> wire::MessageId { + let id = reservation.message_id().clone(); + reservation.commit(); + reservation.activate(); + id +} + +fn charged(session: &IntegrationSession) -> InjectionBytes { + session.injection.state.lock().unwrap().pending_bytes +} + +fn deliver(controller: &InjectionController) -> PendingInject { + match controller.boundary_action(false, false) { + BoundaryAction::Deliver(pending) => pending, + _ => panic!("ready injection should be delivered"), + } +} + +#[test] +fn only_inline_payloads_are_charged_to_media() { + let content = vec![ + wire::ContentBlock::Image( + wire::ImageContent::new("AAAA", "image/png").uri("https://example.test/image"), + ), + wire::ContentBlock::Audio(wire::AudioContent::new("AAAA", "audio/wav")), + wire::ContentBlock::Resource(wire::EmbeddedResource::new( + wire::EmbeddedResourceResource::BlobResourceContents( + wire::BlobResourceContents::new("AAAA", "file:///blob") + .mime_type("application/octet-stream"), + ), + )), + wire::ContentBlock::Resource(wire::EmbeddedResource::new( + wire::EmbeddedResourceResource::TextResourceContents(wire::TextResourceContents::new( + "embedded text", + "file:///text", + )), + )), + wire::ContentBlock::Text(wire::TextContent::new("ordinary text")), + ]; + let (_, bytes) = validate_inject_content(&content).unwrap(); + assert_eq!(bytes.media, 12); + assert_eq!( + bytes.content, + serde_json::to_vec(&content).unwrap().len() - 12 + ); + + // Malformed payload escapes still consume serialized media bytes, not raw length. + let escaped = vec![wire::ContentBlock::Image(wire::ImageContent::new( + "\"\\\n", + "image/png", + ))]; + let (_, bytes) = validate_inject_content(&escaped).unwrap(); + let serialized_payload = serde_json::to_vec("\"\\\n").unwrap().len() - 2; + assert_eq!(bytes.media, serialized_payload); + assert_eq!( + bytes.content, + serde_json::to_vec(&escaped).unwrap().len() - serialized_payload + ); + + let metadata_heavy = vec![wire::ContentBlock::Image( + wire::ImageContent::new("AAAA", "image/png").uri("u".repeat(MAX_PENDING_INJECTION_BYTES)), + )]; + assert!( + validate_inject_content(&metadata_heavy).is_err(), + "URI bytes must not borrow the media allowance" + ); +} + +#[test] +fn content_limit_is_exact_and_independent_of_media() { + let (integration, id, session) = fixture(); + let exact = text_at_size(MAX_PENDING_INJECTION_BYTES); + let (_, bytes) = validate_inject_content(&exact).unwrap(); + assert_eq!( + bytes, + InjectionBytes { + content: MAX_PENDING_INJECTION_BYTES, + media: 0 + } + ); + let reservation = reserve(&integration, &id, exact).unwrap(); + assert_eq!(charged(&session), bytes); + assert!(reserve(&integration, &id, vec![image(4)]).is_err()); + drop(reservation); + assert_eq!(charged(&session), InjectionBytes::default()); + assert!(validate_inject_content(&text_at_size(MAX_PENDING_INJECTION_BYTES + 1)).is_err()); + + // Splitting a legal per-request size across requests must not evade the cap. + let first = reserve( + &integration, + &id, + text_at_size(MAX_PENDING_INJECTION_BYTES / 2), + ) + .unwrap(); + let second = reserve( + &integration, + &id, + text_at_size(MAX_PENDING_INJECTION_BYTES / 2), + ) + .unwrap(); + assert_eq!(charged(&session).content, MAX_PENDING_INJECTION_BYTES); + assert!(reserve(&integration, &id, text_at_size(128)).is_err()); + drop((first, second)); + assert_eq!(charged(&session), InjectionBytes::default()); +} + +#[test] +fn twenty_mib_source_image_remains_media_through_reservation() { + let _large_media = LARGE_MEDIA_TEST.lock().unwrap(); + let (integration, id, session) = fixture(); + let raw_bytes: usize = 20 * 1024 * 1024; + let encoded_bytes = raw_bytes.div_ceil(3) * 4; + // Valid base64 for 20 MiB of zero bytes, without allocating the raw buffer. + let mut payload = "A".repeat(encoded_bytes - 1); + payload.push('='); + let reservation = reserve( + &integration, + &id, + vec![wire::ContentBlock::Image(wire::ImageContent::new( + payload, + "image/png", + ))], + ) + .unwrap(); + assert_eq!(charged(&session).media, encoded_bytes); + assert!(charged(&session).content < MAX_PENDING_INJECTION_BYTES); + let message_id = ready(reservation); + let pending = deliver(&session.injection); + assert_eq!(pending.message_id, message_id); + let wire::ContentBlock::Image(original) = &pending.content[0] else { + panic!("image lost") + }; + let Part::Media(media) = &pending.items[0].parts[0] else { + panic!("image converted to text") + }; + assert_eq!(media.modality, Modality::Image); + let DataRef::InlineText(data) = &media.data else { + panic!("inline payload lost") + }; + let original_payload: &str = original.data.as_ref(); + assert_eq!( + data.strip_prefix("data:image/png;base64,").unwrap(), + original_payload + ); + session.injection.finish_delivery(&pending, true); + assert_eq!(charged(&session), InjectionBytes::default()); +} + +#[test] +fn media_exact_limit_aggregate_inflight_and_failed_replacement() { + let _large_media = LARGE_MEDIA_TEST.lock().unwrap(); + let (integration, id, session) = fixture(); + let half = MAX_PENDING_INJECTION_MEDIA_BYTES / 2; + let first = ready(reserve(&integration, &id, vec![image(half)]).unwrap()); + let second = ready(reserve(&integration, &id, vec![image(half)]).unwrap()); + let before = charged(&session); + assert_eq!(before.media, MAX_PENDING_INJECTION_MEDIA_BYTES); + assert!(reserve(&integration, &id, vec![image(4)]).is_err()); + + // An individually valid replacement fails atomically against the aggregate cap. + let replacement = vec![image(half + 4)]; + let (items, bytes) = validate_inject_content(&replacement).unwrap(); + assert!( + session + .injection + .replace_transition(&first, &replacement, &items, bytes) + .is_err() + ); + drop((replacement, items)); + assert_eq!(charged(&session), before); + let pending = deliver(&session.injection); + assert_eq!(pending.message_id, first); + assert_eq!(pending.commitment, InjectCommitment::Ready); + let wire::ContentBlock::Image(original) = &pending.content[0] else { + panic!("image lost") + }; + assert_eq!(original.data.len(), half); + assert_eq!(pending.bytes.media, half); + assert_eq!( + pending.items, + content_blocks_to_items(&pending.content).unwrap() + ); + assert_eq!(charged(&session), before, "delivery still retains payloads"); + assert!(reserve(&integration, &id, vec![image(4)]).is_err()); + session.injection.finish_delivery(&pending, true); + drop(pending); + let extra = reserve(&integration, &id, vec![image(4)]).unwrap(); + extra.discard(); + let pending = deliver(&session.injection); + assert_eq!( + pending.message_id, second, + "failed replacement preserves FIFO" + ); + session.injection.finish_delivery(&pending, false); + drop(pending); + assert_eq!(charged(&session), InjectionBytes::default()); + + // Exercise validation's per-request boundary independently of admission. + let exact = vec![image(MAX_PENDING_INJECTION_MEDIA_BYTES)]; + let (items, bytes) = validate_inject_content(&exact).unwrap(); + assert_eq!(bytes.media, MAX_PENDING_INJECTION_MEDIA_BYTES); + drop((exact, items)); + assert!(validate_inject_content(&[image(MAX_PENDING_INJECTION_MEDIA_BYTES + 1)]).is_err()); +} + +#[test] +fn failed_content_replacement_preserves_identity_content_and_fifo() { + let (integration, id, session) = fixture(); + let first = ready( + reserve( + &integration, + &id, + text_at_size(MAX_PENDING_INJECTION_BYTES / 2), + ) + .unwrap(), + ); + let second = ready( + reserve( + &integration, + &id, + text_at_size(MAX_PENDING_INJECTION_BYTES / 2), + ) + .unwrap(), + ); + let before = charged(&session); + let replacement = text_at_size(MAX_PENDING_INJECTION_BYTES / 2 + 1); + let (items, bytes) = validate_inject_content(&replacement).unwrap(); + assert!( + session + .injection + .replace_transition(&first, &replacement, &items, bytes) + .is_err() + ); + assert_eq!(charged(&session), before); + for id in [first, second] { + let pending = deliver(&session.injection); + assert_eq!(pending.message_id, id); + assert_eq!( + pending.content, + text_at_size(MAX_PENDING_INJECTION_BYTES / 2) + ); + assert_eq!( + pending.items, + content_blocks_to_items(&pending.content).unwrap() + ); + session.injection.finish_delivery(&pending, true); + } + assert_eq!(charged(&session), InjectionBytes::default()); +} + +#[test] +fn rollback_revoke_discard_and_close_release_both_budgets() { + let (integration, id, session) = fixture(); + let content = || { + vec![ + wire::ContentBlock::Text(wire::TextContent::new("retained text")), + image(512 * 1024), + ] + }; + let reservation = reserve(&integration, &id, content()).unwrap(); + let bytes = charged(&session); + assert!(bytes.content > 0 && bytes.media > 0); + drop(reservation); // Uncommitted request rollback. + assert_eq!(charged(&session), InjectionBytes::default()); + assert_eq!(session.injection.state.lock().unwrap().accepted_count, 0); + let mut reservation = reserve(&integration, &id, content()).unwrap(); + reservation.commit(); + drop(reservation); // Response failure after commitment must also release retention. + assert_eq!(charged(&session), InjectionBytes::default()); + let revoked = ready(reserve(&integration, &id, content()).unwrap()); + assert!(matches!( + session.injection.revoke_transition(&revoked), + PendingTransition::Applied + )); + assert_eq!(charged(&session), InjectionBytes::default()); + reserve(&integration, &id, content()).unwrap().discard(); + assert_eq!(charged(&session), InjectionBytes::default()); + + ready(reserve(&integration, &id, content()).unwrap()); + let pending = deliver(&session.injection); + ready(reserve(&integration, &id, content()).unwrap()); + let reserved = reserve(&integration, &id, content()).unwrap(); + session.injection.close_session(); + assert_eq!( + charged(&session), + bytes.saturating_add(bytes), + "close releases Ready entries, not owned reservations or delivery" + ); + drop(reserved); + assert_eq!(charged(&session), bytes); + session.injection.finish_delivery(&pending, false); + assert_eq!(charged(&session), InjectionBytes::default()); +} + +#[test] +fn cancellation_retains_budget_and_content_until_next_turn() { + let (integration, id, session) = fixture(); + let message_id = ready(reserve(&integration, &id, vec![image(512 * 1024)]).unwrap()); + let before = charged(&session); + session.injection.cancel_turn(); + assert_eq!(charged(&session), before); + assert!(matches!( + session.injection.boundary_action(false, false), + BoundaryAction::Complete(AcpInjectionBoundary::Stopped) + )); + session.injection.reset_cancellation(); + session.injection.start_turn(); + let pending = deliver(&session.injection); + assert_eq!(pending.message_id, message_id); + assert_eq!(pending.bytes, before); + session.injection.finish_delivery(&pending, true); + assert_eq!(charged(&session), InjectionBytes::default()); +} + +#[tokio::test] +async fn media_replacement_grows_shrinks_and_changes_budget_class() { + let (integration, id, session) = fixture(); + let message_id = ready(reserve(&integration, &id, vec![image(512 * 1024)]).unwrap()); + for content in [ + vec![image(1024 * 1024)], + vec![image(4)], + text_at_size(128), + vec![image(512 * 1024)], + ] { + let (_, expected) = validate_inject_content(&content).unwrap(); + integration + .replace_inject(wire::ReplaceInjectSessionRequest::new( + id.clone(), + message_id.clone(), + content.clone(), + )) + .await + .unwrap(); + assert_eq!(charged(&session), expected); + let state = session.injection.state.lock().unwrap(); + assert_eq!(state.pending.len(), 1); + assert_eq!(state.pending[0].content, content); + assert_eq!( + state.pending[0].items, + content_blocks_to_items(&content).unwrap() + ); + } + let pending = deliver(&session.injection); + assert_eq!(pending.message_id, message_id); + session.injection.finish_delivery(&pending, true); + assert_eq!(charged(&session), InjectionBytes::default()); +} + +#[test] +fn poison_recovery_preserves_both_budget_counters() { + let (integration, id, session) = fixture(); + let message_id = ready( + reserve( + &integration, + &id, + vec![ + wire::ContentBlock::Text(wire::TextContent::new("mixed content")), + image(512 * 1024), + ], + ) + .unwrap(), + ); + // Poison an otherwise valid state; no production panic hooks are needed. + let poisoned = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + let _guard = session.injection.state.lock().unwrap(); + panic!("test unwind while owning valid state"); + })); + assert!(poisoned.is_err()); + assert!(session.injection.state.is_poisoned()); + let content = vec![image(4)]; + let (items, bytes) = validate_inject_content(&content).unwrap(); + assert!(matches!( + session + .injection + .replace_transition(&message_id, &content, &items, bytes) + .unwrap(), + PendingTransition::Applied + )); + { + let state = session + .injection + .state + .lock() + .unwrap_or_else(|error| error.into_inner()); + assert_eq!(state.pending_bytes, bytes); + assert_eq!(state.pending[0].content, content); + assert_eq!(state.pending[0].items, items); + } + session.injection.discard(&message_id); + let state = session + .injection + .state + .lock() + .unwrap_or_else(|error| error.into_inner()); + assert_eq!(state.pending_bytes, InjectionBytes::default()); + assert!(state.pending.is_empty()); +} diff --git a/crates/agentkit-acp/src/v2/replace_tests.rs b/crates/agentkit-acp/src/v2/replace_tests.rs index f108f4b..2920b0a 100644 --- a/crates/agentkit-acp/src/v2/replace_tests.rs +++ b/crates/agentkit-acp/src/v2/replace_tests.rs @@ -65,8 +65,9 @@ fn assert_text(pending: &PendingInject, expected: &str) { assert!(matches!(&pending.items[0].parts[0], Part::Text(text) if text.text == expected)); assert_eq!( pending.bytes, - serde_json::to_vec(&pending.content).unwrap().len() + validate_inject_content(&pending.content).unwrap().1 ); + assert_eq!(pending.bytes.media, 0); assert_eq!(pending.commitment, InjectCommitment::Ready); } @@ -108,7 +109,11 @@ async fn replacement_preserves_fifo_identity_and_acceptance_count() { assert_eq!(state.pending.len(), 3); assert_eq!( state.pending_bytes, - state.pending.iter().map(|p| p.bytes).sum::() + state + .pending + .iter() + .map(|p| p.bytes) + .fold(InjectionBytes::default(), InjectionBytes::saturating_add) ); } for (id, text) in [(first, "first"), (second, "final second"), (third, "third")] { @@ -119,7 +124,7 @@ async fn replacement_preserves_fifo_identity_and_acceptance_count() { } let state = injection.state.lock().unwrap(); assert!(state.pending.is_empty()); - assert_eq!(state.pending_bytes, 0); + assert_eq!(state.pending_bytes, InjectionBytes::default()); assert_eq!(state.accepted_count, 3); } @@ -131,15 +136,20 @@ async fn replacement_grows_shrinks_and_rejects_over_budget_atomically() { let injection = &session.session.injection; let in_flight = take_delivery(injection); assert_eq!(in_flight.message_id, first); - let overhead = serde_json::to_vec(&text_content("")).unwrap().len(); - let full = "x".repeat(MAX_PENDING_INJECTION_BYTES - in_flight.bytes - overhead); + let overhead = validate_inject_content(&text_content("")).unwrap().1; + assert_eq!(overhead.media, 0); + assert_eq!(in_flight.bytes.media, 0); + let full = "x".repeat(MAX_PENDING_INJECTION_BYTES - in_flight.bytes.content - overhead.content); integration .replace_inject(request(&session, &second, &full)) .await .unwrap(); assert_eq!( injection.state.lock().unwrap().pending_bytes, - MAX_PENDING_INJECTION_BYTES + InjectionBytes { + content: MAX_PENDING_INJECTION_BYTES, + media: 0 + } ); let error = integration @@ -149,7 +159,13 @@ async fn replacement_grows_shrinks_and_rejects_over_budget_atomically() { assert_eq!(i32::from(error.code), -32602); { let state = injection.state.lock().unwrap(); - assert_eq!(state.pending_bytes, MAX_PENDING_INJECTION_BYTES); + assert_eq!( + state.pending_bytes, + InjectionBytes { + content: MAX_PENDING_INJECTION_BYTES, + media: 0 + } + ); assert_eq!(state.accepted_count, 2); assert_eq!(state.delivering.as_ref(), Some(&first)); assert_eq!(state.pending.len(), 1); @@ -161,17 +177,21 @@ async fn replacement_grows_shrinks_and_rejects_over_budget_atomically() { .replace_inject(request(&session, &second, "small")) .await .unwrap(); - let small_bytes = serde_json::to_vec(&text_content("small")).unwrap().len(); + let small_bytes = validate_inject_content(&text_content("small")).unwrap().1; + assert_eq!(small_bytes.media, 0); assert_eq!( injection.state.lock().unwrap().pending_bytes, - in_flight.bytes + small_bytes + in_flight.bytes.saturating_add(small_bytes) ); injection.finish_delivery(&in_flight, true); assert_eq!(injection.state.lock().unwrap().pending_bytes, small_bytes); let pending = take_delivery(injection); assert_text(&pending, "small"); injection.finish_delivery(&pending, true); - assert_eq!(injection.state.lock().unwrap().pending_bytes, 0); + assert_eq!( + injection.state.lock().unwrap().pending_bytes, + InjectionBytes::default() + ); } #[tokio::test] @@ -223,7 +243,7 @@ async fn replacement_of_abandoned_reservation_wakes_with_unknown_id() { assert_inject_error(error, "unknown_message_id", &id); let state = session.session.injection.state.lock().unwrap(); assert!(state.pending.is_empty()); - assert_eq!(state.pending_bytes, 0); + assert_eq!(state.pending_bytes, InjectionBytes::default()); assert_eq!(state.accepted_count, 0); } @@ -296,7 +316,7 @@ async fn delivery_winner_blocks_replacement_until_delivery_outcome() { .lock() .unwrap() .pending_bytes, - 0 + InjectionBytes::default() ); } } @@ -308,8 +328,8 @@ fn simultaneous_delivery_and_replacement_have_one_linearized_outcome() { let id = accept(&integration, &session, "original"); let injection = &session.session.injection; let content = text_content("replacement"); - let items = content_blocks_to_items(&content).unwrap(); - let bytes = serde_json::to_vec(&content).unwrap().len(); + let (items, bytes) = validate_inject_content(&content).unwrap(); + assert_eq!(bytes.media, 0); let barrier = std::sync::Barrier::new(2); let (transition, pending) = std::thread::scope(|scope| { let replacing = scope.spawn(|| { @@ -331,7 +351,7 @@ fn simultaneous_delivery_and_replacement_have_one_linearized_outcome() { injection.finish_delivery(&pending, true); let state = injection.state.lock().unwrap(); assert!(state.pending.is_empty()); - assert_eq!(state.pending_bytes, 0); + assert_eq!(state.pending_bytes, InjectionBytes::default()); assert_eq!(state.accepted_count, 1); assert_eq!(state.delivered.len(), 1); } @@ -387,7 +407,7 @@ async fn close_wakes_waiting_replacements_and_discards_pending_bytes() { } let state = session.session.injection.state.lock().unwrap(); assert!(state.pending.is_empty()); - assert_eq!(state.pending_bytes, 0); + assert_eq!(state.pending_bytes, InjectionBytes::default()); } let (integration, session) = fixture(); @@ -410,7 +430,7 @@ async fn close_wakes_waiting_replacements_and_discards_pending_bytes() { .lock() .unwrap() .pending_bytes, - 0 + InjectionBytes::default() ); } From 710122fc036edcb8394891c231de50a80651f755 Mon Sep 17 00:00:00 2001 From: daniel Date: Thu, 1 Oct 2026 15:02:57 +0100 Subject: [PATCH 7/7] feat(acp): depend on the published agentkit-acp-sdk fork Replace the unpublishable git patch with agentkit-acp-sdk 2.2.0 so session injection ships behind the unstable-inject feature. ACP v2 prompt responses now carry the accepted user message id, allocated when the prompt is accepted. --- Cargo.lock | 112 ++++++++++-------------- Cargo.toml | 4 - book/src/acp.md | 9 +- crates/agentkit-acp/Cargo.toml | 6 +- crates/agentkit-acp/README.md | 10 +-- crates/agentkit-acp/src/v2.rs | 61 ++++++------- crates/agentkit-loop/Cargo.toml | 2 +- docs/acp.md | 11 +-- examples/openrouter-acp-trio/Cargo.toml | 2 +- 9 files changed, 95 insertions(+), 122 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 632f188..da6c34c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8,57 +8,16 @@ version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" -[[package]] -name = "agent-client-protocol" -version = "2.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6d87bc7769eba641753ba5dc52f73ec3765d51022c6753bf040967125ddc86a8" -dependencies = [ - "agent-client-protocol-derive", - "agent-client-protocol-schema", - "async-io", - "async-process", - "blocking", - "futures", - "futures-concurrency", - "rustc-hash", - "rustix", - "schemars 1.2.2", - "serde", - "serde_json", - "shell-words", - "tracing", - "uuid", - "windows-sys 0.61.2", -] - [[package]] name = "agent-client-protocol-derive" -version = "2.0.0" +version = "2.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3abd4080f51e4f24f5042beb7fb7a66ede29a2dc1c2582c329532e1c27264ddc" +checksum = "f4893bfb11e0e193aa5a64611a126ee87567acb1879556519c902722899229de" dependencies = [ "quote", "syn 3.0.3", ] -[[package]] -name = "agent-client-protocol-schema" -version = "1.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d5c231915b4ab578c722eca2d1bd7df4d300bfd6cac3b8e9f0d1e3ddc95b187c" -dependencies = [ - "anyhow", - "derive_more", - "diffy", - "schemars 1.2.2", - "serde", - "serde_json", - "serde_with", - "strum", - "tracing", -] - [[package]] name = "agentkit" version = "0.10.12" @@ -92,9 +51,9 @@ dependencies = [ [[package]] name = "agentkit-acp" -version = "0.10.11" +version = "0.10.12" dependencies = [ - "agent-client-protocol", + "agentkit-acp-sdk", "agentkit-core", "agentkit-integration-tests", "agentkit-loop", @@ -102,12 +61,53 @@ dependencies = [ "agentkit-tools-core", "async-trait", "base64 0.22.1", + "futures-util", "serde_json", "thiserror 2.0.18", "tokio", "tracing", ] +[[package]] +name = "agentkit-acp-schema" +version = "1.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dfcb839ccca8f4f43bfc9aa26473396f16afbd32071efd00d36801daed62f9d7" +dependencies = [ + "anyhow", + "derive_more", + "schemars 1.2.2", + "serde", + "serde_json", + "serde_with", + "strum", + "tracing", +] + +[[package]] +name = "agentkit-acp-sdk" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b93c10bb458dd1c2c822275a7e1c390b46846a6bcd3cbc84aafa312563e36ad" +dependencies = [ + "agent-client-protocol-derive", + "agentkit-acp-schema", + "async-io", + "async-process", + "blocking", + "futures", + "futures-concurrency", + "rustc-hash", + "rustix", + "schemars 1.2.2", + "serde", + "serde_json", + "shell-words", + "tracing", + "uuid", + "windows-sys 0.61.2", +] + [[package]] name = "agentkit-adapter-completions" version = "0.10.9" @@ -215,7 +215,7 @@ dependencies = [ [[package]] name = "agentkit-loop" -version = "0.10.12" +version = "0.10.13" dependencies = [ "agentkit-core", "agentkit-task-manager", @@ -1245,15 +1245,6 @@ version = "0.1.13" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "56254986775e3233ffa9c4d7d3faaf6d36a2c09d30b20687e9f88bc8bafc16c8" -[[package]] -name = "diffy" -version = "0.5.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3e3dc2f773b6aaa63b1a7684b8589f670a8a0146a510b74d23a401c882364b49" -dependencies = [ - "hashbrown 0.17.1", -] - [[package]] name = "digest" version = "0.10.7" @@ -1424,12 +1415,6 @@ dependencies = [ "miniz_oxide", ] -[[package]] -name = "foldhash" -version = "0.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb" - [[package]] name = "form_urlencoded" version = "1.2.2" @@ -1636,9 +1621,6 @@ name = "hashbrown" version = "0.17.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" -dependencies = [ - "foldhash", -] [[package]] name = "hashlink" @@ -2341,8 +2323,8 @@ checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" name = "openrouter-acp-trio" version = "0.10.5" dependencies = [ - "agent-client-protocol", "agentkit-acp", + "agentkit-acp-sdk", "agentkit-core", "agentkit-loop", "agentkit-provider-openrouter", diff --git a/Cargo.toml b/Cargo.toml index 1206993..c64d394 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -60,10 +60,6 @@ repository = "https://github.com/danielkov/agentkit" rust-version = "1.92" version = "0.10.5" -# Intentionally unpublishable until the upstream SDK releases session injection. -[patch.crates-io] -agent-client-protocol = { git = "https://github.com/danielkov/rust-sdk", rev = "2f039993d1d6ed8da35b38c31f54a7cbb7338c70" } - [workspace.dependencies] async-trait = "0.1.89" dotenvy = "0.15.7" diff --git a/book/src/acp.md b/book/src/acp.md index cac2810..5830054 100644 --- a/book/src/acp.md +++ b/book/src/acp.md @@ -4,10 +4,10 @@ The [Agent Client Protocol (ACP)](https://agentclientprotocol.com) standardizes ## Built on the ACP Rust SDK -Like `agentkit-mcp`, this crate does not define a parallel protocol vocabulary. It currently builds on a commit-pinned [`agent-client-protocol` fork](https://github.com/danielkov/rust-sdk), needed for experimental session injection support, and re-exports the stable v1 wire types (`SessionId`, `ContentBlock`, `SessionUpdate`, `ToolCallUpdate`, `StopReason`, …) at the crate root and under `agentkit_acp::wire`. The full SDK is available as `agentkit_acp::sdk`. Agentkit owns only the host-facing glue: session binding, observer routing, prompt conversion, cancellation handles, and approval resolution. +Like `agentkit-mcp`, this crate does not define a parallel protocol vocabulary. It builds on [`agentkit-acp-sdk`](https://crates.io/crates/agentkit-acp-sdk), a published fork of the `agent-client-protocol` SDK that adds experimental session injection, and re-exports the stable v1 wire types (`SessionId`, `ContentBlock`, `SessionUpdate`, `ToolCallUpdate`, `StopReason`, …) at the crate root and under `agentkit_acp::wire`. The full SDK is available as `agentkit_acp::sdk`. Agentkit owns only the host-facing glue: session binding, observer routing, prompt conversion, cancellation handles, and approval resolution. - **Protocol docs:** [agentclientprotocol.com](https://agentclientprotocol.com/protocol/v1/overview) -- **Rust SDK fork:** [`danielkov/rust-sdk`](https://github.com/danielkov/rust-sdk), pinned in `Cargo.toml` and `Cargo.lock` +- **Rust SDK fork:** [`agentkit-acp-sdk`](https://crates.io/crates/agentkit-acp-sdk) ([source](https://github.com/danielkov/rust-sdk/tree/agentkit-acp-sdk)) ## Opt-in ACP v2 runtime @@ -67,8 +67,7 @@ The initial v2 foundation routes text, reasoning, and tool lifecycle updates. ACP v2 permission callbacks are intentionally deferred; unsupported approval requests are denied while accepted steers remain pending for the next safe boundary. The SDK labels v2 unstable, so opt-in callers should expect the `v2` -namespace to track the pinned fork. The workspace patch is intentionally -unpublishable until the required APIs are released upstream. +namespace to track the SDK fork. ## Two integration shapes @@ -245,4 +244,4 @@ The `agentkit-acp` crate itself has a default `stdio` feature that gates `serve_ > **Example:** [`openrouter-acp-trio`](https://github.com/danielkov/agentkit/tree/main/examples/openrouter-acp-trio) runs three OpenRouter-backed agents (orchestrator, worker, reviewer) that call each other over in-memory ACP endpoints while a REPL drives the orchestrator through a persistent ACP session — session binding, streamed updates, tool call reporting, and agent-to-agent handoffs in one program. > -> **Crate:** [`agentkit-acp`](https://github.com/danielkov/agentkit/tree/main/crates/agentkit-acp) — depends on [`agentkit-core`](https://github.com/danielkov/agentkit/tree/main/crates/agentkit-core), [`agentkit-loop`](https://github.com/danielkov/agentkit/tree/main/crates/agentkit-loop), [`agentkit-tools-core`](https://github.com/danielkov/agentkit/tree/main/crates/agentkit-tools-core), and a commit-pinned [`agent-client-protocol` fork](https://github.com/danielkov/rust-sdk). Design notes: [`docs/acp.md`](https://github.com/danielkov/agentkit/blob/main/docs/acp.md). +> **Crate:** [`agentkit-acp`](https://github.com/danielkov/agentkit/tree/main/crates/agentkit-acp) — depends on [`agentkit-core`](https://github.com/danielkov/agentkit/tree/main/crates/agentkit-core), [`agentkit-loop`](https://github.com/danielkov/agentkit/tree/main/crates/agentkit-loop), [`agentkit-tools-core`](https://github.com/danielkov/agentkit/tree/main/crates/agentkit-tools-core), and [`agentkit-acp-sdk`](https://crates.io/crates/agentkit-acp-sdk), a published `agent-client-protocol` fork. Design notes: [`docs/acp.md`](https://github.com/danielkov/agentkit/blob/main/docs/acp.md). diff --git a/crates/agentkit-acp/Cargo.toml b/crates/agentkit-acp/Cargo.toml index fe158be..865811b 100644 --- a/crates/agentkit-acp/Cargo.toml +++ b/crates/agentkit-acp/Cargo.toml @@ -4,7 +4,7 @@ homepage.workspace = true name = "agentkit-acp" readme = "README.md" repository.workspace = true -version = "0.10.11" +version = "0.10.12" edition.workspace = true license.workspace = true rust-version.workspace = true @@ -20,9 +20,9 @@ unstable-inject = [ ] [dependencies] -agent-client-protocol = "=2.0.0" +agent-client-protocol = { package = "agentkit-acp-sdk", version = "=2.2.0" } agentkit-core = { version = "0.10.5", path = "../agentkit-core" } -agentkit-loop = { version = "0.10.10", path = "../agentkit-loop" } +agentkit-loop = { version = "0.10.13", path = "../agentkit-loop" } agentkit-tools-core = { version = "0.10.5", path = "../agentkit-tools-core" } async-trait.workspace = true base64.workspace = true diff --git a/crates/agentkit-acp/README.md b/crates/agentkit-acp/README.md index f5903e4..af387c8 100644 --- a/crates/agentkit-acp/README.md +++ b/crates/agentkit-acp/README.md @@ -2,7 +2,8 @@ Agent Client Protocol integration for agentkit hosts. -This crate re-exports ACP wire types from a pinned `agent-client-protocol` fork and +This crate re-exports ACP wire types from [`agentkit-acp-sdk`](https://crates.io/crates/agentkit-acp-sdk), a published +`agent-client-protocol` fork, and adds only agentkit-specific glue: session binding, observer routing, prompt conversion, cancellation handles, and approval resolver abstractions. @@ -24,7 +25,7 @@ experimental upstream ACP v2 protocol, enable the additive feature: agentkit-acp = { version = "0.10.11", features = ["protocol-v2"] } ``` -The feature maps directly to the pinned fork's +The feature maps directly to the SDK fork's `agent-client-protocol/unstable_protocol_v2` feature. V2 APIs and v2 wire types live only under `agentkit_acp::v2` (and `agentkit_acp::v2::wire`). Enable `unstable-inject` instead to add the unstable @@ -68,9 +69,8 @@ approval requests are resolved as denials while retaining a provider-valid transcript, so the prompt ends with the custom `_error` stop reason rather than `Refusal`. Already accepted steers remain pending and are delivered at the next safe boundary. Because the SDK marks -protocol v2 unstable, all APIs in the `v2` namespace can evolve with the pinned -SDK fork. The workspace patch is intentionally unpublishable until these SDK -APIs are available in an upstream release. +protocol v2 unstable, all APIs in the `v2` namespace can evolve with the +[`agentkit-acp-sdk`](https://crates.io/crates/agentkit-acp-sdk) fork. Run the stable v1 in-memory end-to-end example with: diff --git a/crates/agentkit-acp/src/v2.rs b/crates/agentkit-acp/src/v2.rs index a70aa40..562b00d 100644 --- a/crates/agentkit-acp/src/v2.rs +++ b/crates/agentkit-acp/src/v2.rs @@ -546,7 +546,7 @@ impl AcpIntegration { session_id: &wire::SessionId, owner: u64, cancellation: TurnCancellation, - ) -> Result { + ) -> Result<(), AcpRuntimeError> { let session = self.session(session_id)?; { let mut prompt_state = session @@ -561,10 +561,7 @@ impl AcpIntegration { cancellation, }); } - let sequence = session.next_message.fetch_add(1, Ordering::Relaxed); - Ok(wire::MessageId::new(format!( - "{session_id}-user-{sequence}" - ))) + Ok(()) } fn finish_prompt_owner(&self, session_id: &wire::SessionId, owner: u64) { @@ -1002,8 +999,8 @@ where let state = Arc::clone(&state); cx.spawn(async move { match state.prompt(request).await { - Ok(start) => { - responder.respond(wire::PromptResponse::new())?; + Ok((start, message_id)) => { + responder.respond(wire::PromptResponse::new(message_id))?; let _ = start.send(()); Ok(()) } @@ -1973,6 +1970,7 @@ enum SessionCommand { Prompt { request: wire::PromptRequest, items: Vec, + message_id: wire::MessageId, prompt_cancellation: TurnCancellation, cancellation_generation: u64, owner: u64, @@ -2200,7 +2198,7 @@ where async fn prompt( &self, request: wire::PromptRequest, - ) -> Result, AcpRuntimeError> { + ) -> Result<(oneshot::Sender<()>, wire::MessageId), AcpRuntimeError> { let items = self.integration.prompt_to_items(&request)?; let entry = self .sessions @@ -2235,6 +2233,7 @@ where "session is already running a prompt".into(), )); } + let message_id = self.integration.next_user_message_id(&request.session_id)?; let prompt_cancellation = entry.cancellation.handle().checkpoint(); #[cfg(feature = "unstable-inject")] entry.session.prepare_injection_turn(); @@ -2244,6 +2243,7 @@ where .send(SessionCommand::Prompt { request, items, + message_id: message_id.clone(), prompt_cancellation, cancellation_generation, owner, @@ -2254,8 +2254,8 @@ where release_prompt(&entry.active_prompt, owner); return Err(AcpRuntimeError::ClientClosed); } + Ok((start_tx, message_id)) } - Ok(start_tx) } async fn cancel( @@ -2493,6 +2493,7 @@ async fn session_worker( let SessionCommand::Prompt { request, items, + message_id: user_message_id, prompt_cancellation, cancellation_generation, owner, @@ -2543,27 +2544,24 @@ async fn session_worker( ); continue; } - let user_message_id = - match integration.begin_prompt_owner(&session_id, owner, prompt_cancellation) { - Ok(message_id) => message_id, - Err(error) => { - tracing::debug!(%error, owner, "failed to begin accepted ACP v2 prompt"); - #[cfg(feature = "unstable-inject")] - session.stop_injection_turn(); - fail_accepted_prompt( - &client, - &integration, - &session_id, - &active_prompt, - &driving_prompt, - &cancelled_prompt, - &lifecycle, - owner, - false, - ); - continue; - } - }; + if let Err(error) = integration.begin_prompt_owner(&session_id, owner, prompt_cancellation) + { + tracing::debug!(%error, owner, "failed to begin accepted ACP v2 prompt"); + #[cfg(feature = "unstable-inject")] + session.stop_injection_turn(); + fail_accepted_prompt( + &client, + &integration, + &session_id, + &active_prompt, + &driving_prompt, + &cancelled_prompt, + &lifecycle, + owner, + false, + ); + continue; + } if let Err(error) = client.update_for( session_id.clone(), @@ -7117,6 +7115,9 @@ mod tests { ) .expect("install prompt state"); let cancellation = CancellationController::new(); + integration + .next_user_message_id(&acp_id) + .expect("accept prompt"); integration .begin_prompt_owner(&acp_id, 1, cancellation.handle().checkpoint()) .expect("begin prompt"); diff --git a/crates/agentkit-loop/Cargo.toml b/crates/agentkit-loop/Cargo.toml index efafdf3..8b93faf 100644 --- a/crates/agentkit-loop/Cargo.toml +++ b/crates/agentkit-loop/Cargo.toml @@ -4,7 +4,7 @@ homepage.workspace = true name = "agentkit-loop" readme = "README.md" repository.workspace = true -version = "0.10.12" +version = "0.10.13" edition.workspace = true license.workspace = true rust-version.workspace = true diff --git a/docs/acp.md b/docs/acp.md index ab8e62a..0bff66d 100644 --- a/docs/acp.md +++ b/docs/acp.md @@ -43,7 +43,7 @@ There is already an official Rust ACP SDK: - `agent-client-protocol-schema` - `agent-client-protocol-rmcp` -So agentkit should not build a parallel protocol crate by default. The current implementation uses a commit-pinned source fork of that SDK for experimental session injection; it still re-exports the SDK wire types and owns only the lifecycle and conversion glue. A separate `racp` crate would make sense only if the SDK API became unsuitable. +So agentkit should not build a parallel protocol crate by default. The current implementation uses `agentkit-acp-sdk`, a published fork of that SDK, for experimental session injection; it still re-exports the SDK wire types and owns only the lifecycle and conversion glue. A separate `racp` crate would make sense only if the SDK API became unsuitable. ## Non-goals @@ -58,7 +58,7 @@ So agentkit should not build a parallel protocol crate by default. The current i - long-term session persistence - policy storage for remembered approval choices -The crate integrates agentkit into ACP. Its pinned SDK source fork must not become a divergent ACP protocol, and the crate must not turn agentkit into a single opinionated coding-agent product. +The crate integrates agentkit into ACP. Its SDK fork must not become a divergent ACP protocol, and the crate must not turn agentkit into a single opinionated coding-agent product. ## Dependencies @@ -66,8 +66,7 @@ Recommended initial crate: ```toml [dependencies] -# Registry-shaped so every workspace consumer resolves one patched type identity. -agent-client-protocol = "=2.0.0" +agent-client-protocol = { package = "agentkit-acp-sdk", version = "=2.2.0" } agentkit-core = { path = "../agentkit-core", version = "0.10.5" } agentkit-loop = { path = "../agentkit-loop", version = "0.10.7" } agentkit-tools-core = { path = "../agentkit-tools-core", version = "0.10.5" } @@ -78,10 +77,6 @@ thiserror = { workspace = true } tokio = { workspace = true, features = ["sync"] } tracing = { workspace = true } -[patch.crates-io] -# Intentionally unpublishable until these SDK APIs are released upstream. -agent-client-protocol = { git = "https://github.com/danielkov/rust-sdk", rev = "2f039993d1d6ed8da35b38c31f54a7cbb7338c70" } - [features] default = ["stdio"] stdio = [] diff --git a/examples/openrouter-acp-trio/Cargo.toml b/examples/openrouter-acp-trio/Cargo.toml index 9624247..f2aec15 100644 --- a/examples/openrouter-acp-trio/Cargo.toml +++ b/examples/openrouter-acp-trio/Cargo.toml @@ -7,7 +7,7 @@ rust-version.workspace = true version.workspace = true [dependencies] -agent-client-protocol = "2.0.0" +agent-client-protocol = { package = "agentkit-acp-sdk", version = "2.2.0" } agentkit-acp = { path = "../../crates/agentkit-acp" } agentkit-core = { path = "../../crates/agentkit-core" } agentkit-loop = { path = "../../crates/agentkit-loop" }