From be95739e987ae366b0b4bed20ea760cf2a1bb50d Mon Sep 17 00:00:00 2001 From: daniel Date: Tue, 25 Aug 2026 17:41:48 +0100 Subject: [PATCH 1/4] feat(acp): add v2 session injection --- Cargo.lock | 37 +- Cargo.toml | 4 + book/src/acp.md | 34 +- crates/agentkit-acp/Cargo.toml | 7 +- crates/agentkit-acp/README.md | 32 +- crates/agentkit-acp/src/v2.rs | 2214 ++++++++++++++++++++++++++++++-- docs/acp.md | 28 +- 7 files changed, 2208 insertions(+), 148 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index edfe1fe..d4fb8c9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11,8 +11,7 @@ checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" [[package]] name = "agent-client-protocol" version = "2.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6d87bc7769eba641753ba5dc52f73ec3765d51022c6753bf040967125ddc86a8" +source = "git+https://github.com/danielkov/rust-sdk?rev=2f039993d1d6ed8da35b38c31f54a7cbb7338c70#2f039993d1d6ed8da35b38c31f54a7cbb7338c70" dependencies = [ "agent-client-protocol-derive", "agent-client-protocol-schema", @@ -35,8 +34,7 @@ dependencies = [ [[package]] name = "agent-client-protocol-derive" version = "2.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3abd4080f51e4f24f5042beb7fb7a66ede29a2dc1c2582c329532e1c27264ddc" +source = "git+https://github.com/danielkov/rust-sdk?rev=2f039993d1d6ed8da35b38c31f54a7cbb7338c70#2f039993d1d6ed8da35b38c31f54a7cbb7338c70" dependencies = [ "quote", "syn 3.0.3", @@ -44,13 +42,11 @@ dependencies = [ [[package]] name = "agent-client-protocol-schema" -version = "1.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d5c231915b4ab578c722eca2d1bd7df4d300bfd6cac3b8e9f0d1e3ddc95b187c" +version = "1.7.0" +source = "git+https://github.com/danielkov/agent-client-protocol?rev=6e7e044f9464c4fd652d90699a09e9edc8b3bbad#6e7e044f9464c4fd652d90699a09e9edc8b3bbad" dependencies = [ "anyhow", "derive_more", - "diffy", "schemars 1.2.2", "serde", "serde_json", @@ -92,7 +88,7 @@ dependencies = [ [[package]] name = "agentkit-acp" -version = "0.10.8" +version = "0.10.9" dependencies = [ "agent-client-protocol", "agentkit-core", @@ -101,6 +97,7 @@ dependencies = [ "agentkit-tools-core", "async-trait", "base64 0.22.1", + "futures-util", "serde_json", "thiserror 2.0.18", "tokio", @@ -1228,15 +1225,6 @@ version = "0.1.13" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "56254986775e3233ffa9c4d7d3faaf6d36a2c09d30b20687e9f88bc8bafc16c8" -[[package]] -name = "diffy" -version = "0.5.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "10aec8f7f9393bd6a4f2762be0ceb012d3cbe2478987258cc9960de148561914" -dependencies = [ - "hashbrown 0.17.1", -] - [[package]] name = "digest" version = "0.10.7" @@ -1407,12 +1395,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" @@ -1619,9 +1601,6 @@ name = "hashbrown" version = "0.17.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" -dependencies = [ - "foldhash", -] [[package]] name = "hashlink" @@ -2301,7 +2280,7 @@ version = "5.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "51e219e79014df21a225b1860a479e2dcd7cbd9130f4defd4bd0e191ea31d67d" dependencies = [ - "base64 0.22.1", + "base64 0.21.7", "chrono", "getrandom 0.2.17", "http", @@ -3364,7 +3343,7 @@ checksum = "f83ad47c2f14654528a89495f8d0dbc64173176f8512c7c72386cbe81009f661" dependencies = [ "ahash", "annotate-snippets", - "base64 0.22.1", + "base64 0.21.7", "encoding_rs_io", "getrandom 0.3.4", "nohash-hasher", diff --git a/Cargo.toml b/Cargo.toml index f932c92..2e06d21 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -60,6 +60,10 @@ repository = "https://github.com/danielkov/agentkit" rust-version = "1.92" version = "0.10.5" +# Intentionally unpublishable until the upstream SDK releases session injection. +[patch.crates-io] +agent-client-protocol = { git = "https://github.com/danielkov/rust-sdk", rev = "2f039993d1d6ed8da35b38c31f54a7cbb7338c70" } + [workspace.dependencies] async-trait = "0.1.89" dotenvy = "0.15.7" diff --git a/book/src/acp.md b/book/src/acp.md index 02d0769..a5db478 100644 --- a/book/src/acp.md +++ b/book/src/acp.md @@ -2,27 +2,28 @@ The [Agent Client Protocol (ACP)](https://agentclientprotocol.com) standardizes communication between clients (code editors, IDEs, desktop apps) and coding agents. Where MCP connects an agent to external tools, ACP connects a client to the agent itself: session lifecycle, prompt turns, streamed updates, tool call reporting, and permission prompts all travel over JSON-RPC — usually with the agent running as an editor child process on stdio. This chapter covers [`agentkit-acp`](https://github.com/danielkov/agentkit/tree/main/crates/agentkit-acp): how an agentkit host becomes ACP-addressable, and how a standalone agent serves ACP directly. -## Built on the official SDK +## Built on the ACP Rust SDK -Like `agentkit-mcp`, this crate does not define a parallel protocol vocabulary. It builds on the official Rust SDK, [`agent-client-protocol`](https://crates.io/crates/agent-client-protocol), and re-exports the stable v1 wire types (`SessionId`, `ContentBlock`, `SessionUpdate`, `ToolCallUpdate`, `StopReason`, …) at the crate root and under `agentkit_acp::wire`. The full upstream SDK is available as `agentkit_acp::sdk`. Agentkit owns only the host-facing glue: session binding, observer routing, prompt conversion, cancellation handles, and approval resolution. +Like `agentkit-mcp`, this crate does not define a parallel protocol vocabulary. It currently builds on a commit-pinned [`agent-client-protocol` fork](https://github.com/danielkov/rust-sdk), needed for experimental session injection support, and re-exports the stable v1 wire types (`SessionId`, `ContentBlock`, `SessionUpdate`, `ToolCallUpdate`, `StopReason`, …) at the crate root and under `agentkit_acp::wire`. The full SDK is available as `agentkit_acp::sdk`. Agentkit owns only the host-facing glue: session binding, observer routing, prompt conversion, cancellation handles, and approval resolution. - **Protocol docs:** [agentclientprotocol.com](https://agentclientprotocol.com/protocol/v1/overview) -- **Rust SDK:** [`agent-client-protocol` on crates.io](https://crates.io/crates/agent-client-protocol) +- **Rust SDK fork:** [`danielkov/rust-sdk`](https://github.com/danielkov/rust-sdk), pinned in `Cargo.toml` and `Cargo.lock` ## Opt-in ACP v2 runtime ACP v2 support is additive and disabled by default. Enable it explicitly: ```toml -agentkit-acp = { version = "0.10.8", features = ["protocol-v2"] } +agentkit-acp = { version = "0.10.9", 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,20 @@ 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. Revoke is serialized with acknowledged user-message +forwarding. Queueing and replacement are not supported. Cancellation and close +discard pending injected messages and never carry them into a later prompt; +revoke returns `already_delivered` when acknowledged delivery won first. + The initial v2 foundation routes text, reasoning, and tool lifecycle updates. -ACP v2 permission callbacks are intentionally deferred; an unsupported approval -interrupt retains the transcript and ends the prompt with the custom `_error` -stop reason rather than `refusal`. Upstream labels the v2 protocol unstable, so -opt-in callers should expect the `v2` namespace to track official SDK changes. +ACP v2 permission callbacks are intentionally deferred; unsupported approval +requests are denied while accepted steers remain pending for the next safe +boundary. The SDK labels v2 unstable, so opt-in callers should expect the `v2` +namespace to track the pinned fork. The workspace patch is intentionally +unpublishable until the required APIs are released upstream. ## Two integration shapes @@ -228,4 +238,4 @@ The `agentkit-acp` crate itself has a default `stdio` feature that gates `serve_ > **Example:** [`openrouter-acp-trio`](https://github.com/danielkov/agentkit/tree/main/examples/openrouter-acp-trio) runs three OpenRouter-backed agents (orchestrator, worker, reviewer) that call each other over in-memory ACP endpoints while a REPL drives the orchestrator through a persistent ACP session — session binding, streamed updates, tool call reporting, and agent-to-agent handoffs in one program. > -> **Crate:** [`agentkit-acp`](https://github.com/danielkov/agentkit/tree/main/crates/agentkit-acp) — depends on [`agentkit-core`](https://github.com/danielkov/agentkit/tree/main/crates/agentkit-core), [`agentkit-loop`](https://github.com/danielkov/agentkit/tree/main/crates/agentkit-loop), [`agentkit-tools-core`](https://github.com/danielkov/agentkit/tree/main/crates/agentkit-tools-core), and [`agent-client-protocol`](https://crates.io/crates/agent-client-protocol). Design notes: [`docs/acp.md`](https://github.com/danielkov/agentkit/blob/main/docs/acp.md). +> **Crate:** [`agentkit-acp`](https://github.com/danielkov/agentkit/tree/main/crates/agentkit-acp) — depends on [`agentkit-core`](https://github.com/danielkov/agentkit/tree/main/crates/agentkit-core), [`agentkit-loop`](https://github.com/danielkov/agentkit/tree/main/crates/agentkit-loop), [`agentkit-tools-core`](https://github.com/danielkov/agentkit/tree/main/crates/agentkit-tools-core), and a commit-pinned [`agent-client-protocol` fork](https://github.com/danielkov/rust-sdk). Design notes: [`docs/acp.md`](https://github.com/danielkov/agentkit/blob/main/docs/acp.md). diff --git a/crates/agentkit-acp/Cargo.toml b/crates/agentkit-acp/Cargo.toml index 5ee9288..ced0f9c 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.8" +version = "0.10.9" edition.workspace = true license.workspace = true rust-version.workspace = true @@ -14,6 +14,10 @@ default = ["stdio"] stdio = [] unstable-acp = ["agent-client-protocol/unstable"] protocol-v2 = ["agent-client-protocol/unstable_protocol_v2"] +unstable-inject = [ + "protocol-v2", + "agent-client-protocol/unstable_session_inject", +] [dependencies] agent-client-protocol = "=2.0.0" @@ -28,5 +32,6 @@ tokio = { workspace = true, features = ["sync", "time"] } tracing.workspace = true [dev-dependencies] +futures-util.workspace = true agentkit-integration-tests = { path = "../agentkit-integration-tests" } 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..1df9298 100644 --- a/crates/agentkit-acp/README.md +++ b/crates/agentkit-acp/README.md @@ -2,7 +2,7 @@ Agent Client Protocol integration for agentkit hosts. -This crate re-exports upstream ACP wire types from `agent-client-protocol` and +This crate re-exports ACP wire types from a pinned `agent-client-protocol` fork and adds only agentkit-specific glue: session binding, observer routing, prompt conversion, cancellation handles, and approval resolver abstractions. @@ -21,13 +21,14 @@ 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.9", features = ["protocol-v2"] } ``` -The feature maps directly to the official -`agent-client-protocol/unstable_protocol_v2` feature. V2 APIs and official v2 -wire types live only under `agentkit_acp::v2` (and -`agentkit_acp::v2::wire`). +The feature maps directly to the pinned fork's +`agent-client-protocol/unstable_protocol_v2` feature. V2 APIs and v2 wire +types live only under `agentkit_acp::v2` (and +`agentkit_acp::v2::wire`). Enable `unstable-inject` instead to add the unstable +ACP v2 session-injection surface; it implies `protocol-v2`. `v2::AcpHeadlessRuntime` supports ACP v2 initialize, new/list/resume session, prompt, cancel, session updates, and close. Listing and resume cover @@ -38,11 +39,24 @@ 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. +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. Batch requests use the same +receipt-backed acceptance path. `session/revoke_inject` is mandatory and is +serialized with acknowledged `UserMessage` forwarding. Queueing, stream +interruption, and replacement are not supported. Cancelling or closing a +session drops every still-pending injected message; no matching `UserMessage` +is emitted and it does not carry into the next prompt. If acknowledged delivery +wins the race first, revoke returns `already_delivered`. + 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; already accepted steers remain +pending and are delivered at the next safe boundary. Because the SDK marks +protocol v2 unstable, all APIs in the `v2` namespace can evolve with the pinned +SDK fork. The workspace patch is intentionally unpublishable until these SDK +APIs are available in an upstream release. Run the stable v1 in-memory end-to-end example with: diff --git a/crates/agentkit-acp/src/v2.rs b/crates/agentkit-acp/src/v2.rs index ace92cb..5bde342 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, @@ -21,19 +23,33 @@ 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::*; @@ -41,6 +57,10 @@ pub mod wire { enum ClientMessage { Update(Box), + AcknowledgedUpdate { + notification: Box, + acknowledged: oneshot::Sender>, + }, Flush(oneshot::Sender<()>), } @@ -67,6 +87,21 @@ impl ClientHandle { .map_err(|_| AcpRuntimeError::ClientClosed) } + async fn update_acknowledged( + &self, + session_id: wire::SessionId, + update: wire::SessionUpdate, + ) -> Result<(), AcpRuntimeError> { + let (tx, rx) = oneshot::channel(); + self.tx + .send(ClientMessage::AcknowledgedUpdate { + notification: Box::new(wire::UpdateSessionNotification::new(session_id, update)), + 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 @@ -78,7 +113,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 { @@ -88,6 +123,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(()); } @@ -189,18 +237,27 @@ impl AcpIntegration { .ok_or_else(|| AcpRuntimeError::SessionNotFound(session_id.to_string())) } - fn begin_prompt( + 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); - finish_model_message(&session); Ok(wire::MessageId::new(format!( "{session_id}-user-{sequence}" ))) } + 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) + } + fn finish_prompt(&self, session_id: &wire::SessionId) { if let Ok(session) = self.session(session_id) { finish_model_message(&session); @@ -419,7 +476,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, @@ -429,7 +486,7 @@ where .ok_or(AcpRuntimeError::MissingField("agent_factory"))?; let state = Arc::new(RuntimeState::new(factory, 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( @@ -508,7 +565,58 @@ where } }, agent_client_protocol::on_receive_request!(), + ); + #[cfg(feature = "unstable-inject")] + let agent = agent + .on_receive_request( + { + let state = Arc::clone(&state); + async move |request: wire::InjectSessionRequest, responder, cx| { + let cancellation = responder.cancellation(); + // Let a cancellation frame already queued behind this request + // linearize before AgentKit reserves the message. + tokio::task::yield_now().await; + if cancellation.is_cancelled() { + return responder.respond_with_result(Err( + agent_client_protocol::Error::request_cancelled(), + )); + } + match state.inject(request).await { + Ok(acceptance) => { + let receipt = match responder.respond_tracked(acceptance.response()) + { + Ok(receipt) => receipt, + Err(error) => { + acceptance.discard(); + return Err(error); + } + }; + cx.spawn(async move { + if receipt.await.is_ok() { + acceptance.activate(&cancellation); + } else { + acceptance.discard(); + } + Ok(()) + })?; + Ok(()) + } + Err(error) => responder.respond_with_result(Err(error)), + } + } + }, + agent_client_protocol::on_receive_request!(), ) + .on_receive_request( + { + let state = Arc::clone(&state); + async move |request: wire::RevokeInjectSessionRequest, responder, _cx| { + responder.respond_with_result(state.revoke_inject(request).await) + } + }, + agent_client_protocol::on_receive_request!(), + ); + let connection = agent .on_receive_notification( { let state = Arc::clone(&state); @@ -554,6 +662,257 @@ where } } +#[cfg(feature = "unstable-inject")] +const MAX_PENDING_INJECTIONS: usize = 64; +#[cfg(feature = "unstable-inject")] +const MAX_PENDING_INJECTION_BYTES: usize = 256 * 1024; +#[cfg(feature = "unstable-inject")] +const MAX_DELIVERED_TOMBSTONES: usize = 64; + +#[cfg(feature = "unstable-inject")] +struct PendingInject { + message_id: wire::MessageId, + content: Vec, + items: Vec, + bytes: usize, + accepted: bool, +} + +#[cfg(feature = "unstable-inject")] +#[derive(Default)] +struct InjectionState { + running: bool, + cancelled: bool, + at_boundary: bool, + pending: VecDeque, + delivering: Option, + pending_bytes: usize, + delivered: VecDeque, +} + +#[cfg(feature = "unstable-inject")] +#[derive(Default)] +struct InjectionController { + state: Mutex, + changed: Notify, +} + +#[cfg(feature = "unstable-inject")] +enum BoundaryAction { + Wait, + Deliver(PendingInject), + Complete(InjectBoundary), +} + +#[cfg(feature = "unstable-inject")] +enum RevokeTransition { + Revoked, + WaitForDelivery, + AlreadyDelivered, + Unknown, +} + +#[cfg(feature = "unstable-inject")] +impl InjectionController { + fn start_turn(&self) { + let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + state.running = !state.cancelled; + state.at_boundary = false; + state.pending.clear(); + state.delivering = None; + state.pending_bytes = 0; + state.delivered.clear(); + } + + fn reset_cancellation(&self) { + self.state + .lock() + .unwrap_or_else(|error| error.into_inner()) + .cancelled = false; + } + + fn reserve(&self, pending: PendingInject) -> Result<(), agent_client_protocol::Error> { + let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + if !state.running { + return Err(sdk_v2_error(wire::Error::inject_no_running_turn())); + } + if state.pending.len() + usize::from(state.delivering.is_some()) >= MAX_PENDING_INJECTIONS + || state.pending_bytes.saturating_add(pending.bytes) > MAX_PENDING_INJECTION_BYTES + { + return Err(agent_client_protocol::Error::new( + -32602, + "pending session injection budget exceeded", + )); + } + state.pending_bytes += pending.bytes; + state.pending.push_back(pending); + Ok(()) + } + + fn activate( + &self, + message_id: &wire::MessageId, + cancellation: &agent_client_protocol::RequestCancellation, + ) { + let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + if cancellation.is_cancelled() || !state.running { + Self::remove_pending(&mut state, message_id); + } else if let Some(pending) = state + .pending + .iter_mut() + .find(|pending| &pending.message_id == message_id) + { + pending.accepted = true; + } + drop(state); + self.changed.notify_waiters(); + } + + fn discard(&self, message_id: &wire::MessageId) { + let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + Self::remove_pending(&mut state, message_id); + drop(state); + self.changed.notify_waiters(); + } + + fn remove_pending(state: &mut InjectionState, message_id: &wire::MessageId) -> bool { + let Some(index) = state + .pending + .iter() + .position(|pending| &pending.message_id == message_id) + else { + return false; + }; + if let Some(pending) = state.pending.remove(index) { + state.pending_bytes = state.pending_bytes.saturating_sub(pending.bytes); + } + true + } + + fn boundary_action(&self, terminal: bool, delivered_any: bool) -> BoundaryAction { + let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + state.at_boundary = true; + if !state.running { + state.at_boundary = false; + return BoundaryAction::Complete(InjectBoundary::Stopped); + } + if state + .pending + .front() + .is_some_and(|pending| !pending.accepted) + { + 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(InjectBoundary::Delivered) + } else if terminal { + state.running = false; + BoundaryAction::Complete(InjectBoundary::Finished) + } else { + BoundaryAction::Complete(InjectBoundary::Continue) + } + } + + fn finish_delivery(&self, pending: &PendingInject, delivered: bool) { + let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + debug_assert_eq!(state.delivering.as_ref(), Some(&pending.message_id)); + state.delivering = None; + state.pending_bytes = state.pending_bytes.saturating_sub(pending.bytes); + if delivered { + if state.delivered.len() == MAX_DELIVERED_TOMBSTONES { + state.delivered.pop_front(); + } + state.delivered.push_back(pending.message_id.clone()); + } + drop(state); + self.changed.notify_waiters(); + } + + fn revoke_transition(&self, message_id: &wire::MessageId) -> RevokeTransition { + let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + if Self::remove_pending(&mut state, message_id) { + drop(state); + self.changed.notify_waiters(); + RevokeTransition::Revoked + } else if state.delivering.as_ref() == Some(message_id) { + RevokeTransition::WaitForDelivery + } else if state.delivered.contains(message_id) { + RevokeTransition::AlreadyDelivered + } else { + RevokeTransition::Unknown + } + } + + fn cancel_turn(&self) { + let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + state.running = false; + state.cancelled = true; + state.at_boundary = false; + let queued_bytes = state + .pending + .iter() + .map(|pending| pending.bytes) + .sum::(); + state.pending.clear(); + state.pending_bytes = state.pending_bytes.saturating_sub(queued_bytes); + drop(state); + self.changed.notify_waiters(); + } + + fn stop_turn(&self) { + let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + state.running = false; + state.at_boundary = false; + state.pending.clear(); + state.delivering = None; + state.pending_bytes = 0; + drop(state); + self.changed.notify_waiters(); + } +} + +#[cfg(feature = "unstable-inject")] +struct InjectAcceptance { + entry: Arc, + message_id: Option, +} + +#[cfg(feature = "unstable-inject")] +impl InjectAcceptance { + fn response(&self) -> wire::InjectSessionResponse { + wire::InjectSessionResponse::new( + self.message_id + .as_ref() + .expect("acceptance has a message id") + .clone(), + ) + } + + fn activate(mut self, cancellation: &agent_client_protocol::RequestCancellation) { + let message_id = self.message_id.take().expect("acceptance has a message id"); + self.entry.injection.activate(&message_id, cancellation); + } + + fn discard(mut self) { + let message_id = self.message_id.take().expect("acceptance has a message id"); + self.entry.injection.discard(&message_id); + } +} + +#[cfg(feature = "unstable-inject")] +impl Drop for InjectAcceptance { + fn drop(&mut self) { + if let Some(message_id) = self.message_id.take() { + self.entry.injection.discard(&message_id); + } + } +} + struct SessionEntry { commands: mpsc::UnboundedSender, cancellation: CancellationController, @@ -561,6 +920,8 @@ struct SessionEntry { busy: Arc, closed: AtomicBool, lifecycle: Mutex<()>, + #[cfg(feature = "unstable-inject")] + injection: Arc, task: Mutex>>, drain_task: Mutex>>, } @@ -622,7 +983,7 @@ where async fn new_session( self: &Arc, request: wire::NewSessionRequest, - cx: ConnectionTo, + cx: V2ConnectionTo, ) -> Result { let sequence = self.next_session.fetch_add(1, Ordering::Relaxed); let acp_session_id = wire::SessionId::new(format!("session-{sequence}")); @@ -675,6 +1036,10 @@ where let integration = Arc::clone(&self.integration); let worker_session_id = acp_session_id.clone(); let worker_cancellation = cancellation.handle(); + #[cfg(feature = "unstable-inject")] + let injection = Arc::new(InjectionController::default()); + #[cfg(feature = "unstable-inject")] + let worker_injection = Arc::clone(&injection); let task = tokio::spawn(async move { session_worker( worker_session_id, @@ -683,6 +1048,8 @@ where integration, worker_cancellation, worker_busy, + #[cfg(feature = "unstable-inject")] + worker_injection, rx, ) .await; @@ -694,6 +1061,8 @@ where busy, closed: AtomicBool::new(false), lifecycle: Mutex::new(()), + #[cfg(feature = "unstable-inject")] + injection, task: Mutex::new(Some(task)), drain_task: Mutex::new(Some(drain_task)), }); @@ -795,6 +1164,8 @@ where )); } let cancellation_generation = entry.cancellation.handle().generation(); + #[cfg(feature = "unstable-inject")] + entry.injection.reset_cancellation(); if entry .commands .send(SessionCommand::Prompt { @@ -812,6 +1183,96 @@ where rx.await.map_err(|_| AcpRuntimeError::ClientClosed)? } + #[cfg(feature = "unstable-inject")] + async fn inject( + &self, + request: wire::InjectSessionRequest, + ) -> Result { + if !matches!(request.mode, wire::SessionInjectMode::Steer) { + return Err(agent_client_protocol::Error::new( + -32602, + "unsupported session injection mode", + )); + } + let items = content_to_items(&request.content).map_err(crate::sdk_error)?; + let bytes = serde_json::to_vec(&request.content) + .map_err(|error| agent_client_protocol::Error::new(-32603, error.to_string()))? + .len(); + let entry = self + .sessions + .lock() + .await + .get(&request.session_id) + .cloned() + .ok_or_else(|| session_not_found_error(&request.session_id))?; + let _lifecycle = entry + .lifecycle + .lock() + .unwrap_or_else(|error| error.into_inner()); + if entry.closed.load(Ordering::Acquire) { + return Err(session_not_found_error(&request.session_id)); + } + let message_id = self + .integration + .next_user_message_id(&request.session_id) + .map_err(crate::sdk_error)?; + entry.injection.reserve(PendingInject { + message_id: message_id.clone(), + content: request.content, + items, + bytes, + accepted: false, + })?; + drop(_lifecycle); + Ok(InjectAcceptance { + entry, + message_id: Some(message_id), + }) + } + + #[cfg(feature = "unstable-inject")] + async fn revoke_inject( + &self, + request: wire::RevokeInjectSessionRequest, + ) -> Result { + let entry = self + .sessions + .lock() + .await + .get(&request.session_id) + .cloned() + .ok_or_else(|| session_not_found_error(&request.session_id))?; + loop { + let changed = entry.injection.changed.notified(); + tokio::pin!(changed); + changed.as_mut().enable(); + let _lifecycle = entry + .lifecycle + .lock() + .unwrap_or_else(|error| error.into_inner()); + if entry.closed.load(Ordering::Acquire) { + return Err(session_not_found_error(&request.session_id)); + } + match entry.injection.revoke_transition(&request.message_id) { + RevokeTransition::Revoked => { + return Ok(wire::RevokeInjectSessionResponse::new()); + } + RevokeTransition::WaitForDelivery => drop(_lifecycle), + RevokeTransition::AlreadyDelivered => { + return Err(sdk_v2_error(wire::Error::inject_already_delivered( + request.message_id, + ))); + } + RevokeTransition::Unknown => { + return Err(sdk_v2_error(wire::Error::inject_unknown_message_id( + request.message_id, + ))); + } + } + changed.await; + } + } + async fn cancel( &self, notification: wire::CancelSessionNotification, @@ -828,6 +1289,8 @@ where .lock() .unwrap_or_else(|error| error.into_inner()); if !entry.closed.load(Ordering::Acquire) && entry.busy.load(Ordering::Acquire) { + #[cfg(feature = "unstable-inject")] + cancel_pending_injections(&entry); entry.cancellation.interrupt(); } } @@ -881,10 +1344,17 @@ fn signal_session_stop(entry: &Arc) { .lock() .unwrap_or_else(|error| error.into_inner()); entry.closed.store(true, Ordering::Release); + #[cfg(feature = "unstable-inject")] + cancel_pending_injections(entry); entry.cancellation.interrupt(); let _ = entry.commands.send(SessionCommand::Shutdown); } +#[cfg(feature = "unstable-inject")] +fn cancel_pending_injections(entry: &SessionEntry) { + entry.injection.cancel_turn(); +} + fn take_task( task: &Mutex>>, ) -> Option> { @@ -928,6 +1398,7 @@ async fn stop_client(entry: Arc) { } } +#[allow(clippy::too_many_arguments)] // Feature-gated injection adds one state/notify pair. async fn session_worker( session_id: wire::SessionId, mut driver: agentkit_loop::LoopDriver, @@ -935,6 +1406,7 @@ async fn session_worker( integration: Arc, cancellation: CancellationHandle, busy: Arc, + #[cfg(feature = "unstable-inject")] injection: Arc, mut commands: mpsc::UnboundedReceiver, ) where S: ModelSession + Send + 'static, @@ -965,8 +1437,12 @@ async fn session_worker( continue; } }; + #[cfg(feature = "unstable-inject")] + injection.start_turn(); let (start_tx, start_rx) = oneshot::channel(); if response.send(Ok(start_tx)).is_err() || start_rx.await.is_err() { + #[cfg(feature = "unstable-inject")] + injection.stop_turn(); integration.finish_prompt(&session_id); busy.store(false, Ordering::Release); continue; @@ -989,12 +1465,25 @@ async fn session_worker( }) .is_err() { + #[cfg(feature = "unstable-inject")] + injection.stop_turn(); integration.finish_prompt(&session_id); busy.store(false, Ordering::Release); continue; } - let stop_reason = drive_prompt(&mut driver, &cancellation, cancellation_generation).await; + let stop_reason = drive_prompt( + &mut driver, + &cancellation, + cancellation_generation, + #[cfg(feature = "unstable-inject")] + &client, + #[cfg(feature = "unstable-inject")] + &session_id, + #[cfg(feature = "unstable-inject")] + &injection, + ) + .await; if let Err(error) = client.flush().await { tracing::debug!(%error, "failed to flush ACP v2 output"); } @@ -1006,6 +1495,61 @@ async fn session_worker( wire::IdleStateUpdate::new().stop_reason(stop_reason), )), ); + #[cfg(feature = "unstable-inject")] + injection.stop_turn(); + } +} + +#[cfg(feature = "unstable-inject")] +#[derive(Debug, Eq, PartialEq)] +enum InjectBoundary { + Continue, + Delivered, + Finished, + Stopped, +} + +#[cfg(feature = "unstable-inject")] +async fn handle_inject_boundary( + driver: &mut agentkit_loop::LoopDriver, + client: &ClientHandle, + session_id: &wire::SessionId, + injection: &InjectionController, + terminal: bool, +) -> Result +where + S: ModelSession + Send + 'static, +{ + 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 = client + .update_acknowledged( + 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; + } + } } } @@ -1013,60 +1557,102 @@ async fn drive_prompt( driver: &mut agentkit_loop::LoopDriver, cancellation: &CancellationHandle, generation: u64, + #[cfg(feature = "unstable-inject")] client: &ClientHandle, + #[cfg(feature = "unstable-inject")] session_id: &wire::SessionId, + #[cfg(feature = "unstable-inject")] injection: &InjectionController, ) -> wire::StopReason where S: ModelSession + Send + 'static, { loop { - // Cancellation is installed on the driver and its model/tool work. Keep - // polling the driver so it can close interrupted tool calls and leave a - // resumable transcript before the session becomes idle. - match driver.next().await { - Ok(LoopStep::Finished(result)) => { - if result.finish_reason == FinishReason::ToolCall { - continue; - } + let step = match driver.next().await { + Ok(step) => step, + Err(error) => { + tracing::debug!(%error, "ACP v2 agent loop failed"); + #[cfg(feature = "unstable-inject")] + injection.stop_turn(); return if cancellation.is_cancelled_since(generation) { wire::StopReason::Cancelled } else { - finish_reason_to_stop_reason(&result.finish_reason) + error_stop_reason() }; } - Ok(LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_))) => { - return if cancellation.is_cancelled_since(generation) { - wire::StopReason::Cancelled - } else { - wire::StopReason::EndTurn - }; + }; + if cancellation.is_cancelled_since(generation) { + #[cfg(feature = "unstable-inject")] + injection.stop_turn(); + return wire::StopReason::Cancelled; + } + match step { + LoopStep::Finished(result) => { + if result.finish_reason == FinishReason::ToolCall { + continue; + } + #[cfg(feature = "unstable-inject")] + match handle_inject_boundary(driver, client, session_id, injection, true).await { + Ok(InjectBoundary::Delivered | InjectBoundary::Continue) => continue, + Ok(InjectBoundary::Stopped) => return wire::StopReason::Cancelled, + Ok(InjectBoundary::Finished) => { + return finish_reason_to_stop_reason(&result.finish_reason); + } + Err(error) => { + tracing::debug!(%error, "failed to deliver ACP v2 injected message"); + injection.stop_turn(); + return error_stop_reason(); + } + } + #[cfg(not(feature = "unstable-inject"))] + return finish_reason_to_stop_reason(&result.finish_reason); } - Ok(LoopStep::Interrupt(LoopInterrupt::AfterToolResult(_))) => continue, - Ok(LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(_))) => { - if let Err(error) = driver.cancel_pending_approvals().await { - tracing::debug!(%error, "failed to cancel unsupported ACP v2 approval"); + LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_)) => { + #[cfg(feature = "unstable-inject")] + match handle_inject_boundary(driver, client, session_id, injection, true).await { + Ok(InjectBoundary::Delivered | InjectBoundary::Continue) => continue, + Ok(InjectBoundary::Stopped) => return wire::StopReason::Cancelled, + Ok(InjectBoundary::Finished) => return wire::StopReason::EndTurn, + Err(error) => { + tracing::debug!(%error, "failed to deliver ACP v2 injected message"); + injection.stop_turn(); + return error_stop_reason(); + } } - return if cancellation.is_cancelled_since(generation) { - wire::StopReason::Cancelled - } else { - error_stop_reason() - }; + #[cfg(not(feature = "unstable-inject"))] + return wire::StopReason::EndTurn; } - Err(error) => { - tracing::debug!(%error, "ACP v2 agent loop failed"); - return if cancellation.is_cancelled_since(generation) { - wire::StopReason::Cancelled - } else { - error_stop_reason() - }; + LoopStep::Interrupt(LoopInterrupt::AfterToolResult(_)) => + { + #[cfg(feature = "unstable-inject")] + match handle_inject_boundary(driver, client, session_id, injection, false).await { + Ok(InjectBoundary::Stopped) => return wire::StopReason::Cancelled, + Err(error) => { + tracing::debug!(%error, "failed to deliver ACP v2 injected message"); + injection.stop_turn(); + return error_stop_reason(); + } + _ => {} + } + } + LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(_)) => { + if let Err(error) = driver.cancel_pending_approvals().await { + tracing::debug!(%error, "failed to resolve unsupported ACP v2 approval"); + #[cfg(feature = "unstable-inject")] + injection.stop_turn(); + return error_stop_reason(); + } } } } } fn prompt_to_items(request: &wire::PromptRequest) -> Result, AcpRuntimeError> { + content_to_items(&request.prompt) +} + +fn content_to_items(content: &[wire::ContentBlock]) -> Result, AcpRuntimeError> { let mut user_parts = Vec::new(); let mut context_items = Vec::new(); - for block in &request.prompt { + for block in content { match block { wire::ContentBlock::Text(text) => { user_parts.push(Part::Text(TextPart::new(text.text.clone()))); @@ -1404,14 +1990,18 @@ fn finish_reason_to_stop_reason(reason: &FinishReason) -> wire::StopReason { } fn headless_capabilities() -> wire::AgentCapabilities { - wire::AgentCapabilities::new().session( - wire::SessionCapabilities::new().prompt( - wire::PromptCapabilities::new() - .image(wire::PromptImageCapabilities::new()) - .audio(wire::PromptAudioCapabilities::new()) - .embedded_context(wire::PromptEmbeddedContextCapabilities::new()), - ), - ) + 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( + wire::SessionInjectCapabilities::new(vec![wire::SessionInjectMode::Steer]) + .steer_in_stream(vec![wire::SessionInjectSteerInStream::Finish]), + ); + wire::AgentCapabilities::new().session(session) } #[cfg(test)] @@ -1421,15 +2011,23 @@ mod tests { use std::time::Duration; 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_loop::{ Agent, LoopError, ModelSession, ModelTurn, ModelTurnEvent, ModelTurnResult, SessionConfig, TurnRequest, }; + #[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 _; #[derive(Clone)] struct TestFactory { @@ -1542,43 +2140,335 @@ mod tests { } } - 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(), - }), - ]) + #[cfg(feature = "unstable-inject")] + struct ApprovalPermissionRequest { + metadata: MetadataMap, } - 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(), + #[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())) + } + } + + 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(), }), ]) } @@ -1997,7 +2887,7 @@ mod tests { wire::SessionUpdate::AgentMessageChunk(chunk) => Some(chunk.message_id), _ => None, }, - ClientMessage::Flush(_) => None, + ClientMessage::AcknowledgedUpdate { .. } | ClientMessage::Flush(_) => None, }) .collect::>(); assert_eq!( @@ -2259,6 +3149,1152 @@ mod tests { } } + #[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)), + } + } + + async fn wait_for_second_start(&self) { + if self.starts.load(Ordering::Acquire) < 2 { + tokio::time::timeout(Duration::from_secs(2), self.second_started.notified()) + .await + .expect("batched sibling did not start"); + } + } + + fn release_second_start(&self) { + self.second_release.add_permits(1); + } + } + + #[cfg(feature = "unstable-inject")] + #[async_trait] + impl AcpAgentFactory for SecondStartGatedFactory { + async fn start( + &self, + ctx: AcpAgentFactoryContext, + ) -> Result, AcpRuntimeError> { + if self.starts.fetch_add(1, Ordering::AcqRel) == 1 { + self.second_started.notify_one(); + Arc::clone(&self.second_release) + .acquire_owned() + .await + .expect("second-start gate stays open") + .forget(); + } + Agent::builder() + .model(self.adapter.clone()) + .observer(ctx.integration.as_ref().clone()) + .cancellation(ctx.cancellation) + .build() + .map_err(|error| AcpRuntimeError::Loop(error.to_string()))? + .start(SessionConfig::new(ctx.agentkit_session_id).with_metadata(ctx.metadata)) + .await + .map_err(|error| AcpRuntimeError::Loop(error.to_string())) + } + } + + #[cfg(feature = "unstable-inject")] + async fn raw_request( + channel: &mut Channel, + id: i64, + method: &str, + params: serde_json::Value, + ) -> serde_json::Value { + channel + .tx + .unbounded_send(TransportFrame::Single( + RawJsonRpcMessage::request(method.to_string(), params, id.into()) + .expect("valid raw request"), + )) + .expect("server channel stays open"); + tokio::time::timeout(Duration::from_secs(2), async { + loop { + let frame = channel.rx.next().await.expect("server channel closed"); + if let TransportFrame::Single(message @ RawJsonRpcMessage::Response(_)) = frame { + let response = serde_json::to_value(message).expect("response serializes"); + if response.get("id") == Some(&json!(id)) { + return response; + } + } + } + }) + .await + .expect("raw request timed out") + } + + #[cfg(feature = "unstable-inject")] + struct GatedSession { + permits: Arc, + next_turn: Arc, + requests: Arc>>, + } + + #[cfg(feature = "unstable-inject")] + struct GatedTurn { + permits: Arc, + events: VecDeque, + started: bool, + } + + #[cfg(feature = "unstable-inject")] + #[async_trait] + impl ModelAdapter for GatedAdapter { + type Session = GatedSession; + + async fn start_session(&self, _config: SessionConfig) -> Result { + Ok(GatedSession { + permits: Arc::clone(&self.permits), + next_turn: Arc::clone(&self.next_turn), + requests: Arc::clone(&self.requests), + }) + } + } + + #[cfg(feature = "unstable-inject")] + #[async_trait] + impl ModelSession for GatedSession { + type Turn = GatedTurn; + + async fn begin_turn( + &mut self, + request: TurnRequest, + _cancellation: Option, + ) -> Result { + self.requests + .lock() + .unwrap_or_else(|error| error.into_inner()) + .push(request); + let turn = self.next_turn.fetch_add(1, Ordering::AcqRel); + let text = format!("turn-{turn}"); + Ok(GatedTurn { + permits: Arc::clone(&self.permits), + events: VecDeque::from([ + ModelTurnEvent::Delta(Delta::BeginPart { + part_id: PartId::new(format!("part-{turn}")), + kind: PartKind::Text, + }), + ModelTurnEvent::Delta(Delta::AppendText { + part_id: PartId::new(format!("part-{turn}")), + chunk: text.clone(), + }), + ModelTurnEvent::Finished(ModelTurnResult { + model: None, + response_id: None, + finish_reason: FinishReason::Completed, + output_items: vec![Item::text(ItemKind::Assistant, text)], + usage: None, + metadata: MetadataMap::new(), + }), + ]), + started: false, + }) + } + } + + #[cfg(feature = "unstable-inject")] + #[async_trait] + impl ModelTurn for GatedTurn { + async fn next_event( + &mut self, + cancellation: Option, + ) -> Result, LoopError> { + if !self.started { + self.started = true; + if let Some(cancellation) = cancellation { + tokio::select! { + permit = Arc::clone(&self.permits).acquire_owned() => { + permit.expect("gate stays open").forget(); + } + _ = cancellation.cancelled() => { + return Ok(Some(ModelTurnEvent::Finished(ModelTurnResult { + model: None, + response_id: None, + finish_reason: FinishReason::Cancelled, + output_items: Vec::new(), + usage: None, + metadata: MetadataMap::new(), + }))); + } + } + } + } + Ok(self.events.pop_front()) + } + } + + #[cfg(feature = "unstable-inject")] + #[tokio::test] + async fn batched_session_inject_activates_after_aggregate_response_enqueue() { + let adapter = GatedAdapter::new(); + let factory = SecondStartGatedFactory::new(adapter.clone()); + let (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 + } + }); + + let initialize = raw_request( + &mut client, + 1, + "initialize", + serde_json::to_value(wire::InitializeRequest::new( + wire::ProtocolVersion::V2, + wire::Implementation::new("batch-test", "1"), + )) + .unwrap(), + ) + .await; + assert!(initialize.get("result").is_some()); + let cwd = std::env::current_dir().unwrap(); + let new_session = raw_request( + &mut client, + 2, + "session/new", + serde_json::to_value(wire::NewSessionRequest::new(cwd.clone())).unwrap(), + ) + .await; + let session_id = new_session["result"]["sessionId"] + .as_str() + .unwrap() + .to_string(); + let prompt = raw_request( + &mut client, + 3, + "session/prompt", + serde_json::to_value(wire::PromptRequest::new( + wire::SessionId::new(session_id.clone()), + vec![wire::ContentBlock::Text(wire::TextContent::new("start"))], + )) + .unwrap(), + ) + .await; + assert!(prompt.get("result").is_some()); + + let inject = RawJsonRpcMessage::request( + "session/inject".to_string(), + serde_json::to_value(wire::InjectSessionRequest::new( + wire::SessionId::new(session_id.clone()), + wire::SessionInjectMode::Steer, + vec![wire::ContentBlock::Text(wire::TextContent::new( + "batched steer", + ))], + )) + .unwrap(), + 10.into(), + ) + .unwrap(); + let second_inject = RawJsonRpcMessage::request( + "session/inject".to_string(), + serde_json::to_value(wire::InjectSessionRequest::new( + wire::SessionId::new(session_id.clone()), + wire::SessionInjectMode::Steer, + vec![wire::ContentBlock::Text(wire::TextContent::new( + "second batched steer", + ))], + )) + .unwrap(), + 12.into(), + ) + .unwrap(); + let slow_sibling = RawJsonRpcMessage::request( + "session/new".to_string(), + serde_json::to_value(wire::NewSessionRequest::new(cwd)).unwrap(), + 11.into(), + ) + .unwrap(); + client + .tx + .unbounded_send(TransportFrame::Batch( + TransportBatch::from_messages([inject, second_inject, slow_sibling]).unwrap(), + )) + .unwrap(); + + factory.wait_for_second_start().await; + adapter.release(1); + tokio::time::sleep(Duration::from_millis(30)).await; + 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!( + entries + .iter() + .find(|entry| entry.get("id") == Some(&json!(12))) + .is_some_and(|entry| entry.get("result").is_some()) + ); + + adapter.release(1); + let delivered = tokio::time::timeout(Duration::from_secs(2), async { + loop { + let frame = client.rx.next().await.expect("server channel stays open"); + let value = match frame { + TransportFrame::Single(message) => serde_json::to_value(message).unwrap(), + TransportFrame::Batch(batch) => serde_json::to_value(batch).unwrap(), + TransportFrame::Malformed { raw, .. } => serde_json::Value::String(raw), + }; + if serde_json::to_string(&value) + .unwrap() + .contains("batched steer") + { + break true; + } + } + }) + .await + .expect("batched injection was not delivered"); + assert!(delivered); + let requests = adapter.requests(); + assert!(requests.len() >= 2); + let steers = requests[1] + .transcript + .iter() + .filter(|item| item.kind == ItemKind::User) + .flat_map(|item| &item.parts) + .filter_map(|part| match part { + Part::Text(text) if text.text.contains("batched steer") => Some(text.text.as_str()), + _ => None, + }) + .collect::>(); + assert_eq!(steers, ["batched steer", "second batched steer"]); + + drop(client); + server.abort(); + let _ = server.await; + } + + #[cfg(feature = "unstable-inject")] + #[tokio::test] + async fn cancel_request_before_acceptance_discards_inject() { + 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, + wire::SessionInjectMode::Steer, + vec![wire::ContentBlock::Text(wire::TextContent::new( + large_steer, + ))], + )) + .unwrap(), + 10.into(), + ) + .unwrap(), + )) + .unwrap(); + client + .tx + .unbounded_send(TransportFrame::Single( + RawJsonRpcMessage::notification( + "$/cancel_request".to_string(), + json!({ "requestId": 10 }), + ) + .unwrap(), + )) + .unwrap(); + adapter.release(1); + + tokio::time::sleep(Duration::from_millis(50)).await; + assert_eq!(adapter.requests().len(), 1); + drop(client); + server.abort(); + let _ = server.await; + } + + #[cfg(feature = "unstable-inject")] + #[tokio::test] + async fn closing_session_discards_pending_inject_and_returns_not_found() { + let adapter = GatedAdapter::new(); + let factory = SecondStartGatedFactory::new(adapter); + let (mut client, agent_transport) = Channel::duplex(); + let server = tokio::spawn({ + let factory = factory.clone(); + async move { + AcpHeadlessRuntime::::builder() + .agent_factory(factory) + .serve(agent_transport) + .await + } + }); + raw_request( + &mut client, + 1, + "initialize", + serde_json::to_value(wire::InitializeRequest::new( + wire::ProtocolVersion::V2, + wire::Implementation::new("close-test", "1"), + )) + .unwrap(), + ) + .await; + let cwd = std::env::current_dir().unwrap(); + let created = raw_request( + &mut client, + 2, + "session/new", + serde_json::to_value(wire::NewSessionRequest::new(cwd.clone())).unwrap(), + ) + .await; + let session_id = + wire::SessionId::new(created["result"]["sessionId"].as_str().unwrap().to_string()); + raw_request( + &mut client, + 3, + "session/prompt", + serde_json::to_value(wire::PromptRequest::new( + session_id.clone(), + vec![wire::ContentBlock::Text(wire::TextContent::new("start"))], + )) + .unwrap(), + ) + .await; + let inject = RawJsonRpcMessage::request( + "session/inject".to_string(), + serde_json::to_value(wire::InjectSessionRequest::new( + session_id.clone(), + wire::SessionInjectMode::Steer, + vec![wire::ContentBlock::Text(wire::TextContent::new("pending"))], + )) + .unwrap(), + 10.into(), + ) + .unwrap(); + let sibling = RawJsonRpcMessage::request( + "session/new".to_string(), + serde_json::to_value(wire::NewSessionRequest::new(cwd)).unwrap(), + 11.into(), + ) + .unwrap(); + client + .tx + .unbounded_send(TransportFrame::Batch( + TransportBatch::from_messages([inject, sibling]).unwrap(), + )) + .unwrap(); + factory.wait_for_second_start().await; + let close = raw_request( + &mut client, + 20, + "session/close", + serde_json::to_value(wire::CloseSessionRequest::new(session_id.clone())).unwrap(), + ) + .await; + assert!(close.get("result").is_some()); + factory.release_second_start(); + tokio::time::sleep(Duration::from_millis(30)).await; + let missing = raw_request( + &mut client, + 21, + "session/inject", + serde_json::to_value(wire::InjectSessionRequest::new( + session_id, + wire::SessionInjectMode::Steer, + vec![wire::ContentBlock::Text(wire::TextContent::new("late"))], + )) + .unwrap(), + ) + .await; + assert_eq!(missing["error"]["code"], json!(-32002)); + + drop(client); + server.abort(); + let _ = server.await; + } + + #[cfg(feature = "unstable-inject")] + #[tokio::test] + async fn accepted_steer_survives_approval_resolution() { + let adapter = ApprovalAdapter::new(); + let updates = Arc::new(Mutex::new(Vec::new())); + let (client_transport, agent_transport) = Channel::duplex(); + let server = tokio::spawn({ + let adapter = adapter.clone(); + async move { + AcpHeadlessRuntime::::builder() + .agent_factory(ApprovalFactory { adapter }) + .serve(agent_transport) + .await + } + }); + let client = agent_client_protocol::Client + .v2() + .on_receive_notification( + { + let updates = Arc::clone(&updates); + async move |notification: wire::UpdateSessionNotification, _cx| { + updates + .lock() + .unwrap() + .push((notification.session_id, notification.update)); + Ok(()) + } + }, + agent_client_protocol::on_receive_notification!(), + ) + .connect_with(client_transport, { + let updates = Arc::clone(&updates); + async move |cx| { + cx.send_request(wire::InitializeRequest::new( + wire::ProtocolVersion::V2, + wire::Implementation::new("approval-test", "1"), + )) + .block_task() + .await?; + let session_id = cx + .send_request(wire::NewSessionRequest::new( + std::env::current_dir() + .map_err(agent_client_protocol::Error::into_internal_error)?, + )) + .block_task() + .await? + .session_id; + cx.send_request(wire::PromptRequest::new( + session_id.clone(), + vec![wire::ContentBlock::Text(wire::TextContent::new("start"))], + )) + .block_task() + .await?; + cx.send_request(wire::InjectSessionRequest::new( + session_id.clone(), + wire::SessionInjectMode::Steer, + vec![wire::ContentBlock::Text(wire::TextContent::new( + "approval steer", + ))], + )) + .block_task() + .await?; + adapter.release(); + adapter.release(); + wait_for_idle(&updates, &session_id).await; + assert!(updates.lock().unwrap().iter().any(|(id, update)| { + id == &session_id + && matches!( + update, + wire::SessionUpdate::UserMessage(message) + if matches!( + &message.content, + agent_client_protocol::schema::MaybeUndefined::Value(content) + if content.iter().any(|block| { + matches!(block, wire::ContentBlock::Text(text) + if text.text == "approval steer") + }) + ) + ) + })); + Ok(()) + } + }); + tokio::time::timeout(Duration::from_secs(4), client) + .await + .expect("approval client timed out") + .expect("approval client failed"); + server.abort(); + let _ = server.await; + } + + #[cfg(feature = "unstable-inject")] + #[tokio::test] + async fn session_inject_accepted_during_tool_runs_after_tool_boundary() { + let adapter = MockAdapter::new(); + adapter.enqueue_many([ + tool_turn("post-tool"), + TurnScript::new([ModelTurnEvent::Finished(ModelTurnResult { + model: None, + response_id: None, + finish_reason: FinishReason::Completed, + output_items: Vec::new(), + usage: None, + metadata: MetadataMap::new(), + })]), + ]); + let tool = CompletionGatedTool::new(); + let (client_transport, agent_transport) = Channel::duplex(); + let server = tokio::spawn({ + let factory = InjectToolFactory { + adapter: adapter.clone(), + tool: tool.clone(), + }; + async move { + AcpHeadlessRuntime::::builder() + .agent_factory(factory) + .serve(agent_transport) + .await + } + }); + + let client = + agent_client_protocol::Client + .v2() + .connect_with(client_transport, async move |cx| { + cx.send_request(wire::InitializeRequest::new( + wire::ProtocolVersion::V2, + wire::Implementation::new("test-client", "1"), + )) + .block_task() + .await?; + let cwd = std::env::current_dir() + .map_err(agent_client_protocol::Error::into_internal_error)?; + let session_id = cx + .send_request(wire::NewSessionRequest::new(cwd)) + .block_task() + .await? + .session_id; + cx.send_request(wire::PromptRequest::new( + session_id.clone(), + vec![wire::ContentBlock::Text(wire::TextContent::new("start"))], + )) + .block_task() + .await?; + tool.wait_for_entered().await; + cx.send_request(wire::InjectSessionRequest::new( + session_id.clone(), + wire::SessionInjectMode::Steer, + vec![wire::ContentBlock::Text(wire::TextContent::new( + "after tool", + ))], + )) + .block_task() + .await?; + tool.release(); + tokio::time::timeout(Duration::from_secs(2), async { + while adapter.observed().len() < 2 { + tokio::task::yield_now().await; + } + }) + .await + .expect("continuation did not start"); + let observed = adapter.observed(); + assert_eq!(observed.len(), 2); + let contains_injection = |item: &Item| { + item.kind == ItemKind::User + && item.parts.iter().any(|part| { + matches!(part, Part::Text(text) if text.text == "after tool") + }) + }; + assert!(!observed[0].transcript.iter().any(contains_injection)); + assert!(observed[1].transcript.iter().any(contains_injection)); + cx.send_request(wire::CloseSessionRequest::new(session_id)) + .block_task() + .await?; + Ok(()) + }); + + tokio::time::timeout(Duration::from_secs(5), client) + .await + .expect("client timed out") + .expect("client run"); + server.abort(); + let _ = server.await; + } + + #[cfg(feature = "unstable-inject")] + #[tokio::test] + async fn delivery_and_revoke_race_has_one_linearized_outcome() { + let adapter = GatedAdapter::new(); + let updates = Arc::new(Mutex::new(Vec::new())); + let (client_transport, agent_transport) = Channel::duplex(); + let server = tokio::spawn({ + let adapter = adapter.clone(); + async move { + AcpHeadlessRuntime::::builder() + .agent_factory(TestFactory { adapter }) + .serve(agent_transport) + .await + } + }); + let client = agent_client_protocol::Client + .v2() + .on_receive_notification( + { + let updates = Arc::clone(&updates); + async move |notification: wire::UpdateSessionNotification, _cx| { + updates + .lock() + .unwrap() + .push((notification.session_id, notification.update)); + Ok(()) + } + }, + agent_client_protocol::on_receive_notification!(), + ) + .connect_with(client_transport, { + let updates = Arc::clone(&updates); + async move |cx| { + cx.send_request(wire::InitializeRequest::new( + wire::ProtocolVersion::V2, + wire::Implementation::new("race-test", "1"), + )) + .block_task() + .await?; + let session_id = cx + .send_request(wire::NewSessionRequest::new( + std::env::current_dir() + .map_err(agent_client_protocol::Error::into_internal_error)?, + )) + .block_task() + .await? + .session_id; + cx.send_request(wire::PromptRequest::new( + session_id.clone(), + vec![wire::ContentBlock::Text(wire::TextContent::new("race"))], + )) + .block_task() + .await?; + let accepted = cx + .send_request(wire::InjectSessionRequest::new( + session_id.clone(), + wire::SessionInjectMode::Steer, + vec![wire::ContentBlock::Text(wire::TextContent::new( + "racing steer", + ))], + )) + .block_task() + .await?; + let revoke = cx + .send_request(wire::RevokeInjectSessionRequest::new( + session_id.clone(), + accepted.message_id.clone(), + )) + .block_task(); + adapter.release(2); + let revoked = match revoke.await { + Ok(_) => true, + Err(error) => { + assert_eq!(i32::from(error.code), -32010); + false + } + }; + wait_for_idle(&updates, &session_id).await; + let was_delivered = updates.lock().unwrap().iter().any(|(id, update)| { + id == &session_id + && matches!(update, wire::SessionUpdate::UserMessage(message) + if message.message_id == accepted.message_id) + }); + assert_ne!(revoked, was_delivered); + Ok(()) + } + }); + tokio::time::timeout(Duration::from_secs(4), client) + .await + .expect("race client timed out") + .expect("race client failed"); + server.abort(); + let _ = server.await; + } + + #[cfg(feature = "unstable-inject")] + #[tokio::test] + async fn session_inject_steers_preserves_content_and_serializes_revoke_cancel() { + let adapter = GatedAdapter::new(); + let updates = Arc::new(Mutex::new(Vec::new())); + let (client_transport, agent_transport) = Channel::duplex(); + let server = tokio::spawn({ + let adapter = adapter.clone(); + async move { + AcpHeadlessRuntime::::builder() + .agent_factory(TestFactory { adapter }) + .serve(agent_transport) + .await + } + }); + + let client = agent_client_protocol::Client + .v2() + .on_receive_notification( + { + let updates = Arc::clone(&updates); + async move |notification: wire::UpdateSessionNotification, _cx| { + updates + .lock() + .unwrap() + .push((notification.session_id, notification.update)); + Ok(()) + } + }, + agent_client_protocol::on_receive_notification!(), + ) + .connect_with(client_transport, { + let adapter = adapter.clone(); + let updates = Arc::clone(&updates); + async move |cx| { + let initialize = cx + .send_request(wire::InitializeRequest::new( + wire::ProtocolVersion::V2, + wire::Implementation::new("test-client", "1"), + )) + .block_task() + .await?; + let inject = initialize + .capabilities + .session + .and_then(|session| session.inject) + .expect("inject capability advertised"); + assert_eq!(inject.modes, vec![wire::SessionInjectMode::Steer]); + assert_eq!( + inject.steer_in_stream, + Some(vec![wire::SessionInjectSteerInStream::Finish]) + ); + assert!(inject.pending.is_none()); + + let cwd = std::env::current_dir() + .map_err(agent_client_protocol::Error::into_internal_error)?; + let session_id = cx + .send_request(wire::NewSessionRequest::new(cwd)) + .block_task() + .await? + .session_id; + let idle_error = cx + .send_request(wire::InjectSessionRequest::new( + session_id.clone(), + wire::SessionInjectMode::Steer, + vec![wire::ContentBlock::Text(wire::TextContent::new("idle"))], + )) + .block_task() + .await + .expect_err("idle injection must fail"); + assert_eq!(i32::from(idle_error.code), -32010); + assert_eq!( + idle_error.data, + Some(json!({ "reason": "no_running_turn" })) + ); + let replace_error = cx + .send_request(wire::ReplaceInjectSessionRequest::new( + session_id.clone(), + "not-pending", + vec![wire::ContentBlock::Text(wire::TextContent::new("replace"))], + )) + .block_task() + .await + .expect_err("replace handler must not be registered"); + assert_eq!(i32::from(replace_error.code), -32601); + + cx.send_request(wire::PromptRequest::new( + session_id.clone(), + vec![wire::ContentBlock::Text(wire::TextContent::new("start"))], + )) + .block_task() + .await?; + let content = vec![ + wire::ContentBlock::Text(wire::TextContent::new("steer")), + wire::ContentBlock::ResourceLink(wire::ResourceLink::new( + "notes", + "file:///tmp/notes.txt", + )), + ]; + let accepted = tokio::time::timeout( + Duration::from_millis(250), + cx.send_request(wire::InjectSessionRequest::new( + session_id.clone(), + wire::SessionInjectMode::Steer, + content.clone(), + )) + .block_task(), + ) + .await + .expect("inject acceptance waited for the model")?; + assert!( + !updates.lock().unwrap().iter().any(|(id, update)| { + id == &session_id + && matches!(update, wire::SessionUpdate::UserMessage(message) + if message.message_id == accepted.message_id) + }), + "UserMessage was emitted before the inject response" + ); + adapter.release(2); + wait_for_idle_count(&updates, &session_id, 1).await; + let delivered = + updates + .lock() + .unwrap() + .iter() + .find_map(|(id, update)| match update { + wire::SessionUpdate::UserMessage(message) + if id == &session_id + && message.message_id == accepted.message_id => + { + Some(message.content.clone()) + } + _ => None, + }); + assert_eq!( + delivered, + Some(agent_client_protocol::schema::MaybeUndefined::Value( + content + )) + ); + { + let ordered = updates.lock().unwrap(); + let first_agent = ordered + .iter() + .position(|(id, update)| { + id == &session_id + && matches!(update, wire::SessionUpdate::AgentMessageChunk(_)) + }) + .expect("first model stream"); + let injected = ordered + .iter() + .position(|(id, update)| { + id == &session_id + && matches!(update, wire::SessionUpdate::UserMessage(message) + if message.message_id == accepted.message_id) + }) + .expect("injected user message"); + let second_agent = ordered + .iter() + .rposition(|(id, update)| { + id == &session_id + && matches!(update, wire::SessionUpdate::AgentMessageChunk(_)) + }) + .expect("continuation model stream"); + assert!(first_agent < injected && injected < second_agent); + assert_eq!( + ordered + .iter() + .filter(|(id, update)| id == &session_id + && matches!( + update, + wire::SessionUpdate::StateUpdate( + wire::StateUpdate::Running(_) + ) + )) + .count(), + 1 + ); + } + let delivered_error = cx + .send_request(wire::RevokeInjectSessionRequest::new( + session_id.clone(), + accepted.message_id.clone(), + )) + .block_task() + .await + .expect_err("delivered injection cannot be revoked"); + assert_eq!(i32::from(delivered_error.code), -32010); + assert_eq!( + delivered_error.data, + Some(json!({ + "reason": "already_delivered", + "messageId": accepted.message_id, + })) + ); + + cx.send_request(wire::PromptRequest::new( + session_id.clone(), + vec![wire::ContentBlock::Text(wire::TextContent::new("revoke"))], + )) + .block_task() + .await?; + let queue_error = cx + .send_request(wire::InjectSessionRequest::new( + session_id.clone(), + wire::SessionInjectMode::Queue, + vec![wire::ContentBlock::Text(wire::TextContent::new("queue"))], + )) + .block_task() + .await + .expect_err("queue mode must not be accepted"); + assert_eq!(i32::from(queue_error.code), -32602); + let revoked = cx + .send_request(wire::InjectSessionRequest::new( + session_id.clone(), + wire::SessionInjectMode::Steer, + vec![wire::ContentBlock::Text(wire::TextContent::new("remove"))], + )) + .block_task() + .await?; + cx.send_request(wire::RevokeInjectSessionRequest::new( + session_id.clone(), + revoked.message_id.clone(), + )) + .block_task() + .await?; + adapter.release(1); + wait_for_idle_count(&updates, &session_id, 2).await; + assert!(!updates.lock().unwrap().iter().any(|(id, update)| { + id == &session_id + && matches!(update, wire::SessionUpdate::UserMessage(message) + if message.message_id == revoked.message_id) + })); + + cx.send_request(wire::PromptRequest::new( + session_id.clone(), + vec![wire::ContentBlock::Text(wire::TextContent::new("cancel"))], + )) + .block_task() + .await?; + let cancelled = cx + .send_request(wire::InjectSessionRequest::new( + session_id.clone(), + wire::SessionInjectMode::Steer, + vec![wire::ContentBlock::Text(wire::TextContent::new("drop"))], + )) + .block_task() + .await?; + cx.send_notification(wire::CancelSessionNotification::new(session_id.clone()))?; + wait_for_idle_count(&updates, &session_id, 3).await; + assert!(!updates.lock().unwrap().iter().any(|(id, update)| { + id == &session_id + && matches!(update, wire::SessionUpdate::UserMessage(message) + if message.message_id == cancelled.message_id) + })); + let stop_reason = updates.lock().unwrap().iter().rev().find_map( + |(id, update)| match update { + wire::SessionUpdate::StateUpdate(wire::StateUpdate::Idle(idle)) + if id == &session_id => + { + idle.stop_reason.clone() + } + _ => None, + }, + ); + assert_eq!(stop_reason, Some(wire::StopReason::Cancelled)); + + cx.send_request(wire::PromptRequest::new( + session_id.clone(), + vec![wire::ContentBlock::Text(wire::TextContent::new( + "after cancel", + ))], + )) + .block_task() + .await?; + adapter.release(1); + 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) + }), + "cancelled injection carried into the next prompt" + ); + cx.send_request(wire::CloseSessionRequest::new(session_id)) + .block_task() + .await?; + Ok(()) + } + }); + + tokio::time::timeout(Duration::from_secs(8), client) + .await + .expect("client timed out") + .expect("client run"); + server.abort(); + let _ = server.await; + } + #[tokio::test] async fn independent_prompts_are_accepted_immediately_and_cancel_separately() { let updates = Arc::new(Mutex::new(Vec::new())); diff --git a/docs/acp.md b/docs/acp.md index bfa4c3e..1b3fa00 100644 --- a/docs/acp.md +++ b/docs/acp.md @@ -1,10 +1,13 @@ # 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.9 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; +> cancellation discards pending injections. Receipt-backed acceptance supports +> individual and batched inject requests. See the crate README or book chapter +> for the supported v2 lifecycle and current limits. ## Purpose @@ -37,7 +40,7 @@ There is already an official Rust ACP SDK: - `agent-client-protocol-schema` - `agent-client-protocol-rmcp` -So agentkit should not build a parallel protocol crate by default. A separate `racp` crate only makes sense if the official SDK becomes unsuitable. The first implementation should mirror `agentkit-mcp`: re-export upstream wire types and make agentkit own only the lifecycle and conversion glue. +So agentkit should not build a parallel protocol crate by default. The current implementation uses a commit-pinned source fork of that SDK for experimental session injection; it still re-exports the SDK wire types and owns only the lifecycle and conversion glue. A separate `racp` crate would make sense only if the SDK API became unsuitable. ## Non-goals @@ -52,7 +55,7 @@ So agentkit should not build a parallel protocol crate by default. A separate `r - long-term session persistence - policy storage for remembered approval choices -The crate integrates agentkit into ACP. It should not fork ACP or turn agentkit into a single opinionated coding-agent product. +The crate integrates agentkit into ACP. Its pinned SDK source fork must not become a divergent ACP protocol, and the crate must not turn agentkit into a single opinionated coding-agent product. ## Dependencies @@ -60,7 +63,8 @@ Recommended initial crate: ```toml [dependencies] -agent-client-protocol = "2.0.0" +# Registry-shaped so every workspace consumer resolves one patched type identity. +agent-client-protocol = "=2.0.0" agentkit-core = { path = "../agentkit-core", version = "0.10.5" } agentkit-loop = { path = "../agentkit-loop", version = "0.10.7" } agentkit-tools-core = { path = "../agentkit-tools-core", version = "0.10.5" } @@ -71,11 +75,19 @@ thiserror = { workspace = true } tokio = { workspace = true, features = ["sync"] } tracing = { workspace = true } +[patch.crates-io] +# Intentionally unpublishable until these SDK APIs are released upstream. +agent-client-protocol = { git = "https://github.com/danielkov/rust-sdk", rev = "2f039993d1d6ed8da35b38c31f54a7cbb7338c70" } + [features] default = ["stdio"] stdio = [] unstable-acp = ["agent-client-protocol/unstable"] protocol-v2 = ["agent-client-protocol/unstable_protocol_v2"] +unstable-inject = [ + "protocol-v2", + "agent-client-protocol/unstable_session_inject", +] ``` The umbrella crate should later add: @@ -727,6 +739,6 @@ The key should be explicit and policy-owned. A good default key should include: - ACP introduction: - ACP Rust SDK: - ACP protocol docs: -- Rust crate docs: +- ACP Rust SDK fork used by this crate: - Existing MCP integration design: [`mcp.md`](./mcp.md) - Agentkit permission design: [`permissions.md`](./permissions.md) From 5332c8a0b2910c608e0109c7541f9b9ff5340685 Mon Sep 17 00:00:00 2001 From: daniel Date: Tue, 25 Aug 2026 19:48:41 +0100 Subject: [PATCH 2/4] fix(acp): make v2 inject acceptance durable --- Cargo.lock | 2 +- book/src/acp.md | 12 +- crates/agentkit-acp/Cargo.toml | 2 +- crates/agentkit-acp/README.md | 19 +- crates/agentkit-acp/src/v2.rs | 510 +++++++++++++++++++++++++++------ docs/acp.md | 9 +- 6 files changed, 444 insertions(+), 110 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index d4fb8c9..6d6161b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -88,7 +88,7 @@ dependencies = [ [[package]] name = "agentkit-acp" -version = "0.10.9" +version = "0.10.10" dependencies = [ "agent-client-protocol", "agentkit-core", diff --git a/book/src/acp.md b/book/src/acp.md index a5db478..cdd1bd2 100644 --- a/book/src/acp.md +++ b/book/src/acp.md @@ -14,7 +14,7 @@ Like `agentkit-mcp`, this crate does not define a parallel protocol vocabulary. ACP v2 support is additive and disabled by default. Enable it explicitly: ```toml -agentkit-acp = { version = "0.10.9", features = ["protocol-v2"] } +agentkit-acp = { version = "0.10.10", features = ["protocol-v2"] } ``` `protocol-v2` enables the fork's @@ -51,10 +51,12 @@ 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. Revoke is serialized with acknowledged user-message -forwarding. Queueing and replacement are not supported. Cancellation and close -discard pending injected messages and never carry them into a later prompt; -revoke returns `already_delivered` when acknowledged delivery won first. +safe model/tool boundary. Request or session cancellation cannot remove a +response-committed steer, so it carries to the next valid boundary when the +current turn is cancelled. Close may discard it. Revoke is serialized with +acknowledged user-message forwarding. Queueing and replacement are not +supported. Delivered IDs retain `already_delivered` classification for the +session lifetime; a 4,096-accept lifetime cap bounds that history. The initial v2 foundation routes text, reasoning, and tool lifecycle updates. ACP v2 permission callbacks are intentionally deferred; unsupported approval diff --git a/crates/agentkit-acp/Cargo.toml b/crates/agentkit-acp/Cargo.toml index ced0f9c..7c745db 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.9" +version = "0.10.10" edition.workspace = true license.workspace = true rust-version.workspace = true diff --git a/crates/agentkit-acp/README.md b/crates/agentkit-acp/README.md index 1df9298..ab4f2a2 100644 --- a/crates/agentkit-acp/README.md +++ b/crates/agentkit-acp/README.md @@ -21,7 +21,7 @@ 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.9", features = ["protocol-v2"] } +agentkit-acp = { version = "0.10.10", features = ["protocol-v2"] } ``` The feature maps directly to the pinned fork's @@ -42,13 +42,16 @@ 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. Batch requests use the same -receipt-backed acceptance path. `session/revoke_inject` is mandatory and is -serialized with acknowledged `UserMessage` forwarding. Queueing, stream -interruption, and replacement are not supported. Cancelling or closing a -session drops every still-pending injected message; no matching `UserMessage` -is emitted and it does not carry into the next prompt. If acknowledged delivery -wins the race first, revoke returns `already_delivered`. +delivers at the next safe model/tool boundary. Single and batch requests use the +same receipt-backed acceptance path. Once response commitment succeeds, request +or session cancellation cannot remove the reservation; a committed steer that +misses the cancelled turn carries into the next prompt. Closing the session may +discard it. `session/revoke_inject` is mandatory and is serialized with +acknowledged `UserMessage` forwarding. Queueing, stream interruption, and +replacement are not supported. Delivered IDs retain their `already_delivered` +classification for the session lifetime. To bound that history safely, each +session accepts at most 4,096 injections and rejects later accepts with a +lifetime-limit error. This first v2 foundation streams text, reasoning, and tool lifecycle updates. The v1 permission bridge is not exposed through v2 wire types. Unsupported diff --git a/crates/agentkit-acp/src/v2.rs b/crates/agentkit-acp/src/v2.rs index 5bde342..ce4dd9b 100644 --- a/crates/agentkit-acp/src/v2.rs +++ b/crates/agentkit-acp/src/v2.rs @@ -572,37 +572,71 @@ where { let state = Arc::clone(&state); async move |request: wire::InjectSessionRequest, responder, cx| { - let cancellation = responder.cancellation(); - // Let a cancellation frame already queued behind this request - // linearize before AgentKit reserves the message. - tokio::task::yield_now().await; - if cancellation.is_cancelled() { - return responder.respond_with_result(Err( - agent_client_protocol::Error::request_cancelled(), - )); - } - match state.inject(request).await { - Ok(acceptance) => { - let receipt = match responder.respond_tracked(acceptance.response()) - { - Ok(receipt) => receipt, - Err(error) => { - acceptance.discard(); - return Err(error); - } - }; - cx.spawn(async move { - if receipt.await.is_ok() { - acceptance.activate(&cancellation); - } else { - acceptance.discard(); + let sequence = state.next_inject_request.fetch_add(1, Ordering::Relaxed); + let state = Arc::clone(&state); + cx.spawn(async move { + state.wait_for_inject_request(sequence).await; + let result = async { + let cancellation = responder.cancellation(); + let session_id = request.session_id.clone(); + // Return the receive callback first so a queued cancellation can + // commit before the injection response does. + tokio::task::yield_now().await; + if cancellation.is_cancelled() { + return responder.respond_with_result(Err( + agent_client_protocol::Error::request_cancelled(), + )); + } + match state.inject(request).await { + Ok(mut acceptance) => { + let entry = Arc::clone(&acceptance.entry); + let lifecycle = entry + .lifecycle + .lock() + .unwrap_or_else(|error| error.into_inner()); + if entry.closed.load(Ordering::Acquire) { + drop(lifecycle); + acceptance.discard(); + return responder.respond_with_result(Err( + session_not_found_error(&session_id), + )); + } + if cancellation.is_cancelled() { + drop(lifecycle); + acceptance.discard(); + return responder.respond_with_result(Err( + agent_client_protocol::Error::request_cancelled(), + )); + } + let receipt = match responder + .respond_tracked(acceptance.response()) + { + Ok(receipt) => receipt, + Err(error) => { + drop(lifecycle); + acceptance.discard(); + return Err(error); + } + }; + acceptance.commit(); + drop(lifecycle); + let _receipt_task = tokio::spawn(async move { + if receipt.await.is_ok() { + acceptance.activate(); + } else { + acceptance.discard(); + } + }); + Ok(()) } - Ok(()) - })?; - Ok(()) + Err(error) => responder.respond_with_result(Err(error)), + } } - Err(error) => responder.respond_with_result(Err(error)), - } + .await; + state.finish_inject_request(sequence); + result + })?; + Ok(()) } }, agent_client_protocol::on_receive_request!(), @@ -666,8 +700,11 @@ where const MAX_PENDING_INJECTIONS: usize = 64; #[cfg(feature = "unstable-inject")] const MAX_PENDING_INJECTION_BYTES: usize = 256 * 1024; +/// Caps every accepted injection ID retained for `already_delivered` +/// classification during a session. At 4,096 compact IDs, lifetime tracking +/// remains bounded without evicting classifications before the session closes. #[cfg(feature = "unstable-inject")] -const MAX_DELIVERED_TOMBSTONES: usize = 64; +const MAX_ACCEPTED_INJECTIONS: usize = 4_096; #[cfg(feature = "unstable-inject")] struct PendingInject { @@ -675,7 +712,15 @@ struct PendingInject { content: Vec, items: Vec, bytes: usize, - accepted: bool, + commitment: InjectCommitment, +} + +#[cfg(feature = "unstable-inject")] +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum InjectCommitment { + Reserved, + Committed, + Ready, } #[cfg(feature = "unstable-inject")] @@ -687,6 +732,7 @@ struct InjectionState { pending: VecDeque, delivering: Option, pending_bytes: usize, + accepted_count: usize, delivered: VecDeque, } @@ -718,10 +764,6 @@ impl InjectionController { let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); state.running = !state.cancelled; state.at_boundary = false; - state.pending.clear(); - state.delivering = None; - state.pending_bytes = 0; - state.delivered.clear(); } fn reset_cancellation(&self) { @@ -736,6 +778,21 @@ impl InjectionController { if !state.running { return Err(sdk_v2_error(wire::Error::inject_no_running_turn())); } + let reserved_count = state + .pending + .iter() + .filter(|pending| pending.commitment == InjectCommitment::Reserved) + .count(); + if state.accepted_count.saturating_add(reserved_count) >= MAX_ACCEPTED_INJECTIONS { + return Err(agent_client_protocol::Error::new( + -32000, + "session injection lifetime limit exceeded", + ) + .data(serde_json::json!({ + "reason": "lifetime_limit_exceeded", + "limit": MAX_ACCEPTED_INJECTIONS, + }))); + } if state.pending.len() + usize::from(state.delivering.is_some()) >= MAX_PENDING_INJECTIONS || state.pending_bytes.saturating_add(pending.bytes) > MAX_PENDING_INJECTION_BYTES { @@ -749,20 +806,27 @@ impl InjectionController { Ok(()) } - fn activate( - &self, - message_id: &wire::MessageId, - cancellation: &agent_client_protocol::RequestCancellation, - ) { + 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 cancellation.is_cancelled() || !state.running { - Self::remove_pending(&mut state, message_id); - } else if let Some(pending) = state + if let Some(pending) = state .pending .iter_mut() .find(|pending| &pending.message_id == message_id) { - pending.accepted = true; + debug_assert_eq!(pending.commitment, InjectCommitment::Committed); + pending.commitment = InjectCommitment::Ready; } drop(state); self.changed.notify_waiters(); @@ -799,7 +863,7 @@ impl InjectionController { if state .pending .front() - .is_some_and(|pending| !pending.accepted) + .is_some_and(|pending| pending.commitment != InjectCommitment::Ready) { return BoundaryAction::Wait; } @@ -824,9 +888,6 @@ impl InjectionController { state.delivering = None; state.pending_bytes = state.pending_bytes.saturating_sub(pending.bytes); if delivered { - if state.delivered.len() == MAX_DELIVERED_TOMBSTONES { - state.delivered.pop_front(); - } state.delivered.push_back(pending.message_id.clone()); } drop(state); @@ -835,11 +896,19 @@ impl InjectionController { fn revoke_transition(&self, message_id: &wire::MessageId) -> RevokeTransition { let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); - if Self::remove_pending(&mut state, message_id) { + let ready = state.pending.iter().any(|pending| { + &pending.message_id == message_id && pending.commitment == InjectCommitment::Ready + }); + if ready && Self::remove_pending(&mut state, message_id) { drop(state); self.changed.notify_waiters(); RevokeTransition::Revoked - } else if state.delivering.as_ref() == Some(message_id) { + } else if state.delivering.as_ref() == Some(message_id) + || state + .pending + .iter() + .any(|pending| &pending.message_id == message_id) + { RevokeTransition::WaitForDelivery } else if state.delivered.contains(message_id) { RevokeTransition::AlreadyDelivered @@ -853,13 +922,8 @@ impl InjectionController { state.running = false; state.cancelled = true; state.at_boundary = false; - let queued_bytes = state - .pending - .iter() - .map(|pending| pending.bytes) - .sum::(); - state.pending.clear(); - state.pending_bytes = state.pending_bytes.saturating_sub(queued_bytes); + // Accepted steers survive cancellation and become deliverable at the + // next valid boundary, potentially in the next prompt. drop(state); self.changed.notify_waiters(); } @@ -868,9 +932,25 @@ impl InjectionController { let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); state.running = false; state.at_boundary = false; - state.pending.clear(); - state.delivering = None; - state.pending_bytes = 0; + drop(state); + self.changed.notify_waiters(); + } + + fn close_session(&self) { + let mut state = self.state.lock().unwrap_or_else(|error| error.into_inner()); + state.running = false; + state.cancelled = true; + state.at_boundary = false; + let discarded_bytes = state + .pending + .iter() + .filter(|pending| pending.commitment == InjectCommitment::Ready) + .map(|pending| pending.bytes) + .sum::(); + state + .pending + .retain(|pending| pending.commitment != InjectCommitment::Ready); + state.pending_bytes = state.pending_bytes.saturating_sub(discarded_bytes); drop(state); self.changed.notify_waiters(); } @@ -893,9 +973,27 @@ impl InjectAcceptance { ) } - fn activate(mut self, cancellation: &agent_client_protocol::RequestCancellation) { + fn commit(&mut self) { + let message_id = self + .message_id + .as_ref() + .expect("acceptance has a message id"); + self.entry.injection.commit(message_id); + } + + fn activate(mut self) { let message_id = self.message_id.take().expect("acceptance has a message id"); - self.entry.injection.activate(&message_id, cancellation); + let lifecycle = self + .entry + .lifecycle + .lock() + .unwrap_or_else(|error| error.into_inner()); + if self.entry.closed.load(Ordering::Acquire) { + self.entry.injection.discard(&message_id); + } else { + self.entry.injection.activate(&message_id); + } + drop(lifecycle); } fn discard(mut self) { @@ -944,6 +1042,12 @@ where integration: Arc, sessions: AsyncMutex>>, next_session: AtomicU64, + #[cfg(feature = "unstable-inject")] + next_inject_request: AtomicU64, + #[cfg(feature = "unstable-inject")] + current_inject_request: AtomicU64, + #[cfg(feature = "unstable-inject")] + inject_request_changed: Notify, name: String, version: String, } @@ -959,11 +1063,42 @@ where integration: Arc::new(AcpIntegration::default()), sessions: AsyncMutex::new(HashMap::new()), next_session: AtomicU64::new(1), + #[cfg(feature = "unstable-inject")] + next_inject_request: AtomicU64::new(0), + #[cfg(feature = "unstable-inject")] + current_inject_request: AtomicU64::new(0), + #[cfg(feature = "unstable-inject")] + inject_request_changed: Notify::new(), name, version, } } + #[cfg(feature = "unstable-inject")] + async fn wait_for_inject_request(&self, sequence: u64) { + loop { + let changed = self.inject_request_changed.notified(); + tokio::pin!(changed); + changed.as_mut().enable(); + if self.current_inject_request.load(Ordering::Acquire) == sequence { + return; + } + changed.await; + } + } + + #[cfg(feature = "unstable-inject")] + fn finish_inject_request(&self, sequence: u64) { + let advanced = self.current_inject_request.compare_exchange( + sequence, + sequence + 1, + Ordering::AcqRel, + Ordering::Acquire, + ); + debug_assert_eq!(advanced, Ok(sequence)); + self.inject_request_changed.notify_waiters(); + } + fn initialize( &self, request: wire::InitializeRequest, @@ -1221,7 +1356,7 @@ where content: request.content, items, bytes, - accepted: false, + commitment: InjectCommitment::Reserved, })?; drop(_lifecycle); Ok(InjectAcceptance { @@ -1290,7 +1425,7 @@ where .unwrap_or_else(|error| error.into_inner()); if !entry.closed.load(Ordering::Acquire) && entry.busy.load(Ordering::Acquire) { #[cfg(feature = "unstable-inject")] - cancel_pending_injections(&entry); + cancel_injection_turn(&entry); entry.cancellation.interrupt(); } } @@ -1345,13 +1480,13 @@ fn signal_session_stop(entry: &Arc) { .unwrap_or_else(|error| error.into_inner()); entry.closed.store(true, Ordering::Release); #[cfg(feature = "unstable-inject")] - cancel_pending_injections(entry); + entry.injection.close_session(); entry.cancellation.interrupt(); let _ = entry.commands.send(SessionCommand::Shutdown); } #[cfg(feature = "unstable-inject")] -fn cancel_pending_injections(entry: &SessionEntry) { +fn cancel_injection_turn(entry: &SessionEntry) { entry.injection.cancel_turn(); } @@ -3253,6 +3388,11 @@ mod tests { .expect("valid raw request"), )) .expect("server channel stays open"); + raw_response(channel, id).await + } + + #[cfg(feature = "unstable-inject")] + async fn raw_response(channel: &mut Channel, id: i64) -> serde_json::Value { tokio::time::timeout(Duration::from_secs(2), async { loop { let frame = channel.rx.next().await.expect("server channel closed"); @@ -3463,6 +3603,11 @@ mod tests { factory.wait_for_second_start().await; adapter.release(1); tokio::time::sleep(Duration::from_millis(30)).await; + assert_eq!( + adapter.requests().len(), + 1, + "batch injections became deliverable before their aggregate response receipt" + ); factory.release_second_start(); let response = tokio::time::timeout(Duration::from_secs(2), async { @@ -3529,7 +3674,7 @@ mod tests { #[cfg(feature = "unstable-inject")] #[tokio::test] - async fn cancel_request_before_acceptance_discards_inject() { + async fn request_cancellation_respects_response_commit() { let adapter = GatedAdapter::new(); let (mut client, agent_transport) = Channel::duplex(); let server = tokio::spawn({ @@ -3583,7 +3728,7 @@ mod tests { RawJsonRpcMessage::request( "session/inject".to_string(), serde_json::to_value(wire::InjectSessionRequest::new( - session_id, + session_id.clone(), wire::SessionInjectMode::Steer, vec![wire::ContentBlock::Text(wire::TextContent::new( large_steer, @@ -3605,10 +3750,49 @@ mod tests { .unwrap(), )) .unwrap(); - adapter.release(1); + let cancelled_response = raw_response(&mut client, 10).await; + assert!(cancelled_response.get("result").is_none()); + assert_eq!(cancelled_response["error"]["code"], json!(-32800)); - tokio::time::sleep(Duration::from_millis(50)).await; - assert_eq!(adapter.requests().len(), 1); + 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; @@ -3714,6 +3898,151 @@ mod tests { let _ = server.await; } + #[cfg(feature = "unstable-inject")] + #[tokio::test] + async fn same_batch_close_prevents_inject_acceptance() { + let adapter = GatedAdapter::new(); + let (mut client, agent_transport) = Channel::duplex(); + let server = tokio::spawn({ + let adapter = adapter.clone(); + async move { + AcpHeadlessRuntime::::builder() + .agent_factory(TestFactory { adapter }) + .serve(agent_transport) + .await + } + }); + raw_request( + &mut client, + 1, + "initialize", + serde_json::to_value(wire::InitializeRequest::new( + wire::ProtocolVersion::V2, + wire::Implementation::new("same-batch-close-test", "1"), + )) + .unwrap(), + ) + .await; + let created = raw_request( + &mut client, + 2, + "session/new", + serde_json::to_value(wire::NewSessionRequest::new( + std::env::current_dir().unwrap(), + )) + .unwrap(), + ) + .await; + let session_id = + wire::SessionId::new(created["result"]["sessionId"].as_str().unwrap().to_string()); + raw_request( + &mut client, + 3, + "session/prompt", + serde_json::to_value(wire::PromptRequest::new( + session_id.clone(), + vec![wire::ContentBlock::Text(wire::TextContent::new("start"))], + )) + .unwrap(), + ) + .await; + + let inject = RawJsonRpcMessage::request( + "session/inject".to_string(), + serde_json::to_value(wire::InjectSessionRequest::new( + session_id.clone(), + wire::SessionInjectMode::Steer, + vec![wire::ContentBlock::Text(wire::TextContent::new("discard"))], + )) + .unwrap(), + 10.into(), + ) + .unwrap(); + let close = RawJsonRpcMessage::request( + "session/close".to_string(), + serde_json::to_value(wire::CloseSessionRequest::new(session_id)).unwrap(), + 11.into(), + ) + .unwrap(); + client + .tx + .unbounded_send(TransportFrame::Batch( + TransportBatch::from_messages([inject, close]).unwrap(), + )) + .unwrap(); + + let response = tokio::time::timeout(Duration::from_secs(2), async { + loop { + if let Some(TransportFrame::Batch(batch)) = client.rx.next().await { + break serde_json::to_value(batch).unwrap(); + } + } + }) + .await + .expect("same-batch close response timed out"); + let entries = response.as_array().expect("aggregate batch response"); + let inject_response = entries + .iter() + .find(|entry| entry.get("id") == Some(&json!(10))) + .expect("inject response slot"); + assert!(inject_response.get("result").is_none()); + assert_eq!(inject_response["error"]["code"], json!(-32002)); + let close_response = entries + .iter() + .find(|entry| entry.get("id") == Some(&json!(11))) + .expect("close response slot"); + assert!(close_response.get("result").is_some()); + assert!(close_response.get("error").is_none()); + assert_eq!(adapter.requests().len(), 1); + + drop(client); + server.abort(); + let _ = server.await; + } + + #[cfg(feature = "unstable-inject")] + #[test] + fn accepted_injection_lifetime_cap_is_enforced() { + let injection = InjectionController::default(); + injection.start_turn(); + for index in 0..MAX_ACCEPTED_INJECTIONS { + let message_id = wire::MessageId::new(format!("accepted-{index}")); + injection + .reserve(PendingInject { + message_id: message_id.clone(), + content: Vec::new(), + items: Vec::new(), + bytes: 0, + commitment: InjectCommitment::Reserved, + }) + .expect("injection below lifetime cap"); + injection.commit(&message_id); + injection.activate(&message_id); + assert!(matches!( + injection.revoke_transition(&message_id), + RevokeTransition::Revoked + )); + } + + let error = injection + .reserve(PendingInject { + message_id: wire::MessageId::new("over-limit"), + content: Vec::new(), + items: Vec::new(), + bytes: 0, + commitment: InjectCommitment::Reserved, + }) + .expect_err("accepted injection lifetime cap must reject new reservations"); + assert_eq!(i32::from(error.code), -32000); + assert_eq!( + error.data, + Some(json!({ + "reason": "lifetime_limit_exceeded", + "limit": MAX_ACCEPTED_INJECTIONS, + })) + ); + } + #[cfg(feature = "unstable-inject")] #[tokio::test] async fn accepted_steer_survives_approval_resolution() { @@ -4174,23 +4503,6 @@ mod tests { 1 ); } - let delivered_error = cx - .send_request(wire::RevokeInjectSessionRequest::new( - session_id.clone(), - accepted.message_id.clone(), - )) - .block_task() - .await - .expect_err("delivered injection cannot be revoked"); - assert_eq!(i32::from(delivered_error.code), -32010); - assert_eq!( - delivered_error.data, - Some(json!({ - "reason": "already_delivered", - "messageId": accepted.message_id, - })) - ); - cx.send_request(wire::PromptRequest::new( session_id.clone(), vec![wire::ContentBlock::Text(wire::TextContent::new("revoke"))], @@ -4228,6 +4540,22 @@ mod tests { && matches!(update, wire::SessionUpdate::UserMessage(message) if message.message_id == revoked.message_id) })); + let delivered_error = cx + .send_request(wire::RevokeInjectSessionRequest::new( + session_id.clone(), + accepted.message_id.clone(), + )) + .block_task() + .await + .expect_err("old delivered injection cannot be revoked in a later turn"); + assert_eq!(i32::from(delivered_error.code), -32010); + assert_eq!( + delivered_error.data, + Some(json!({ + "reason": "already_delivered", + "messageId": accepted.message_id, + })) + ); cx.send_request(wire::PromptRequest::new( session_id.clone(), @@ -4270,15 +4598,15 @@ mod tests { )) .block_task() .await?; - adapter.release(1); + adapter.release(2); wait_for_idle_count(&updates, &session_id, 4).await; assert!( - !updates.lock().unwrap().iter().any(|(id, update)| { + updates.lock().unwrap().iter().any(|(id, update)| { id == &session_id && matches!(update, wire::SessionUpdate::UserMessage(message) if message.message_id == cancelled.message_id) }), - "cancelled injection carried into the next prompt" + "committed injection did not survive session cancellation" ); cx.send_request(wire::CloseSessionRequest::new(session_id)) .block_task() diff --git a/docs/acp.md b/docs/acp.md index 1b3fa00..bd72b2f 100644 --- a/docs/acp.md +++ b/docs/acp.md @@ -1,12 +1,13 @@ # agentkit-acp design > **Implementation status:** Root APIs and default features implement stable ACP -> v1. Version 0.10.9 also provides an opt-in ACP v2 runtime foundation under +> v1. Version 0.10.10 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; -> cancellation discards pending injections. Receipt-backed acceptance supports -> individual and batched inject requests. See the crate README or book chapter +> `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. See the crate README or book chapter > for the supported v2 lifecycle and current limits. ## Purpose From 58a4780def8a11d6a50e8daf8e78c53d5d7e4c60 Mon Sep 17 00:00:00 2001 From: daniel Date: Tue, 25 Aug 2026 19:49:28 +0100 Subject: [PATCH 3/4] chore(release): keep stacked injection feature at 0.10.9 --- Cargo.lock | 2 +- book/src/acp.md | 2 +- crates/agentkit-acp/Cargo.toml | 2 +- crates/agentkit-acp/README.md | 2 +- docs/acp.md | 2 +- 5 files changed, 5 insertions(+), 5 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 6d6161b..d4fb8c9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -88,7 +88,7 @@ dependencies = [ [[package]] name = "agentkit-acp" -version = "0.10.10" +version = "0.10.9" dependencies = [ "agent-client-protocol", "agentkit-core", diff --git a/book/src/acp.md b/book/src/acp.md index cdd1bd2..42a4f3b 100644 --- a/book/src/acp.md +++ b/book/src/acp.md @@ -14,7 +14,7 @@ Like `agentkit-mcp`, this crate does not define a parallel protocol vocabulary. ACP v2 support is additive and disabled by default. Enable it explicitly: ```toml -agentkit-acp = { version = "0.10.10", features = ["protocol-v2"] } +agentkit-acp = { version = "0.10.9", features = ["protocol-v2"] } ``` `protocol-v2` enables the fork's diff --git a/crates/agentkit-acp/Cargo.toml b/crates/agentkit-acp/Cargo.toml index 7c745db..ced0f9c 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.10" +version = "0.10.9" edition.workspace = true license.workspace = true rust-version.workspace = true diff --git a/crates/agentkit-acp/README.md b/crates/agentkit-acp/README.md index ab4f2a2..e2a1276 100644 --- a/crates/agentkit-acp/README.md +++ b/crates/agentkit-acp/README.md @@ -21,7 +21,7 @@ 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.10", features = ["protocol-v2"] } +agentkit-acp = { version = "0.10.9", features = ["protocol-v2"] } ``` The feature maps directly to the pinned fork's diff --git a/docs/acp.md b/docs/acp.md index bd72b2f..0b86b8a 100644 --- a/docs/acp.md +++ b/docs/acp.md @@ -1,7 +1,7 @@ # agentkit-acp design > **Implementation status:** Root APIs and default features implement stable ACP -> v1. Version 0.10.10 also provides an opt-in ACP v2 runtime foundation under +> v1. Version 0.10.9 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. From 4299f31ec1455ea8ee4c46b4dce8e5bec4fdc597 Mon Sep 17 00:00:00 2001 From: daniel Date: Tue, 25 Aug 2026 21:04:27 +0100 Subject: [PATCH 4/4] feat(acp): expose v2 host session coordinator --- crates/agentkit-acp/src/v2.rs | 1433 +++++++++++++++++++++++---------- 1 file changed, 1013 insertions(+), 420 deletions(-) diff --git a/crates/agentkit-acp/src/v2.rs b/crates/agentkit-acp/src/v2.rs index ce4dd9b..162d29d 100644 --- a/crates/agentkit-acp/src/v2.rs +++ b/crates/agentkit-acp/src/v2.rs @@ -55,6 +55,29 @@ pub mod wire { 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 { @@ -75,27 +98,31 @@ impl ClientHandle { (Self { tx }, rx) } - fn update( + fn update_for( &self, session_id: wire::SessionId, update: wire::SessionUpdate, ) -> Result<(), AcpRuntimeError> { + 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, - session_id: wire::SessionId, - update: wire::SessionUpdate, + notification: wire::UpdateSessionNotification, ) -> Result<(), AcpRuntimeError> { let (tx, rx) = oneshot::channel(); self.tx .send(ClientMessage::AcknowledgedUpdate { - notification: Box::new(wire::UpdateSessionNotification::new(session_id, update)), + notification: Box::new(notification), acknowledged: tx, }) .map_err(|_| AcpRuntimeError::ClientClosed)?; @@ -151,34 +178,173 @@ struct CurrentMessageIds { struct IntegrationSession { acp_session_id: wire::SessionId, - client: ClientHandle, + agentkit_session_id: AgentkitSessionId, + sink: Arc, + cancellation: CancellationController, + closed: AtomicBool, + lifecycle: Mutex<()>, + #[cfg(feature = "unstable-inject")] + injection: Arc, next_message: AtomicU64, current_messages: Mutex>, part_kinds: Mutex>, } +/// 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() @@ -195,21 +361,26 @@ 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()), - }), - ); - 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: 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()), + }); + inner.by_acp.insert(acp_session_id, Arc::clone(&session)); + Ok(AcpSessionHandle { session }) } - 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() @@ -221,6 +392,8 @@ impl AcpIntegration { inner .by_agentkit .retain(|_, mapped| mapped != &session.acp_session_id); + drop(inner); + AcpSessionHandle { session }.close(); Ok(()) } @@ -248,7 +421,17 @@ impl AcpIntegration { ))) } - fn begin_prompt( + /// 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 { @@ -258,12 +441,21 @@ impl AcpIntegration { Ok(message_id) } - fn finish_prompt(&self, session_id: &wire::SessionId) { + /// 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 route_event(&self, session_id: &AgentkitSessionId, event: AgentEvent) { let session = { let inner = self.inner.read().unwrap_or_else(|error| error.into_inner()); @@ -313,10 +505,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"); } } @@ -414,6 +606,7 @@ where M: ModelAdapter, { factory: Option>>, + integration: AcpIntegration, name: String, version: String, } @@ -425,6 +618,7 @@ where fn default() -> Self { Self { factory: None, + integration: AcpIntegration::default(), name: "agentkit".into(), version: env!("CARGO_PKG_VERSION").into(), } @@ -456,6 +650,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 { @@ -484,7 +688,12 @@ 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 agent = agent_client_protocol::Agent .v2() @@ -570,71 +779,11 @@ where let agent = agent .on_receive_request( { - let state = Arc::clone(&state); + let integration = Arc::clone(&state.integration); async move |request: wire::InjectSessionRequest, responder, cx| { - let sequence = state.next_inject_request.fetch_add(1, Ordering::Relaxed); - let state = Arc::clone(&state); + let integration = Arc::clone(&integration); cx.spawn(async move { - state.wait_for_inject_request(sequence).await; - let result = async { - let cancellation = responder.cancellation(); - let session_id = request.session_id.clone(); - // Return the receive callback first so a queued cancellation can - // commit before the injection response does. - tokio::task::yield_now().await; - if cancellation.is_cancelled() { - return responder.respond_with_result(Err( - agent_client_protocol::Error::request_cancelled(), - )); - } - match state.inject(request).await { - Ok(mut acceptance) => { - let entry = Arc::clone(&acceptance.entry); - let lifecycle = entry - .lifecycle - .lock() - .unwrap_or_else(|error| error.into_inner()); - if entry.closed.load(Ordering::Acquire) { - drop(lifecycle); - acceptance.discard(); - return responder.respond_with_result(Err( - session_not_found_error(&session_id), - )); - } - if cancellation.is_cancelled() { - drop(lifecycle); - acceptance.discard(); - return responder.respond_with_result(Err( - agent_client_protocol::Error::request_cancelled(), - )); - } - let receipt = match responder - .respond_tracked(acceptance.response()) - { - Ok(receipt) => receipt, - Err(error) => { - drop(lifecycle); - acceptance.discard(); - return Err(error); - } - }; - acceptance.commit(); - drop(lifecycle); - let _receipt_task = tokio::spawn(async move { - if receipt.await.is_ok() { - acceptance.activate(); - } else { - acceptance.discard(); - } - }); - Ok(()) - } - Err(error) => responder.respond_with_result(Err(error)), - } - } - .await; - state.finish_inject_request(sequence); - result + integration.handle_inject_request(request, responder).await })?; Ok(()) } @@ -643,9 +792,9 @@ where ) .on_receive_request( { - let state = Arc::clone(&state); + let integration = Arc::clone(&state.integration); async move |request: wire::RevokeInjectSessionRequest, responder, _cx| { - responder.respond_with_result(state.revoke_inject(request).await) + responder.respond_with_result(integration.revoke_inject(request).await) } }, agent_client_protocol::on_receive_request!(), @@ -696,15 +845,19 @@ where } } +/// Maximum number of pending unstable ACP v2 injections per session. #[cfg(feature = "unstable-inject")] -const MAX_PENDING_INJECTIONS: usize = 64; +pub const MAX_PENDING_INJECTIONS: usize = 64; +/// Maximum serialized content bytes retained by pending unstable injections. #[cfg(feature = "unstable-inject")] -const MAX_PENDING_INJECTION_BYTES: usize = 256 * 1024; -/// Caps every accepted injection ID retained for `already_delivered` -/// classification during a session. At 4,096 compact IDs, lifetime tracking -/// remains bounded without evicting classifications before the session closes. +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")] -const MAX_ACCEPTED_INJECTIONS: usize = 4_096; +pub const MAX_ACCEPTED_INJECTIONS: usize = 4_096; #[cfg(feature = "unstable-inject")] struct PendingInject { @@ -747,7 +900,7 @@ struct InjectionController { enum BoundaryAction { Wait, Deliver(PendingInject), - Complete(InjectBoundary), + Complete(AcpInjectionBoundary), } #[cfg(feature = "unstable-inject")] @@ -858,7 +1011,7 @@ impl InjectionController { state.at_boundary = true; if !state.running { state.at_boundary = false; - return BoundaryAction::Complete(InjectBoundary::Stopped); + return BoundaryAction::Complete(AcpInjectionBoundary::Stopped); } if state .pending @@ -873,12 +1026,12 @@ impl InjectionController { } state.at_boundary = false; if delivered_any { - BoundaryAction::Complete(InjectBoundary::Delivered) + BoundaryAction::Complete(AcpInjectionBoundary::Delivered) } else if terminal { state.running = false; - BoundaryAction::Complete(InjectBoundary::Finished) + BoundaryAction::Complete(AcpInjectionBoundary::Finished) } else { - BoundaryAction::Complete(InjectBoundary::Continue) + BoundaryAction::Complete(AcpInjectionBoundary::Continue) } } @@ -957,69 +1110,389 @@ impl InjectionController { } #[cfg(feature = "unstable-inject")] -struct InjectAcceptance { - entry: Arc, +struct ReservedInject { + session: Arc, message_id: Option, } #[cfg(feature = "unstable-inject")] -impl InjectAcceptance { +impl ReservedInject { fn response(&self) -> wire::InjectSessionResponse { wire::InjectSessionResponse::new( self.message_id .as_ref() - .expect("acceptance has a message id") + .expect("reserved injection has a message id") .clone(), ) } - fn commit(&mut self) { - let message_id = self - .message_id + fn message_id(&self) -> &wire::MessageId { + self.message_id .as_ref() - .expect("acceptance has a message id"); - self.entry.injection.commit(message_id); + .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("acceptance has a message id"); + let message_id = self + .message_id + .take() + .expect("reserved injection has a message id"); let lifecycle = self - .entry + .session .lifecycle .lock() .unwrap_or_else(|error| error.into_inner()); - if self.entry.closed.load(Ordering::Acquire) { - self.entry.injection.discard(&message_id); + if self.session.closed.load(Ordering::Acquire) { + self.session.injection.discard(&message_id); } else { - self.entry.injection.activate(&message_id); + self.session.injection.activate(&message_id); } drop(lifecycle); } fn discard(mut self) { - let message_id = self.message_id.take().expect("acceptance has a message id"); - self.entry.injection.discard(&message_id); + let message_id = self + .message_id + .take() + .expect("reserved injection has a message id"); + self.session.injection.discard(&message_id); } } #[cfg(feature = "unstable-inject")] -impl Drop for InjectAcceptance { +impl Drop for ReservedInject { fn drop(&mut self) { if let Some(message_id) = self.message_id.take() { - self.entry.injection.discard(&message_id); + self.session.injection.discard(&message_id); + } + } +} + +#[cfg(feature = "unstable-inject")] +struct InjectRequestPermit { + integration: AcpIntegration, + sequence: Option, +} + +#[cfg(feature = "unstable-inject")] +impl Drop for InjectRequestPermit { + fn drop(&mut self) { + if let Some(sequence) = self.sequence.take() { + self.integration.finish_inject_request(sequence); + } + } +} + +/// Opaque reservation for one globally ordered unstable ACP v2 inject request. +/// +/// The handle owns the SDK responder so the response cannot be committed +/// outside the session lifecycle lock or against a different cancellation +/// token. Dropping it discards the reservation and advances request ordering. +#[cfg(feature = "unstable-inject")] +#[must_use = "the reserved inject request must be responded to or dropped"] +pub struct AcpInjectRequest { + reserved: Option, + responder: Option>, + cancellation: agent_client_protocol::RequestCancellation, + _permit: InjectRequestPermit, +} + +#[cfg(feature = "unstable-inject")] +impl AcpInjectRequest { + /// Returns the response that will be committed by [`respond_tracked`](Self::respond_tracked). + #[must_use] + pub fn response(&self) -> wire::InjectSessionResponse { + self.reserved + .as_ref() + .expect("inject request has a reservation") + .response() + } + + /// Commits the response under the session lifecycle lock. + /// + /// This checks close and request cancellation immediately before calling the + /// SDK's tracked responder. `Ok(None)` means the coordinator sent a close or + /// cancellation error instead. `Ok(Some(_))` returns an acceptance handle + /// bound to the resulting [`agent_client_protocol::ResponseReceipt`]. + pub fn respond_tracked( + mut self, + ) -> Result, agent_client_protocol::Error> { + let session = Arc::clone( + &self + .reserved + .as_ref() + .expect("inject request has a reservation") + .session, + ); + let lifecycle = session + .lifecycle + .lock() + .unwrap_or_else(|error| error.into_inner()); + let rejection = if session.closed.load(Ordering::Acquire) { + Some(session_not_found_error(&session.acp_session_id)) + } else if self.cancellation.is_cancelled() { + Some(agent_client_protocol::Error::request_cancelled()) + } else { + None + }; + if let Some(error) = rejection { + drop(lifecycle); + self.reserved.take().expect("inject reservation").discard(); + self.responder + .take() + .expect("inject request has a responder") + .respond_with_result(Err(error))?; + return Ok(None); + } + + let response = self.response(); + let receipt = match self + .responder + .take() + .expect("inject request has a responder") + .respond_tracked(response) + { + Ok(receipt) => receipt, + Err(error) => { + drop(lifecycle); + self.reserved.take().expect("inject reservation").discard(); + return Err(error); + } + }; + let mut reserved = self.reserved.take().expect("inject reservation"); + reserved.commit(); + drop(lifecycle); + Ok(Some(AcpInjectAcceptance { + reserved: Some(reserved), + receipt: Some(receipt), + })) + } +} + +/// Opaque response-committed unstable ACP v2 injection acceptance. +/// +/// This handle owns both the reservation and the SDK response receipt. Dropping +/// it before activation discards the reservation. +#[cfg(feature = "unstable-inject")] +#[must_use = "the response receipt must be activated or the injection is discarded"] +pub struct AcpInjectAcceptance { + reserved: Option, + receipt: Option, +} + +#[cfg(feature = "unstable-inject")] +impl AcpInjectAcceptance { + /// Returns the agent-owned message ID reserved for this acceptance. + #[must_use] + pub fn message_id(&self) -> &wire::MessageId { + self.reserved + .as_ref() + .expect("inject acceptance has a reservation") + .message_id() + } + + /// Waits for the SDK response receipt, then makes the steer deliverable. + /// + /// Batch safety requires calling this only after the receive callback has + /// returned, normally in a spawned task. Receipt failure or future + /// cancellation drops and discards the reservation. + pub async fn activate_after_response(mut self) -> Result<(), agent_client_protocol::Error> { + self.receipt + .take() + .expect("inject acceptance has a response receipt") + .await?; + self.reserved + .take() + .expect("inject acceptance has a reservation") + .activate(); + Ok(()) + } +} + +#[cfg(feature = "unstable-inject")] +impl AcpIntegration { + async fn wait_for_inject_request(&self, sequence: u64) { + loop { + let changed = self.inject_requests.changed.notified(); + tokio::pin!(changed); + changed.as_mut().enable(); + if self.inject_requests.current.load(Ordering::Acquire) == sequence { + return; + } + changed.await; + } + } + + fn finish_inject_request(&self, sequence: u64) { + let advanced = self.inject_requests.current.compare_exchange( + sequence, + sequence + 1, + Ordering::AcqRel, + Ordering::Acquire, + ); + debug_assert_eq!(advanced, Ok(sequence)); + self.inject_requests.changed.notify_waiters(); + } + + fn reserve_inject( + &self, + request: wire::InjectSessionRequest, + ) -> Result { + if !matches!(request.mode, wire::SessionInjectMode::Steer) { + return Err(agent_client_protocol::Error::new( + -32602, + "unsupported session injection mode", + )); + } + let items = content_blocks_to_items(&request.content).map_err(crate::sdk_error)?; + let bytes = serde_json::to_vec(&request.content) + .map_err(|error| agent_client_protocol::Error::new(-32603, error.to_string()))? + .len(); + let session = self + .session(&request.session_id) + .map_err(|_| session_not_found_error(&request.session_id))?; + let _lifecycle = session + .lifecycle + .lock() + .unwrap_or_else(|error| error.into_inner()); + if session.closed.load(Ordering::Acquire) { + return Err(session_not_found_error(&request.session_id)); + } + let message_id = self + .next_user_message_id(&request.session_id) + .map_err(crate::sdk_error)?; + session.injection.reserve(PendingInject { + message_id: message_id.clone(), + content: request.content, + items, + bytes, + commitment: InjectCommitment::Reserved, + })?; + drop(_lifecycle); + Ok(ReservedInject { + session, + message_id: Some(message_id), + }) + } + + /// Reserves one globally ordered unstable ACP v2 inject request. + /// + /// The SDK responder is intentionally consumed here: its cancellation token + /// and tracked response must remain paired with this reservation. This method + /// responds to pre-reservation cancellation or validation errors itself and + /// returns `Ok(None)`. A returned handle owns FIFO ordering and Drop cleanup. + pub async fn reserve_inject_request( + &self, + request: wire::InjectSessionRequest, + responder: agent_client_protocol::Responder, + ) -> Result, agent_client_protocol::Error> { + let sequence = self.inject_requests.next.fetch_add(1, Ordering::Relaxed); + self.wait_for_inject_request(sequence).await; + let permit = InjectRequestPermit { + integration: self.clone(), + sequence: Some(sequence), + }; + let cancellation = responder.cancellation(); + // Return the receive callback first so a queued cancellation can commit + // before reservation or response commitment. + tokio::task::yield_now().await; + if cancellation.is_cancelled() { + responder + .respond_with_result(Err(agent_client_protocol::Error::request_cancelled()))?; + drop(permit); + return Ok(None); + } + let reserved = match self.reserve_inject(request) { + Ok(reserved) => reserved, + Err(error) => { + responder.respond_with_result(Err(error))?; + drop(permit); + return Ok(None); + } + }; + Ok(Some(AcpInjectRequest { + reserved: Some(reserved), + responder: Some(responder), + cancellation, + _permit: permit, + })) + } + + /// Handles one unstable ACP v2 `session/inject` request end to end. + /// + /// This convenience path uses [`reserve_inject_request`](Self::reserve_inject_request), + /// [`AcpInjectRequest::respond_tracked`], and + /// [`AcpInjectAcceptance::activate_after_response`]. Hosts that need staged + /// control can call those same methods without reimplementing the race. + pub async fn handle_inject_request( + &self, + request: wire::InjectSessionRequest, + responder: agent_client_protocol::Responder, + ) -> Result<(), agent_client_protocol::Error> { + let Some(request) = self.reserve_inject_request(request, responder).await? else { + return Ok(()); + }; + let Some(acceptance) = request.respond_tracked()? else { + return Ok(()); + }; + let _receipt_task = tokio::spawn(async move { + if let Err(error) = acceptance.activate_after_response().await { + tracing::debug!(%error, "ACP v2 inject response was not accepted"); + } + }); + Ok(()) + } + + /// Handles the state transition for unstable ACP v2 `session/revoke_inject`. + pub async fn revoke_inject( + &self, + request: wire::RevokeInjectSessionRequest, + ) -> Result { + let session = self + .session(&request.session_id) + .map_err(|_| session_not_found_error(&request.session_id))?; + loop { + let changed = session.injection.changed.notified(); + tokio::pin!(changed); + changed.as_mut().enable(); + let _lifecycle = session + .lifecycle + .lock() + .unwrap_or_else(|error| error.into_inner()); + if session.closed.load(Ordering::Acquire) { + return Err(session_not_found_error(&request.session_id)); + } + match session.injection.revoke_transition(&request.message_id) { + RevokeTransition::Revoked => { + return Ok(wire::RevokeInjectSessionResponse::new()); + } + RevokeTransition::WaitForDelivery => drop(_lifecycle), + RevokeTransition::AlreadyDelivered => { + return Err(sdk_v2_error(wire::Error::inject_already_delivered( + request.message_id, + ))); + } + RevokeTransition::Unknown => { + return Err(sdk_v2_error(wire::Error::inject_unknown_message_id( + request.message_id, + ))); + } + } + changed.await; } } } struct SessionEntry { - commands: mpsc::UnboundedSender, - cancellation: CancellationController, - info: wire::SessionInfo, - busy: Arc, - closed: AtomicBool, - lifecycle: Mutex<()>, - #[cfg(feature = "unstable-inject")] - injection: Arc, + commands: mpsc::UnboundedSender, + session: AcpSessionHandle, + info: wire::SessionInfo, + busy: Arc, task: Mutex>>, drain_task: Mutex>>, } @@ -1042,12 +1515,6 @@ where integration: Arc, sessions: AsyncMutex>>, next_session: AtomicU64, - #[cfg(feature = "unstable-inject")] - next_inject_request: AtomicU64, - #[cfg(feature = "unstable-inject")] - current_inject_request: AtomicU64, - #[cfg(feature = "unstable-inject")] - inject_request_changed: Notify, name: String, version: String, } @@ -1057,48 +1524,22 @@ where M: ModelAdapter + Send + Sync + 'static, M::Session: Send + 'static, { - fn new(factory: Arc>, name: String, version: String) -> Self { + fn new( + factory: Arc>, + integration: Arc, + name: String, + version: String, + ) -> Self { Self { factory, - integration: Arc::new(AcpIntegration::default()), + integration, sessions: AsyncMutex::new(HashMap::new()), next_session: AtomicU64::new(1), - #[cfg(feature = "unstable-inject")] - next_inject_request: AtomicU64::new(0), - #[cfg(feature = "unstable-inject")] - current_inject_request: AtomicU64::new(0), - #[cfg(feature = "unstable-inject")] - inject_request_changed: Notify::new(), name, version, } } - #[cfg(feature = "unstable-inject")] - async fn wait_for_inject_request(&self, sequence: u64) { - loop { - let changed = self.inject_request_changed.notified(); - tokio::pin!(changed); - changed.as_mut().enable(); - if self.current_inject_request.load(Ordering::Acquire) == sequence { - return; - } - changed.await; - } - } - - #[cfg(feature = "unstable-inject")] - fn finish_inject_request(&self, sequence: u64) { - let advanced = self.current_inject_request.compare_exchange( - sequence, - sequence + 1, - Ordering::AcqRel, - Ordering::Acquire, - ); - debug_assert_eq!(advanced, Ok(sequence)); - self.inject_request_changed.notify_waiters(); - } - fn initialize( &self, request: wire::InitializeRequest, @@ -1112,7 +1553,7 @@ where wire::ProtocolVersion::V2, wire::Implementation::new(self.name.clone(), self.version.clone()), ) - .capabilities(headless_capabilities())) + .capabilities(agent_capabilities())) } async fn new_session( @@ -1136,15 +1577,18 @@ where json!(request.additional_directories), ); - self.integration.bind( - acp_session_id.clone(), - agentkit_session_id.clone(), - client.clone(), + let session = self.integration.bind_session( + AcpSessionBinding::new( + acp_session_id.clone(), + agentkit_session_id.clone(), + client.clone(), + ) + .cancellation(cancellation), )?; let drain_task = tokio::spawn(drain_client_messages(client_messages, cx)); let ctx = AcpAgentFactoryContext { acp_session_id: acp_session_id.clone(), - agentkit_session_id, + agentkit_session_id: agentkit_session_id.clone(), cwd: request.cwd.into_inner(), additional_directories: request .additional_directories @@ -1152,13 +1596,13 @@ where .map(wire::AbsolutePath::into_inner) .collect(), integration: Arc::clone(&self.integration), - cancellation: cancellation.handle(), + cancellation: session.cancellation_handle(), metadata, }; let driver = match self.factory.start(ctx).await { Ok(driver) => driver, Err(error) => { - let _ = self.integration.unbind(&acp_session_id); + let _ = self.integration.unbind_session(&acp_session_id); drain_task.abort(); let _ = drain_task.await; return Err(error); @@ -1169,35 +1613,15 @@ where let busy = Arc::new(AtomicBool::new(false)); let worker_busy = Arc::clone(&busy); let integration = Arc::clone(&self.integration); - let worker_session_id = acp_session_id.clone(); - let worker_cancellation = cancellation.handle(); - #[cfg(feature = "unstable-inject")] - let injection = Arc::new(InjectionController::default()); - #[cfg(feature = "unstable-inject")] - let worker_injection = Arc::clone(&injection); + let worker_session = session.clone(); let task = tokio::spawn(async move { - session_worker( - worker_session_id, - driver, - client, - integration, - worker_cancellation, - worker_busy, - #[cfg(feature = "unstable-inject")] - worker_injection, - rx, - ) - .await; + session_worker(worker_session, driver, client, integration, worker_busy, rx).await; }); let entry = Arc::new(SessionEntry { commands, - cancellation, + session, info, busy, - closed: AtomicBool::new(false), - lifecycle: Mutex::new(()), - #[cfg(feature = "unstable-inject")] - injection, task: Mutex::new(Some(task)), drain_task: Mutex::new(Some(drain_task)), }); @@ -1221,7 +1645,7 @@ where let mut infos = sessions .values() .filter(|entry| { - !entry.closed.load(Ordering::Acquire) + !entry.session.is_closed() && request .cwd .as_ref() @@ -1253,7 +1677,7 @@ where .get(&request.session_id) .cloned() .ok_or_else(|| AcpRuntimeError::SessionNotFound(request.session_id.to_string()))?; - if entry.closed.load(Ordering::Acquire) || entry.info.cwd != request.cwd { + if entry.session.is_closed() || entry.info.cwd != request.cwd { return Err(AcpRuntimeError::SessionNotFound( request.session_id.to_string(), )); @@ -1270,7 +1694,7 @@ where &self, request: wire::PromptRequest, ) -> Result, AcpRuntimeError> { - let items = prompt_to_items(&request)?; + let items = self.integration.prompt_to_items(&request)?; let entry = self .sessions .lock() @@ -1281,10 +1705,12 @@ where let (tx, rx) = oneshot::channel(); { let _lifecycle = entry + .session + .session .lifecycle .lock() .unwrap_or_else(|error| error.into_inner()); - if entry.closed.load(Ordering::Acquire) { + if entry.session.is_closed() { return Err(AcpRuntimeError::SessionNotFound( request.session_id.to_string(), )); @@ -1298,9 +1724,9 @@ where "session is already running a prompt".into(), )); } - let cancellation_generation = entry.cancellation.handle().generation(); + let cancellation_generation = entry.session.cancellation_handle().generation(); #[cfg(feature = "unstable-inject")] - entry.injection.reset_cancellation(); + entry.session.prepare_injection_turn(); if entry .commands .send(SessionCommand::Prompt { @@ -1318,96 +1744,6 @@ where rx.await.map_err(|_| AcpRuntimeError::ClientClosed)? } - #[cfg(feature = "unstable-inject")] - async fn inject( - &self, - request: wire::InjectSessionRequest, - ) -> Result { - if !matches!(request.mode, wire::SessionInjectMode::Steer) { - return Err(agent_client_protocol::Error::new( - -32602, - "unsupported session injection mode", - )); - } - let items = content_to_items(&request.content).map_err(crate::sdk_error)?; - let bytes = serde_json::to_vec(&request.content) - .map_err(|error| agent_client_protocol::Error::new(-32603, error.to_string()))? - .len(); - let entry = self - .sessions - .lock() - .await - .get(&request.session_id) - .cloned() - .ok_or_else(|| session_not_found_error(&request.session_id))?; - let _lifecycle = entry - .lifecycle - .lock() - .unwrap_or_else(|error| error.into_inner()); - if entry.closed.load(Ordering::Acquire) { - return Err(session_not_found_error(&request.session_id)); - } - let message_id = self - .integration - .next_user_message_id(&request.session_id) - .map_err(crate::sdk_error)?; - entry.injection.reserve(PendingInject { - message_id: message_id.clone(), - content: request.content, - items, - bytes, - commitment: InjectCommitment::Reserved, - })?; - drop(_lifecycle); - Ok(InjectAcceptance { - entry, - message_id: Some(message_id), - }) - } - - #[cfg(feature = "unstable-inject")] - async fn revoke_inject( - &self, - request: wire::RevokeInjectSessionRequest, - ) -> Result { - let entry = self - .sessions - .lock() - .await - .get(&request.session_id) - .cloned() - .ok_or_else(|| session_not_found_error(&request.session_id))?; - loop { - let changed = entry.injection.changed.notified(); - tokio::pin!(changed); - changed.as_mut().enable(); - let _lifecycle = entry - .lifecycle - .lock() - .unwrap_or_else(|error| error.into_inner()); - if entry.closed.load(Ordering::Acquire) { - return Err(session_not_found_error(&request.session_id)); - } - match entry.injection.revoke_transition(&request.message_id) { - RevokeTransition::Revoked => { - return Ok(wire::RevokeInjectSessionResponse::new()); - } - RevokeTransition::WaitForDelivery => drop(_lifecycle), - RevokeTransition::AlreadyDelivered => { - return Err(sdk_v2_error(wire::Error::inject_already_delivered( - request.message_id, - ))); - } - RevokeTransition::Unknown => { - return Err(sdk_v2_error(wire::Error::inject_unknown_message_id( - request.message_id, - ))); - } - } - changed.await; - } - } - async fn cancel( &self, notification: wire::CancelSessionNotification, @@ -1418,16 +1754,10 @@ where .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) && entry.busy.load(Ordering::Acquire) { - #[cfg(feature = "unstable-inject")] - cancel_injection_turn(&entry); - entry.cancellation.interrupt(); - } + if let Some(entry) = entry + && entry.busy.load(Ordering::Acquire) + { + entry.session.interrupt(); } Ok(()) } @@ -1443,7 +1773,7 @@ where .remove(&request.session_id) .ok_or_else(|| AcpRuntimeError::SessionNotFound(request.session_id.to_string()))?; stop_session(Arc::clone(&entry)).await; - self.integration.unbind(&request.session_id)?; + self.integration.unbind_session(&request.session_id)?; stop_client(entry).await; Ok(wire::CloseSessionResponse::new()) } @@ -1460,7 +1790,7 @@ where if let Some(task) = take_task(&entry.task) { session_tasks.push(task); } - let _ = self.integration.unbind(session_id); + let _ = self.integration.unbind_session(session_id); } join_tasks_until(deadline, session_tasks).await; @@ -1474,22 +1804,10 @@ where } fn signal_session_stop(entry: &Arc) { - let _lifecycle = entry - .lifecycle - .lock() - .unwrap_or_else(|error| error.into_inner()); - entry.closed.store(true, Ordering::Release); - #[cfg(feature = "unstable-inject")] - entry.injection.close_session(); - entry.cancellation.interrupt(); + entry.session.close(); let _ = entry.commands.send(SessionCommand::Shutdown); } -#[cfg(feature = "unstable-inject")] -fn cancel_injection_turn(entry: &SessionEntry) { - entry.injection.cancel_turn(); -} - fn take_task( task: &Mutex>>, ) -> Option> { @@ -1528,24 +1846,25 @@ async fn stop_session(entry: Arc) { } async fn stop_client(entry: Arc) { - if let Some(task) = take_task(&entry.drain_task) { + let task = take_task(&entry.drain_task); + drop(entry); + if let Some(task) = task { let _ = task.await; } } -#[allow(clippy::too_many_arguments)] // Feature-gated injection adds one state/notify pair. async fn session_worker( - session_id: wire::SessionId, + session: AcpSessionHandle, mut driver: agentkit_loop::LoopDriver, client: ClientHandle, integration: Arc, - cancellation: CancellationHandle, busy: Arc, - #[cfg(feature = "unstable-inject")] injection: Arc, mut commands: mpsc::UnboundedReceiver, ) where S: ModelSession + Send + 'static, { + let session_id = session.acp_session_id().clone(); + let cancellation = session.cancellation_handle(); while let Some(command) = commands.recv().await { let SessionCommand::Prompt { request, @@ -1573,25 +1892,25 @@ async fn session_worker( } }; #[cfg(feature = "unstable-inject")] - injection.start_turn(); + session.start_injection_turn(); let (start_tx, start_rx) = oneshot::channel(); if response.send(Ok(start_tx)).is_err() || start_rx.await.is_err() { #[cfg(feature = "unstable-inject")] - injection.stop_turn(); + session.stop_injection_turn(); integration.finish_prompt(&session_id); busy.store(false, Ordering::Release); continue; } if client - .update( + .update_for( session_id.clone(), wire::SessionUpdate::UserMessage( wire::UserMessage::new(user_message_id).content(request.prompt), ), ) .and_then(|()| { - client.update( + client.update_for( session_id.clone(), wire::SessionUpdate::StateUpdate(wire::StateUpdate::Running( wire::RunningStateUpdate::new(), @@ -1601,7 +1920,7 @@ async fn session_worker( .is_err() { #[cfg(feature = "unstable-inject")] - injection.stop_turn(); + session.stop_injection_turn(); integration.finish_prompt(&session_id); busy.store(false, Ordering::Release); continue; @@ -1612,11 +1931,7 @@ async fn session_worker( &cancellation, cancellation_generation, #[cfg(feature = "unstable-inject")] - &client, - #[cfg(feature = "unstable-inject")] - &session_id, - #[cfg(feature = "unstable-inject")] - &injection, + &session, ) .await; if let Err(error) = client.flush().await { @@ -1624,65 +1939,78 @@ async fn session_worker( } integration.finish_prompt(&session_id); busy.store(false, Ordering::Release); - let _ = client.update( + let _ = client.update_for( session_id.clone(), wire::SessionUpdate::StateUpdate(wire::StateUpdate::Idle( wire::IdleStateUpdate::new().stop_reason(stop_reason), )), ); #[cfg(feature = "unstable-inject")] - injection.stop_turn(); + session.stop_injection_turn(); } } +/// Outcome of an unstable ACP v2 injection boundary. #[cfg(feature = "unstable-inject")] #[derive(Debug, Eq, PartialEq)] -enum InjectBoundary { +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")] -async fn handle_inject_boundary( - driver: &mut agentkit_loop::LoopDriver, - client: &ClientHandle, - session_id: &wire::SessionId, - injection: &InjectionController, - terminal: bool, -) -> Result -where - S: ModelSession + Send + 'static, -{ - 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); +impl AcpSessionHandle { + /// Delivers ready unstable ACP v2 steers at one safe loop boundary. + /// + /// `terminal` must be true for terminal model/input boundaries and false + /// after tool results. The returned outcome tells the host whether to + /// continue, finish, or report cancellation. + pub async fn handle_injection_boundary( + &self, + driver: &mut agentkit_loop::LoopDriver, + terminal: bool, + ) -> Result + where + S: ModelSession + Send + 'static, + { + let injection = &self.session.injection; + let mut delivered_any = false; + loop { + let changed = injection.changed.notified(); + tokio::pin!(changed); + changed.as_mut().enable(); + match injection.boundary_action(terminal, delivered_any) { + BoundaryAction::Wait => changed.await, + BoundaryAction::Complete(outcome) => return Ok(outcome), + BoundaryAction::Deliver(pending) => { + if let Err(error) = driver + .submit_input(pending.items.clone()) + .map_err(|error| AcpRuntimeError::Loop(error.to_string())) + { + injection.finish_delivery(&pending, false); + return Err(error); + } + let result = self + .session + .sink + .update_acknowledged(wire::UpdateSessionNotification::new( + self.session.acp_session_id.clone(), + wire::SessionUpdate::UserMessage( + wire::UserMessage::new(pending.message_id.clone()) + .content(pending.content.clone()), + ), + )) + .await; + injection.finish_delivery(&pending, result.is_ok()); + result?; + delivered_any = true; } - let result = client - .update_acknowledged( - 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; } } } @@ -1692,9 +2020,7 @@ async fn drive_prompt( driver: &mut agentkit_loop::LoopDriver, cancellation: &CancellationHandle, generation: u64, - #[cfg(feature = "unstable-inject")] client: &ClientHandle, - #[cfg(feature = "unstable-inject")] session_id: &wire::SessionId, - #[cfg(feature = "unstable-inject")] injection: &InjectionController, + #[cfg(feature = "unstable-inject")] session: &AcpSessionHandle, ) -> wire::StopReason where S: ModelSession + Send + 'static, @@ -1705,7 +2031,7 @@ where Err(error) => { tracing::debug!(%error, "ACP v2 agent loop failed"); #[cfg(feature = "unstable-inject")] - injection.stop_turn(); + session.stop_injection_turn(); return if cancellation.is_cancelled_since(generation) { wire::StopReason::Cancelled } else { @@ -1715,7 +2041,7 @@ where }; if cancellation.is_cancelled_since(generation) { #[cfg(feature = "unstable-inject")] - injection.stop_turn(); + session.stop_injection_turn(); return wire::StopReason::Cancelled; } match step { @@ -1724,15 +2050,17 @@ where continue; } #[cfg(feature = "unstable-inject")] - match handle_inject_boundary(driver, client, session_id, injection, true).await { - Ok(InjectBoundary::Delivered | InjectBoundary::Continue) => continue, - Ok(InjectBoundary::Stopped) => return wire::StopReason::Cancelled, - Ok(InjectBoundary::Finished) => { + 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"); - injection.stop_turn(); + session.stop_injection_turn(); return error_stop_reason(); } } @@ -1741,27 +2069,28 @@ where } LoopStep::Interrupt(LoopInterrupt::AwaitingInput(_)) => { #[cfg(feature = "unstable-inject")] - match handle_inject_boundary(driver, client, session_id, injection, true).await { - Ok(InjectBoundary::Delivered | InjectBoundary::Continue) => continue, - Ok(InjectBoundary::Stopped) => return wire::StopReason::Cancelled, - Ok(InjectBoundary::Finished) => return wire::StopReason::EndTurn, + 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"); - injection.stop_turn(); + session.stop_injection_turn(); return error_stop_reason(); } } #[cfg(not(feature = "unstable-inject"))] return wire::StopReason::EndTurn; } - LoopStep::Interrupt(LoopInterrupt::AfterToolResult(_)) => - { + LoopStep::Interrupt(LoopInterrupt::AfterToolResult(_)) => { #[cfg(feature = "unstable-inject")] - match handle_inject_boundary(driver, client, session_id, injection, false).await { - Ok(InjectBoundary::Stopped) => return wire::StopReason::Cancelled, + 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"); - injection.stop_turn(); + session.stop_injection_turn(); return error_stop_reason(); } _ => {} @@ -1771,7 +2100,7 @@ where if let Err(error) = driver.cancel_pending_approvals().await { tracing::debug!(%error, "failed to resolve unsupported ACP v2 approval"); #[cfg(feature = "unstable-inject")] - injection.stop_turn(); + session.stop_injection_turn(); return error_stop_reason(); } } @@ -1780,10 +2109,16 @@ where } fn prompt_to_items(request: &wire::PromptRequest) -> Result, AcpRuntimeError> { - content_to_items(&request.prompt) + content_blocks_to_items(&request.prompt) } -fn content_to_items(content: &[wire::ContentBlock]) -> Result, AcpRuntimeError> { +/// 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(); @@ -2124,7 +2459,22 @@ fn finish_reason_to_stop_reason(reason: &FinishReason) -> wire::StopReason { } } -fn headless_capabilities() -> wire::AgentCapabilities { +/// Constructs the exact unstable ACP v2 injection capability implemented here. +/// +/// This function is available only with `unstable-inject`, matching the request +/// handlers and turn-boundary APIs it advertises. +#[cfg(feature = "unstable-inject")] +pub fn session_inject_capabilities() -> wire::SessionInjectCapabilities { + wire::SessionInjectCapabilities::new(vec![wire::SessionInjectMode::Steer]) + .steer_in_stream(vec![wire::SessionInjectSteerInStream::Finish]) +} + +/// Constructs the honest ACP v2 capabilities supported by this build. +/// +/// Injection is advertised only when `unstable-inject` is enabled, and uses +/// [`session_inject_capabilities`] so hosts and [`AcpHeadlessRuntime`] report +/// the same behavior. +pub fn agent_capabilities() -> wire::AgentCapabilities { let session = wire::SessionCapabilities::new().prompt( wire::PromptCapabilities::new() .image(wire::PromptImageCapabilities::new()) @@ -2132,10 +2482,7 @@ fn headless_capabilities() -> wire::AgentCapabilities { .embedded_context(wire::PromptEmbeddedContextCapabilities::new()), ); #[cfg(feature = "unstable-inject")] - let session = session.inject( - wire::SessionInjectCapabilities::new(vec![wire::SessionInjectMode::Steer]) - .steer_in_stream(vec![wire::SessionInjectSteerInStream::Finish]), - ); + let session = session.inject(session_inject_capabilities()); wire::AgentCapabilities::new().session(session) } @@ -2164,6 +2511,38 @@ mod tests { #[cfg(feature = "unstable-inject")] use futures_util::StreamExt as _; + #[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, @@ -2979,14 +3358,209 @@ mod tests { 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 tool_execution_boundary_rotates_message_ids_without_a_terminal_result() { + 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(acp_id.clone(), agentkit_id.clone(), client) + .bind_session(AcpSessionBinding::new( + acp_id.clone(), + agentkit_id.clone(), + client, + )) .expect("bind session"); integration.begin_prompt(&acp_id).expect("begin prompt"); @@ -3035,7 +3609,7 @@ mod tests { } #[test] - fn integration_rejects_duplicate_agentkit_session_ids() { + fn public_host_binding_rejects_duplicate_agentkit_session_ids() { let integration = AcpIntegration::default(); let (first_client, _first_rx) = ClientHandle::channel(); let (second_client, _second_rx) = ClientHandle::channel(); @@ -3043,17 +3617,28 @@ mod tests { let first_acp = wire::SessionId::new("first-acp-session"); let second_acp = wire::SessionId::new("second-acp-session"); integration - .bind(first_acp.clone(), agentkit_id.clone(), first_client) + .bind_session(AcpSessionBinding::new( + first_acp.clone(), + agentkit_id.clone(), + first_client, + )) .expect("first binding"); let error = integration - .bind(second_acp.clone(), agentkit_id, second_client) - .expect_err("duplicate AgentKit session id must be rejected"); + .bind_session(AcpSessionBinding::new( + second_acp.clone(), + agentkit_id, + second_client, + )) + .err() + .expect("duplicate AgentKit session id must be rejected"); assert!(matches!(error, AcpRuntimeError::SessionAlreadyBound(_))); assert!(matches!( integration.session(&second_acp), Err(AcpRuntimeError::SessionNotFound(_)) )); - integration.unbind(&first_acp).expect("unbind first"); + integration + .unbind_session(&first_acp) + .expect("unbind first"); } #[test] @@ -3513,12 +4098,15 @@ mod tests { 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 } @@ -3626,6 +4214,11 @@ mod tests { .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()