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/book/src/acp.md b/book/src/acp.md index 02d0769..5830054 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 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:** [`agent-client-protocol` on crates.io](https://crates.io/crates/agent-client-protocol) +- **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 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,26 @@ 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. 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; 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 SDK fork. ## Two integration shapes @@ -214,7 +230,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 +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 [`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 [`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 993b3f6..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 @@ -14,11 +14,15 @@ 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" +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 @@ -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..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 upstream ACP wire types from `agent-client-protocol` 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. @@ -21,28 +22,55 @@ 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 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 +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. 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. 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 +[`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 ba5b8b3..562b00d 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,12 +541,12 @@ impl AcpIntegration { Ok(()) } - fn begin_prompt( + fn begin_prompt_owner( &self, session_id: &wire::SessionId, owner: u64, cancellation: TurnCancellation, - ) -> Result { + ) -> Result<(), AcpRuntimeError> { let session = self.session(session_id)?; { let mut prompt_state = session @@ -268,14 +561,10 @@ impl AcpIntegration { cancellation, }); } - let sequence = session.next_message.fetch_add(1, Ordering::Relaxed); - finish_model_message(&session); - Ok(wire::MessageId::new(format!( - "{session_id}-user-{sequence}" - ))) + Ok(()) } - 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 +644,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 +681,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 +744,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 +845,7 @@ where M: ModelAdapter, { factory: Option>>, + integration: AcpIntegration, name: String, version: String, } @@ -567,6 +857,7 @@ where fn default() -> Self { Self { factory: None, + integration: AcpIntegration::default(), name: "agentkit".into(), version: env!("CARGO_PKG_VERSION").into(), } @@ -598,6 +889,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 +919,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 +927,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( @@ -693,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(()) } @@ -707,7 +1013,45 @@ 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!(), + ) + .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( { let state = Arc::clone(&state); @@ -753,1886 +1097,4199 @@ 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>>, -} - -enum SessionCommand { - Prompt { - request: wire::PromptRequest, - items: Vec, - prompt_cancellation: TurnCancellation, - cancellation_generation: u64, - owner: u64, - start: oneshot::Receiver<()>, - }, - Shutdown, +/// Maximum number of pending unstable ACP v2 injections per session. +#[cfg(feature = "unstable-inject")] +pub const MAX_PENDING_INJECTIONS: usize = 64; +/// 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. +/// +/// 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; + +/// 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, } -struct RuntimeState -where - M: ModelAdapter, -{ - factory: Arc>, - integration: Arc, - sessions: AsyncMutex>>, - next_session: AtomicU64, - name: String, - version: String, -} +#[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), + } + } -impl RuntimeState -where - M: ModelAdapter + Send + Sync + 'static, - M::Session: Send + 'static, -{ - fn new(factory: Arc>, name: String, version: String) -> Self { + fn saturating_sub(self, other: Self) -> Self { Self { - factory, - integration: Arc::new(AcpIntegration::default()), - sessions: AsyncMutex::new(HashMap::new()), - next_session: AtomicU64::new(1), - name, - version, + content: self.content.saturating_sub(other.content), + media: self.media.saturating_sub(other.media), } } - 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 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(wire::InitializeResponse::new( - wire::ProtocolVersion::V2, - wire::Implementation::new(self.name.clone(), self.version.clone()), - ) - .capabilities(headless_capabilities())) + Ok(bytes.len()) } - 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()); - let cancellation = CancellationController::new(); - let (client, client_messages) = ClientHandle::channel(); - let info = wire::SessionInfo::new(acp_session_id.clone(), request.cwd.clone()) - .additional_directories(request.additional_directories.clone()); + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } +} - let mut metadata = MetadataMap::new(); - metadata.insert("acp.protocol_version".into(), json!(2)); - metadata.insert("acp.cwd".into(), json!(request.cwd)); - metadata.insert( - "acp.additional_directories".into(), - json!(request.additional_directories), - ); +#[cfg(feature = "unstable-inject")] +fn injection_budget_error() -> agent_client_protocol::Error { + agent_client_protocol::Error::new(-32602, "pending session injection budget exceeded") +} - 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(), - )?; - 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); - 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, - cwd: request.cwd.into_inner(), - additional_directories: request - .additional_directories - .into_iter() - .map(wire::AbsolutePath::into_inner) - .collect(), - integration: Arc::clone(&self.integration), - cancellation: cancellation.handle(), - metadata, - }; - let driver = match self.factory.start(ctx).await { - Ok(driver) => driver, - Err(error) => { - let _ = self.integration.unbind(&acp_session_id); - drain_task.abort(); - let _ = drain_task.await; - return Err(error); - } +#[cfg(feature = "unstable-inject")] +fn validate_inject_content( + content: &[wire::ContentBlock], +) -> 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)?; + Ok((items, bytes)) +} - let (commands, rx) = mpsc::unbounded_channel(); - let worker_active_prompt = Arc::clone(&active_prompt); - let worker_driving_prompt = Arc::clone(&driving_prompt); - 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_cancellation = cancellation.handle(); - let task = tokio::spawn(async move { - session_worker( - worker_session_id, - driver, - client, - integration, - worker_cancellation, - worker_active_prompt, - worker_driving_prompt, - worker_cancelled_prompt, - worker_lifecycle, - rx, - ) - .await; - }); - let entry = Arc::new(SessionEntry { - commands, - cancellation, - info, - active_prompt, - driving_prompt, - cancelled_prompt, - next_prompt_owner: AtomicU64::new(1), - closed: AtomicBool::new(false), - lifecycle, - task: Mutex::new(Some(task)), - drain_task: Mutex::new(Some(drain_task)), - }); - self.sessions +#[cfg(feature = "unstable-inject")] +struct PendingInject { + message_id: wire::MessageId, + content: Vec, + items: Vec, + bytes: InjectionBytes, + commitment: InjectCommitment, +} + +#[cfg(feature = "unstable-inject")] +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum InjectCommitment { + Reserved, + Committed, + Ready, +} + +#[cfg(feature = "unstable-inject")] +#[derive(Default)] +struct InjectionState { + running: bool, + cancelled: bool, + at_boundary: bool, + pending: VecDeque, + delivering: Option, + pending_bytes: InjectionBytes, + accepted_count: usize, + delivered: VecDeque, +} + +#[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, +} + +#[cfg(feature = "unstable-inject")] +enum BoundaryAction { + Wait, + Deliver(PendingInject), + Complete(AcpInjectionBoundary), +} + +#[cfg(feature = "unstable-inject")] +enum PendingTransition { + Applied, + 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 reset_cancellation(&self) { + self.state .lock() - .await - .insert(acp_session_id.clone(), entry); - Ok(wire::NewSessionResponse::new(acp_session_id)) + .unwrap_or_else(|error| error.into_inner()) + .cancelled = false; } - async fn list_sessions( - &self, - request: wire::ListSessionsRequest, - ) -> Result { - if request.cursor.is_some() { - return Err(AcpRuntimeError::Unsupported( - "ACP v2 session list cursors are not supported".into(), + 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) + .exceeds_budget() + { + return Err(agent_client_protocol::Error::new( + -32602, + "pending session injection budget exceeded", )); } - let sessions = self.sessions.lock().await; - let mut infos = sessions - .values() - .filter(|entry| { - !entry.closed.load(Ordering::Acquire) - && request - .cwd - .as_ref() - .is_none_or(|cwd| cwd == &entry.info.cwd) - }) - .map(|entry| entry.info.clone()) - .collect::>(); - infos.sort_by(|left, right| { - left.session_id - .to_string() - .cmp(&right.session_id.to_string()) - }); - Ok(wire::ListSessionsResponse::new(infos)) + let pending_bytes = state.pending_bytes.saturating_add(pending.bytes); + state.pending.push_back(pending); + state.pending_bytes = pending_bytes; + Ok(()) } - async fn resume_session( - &self, - request: wire::ResumeSessionRequest, - ) -> Result { - if request.replay_from.is_some() { - return Err(AcpRuntimeError::Unsupported( - "ACP v2 session replay is not supported".into(), - )); + 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; } - let entry = self - .sessions - .lock() - .await - .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 { - return Err(AcpRuntimeError::SessionNotFound( - request.session_id.to_string(), - )); + 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); } - if entry.info.additional_directories != request.additional_directories { - return Err(AcpRuntimeError::Unsupported( - "changing ACP v2 session directories on resume is not supported".into(), - )); + 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) } - Ok(wire::ResumeSessionResponse::new()) } - async fn prompt( + 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 replace_transition( &self, - request: wire::PromptRequest, - ) -> Result, AcpRuntimeError> { - let items = prompt_to_items(&request)?; - let entry = self - .sessions - .lock() - .await - .get(&request.session_id) - .cloned() - .ok_or_else(|| AcpRuntimeError::SessionNotFound(request.session_id.to_string()))?; - let (start_tx, start_rx) = oneshot::channel(); + message_id: &wire::MessageId, + content: &[wire::ContentBlock], + items: &[Item], + bytes: InjectionBytes, + ) -> 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) { - let _lifecycle = entry - .lifecycle - .lock() - .unwrap_or_else(|error| error.into_inner()); - if entry.closed.load(Ordering::Acquire) { - return Err(AcpRuntimeError::SessionNotFound( - request.session_id.to_string(), - )); - } - let mut owner = entry.next_prompt_owner.fetch_add(1, Ordering::Relaxed); - if owner == 0 { - owner = entry.next_prompt_owner.fetch_add(1, Ordering::Relaxed); + if state.pending[index].commitment != InjectCommitment::Ready { + return Ok(PendingTransition::WaitForDelivery); } - if entry - .active_prompt - .compare_exchange(0, owner, Ordering::AcqRel, Ordering::Acquire) - .is_err() - { - return Err(AcpRuntimeError::Unsupported( - "session is already running a prompt".into(), + // 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 + .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", )); } - let prompt_cancellation = entry.cancellation.handle().checkpoint(); - let cancellation_generation = prompt_cancellation.generation(); - if entry - .commands - .send(SessionCommand::Prompt { - request, - items, - prompt_cancellation, - cancellation_generation, - owner, - start: start_rx, - }) - .is_err() - { - release_prompt(&entry.active_prompt, owner); - return Err(AcpRuntimeError::ClientClosed); - } + // 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]; + 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.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) + } 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) } - Ok(start_tx) } - async fn cancel( - &self, - notification: wire::CancelSessionNotification, - ) -> Result<(), AcpRuntimeError> { - let entry = self - .sessions - .lock() - .await - .get(¬ification.session_id) - .cloned(); - if let Some(entry) = entry { - let _lifecycle = entry - .lifecycle - .lock() - .unwrap_or_else(|error| error.into_inner()); - if !entry.closed.load(Ordering::Acquire) { - let driving = entry.driving_prompt.load(Ordering::Acquire); - if driving != 0 { - entry.cancellation.interrupt(); - } else { - let queued = entry.active_prompt.load(Ordering::Acquire); - if queued != 0 { - entry.cancelled_prompt.store(queued, Ordering::Release); - } else { - entry.cancellation.interrupt(); - } - } - } + 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 + }); + if ready && Self::remove_pending(&mut state, message_id) { + drop(state); + self.changed.notify_waiters(); + PendingTransition::Applied + } else if state.delivering.as_ref() == Some(message_id) + || state + .pending + .iter() + .any(|pending| &pending.message_id == message_id) + { + PendingTransition::WaitForDelivery + } else if state.delivered.contains(message_id) { + PendingTransition::AlreadyDelivered + } else { + PendingTransition::Unknown } - Ok(()) } - async fn close( - &self, - request: wire::CloseSessionRequest, - ) -> Result { - let entry = self - .sessions - .lock() - .await - .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)?; - stop_client(entry).await; - Ok(wire::CloseSessionResponse::new()) + 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(); } - async fn shutdown(&self) { - let sessions = { - let mut sessions = self.sessions.lock().await; - sessions.drain().collect::>() - }; - let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(2); - let mut session_tasks = Vec::with_capacity(sessions.len()); - for (session_id, entry) in &sessions { - signal_session_stop(entry); - if let Some(task) = take_task(&entry.task) { - session_tasks.push(task); - } - let _ = self.integration.unbind(session_id); - } - join_tasks_until(deadline, session_tasks).await; + 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(); + } - let drain_tasks = sessions + 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_map(|(_, entry)| take_task(&entry.drain_task)) - .collect(); - drop(sessions); - join_tasks_until(deadline, drain_tasks).await; + .filter(|pending| pending.commitment == InjectCommitment::Ready) + .map(|pending| pending.bytes) + .fold(InjectionBytes::default(), InjectionBytes::saturating_add); + state + .pending + .retain(|pending| pending.commitment != InjectCommitment::Ready); + state.pending_bytes = state.pending_bytes.saturating_sub(discarded_bytes); + drop(state); + self.changed.notify_waiters(); } } -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(); - let _ = entry.commands.send(SessionCommand::Shutdown); +#[cfg(feature = "unstable-inject")] +struct ReservedInject { + session: Arc, + message_id: Option, } -fn take_task( - task: &Mutex>>, -) -> Option> { - task.lock() - .unwrap_or_else(|error| error.into_inner()) - .take() -} +#[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(), + ) + } -async fn join_tasks_until( - deadline: tokio::time::Instant, - mut tasks: Vec>, -) { - if tokio::time::timeout_at(deadline, async { - for task in &mut tasks { - let _ = task.await; - } - }) - .await - .is_err() - { - for task in tasks { - task.abort(); + 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); } -} -async fn stop_session(entry: Arc) { - signal_session_stop(&entry); - if let Some(task) = take_task(&entry.task) { - join_tasks_until( - tokio::time::Instant::now() + std::time::Duration::from_secs(2), - vec![task], - ) - .await; + fn discard(mut self) { + let message_id = self + .message_id + .take() + .expect("reserved injection has a message id"); + self.session.injection.discard(&message_id); } } -async fn stop_client(entry: Arc) { - if let Some(task) = take_task(&entry.drain_task) { - let _ = task.await; +#[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); + } } } -fn release_prompt(active_prompt: &AtomicU64, owner: u64) { - let _ = active_prompt.compare_exchange(owner, 0, Ordering::AcqRel, Ordering::Acquire); +#[cfg(feature = "unstable-inject")] +struct InjectRequestPermit { + integration: AcpIntegration, + sequence: Option, } -fn clear_prompt_tracking( - active_prompt: &AtomicU64, - driving_prompt: &AtomicU64, - cancelled_prompt: &AtomicU64, - owner: u64, -) { - let _ = driving_prompt.compare_exchange(owner, 0, Ordering::AcqRel, Ordering::Acquire); - let _ = cancelled_prompt.compare_exchange(owner, 0, Ordering::AcqRel, Ordering::Acquire); - release_prompt(active_prompt, owner); +#[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); + } + } } -#[allow(clippy::too_many_arguments)] -fn fail_accepted_prompt( - client: &ClientHandle, - integration: &AcpIntegration, - session_id: &wire::SessionId, - active_prompt: &AtomicU64, - driving_prompt: &AtomicU64, - cancelled_prompt: &AtomicU64, - lifecycle: &Mutex<()>, - owner: u64, - prompt_began: bool, -) { - let _lifecycle = lifecycle.lock().unwrap_or_else(|error| error.into_inner()); - if prompt_began { - integration.finish_prompt(session_id, owner); - } - clear_prompt_tracking(active_prompt, driving_prompt, cancelled_prompt, owner); - if let Err(error) = client.update( - session_id.clone(), - wire::SessionUpdate::StateUpdate(wire::StateUpdate::Idle( - wire::IdleStateUpdate::new().stop_reason(error_stop_reason()), - )), - ) { - tracing::debug!(%error, owner, "failed to queue ACP v2 error idle update"); - } +/// 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, } -async fn session_worker( - session_id: wire::SessionId, - mut driver: agentkit_loop::LoopDriver, - client: ClientHandle, - integration: Arc, - cancellation: CancellationHandle, - active_prompt: Arc, - driving_prompt: Arc, - cancelled_prompt: Arc, - lifecycle: Arc>, - mut commands: mpsc::UnboundedReceiver, -) where - S: ModelSession + Send + 'static, -{ - enum SessionAction { - Command(Option), - LoopUpdate(Result<(), LoopError>), +#[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() } - // Let one fresh prompt win a race with unrelated background work, then - // prefer the deferred update so repeated prompts cannot starve its delivery. - let mut prefer_loop_update = false; - loop { - let action = if prefer_loop_update { - tokio::select! { - biased; - wake = driver.wait_for_loop_update() => SessionAction::LoopUpdate(wake), - command = commands.recv() => SessionAction::Command(command), - } + /// 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 { - tokio::select! { - biased; - command = commands.recv() => SessionAction::Command(command), - wake = driver.wait_for_loop_update() => SessionAction::LoopUpdate(wake), - } + None }; - let command = match action { - SessionAction::Command(Some(command)) => command, - SessionAction::Command(None) => break, - SessionAction::LoopUpdate(wake) => { - if let Err(error) = wake { - tracing::debug!(%error, "failed waiting for an ACP v2 loop update"); - break; - } - prefer_loop_update = false; - drive_prompt( - &mut driver, - &integration, - &session_id, - &cancellation, - cancellation.generation(), - ) - .await; - if let Err(error) = client.flush().await { - tracing::debug!(%error, "failed to flush idle ACP v2 loop update"); - } - continue; + 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 SessionCommand::Prompt { - request, + 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, bytes) = validate_inject_content(&request.content)?; + 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, - prompt_cancellation, - cancellation_generation, - owner, - start, - } = command - else { - break; + 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), }; - if start.await.is_err() { - let _lifecycle = lifecycle.lock().unwrap_or_else(|error| error.into_inner()); - clear_prompt_tracking(&active_prompt, &driving_prompt, &cancelled_prompt, owner); - continue; + 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 should_start = { - let _lifecycle = lifecycle.lock().unwrap_or_else(|error| error.into_inner()); - let still_active = active_prompt.load(Ordering::Acquire) == owner; - let was_cancelled = cancelled_prompt.load(Ordering::Acquire) == owner; - if !still_active || was_cancelled { - clear_prompt_tracking(&active_prompt, &driving_prompt, &cancelled_prompt, owner); - false - } else { - driving_prompt.store(owner, Ordering::Release); - true + let reserved = match self.reserve_inject(request) { + Ok(reserved) => reserved, + Err(error) => { + responder.respond_with_result(Err(error))?; + drop(permit); + return Ok(None); } }; - if !should_start { - continue; - } - - if let Err(error) = driver.submit_input(items) { - tracing::debug!(%error, owner, "failed to submit accepted ACP v2 prompt"); - fail_accepted_prompt( - &client, - &integration, - &session_id, - &active_prompt, - &driving_prompt, - &cancelled_prompt, - &lifecycle, - owner, - false, - ); - continue; - } - let user_message_id = - match integration.begin_prompt(&session_id, owner, prompt_cancellation) { - Ok(message_id) => message_id, - Err(error) => { - tracing::debug!(%error, owner, "failed to begin accepted ACP v2 prompt"); - fail_accepted_prompt( - &client, - &integration, - &session_id, - &active_prompt, - &driving_prompt, - &cancelled_prompt, - &lifecycle, - owner, - false, - ); - continue; - } - }; - - if let Err(error) = client.update( - 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"); - fail_accepted_prompt( - &client, - &integration, - &session_id, - &active_prompt, - &driving_prompt, - &cancelled_prompt, - &lifecycle, - owner, - true, - ); - continue; - } + Ok(Some(AcpInjectRequest { + reserved: Some(reserved), + responder: Some(responder), + cancellation, + _permit: permit, + })) + } - drive_prompt( - &mut driver, - &integration, - &session_id, - &cancellation, - cancellation_generation, - ) - .await; - { - let _lifecycle = lifecycle.lock().unwrap_or_else(|error| error.into_inner()); - integration.finish_prompt(&session_id, owner); - clear_prompt_tracking(&active_prompt, &driving_prompt, &cancelled_prompt, owner); - } - if let Err(error) = client.flush().await { - tracing::debug!(%error, "failed to flush ACP v2 output"); - } - prefer_loop_update = true; + /// 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(()) } -} -async fn drive_prompt( - driver: &mut agentkit_loop::LoopDriver, - integration: &AcpIntegration, - session_id: &wire::SessionId, - cancellation: &CancellationHandle, - generation: u64, -) 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; + /// 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)); } - Ok(LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_))) => return, - Ok(LoopStep::Interrupt(LoopInterrupt::AfterToolResult(_))) => continue, - Ok(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"); + 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, + ))); } - return; - } - Err(error) => { - tracing::debug!(%error, "ACP v2 agent loop failed"); - return; } + changed.await; } } -} - -fn prompt_to_items(request: &wire::PromptRequest) -> Result, AcpRuntimeError> { - let mut user_parts = Vec::new(); - let mut context_items = Vec::new(); - for block in &request.prompt { - match block { - wire::ContentBlock::Text(text) => { - user_parts.push(Part::Text(TextPart::new(text.text.clone()))); - } - wire::ContentBlock::Image(image) => { - let mime_type = image.mime_type.as_ref(); - user_parts.push(Part::media( - Modality::Image, - mime_type, - crate::data_url_ref(mime_type, &image.data), - )); - } - wire::ContentBlock::Audio(audio) => { - let mime_type = audio.mime_type.as_ref(); - user_parts.push(Part::media( - Modality::Audio, - mime_type, - crate::data_url_ref(mime_type, &audio.data), - )); - } - wire::ContentBlock::ResourceLink(link) => { - context_items.push(resource_link_item(link)); - } - wire::ContentBlock::Resource(resource) => { - context_items.push(resource_item(resource)?); - } - wire::ContentBlock::Other(other) => { - return Err(AcpRuntimeError::UnsupportedContent(format!( - "unknown ACP v2 content block: {}", - other.type_ - ))); + /// 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)); } - _ => { - return Err(AcpRuntimeError::UnsupportedContent( - "unknown ACP v2 content block".into(), - )); + match session.injection.revoke_transition(&request.message_id) { + PendingTransition::Applied => { + return Ok(wire::RevokeInjectSessionResponse::new()); + } + 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; } } +} - let mut items = context_items; - if !user_parts.is_empty() { - items.push(Item::new(ItemKind::User, user_parts)); - } else if !items.is_empty() { - items.push(Item::text(ItemKind::User, "Use the provided context.")); - } - Ok(items) +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>>, } -fn resource_link_item(link: &wire::ResourceLink) -> Item { - let mut metadata = MetadataMap::new(); - metadata.insert("acp.resource.uri".into(), json!(link.uri)); - metadata.insert("acp.resource.name".into(), json!(link.name)); - if let Some(description) = &link.description { - metadata.insert("acp.resource.description".into(), json!(description)); - } - if let Some(mime_type) = &link.mime_type { - metadata.insert("acp.resource.mime_type".into(), json!(mime_type)); - } - Item::new( - ItemKind::Context, - vec![Part::file(DataRef::uri(link.uri.clone()))], - ) - .with_metadata(metadata) +enum SessionCommand { + Prompt { + request: wire::PromptRequest, + items: Vec, + message_id: wire::MessageId, + prompt_cancellation: TurnCancellation, + cancellation_generation: u64, + owner: u64, + start: oneshot::Receiver<()>, + }, + Shutdown, } -fn resource_item(resource: &wire::EmbeddedResource) -> Result { - match &resource.resource { - wire::EmbeddedResourceResource::TextResourceContents(text) => { - let mut metadata = MetadataMap::new(); - metadata.insert("acp.resource.uri".into(), json!(text.uri)); - if let Some(mime_type) = &text.mime_type { - metadata.insert("acp.resource.mime_type".into(), json!(mime_type)); - } - Ok(Item::text(ItemKind::Context, text.text.clone()).with_metadata(metadata)) - } - wire::EmbeddedResourceResource::BlobResourceContents(blob) => { - let mime_type = blob - .mime_type - .as_ref() - .map(|mime_type| mime_type.as_ref()) - .unwrap_or("application/octet-stream"); - let mut metadata = MetadataMap::new(); - metadata.insert("acp.resource.uri".into(), json!(blob.uri)); - metadata.insert("acp.resource.mime_type".into(), json!(mime_type)); - Ok(Item::new( - ItemKind::Context, - vec![Part::media( - Modality::Binary, - mime_type, - crate::data_url_ref(mime_type, &blob.blob), - )], - ) - .with_metadata(metadata)) - } - _ => Err(AcpRuntimeError::UnsupportedContent( - "unknown ACP v2 embedded resource".into(), - )), - } +struct RuntimeState +where + M: ModelAdapter, +{ + factory: Arc>, + integration: Arc, + sessions: AsyncMutex>>, + next_session: AtomicU64, + name: String, + version: String, } -fn event_to_update( - event: &AgentEvent, - message_ids: Option<&CurrentMessageIds>, - part_kinds: &mut HashMap, -) -> Option { - match event { - AgentEvent::ContentDelta(delta) => delta_to_update(delta, message_ids, part_kinds), - AgentEvent::ToolCallRequested(call) => { - Some(wire::SessionUpdate::ToolCallUpdate(tool_call_update(call))) +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, } - AgentEvent::ToolExecutionStarted(call) => Some(wire::SessionUpdate::ToolCallUpdate( - tool_status_update(&call.id, wire::ToolCallStatus::InProgress), - )), - AgentEvent::ToolExecutionProgress(result) => Some(wire::SessionUpdate::ToolCallUpdate( - tool_result_update(result, wire::ToolCallStatus::InProgress), - )), - AgentEvent::ToolResultReceived(result) => { - Some(wire::SessionUpdate::ToolCallUpdate(tool_result_update( - result, - if result.is_error { - wire::ToolCallStatus::Failed - } else { - wire::ToolCallStatus::Completed - }, - ))) - } - AgentEvent::Warning { message } => { - tracing::warn!(%message, "agentkit warning while routing ACP v2 event"); - None - } - AgentEvent::RunFailed { message } => { - tracing::debug!(%message, "agentkit run failed while routing ACP v2 event"); - None - } - _ => None, } -} -fn delta_to_update( - delta: &Delta, - message_ids: Option<&CurrentMessageIds>, - part_kinds: &mut HashMap, -) -> Option { - match delta { - Delta::BeginPart { part_id, kind } => { - part_kinds.insert(part_id.clone(), *kind); - None - } - Delta::AppendText { part_id, chunk } => { - let message_ids = message_ids?; - let content = wire::ContentBlock::Text(wire::TextContent::new(chunk.clone())); - match part_kinds.get(part_id) { - Some(PartKind::Reasoning) => Some(wire::SessionUpdate::AgentThoughtChunk( - wire::ContentChunk::new(content, message_ids.thought.clone()), - )), - Some(PartKind::Text) | None => Some(wire::SessionUpdate::AgentMessageChunk( - wire::ContentChunk::new(content, message_ids.agent.clone()), - )), - Some(_) => None, - } + 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(), + )); } - Delta::CommitPart { .. } - | Delta::AppendBytes { .. } - | Delta::ReplaceStructured { .. } - | Delta::SetMetadata { .. } => None, + Ok(wire::InitializeResponse::new( + wire::ProtocolVersion::V2, + wire::Implementation::new(self.name.clone(), self.version.clone()), + ) + .capabilities(agent_capabilities())) } -} - -fn tool_call_update(call: &ToolCallPart) -> wire::ToolCallUpdate { - wire::ToolCallUpdate::new(call.id.to_string()) - .title(call.name.clone()) - .status(wire::ToolCallStatus::Pending) - .raw_input(call.input.clone()) -} -fn tool_status_update( - call_id: &agentkit_core::ToolCallId, - status: wire::ToolCallStatus, -) -> wire::ToolCallUpdate { - wire::ToolCallUpdate::new(call_id.to_string()).status(status) -} + 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()) + .additional_directories(request.additional_directories.clone()); -fn tool_result_update( - result: &ToolResultPart, - status: wire::ToolCallStatus, -) -> wire::ToolCallUpdate { - wire::ToolCallUpdate::new(result.call_id.to_string()) - .status(status) - .raw_output(crate::tool_output_raw(&result.output)) - .content(tool_output_content(&result.output)) -} + let mut metadata = MetadataMap::new(); + metadata.insert("acp.protocol_version".into(), json!(2)); + metadata.insert("acp.cwd".into(), json!(request.cwd)); + metadata.insert( + "acp.additional_directories".into(), + json!(request.additional_directories), + ); -fn tool_output_content(output: &ToolOutput) -> Option> { - let content = match output { - ToolOutput::Text(text) => vec![text_to_tool_content(text.clone())], - ToolOutput::Structured(value) => vec![text_to_tool_content(value.to_string())], - ToolOutput::Parts(parts) => parts.iter().filter_map(part_to_tool_content).collect(), - ToolOutput::Files(files) => files.iter().map(file_to_tool_content).collect(), - }; - (!content.is_empty()).then_some(content) -} + 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 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_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.clone(), + cwd: request.cwd.into_inner(), + additional_directories: request + .additional_directories + .into_iter() + .map(wire::AbsolutePath::into_inner) + .collect(), + integration: Arc::clone(&self.integration), + cancellation: cancellation.handle(), + metadata, + }; + let driver = match self.factory.start(ctx).await { + Ok(driver) => driver, + Err(error) => { + let _ = self.integration.unbind_session(&acp_session_id); + drain_task.abort(); + let _ = drain_task.await; + return Err(error); + } + }; -fn part_to_tool_content(part: &Part) -> Option { - match part { - Part::Text(text) => Some(text_to_tool_content(text.text.clone())), - Part::Structured(value) => Some(structured_to_tool_content(value)), - Part::Media(media) => Some(media_to_tool_content(media)), - Part::File(file) => Some(file_to_tool_content(file)), - Part::Reasoning(reasoning) => reasoning - .summary - .as_ref() - .map(|summary| text_to_tool_content(summary.clone())), - Part::Custom(custom) => Some(text_to_tool_content( - custom - .value - .as_ref() - .map(ToString::to_string) - .or_else(|| custom.data.as_ref().map(crate::data_ref_payload)) - .unwrap_or_else(|| custom.kind.clone()), - )), - Part::ToolCall(_) | Part::ToolResult(_) => None, + let (commands, rx) = mpsc::unbounded_channel(); + let worker_active_prompt = Arc::clone(&active_prompt); + let worker_driving_prompt = Arc::clone(&driving_prompt); + let worker_cancelled_prompt = Arc::clone(&cancelled_prompt); + let worker_lifecycle = Arc::clone(&lifecycle); + let integration = Arc::clone(&self.integration); + let worker_session = session.clone(); + let worker_cancellation = cancellation.handle(); + let task = tokio::spawn(async move { + session_worker( + worker_session, + driver, + client, + integration, + worker_cancellation, + worker_active_prompt, + worker_driving_prompt, + worker_cancelled_prompt, + worker_lifecycle, + rx, + ) + .await; + }); + let entry = Arc::new(SessionEntry { + commands, + session, + cancellation, + info, + active_prompt, + driving_prompt, + cancelled_prompt, + next_prompt_owner: AtomicU64::new(1), + closed: AtomicBool::new(false), + lifecycle, + task: Mutex::new(Some(task)), + drain_task: Mutex::new(Some(drain_task)), + }); + self.sessions + .lock() + .await + .insert(acp_session_id.clone(), entry); + Ok(wire::NewSessionResponse::new(acp_session_id)) } -} - -fn text_to_tool_content(text: String) -> wire::ToolCallContent { - wire::ToolCallContent::Content(Box::new(wire::Content::new(wire::ContentBlock::Text( - wire::TextContent::new(text), - )))) -} - -fn structured_to_tool_content(part: &StructuredPart) -> wire::ToolCallContent { - text_to_tool_content(part.value.to_string()) -} -fn media_to_tool_content(media: &MediaPart) -> wire::ToolCallContent { - match media.modality { - Modality::Image - if matches!(media.data, DataRef::InlineText(_) | DataRef::InlineBytes(_)) => - { - let mut image = wire::ImageContent::new( - crate::data_ref_base64_payload(&media.data), - media.mime_type.clone(), - ); - if let Some(uri) = crate::data_ref_uri(&media.data) { - image = image.uri(uri); - } - wire::ToolCallContent::Content(Box::new(wire::Content::new(wire::ContentBlock::Image( - image, - )))) - } - Modality::Audio - if matches!(media.data, DataRef::InlineText(_) | DataRef::InlineBytes(_)) => - { - wire::ToolCallContent::Content(Box::new(wire::Content::new(wire::ContentBlock::Audio( - wire::AudioContent::new( - crate::data_ref_base64_payload(&media.data), - media.mime_type.clone(), - ), - )))) - } - Modality::Image | Modality::Audio | Modality::Video | Modality::Binary => { - data_ref_to_resource_content(None, Some(&media.mime_type), &media.data) + async fn list_sessions( + &self, + request: wire::ListSessionsRequest, + ) -> Result { + if request.cursor.is_some() { + return Err(AcpRuntimeError::Unsupported( + "ACP v2 session list cursors are not supported".into(), + )); } - } -} - -fn file_to_tool_content(file: &FilePart) -> wire::ToolCallContent { - data_ref_to_resource_content(file.name.as_deref(), file.mime_type.as_deref(), &file.data) -} - -fn data_ref_to_resource_content( - name: Option<&str>, - mime_type: Option<&str>, - data: &DataRef, -) -> wire::ToolCallContent { - let content = match data { - DataRef::Uri(uri) => { - let mut link = wire::ResourceLink::new(name.unwrap_or(uri), uri.clone()); - if let Some(mime_type) = mime_type { - link = link.mime_type(mime_type.to_string()); - } - wire::ContentBlock::ResourceLink(link) + let sessions = self.sessions.lock().await; + let mut infos = sessions + .values() + .filter(|entry| { + !entry.session.is_closed() + && request + .cwd + .as_ref() + .is_none_or(|cwd| cwd == &entry.info.cwd) + }) + .map(|entry| entry.info.clone()) + .collect::>(); + infos.sort_by(|left, right| { + left.session_id + .to_string() + .cmp(&right.session_id.to_string()) + }); + Ok(wire::ListSessionsResponse::new(infos)) + } + + async fn resume_session( + &self, + request: wire::ResumeSessionRequest, + ) -> Result { + if request.replay_from.is_some() { + return Err(AcpRuntimeError::Unsupported( + "ACP v2 session replay is not supported".into(), + )); } - DataRef::Handle(handle) => { - let uri = format!("artifact://{handle}"); - let link_name = name.map(str::to_owned).unwrap_or_else(|| uri.clone()); - let mut link = wire::ResourceLink::new(link_name, uri); - if let Some(mime_type) = mime_type { - link = link.mime_type(mime_type.to_string()); + let entry = self + .sessions + .lock() + .await + .get(&request.session_id) + .cloned() + .ok_or_else(|| AcpRuntimeError::SessionNotFound(request.session_id.to_string()))?; + if entry.session.is_closed() || entry.info.cwd != request.cwd { + return Err(AcpRuntimeError::SessionNotFound( + request.session_id.to_string(), + )); + } + if entry.info.additional_directories != request.additional_directories { + return Err(AcpRuntimeError::Unsupported( + "changing ACP v2 session directories on resume is not supported".into(), + )); + } + Ok(wire::ResumeSessionResponse::new()) + } + + async fn prompt( + &self, + request: wire::PromptRequest, + ) -> Result<(oneshot::Sender<()>, wire::MessageId), AcpRuntimeError> { + let items = self.integration.prompt_to_items(&request)?; + let entry = self + .sessions + .lock() + .await + .get(&request.session_id) + .cloned() + .ok_or_else(|| AcpRuntimeError::SessionNotFound(request.session_id.to_string()))?; + let (start_tx, start_rx) = oneshot::channel(); + { + let _lifecycle = entry + .session + .session + .lifecycle + .lock() + .unwrap_or_else(|error| error.into_inner()); + if entry.session.is_closed() { + return Err(AcpRuntimeError::SessionNotFound( + request.session_id.to_string(), + )); } - wire::ContentBlock::ResourceLink(link) + let mut owner = entry.next_prompt_owner.fetch_add(1, Ordering::Relaxed); + if owner == 0 { + owner = entry.next_prompt_owner.fetch_add(1, Ordering::Relaxed); + } + if entry + .active_prompt + .compare_exchange(0, owner, Ordering::AcqRel, Ordering::Acquire) + .is_err() + { + return Err(AcpRuntimeError::Unsupported( + "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(); + let cancellation_generation = prompt_cancellation.generation(); + if entry + .commands + .send(SessionCommand::Prompt { + request, + items, + message_id: message_id.clone(), + prompt_cancellation, + cancellation_generation, + owner, + start: start_rx, + }) + .is_err() + { + release_prompt(&entry.active_prompt, owner); + return Err(AcpRuntimeError::ClientClosed); + } + Ok((start_tx, message_id)) } - DataRef::InlineText(text) if mime_type.is_none_or(|mime| mime.starts_with("text/")) => { - let mut resource = wire::TextResourceContents::new( - text.clone(), - crate::inline_resource_uri(name.unwrap_or("tool-output")), - ); - if let Some(mime_type) = mime_type { - resource = resource.mime_type(mime_type.to_string()); + } + + async fn cancel( + &self, + notification: wire::CancelSessionNotification, + ) -> Result<(), AcpRuntimeError> { + let entry = self + .sessions + .lock() + .await + .get(¬ification.session_id) + .cloned(); + if let Some(entry) = entry { + let _lifecycle = entry + .lifecycle + .lock() + .unwrap_or_else(|error| error.into_inner()); + 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(); + } + } } - wire::ContentBlock::Resource(wire::EmbeddedResource::new( - wire::EmbeddedResourceResource::TextResourceContents(resource), - )) } - _ => { - let mut resource = wire::BlobResourceContents::new( - crate::data_ref_base64_payload(data), - crate::inline_resource_uri(name.unwrap_or("tool-output")), - ); - if let Some(mime_type) = mime_type { - resource = resource.mime_type(mime_type.to_string()); + Ok(()) + } + + async fn close( + &self, + request: wire::CloseSessionRequest, + ) -> Result { + let entry = self + .sessions + .lock() + .await + .remove(&request.session_id) + .ok_or_else(|| AcpRuntimeError::SessionNotFound(request.session_id.to_string()))?; + stop_session(Arc::clone(&entry)).await; + self.integration.unbind_session(&request.session_id)?; + stop_client(entry).await; + Ok(wire::CloseSessionResponse::new()) + } + + async fn shutdown(&self) { + let sessions = { + let mut sessions = self.sessions.lock().await; + sessions.drain().collect::>() + }; + let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(2); + let mut session_tasks = Vec::with_capacity(sessions.len()); + for (session_id, entry) in &sessions { + signal_session_stop(entry); + if let Some(task) = take_task(&entry.task) { + session_tasks.push(task); } - wire::ContentBlock::Resource(wire::EmbeddedResource::new( - wire::EmbeddedResourceResource::BlobResourceContents(resource), - )) + let _ = self.integration.unbind_session(session_id); } - }; - wire::ToolCallContent::Content(Box::new(wire::Content::new(content))) + join_tasks_until(deadline, session_tasks).await; + + let drain_tasks = sessions + .iter() + .filter_map(|(_, entry)| take_task(&entry.drain_task)) + .collect(); + drop(sessions); + join_tasks_until(deadline, drain_tasks).await; + } } -fn error_stop_reason() -> wire::StopReason { - wire::StopReason::Other("_error".into()) +fn signal_session_stop(entry: &Arc) { + entry.closed.store(true, Ordering::Release); + entry.session.close(); + let _ = entry.commands.send(SessionCommand::Shutdown); } -fn finish_reason_to_stop_reason(reason: &FinishReason) -> wire::StopReason { - match reason { - FinishReason::Completed | FinishReason::ToolCall | FinishReason::Other(_) => { - wire::StopReason::EndTurn +fn take_task( + task: &Mutex>>, +) -> Option> { + task.lock() + .unwrap_or_else(|error| error.into_inner()) + .take() +} + +async fn join_tasks_until( + deadline: tokio::time::Instant, + mut tasks: Vec>, +) { + if tokio::time::timeout_at(deadline, async { + for task in &mut tasks { + let _ = task.await; + } + }) + .await + .is_err() + { + for task in tasks { + task.abort(); } - FinishReason::MaxTokens => wire::StopReason::MaxTokens, - FinishReason::Cancelled => wire::StopReason::Cancelled, - FinishReason::Blocked | FinishReason::Error => error_stop_reason(), } } -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()), - ), - ) +async fn stop_session(entry: Arc) { + signal_session_stop(&entry); + if let Some(task) = take_task(&entry.task) { + join_tasks_until( + tokio::time::Instant::now() + std::time::Duration::from_secs(2), + vec![task], + ) + .await; + } } -#[cfg(test)] -mod tests { - use super::*; - use std::collections::VecDeque; - use std::sync::atomic::AtomicUsize; - use std::time::Duration; - - use agent_client_protocol::Channel; - use agentkit_core::{ - ItemKind, ToolCallId, ToolOutput, ToolResultPart, TurnCancellation, TurnId, - }; - use agentkit_integration_tests::mock_model::{MockAdapter, TurnScript}; - use agentkit_integration_tests::mock_tool::BlockingTool; - use agentkit_loop::{ - Agent, ModelSession, ModelTurn, ModelTurnEvent, ModelTurnResult, SessionConfig, - TurnRequest, TurnResult, - }; - use agentkit_task_manager::{AsyncTaskManager, RoutingDecision}; - use agentkit_tools_core::{ - Tool, ToolContext, ToolError, ToolRegistry, ToolRequest, ToolResult, ToolSpec, - }; - use tokio::sync::Notify; - - #[derive(Clone)] - struct TestFactory { - adapter: A, +async fn stop_client(entry: Arc) { + let task = take_task(&entry.drain_task); + drop(entry); + if let Some(task) = task { + let _ = task.await; } +} - #[async_trait] - impl AcpAgentFactory for TestFactory - where - A: ModelAdapter + Clone + Send + Sync + 'static, - A::Session: Send + 'static, - { - async fn start( - &self, - ctx: AcpAgentFactoryContext, - ) -> Result, AcpRuntimeError> { - 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())) - } +fn release_prompt(active_prompt: &AtomicU64, owner: u64) { + let _ = active_prompt.compare_exchange(owner, 0, Ordering::AcqRel, Ordering::Acquire); +} + +fn clear_prompt_tracking( + active_prompt: &AtomicU64, + driving_prompt: &AtomicU64, + cancelled_prompt: &AtomicU64, + owner: u64, +) { + let _ = driving_prompt.compare_exchange(owner, 0, Ordering::AcqRel, Ordering::Acquire); + let _ = cancelled_prompt.compare_exchange(owner, 0, Ordering::AcqRel, Ordering::Acquire); + release_prompt(active_prompt, owner); +} + +#[allow(clippy::too_many_arguments)] +fn fail_accepted_prompt( + client: &ClientHandle, + integration: &AcpIntegration, + session_id: &wire::SessionId, + active_prompt: &AtomicU64, + driving_prompt: &AtomicU64, + cancelled_prompt: &AtomicU64, + lifecycle: &Mutex<()>, + owner: u64, + prompt_began: bool, +) { + let _lifecycle = lifecycle.lock().unwrap_or_else(|error| error.into_inner()); + if prompt_began { + integration.finish_prompt_owner(session_id, owner); + } + clear_prompt_tracking(active_prompt, driving_prompt, cancelled_prompt, owner); + if let Err(error) = client.update_for( + session_id.clone(), + wire::SessionUpdate::StateUpdate(wire::StateUpdate::Idle( + wire::IdleStateUpdate::new().stop_reason(error_stop_reason()), + )), + ) { + tracing::debug!(%error, owner, "failed to queue ACP v2 error idle update"); } +} - #[derive(Clone)] - struct CancellationAwareTool { - spec: ToolSpec, - entered: Arc, - cleaned: Arc, +async fn session_worker( + session: AcpSessionHandle, + mut driver: agentkit_loop::LoopDriver, + client: ClientHandle, + integration: Arc, + cancellation: CancellationHandle, + active_prompt: Arc, + driving_prompt: Arc, + cancelled_prompt: Arc, + lifecycle: Arc>, + mut commands: mpsc::UnboundedReceiver, +) where + S: ModelSession + Send + 'static, +{ + let session_id = session.acp_session_id().clone(); + enum SessionAction { + Command(Option), + LoopUpdate(Result<(), LoopError>), } - impl CancellationAwareTool { - fn new() -> Self { - Self { - spec: ToolSpec::new( - "blocking_tool", - "waits for turn cancellation", - json!({ "type": "object" }), - ), - entered: Arc::new(AtomicUsize::new(0)), - cleaned: Arc::new(AtomicUsize::new(0)), + // Let one fresh prompt win a race with unrelated background work, then + // prefer the deferred update so repeated prompts cannot starve its delivery. + let mut prefer_loop_update = false; + loop { + let action = if prefer_loop_update { + tokio::select! { + biased; + wake = driver.wait_for_loop_update() => SessionAction::LoopUpdate(wake), + command = commands.recv() => SessionAction::Command(command), } - } - - async fn wait_for_entered(&self, count: usize) { - tokio::time::timeout(Duration::from_secs(2), async { - while self.entered.load(Ordering::Acquire) < count { - tokio::task::yield_now().await; + } else { + tokio::select! { + biased; + command = commands.recv() => SessionAction::Command(command), + wake = driver.wait_for_loop_update() => SessionAction::LoopUpdate(wake), + } + }; + let command = match action { + SessionAction::Command(Some(command)) => command, + SessionAction::Command(None) => break, + SessionAction::LoopUpdate(wake) => { + if let Err(error) = wake { + tracing::debug!(%error, "failed waiting for an ACP v2 loop update"); + break; } - }) - .await - .expect("tool did not start"); + prefer_loop_update = false; + drive_prompt( + &mut driver, + &integration, + &session_id, + &cancellation, + cancellation.generation(), + #[cfg(feature = "unstable-inject")] + &session, + ) + .await; + if let Err(error) = client.flush().await { + tracing::debug!(%error, "failed to flush idle ACP v2 loop update"); + } + continue; + } + }; + let SessionCommand::Prompt { + request, + items, + message_id: user_message_id, + prompt_cancellation, + cancellation_generation, + owner, + start, + } = command + else { + break; + }; + if start.await.is_err() { + let _lifecycle = lifecycle.lock().unwrap_or_else(|error| error.into_inner()); + clear_prompt_tracking(&active_prompt, &driving_prompt, &cancelled_prompt, owner); + continue; } - } + let should_start = { + let _lifecycle = lifecycle.lock().unwrap_or_else(|error| error.into_inner()); + let still_active = active_prompt.load(Ordering::Acquire) == owner; + let was_cancelled = cancelled_prompt.load(Ordering::Acquire) == owner; + if !still_active || was_cancelled { + clear_prompt_tracking(&active_prompt, &driving_prompt, &cancelled_prompt, owner); + false + } else { + driving_prompt.store(owner, Ordering::Release); + true + } + }; + if !should_start { + #[cfg(feature = "unstable-inject")] + session.stop_injection_turn(); + continue; + } + #[cfg(feature = "unstable-inject")] + session.start_injection_turn(); - #[async_trait] - impl Tool for CancellationAwareTool { - fn spec(&self) -> &ToolSpec { - &self.spec + 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, + &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; } - async fn invoke( - &self, - request: ToolRequest, - ctx: &mut ToolContext<'_>, - ) -> Result { - self.entered.fetch_add(1, Ordering::AcqRel); - ctx.cancellation - .as_ref() - .expect("turn cancellation installed") - .cancelled() - .await; - self.cleaned.fetch_add(1, Ordering::AcqRel); - Ok(ToolResult::new(ToolResultPart::error( - request.call_id, - ToolOutput::text("cancelled"), - ))) + 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, + &session_id, + &active_prompt, + &driving_prompt, + &cancelled_prompt, + &lifecycle, + owner, + true, + ); + continue; } - } - #[derive(Clone)] - struct ToolTestFactory { - adapter: MockAdapter, - tool: CancellationAwareTool, - } - - #[derive(Clone)] - struct BackgroundToolTestFactory { - adapter: A, - tool: BlockingTool, - } - - #[async_trait] - impl AcpAgentFactory for ToolTestFactory { - 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())) + drive_prompt( + &mut driver, + &integration, + &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_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"); } + prefer_loop_update = true; } +} - #[async_trait] - impl AcpAgentFactory for BackgroundToolTestFactory +/// 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 - A: ModelAdapter + Clone + Send + Sync + 'static, - A::Session: Send + 'static, + S: ModelSession + Send + 'static, { - async fn start( - &self, - ctx: AcpAgentFactoryContext, - ) -> Result, AcpRuntimeError> { - Agent::builder() - .model(self.adapter.clone()) - .add_tool_source(ToolRegistry::new().with(self.tool.clone())) - .task_manager( - AsyncTaskManager::new().routing(|_: &ToolRequest| RoutingDecision::Background), - ) - .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())) + 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; + } + } } } +} - struct BlockingInferenceState { - scripts: Mutex>, - next_turn: AtomicUsize, - blocked_turn: usize, - entered: AtomicBool, - release: Notify, - cancelled: AtomicBool, - } - - #[derive(Clone)] - struct BlockingInferenceAdapter { - state: Arc, - } - - impl BlockingInferenceAdapter { - fn new(scripts: impl IntoIterator, blocked_turn: usize) -> Self { - Self { - state: Arc::new(BlockingInferenceState { - scripts: Mutex::new(scripts.into_iter().collect()), - next_turn: AtomicUsize::new(0), - blocked_turn, - entered: AtomicBool::new(false), - release: Notify::new(), - cancelled: AtomicBool::new(false), - }), +async fn drive_prompt( + driver: &mut agentkit_loop::LoopDriver, + integration: &AcpIntegration, + session_id: &wire::SessionId, + cancellation: &CancellationHandle, + generation: u64, + #[cfg(feature = "unstable-inject")] session: &AcpSessionHandle, +) -> wire::StopReason +where + S: ModelSession + Send + 'static, +{ + loop { + 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; } - - async fn wait_until_blocked(&self) { - tokio::time::timeout(Duration::from_secs(2), async { - while !self.state.entered.load(Ordering::Acquire) { - tokio::task::yield_now().await; + match step { + LoopStep::Finished(result) => { + if result.finish_reason == FinishReason::ToolCall { + continue; } - }) - .await - .expect("model inference did not block"); - } - - fn release(&self) { - self.state.release.notify_one(); - } - - fn was_cancelled(&self) -> bool { - self.state.cancelled.load(Ordering::Acquire) + #[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(); + } + _ => {} + } + } + 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(); + } + #[cfg(not(feature = "unstable-inject"))] + return error_stop_reason(); + } } } +} - struct BlockingInferenceSession { - state: Arc, - } - - struct BlockingInferenceTurn { - state: Arc, - events: VecDeque, - blocked: bool, - } +fn prompt_to_items(request: &wire::PromptRequest) -> Result, AcpRuntimeError> { + content_blocks_to_items(&request.prompt) +} - #[async_trait] - impl ModelAdapter for BlockingInferenceAdapter { - type Session = BlockingInferenceSession; +/// 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(); - async fn start_session(&self, _config: SessionConfig) -> Result { - Ok(BlockingInferenceSession { - state: Arc::clone(&self.state), - }) + for block in content { + match block { + wire::ContentBlock::Text(text) => { + user_parts.push(Part::Text(TextPart::new(text.text.clone()))); + } + wire::ContentBlock::Image(image) => { + let mime_type = image.mime_type.as_ref(); + user_parts.push(Part::media( + Modality::Image, + mime_type, + crate::data_url_ref(mime_type, &image.data), + )); + } + wire::ContentBlock::Audio(audio) => { + let mime_type = audio.mime_type.as_ref(); + user_parts.push(Part::media( + Modality::Audio, + mime_type, + crate::data_url_ref(mime_type, &audio.data), + )); + } + wire::ContentBlock::ResourceLink(link) => { + context_items.push(resource_link_item(link)); + } + wire::ContentBlock::Resource(resource) => { + context_items.push(resource_item(resource)?); + } + wire::ContentBlock::Other(other) => { + return Err(AcpRuntimeError::UnsupportedContent(format!( + "unknown ACP v2 content block: {}", + other.type_ + ))); + } + _ => { + return Err(AcpRuntimeError::UnsupportedContent( + "unknown ACP v2 content block".into(), + )); + } } } - #[async_trait] - impl ModelSession for BlockingInferenceSession { - type Turn = BlockingInferenceTurn; + let mut items = context_items; + if !user_parts.is_empty() { + items.push(Item::new(ItemKind::User, user_parts)); + } else if !items.is_empty() { + items.push(Item::text(ItemKind::User, "Use the provided context.")); + } + Ok(items) +} - async fn begin_turn( - &mut self, - _request: TurnRequest, - _cancellation: Option, - ) -> Result { - let turn = self.state.next_turn.fetch_add(1, Ordering::AcqRel); - let script = self - .state - .scripts - .lock() - .unwrap() - .pop_front() - .ok_or_else(|| LoopError::InvalidState("missing blocking turn script".into()))?; - Ok(BlockingInferenceTurn { - state: Arc::clone(&self.state), - events: script.events.into(), - blocked: turn == self.state.blocked_turn, - }) - } +fn resource_link_item(link: &wire::ResourceLink) -> Item { + let mut metadata = MetadataMap::new(); + metadata.insert("acp.resource.uri".into(), json!(link.uri)); + metadata.insert("acp.resource.name".into(), json!(link.name)); + if let Some(description) = &link.description { + metadata.insert("acp.resource.description".into(), json!(description)); } + if let Some(mime_type) = &link.mime_type { + metadata.insert("acp.resource.mime_type".into(), json!(mime_type)); + } + Item::new( + ItemKind::Context, + vec![Part::file(DataRef::uri(link.uri.clone()))], + ) + .with_metadata(metadata) +} - #[async_trait] - impl ModelTurn for BlockingInferenceTurn { - async fn next_event( - &mut self, - cancellation: Option, - ) -> Result, LoopError> { - if self.blocked { - self.blocked = false; - self.state.entered.store(true, Ordering::Release); - if let Some(cancellation) = cancellation { - tokio::select! { - _ = self.state.release.notified() => {} - _ = cancellation.cancelled() => { - self.state.cancelled.store(true, Ordering::Release); - return Err(LoopError::Cancelled); - } - } - } else { - self.state.release.notified().await; - } +fn resource_item(resource: &wire::EmbeddedResource) -> Result { + match &resource.resource { + wire::EmbeddedResourceResource::TextResourceContents(text) => { + let mut metadata = MetadataMap::new(); + metadata.insert("acp.resource.uri".into(), json!(text.uri)); + if let Some(mime_type) = &text.mime_type { + metadata.insert("acp.resource.mime_type".into(), json!(mime_type)); } - Ok(self.events.pop_front()) + Ok(Item::text(ItemKind::Context, text.text.clone()).with_metadata(metadata)) + } + wire::EmbeddedResourceResource::BlobResourceContents(blob) => { + let mime_type = blob + .mime_type + .as_ref() + .map(|mime_type| mime_type.as_ref()) + .unwrap_or("application/octet-stream"); + let mut metadata = MetadataMap::new(); + metadata.insert("acp.resource.uri".into(), json!(blob.uri)); + metadata.insert("acp.resource.mime_type".into(), json!(mime_type)); + Ok(Item::new( + ItemKind::Context, + vec![Part::media( + Modality::Binary, + mime_type, + crate::data_url_ref(mime_type, &blob.blob), + )], + ) + .with_metadata(metadata)) } + _ => Err(AcpRuntimeError::UnsupportedContent( + "unknown ACP v2 embedded resource".into(), + )), } +} - fn tool_turn(call_id: &str) -> TurnScript { - let call = ToolCallPart::new(ToolCallId::new(call_id), "blocking_tool", json!({})); - TurnScript::new([ - 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(), - }), - ]) +fn event_to_update( + event: &AgentEvent, + message_ids: Option<&CurrentMessageIds>, + part_kinds: &mut HashMap, +) -> Option { + match event { + AgentEvent::ContentDelta(delta) => delta_to_update(delta, message_ids, part_kinds), + AgentEvent::ToolCallRequested(call) => { + Some(wire::SessionUpdate::ToolCallUpdate(tool_call_update(call))) + } + AgentEvent::ToolExecutionStarted(call) => Some(wire::SessionUpdate::ToolCallUpdate( + tool_status_update(&call.id, wire::ToolCallStatus::InProgress), + )), + AgentEvent::ToolExecutionProgress(result) => Some(wire::SessionUpdate::ToolCallUpdate( + tool_result_update(result, wire::ToolCallStatus::InProgress), + )), + AgentEvent::ToolResultReceived(result) => { + Some(wire::SessionUpdate::ToolCallUpdate(tool_result_update( + result, + if result.is_error { + wire::ToolCallStatus::Failed + } else { + wire::ToolCallStatus::Completed + }, + ))) + } + AgentEvent::Warning { message } => { + tracing::warn!(%message, "agentkit warning while routing ACP v2 event"); + None + } + AgentEvent::RunFailed { message } => { + tracing::debug!(%message, "agentkit run failed while routing ACP v2 event"); + None + } + _ => None, } +} - fn streamed_text_and_tool(text: &str, call_id: &str) -> TurnScript { - let call = ToolCallPart::new(ToolCallId::new(call_id), "missing_tool", json!({})); - TurnScript::new([ - ModelTurnEvent::Delta(Delta::BeginPart { - part_id: PartId::new("part-1"), - kind: PartKind::Text, - }), - ModelTurnEvent::Delta(Delta::AppendText { - part_id: PartId::new("part-1"), - chunk: text.to_string(), - }), - 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::text(text), Part::ToolCall(call)], - )], - usage: None, - metadata: MetadataMap::new(), - }), - ]) +fn delta_to_update( + delta: &Delta, + message_ids: Option<&CurrentMessageIds>, + part_kinds: &mut HashMap, +) -> Option { + match delta { + Delta::BeginPart { part_id, kind } => { + part_kinds.insert(part_id.clone(), *kind); + None + } + Delta::AppendText { part_id, chunk } => { + let message_ids = message_ids?; + let content = wire::ContentBlock::Text(wire::TextContent::new(chunk.clone())); + match part_kinds.get(part_id) { + Some(PartKind::Reasoning) => Some(wire::SessionUpdate::AgentThoughtChunk( + wire::ContentChunk::new(content, message_ids.thought.clone()), + )), + Some(PartKind::Text) | None => Some(wire::SessionUpdate::AgentMessageChunk( + wire::ContentChunk::new(content, message_ids.agent.clone()), + )), + Some(_) => None, + } + } + Delta::CommitPart { .. } + | Delta::AppendBytes { .. } + | Delta::ReplaceStructured { .. } + | Delta::SetMetadata { .. } => None, } +} - fn streamed_text(text: &str) -> TurnScript { - TurnScript::new([ - ModelTurnEvent::Delta(Delta::BeginPart { - part_id: PartId::new("part-1"), - kind: PartKind::Text, - }), - ModelTurnEvent::Delta(Delta::AppendText { - part_id: PartId::new("part-1"), - chunk: text.to_string(), - }), - 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(), - }), - ]) - } +fn tool_call_update(call: &ToolCallPart) -> wire::ToolCallUpdate { + wire::ToolCallUpdate::new(call.id.to_string()) + .title(call.name.clone()) + .status(wire::ToolCallStatus::Pending) + .raw_input(call.input.clone()) +} - fn streamed_thought_and_text(thought: &str, text: &str) -> TurnScript { - TurnScript::new([ - ModelTurnEvent::Delta(Delta::BeginPart { - part_id: PartId::new("thought-part"), - kind: PartKind::Reasoning, - }), - ModelTurnEvent::Delta(Delta::AppendText { - part_id: PartId::new("thought-part"), - chunk: thought.to_string(), - }), - ModelTurnEvent::Delta(Delta::BeginPart { - part_id: PartId::new("text-part"), - kind: PartKind::Text, - }), - ModelTurnEvent::Delta(Delta::AppendText { - part_id: PartId::new("text-part"), - chunk: text.to_string(), - }), - 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(), - }), - ]) - } +fn tool_status_update( + call_id: &agentkit_core::ToolCallId, + status: wire::ToolCallStatus, +) -> wire::ToolCallUpdate { + wire::ToolCallUpdate::new(call_id.to_string()).status(status) +} - async fn wait_for_idle( - updates: &Arc>>, - session_id: &wire::SessionId, - ) { - wait_for_idle_count(updates, session_id, 1).await; +fn tool_result_update( + result: &ToolResultPart, + status: wire::ToolCallStatus, +) -> wire::ToolCallUpdate { + wire::ToolCallUpdate::new(result.call_id.to_string()) + .status(status) + .raw_output(crate::tool_output_raw(&result.output)) + .content(tool_output_content(&result.output)) +} + +fn tool_output_content(output: &ToolOutput) -> Option> { + let content = match output { + ToolOutput::Text(text) => vec![text_to_tool_content(text.clone())], + ToolOutput::Structured(value) => vec![text_to_tool_content(value.to_string())], + ToolOutput::Parts(parts) => parts.iter().filter_map(part_to_tool_content).collect(), + ToolOutput::Files(files) => files.iter().map(file_to_tool_content).collect(), + }; + (!content.is_empty()).then_some(content) +} + +fn part_to_tool_content(part: &Part) -> Option { + match part { + Part::Text(text) => Some(text_to_tool_content(text.text.clone())), + Part::Structured(value) => Some(structured_to_tool_content(value)), + Part::Media(media) => Some(media_to_tool_content(media)), + Part::File(file) => Some(file_to_tool_content(file)), + Part::Reasoning(reasoning) => reasoning + .summary + .as_ref() + .map(|summary| text_to_tool_content(summary.clone())), + Part::Custom(custom) => Some(text_to_tool_content( + custom + .value + .as_ref() + .map(ToString::to_string) + .or_else(|| custom.data.as_ref().map(crate::data_ref_payload)) + .unwrap_or_else(|| custom.kind.clone()), + )), + Part::ToolCall(_) | Part::ToolResult(_) => None, } +} - async fn wait_for_idle_count( - updates: &Arc>>, - session_id: &wire::SessionId, - count: usize, - ) { - tokio::time::timeout(Duration::from_secs(2), async { - loop { - let idle_count = updates - .lock() - .unwrap() - .iter() - .filter(|(id, update)| { - id == session_id - && matches!( - update, - wire::SessionUpdate::StateUpdate(wire::StateUpdate::Idle(_)) - ) - }) - .count(); - if idle_count >= count { - return; - } - tokio::time::sleep(Duration::from_millis(5)).await; +fn text_to_tool_content(text: String) -> wire::ToolCallContent { + wire::ToolCallContent::Content(Box::new(wire::Content::new(wire::ContentBlock::Text( + wire::TextContent::new(text), + )))) +} + +fn structured_to_tool_content(part: &StructuredPart) -> wire::ToolCallContent { + text_to_tool_content(part.value.to_string()) +} + +fn media_to_tool_content(media: &MediaPart) -> wire::ToolCallContent { + match media.modality { + Modality::Image + if matches!(media.data, DataRef::InlineText(_) | DataRef::InlineBytes(_)) => + { + let mut image = wire::ImageContent::new( + crate::data_ref_base64_payload(&media.data), + media.mime_type.clone(), + ); + if let Some(uri) = crate::data_ref_uri(&media.data) { + image = image.uri(uri); } - }) - .await - .expect("idle update timed out"); + wire::ToolCallContent::Content(Box::new(wire::Content::new(wire::ContentBlock::Image( + image, + )))) + } + Modality::Audio + if matches!(media.data, DataRef::InlineText(_) | DataRef::InlineBytes(_)) => + { + wire::ToolCallContent::Content(Box::new(wire::Content::new(wire::ContentBlock::Audio( + wire::AudioContent::new( + crate::data_ref_base64_payload(&media.data), + media.mime_type.clone(), + ), + )))) + } + Modality::Image | Modality::Audio | Modality::Video | Modality::Binary => { + data_ref_to_resource_content(None, Some(&media.mime_type), &media.data) + } } +} - #[tokio::test] - async fn runtime_rotates_message_ids_per_model_turn_and_session() { - let adapter = MockAdapter::new(); - adapter.enqueue(streamed_text_and_tool("before tool", "call-1")); - adapter.enqueue(streamed_thought_and_text("first thought", "first output")); - adapter.enqueue(streamed_thought_and_text("second thought", "second output")); - let updates = Arc::new(Mutex::new(Vec::new())); - let (client_transport, agent_transport) = Channel::duplex(); +fn file_to_tool_content(file: &FilePart) -> wire::ToolCallContent { + data_ref_to_resource_content(file.name.as_deref(), file.mime_type.as_deref(), &file.data) +} - let server = tokio::spawn({ - let factory = TestFactory { - adapter: adapter.clone(), - }; - async move { - AcpHeadlessRuntime::::builder() - .name("agentkit-v2-test") - .agent_factory(factory) - .serve(agent_transport) - .await +fn data_ref_to_resource_content( + name: Option<&str>, + mime_type: Option<&str>, + data: &DataRef, +) -> wire::ToolCallContent { + let content = match data { + DataRef::Uri(uri) => { + let mut link = wire::ResourceLink::new(name.unwrap_or(uri), uri.clone()); + if let Some(mime_type) = mime_type { + link = link.mime_type(mime_type.to_string()); } - }); + wire::ContentBlock::ResourceLink(link) + } + DataRef::Handle(handle) => { + let uri = format!("artifact://{handle}"); + let link_name = name.map(str::to_owned).unwrap_or_else(|| uri.clone()); + let mut link = wire::ResourceLink::new(link_name, uri); + if let Some(mime_type) = mime_type { + link = link.mime_type(mime_type.to_string()); + } + wire::ContentBlock::ResourceLink(link) + } + DataRef::InlineText(text) if mime_type.is_none_or(|mime| mime.starts_with("text/")) => { + let mut resource = wire::TextResourceContents::new( + text.clone(), + crate::inline_resource_uri(name.unwrap_or("tool-output")), + ); + if let Some(mime_type) = mime_type { + resource = resource.mime_type(mime_type.to_string()); + } + wire::ContentBlock::Resource(wire::EmbeddedResource::new( + wire::EmbeddedResourceResource::TextResourceContents(resource), + )) + } + _ => { + let mut resource = wire::BlobResourceContents::new( + crate::data_ref_base64_payload(data), + crate::inline_resource_uri(name.unwrap_or("tool-output")), + ); + if let Some(mime_type) = mime_type { + resource = resource.mime_type(mime_type.to_string()); + } + wire::ContentBlock::Resource(wire::EmbeddedResource::new( + wire::EmbeddedResourceResource::BlobResourceContents(resource), + )) + } + }; + wire::ToolCallContent::Content(Box::new(wire::Content::new(content))) +} - 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| { - let initialize = cx - .send_request(wire::InitializeRequest::new( - wire::ProtocolVersion::V2, - wire::Implementation::new("test-client", "1"), - )) - .block_task() - .await?; - assert_eq!(initialize.protocol_version, wire::ProtocolVersion::V2); - assert_eq!(initialize.info.name, "agentkit-v2-test"); +fn error_stop_reason() -> wire::StopReason { + wire::StopReason::Other("_error".into()) +} - 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?; - assert_ne!(first.session_id, second.session_id); +fn finish_reason_to_stop_reason(reason: &FinishReason) -> wire::StopReason { + match reason { + FinishReason::Completed | FinishReason::ToolCall | FinishReason::Other(_) => { + wire::StopReason::EndTurn + } + FinishReason::MaxTokens => wire::StopReason::MaxTokens, + FinishReason::Cancelled => wire::StopReason::Cancelled, + FinishReason::Blocked | FinishReason::Error => error_stop_reason(), + } +} - cx.send_request(wire::PromptRequest::new( - first.session_id.clone(), - vec![wire::ContentBlock::Text(wire::TextContent::new("first"))], - )) - .block_task() - .await?; - wait_for_idle(&updates, &first.session_id).await; +/// 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]) + .pending(wire::SessionInjectPendingCapabilities::new().replace(true)) +} - cx.send_request(wire::PromptRequest::new( - second.session_id.clone(), - vec![wire::ContentBlock::Text(wire::TextContent::new("second"))], - )) - .block_task() - .await?; - wait_for_idle(&updates, &second.session_id).await; +/// 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) +} - 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(test)] +mod tests { + use super::*; + use std::collections::VecDeque; + use std::sync::atomic::AtomicUsize; + use std::time::Duration; - tokio::time::timeout(Duration::from_secs(5), client) - .await - .expect("client timed out") - .expect("client run"); - server.abort(); - let _ = server.await; + use agent_client_protocol::Channel; + #[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::{ + Agent, ModelSession, ModelTurn, ModelTurnEvent, ModelTurnResult, SessionConfig, + 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; - let updates = updates.lock().unwrap(); - for session_id in ["session-1", "session-2"] { - let session_updates = updates - .iter() - .filter(|(id, _)| id.to_string() == session_id) - .map(|(_, update)| update) - .collect::>(); - assert!(matches!( - session_updates[0], - wire::SessionUpdate::UserMessage(_) - )); - assert!(matches!( - session_updates[1], - wire::SessionUpdate::StateUpdate(wire::StateUpdate::Running(_)) - )); - assert_eq!( - session_updates - .iter() + #[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, + } + + #[async_trait] + impl AcpAgentFactory for TestFactory + where + A: ModelAdapter + Clone + Send + Sync + 'static, + A::Session: Send + 'static, + { + async fn start( + &self, + ctx: AcpAgentFactoryContext, + ) -> Result, AcpRuntimeError> { + 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())) + } + } + + #[derive(Clone)] + struct CancellationAwareTool { + spec: ToolSpec, + entered: Arc, + cleaned: Arc, + } + + impl CancellationAwareTool { + fn new() -> Self { + Self { + spec: ToolSpec::new( + "blocking_tool", + "waits for turn cancellation", + json!({ "type": "object" }), + ), + entered: Arc::new(AtomicUsize::new(0)), + cleaned: Arc::new(AtomicUsize::new(0)), + } + } + + async fn wait_for_entered(&self, count: usize) { + tokio::time::timeout(Duration::from_secs(2), async { + while self.entered.load(Ordering::Acquire) < count { + tokio::task::yield_now().await; + } + }) + .await + .expect("tool did not start"); + } + } + + #[async_trait] + impl Tool for CancellationAwareTool { + fn spec(&self) -> &ToolSpec { + &self.spec + } + + async fn invoke( + &self, + request: ToolRequest, + ctx: &mut ToolContext<'_>, + ) -> Result { + self.entered.fetch_add(1, Ordering::AcqRel); + ctx.cancellation + .as_ref() + .expect("turn cancellation installed") + .cancelled() + .await; + self.cleaned.fetch_add(1, Ordering::AcqRel); + Ok(ToolResult::new(ToolResultPart::error( + request.call_id, + ToolOutput::text("cancelled"), + ))) + } + } + + #[derive(Clone)] + struct ToolTestFactory { + adapter: MockAdapter, + tool: CancellationAwareTool, + } + + #[derive(Clone)] + struct BackgroundToolTestFactory { + adapter: A, + tool: BlockingTool, + } + + #[async_trait] + impl AcpAgentFactory for ToolTestFactory { + 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())) + } + } + + #[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 + A: ModelAdapter + Clone + Send + Sync + 'static, + A::Session: Send + 'static, + { + async fn start( + &self, + ctx: AcpAgentFactoryContext, + ) -> Result, AcpRuntimeError> { + Agent::builder() + .model(self.adapter.clone()) + .add_tool_source(ToolRegistry::new().with(self.tool.clone())) + .task_manager( + AsyncTaskManager::new().routing(|_: &ToolRequest| RoutingDecision::Background), + ) + .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())) + } + } + + struct BlockingInferenceState { + scripts: Mutex>, + next_turn: AtomicUsize, + blocked_turn: usize, + entered: AtomicBool, + release: Notify, + cancelled: AtomicBool, + } + + #[derive(Clone)] + struct BlockingInferenceAdapter { + state: Arc, + } + + impl BlockingInferenceAdapter { + fn new(scripts: impl IntoIterator, blocked_turn: usize) -> Self { + Self { + state: Arc::new(BlockingInferenceState { + scripts: Mutex::new(scripts.into_iter().collect()), + next_turn: AtomicUsize::new(0), + blocked_turn, + entered: AtomicBool::new(false), + release: Notify::new(), + cancelled: AtomicBool::new(false), + }), + } + } + + async fn wait_until_blocked(&self) { + tokio::time::timeout(Duration::from_secs(2), async { + while !self.state.entered.load(Ordering::Acquire) { + tokio::task::yield_now().await; + } + }) + .await + .expect("model inference did not block"); + } + + fn release(&self) { + self.state.release.notify_one(); + } + + fn was_cancelled(&self) -> bool { + self.state.cancelled.load(Ordering::Acquire) + } + } + + struct BlockingInferenceSession { + state: Arc, + } + + struct BlockingInferenceTurn { + state: Arc, + events: VecDeque, + blocked: bool, + } + + #[async_trait] + impl ModelAdapter for BlockingInferenceAdapter { + type Session = BlockingInferenceSession; + + async fn start_session(&self, _config: SessionConfig) -> Result { + Ok(BlockingInferenceSession { + state: Arc::clone(&self.state), + }) + } + } + + #[async_trait] + impl ModelSession for BlockingInferenceSession { + type Turn = BlockingInferenceTurn; + + async fn begin_turn( + &mut self, + _request: TurnRequest, + _cancellation: Option, + ) -> Result { + let turn = self.state.next_turn.fetch_add(1, Ordering::AcqRel); + let script = self + .state + .scripts + .lock() + .unwrap() + .pop_front() + .ok_or_else(|| LoopError::InvalidState("missing blocking turn script".into()))?; + Ok(BlockingInferenceTurn { + state: Arc::clone(&self.state), + events: script.events.into(), + blocked: turn == self.state.blocked_turn, + }) + } + } + + #[async_trait] + impl ModelTurn for BlockingInferenceTurn { + async fn next_event( + &mut self, + cancellation: Option, + ) -> Result, LoopError> { + if self.blocked { + self.blocked = false; + self.state.entered.store(true, Ordering::Release); + if let Some(cancellation) = cancellation { + tokio::select! { + _ = self.state.release.notified() => {} + _ = cancellation.cancelled() => { + self.state.cancelled.store(true, Ordering::Release); + return Err(LoopError::Cancelled); + } + } + } else { + self.state.release.notified().await; + } + } + Ok(self.events.pop_front()) + } + } + + fn tool_turn(call_id: &str) -> TurnScript { + let call = ToolCallPart::new(ToolCallId::new(call_id), "blocking_tool", json!({})); + TurnScript::new([ + 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(), + }), + ]) + } + + fn streamed_text_and_tool(text: &str, call_id: &str) -> TurnScript { + let call = ToolCallPart::new(ToolCallId::new(call_id), "missing_tool", json!({})); + TurnScript::new([ + ModelTurnEvent::Delta(Delta::BeginPart { + part_id: PartId::new("part-1"), + kind: PartKind::Text, + }), + ModelTurnEvent::Delta(Delta::AppendText { + part_id: PartId::new("part-1"), + chunk: text.to_string(), + }), + 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::text(text), Part::ToolCall(call)], + )], + usage: None, + metadata: MetadataMap::new(), + }), + ]) + } + + fn streamed_text(text: &str) -> TurnScript { + TurnScript::new([ + ModelTurnEvent::Delta(Delta::BeginPart { + part_id: PartId::new("part-1"), + kind: PartKind::Text, + }), + ModelTurnEvent::Delta(Delta::AppendText { + part_id: PartId::new("part-1"), + chunk: text.to_string(), + }), + 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(), + }), + ]) + } + + fn streamed_thought_and_text(thought: &str, text: &str) -> TurnScript { + TurnScript::new([ + ModelTurnEvent::Delta(Delta::BeginPart { + part_id: PartId::new("thought-part"), + kind: PartKind::Reasoning, + }), + ModelTurnEvent::Delta(Delta::AppendText { + part_id: PartId::new("thought-part"), + chunk: thought.to_string(), + }), + ModelTurnEvent::Delta(Delta::BeginPart { + part_id: PartId::new("text-part"), + kind: PartKind::Text, + }), + ModelTurnEvent::Delta(Delta::AppendText { + part_id: PartId::new("text-part"), + chunk: text.to_string(), + }), + 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(), + }), + ]) + } + + async fn wait_for_idle( + updates: &Arc>>, + session_id: &wire::SessionId, + ) { + wait_for_idle_count(updates, session_id, 1).await; + } + + async fn wait_for_idle_count( + updates: &Arc>>, + session_id: &wire::SessionId, + count: usize, + ) { + tokio::time::timeout(Duration::from_secs(2), async { + loop { + let idle_count = updates + .lock() + .unwrap() + .iter() + .filter(|(id, update)| { + id == session_id + && matches!( + update, + wire::SessionUpdate::StateUpdate(wire::StateUpdate::Idle(_)) + ) + }) + .count(); + if idle_count >= count { + return; + } + tokio::time::sleep(Duration::from_millis(5)).await; + } + }) + .await + .expect("idle update timed out"); + } + + #[tokio::test] + async fn runtime_rotates_message_ids_per_model_turn_and_session() { + let adapter = MockAdapter::new(); + adapter.enqueue(streamed_text_and_tool("before tool", "call-1")); + adapter.enqueue(streamed_thought_and_text("first thought", "first output")); + adapter.enqueue(streamed_thought_and_text("second thought", "second output")); + let updates = Arc::new(Mutex::new(Vec::new())); + let (client_transport, agent_transport) = Channel::duplex(); + + let server = tokio::spawn({ + let factory = TestFactory { + adapter: adapter.clone(), + }; + async move { + AcpHeadlessRuntime::::builder() + .name("agentkit-v2-test") + .agent_factory(factory) + .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| { + let initialize = cx + .send_request(wire::InitializeRequest::new( + wire::ProtocolVersion::V2, + wire::Implementation::new("test-client", "1"), + )) + .block_task() + .await?; + assert_eq!(initialize.protocol_version, wire::ProtocolVersion::V2); + assert_eq!(initialize.info.name, "agentkit-v2-test"); + + 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?; + assert_ne!(first.session_id, second.session_id); + + cx.send_request(wire::PromptRequest::new( + first.session_id.clone(), + vec![wire::ContentBlock::Text(wire::TextContent::new("first"))], + )) + .block_task() + .await?; + wait_for_idle(&updates, &first.session_id).await; + + cx.send_request(wire::PromptRequest::new( + second.session_id.clone(), + vec![wire::ContentBlock::Text(wire::TextContent::new("second"))], + )) + .block_task() + .await?; + wait_for_idle(&updates, &second.session_id).await; + + 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; + + let updates = updates.lock().unwrap(); + for session_id in ["session-1", "session-2"] { + let session_updates = updates + .iter() + .filter(|(id, _)| id.to_string() == session_id) + .map(|(_, update)| update) + .collect::>(); + assert!(matches!( + session_updates[0], + wire::SessionUpdate::UserMessage(_) + )); + assert!(matches!( + session_updates[1], + wire::SessionUpdate::StateUpdate(wire::StateUpdate::Running(_)) + )); + assert_eq!( + session_updates + .iter() .filter(|update| matches!( update, wire::SessionUpdate::StateUpdate(wire::StateUpdate::Running(_)) )) - .count(), - 1, - "Running must come only from TurnStarted" - ); + .count(), + 1, + "Running must come only from TurnStarted" + ); + assert_eq!( + session_updates + .iter() + .filter(|update| matches!( + update, + wire::SessionUpdate::StateUpdate(wire::StateUpdate::Idle(_)) + )) + .count(), + 1, + "Idle must come only from TurnFinished" + ); + let user_id = match session_updates[0] { + wire::SessionUpdate::UserMessage(message) => message.message_id.to_string(), + _ => unreachable!(), + }; + let agent_ids = session_updates + .iter() + .filter_map(|update| match update { + wire::SessionUpdate::AgentMessageChunk(chunk) => { + Some(chunk.message_id.to_string()) + } + _ => None, + }) + .collect::>(); + let thought_ids = session_updates + .iter() + .filter_map(|update| match update { + wire::SessionUpdate::AgentThoughtChunk(chunk) => { + Some(chunk.message_id.to_string()) + } + _ => None, + }) + .collect::>(); + assert_eq!(user_id, format!("{session_id}-user-1")); + if session_id == "session-1" { + assert_eq!( + agent_ids, + [ + "session-1-agent-2".to_string(), + "session-1-agent-3".to_string() + ] + ); + assert_eq!(thought_ids, ["session-1-thought-3"]); + assert_ne!(agent_ids[0], agent_ids[1]); + } else { + assert_eq!(agent_ids, ["session-2-agent-2"]); + assert_eq!(thought_ids, ["session-2-thought-2"]); + } + assert_ne!(agent_ids.last(), thought_ids.last()); + assert!(matches!( + session_updates.last(), + Some(wire::SessionUpdate::StateUpdate(wire::StateUpdate::Idle(idle))) + if idle.stop_reason == Some(wire::StopReason::EndTurn) + )); + } + } + + #[tokio::test] + async fn idle_session_drives_background_completion_without_another_prompt() { + let adapter = MockAdapter::new(); + adapter.enqueue(tool_turn("background-call")); + adapter.enqueue(streamed_text("background observed")); + let tool = BlockingTool::text("blocking_tool", "background done"); + let updates = Arc::new(Mutex::new(Vec::new())); + let (client_transport, agent_transport) = Channel::duplex(); + + let server = tokio::spawn({ + let factory = BackgroundToolTestFactory { + adapter, + tool: tool.clone(), + }; + async move { + AcpHeadlessRuntime::::builder() + .agent_factory(factory) + .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 session = cx + .send_request(wire::NewSessionRequest::new(cwd)) + .block_task() + .await?; + cx.send_request(wire::PromptRequest::new( + session.session_id.clone(), + vec![wire::ContentBlock::Text(wire::TextContent::new( + "start background work", + ))], + )) + .block_task() + .await?; + wait_for_idle(&updates, &session.session_id).await; + tool.wait_until_entered().await; + tool.release(); + + tokio::time::timeout(Duration::from_secs(2), async { + loop { + let updates = updates.lock().unwrap(); + let completed = updates.iter().any(|(id, update)| { + id == &session.session_id + && matches!( + update, + wire::SessionUpdate::ToolCallUpdate(update) + if update.tool_call_id.to_string() == "background-call" + && matches!( + update.status, + agent_client_protocol::schema::MaybeUndefined::Value( + wire::ToolCallStatus::Completed + ) + ) + ) + }); + let message_delivered = updates.iter().any(|(id, update)| { + id == &session.session_id + && matches!( + update, + wire::SessionUpdate::AgentMessageChunk(chunk) + if matches!( + &chunk.content, + wire::ContentBlock::Text(text) + if text.text == "background observed" + ) + ) + }); + if completed && message_delivered { + break; + } + drop(updates); + tokio::time::sleep(Duration::from_millis(5)).await; + } + }) + .await + .expect("background completion was not delivered while idle"); + cx.send_request(wire::CloseSessionRequest::new(session.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; + + let updates = updates.lock().unwrap(); + assert_eq!( + updates + .iter() + .filter(|(_, update)| matches!(update, wire::SessionUpdate::UserMessage(_))) + .count(), + 1, + "the background continuation must not require or replay a prompt" + ); + assert!(updates.iter().any(|(_, update)| matches!( + update, + wire::SessionUpdate::AgentMessageChunk(chunk) + if matches!( + &chunk.content, + wire::ContentBlock::Text(text) if text.text == "background observed" + ) + ))); + assert_eq!( + updates + .iter() + .filter(|(_, update)| matches!( + update, + wire::SessionUpdate::StateUpdate(wire::StateUpdate::Running(_)) + )) + .count(), + 1, + "idle background completion must not re-enter Running" + ); + assert_eq!( + updates + .iter() + .filter(|(_, update)| matches!( + update, + wire::SessionUpdate::StateUpdate(wire::StateUpdate::Idle(_)) + )) + .count(), + 1, + "idle background completion must not emit another Idle" + ); + } + + #[tokio::test] + async fn cancel_without_prompt_owner_interrupts_blocked_autonomous_inference() { + let adapter = BlockingInferenceAdapter::new( + [ + tool_turn("background-call"), + streamed_text("autonomous complete"), + ], + 1, + ); + let tool = BlockingTool::text("blocking_tool", "background done"); + let updates = Arc::new(Mutex::new(Vec::new())); + let (client_transport, agent_transport) = Channel::duplex(); + + let server = tokio::spawn({ + let factory = BackgroundToolTestFactory { + adapter: adapter.clone(), + tool: tool.clone(), + }; + async move { + AcpHeadlessRuntime::::builder() + .agent_factory(factory) + .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 tool = tool.clone(); + 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 session = cx + .send_request(wire::NewSessionRequest::new(cwd)) + .block_task() + .await?; + cx.send_request(wire::PromptRequest::new( + session.session_id.clone(), + vec![wire::ContentBlock::Text(wire::TextContent::new("start"))], + )) + .block_task() + .await?; + wait_for_idle(&updates, &session.session_id).await; + tool.wait_until_entered().await; + tool.release(); + adapter.wait_until_blocked().await; + + cx.send_notification(wire::CancelSessionNotification::new( + session.session_id.clone(), + ))?; + cx.send_request(wire::ListSessionsRequest::new()) + .block_task() + .await?; + tokio::time::timeout(Duration::from_secs(2), async { + while !adapter.was_cancelled() { + tokio::task::yield_now().await; + } + }) + .await + .expect("session cancel did not interrupt autonomous inference"); + + cx.send_request(wire::CloseSessionRequest::new(session.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 queued_prompt_is_acknowledged_without_cancelling_blocked_autonomous_inference() { + let adapter = BlockingInferenceAdapter::new( + [ + tool_turn("background-call"), + streamed_text("autonomous complete"), + streamed_text("queued prompt complete"), + ], + 1, + ); + let tool = BlockingTool::text("blocking_tool", "background done"); + let updates = Arc::new(Mutex::new(Vec::new())); + let (client_transport, agent_transport) = Channel::duplex(); + + let server = tokio::spawn({ + let factory = BackgroundToolTestFactory { + adapter: adapter.clone(), + tool: tool.clone(), + }; + async move { + AcpHeadlessRuntime::::builder() + .agent_factory(factory) + .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 tool = tool.clone(); + 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 session = cx + .send_request(wire::NewSessionRequest::new(cwd)) + .block_task() + .await?; + cx.send_request(wire::PromptRequest::new( + session.session_id.clone(), + vec![wire::ContentBlock::Text(wire::TextContent::new("start"))], + )) + .block_task() + .await?; + wait_for_idle(&updates, &session.session_id).await; + tool.wait_until_entered().await; + tool.release(); + adapter.wait_until_blocked().await; + + tokio::time::timeout( + Duration::from_millis(500), + cx.send_request(wire::PromptRequest::new( + session.session_id.clone(), + vec![wire::ContentBlock::Text(wire::TextContent::new( + "queued while idle", + ))], + )) + .block_task(), + ) + .await + .expect("prompt response waited for autonomous inference")?; + assert_eq!( + updates + .lock() + .unwrap() + .iter() + .filter(|(_, update)| matches!( + update, + wire::SessionUpdate::UserMessage(_) + )) + .count(), + 1, + "queued prompt update crossed the response gate" + ); + + cx.send_notification(wire::CancelSessionNotification::new( + session.session_id.clone(), + ))?; + cx.send_request(wire::ListSessionsRequest::new()) + .block_task() + .await?; + assert!( + tokio::time::timeout(Duration::from_millis(50), async { + while !adapter.was_cancelled() { + tokio::task::yield_now().await; + } + }) + .await + .is_err(), + "queued prompt cancellation interrupted autonomous inference" + ); + + adapter.release(); + tokio::time::timeout(Duration::from_secs(2), async { + loop { + if updates.lock().unwrap().iter().any(|(_, update)| { + matches!( + update, + wire::SessionUpdate::AgentMessageChunk(chunk) + if matches!( + &chunk.content, + wire::ContentBlock::Text(text) + if text.text == "autonomous complete" + ) + ) + }) { + break; + } + tokio::time::sleep(Duration::from_millis(5)).await; + } + }) + .await + .expect("autonomous inference did not complete"); + assert!(!adapter.was_cancelled()); + assert_eq!( + updates + .lock() + .unwrap() + .iter() + .filter(|(_, update)| matches!( + update, + wire::SessionUpdate::UserMessage(_) + )) + .count(), + 1, + "cancelled queued prompt was driven" + ); + + cx.send_request(wire::CloseSessionRequest::new(session.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 v2_tool_updates_include_visible_text_structured_parts_and_files() { + let outputs = [ + ToolOutput::text("plain text"), + ToolOutput::structured(json!({ "ok": true })), + ToolOutput::parts(vec![ + Part::text("part text"), + Part::structured(json!({ "part": true })), + ]), + ToolOutput::files(vec![ + FilePart::named("artifact.txt", DataRef::inline_text("artifact body")) + .with_mime_type("text/plain"), + FilePart::named("remote.txt", DataRef::uri("file:///tmp/remote.txt")), + ]), + ]; + let contents = outputs + .into_iter() + .map(|output| { + let result = ToolResultPart::success(ToolCallId::new("call"), output); + let update = tool_result_update(&result, wire::ToolCallStatus::Completed); + serde_json::to_value(update).expect("serialize tool update")["content"].clone() + }) + .collect::>(); + + assert_eq!(contents[0][0]["content"]["text"], "plain text"); + assert_eq!(contents[1][0]["content"]["text"], r#"{"ok":true}"#); + assert_eq!(contents[2].as_array().map(Vec::len), Some(2)); + assert_eq!(contents[2][0]["content"]["text"], "part text"); + assert_eq!(contents[3].as_array().map(Vec::len), Some(2)); + assert_eq!( + contents[3][0]["content"]["resource"]["text"], + "artifact body" + ); + assert_eq!(contents[3][1]["content"]["uri"], "file:///tmp/remote.txt"); + } + + #[tokio::test] + async fn joining_stuck_session_tasks_uses_one_shared_deadline() { + struct DropMarker(Arc); + impl Drop for DropMarker { + fn drop(&mut self) { + self.0.fetch_add(1, Ordering::AcqRel); + } + } + + let dropped = Arc::new(AtomicUsize::new(0)); + let tasks = (0..3) + .map(|_| { + let dropped = Arc::clone(&dropped); + tokio::spawn(async move { + let _marker = DropMarker(dropped); + std::future::pending::<()>().await; + }) + }) + .collect(); + tokio::task::yield_now().await; + let started = std::time::Instant::now(); + join_tasks_until( + tokio::time::Instant::now() + Duration::from_millis(50), + tasks, + ) + .await; + + assert!(started.elapsed() < Duration::from_millis(250)); + tokio::time::timeout(Duration::from_millis(250), async { + while dropped.load(Ordering::Acquire) != 3 { + tokio::task::yield_now().await; + } + }) + .await + .expect("aborted session tasks were not dropped"); + } + + #[tokio::test] + async fn aborting_serve_cleans_up_active_sessions() { + let adapter = MockAdapter::new(); + adapter.enqueue(tool_turn("serve-cancel")); + let tool = CancellationAwareTool::new(); + let (client_transport, agent_transport) = Channel::duplex(); + let server = tokio::spawn({ + let factory = ToolTestFactory { + adapter, + tool: tool.clone(), + }; + async move { + AcpHeadlessRuntime::::builder() + .agent_factory(factory) + .serve(agent_transport) + .await + } + }); + let client = tokio::spawn(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 = cx + .send_request(wire::NewSessionRequest::new(cwd)) + .block_task() + .await?; + cx.send_request(wire::PromptRequest::new( + session.session_id, + vec![wire::ContentBlock::Text(wire::TextContent::new("run"))], + )) + .block_task() + .await?; + std::future::pending::>().await + }, + )); + + tool.wait_for_entered(1).await; + server.abort(); + let _ = server.await; + tokio::time::timeout(Duration::from_secs(2), async { + while tool.cleaned.load(Ordering::Acquire) != 1 { + tokio::task::yield_now().await; + } + }) + .await + .expect("serve drop did not clean the active tool"); + client.abort(); + let _ = client.await; + } + + #[cfg(not(feature = "unstable-inject"))] + #[test] + 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 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(); + + 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(), + }), + 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 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(); + 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_session(AcpSessionBinding::new( + first_acp.clone(), + agentkit_id.clone(), + first_client, + )) + .expect("first binding"); + let error = integration + .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_session(&first_acp) + .expect("unbind first"); + } + + #[test] + fn retained_transcript_failures_use_custom_error_stop_reason() { + for reason in [FinishReason::Blocked, FinishReason::Error] { assert_eq!( - session_updates - .iter() - .filter(|update| matches!( - update, - wire::SessionUpdate::StateUpdate(wire::StateUpdate::Idle(_)) - )) - .count(), - 1, - "Idle must come only from TurnFinished" + finish_reason_to_stop_reason(&reason), + wire::StopReason::Other("_error".into()) ); - let user_id = match session_updates[0] { - wire::SessionUpdate::UserMessage(message) => message.message_id.to_string(), - _ => unreachable!(), - }; - let agent_ids = session_updates - .iter() - .filter_map(|update| match update { - wire::SessionUpdate::AgentMessageChunk(chunk) => { - Some(chunk.message_id.to_string()) - } - _ => None, - }) - .collect::>(); - let thought_ids = session_updates - .iter() - .filter_map(|update| match update { - wire::SessionUpdate::AgentThoughtChunk(chunk) => { - Some(chunk.message_id.to_string()) - } - _ => None, - }) - .collect::>(); - assert_eq!(user_id, format!("{session_id}-user-1")); - if session_id == "session-1" { - assert_eq!( - agent_ids, - [ - "session-1-agent-2".to_string(), - "session-1-agent-3".to_string() - ] - ); - assert_eq!(thought_ids, ["session-1-thought-3"]); - assert_ne!(agent_ids[0], agent_ids[1]); - } else { - assert_eq!(agent_ids, ["session-2-agent-2"]); - assert_eq!(thought_ids, ["session-2-thought-2"]); + } + } + + #[test] + fn unsupported_approval_keeps_error_stop_reason_on_cancelled_turn() { + 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"); + let cancellation = CancellationController::new(); + let generation = cancellation.handle().generation(); + 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"); + integration + .begin_prompt_owner(&acp_id, 1, cancellation.handle().checkpoint()) + .expect("begin prompt"); + + integration.route_event( + &agentkit_id, + AgentEvent::TurnStarted { + session_id: agentkit_id.clone(), + turn_id: TurnId::new("turn-1"), + }, + ); + integration.mark_unsupported_approval(&acp_id, cancellation.handle(), generation); + integration.route_event( + &agentkit_id, + AgentEvent::TurnFinished(TurnResult { + turn_id: TurnId::new("turn-1"), + finish_reason: FinishReason::Cancelled, + items: Vec::new(), + usage: None, + metadata: MetadataMap::new(), + }), + ); + + assert!(matches!( + messages.try_recv(), + Ok(ClientMessage::Update(notification)) + if matches!( + ¬ification.update, + wire::SessionUpdate::StateUpdate(wire::StateUpdate::Running(_)) + ) + )); + assert!(matches!( + messages.try_recv(), + Ok(ClientMessage::Update(notification)) + if matches!( + ¬ification.update, + wire::SessionUpdate::StateUpdate(wire::StateUpdate::Idle(idle)) + if idle.stop_reason.as_ref() + == Some(&wire::StopReason::Other("_error".into())) + ) + )); + assert!(messages.try_recv().is_err()); + } + + #[test] + fn prompt_cancellation_overrides_terminal_and_unsupported_stop_reasons() { + for (index, finish_reason, unsupported_approval) in [ + (1, FinishReason::Completed, false), + (2, FinishReason::Error, false), + (3, FinishReason::Completed, true), + ] { + let integration = AcpIntegration::default(); + let (client, mut messages) = ClientHandle::channel(); + let acp_id = wire::SessionId::new(format!("acp-session-{index}")); + let agentkit_id = AgentkitSessionId::new(format!("agentkit-session-{index}")); + let active_prompt = Arc::new(AtomicU64::new(index)); + let cancellation = CancellationController::new(); + let generation = cancellation.handle().generation(); + integration + .bind(acp_id.clone(), agentkit_id.clone(), client) + .expect("bind session"); + integration + .install_prompt_state(&acp_id, active_prompt, Arc::new(Mutex::new(()))) + .expect("install prompt state"); + integration + .begin_prompt_owner(&acp_id, index, cancellation.handle().checkpoint()) + .expect("begin prompt"); + let turn_id = TurnId::new(format!("turn-{index}")); + integration.route_event( + &agentkit_id, + AgentEvent::TurnStarted { + session_id: agentkit_id.clone(), + turn_id: turn_id.clone(), + }, + ); + if unsupported_approval { + integration.mark_unsupported_approval(&acp_id, cancellation.handle(), generation); } - assert_ne!(agent_ids.last(), thought_ids.last()); + cancellation.interrupt(); + integration.route_event( + &agentkit_id, + AgentEvent::TurnFinished(agentkit_loop::TurnResult { + turn_id, + finish_reason, + items: Vec::new(), + usage: None, + metadata: MetadataMap::new(), + }), + ); + assert!(matches!( - session_updates.last(), - Some(wire::SessionUpdate::StateUpdate(wire::StateUpdate::Idle(idle))) - if idle.stop_reason == Some(wire::StopReason::EndTurn) + messages.try_recv(), + Ok(ClientMessage::Update(notification)) + if matches!( + notification.update, + wire::SessionUpdate::StateUpdate(wire::StateUpdate::Running(_)) + ) + )); + assert!(matches!( + messages.try_recv(), + Ok(ClientMessage::Update(notification)) + if matches!( + ¬ification.update, + wire::SessionUpdate::StateUpdate(wire::StateUpdate::Idle(idle)) + if idle.stop_reason == Some(wire::StopReason::Cancelled) + ) )); } } - #[tokio::test] - async fn idle_session_drives_background_completion_without_another_prompt() { - let adapter = MockAdapter::new(); - adapter.enqueue(tool_turn("background-call")); - adapter.enqueue(streamed_text("background observed")); - let tool = BlockingTool::text("blocking_tool", "background done"); - let updates = Arc::new(Mutex::new(Vec::new())); - let (client_transport, agent_transport) = Channel::duplex(); - - let server = tokio::spawn({ - let factory = BackgroundToolTestFactory { - adapter, - tool: tool.clone(), - }; - async move { - AcpHeadlessRuntime::::builder() - .agent_factory(factory) - .serve(agent_transport) - .await + #[test] + fn accepted_prompt_failure_queues_error_idle_after_cleanup() { + let integration = AcpIntegration::default(); + let owner = 7; + let active_prompt = Arc::new(AtomicU64::new(owner)); + let driving_prompt = Arc::new(AtomicU64::new(owner)); + let cancelled_prompt = Arc::new(AtomicU64::new(owner)); + let lifecycle = Arc::new(Mutex::new(())); + let prompt_at_enqueue = Arc::clone(&active_prompt); + let lifecycle_at_enqueue = Arc::clone(&lifecycle); + let (client, mut messages) = ClientHandle::channel_with_update_hook(move |update| { + if matches!( + update, + wire::SessionUpdate::StateUpdate(wire::StateUpdate::Idle(_)) + ) { + assert_eq!(prompt_at_enqueue.load(Ordering::Acquire), 0); + assert!(lifecycle_at_enqueue.try_lock().is_err()); } }); - - 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!(), + let acp_id = wire::SessionId::new("acp-failed-prompt"); + integration + .bind( + acp_id.clone(), + AgentkitSessionId::new("agentkit-failed-prompt"), + client.clone(), ) - .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 session = cx - .send_request(wire::NewSessionRequest::new(cwd)) - .block_task() - .await?; - cx.send_request(wire::PromptRequest::new( - session.session_id.clone(), - vec![wire::ContentBlock::Text(wire::TextContent::new( - "start background work", - ))], - )) - .block_task() - .await?; - wait_for_idle(&updates, &session.session_id).await; - tool.wait_until_entered().await; - tool.release(); - - tokio::time::timeout(Duration::from_secs(2), async { - loop { - let updates = updates.lock().unwrap(); - let completed = updates.iter().any(|(id, update)| { - id == &session.session_id - && matches!( - update, - wire::SessionUpdate::ToolCallUpdate(update) - if update.tool_call_id.to_string() == "background-call" - && matches!( - update.status, - agent_client_protocol::schema::MaybeUndefined::Value( - wire::ToolCallStatus::Completed - ) - ) - ) - }); - let message_delivered = updates.iter().any(|(id, update)| { - id == &session.session_id - && matches!( - update, - wire::SessionUpdate::AgentMessageChunk(chunk) - if matches!( - &chunk.content, - wire::ContentBlock::Text(text) - if text.text == "background observed" - ) - ) - }); - if completed && message_delivered { - break; - } - drop(updates); - tokio::time::sleep(Duration::from_millis(5)).await; - } - }) - .await - .expect("background completion was not delivered while idle"); - cx.send_request(wire::CloseSessionRequest::new(session.session_id)) - .block_task() - .await?; - Ok(()) - } - }); + .expect("bind session"); + integration + .install_prompt_state(&acp_id, Arc::clone(&active_prompt), Arc::clone(&lifecycle)) + .expect("install prompt state"); + let cancellation = CancellationController::new(); + integration + .begin_prompt_owner(&acp_id, owner, cancellation.handle().checkpoint()) + .expect("begin prompt"); - tokio::time::timeout(Duration::from_secs(5), client) - .await - .expect("client timed out") - .expect("client run"); - server.abort(); - let _ = server.await; + fail_accepted_prompt( + &client, + &integration, + &acp_id, + &active_prompt, + &driving_prompt, + &cancelled_prompt, + &lifecycle, + owner, + true, + ); - let updates = updates.lock().unwrap(); - assert_eq!( - updates - .iter() - .filter(|(_, update)| matches!(update, wire::SessionUpdate::UserMessage(_))) - .count(), - 1, - "the background continuation must not require or replay a prompt" + assert_eq!(active_prompt.load(Ordering::Acquire), 0); + assert_eq!(driving_prompt.load(Ordering::Acquire), 0); + assert_eq!(cancelled_prompt.load(Ordering::Acquire), 0); + let session = integration.session(&acp_id).expect("bound session"); + assert!( + session + .prompt_state + .lock() + .unwrap() + .as_ref() + .is_some_and(|state| state.pending_owner.is_none()) ); - assert!(updates.iter().any(|(_, update)| matches!( - update, - wire::SessionUpdate::AgentMessageChunk(chunk) + assert!(matches!( + messages.try_recv(), + Ok(ClientMessage::Update(notification)) if matches!( - &chunk.content, - wire::ContentBlock::Text(text) if text.text == "background observed" + ¬ification.update, + wire::SessionUpdate::StateUpdate(wire::StateUpdate::Idle(idle)) + if idle.stop_reason == Some(error_stop_reason()) ) - ))); - assert_eq!( - updates - .iter() - .filter(|(_, update)| matches!( - update, - wire::SessionUpdate::StateUpdate(wire::StateUpdate::Running(_)) - )) - .count(), - 1, - "idle background completion must not re-enter Running" + )); + } + + #[test] + fn idle_is_queued_after_releasing_its_prompt_owner() { + let integration = AcpIntegration::default(); + let active_prompt = Arc::new(AtomicU64::new(7)); + let lifecycle = Arc::new(Mutex::new(())); + let prompt_at_enqueue = Arc::clone(&active_prompt); + let lifecycle_at_enqueue = Arc::clone(&lifecycle); + let (client, mut messages) = ClientHandle::channel_with_update_hook(move |update| { + if matches!( + update, + wire::SessionUpdate::StateUpdate(wire::StateUpdate::Idle(_)) + ) { + assert_eq!( + prompt_at_enqueue.load(Ordering::Acquire), + 0, + "Idle was enqueued before prompt ownership was released" + ); + assert!( + lifecycle_at_enqueue.try_lock().is_err(), + "Idle was enqueued outside the prompt lifecycle lock" + ); + } + }); + 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::clone(&active_prompt), Arc::clone(&lifecycle)) + .expect("install prompt state"); + let cancellation = CancellationController::new(); + integration + .begin_prompt_owner(&acp_id, 7, cancellation.handle().checkpoint()) + .expect("begin prompt"); + let turn_id = TurnId::new("turn-1"); + integration.route_event( + &agentkit_id, + AgentEvent::TurnStarted { + session_id: agentkit_id.clone(), + turn_id: turn_id.clone(), + }, ); - assert_eq!( - updates - .iter() - .filter(|(_, update)| matches!( - update, - wire::SessionUpdate::StateUpdate(wire::StateUpdate::Idle(_)) - )) - .count(), - 1, - "idle background completion must not emit another Idle" + integration.route_event( + &agentkit_id, + AgentEvent::TurnFinished(agentkit_loop::TurnResult { + turn_id, + finish_reason: FinishReason::Completed, + items: Vec::new(), + usage: None, + metadata: MetadataMap::new(), + }), ); - } - #[tokio::test] - async fn cancel_without_prompt_owner_interrupts_blocked_autonomous_inference() { - let adapter = BlockingInferenceAdapter::new( - [ - tool_turn("background-call"), - streamed_text("autonomous complete"), - ], - 1, + assert_eq!(active_prompt.load(Ordering::Acquire), 0); + assert!( + std::iter::from_fn(|| messages.try_recv().ok()).any(|message| { + matches!( + message, + ClientMessage::Update(notification) + if matches!( + notification.update, + wire::SessionUpdate::StateUpdate(wire::StateUpdate::Idle(_)) + ) + ) + }) ); - let tool = BlockingTool::text("blocking_tool", "background done"); - let updates = Arc::new(Mutex::new(Vec::new())); - let (client_transport, agent_transport) = Channel::duplex(); + } - let server = tokio::spawn({ - let factory = BackgroundToolTestFactory { - adapter: adapter.clone(), - tool: tool.clone(), - }; - async move { - AcpHeadlessRuntime::::builder() - .agent_factory(factory) - .serve(agent_transport) - .await + #[test] + fn cancellation_while_turn_finish_waits_for_lifecycle_lock_wins() { + let integration = AcpIntegration::default(); + let active_prompt = Arc::new(AtomicU64::new(9)); + let lifecycle = Arc::new(Mutex::new(())); + let prompt_at_enqueue = Arc::clone(&active_prompt); + let lifecycle_at_enqueue = Arc::clone(&lifecycle); + let (client, mut messages) = ClientHandle::channel_with_update_hook(move |update| { + if matches!( + update, + wire::SessionUpdate::StateUpdate(wire::StateUpdate::Idle(_)) + ) { + assert_eq!(prompt_at_enqueue.load(Ordering::Acquire), 0); + assert!(lifecycle_at_enqueue.try_lock().is_err()); } }); + let acp_id = wire::SessionId::new("acp-race-session"); + let agentkit_id = AgentkitSessionId::new("agentkit-race-session"); + integration + .bind(acp_id.clone(), agentkit_id.clone(), client) + .expect("bind session"); + integration + .install_prompt_state(&acp_id, Arc::clone(&active_prompt), Arc::clone(&lifecycle)) + .expect("install prompt state"); + let cancellation = CancellationController::new(); + integration + .begin_prompt_owner(&acp_id, 9, cancellation.handle().checkpoint()) + .expect("begin prompt"); + let turn_id = TurnId::new("turn-race"); + integration.route_event( + &agentkit_id, + AgentEvent::TurnStarted { + session_id: agentkit_id.clone(), + turn_id: turn_id.clone(), + }, + ); - 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 tool = tool.clone(); - 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 session = cx - .send_request(wire::NewSessionRequest::new(cwd)) - .block_task() - .await?; - cx.send_request(wire::PromptRequest::new( - session.session_id.clone(), - vec![wire::ContentBlock::Text(wire::TextContent::new("start"))], - )) - .block_task() - .await?; - wait_for_idle(&updates, &session.session_id).await; - tool.wait_until_entered().await; - tool.release(); - adapter.wait_until_blocked().await; + let lifecycle_guard = lifecycle.lock().unwrap(); + let ready = Arc::new(std::sync::Barrier::new(2)); + let finish_ready = Arc::clone(&ready); + let finish_integration = integration.clone(); + let finish_session = agentkit_id.clone(); + let finish = std::thread::spawn(move || { + finish_ready.wait(); + finish_integration.route_event( + &finish_session, + AgentEvent::TurnFinished(agentkit_loop::TurnResult { + turn_id, + finish_reason: FinishReason::Completed, + items: Vec::new(), + usage: None, + metadata: MetadataMap::new(), + }), + ); + }); + ready.wait(); + cancellation.interrupt(); + drop(lifecycle_guard); + finish.join().expect("turn finish thread"); - cx.send_notification(wire::CancelSessionNotification::new( - session.session_id.clone(), - ))?; - cx.send_request(wire::ListSessionsRequest::new()) - .block_task() - .await?; - tokio::time::timeout(Duration::from_secs(2), async { - while !adapter.was_cancelled() { - tokio::task::yield_now().await; - } - }) - .await - .expect("session cancel did not interrupt autonomous inference"); + assert!(matches!( + messages.try_recv(), + Ok(ClientMessage::Update(notification)) + if matches!( + notification.update, + wire::SessionUpdate::StateUpdate(wire::StateUpdate::Running(_)) + ) + )); + assert!(matches!( + messages.try_recv(), + Ok(ClientMessage::Update(notification)) + if matches!( + ¬ification.update, + wire::SessionUpdate::StateUpdate(wire::StateUpdate::Idle(idle)) + if idle.stop_reason == Some(wire::StopReason::Cancelled) + ) + )); + } - cx.send_request(wire::CloseSessionRequest::new(session.session_id)) - .block_task() - .await?; - Ok(()) - } - }); + #[test] + fn stale_prompt_owner_cannot_release_new_prompt() { + let active_prompt = AtomicU64::new(1); + release_prompt(&active_prompt, 1); + assert_eq!(active_prompt.load(Ordering::Acquire), 0); - tokio::time::timeout(Duration::from_secs(5), client) - .await - .expect("client timed out") - .expect("client run"); - server.abort(); - let _ = server.await; + active_prompt.store(2, Ordering::Release); + release_prompt(&active_prompt, 1); + assert_eq!(active_prompt.load(Ordering::Acquire), 2); + release_prompt(&active_prompt, 2); + assert_eq!(active_prompt.load(Ordering::Acquire), 0); } #[tokio::test] - async fn queued_prompt_is_acknowledged_without_cancelling_blocked_autonomous_inference() { - let adapter = BlockingInferenceAdapter::new( - [ - tool_turn("background-call"), - streamed_text("autonomous complete"), - streamed_text("queued prompt complete"), - ], - 1, - ); - let tool = BlockingTool::text("blocking_tool", "background done"); + async fn cancel_then_prompt_and_close_cleanup_active_tools() { + let adapter = MockAdapter::new(); + adapter.enqueue(tool_turn("call-cancel")); + adapter.enqueue(streamed_text("ready again")); + adapter.enqueue(tool_turn("call-close")); + let tool = CancellationAwareTool::new(); let updates = Arc::new(Mutex::new(Vec::new())); let (client_transport, agent_transport) = Channel::duplex(); let server = tokio::spawn({ - let factory = BackgroundToolTestFactory { + let factory = ToolTestFactory { adapter: adapter.clone(), tool: tool.clone(), }; async move { - AcpHeadlessRuntime::::builder() + AcpHeadlessRuntime::::builder() .agent_factory(factory) .serve(agent_transport) .await @@ -2655,731 +5312,1244 @@ mod tests { agent_client_protocol::on_receive_notification!(), ) .connect_with(client_transport, { - let adapter = adapter.clone(); - let tool = tool.clone(); let updates = Arc::clone(&updates); + let tool = tool.clone(); async move |cx| { - cx.send_request(wire::InitializeRequest::new( - wire::ProtocolVersion::V2, - wire::Implementation::new("test-client", "1"), - )) - .block_task() - .await?; + let initialize = cx + .send_request(wire::InitializeRequest::new( + wire::ProtocolVersion::V2, + wire::Implementation::new("test-client", "1"), + )) + .block_task() + .await?; + assert!(initialize.capabilities.session.is_some()); let cwd = std::env::current_dir() .map_err(agent_client_protocol::Error::into_internal_error)?; - let session = cx - .send_request(wire::NewSessionRequest::new(cwd)) + let first = cx + .send_request(wire::NewSessionRequest::new(cwd.clone())) + .block_task() + .await?; + + let listed = cx + .send_request(wire::ListSessionsRequest::new()) .block_task() .await?; + assert_eq!(listed.sessions.len(), 1); + assert_eq!(listed.sessions[0].session_id, first.session_id); + cx.send_request(wire::ResumeSessionRequest::new( + first.session_id.clone(), + cwd.clone(), + )) + .block_task() + .await?; + cx.send_request(wire::PromptRequest::new( - session.session_id.clone(), - vec![wire::ContentBlock::Text(wire::TextContent::new("start"))], + first.session_id.clone(), + vec![wire::ContentBlock::Text(wire::TextContent::new( + "start tool", + ))], )) .block_task() .await?; - wait_for_idle(&updates, &session.session_id).await; - tool.wait_until_entered().await; - tool.release(); - adapter.wait_until_blocked().await; + tool.wait_for_entered(1).await; + cx.send_notification(wire::CancelSessionNotification::new( + first.session_id.clone(), + ))?; + wait_for_idle_count(&updates, &first.session_id, 1).await; + assert_eq!(tool.cleaned.load(Ordering::Acquire), 1); tokio::time::timeout( - Duration::from_millis(500), + Duration::from_millis(250), cx.send_request(wire::PromptRequest::new( - session.session_id.clone(), - vec![wire::ContentBlock::Text(wire::TextContent::new( - "queued while idle", - ))], + first.session_id.clone(), + vec![wire::ContentBlock::Text(wire::TextContent::new("again"))], )) .block_task(), ) .await - .expect("prompt response waited for autonomous inference")?; - assert_eq!( - updates - .lock() - .unwrap() - .iter() - .filter(|(_, update)| matches!( - update, - wire::SessionUpdate::UserMessage(_) - )) - .count(), - 1, - "queued prompt update crossed the response gate" - ); + .expect("session was not ready when Idle was emitted")?; + wait_for_idle_count(&updates, &first.session_id, 2).await; - cx.send_notification(wire::CancelSessionNotification::new( - session.session_id.clone(), - ))?; - cx.send_request(wire::ListSessionsRequest::new()) + let second = cx + .send_request(wire::NewSessionRequest::new(cwd)) .block_task() .await?; - assert!( - tokio::time::timeout(Duration::from_millis(50), async { - while !adapter.was_cancelled() { - tokio::task::yield_now().await; - } - }) - .await - .is_err(), - "queued prompt cancellation interrupted autonomous inference" - ); - - adapter.release(); - tokio::time::timeout(Duration::from_secs(2), async { - loop { - if updates.lock().unwrap().iter().any(|(_, update)| { - matches!( - update, - wire::SessionUpdate::AgentMessageChunk(chunk) - if matches!( - &chunk.content, - wire::ContentBlock::Text(text) - if text.text == "autonomous complete" - ) - ) - }) { - break; - } - tokio::time::sleep(Duration::from_millis(5)).await; - } - }) + cx.send_request(wire::PromptRequest::new( + second.session_id.clone(), + vec![wire::ContentBlock::Text(wire::TextContent::new( + "close tool", + ))], + )) + .block_task() + .await?; + tool.wait_for_entered(2).await; + tokio::time::timeout( + Duration::from_secs(3), + cx.send_request(wire::CloseSessionRequest::new(second.session_id.clone())) + .block_task(), + ) .await - .expect("autonomous inference did not complete"); - assert!(!adapter.was_cancelled()); - assert_eq!( - updates - .lock() - .unwrap() - .iter() - .filter(|(_, update)| matches!( - update, - wire::SessionUpdate::UserMessage(_) - )) - .count(), - 1, - "cancelled queued prompt was driven" - ); - - cx.send_request(wire::CloseSessionRequest::new(session.session_id)) + .expect("close did not complete after cooperative cleanup")?; + assert_eq!(tool.cleaned.load(Ordering::Acquire), 2); + cx.send_request(wire::CloseSessionRequest::new(first.session_id)) .block_task() .await?; Ok(()) } }); - tokio::time::timeout(Duration::from_secs(5), client) + tokio::time::timeout(Duration::from_secs(8), client) .await .expect("client timed out") .expect("client run"); server.abort(); let _ = server.await; + + let updates = updates.lock().unwrap(); + for call_id in ["call-cancel", "call-close"] { + let statuses = updates + .iter() + .filter_map(|(_, update)| match update { + wire::SessionUpdate::ToolCallUpdate(update) + if update.tool_call_id.to_string() == call_id => + { + Some(update.status.clone()) + } + _ => None, + }) + .collect::>(); + assert!(statuses.iter().any(|status| matches!( + status, + agent_client_protocol::schema::MaybeUndefined::Value(wire::ToolCallStatus::Pending) + ))); + assert!(statuses.iter().any(|status| matches!( + status, + agent_client_protocol::schema::MaybeUndefined::Value( + wire::ToolCallStatus::InProgress + ) + ))); + assert!(statuses.iter().any(|status| matches!( + status, + agent_client_protocol::schema::MaybeUndefined::Value(wire::ToolCallStatus::Failed) + ))); + } } - #[test] - fn v2_tool_updates_include_visible_text_structured_parts_and_files() { - let outputs = [ - ToolOutput::text("plain text"), - ToolOutput::structured(json!({ "ok": true })), - ToolOutput::parts(vec![ - Part::text("part text"), - Part::structured(json!({ "part": true })), - ]), - ToolOutput::files(vec![ - FilePart::named("artifact.txt", DataRef::inline_text("artifact body")) - .with_mime_type("text/plain"), - FilePart::named("remote.txt", DataRef::uri("file:///tmp/remote.txt")), - ]), - ]; - let contents = outputs - .into_iter() - .map(|output| { - let result = ToolResultPart::success(ToolCallId::new("call"), output); - let update = tool_result_update(&result, wire::ToolCallStatus::Completed); - serde_json::to_value(update).expect("serialize tool update")["content"].clone() - }) - .collect::>(); + #[derive(Clone, Default)] + struct BlockingAdapter; - assert_eq!(contents[0][0]["content"]["text"], "plain text"); - assert_eq!(contents[1][0]["content"]["text"], r#"{"ok":true}"#); - assert_eq!(contents[2].as_array().map(Vec::len), Some(2)); - assert_eq!(contents[2][0]["content"]["text"], "part text"); - assert_eq!(contents[3].as_array().map(Vec::len), Some(2)); - assert_eq!( - contents[3][0]["content"]["resource"]["text"], - "artifact body" - ); - assert_eq!(contents[3][1]["content"]["uri"], "file:///tmp/remote.txt"); + struct BlockingSession; + + struct BlockingTurn { + finished: bool, } - #[tokio::test] - async fn joining_stuck_session_tasks_uses_one_shared_deadline() { - struct DropMarker(Arc); - impl Drop for DropMarker { - fn drop(&mut self) { - self.0.fetch_add(1, Ordering::AcqRel); + #[async_trait] + impl ModelAdapter for BlockingAdapter { + type Session = BlockingSession; + + async fn start_session(&self, _config: SessionConfig) -> Result { + Ok(BlockingSession) + } + } + + #[async_trait] + impl ModelSession for BlockingSession { + type Turn = BlockingTurn; + + async fn begin_turn( + &mut self, + _request: TurnRequest, + _cancellation: Option, + ) -> Result { + Ok(BlockingTurn { finished: false }) + } + } + + #[async_trait] + impl ModelTurn for BlockingTurn { + async fn next_event( + &mut self, + cancellation: Option, + ) -> Result, LoopError> { + if self.finished { + return Ok(None); + } + cancellation + .expect("cancellation installed") + .cancelled() + .await; + self.finished = true; + Ok(Some(ModelTurnEvent::Finished(ModelTurnResult { + model: None, + response_id: None, + finish_reason: FinishReason::Cancelled, + output_items: Vec::new(), + usage: None, + metadata: MetadataMap::new(), + }))) + } + } + + #[cfg(feature = "unstable-inject")] + #[derive(Clone)] + struct GatedAdapter { + permits: Arc, + next_turn: Arc, + requests: Arc>>, + } + + #[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())), + } + } + + fn release(&self, count: usize) { + self.permits.add_permits(count); + } + + fn requests(&self) -> Vec { + self.requests + .lock() + .unwrap_or_else(|error| error.into_inner()) + .clone() + } + } + + #[cfg(feature = "unstable-inject")] + #[derive(Clone)] + struct SecondStartGatedFactory { + adapter: GatedAdapter, + starts: Arc, + second_started: Arc, + second_release: Arc, + } + + #[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)), } } - let dropped = Arc::new(AtomicUsize::new(0)); - let tasks = (0..3) - .map(|_| { - let dropped = Arc::clone(&dropped); - tokio::spawn(async move { - let _marker = DropMarker(dropped); - std::future::pending::<()>().await; - }) - }) - .collect(); - tokio::task::yield_now().await; - let started = std::time::Instant::now(); - join_tasks_until( - tokio::time::Instant::now() + Duration::from_millis(50), - tasks, - ) - .await; - - assert!(started.elapsed() < Duration::from_millis(250)); - tokio::time::timeout(Duration::from_millis(250), async { - while dropped.load(Ordering::Acquire) != 3 { - tokio::task::yield_now().await; + 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"); } - }) - .await - .expect("aborted session tasks were not dropped"); + } + + fn release_second_start(&self) { + self.second_release.add_permits(1); + } } - #[tokio::test] - async fn aborting_serve_cleans_up_active_sessions() { - let adapter = MockAdapter::new(); - adapter.enqueue(tool_turn("serve-cancel")); - let tool = CancellationAwareTool::new(); - let (client_transport, agent_transport) = Channel::duplex(); - let server = tokio::spawn({ - let factory = ToolTestFactory { - adapter, - tool: tool.clone(), - }; - async move { - AcpHeadlessRuntime::::builder() - .agent_factory(factory) - .serve(agent_transport) + #[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(); } - }); - let client = tokio::spawn(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 = cx - .send_request(wire::NewSessionRequest::new(cwd)) - .block_task() - .await?; - cx.send_request(wire::PromptRequest::new( - session.session_id, - vec![wire::ContentBlock::Text(wire::TextContent::new("run"))], - )) - .block_task() - .await?; - std::future::pending::>().await - }, - )); + 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())) + } + } - tool.wait_for_entered(1).await; - server.abort(); - let _ = server.await; + #[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 { - while tool.cleaned.load(Ordering::Acquire) != 1 { - tokio::task::yield_now().await; + 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("serve drop did not clean the active tool"); - client.abort(); - let _ = client.await; + .expect("raw request timed out") } - #[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(&acp_id, 1, cancellation.handle().checkpoint()) - .expect("begin prompt"); + #[cfg(feature = "unstable-inject")] + struct GatedSession { + permits: Arc, + next_turn: Arc, + requests: Arc>>, + } - 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); - } + #[cfg(feature = "unstable-inject")] + struct GatedTurn { + permits: Arc, + events: VecDeque, + started: bool, + } - 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::Flush(_) => None, + #[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), }) - .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(&first_acp).expect("unbind first"); + #[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, + }) + } } - #[test] - fn retained_transcript_failures_use_custom_error_stop_reason() { - for reason in [FinishReason::Blocked, FinishReason::Error] { - assert_eq!( - finish_reason_to_stop_reason(&reason), - wire::StopReason::Other("_error".into()) - ); + #[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()) } } - #[test] - fn unsupported_approval_keeps_error_stop_reason_on_cancelled_turn() { - 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"); - let cancellation = CancellationController::new(); - let generation = cancellation.handle().generation(); - 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"); - integration - .begin_prompt(&acp_id, 1, cancellation.handle().checkpoint()) - .expect("begin prompt"); + #[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(); - integration.route_event( - &agentkit_id, - AgentEvent::TurnStarted { - session_id: agentkit_id.clone(), - turn_id: TurnId::new("turn-1"), - }, + 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" ); - integration.mark_unsupported_approval(&acp_id, cancellation.handle(), generation); - integration.route_event( - &agentkit_id, - AgentEvent::TurnFinished(TurnResult { - turn_id: TurnId::new("turn-1"), - finish_reason: FinishReason::Cancelled, - items: Vec::new(), - usage: None, - metadata: MetadataMap::new(), - }), + 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()) ); - assert!(matches!( - messages.try_recv(), - Ok(ClientMessage::Update(notification)) - if matches!( - ¬ification.update, - wire::SessionUpdate::StateUpdate(wire::StateUpdate::Running(_)) + 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(), ) - )); - assert!(matches!( - messages.try_recv(), - Ok(ClientMessage::Update(notification)) - if matches!( - ¬ification.update, - wire::SessionUpdate::StateUpdate(wire::StateUpdate::Idle(idle)) - if idle.stop_reason.as_ref() - == Some(&wire::StopReason::Other("_error".into())) + .unwrap(), + )) + .unwrap(); + client + .tx + .unbounded_send(TransportFrame::Single( + RawJsonRpcMessage::notification( + "$/cancel_request".to_string(), + json!({ "requestId": 10 }), ) - )); - assert!(messages.try_recv().is_err()); + .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; } - #[test] - fn prompt_cancellation_overrides_terminal_and_unsupported_stop_reasons() { - for (index, finish_reason, unsupported_approval) in [ - (1, FinishReason::Completed, false), - (2, FinishReason::Error, false), - (3, FinishReason::Completed, true), - ] { - let integration = AcpIntegration::default(); - let (client, mut messages) = ClientHandle::channel(); - let acp_id = wire::SessionId::new(format!("acp-session-{index}")); - let agentkit_id = AgentkitSessionId::new(format!("agentkit-session-{index}")); - let active_prompt = Arc::new(AtomicU64::new(index)); - let cancellation = CancellationController::new(); - let generation = cancellation.handle().generation(); - integration - .bind(acp_id.clone(), agentkit_id.clone(), client) - .expect("bind session"); - integration - .install_prompt_state(&acp_id, active_prompt, Arc::new(Mutex::new(()))) - .expect("install prompt state"); - integration - .begin_prompt(&acp_id, index, cancellation.handle().checkpoint()) - .expect("begin prompt"); - let turn_id = TurnId::new(format!("turn-{index}")); - integration.route_event( - &agentkit_id, - AgentEvent::TurnStarted { - session_id: agentkit_id.clone(), - turn_id: turn_id.clone(), - }, - ); - if unsupported_approval { - integration.mark_unsupported_approval(&acp_id, cancellation.handle(), generation); + #[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 } - cancellation.interrupt(); - integration.route_event( - &agentkit_id, - AgentEvent::TurnFinished(agentkit_loop::TurnResult { - turn_id, - finish_reason, - items: Vec::new(), - usage: None, - metadata: MetadataMap::new(), - }), - ); + }); + 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)); - assert!(matches!( - messages.try_recv(), - Ok(ClientMessage::Update(notification)) - if matches!( - notification.update, - wire::SessionUpdate::StateUpdate(wire::StateUpdate::Running(_)) - ) - )); - assert!(matches!( - messages.try_recv(), - Ok(ClientMessage::Update(notification)) - if matches!( - ¬ification.update, - wire::SessionUpdate::StateUpdate(wire::StateUpdate::Idle(idle)) - if idle.stop_reason == Some(wire::StopReason::Cancelled) - ) - )); - } + drop(client); + server.abort(); + let _ = server.await; } - #[test] - fn accepted_prompt_failure_queues_error_idle_after_cleanup() { - let integration = AcpIntegration::default(); - let owner = 7; - let active_prompt = Arc::new(AtomicU64::new(owner)); - let driving_prompt = Arc::new(AtomicU64::new(owner)); - let cancelled_prompt = Arc::new(AtomicU64::new(owner)); - let lifecycle = Arc::new(Mutex::new(())); - let prompt_at_enqueue = Arc::clone(&active_prompt); - let lifecycle_at_enqueue = Arc::clone(&lifecycle); - let (client, mut messages) = ClientHandle::channel_with_update_hook(move |update| { - if matches!( - update, - wire::SessionUpdate::StateUpdate(wire::StateUpdate::Idle(_)) - ) { - assert_eq!(prompt_at_enqueue.load(Ordering::Acquire), 0); - assert!(lifecycle_at_enqueue.try_lock().is_err()); + #[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 } }); - let acp_id = wire::SessionId::new("acp-failed-prompt"); - integration - .bind( - acp_id.clone(), - AgentkitSessionId::new("agentkit-failed-prompt"), - client.clone(), - ) - .expect("bind session"); - integration - .install_prompt_state(&acp_id, Arc::clone(&active_prompt), Arc::clone(&lifecycle)) - .expect("install prompt state"); - let cancellation = CancellationController::new(); - integration - .begin_prompt(&acp_id, owner, cancellation.handle().checkpoint()) - .expect("begin prompt"); + 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; - fail_accepted_prompt( - &client, - &integration, - &acp_id, - &active_prompt, - &driving_prompt, - &cancelled_prompt, - &lifecycle, - owner, - true, - ); + 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); - assert_eq!(active_prompt.load(Ordering::Acquire), 0); - assert_eq!(driving_prompt.load(Ordering::Acquire), 0); - assert_eq!(cancelled_prompt.load(Ordering::Acquire), 0); - let session = integration.session(&acp_id).expect("bound session"); - assert!( - session - .prompt_state - .lock() - .unwrap() - .as_ref() - .is_some_and(|state| state.pending_owner.is_none()) - ); - assert!(matches!( - messages.try_recv(), - Ok(ClientMessage::Update(notification)) - if matches!( - ¬ification.update, - wire::SessionUpdate::StateUpdate(wire::StateUpdate::Idle(idle)) - if idle.stop_reason == Some(error_stop_reason()) - ) - )); + drop(client); + server.abort(); + let _ = server.await; } + #[cfg(feature = "unstable-inject")] #[test] - fn idle_is_queued_after_releasing_its_prompt_owner() { - let integration = AcpIntegration::default(); - let active_prompt = Arc::new(AtomicU64::new(7)); - let lifecycle = Arc::new(Mutex::new(())); - let prompt_at_enqueue = Arc::clone(&active_prompt); - let lifecycle_at_enqueue = Arc::clone(&lifecycle); - let (client, mut messages) = ClientHandle::channel_with_update_hook(move |update| { - if matches!( - update, - wire::SessionUpdate::StateUpdate(wire::StateUpdate::Idle(_)) - ) { - assert_eq!( - prompt_at_enqueue.load(Ordering::Acquire), - 0, - "Idle was enqueued before prompt ownership was released" - ); - assert!( - lifecycle_at_enqueue.try_lock().is_err(), - "Idle was enqueued outside the prompt lifecycle lock" - ); - } - }); - 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::clone(&active_prompt), Arc::clone(&lifecycle)) - .expect("install prompt state"); - let cancellation = CancellationController::new(); - integration - .begin_prompt(&acp_id, 7, cancellation.handle().checkpoint()) - .expect("begin prompt"); - let turn_id = TurnId::new("turn-1"); - integration.route_event( - &agentkit_id, - AgentEvent::TurnStarted { - session_id: agentkit_id.clone(), - turn_id: turn_id.clone(), - }, - ); - integration.route_event( - &agentkit_id, - AgentEvent::TurnFinished(agentkit_loop::TurnResult { - turn_id, - finish_reason: FinishReason::Completed, - items: Vec::new(), - usage: None, - metadata: MetadataMap::new(), - }), - ); + 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: InjectionBytes::default(), + commitment: InjectCommitment::Reserved, + }) + .expect("injection below lifetime cap"); + injection.commit(&message_id); + injection.activate(&message_id); + assert!(matches!( + injection.revoke_transition(&message_id), + PendingTransition::Applied + )); + } - assert_eq!(active_prompt.load(Ordering::Acquire), 0); - assert!( - std::iter::from_fn(|| messages.try_recv().ok()).any(|message| { - matches!( - message, - ClientMessage::Update(notification) - if matches!( - notification.update, - wire::SessionUpdate::StateUpdate(wire::StateUpdate::Idle(_)) - ) - ) + let error = injection + .reserve(PendingInject { + message_id: wire::MessageId::new("over-limit"), + content: Vec::new(), + items: Vec::new(), + bytes: InjectionBytes::default(), + 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, + })) ); } - #[test] - fn cancellation_while_turn_finish_waits_for_lifecycle_lock_wins() { - let integration = AcpIntegration::default(); - let active_prompt = Arc::new(AtomicU64::new(9)); - let lifecycle = Arc::new(Mutex::new(())); - let prompt_at_enqueue = Arc::clone(&active_prompt); - let lifecycle_at_enqueue = Arc::clone(&lifecycle); - let (client, mut messages) = ClientHandle::channel_with_update_hook(move |update| { - if matches!( - update, - wire::SessionUpdate::StateUpdate(wire::StateUpdate::Idle(_)) - ) { - assert_eq!(prompt_at_enqueue.load(Ordering::Acquire), 0); - assert!(lifecycle_at_enqueue.try_lock().is_err()); + #[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 acp_id = wire::SessionId::new("acp-race-session"); - let agentkit_id = AgentkitSessionId::new("agentkit-race-session"); - integration - .bind(acp_id.clone(), agentkit_id.clone(), client) - .expect("bind session"); - integration - .install_prompt_state(&acp_id, Arc::clone(&active_prompt), Arc::clone(&lifecycle)) - .expect("install prompt state"); - let cancellation = CancellationController::new(); - integration - .begin_prompt(&acp_id, 9, cancellation.handle().checkpoint()) - .expect("begin prompt"); - let turn_id = TurnId::new("turn-race"); - integration.route_event( - &agentkit_id, - AgentEvent::TurnStarted { - session_id: agentkit_id.clone(), - turn_id: turn_id.clone(), - }, - ); + 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; + } - let lifecycle_guard = lifecycle.lock().unwrap(); - let ready = Arc::new(std::sync::Barrier::new(2)); - let finish_ready = Arc::clone(&ready); - let finish_integration = integration.clone(); - let finish_session = agentkit_id.clone(); - let finish = std::thread::spawn(move || { - finish_ready.wait(); - finish_integration.route_event( - &finish_session, - AgentEvent::TurnFinished(agentkit_loop::TurnResult { - turn_id, - finish_reason: FinishReason::Completed, - items: Vec::new(), - usage: None, - metadata: MetadataMap::new(), - }), - ); + #[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 + } }); - ready.wait(); - cancellation.interrupt(); - drop(lifecycle_guard); - finish.join().expect("turn finish thread"); - assert!(matches!( - messages.try_recv(), - Ok(ClientMessage::Update(notification)) - if matches!( - notification.update, - wire::SessionUpdate::StateUpdate(wire::StateUpdate::Running(_)) - ) - )); - assert!(matches!( - messages.try_recv(), - Ok(ClientMessage::Update(notification)) - if matches!( - ¬ification.update, - wire::SessionUpdate::StateUpdate(wire::StateUpdate::Idle(idle)) - if idle.stop_reason == Some(wire::StopReason::Cancelled) - ) - )); + 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; } - #[test] - fn stale_prompt_owner_cannot_release_new_prompt() { - let active_prompt = AtomicU64::new(1); - release_prompt(&active_prompt, 1); - assert_eq!(active_prompt.load(Ordering::Acquire), 0); - - active_prompt.store(2, Ordering::Release); - release_prompt(&active_prompt, 1); - assert_eq!(active_prompt.load(Ordering::Acquire), 2); - release_prompt(&active_prompt, 2); - assert_eq!(active_prompt.load(Ordering::Acquire), 0); + #[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 cancel_then_prompt_and_close_cleanup_active_tools() { - let adapter = MockAdapter::new(); - adapter.enqueue(tool_turn("call-cancel")); - adapter.enqueue(streamed_text("ready again")); - adapter.enqueue(tool_turn("call-close")); - let tool = CancellationAwareTool::new(); + 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 factory = ToolTestFactory { - adapter: adapter.clone(), - tool: tool.clone(), - }; + let adapter = adapter.clone(); async move { - AcpHeadlessRuntime::::builder() - .agent_factory(factory) + AcpHeadlessRuntime::::builder() + .agent_factory(TestFactory { adapter }) .serve(agent_transport) .await } @@ -3401,8 +6571,8 @@ mod tests { agent_client_protocol::on_receive_notification!(), ) .connect_with(client_transport, { + let adapter = adapter.clone(); let updates = Arc::clone(&updates); - let tool = tool.clone(); async move |cx| { let initialize = cx .send_request(wire::InitializeRequest::new( @@ -3411,173 +6581,284 @@ mod tests { )) .block_task() .await?; - assert!(initialize.capabilities.session.is_some()); + 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_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)?; - let first = cx - .send_request(wire::NewSessionRequest::new(cwd.clone())) + let session_id = cx + .send_request(wire::NewSessionRequest::new(cwd)) .block_task() - .await?; - - let listed = cx - .send_request(wire::ListSessionsRequest::new()) + .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?; - assert_eq!(listed.sessions.len(), 1); - assert_eq!(listed.sessions[0].session_id, first.session_id); - cx.send_request(wire::ResumeSessionRequest::new( - first.session_id.clone(), - cwd.clone(), - )) - .block_task() - .await?; + .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("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( - first.session_id.clone(), - vec![wire::ContentBlock::Text(wire::TextContent::new( - "start tool", - ))], + session_id.clone(), + vec![wire::ContentBlock::Text(wire::TextContent::new("start"))], )) .block_task() .await?; - tool.wait_for_entered(1).await; - cx.send_notification(wire::CancelSessionNotification::new( - first.session_id.clone(), - ))?; - wait_for_idle_count(&updates, &first.session_id, 1).await; - assert_eq!(tool.cleaned.load(Ordering::Acquire), 1); - - tokio::time::timeout( + 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::PromptRequest::new( - first.session_id.clone(), - vec![wire::ContentBlock::Text(wire::TextContent::new("again"))], + cx.send_request(wire::InjectSessionRequest::new( + session_id.clone(), + wire::SessionInjectMode::Steer, + content.clone(), )) .block_task(), ) .await - .expect("session was not ready when Idle was emitted")?; - wait_for_idle_count(&updates, &first.session_id, 2).await; - - let second = cx - .send_request(wire::NewSessionRequest::new(cwd)) + .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" + ); + 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 = + 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( - second.session_id.clone(), - vec![wire::ContentBlock::Text(wire::TextContent::new( - "close tool", - ))], + session_id.clone(), + vec![wire::ContentBlock::Text(wire::TextContent::new("revoke"))], )) .block_task() .await?; - tool.wait_for_entered(2).await; - tokio::time::timeout( - Duration::from_secs(3), - cx.send_request(wire::CloseSessionRequest::new(second.session_id.clone())) - .block_task(), - ) - .await - .expect("close did not complete after cooperative cleanup")?; - assert_eq!(tool.cleaned.load(Ordering::Acquire), 2); - cx.send_request(wire::CloseSessionRequest::new(first.session_id)) + 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?; - Ok(()) - } - }); - - tokio::time::timeout(Duration::from_secs(8), client) - .await - .expect("client timed out") - .expect("client run"); - server.abort(); - let _ = server.await; - - let updates = updates.lock().unwrap(); - for call_id in ["call-cancel", "call-close"] { - let statuses = updates - .iter() - .filter_map(|(_, update)| match update { - wire::SessionUpdate::ToolCallUpdate(update) - if update.tool_call_id.to_string() == call_id => - { - Some(update.status.clone()) - } - _ => None, - }) - .collect::>(); - assert!(statuses.iter().any(|status| matches!( - status, - agent_client_protocol::schema::MaybeUndefined::Value(wire::ToolCallStatus::Pending) - ))); - assert!(statuses.iter().any(|status| matches!( - status, - agent_client_protocol::schema::MaybeUndefined::Value( - wire::ToolCallStatus::InProgress - ) - ))); - assert!(statuses.iter().any(|status| matches!( - status, - agent_client_protocol::schema::MaybeUndefined::Value(wire::ToolCallStatus::Failed) - ))); - } - } - - #[derive(Clone, Default)] - struct BlockingAdapter; - - struct BlockingSession; - - struct BlockingTurn { - finished: bool, - } - - #[async_trait] - impl ModelAdapter for BlockingAdapter { - type Session = BlockingSession; - - async fn start_session(&self, _config: SessionConfig) -> Result { - Ok(BlockingSession) - } - } + .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, + })) + ); - #[async_trait] - impl ModelSession for BlockingSession { - type Turn = BlockingTurn; + 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)); - async fn begin_turn( - &mut self, - _request: TurnRequest, - _cancellation: Option, - ) -> Result { - Ok(BlockingTurn { finished: false }) - } - } + 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(()) + } + }); - #[async_trait] - impl ModelTurn for BlockingTurn { - async fn next_event( - &mut self, - cancellation: Option, - ) -> Result, LoopError> { - if self.finished { - return Ok(None); - } - cancellation - .expect("cancellation installed") - .cancelled() - .await; - self.finished = true; - Ok(Some(ModelTurnEvent::Finished(ModelTurnResult { - model: None, - response_id: None, - finish_reason: FinishReason::Cancelled, - output_items: Vec::new(), - usage: None, - metadata: MetadataMap::new(), - }))) - } + tokio::time::timeout(Duration::from_secs(8), client) + .await + .expect("client timed out") + .expect("client run"); + server.abort(); + let _ = server.await; } #[tokio::test] @@ -3710,4 +6991,210 @@ mod tests { 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 + .next_user_message_id(&acp_id) + .expect("accept prompt"); + 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"); + } } + +#[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 new file mode 100644 index 0000000..2920b0a --- /dev/null +++ b/crates/agentkit-acp/src/v2/replace_tests.rs @@ -0,0 +1,468 @@ +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, + validate_inject_content(&pending.content).unwrap().1 + ); + assert_eq!(pending.bytes.media, 0); + 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) + .fold(InjectionBytes::default(), InjectionBytes::saturating_add) + ); + } + 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, InjectionBytes::default()); + 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 = 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, + InjectionBytes { + content: MAX_PENDING_INJECTION_BYTES, + media: 0 + } + ); + + 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, + 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); + 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 = 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.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, + InjectionBytes::default() + ); +} + +#[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, InjectionBytes::default()); + 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, + InjectionBytes::default() + ); + } +} + +#[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, 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(|| { + 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, InjectionBytes::default()); + 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, InjectionBytes::default()); + } + + 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, + InjectionBytes::default() + ); +} + +#[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); +} 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/crates/agentkit-loop/src/lib.rs b/crates/agentkit-loop/src/lib.rs index ec8aef7..33f7c43 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, @@ -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"; @@ -2398,6 +2400,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 +2415,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.reset(); + } + self.emit(AgentEvent::ResponseAttemptSuperseded); + } if cancellation .as_ref() .is_some_and(TurnCancellation::is_cancelled) @@ -2423,22 +2440,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); + } + 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 +2732,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, @@ -2915,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 { @@ -2990,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, @@ -3005,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() @@ -3026,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 { @@ -4141,20 +4215,501 @@ fn main() { "#; } +#[derive(Default)] +struct StreamedAssistantContent { + parts: Vec, + retained_bytes: usize, + overflowed: bool, +} + +enum StreamedPart { + Open { + id: PartId, + kind: PartKind, + text: String, + bytes: Vec, + structured: Option, + structured_bytes: usize, + metadata: MetadataMap, + metadata_bytes: usize, + retained_bytes: usize, + }, + Committed { + part: Part, + retained_bytes: usize, + }, +} + +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(index) = index { + self.parts[index] = open; + } else { + self.parts.push(open); + } + } + Delta::AppendText { part_id, chunk } => { + 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 } => { + 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 } => { + 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 } => { + 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), + } + } + + 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 self.overflowed { + return; + } + + 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)); + 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, + }; + 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(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 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 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() + .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 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) { + ( + 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, + }, + } + } +} + +#[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, + 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 +4952,10 @@ mod tests { struct FakeAdapter; struct SupersedingAdapter; + struct InterruptedStreamAdapter { + controller: CancellationController, + scenario: InterruptedStreamScenario, + } struct SlowAdapter; struct RecordingAdapter { seen_descriptions: StdArc>>>, @@ -4409,6 +4968,10 @@ mod tests { struct SupersedingSession { supersession_enabled: bool, } + struct InterruptedStreamSession { + controller: CancellationController, + scenario: InterruptedStreamScenario, + } struct SlowSession; struct RecordingSession { seen_descriptions: StdArc>>>, @@ -4425,6 +4988,19 @@ mod tests { events: VecDeque, } + #[derive(Clone, Copy)] + enum InterruptedStreamScenario { + PreserveContent, + CancelOnSupersession, + OverflowThenSupersession, + } + + struct InterruptedStreamTurn { + controller: CancellationController, + events: VecDeque, + cancel_on_supersession: bool, + } + struct SlowTurn { emitted: bool, } @@ -4669,6 +5245,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 +5328,108 @@ 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, + ]), + 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(), + events, + cancel_on_supersession: matches!( + self.scenario, + InterruptedStreamScenario::CancelOnSupersession + ), + }) + } + } + #[async_trait] impl ModelSession for FakeSession { type Turn = FakeTurn; @@ -5059,6 +5749,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( @@ -5432,97 +6141,447 @@ mod tests { mutator: "keep-recent".into(), point: ctx.point, }); - cursor.drain(..drop); - ctx.emitter.emit(AgentEvent::MutationFinished { - session_id: ctx.session_id.clone(), - turn_id: ctx.turn_id.cloned(), - mutator: "keep-recent".into(), - dirty: true, - metadata: MetadataMap::new(), + cursor.drain(..drop); + ctx.emitter.emit(AgentEvent::MutationFinished { + session_id: ctx.session_id.clone(), + turn_id: ctx.turn_id.cloned(), + mutator: "keep-recent".into(), + dirty: true, + metadata: MetadataMap::new(), + }); + Ok(()) + } + } + + /// No-op mutator that records the [`MutationPoint`] it is invoked with at + /// each mutation site, so a test can assert which point the loop reports. + struct PointRecordingMutator { + points: StdArc>>, + } + + #[async_trait] + impl LoopMutator for PointRecordingMutator { + async fn mutate( + &self, + _cursor: &mut TranscriptCursor<'_>, + ctx: LoopCtx<'_>, + ) -> Result<(), LoopError> { + self.points.lock().unwrap().push(ctx.point); + Ok(()) + } + } + + struct RecordingObserver { + events: StdArc>>, + } + + impl LoopObserver for RecordingObserver { + fn handle_event(&self, event: ObservedEvent) { + let event = event.event; + self.events.lock().unwrap().push(event); + } + } + + #[test] + fn session_consumer_capabilities_are_typed_and_serde_defaulted() { + let config = SessionConfig::new("session").with_response_attempt_supersession(); + assert!(config.consumer_capabilities.response_attempt_supersession); + + let decoded: SessionConfig = serde_json::from_value(json!({ + "session_id": "session", + "metadata": {}, + "cache": null + })) + .unwrap(); + assert_eq!( + decoded.consumer_capabilities, + SessionConsumerCapabilities::default() + ); + } + + #[tokio::test] + async fn response_attempt_supersession_is_forwarded_and_resets_attempt_state() { + let events = StdArc::new(StdMutex::new(Vec::new())); + let agent = Agent::builder() + .model(SupersedingAdapter) + .observer(RecordingObserver { + events: events.clone(), + }) + .build() + .unwrap(); + let mut driver = agent + .start(SessionConfig::new("supersession-session").with_response_attempt_supersession()) + .await + .unwrap(); + driver + .submit_input(vec![Item::text(ItemKind::User, "hello")]) + .unwrap(); + + let LoopStep::Finished(result) = run_until_finished(&mut driver).await else { + panic!("turn did not finish"); + }; + assert!(result.usage.is_none()); + + let events = events.lock().unwrap(); + let tool_call = events + .iter() + .position(|event| matches!(event, AgentEvent::ToolCallRequested(_))) + .unwrap(); + let superseded = events + .iter() + .position(|event| matches!(event, AgentEvent::ResponseAttemptSuperseded)) + .unwrap(); + 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(), }); - Ok(()) } + 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")) + ); } - /// No-op mutator that records the [`MutationPoint`] it is invoked with at - /// each mutation site, so a test can assert which point the loop reports. - struct PointRecordingMutator { - points: StdArc>>, + #[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") + )); } - #[async_trait] - impl LoopMutator for PointRecordingMutator { - async fn mutate( - &self, - _cursor: &mut TranscriptCursor<'_>, - ctx: LoopCtx<'_>, - ) -> Result<(), LoopError> { - self.points.lock().unwrap().push(ctx.point); - Ok(()) + #[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 } - } - struct RecordingObserver { - events: StdArc>>, + 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); } - impl LoopObserver for RecordingObserver { - fn handle_event(&self, event: ObservedEvent) { - let event = event.event; - self.events.lock().unwrap().push(event); + #[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 session_consumer_capabilities_are_typed_and_serde_defaulted() { - let config = SessionConfig::new("session").with_response_attempt_supersession(); - assert!(config.consumer_capabilities.response_attempt_supersession); + 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, + }); - let decoded: SessionConfig = serde_json::from_value(json!({ - "session_id": "session", - "metadata": {}, - "cache": null - })) - .unwrap(); - assert_eq!( - decoded.consumer_capabilities, - SessionConsumerCapabilities::default() - ); + 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 response_attempt_supersession_is_forwarded_and_resets_attempt_state() { + 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(SupersedingAdapter) + .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("supersession-session").with_response_attempt_supersession()) + .start(SessionConfig::new("supersede-then-cancel").with_response_attempt_supersession()) .await .unwrap(); driver - .submit_input(vec![Item::text(ItemKind::User, "hello")]) + .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!(result.usage.is_none()); + 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 tool_call = events - .iter() - .position(|event| matches!(event, AgentEvent::ToolCallRequested(_))) - .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); + } + + #[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(); - assert!(tool_call < superseded); + 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( @@ -8593,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()); diff --git a/docs/acp.md b/docs/acp.md index bfa4c3e..0bff66d 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 `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 @@ -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 SDK 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,7 @@ Recommended initial crate: ```toml [dependencies] -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" } @@ -76,6 +82,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", +] ``` The umbrella crate should later add: @@ -727,6 +737,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) 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" }