diff --git a/README.md b/README.md index c90563a..0dbeb4f 100644 --- a/README.md +++ b/README.md @@ -22,7 +22,8 @@ This repository is the official **Rust SDK** for ACP. It provides crates for bui Native MCP-over-ACP support is currently opt-in through the core crate's `unstable_mcp_over_acp` feature. Standalone MCP servers need no ACP transport feature; the rmcp integration exposes a matching passthrough feature when those -servers are attached to ACP. +servers are attached to ACP. MCP attachment and proxy-session helpers currently +use stable protocol v1. **Proxy orchestration** @@ -40,6 +41,9 @@ servers are attached to ACP. - **API reference** for individual crates is on [docs.rs/agent-client-protocol](https://docs.rs/agent-client-protocol). - **Design and architecture documentation** lives in the mdbook at [agentclientprotocol.github.io/rust-sdk](https://agentclientprotocol.github.io/rust-sdk/). Source is in [`md/`](./md/). +- **Draft protocol v2** setup, version-typed connections, and high-level + session usage are covered in + [Protocol V2](./md/protocol-v2.md). ## Integrations diff --git a/md/protocol-v2.md b/md/protocol-v2.md index f7d9bfc..1cd0680 100644 --- a/md/protocol-v2.md +++ b/md/protocol-v2.md @@ -21,7 +21,10 @@ rules. By default, `Client.builder()` and `Agent.builder()` continue to expose the stable v1 API and advertise protocol v1. To use the v2 API for a connection, -construct the builder with `Client.v2()` or `Agent.v2()`: +construct the builder with `Client.v2()` or `Agent.v2()`. Fluent typed handlers, +spawned tasks, close callbacks, and `connect_with` receive +`V2ConnectionTo<_>`, so the protocol version is reflected in the high-level +Rust API as well as on the wire: ```rust use agent_client_protocol::schema::{ProtocolVersion, v2}; @@ -73,6 +76,142 @@ When v2 mode is enabled, application code should use types from exports remain the stable v1 schema. This will likely change as v2 gets closer to release. +## High-level v2 sessions + +Stable callbacks receive `ConnectionTo<_>` and expose the protocol v1 +`build_session*`, `SessionBuilder`, `ActiveSession`, and `SessionMessage` APIs. +Callbacks installed through `Client.v2()` receive `V2ConnectionTo<_>` and expose +the v2 `build_session*` and `resume_session*` helpers. The shared names describe +the same lifecycle operations while the connection type selects their schema +and return types at compile time. + +Low-level custom `with_handler` and `with_runner` implementations continue to +receive the protocol-neutral `ConnectionTo<_>`, and generic `send_request` +remains schema-agnostic. Runtime compatibility checks remain at those explicit +escape hatches. Dynamic handlers registered through +`V2ConnectionTo::add_dynamic_handler` use the same low-level +`HandleDispatchFrom` interface and therefore also receive `ConnectionTo<_>`. + +Nested connections preserve the stable `ConnectionTo` API while still typing +the child implementation's callbacks. On a raw `ConnectionTo<_>`, +`spawn_connection(Client.v2(), transport)` returns a raw `ConnectionTo<_>`; +callbacks installed on that v2 child builder still receive +`V2ConnectionTo<_>`. Existing `spawn_connection::` calls therefore remain +source-compatible. + +When a raw parent also needs a typed handle to the v2 child, the +`unstable_protocol_v2` feature exposes +`ConnectionTo::spawn_connection_with_context`, which returns the context +selected by the child builder. `V2ConnectionTo::spawn_connection` likewise +follows the child builder naturally, so spawning `Client.v2()` through an +already-typed v2 connection returns another `V2ConnectionTo<_>`. + +```rust,ignore +use agent_client_protocol::schema::{ProtocolVersion, v2}; +use agent_client_protocol::{Client, Responder}; + +Client + .v2() + .on_receive_notification( + async |update: v2::UpdateSessionNotification, _cx| { + apply_session_update(update)?; + Ok(()) + }, + agent_client_protocol::on_receive_notification!(), + ) + .on_receive_request( + async |request: v2::RequestPermissionRequest, + responder: Responder, + _cx| { + // Transfer the responder to application-owned permission handling + // without waiting for user input in the dispatch callback. + queue_permission_request(request, responder)?; + Ok(()) + }, + agent_client_protocol::on_receive_request!(), + ) + .connect_with(agent_transport, async |cx| { + let initialize = cx + .send_request(v2::InitializeRequest::new( + ProtocolVersion::V2, + v2::Implementation::new("example", "0.1.0"), + )) + .block_task() + .await?; + assert!(initialize.capabilities.session.is_some()); + + let opened = cx + .build_session_cwd()? + .start_session() + .block_task() + .await?; + let (session, new_session_response) = opened.into_parts(); + assert_eq!(session.session_id(), &new_session_response.session_id); + + session + .send_prompt("What is 2 + 2?") + .block_task() + .await?; + println!("prompt accepted"); + + Ok(()) + }) + .await?; +``` + +Here `apply_session_update` updates application-owned state, while +`queue_permission_request` transfers the request and its responder to a +separate permission workflow. + +V2 deliberately separates prompt submission from session observation: + +- `session/prompt` returns a `PromptResponse` as soon as the agent accepts the + prompt. `V2Session::send_prompt` returns that request as a + `SentRequest`; callers must explicitly await it, register a + response callback, or detach it. +- `V2SessionBuilder::start_session` likewise returns a mapped `SentRequest`. + Its `OpenedV2Session` result keeps the command handle separate from the + complete `NewSessionResponse` represented by the linked schema, rather than + reconstructing a selected subset of its fields. +- `V2Session` is a cloneable command handle containing only the session ID and + connection. It does not own, buffer, or unregister inbound messages. +- Register typed `UpdateSessionNotification` and `RequestPermissionRequest` + handlers on `Client.v2()` before connecting. Updates and interactive requests + are separate protocol lanes; permission handlers should transfer responders + to application-owned work rather than waiting for user input inside the + connection dispatch loop. Without matching handlers, unhandled v2 + notifications are ignored and unhandled requests receive a method-not-found + response; they are not retained for a later per-session receiver. +- `session/update` events can arrive before, during, or after a prompt request. + They carry a session ID and entity IDs, but no prompt or turn ID. The SDK + therefore does not attribute intervening events to a locally submitted + prompt or provide a prompt-scoped text accumulator. +- `state_update` describes the session-wide foreground state. `idle` means the + session can accept ordinary new foreground work; it is not a wire-level + boundary assigning previous events to one prompt, and background updates may + continue while idle. +- `cancel_active_work` sends session-wide `session/cancel`. Cancellation + completes after the required `idle` update with stop reason `cancelled`. The + client should immediately mark unfinished tool calls for the active work as + cancelled and must resolve every pending permission request with the + cancelled outcome. Cancelling or dropping the prompt's `SentRequest` is the + separate JSON-RPC request-cancellation mechanism. +- `set_config_option` returns the authoritative replacement option set, and + `close` returns the complete close response. Mutable configuration is not + cached on the command handle. + +Install connection handlers before `session/new` and `session/resume` requests. +This is especially important for `resume_session_from`: replay updates +precede the resume response on the wire, so preinstalled typed handlers observe +them in order. If a handler forwards updates to another task, the application +is responsible for any additional projection-drained barrier it needs before +treating replay as locally applied. + +Dropping command handles has no network or inbound-routing side effect. If an +application wants stream ergonomics, it can fan typed updates out from the +connection handler with an explicit buffering and subscriber policy. MCP +attachment and proxy-session helpers are still v1-only. + The SDK handles the `initialize` negotiation at the JSON-RPC boundary: - A v2 client advertises protocol v2 as its latest supported version. diff --git a/src/agent-client-protocol/CHANGELOG.md b/src/agent-client-protocol/CHANGELOG.md index 1c55a66..765a072 100644 --- a/src/agent-client-protocol/CHANGELOG.md +++ b/src/agent-client-protocol/CHANGELOG.md @@ -8,11 +8,32 @@ `unstable_tool_call_name` feature. - *(unstable)* Expose v1 and draft-v2 plan operations through the `unstable_plan_operations` feature. +- *(unstable-v2)* Add high-level protocol v2 session builders and handles with + independent prompt-acceptance requests, create and resume setup responses, + cloneable command handles, configuration, close, and session-wide + cancellation. `Client.v2()` and `Agent.v2()` callbacks receive a + `V2ConnectionTo` whose unversioned `build_session*` and `resume_session*` + methods expose only the v2 lifecycle. Session updates and interactive + requests remain on typed connection handlers; MCP attachment and + proxy-session helpers remain v1-only. +- *(unstable-v2)* Expose protocol-neutral dynamic handler registration on + `V2ConnectionTo`. +- *(unstable-v2)* Add `ConnectionTo::spawn_connection_with_context` for raw + parents that need the context selected by a child builder. + `V2ConnectionTo::spawn_connection` also returns that natural typed context. + The existing `ConnectionTo::spawn_connection::` API remains + source-compatible and returns a raw `ConnectionTo`, including when the v2 + child builder's callbacks receive `V2ConnectionTo`. ### Fixed - *(unstable-v2)* Preserve unknown initialize fields when the protocol router hands a same-version connection to its selected implementation. +- *(unstable-v2)* Do not retain unhandled v2 session messages for a dynamic v1 + session route that will never be installed. +- Register framework-owned ordered response handling before outbound requests + become visible to the peer, preventing fast responses from overtaking session + routing or proxy forwarding. - Allow the portable protocol engine and transport abstractions to build for the `wasm32-wasip1` and `wasm32-wasip2` targets. The native process-backed `AcpAgent`, thread-backed `Stdio`, and associated `LineDirection` type are not diff --git a/src/agent-client-protocol/README.md b/src/agent-client-protocol/README.md index 43ff867..309cbd1 100644 --- a/src/agent-client-protocol/README.md +++ b/src/agent-client-protocol/README.md @@ -17,7 +17,8 @@ enabling features like tool use, permission requests, and streaming responses. ## Quick Start: Connecting to an Agent -The most common use case is connecting to an existing ACP agent as a client: +The most common use case is connecting to an existing ACP agent as a client. +This quick start uses stable protocol v1: ```rust,no_run use agent_client_protocol::{AcpAgent, Client, Result}; @@ -39,6 +40,13 @@ Client.builder() # } ``` +Draft protocol v2 is opt-in through `unstable_protocol_v2`. `Client.v2()` and +`Agent.v2()` callbacks receive a version-typed `V2ConnectionTo` with +high-level, command-only session helpers because prompt acceptance and inbound +traffic are independent. Session updates and interactive requests use typed +connection handlers. See [Protocol V2](https://agentclientprotocol.github.io/rust-sdk/protocol-v2.html#high-level-v2-sessions). +MCP attachment and proxy-session helpers remain v1-only. + ## MCP Server Attachment The runtime-agnostic `mcp_server` module can build and directly serve standalone @@ -47,6 +55,7 @@ the `with_mcp_server` builder methods requires `unstable_mcp_over_acp`. Attached servers are advertised with native `McpServer::Acp` declarations and communicate through `mcp/connect`, `mcp/message`, and `mcp/disconnect`. Use `agent-client-protocol-polyfill` immediately before an HTTP-capable agent. +These ACP attachment and proxy paths currently use stable protocol v1. ## Learning More diff --git a/src/agent-client-protocol/src/concepts/sessions.rs b/src/agent-client-protocol/src/concepts/sessions.rs index 9473e0c..dc08ced 100644 --- a/src/agent-client-protocol/src/concepts/sessions.rs +++ b/src/agent-client-protocol/src/concepts/sessions.rs @@ -4,6 +4,13 @@ //! session, you can send prompts, receive responses, and the agent maintains //! context across turns. //! +//! The examples below use the stable protocol v1 `SessionBuilder` and +//! `ActiveSession`. With the `unstable_protocol_v2` feature, callbacks created +//! through `Client.v2()` receive `V2ConnectionTo` and its `build_session*`, +//! `V2SessionBuilder`, and command-only `V2Session` APIs. V2 +//! prompt responses acknowledge acceptance independently; receive session-wide +//! updates and interactive requests through typed connection handlers. +//! //! # Creating a Session //! //! Use the session builder to create a new session: diff --git a/src/agent-client-protocol/src/jsonrpc.rs b/src/agent-client-protocol/src/jsonrpc.rs index 499e795..e5fbd23 100644 --- a/src/agent-client-protocol/src/jsonrpc.rs +++ b/src/agent-client-protocol/src/jsonrpc.rs @@ -11,6 +11,7 @@ use serde::{Deserialize, Serialize}; use std::any::TypeId; use std::collections::HashMap; use std::fmt::Debug; +use std::marker::PhantomData; use std::panic::Location; use std::pin::pin; use std::sync::{ @@ -631,6 +632,78 @@ where } } +/// Selects the connection context exposed by a [`Builder`]'s callbacks. +/// +/// This trait is an implementation detail of the typed builder aliases. It is +/// public so the callback connection type remains expressible in public API +/// signatures. +#[doc(hidden)] +#[allow(private_bounds)] +pub trait ConnectionContext: connection_context::Sealed + Send + Sync + 'static { + /// The connection type exposed to callbacks for `Counterpart`. + type Connection: Clone + Send + Sync + 'static; +} + +mod connection_context { + use super::{ConnectionContext, ConnectionTo, Role}; + + pub trait Sealed { + fn from_raw( + connection: ConnectionTo, + ) -> ::Connection + where + Self: ConnectionContext; + } + + pub(crate) fn from_raw( + connection: ConnectionTo, + ) -> Context::Connection { + ::from_raw(connection) + } +} + +/// The default callback context used by stable and low-level builders. +#[doc(hidden)] +#[derive(Copy, Clone, Debug, Default)] +pub struct RawConnectionContext; + +impl connection_context::Sealed for RawConnectionContext { + fn from_raw( + connection: ConnectionTo, + ) -> ::Connection { + connection + } +} + +impl ConnectionContext for RawConnectionContext { + type Connection = ConnectionTo; +} + +/// The callback context used by ACP protocol v2 builders. +#[cfg(feature = "unstable_protocol_v2")] +#[doc(hidden)] +#[derive(Copy, Clone, Debug, Default)] +pub struct V2ConnectionContext; + +#[cfg(feature = "unstable_protocol_v2")] +impl connection_context::Sealed for V2ConnectionContext { + fn from_raw( + connection: ConnectionTo, + ) -> ::Connection { + V2ConnectionTo { inner: connection } + } +} + +#[cfg(feature = "unstable_protocol_v2")] +impl ConnectionContext for V2ConnectionContext { + type Connection = V2ConnectionTo; +} + +/// A JSON-RPC connection builder whose callbacks receive [`V2ConnectionTo`]. +#[cfg(feature = "unstable_protocol_v2")] +pub type V2Builder = + Builder; + /// A JSON-RPC connection that can act as either a server, client, or both. /// /// [`Builder`] provides a builder-style API for creating JSON-RPC servers and clients. @@ -928,11 +1001,17 @@ where /// ``` #[must_use] #[derive(Debug)] -pub struct Builder -where +pub struct Builder< + Host: Role, + Handler = NullHandler, + Runner = NullRun, + Close = NullClose, + Context = RawConnectionContext, +> where Handler: HandleDispatchFrom, Runner: RunWithConnectionTo, Close: HandleConnectionClose, + Context: ConnectionContext, { /// My role. host: Host, @@ -951,6 +1030,9 @@ where /// Handler run when the incoming transport reaches clean EOF. on_close: Close, + + /// Selects the connection type exposed to user callbacks. + context: PhantomData Context>, } fn default_protocol_mode() -> ProtocolMode { @@ -977,6 +1059,7 @@ impl Builder { runner: NullRun, protocol_mode: default_protocol_mode::(), on_close: NullClose, + context: PhantomData, } } } @@ -994,16 +1077,51 @@ where runner: NullRun, protocol_mode: default_protocol_mode::(), on_close: NullClose, + context: PhantomData, } } } +#[cfg(feature = "unstable_protocol_v2")] impl< Host: Role, Handler: HandleDispatchFrom, Runner: RunWithConnectionTo, Close: HandleConnectionClose, > Builder +{ + pub(crate) fn v2_agent(self) -> V2Builder { + Builder { + host: self.host, + name: self.name, + handler: self.handler, + runner: self.runner, + protocol_mode: ProtocolMode::v2_agent(), + on_close: self.on_close, + context: PhantomData, + } + } + + pub(crate) fn v2_client(self) -> V2Builder { + Builder { + host: self.host, + name: self.name, + handler: self.handler, + runner: self.runner, + protocol_mode: ProtocolMode::v2_client(), + on_close: self.on_close, + context: PhantomData, + } + } +} + +impl< + Host: Role, + Handler: HandleDispatchFrom, + Runner: RunWithConnectionTo, + Close: HandleConnectionClose, + Context: ConnectionContext, +> Builder { /// Set the "name" of this connection -- used only for debugging logs. pub fn name(mut self, name: impl ToString) -> Self { @@ -1021,18 +1139,6 @@ impl< self } - #[cfg(feature = "unstable_protocol_v2")] - pub(crate) fn v2_agent(mut self) -> Self { - self.protocol_mode = ProtocolMode::v2_agent(); - self - } - - #[cfg(feature = "unstable_protocol_v2")] - pub(crate) fn v2_client(mut self) -> Self { - self.protocol_mode = ProtocolMode::v2_client(); - self - } - /// Merge another [`Builder`] into this one. /// /// Prefer [`Self::on_receive_request`] or [`Self::on_receive_notification`]. @@ -1044,12 +1150,14 @@ impl< impl HandleDispatchFrom, impl RunWithConnectionTo, impl HandleConnectionClose, + Context, >, ) -> Builder< Host, impl HandleDispatchFrom, impl RunWithConnectionTo, impl HandleConnectionClose, + Context, > { let Builder { name: other_name, @@ -1057,6 +1165,7 @@ impl< runner: other_runner, protocol_mode: other_protocol_mode, on_close: other_on_close, + context: _, host: _, } = other; Builder { @@ -1069,6 +1178,7 @@ impl< runner: ChainRun::new(self.runner, other_runner), protocol_mode: self.protocol_mode.merge(other_protocol_mode), on_close: ChainedClose::new(self.on_close, other_on_close), + context: PhantomData, } } @@ -1079,7 +1189,7 @@ impl< pub fn with_handler( self, handler: impl HandleDispatchFrom, - ) -> Builder, Runner, Close> { + ) -> Builder, Runner, Close, Context> { Builder { host: self.host, name: self.name, @@ -1087,6 +1197,7 @@ impl< runner: self.runner, protocol_mode: self.protocol_mode, on_close: self.on_close, + context: PhantomData, } } @@ -1094,7 +1205,7 @@ impl< pub fn with_runner( self, runner: Run1, - ) -> Builder, Close> + ) -> Builder, Close, Context> where Run1: RunWithConnectionTo, { @@ -1105,6 +1216,7 @@ impl< runner: ChainRun::new(self.runner, runner), protocol_mode: self.protocol_mode, on_close: self.on_close, + context: PhantomData, } } @@ -1113,13 +1225,13 @@ impl< pub fn with_spawned( self, task: F, - ) -> Builder, Close> + ) -> Builder, Close, Context> where - F: FnOnce(ConnectionTo) -> Fut + Send, + F: FnOnce(Context::Connection) -> Fut + Send, Fut: Future> + Send, { let location = Location::caller(); - self.with_runner(SpawnedRun::new(location, task)) + self.with_runner(SpawnedRun::<_, Context>::new(location, task)) } /// Run a callback when the incoming transport reaches clean EOF. @@ -1132,9 +1244,10 @@ impl< /// /// Multiple callbacks run sequentially in registration order. All of them /// run even if an earlier callback fails, after which the first error is - /// returned. Pending requests are failed before callbacks begin, while - /// [`ConnectionTo::incoming_closed`] completes only after they finish. A - /// callback must therefore not await that close future itself. + /// returned. Pending requests are failed before callbacks begin, while the + /// selected connection context's `incoming_closed` future completes only + /// after they finish. A callback must therefore not await that close + /// future itself. /// /// This separation lets applications choose their cancellation policy. A /// callback can notify application-owned tasks and return `Ok(())` for @@ -1157,9 +1270,9 @@ impl< pub fn on_close( self, callback: F, - ) -> Builder> + ) -> Builder, Context> where - F: FnOnce(ConnectionTo) -> Fut + Send, + F: FnOnce(Context::Connection) -> Fut + Send, Fut: Future> + Send, { Builder { @@ -1168,7 +1281,8 @@ impl< handler: self.handler, runner: self.runner, protocol_mode: self.protocol_mode, - on_close: ChainedClose::new(self.on_close, CloseCallback::new(callback)), + on_close: ChainedClose::new(self.on_close, CloseCallback::<_, Context>::new(callback)), + context: PhantomData, } } @@ -1222,26 +1336,26 @@ impl< self, op: F, to_future_hack: ToFut, - ) -> Builder, Runner, Close> + ) -> Builder, Runner, Close, Context> where Host::Counterpart: HasPeer, Req: JsonRpcRequest, Notif: JsonRpcNotification, F: AsyncFnMut( Dispatch, - ConnectionTo, + Context::Connection, ) -> Result + Send, T: IntoHandled>, ToFut: Fn( &mut F, Dispatch, - ConnectionTo, + Context::Connection, ) -> crate::BoxFuture<'_, Result> + Send + Sync, { - let handler = MessageHandler::new( + let handler = MessageHandler::<_, _, _, _, _, _, Context>::new( self.host.counterpart(), self.host.counterpart(), op, @@ -1255,12 +1369,15 @@ impl< /// Your handler receives three arguments: /// 1. The request (type `Req`) /// 2. A [`Responder`] for sending the response - /// 3. A [`ConnectionTo`] for the peer that sent the request + /// 3. The builder-selected connection context for the peer that sent the + /// request (`ConnectionTo` by default, or `V2ConnectionTo` for a + /// `V2Builder`) /// /// The request context allows you to: /// - Send the response with [`Responder::respond`] - /// - Send notifications to the client with [`ConnectionTo::send_notification`] - /// - Send requests to the client with [`ConnectionTo::send_request`] + /// - Send notifications to the client with the context's + /// `send_notification` method + /// - Send requests to the client with the context's `send_request` method /// /// # Example /// @@ -1296,13 +1413,13 @@ impl< self, op: F, to_future_hack: ToFut, - ) -> Builder, Runner, Close> + ) -> Builder, Runner, Close, Context> where Host::Counterpart: HasPeer, F: AsyncFnMut( Req, Responder, - ConnectionTo, + Context::Connection, ) -> Result + Send, T: IntoHandled<(Req, Responder)>, @@ -1310,12 +1427,12 @@ impl< &mut F, Req, Responder, - ConnectionTo, + Context::Connection, ) -> crate::BoxFuture<'_, Result> + Send + Sync, { - let handler = RequestHandler::new( + let handler = RequestHandler::<_, _, _, _, _, Context>::new( self.host.counterpart(), self.host.counterpart(), op, @@ -1329,7 +1446,8 @@ impl< /// Notifications are fire-and-forget messages that don't expect a response. /// Your handler receives: /// 1. The notification (type `Notif`) - /// 2. A [`ConnectionTo`] for sending messages to the other side + /// 2. The builder-selected connection context for sending messages to the + /// other side /// /// Unlike request handlers, you cannot send a response (notifications don't have IDs), /// but you can still send your own requests and notifications using the context. @@ -1370,21 +1488,22 @@ impl< self, op: F, to_future_hack: ToFut, - ) -> Builder, Runner, Close> + ) -> Builder, Runner, Close, Context> where Host::Counterpart: HasPeer, Notif: JsonRpcNotification, - F: AsyncFnMut(Notif, ConnectionTo) -> Result + Send, - T: IntoHandled<(Notif, ConnectionTo)>, + F: AsyncFnMut(Notif, Context::Connection) -> Result + + Send, + T: IntoHandled<(Notif, Context::Connection)>, ToFut: Fn( &mut F, Notif, - ConnectionTo, + Context::Connection, ) -> crate::BoxFuture<'_, Result> + Send + Sync, { - let handler = NotificationHandler::new( + let handler = NotificationHandler::<_, _, _, _, _, Context>::new( self.host.counterpart(), self.host.counterpart(), op, @@ -1420,24 +1539,29 @@ impl< peer: Peer, op: F, to_future_hack: ToFut, - ) -> Builder, Runner, Close> + ) -> Builder, Runner, Close, Context> where Host::Counterpart: HasPeer, F: AsyncFnMut( Dispatch, - ConnectionTo, + Context::Connection, ) -> Result + Send, T: IntoHandled>, ToFut: Fn( &mut F, Dispatch, - ConnectionTo, + Context::Connection, ) -> crate::BoxFuture<'_, Result> + Send + Sync, { - let handler = MessageHandler::new(self.host.counterpart(), peer, op, to_future_hack); + let handler = MessageHandler::<_, _, _, _, _, _, Context>::new( + self.host.counterpart(), + peer, + op, + to_future_hack, + ); self.with_handler(handler) } @@ -1474,13 +1598,13 @@ impl< peer: Peer, op: F, to_future_hack: ToFut, - ) -> Builder, Runner, Close> + ) -> Builder, Runner, Close, Context> where Host::Counterpart: HasPeer, F: AsyncFnMut( Req, Responder, - ConnectionTo, + Context::Connection, ) -> Result + Send, T: IntoHandled<(Req, Responder)>, @@ -1488,12 +1612,17 @@ impl< &mut F, Req, Responder, - ConnectionTo, + Context::Connection, ) -> crate::BoxFuture<'_, Result> + Send + Sync, { - let handler = RequestHandler::new(self.host.counterpart(), peer, op, to_future_hack); + let handler = RequestHandler::<_, _, _, _, _, Context>::new( + self.host.counterpart(), + peer, + op, + to_future_hack, + ); self.with_handler(handler) } @@ -1517,20 +1646,26 @@ impl< peer: Peer, op: F, to_future_hack: ToFut, - ) -> Builder, Runner, Close> + ) -> Builder, Runner, Close, Context> where Host::Counterpart: HasPeer, - F: AsyncFnMut(Notif, ConnectionTo) -> Result + Send, - T: IntoHandled<(Notif, ConnectionTo)>, + F: AsyncFnMut(Notif, Context::Connection) -> Result + + Send, + T: IntoHandled<(Notif, Context::Connection)>, ToFut: Fn( &mut F, Notif, - ConnectionTo, + Context::Connection, ) -> crate::BoxFuture<'_, Result> + Send + Sync, { - let handler = NotificationHandler::new(self.host.counterpart(), peer, op, to_future_hack); + let handler = NotificationHandler::<_, _, _, _, _, Context>::new( + self.host.counterpart(), + peer, + op, + to_future_hack, + ); self.with_handler(handler) } @@ -1549,6 +1684,7 @@ impl< impl HandleDispatchFrom, impl RunWithConnectionTo, Close, + Context, > where Host::Counterpart: HasPeer + HasPeer, @@ -1602,25 +1738,26 @@ impl< self, transport: impl ConnectTo + 'static, ) -> Result<(), crate::Error> { - self.connect_with(transport, async move |cx| { + let (_, future) = self.into_connection_and_future(transport, async move |cx| { cx.incoming_closed().await; cx.drain_outgoing().await - }) - .await + }); + future.await } /// Run the connection until the provided closure completes. /// /// This drives the connection by: /// 1. Running your registered handlers in the background to process incoming messages - /// 2. Executing your `main_fn` closure with a [`ConnectionTo`] for sending requests/notifications + /// 2. Executing your `main_fn` closure with the builder-selected connection + /// context for sending requests and notifications /// /// The connection stays active until your `main_fn` returns, then shuts down. /// Clean incoming EOF fails every pending request and makes future /// requests fail immediately. It does not cancel unrelated work in - /// `main_fn`: that future may observe [`ConnectionTo::incoming_closed`], or - /// the builder can use [`on_close`](Self::on_close) to notify it or return - /// an error and stop it. + /// `main_fn`: that future may observe the context's `incoming_closed` + /// future, or the builder can use [`on_close`](Self::on_close) to notify it + /// or return an error and stop it. /// /// Use this mode when you need to initiate communication (send requests/notifications) /// in addition to responding to incoming messages. For server-only mode where you just @@ -1664,20 +1801,23 @@ impl< /// /// # Parameters /// - /// - `main_fn`: Your client logic. Receives a [`ConnectionTo`] for sending messages. + /// - `main_fn`: Your client logic. Receives the builder-selected connection + /// context for sending messages. /// /// # Errors /// /// Returns an error if a handler, background task, transport, or close /// callback fails, or if `main_fn` returns an error. Clean incoming EOF is - /// observable through [`ConnectionTo::incoming_closed`] and is not itself - /// an error in this mode. + /// observable through the context's `incoming_closed` future and is not + /// itself an error in this mode. pub async fn connect_with( self, transport: impl ConnectTo + 'static, - main_fn: impl AsyncFnOnce(ConnectionTo) -> Result, + main_fn: impl AsyncFnOnce(Context::Connection) -> Result, ) -> Result { - let (_, future) = self.into_connection_and_future(transport, main_fn); + let (_, future) = self.into_connection_and_future(transport, async move |connection| { + main_fn(connection_context::from_raw::(connection)).await + }); future.await } @@ -1697,6 +1837,7 @@ impl< host: me, protocol_mode, on_close, + context: _, } = self; let (outgoing_tx, outgoing_rx) = mpsc::unbounded(); @@ -1727,6 +1868,7 @@ impl< dynamic_handler_tx, transport_completion, pending_replies.registrar(), + protocol_mode, ); let spawn_result = connection.spawn(async move { let result = transport_future.await; @@ -1797,12 +1939,13 @@ impl< } } -impl ConnectTo for Builder +impl ConnectTo for Builder where R: Role, H: HandleDispatchFrom + 'static, Run: RunWithConnectionTo + 'static, Close: HandleConnectionClose + 'static, + Context: ConnectionContext, { async fn connect_to(self, client: impl ConnectTo) -> Result<(), crate::Error> { Builder::connect_to(self, client).await @@ -2882,6 +3025,210 @@ impl IntoHandled for Handled { } } +/// A protocol-v2 connection context. +/// +/// Values of this type are supplied to callbacks registered on a +/// [`V2Builder`]. It exposes the general connection operations that are valid +/// for protocol v2 while keeping version-specific high-level helpers for other +/// protocol versions out of the typed context. The generic JSON-RPC send +/// methods remain intentionally schema-agnostic. +/// +/// This is a thin, cheaply cloneable handle to the underlying JSON-RPC +/// connection. It intentionally does not implement [`Deref`](std::ops::Deref) +/// to [`ConnectionTo`]. +#[cfg(feature = "unstable_protocol_v2")] +#[derive(Clone, Debug)] +pub struct V2ConnectionTo { + inner: ConnectionTo, +} + +#[cfg(feature = "unstable_protocol_v2")] +impl V2ConnectionTo { + /// Return the counterpart role this connection is talking to. + pub fn counterpart(&self) -> Counterpart { + self.inner.counterpart() + } + + /// Wait until the incoming transport reaches clean EOF. + pub async fn incoming_closed(&self) { + self.inner.incoming_closed().await; + } + + /// Return whether clean incoming-EOF processing has completed. + #[must_use] + pub fn is_incoming_closed(&self) -> bool { + self.inner.is_incoming_closed() + } + + /// Spawn a task that runs for as long as the JSON-RPC connection is served. + #[track_caller] + pub fn spawn( + &self, + task: impl IntoFuture, IntoFuture: Send + 'static>, + ) -> Result<(), crate::Error> { + self.inner.spawn(task) + } + + /// Spawn a JSON-RPC connection in the background. + /// + /// The returned connection context is selected by `builder`; spawning a + /// [`V2Builder`] therefore returns another [`V2ConnectionTo`]. + /// + /// ```no_run + /// # use agent_client_protocol::{ + /// # Agent, Client, ConnectTo, Error, V2ConnectionTo, + /// # }; + /// # fn example( + /// # connection: V2ConnectionTo, + /// # transport: impl ConnectTo + 'static, + /// # ) -> Result<(), Error> { + /// let child: V2ConnectionTo = + /// connection.spawn_connection(Client.v2(), transport)?; + /// # drop(child); + /// # Ok(()) + /// # } + /// ``` + #[track_caller] + pub fn spawn_connection( + &self, + builder: Builder< + R, + impl HandleDispatchFrom + 'static, + impl RunWithConnectionTo + 'static, + impl HandleConnectionClose + 'static, + Context, + >, + transport: impl ConnectTo + 'static, + ) -> Result, crate::Error> { + self.inner.spawn_connection_with_context(builder, transport) + } + + /// Send a request or notification and forward its response appropriately. + pub fn send_proxied_message, Notif: JsonRpcNotification>( + &self, + message: Dispatch, + ) -> Result<(), crate::Error> + where + Counterpart: HasPeer, + { + self.inner.send_proxied_message(message) + } + + /// Send a request or notification to a specific peer and forward its + /// response appropriately. + pub fn send_proxied_message_to< + Peer: Role, + Req: JsonRpcRequest, + Notif: JsonRpcNotification, + >( + &self, + peer: Peer, + message: Dispatch, + ) -> Result<(), crate::Error> + where + Counterpart: HasPeer, + { + self.inner.send_proxied_message_to(peer, message) + } + + /// Send an outgoing request to the default counterpart peer. + pub fn send_request(&self, request: Req) -> SentRequest + where + Counterpart: HasPeer, + { + self.inner.send_request(request) + } + + /// Send an outgoing request to a specific peer. + pub fn send_request_to( + &self, + peer: Peer, + request: Req, + ) -> SentRequest + where + Counterpart: HasPeer, + { + self.inner.send_request_to(peer, request) + } + + /// Send an outgoing notification to the default counterpart peer. + pub fn send_notification( + &self, + notification: N, + ) -> Result<(), crate::Error> + where + Counterpart: HasPeer, + { + self.inner.send_notification(notification) + } + + /// Send an outgoing notification to a specific peer. + pub fn send_notification_to( + &self, + peer: Peer, + notification: N, + ) -> Result<(), crate::Error> + where + Counterpart: HasPeer, + { + self.inner.send_notification_to(peer, notification) + } + + /// Send a `$/cancel_request` notification to the default counterpart peer. + pub fn send_cancel_request( + &self, + request_id: impl Into, + ) -> Result<(), crate::Error> + where + Counterpart: HasPeer, + { + self.inner.send_cancel_request(request_id) + } + + /// Send a `$/cancel_request` notification to a specific peer. + pub fn send_cancel_request_to( + &self, + peer: Peer, + request_id: impl Into, + ) -> Result<(), crate::Error> + where + Counterpart: HasPeer, + { + self.inner.send_cancel_request_to(peer, request_id) + } + + /// Register a low-level dynamic message handler. + /// + /// Dynamic handlers use the version-neutral [`HandleDispatchFrom`] trait, + /// so their callback receives the underlying [`ConnectionTo`]. Prefer + /// typed handlers on [`V2Builder`] when registration can happen before the + /// connection starts. + /// + /// ```no_run + /// # use agent_client_protocol::{ + /// # Agent, Client, ConnectTo, DynamicHandlerGuard, Error, NullHandler, + /// # }; + /// # async fn example( + /// # transport: impl ConnectTo + 'static, + /// # ) -> Result<(), Error> { + /// Client.v2().connect_with(transport, async |connection| { + /// let guard: DynamicHandlerGuard = + /// connection.add_dynamic_handler(NullHandler)?; + /// + /// // Keep `guard` alive for as long as the modal handler is needed. + /// drop(guard); + /// Ok(()) + /// }).await + /// # } + /// ``` + pub fn add_dynamic_handler( + &self, + handler: impl HandleDispatchFrom + 'static, + ) -> Result, crate::Error> { + self.inner.add_dynamic_handler(handler) + } +} + /// Connection context for sending messages and spawning tasks. /// /// This is the primary handle for interacting with the JSON-RPC connection from @@ -2912,6 +3259,14 @@ pub struct ConnectionTo { dynamic_handler_tx: mpsc::UnboundedSender>, transport_completion: SharedTransportCompletion, pending_replies: PendingRepliesRegistrar, + #[cfg_attr( + not(feature = "unstable_protocol_v2"), + allow( + dead_code, + reason = "retained so ConnectionTo has one constructor shape" + ) + )] + protocol_mode: ProtocolMode, incoming_closed: IncomingClosed, } @@ -3058,6 +3413,7 @@ impl ConnectionTo { dynamic_handler_tx: mpsc::UnboundedSender>, transport_completion: SharedTransportCompletion, pending_replies: PendingRepliesRegistrar, + protocol_mode: ProtocolMode, ) -> Self { Self { counterpart, @@ -3066,10 +3422,16 @@ impl ConnectionTo { dynamic_handler_tx, transport_completion, pending_replies, + protocol_mode, incoming_closed: IncomingClosed::new(), } } + #[cfg(feature = "unstable_protocol_v2")] + pub(crate) fn acp_protocol_version(&self) -> Option { + self.protocol_mode.api_protocol_version() + } + /// Return the counterpart role this connection is talking to. pub fn counterpart(&self) -> Counterpart { self.counterpart.clone() @@ -3178,7 +3540,8 @@ impl ConnectionTo { Task::new(location, task).spawn(&self.task_tx) } - /// Spawn a JSON-RPC connection in the background and return a [`ConnectionTo`] for sending messages to it. + /// Spawn a JSON-RPC connection in the background and return a raw + /// [`ConnectionTo`] for it. /// /// This is useful for creating multiple connections that communicate with each other, /// such as implementing proxy patterns or connecting to multiple backend services. @@ -3190,7 +3553,13 @@ impl ConnectionTo { /// /// # Returns /// - /// A `ConnectionTo` that you can use to send requests and notifications to the spawned connection. + /// The child builder may select any callback context. For example, this + /// method can spawn a `V2Builder`, whose callbacks receive + /// `V2ConnectionTo`, while preserving this method's existing raw return + /// type and single explicit role parameter. + /// + /// When a raw parent also needs the builder-selected child handle, use the + /// protocol-v2 `spawn_connection_with_context` method. /// /// # Example: Proxying to a backend connection /// @@ -3206,7 +3575,7 @@ impl ConnectionTo { /// }, agent_client_protocol::on_receive_request!()); /// /// // Spawn it and get a context to send requests to it - /// let backend_connection = cx.spawn_connection(backend, MockTransport)?; + /// let backend_connection = cx.spawn_connection::(backend, MockTransport)?; /// /// // Now you can forward requests to the backend /// let response = backend_connection.send_request(MyRequest {}).block_task().await?; @@ -3221,6 +3590,63 @@ impl ConnectionTo { impl HandleDispatchFrom + 'static, impl RunWithConnectionTo + 'static, impl HandleConnectionClose + 'static, + impl ConnectionContext, + >, + transport: impl ConnectTo + 'static, + ) -> Result, crate::Error> { + self.spawn_connection_raw(builder, transport) + } + + /// Spawn a JSON-RPC connection and return the connection context selected + /// by its builder. + /// + /// This is the low-level counterpart to + /// [`V2ConnectionTo::spawn_connection`] for code that intentionally works + /// with a raw [`ConnectionTo`], such as custom [`HandleDispatchFrom`] or + /// [`RunWithConnectionTo`] implementations. Prefer [`Self::spawn_connection`] + /// when a raw child handle is sufficient. + /// + /// ```no_run + /// # use agent_client_protocol::{ + /// # Agent, Client, ConnectTo, ConnectionTo, Error, UntypedRole, + /// # V2ConnectionTo, + /// # }; + /// # fn example( + /// # connection: ConnectionTo, + /// # transport: impl ConnectTo + 'static, + /// # ) -> Result<(), Error> { + /// let child: V2ConnectionTo = + /// connection.spawn_connection_with_context(Client.v2(), transport)?; + /// # drop(child); + /// # Ok(()) + /// # } + /// ``` + #[cfg(feature = "unstable_protocol_v2")] + #[track_caller] + pub fn spawn_connection_with_context( + &self, + builder: Builder< + R, + impl HandleDispatchFrom + 'static, + impl RunWithConnectionTo + 'static, + impl HandleConnectionClose + 'static, + Context, + >, + transport: impl ConnectTo + 'static, + ) -> Result, crate::Error> { + let connection = self.spawn_connection_raw(builder, transport)?; + Ok(connection_context::from_raw::(connection)) + } + + #[track_caller] + fn spawn_connection_raw( + &self, + builder: Builder< + R, + impl HandleDispatchFrom + 'static, + impl RunWithConnectionTo + 'static, + impl HandleConnectionClose + 'static, + Context, >, transport: impl ConnectTo + 'static, ) -> Result, crate::Error> { @@ -3269,7 +3695,7 @@ impl ConnectionTo { { match message { Dispatch::Request(request, responder) => self - .send_request_to(peer, request) + .send_ordered_request_to(peer, request) .forward_response_to(responder), Dispatch::Notification(notification) => { // `$/cancel_request` is connection-scoped: its `requestId` was @@ -3362,6 +3788,35 @@ impl ConnectionTo { peer: Peer, request: Req, ) -> SentRequest + where + Counterpart: HasPeer, + { + self.send_request_to_with_ordering(peer, request, false) + } + + /// Send a request whose callback must run before later inbound messages. + /// + /// The ordering marker is installed before the request enters the outgoing + /// queue, closing the race between a fast peer response and the immediate + /// [`SentRequest::on_receiving_result`] call. Callers must consume the + /// returned request with a callback-style method without yielding. + pub(crate) fn send_ordered_request_to( + &self, + peer: Peer, + request: Req, + ) -> SentRequest + where + Counterpart: HasPeer, + { + self.send_request_to_with_ordering(peer, request, true) + } + + fn send_request_to_with_ordering( + &self, + peer: Peer, + request: Req, + ordered: bool, + ) -> SentRequest where Counterpart: HasPeer, { @@ -3369,6 +3824,9 @@ impl ConnectionTo { let id = RequestId::Str(uuid::Uuid::new_v4().to_string()); let (response_tx, response_rx) = oneshot::channel(); let response_ordering = ResponseOrdering::default(); + if ordered { + response_ordering.mark_ordered(); + } let role_id = peer.role_id(); let remote_style = self.counterpart.remote_style(peer); let cancellation = @@ -5742,6 +6200,149 @@ impl ConnectTo for Channel { mod tests { use super::*; + #[cfg(feature = "unstable_protocol_v2")] + fn connection_with_task_receiver() -> ( + ConnectionTo, + mpsc::UnboundedReceiver, + ) { + let (message_tx, _message_rx) = mpsc::unbounded(); + let (task_tx, task_rx) = mpsc::unbounded(); + let (dynamic_handler_tx, _dynamic_handler_rx) = mpsc::unbounded(); + let transport_completion: SharedTransportCompletion = + future::ready(Ok::<(), crate::Error>(())).boxed().shared(); + let pending_replies = PendingReplies::default(); + + ( + ConnectionTo::new( + crate::role::UntypedRole, + message_tx, + task_tx, + dynamic_handler_tx, + transport_completion, + pending_replies.registrar(), + ProtocolMode::disabled(), + ), + task_rx, + ) + } + + #[cfg(feature = "unstable_protocol_v2")] + #[test] + fn v2_builder_exposes_typed_context_to_user_callbacks() { + fn assert_v2_context(_connection: &V2ConnectionTo) {} + + let _builder = Client + .v2() + .on_receive_request( + async |_request: UntypedMessage, _responder, connection| { + assert_v2_context(&connection); + Ok(()) + }, + crate::on_receive_request!(), + ) + .on_receive_notification( + async |_notification: UntypedMessage, connection| { + assert_v2_context(&connection); + Ok(()) + }, + crate::on_receive_notification!(), + ) + .on_receive_dispatch( + async |_dispatch: Dispatch, connection| { + assert_v2_context(&connection); + Ok(()) + }, + crate::on_receive_dispatch!(), + ) + .on_receive_request_from( + Agent, + async |_request: UntypedMessage, _responder, connection| { + assert_v2_context(&connection); + Ok(()) + }, + crate::on_receive_request!(), + ) + .on_receive_notification_from( + Agent, + async |_notification: UntypedMessage, connection| { + assert_v2_context(&connection); + Ok(()) + }, + crate::on_receive_notification!(), + ) + .on_receive_dispatch_from( + Agent, + async |_dispatch: Dispatch, connection| { + assert_v2_context(&connection); + Ok(()) + }, + crate::on_receive_dispatch!(), + ) + .with_spawned(async |connection| { + assert_v2_context(&connection); + Ok(()) + }) + .on_close(async |connection| { + assert_v2_context(&connection); + Ok(()) + }); + } + + #[cfg(feature = "unstable_protocol_v2")] + #[test] + fn raw_connection_spawns_v2_builder_with_typed_child_callback() { + let (parent, mut task_rx) = connection_with_task_receiver(); + let (transport, _peer) = Channel::duplex(); + let (callback_tx, callback_rx) = oneshot::channel(); + + let child: ConnectionTo = parent + .spawn_connection::( + Client + .v2() + .with_spawned(async move |_connection: V2ConnectionTo| { + callback_tx.send(()).map_err(|()| { + crate::util::internal_error("typed child callback receiver was dropped") + }) + }), + transport, + ) + .expect("v2 child connection should be spawned"); + + let task = futures::FutureExt::now_or_never(futures::StreamExt::next(&mut task_rx)) + .expect("child connection task should already be queued") + .expect("parent task queue should remain open"); + futures::executor::block_on(async { + match future::select(Box::pin(task.run_for_test()), Box::pin(callback_rx)).await { + Either::Right((Ok(()), child_task)) => drop(child_task), + Either::Right((Err(error), _)) => { + panic!("typed child callback sender was dropped: {error}") + } + Either::Left((result, _)) => { + panic!("child connection stopped before its typed callback ran: {result:?}") + } + } + }); + + drop(child); + } + + #[cfg(feature = "unstable_protocol_v2")] + #[test] + fn raw_connection_can_return_v2_context_for_spawned_builder() { + let (parent, mut task_rx) = connection_with_task_receiver(); + let (transport, _peer) = Channel::duplex(); + + let child: V2ConnectionTo = parent + .spawn_connection_with_context(Client.v2(), transport) + .expect("v2 child connection should be spawned"); + + let child_task = futures::FutureExt::now_or_never(futures::StreamExt::next(&mut task_rx)) + .expect("child connection task should already be queued") + .expect("parent task queue should remain open"); + + drop((child, child_task)); + } + fn connection_with_dynamic_handler_receiver() -> ( ConnectionTo, mpsc::UnboundedReceiver>, @@ -5761,18 +6362,136 @@ mod tests { dynamic_handler_tx, transport_completion, pending_replies.registrar(), + ProtocolMode::disabled(), ), dynamic_handler_rx, ) } - fn next_dynamic_handler_message( - receiver: &mut mpsc::UnboundedReceiver>, - ) -> Option> { + #[test] + fn ordered_request_is_marked_before_entering_outgoing_queue() { + let (message_tx, mut message_rx) = mpsc::unbounded(); + let (task_tx, mut task_rx) = mpsc::unbounded(); + let (dynamic_handler_tx, _dynamic_handler_rx) = mpsc::unbounded(); + let transport_completion: SharedTransportCompletion = + future::ready(Ok::<(), crate::Error>(())).boxed().shared(); + let pending_replies = PendingReplies::default(); + let connection = ConnectionTo::new( + crate::role::UntypedRole, + message_tx, + task_tx, + dynamic_handler_tx, + transport_completion, + pending_replies.registrar(), + ProtocolMode::disabled(), + ); + + let sent = connection.send_ordered_request_to( + crate::role::UntypedRole, + UntypedMessage::new("ordered", serde_json::json!({})) + .expect("test request should serialize"), + ); + let request_id = sent.id().clone(); + let message = futures::FutureExt::now_or_never(futures::StreamExt::next(&mut message_rx)) + .expect("outgoing request should already be queued") + .expect("outgoing request queue should remain open"); + let OutgoingMessage::Request { id, .. } = message else { + panic!("expected an outgoing request"); + }; + assert_eq!(id, request_id); + + let pending_reply = pending_replies + .remove(&request_id) + .expect("the request should have a pending reply"); + assert!( + pending_reply.ordering.is_ordered(), + "the response ordering barrier must be installed before publication" + ); + + // Route the response before the callback is registered. The pre-set + // ordering marker must hold dispatch until the callback task is + // subsequently installed and completes. + let (dispatch, response_dispatch) = incoming_actor::dispatch_from_response( + request_id, + pending_reply, + Ok(serde_json::json!({"ok": true})), + ); + let Dispatch::Response(result, router) = dispatch else { + panic!("expected a response dispatch"); + }; + router + .route_with_result(result) + .expect("response should route to the pending request"); + let acknowledgment = response_dispatch + .complete() + .expect("an ordered response should require acknowledgment"); + + let callback_ran = Arc::new(AtomicBool::new(false)); + sent.on_receiving_result({ + let callback_ran = callback_ran.clone(); + async move |result| { + assert_eq!(result?, serde_json::json!({"ok": true})); + callback_ran.store(true, Ordering::Release); + Ok(()) + } + }) + .expect("ordered callback should be scheduled"); + + let task = futures::FutureExt::now_or_never(futures::StreamExt::next(&mut task_rx)) + .expect("callback task should already be queued") + .expect("callback task queue should remain open"); + futures::executor::block_on(task.run_for_test()).expect("callback task should succeed"); + futures::executor::block_on(acknowledgment) + .expect("callback completion should acknowledge dispatch"); + assert!(callback_ran.load(Ordering::Acquire)); + } + + fn next_dynamic_handler_message( + receiver: &mut mpsc::UnboundedReceiver>, + ) -> Option> { futures::FutureExt::now_or_never(futures::StreamExt::next(receiver)) .expect("dynamic-handler receiver should be ready") } + #[cfg(feature = "unstable_protocol_v2")] + #[test] + fn v2_dynamic_handler_guard_registers_and_removes_handler() { + let (message_tx, _message_rx) = mpsc::unbounded(); + let (task_tx, _task_rx) = mpsc::unbounded(); + let (dynamic_handler_tx, mut dynamic_handler_rx) = mpsc::unbounded(); + let transport_completion: SharedTransportCompletion = + future::ready(Ok::<(), crate::Error>(())).boxed().shared(); + let pending_replies = PendingReplies::default(); + let connection = V2ConnectionTo { + inner: ConnectionTo::new( + Agent, + message_tx, + task_tx, + dynamic_handler_tx, + transport_completion, + pending_replies.registrar(), + ProtocolMode::v2_client(), + ), + }; + + let guard = connection + .add_dynamic_handler(NullHandler) + .expect("v2 dynamic handler should register"); + let added_uuid = match next_dynamic_handler_message(&mut dynamic_handler_rx) { + Some(DynamicHandlerMessage::AddDynamicHandler(uuid, _)) => uuid, + other => panic!("expected v2 handler registration, got {other:?}"), + }; + + drop(guard); + + match next_dynamic_handler_message(&mut dynamic_handler_rx) { + Some(DynamicHandlerMessage::RemoveDynamicHandler(uuid)) => { + assert_eq!(uuid, added_uuid); + } + other => panic!("expected v2 handler removal, got {other:?}"), + } + } + #[test] fn dropping_dynamic_handler_guard_unregisters_handler() { let (connection, mut receiver) = connection_with_dynamic_handler_receiver(); diff --git a/src/agent-client-protocol/src/jsonrpc/close.rs b/src/agent-client-protocol/src/jsonrpc/close.rs index 22747f3..dc61477 100644 --- a/src/agent-client-protocol/src/jsonrpc/close.rs +++ b/src/agent-client-protocol/src/jsonrpc/close.rs @@ -2,8 +2,13 @@ use std::fmt::Debug; use std::future::Future; +use std::marker::PhantomData; -use crate::{ConnectionTo, role::Role}; +use crate::{ + ConnectionTo, + jsonrpc::{ConnectionContext, RawConnectionContext, connection_context}, + role::Role, +}; /// A handler that runs after the incoming transport reaches clean EOF. /// @@ -32,17 +37,21 @@ impl HandleConnectionClose for NullClose { } } -pub(crate) struct CloseCallback { +pub(crate) struct CloseCallback { callback: F, + context: PhantomData Context>, } -impl CloseCallback { +impl CloseCallback { pub(crate) fn new(callback: F) -> Self { - Self { callback } + Self { + callback, + context: PhantomData, + } } } -impl Debug for CloseCallback { +impl Debug for CloseCallback { fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { formatter .debug_struct("CloseCallback") @@ -50,17 +59,18 @@ impl Debug for CloseCallback { } } -impl HandleConnectionClose for CloseCallback +impl HandleConnectionClose for CloseCallback where Counterpart: Role, - F: FnOnce(ConnectionTo) -> Fut + Send, + Context: ConnectionContext, + F: FnOnce(Context::Connection) -> Fut + Send, Fut: Future> + Send, { async fn handle_connection_close( self, connection: ConnectionTo, ) -> Result<(), crate::Error> { - let result = (self.callback)(connection).await; + let result = (self.callback)(connection_context::from_raw::(connection)).await; if let Err(error) = &result { tracing::warn!(?error, "Connection close callback failed"); } diff --git a/src/agent-client-protocol/src/jsonrpc/handlers.rs b/src/agent-client-protocol/src/jsonrpc/handlers.rs index cce477e..efc904b 100644 --- a/src/agent-client-protocol/src/jsonrpc/handlers.rs +++ b/src/agent-client-protocol/src/jsonrpc/handlers.rs @@ -1,4 +1,7 @@ -use crate::jsonrpc::{HandleDispatchFrom, Handled, IntoHandled, JsonRpcResponse}; +use crate::jsonrpc::{ + ConnectionContext, HandleDispatchFrom, Handled, IntoHandled, JsonRpcResponse, + RawConnectionContext, connection_context, +}; use crate::role::{HasPeer, Role, handle_incoming_dispatch}; use crate::{ConnectionTo, Dispatch, JsonRpcNotification, JsonRpcRequest, UntypedMessage}; @@ -41,16 +44,17 @@ pub struct RequestHandler< Req: JsonRpcRequest = UntypedMessage, F = (), ToFut = (), + Context = RawConnectionContext, > { counterpart: Counterpart, peer: Peer, handler: F, to_future_hack: ToFut, - phantom: PhantomData, + phantom: PhantomData, } -impl - RequestHandler +impl + RequestHandler { /// Creates a new request handler pub fn new(counterpart: Counterpart, peer: Peer, handler: F, to_future_hack: ToFut) -> Self { @@ -64,15 +68,16 @@ impl } } -impl HandleDispatchFrom - for RequestHandler +impl HandleDispatchFrom + for RequestHandler where Counterpart: HasPeer, Req: JsonRpcRequest, + Context: ConnectionContext, F: AsyncFnMut( Req, Responder, - ConnectionTo, + Context::Connection, ) -> Result + Send, T: crate::IntoHandled<(Req, Responder)>, @@ -80,7 +85,7 @@ where &mut F, Req, Responder, - ConnectionTo, + Context::Connection, ) -> crate::BoxFuture<'_, Result> + Send + Sync, @@ -119,7 +124,7 @@ where &mut self.handler, req, typed_responder, - connection, + connection_context::from_raw::(connection), ) .await?; match result.into_handled() { @@ -175,16 +180,17 @@ pub struct NotificationHandler< Notif: JsonRpcNotification = UntypedMessage, F = (), ToFut = (), + Context = RawConnectionContext, > { counterpart: Counterpart, peer: Peer, handler: F, to_future_hack: ToFut, - phantom: PhantomData, + phantom: PhantomData, } -impl - NotificationHandler +impl + NotificationHandler { /// Creates a new notification handler pub fn new(counterpart: Counterpart, peer: Peer, handler: F, to_future_hack: ToFut) -> Self { @@ -198,17 +204,18 @@ impl } } -impl HandleDispatchFrom - for NotificationHandler +impl HandleDispatchFrom + for NotificationHandler where Counterpart: HasPeer, Notif: JsonRpcNotification, - F: AsyncFnMut(Notif, ConnectionTo) -> Result + Send, - T: crate::IntoHandled<(Notif, ConnectionTo)>, + Context: ConnectionContext, + F: AsyncFnMut(Notif, Context::Connection) -> Result + Send, + T: crate::IntoHandled<(Notif, Context::Connection)>, ToFut: Fn( &mut F, Notif, - ConnectionTo, + Context::Connection, ) -> crate::BoxFuture<'_, Result> + Send + Sync, @@ -242,9 +249,12 @@ where ?notif, "NotificationHandler::handle_notification: parse completed" ); - let result = - (self.to_future_hack)(&mut self.handler, notif, connection) - .await?; + let result = (self.to_future_hack)( + &mut self.handler, + notif, + connection_context::from_raw::(connection), + ) + .await?; match result.into_handled() { Handled::Yes => Ok(Handled::Yes), Handled::No { @@ -298,16 +308,24 @@ pub struct MessageHandler< Notif: JsonRpcNotification = UntypedMessage, F = (), ToFut = (), + Context = RawConnectionContext, > { counterpart: Counterpart, peer: Peer, handler: F, to_future_hack: ToFut, - phantom: PhantomData)>, + phantom: PhantomData, Context)>, } -impl - MessageHandler +impl< + Counterpart: Role, + Peer: Role, + Req: JsonRpcRequest, + Notif: JsonRpcNotification, + F, + ToFut, + Context, +> MessageHandler { /// Creates a new message handler pub fn new(counterpart: Counterpart, peer: Peer, handler: F, to_future_hack: ToFut) -> Self { @@ -321,17 +339,30 @@ impl - HandleDispatchFrom for MessageHandler +impl< + Counterpart: Role, + Peer: Role, + Req: JsonRpcRequest, + Notif: JsonRpcNotification, + F, + T, + ToFut, + Context, +> HandleDispatchFrom + for MessageHandler where Counterpart: HasPeer, - F: AsyncFnMut(Dispatch, ConnectionTo) -> Result + Context: ConnectionContext, + F: AsyncFnMut( + Dispatch, + Context::Connection, + ) -> Result + Send, T: IntoHandled>, ToFut: Fn( &mut F, Dispatch, - ConnectionTo, + Context::Connection, ) -> crate::BoxFuture<'_, Result> + Send + Sync, @@ -356,9 +387,12 @@ where connection, async |dispatch, connection| match dispatch.into_typed_dispatch::()? { Ok(typed_dispatch) => { - let result = - (self.to_future_hack)(&mut self.handler, typed_dispatch, connection) - .await?; + let result = (self.to_future_hack)( + &mut self.handler, + typed_dispatch, + connection_context::from_raw::(connection), + ) + .await?; match result.into_handled() { Handled::Yes => Ok(Handled::Yes), Handled::No { diff --git a/src/agent-client-protocol/src/jsonrpc/incoming_actor.rs b/src/agent-client-protocol/src/jsonrpc/incoming_actor.rs index 6ead0c0..6291e3c 100644 --- a/src/agent-client-protocol/src/jsonrpc/incoming_actor.rs +++ b/src/agent-client-protocol/src/jsonrpc/incoming_actor.rs @@ -465,7 +465,7 @@ fn dispatch_from_message( /// This allows handlers to intercept and process responses before they reach /// the awaiting code. The default behavior is to forward the response to the /// local awaiter via the oneshot channel. -fn dispatch_from_response( +pub(super) fn dispatch_from_response( id: RequestId, pending_reply: PendingReply, result: Result, diff --git a/src/agent-client-protocol/src/jsonrpc/protocol_compat.rs b/src/agent-client-protocol/src/jsonrpc/protocol_compat.rs index 369d4e1..b5c788f 100644 --- a/src/agent-client-protocol/src/jsonrpc/protocol_compat.rs +++ b/src/agent-client-protocol/src/jsonrpc/protocol_compat.rs @@ -139,6 +139,13 @@ mod imp { } } } + + pub(crate) fn api_protocol_version(self) -> Option { + match self { + Self::Disabled => None, + Self::Acp(mode) => Some(mode.api.as_protocol_version()), + } + } } #[derive(Clone, Debug)] diff --git a/src/agent-client-protocol/src/jsonrpc/run.rs b/src/agent-client-protocol/src/jsonrpc/run.rs index ff07f62..e8adb79 100644 --- a/src/agent-client-protocol/src/jsonrpc/run.rs +++ b/src/agent-client-protocol/src/jsonrpc/run.rs @@ -5,8 +5,13 @@ //! channels and invoke user-provided closures. use std::future::Future; +use std::marker::PhantomData; -use crate::{ConnectionTo, role::Role}; +use crate::{ + ConnectionTo, + jsonrpc::{ConnectionContext, RawConnectionContext, connection_context}, + role::Role, +}; /// A background task that runs alongside a connection. /// @@ -68,22 +73,28 @@ where } /// A RunIn created from a closure via [`with_spawned`](crate::Builder::with_spawned). -pub struct SpawnedRun { +pub struct SpawnedRun { task_fn: F, location: &'static std::panic::Location<'static>, + context: PhantomData Context>, } -impl SpawnedRun { +impl SpawnedRun { /// Create a new spawned RunIn from a closure. pub fn new(location: &'static std::panic::Location<'static>, task_fn: F) -> Self { - Self { task_fn, location } + Self { + task_fn, + location, + context: PhantomData, + } } } -impl RunWithConnectionTo for SpawnedRun +impl RunWithConnectionTo for SpawnedRun where Counterpart: Role, - F: FnOnce(ConnectionTo) -> Fut + Send, + Context: ConnectionContext, + F: FnOnce(Context::Connection) -> Fut + Send, Fut: Future> + Send, { async fn run_with_connection_to( @@ -91,14 +102,14 @@ where connection: ConnectionTo, ) -> Result<(), crate::Error> { let location = self.location; - (self.task_fn)(connection).await.map_err(|err| { - let data = err.data.clone(); - err.data(serde_json::json! { - { + (self.task_fn)(connection_context::from_raw::(connection)) + .await + .map_err(|err| { + let data = err.data.clone(); + err.data(serde_json::json!({ "spawned_at": format!("{}:{}:{}", location.file(), location.line(), location.column()), "data": data, - } + })) }) - }) } } diff --git a/src/agent-client-protocol/src/jsonrpc/task_actor.rs b/src/agent-client-protocol/src/jsonrpc/task_actor.rs index a885f15..92a05a6 100644 --- a/src/agent-client-protocol/src/jsonrpc/task_actor.rs +++ b/src/agent-client-protocol/src/jsonrpc/task_actor.rs @@ -45,6 +45,11 @@ impl Task { .map_err(crate::util::internal_error)?; Ok(()) } + + #[cfg(test)] + pub(super) async fn run_for_test(self) -> Result<(), crate::Error> { + self.future.await + } } /// The "task actor" manages dynamically spawned tasks. diff --git a/src/agent-client-protocol/src/lib.rs b/src/agent-client-protocol/src/lib.rs index ceafa7d..87d8177 100644 --- a/src/agent-client-protocol/src/lib.rs +++ b/src/agent-client-protocol/src/lib.rs @@ -18,8 +18,14 @@ //! ## Quick Start: Connecting to an Agent //! //! The most common use case is connecting to an existing ACP agent as a client. -//! Here's a minimal example that initializes a connection, creates a session, -//! and sends a prompt: +//! This example uses stable ACP protocol v1. The draft protocol v2 feature +//! provides a command-only `V2Session` API and receives updates and interactive +//! requests through typed connection handlers because prompt acceptance and +//! inbound traffic are independent. MCP attachment and proxy-session helpers +//! remain v1-only. +//! +//! Here's a minimal example that initializes a v1 connection, creates a +//! session, and sends a prompt: //! //! ```no_run //! use agent_client_protocol::Client; @@ -109,15 +115,17 @@ pub mod util; pub use capabilities::*; pub use jsonrpc::{ - Builder, ByteStreams, Channel, ConnectionTo, Dispatch, DynamicHandlerGuard, + Builder, ByteStreams, Channel, ConnectionContext, ConnectionTo, Dispatch, DynamicHandlerGuard, HandleConnectionClose, HandleDispatchFrom, Handled, INCOMING_TRANSPORT_CLOSED_REASON, IntoHandled, JsonRpcMessage, JsonRpcNotification, JsonRpcRequest, JsonRpcResponse, Lines, - NullClose, NullHandler, RawJsonRpcMessage, RawJsonRpcParams, Responder, ResponseRouter, - SentRequest, TransportBatch, TransportBatchEntry, TransportFrame, UntypedMessage, - is_incoming_transport_closed, + NullClose, NullHandler, RawConnectionContext, RawJsonRpcMessage, RawJsonRpcParams, Responder, + ResponseRouter, SentRequest, TransportBatch, TransportBatchEntry, TransportFrame, + UntypedMessage, is_incoming_transport_closed, run::{ChainRun, NullRun, RunWithConnectionTo}, }; pub use jsonrpc::{RequestCancellation, is_cancel_request_notification}; +#[cfg(feature = "unstable_protocol_v2")] +pub use jsonrpc::{V2Builder, V2ConnectionContext, V2ConnectionTo}; #[cfg(feature = "unstable_protocol_v2")] pub use role::acp::AgentProtocolRouter; diff --git a/src/agent-client-protocol/src/role/acp.rs b/src/agent-client-protocol/src/role/acp.rs index d47ce31..bfe19d1 100644 --- a/src/agent-client-protocol/src/role/acp.rs +++ b/src/agent-client-protocol/src/role/acp.rs @@ -10,7 +10,7 @@ use crate::DynConnectTo; use crate::jsonrpc::{Builder, handlers::NullHandler, run::NullRun}; #[cfg(feature = "unstable_protocol_v2")] use crate::jsonrpc::{ - TransportBatch, TransportBatchEntry, TransportFrame, is_response_only_shape, + TransportBatch, TransportBatchEntry, TransportFrame, V2Builder, is_response_only_shape, raw_is_response_only_shape, }; use crate::role::{HasPeer, RemoteStyle}; @@ -72,7 +72,7 @@ impl Client { /// /// Requires the `unstable_protocol_v2` crate feature. #[cfg(feature = "unstable_protocol_v2")] - pub fn v2(self) -> Builder { + pub fn v2(self) -> V2Builder { self.builder().v2_client() } @@ -298,13 +298,18 @@ impl Role for Agent { ) -> Result, crate::Error> { MatchDispatchFrom::new(message, &connection) .if_dispatch_from(Agent, async |message: Dispatch| { - // Subtle: messages that have a session-id field - // should be captured by a dynamic message handler - // for that session -- but there is a race condition - // between the dynamic handler being added and - // possible updates. Therefore, we "retry" all such - // messages, so that they will be resent as new handlers - // are added. + // Stable v1 session helpers install a dynamic handler after + // `session/new`. Retry session messages to close the race + // between the response and that registration. + // + // V2 uses typed handlers installed before the connection + // starts. Retrying an unhandled v2 message would retain it + // forever because no per-session dynamic handler is expected. + #[cfg(feature = "unstable_protocol_v2")] + let retry = message.has_session_id() + && connection.acp_protocol_version() + != Some(crate::schema::ProtocolVersion::V2); + #[cfg(not(feature = "unstable_protocol_v2"))] let retry = message.has_session_id(); Ok(Handled::No { message, retry }) }) @@ -326,7 +331,7 @@ impl Agent { /// /// Requires the `unstable_protocol_v2` crate feature. #[cfg(feature = "unstable_protocol_v2")] - pub fn v2(self) -> Builder { + pub fn v2(self) -> V2Builder { self.builder().v2_agent() } @@ -1131,7 +1136,7 @@ impl Role for Conductor { Client, async |request: InitializeProxyRequest, responder| { let InitializeProxyRequest { initialize } = request; - cx.send_request_to(Agent, initialize) + cx.send_ordered_request_to(Agent, initialize) .forward_response_to(responder) }, ) @@ -1139,7 +1144,7 @@ impl Role for Conductor { // New session coming from the client -- proxy to the agent // and add a dynamic handler for that session-id. .if_request_from(Client, async |request: NewSessionRequest, responder| { - let sent = cx.send_request_to(Agent, request); + let sent = cx.send_ordered_request_to(Agent, request); // The dynamic-handler hook below means we cannot use // `forward_response_to`, so wire up cancellation forwarding // explicitly to keep `session/new` cancellable like every diff --git a/src/agent-client-protocol/src/session.rs b/src/agent-client-protocol/src/session.rs index 8c5cb82..47200c0 100644 --- a/src/agent-client-protocol/src/session.rs +++ b/src/agent-client-protocol/src/session.rs @@ -20,6 +20,11 @@ use crate::{ #[cfg(feature = "unstable_mcp_over_acp")] use crate::{jsonrpc::run::ChainRun, mcp_server::McpServer}; +#[cfg(feature = "unstable_protocol_v2")] +mod v2; +#[cfg(feature = "unstable_protocol_v2")] +pub use v2::*; + /// Marker type indicating the session builder will block the current task. #[derive(Debug)] pub struct Blocking; @@ -38,12 +43,15 @@ impl ConnectionTo where Counterpart: HasPeer, { - /// Session builder for a new session request. + /// Stable protocol v1 session builder for a new session request. + /// + /// With `unstable_protocol_v2`, a `Client.v2()` callback receives a + /// `V2ConnectionTo` with its own v2 `build_session` helper. pub fn build_session(&self, cwd: impl AsRef) -> SessionBuilder { SessionBuilder::new(self, NewSessionRequest::new(cwd.as_ref())) } - /// Session builder using the current working directory. + /// Stable protocol v1 session builder using the current working directory. /// /// This is a convenience wrapper around [`build_session`](Self::build_session) /// that uses [`std::env::current_dir`] to get the working directory. @@ -56,7 +64,7 @@ where Ok(self.build_session(cwd)) } - /// Session builder starting from an existing request. + /// Stable protocol v1 session builder starting from an existing request. /// /// Use this when you've intercepted a `session.new` request and want to /// modify it (e.g., inject MCP servers) before forwarding. @@ -107,7 +115,7 @@ where } } -/// Session builder for a new session request. +/// Stable protocol v1 session builder for a new session request. /// Allows you to add MCP servers or set other details for this session. /// /// The `BlockState` type parameter tracks whether blocking methods are available: @@ -222,6 +230,8 @@ where F: FnOnce(ActiveSession<'static, Counterpart>) -> Fut + Send + 'static, Fut: Future> + Send, { + ensure_v1_session_protocol(&self.connection)?; + let Self { connection, request, @@ -231,7 +241,7 @@ where } = self; connection - .send_request_to(Agent, request) + .send_ordered_request_to(Agent, request) .on_receiving_result({ let connection = connection.clone(); async move |result| { @@ -310,6 +320,8 @@ where Counterpart: HasPeer, R: 'static, { + ensure_v1_session_protocol(&self.connection)?; + let Self { connection, request, @@ -319,7 +331,7 @@ where } = self; // Send the "new session" request to the agent. - let sent = connection.send_request_to(Agent, request); + let sent = connection.send_ordered_request_to(Agent, request); let sent = sent.forward_cancellation_from(responder.cancellation()); sent.on_receiving_ok_result(responder, { @@ -392,6 +404,8 @@ where ActiveSession<'runner, Counterpart>, ) -> Result, ) -> Result { + ensure_v1_session_protocol(&self.connection)?; + let Self { connection, request, @@ -427,6 +441,8 @@ where where R: 'static, { + ensure_v1_session_protocol(&self.connection)?; + let Self { connection, request, @@ -493,7 +509,7 @@ where } } -/// Active session struct that lets you send prompts and receive updates. +/// Stable protocol v1 active session that lets you send prompts and receive updates. /// /// The `'runner` lifetime represents the span during which session support runners /// (such as MCP servers) are active. When created via [`SessionBuilder::start_session`], @@ -527,7 +543,7 @@ where _runner: PhantomData<&'runner ()>, } -/// Incoming message from the agent +/// Incoming stable protocol v1 message from the agent. #[non_exhaustive] #[derive(Debug)] #[allow( @@ -581,7 +597,7 @@ where pub fn send_prompt(&mut self, prompt: impl ToString) -> Result<(), crate::Error> { let update_tx = self.update_tx.clone(); self.connection - .send_request_to( + .send_ordered_request_to( Agent, PromptRequest::new(self.session_id.clone(), vec![prompt.to_string().into()]), ) @@ -785,3 +801,28 @@ where format!("ActiveSessionHandler({})", self.session_id) } } + +#[cfg(not(feature = "unstable_protocol_v2"))] +#[allow( + clippy::unnecessary_wraps, + reason = "signature matches the feature-enabled protocol guard" +)] +fn ensure_v1_session_protocol( + _connection: &ConnectionTo, +) -> Result<(), crate::Error> { + Ok(()) +} + +#[cfg(feature = "unstable_protocol_v2")] +fn ensure_v1_session_protocol( + connection: &ConnectionTo, +) -> Result<(), crate::Error> { + if connection.acp_protocol_version() != Some(crate::schema::ProtocolVersion::V2) { + return Ok(()); + } + + Err(crate::Error::invalid_request().data( + "`build_session` uses ACP protocol v1 types, but this is a protocol v2 connection; \ + use the `V2ConnectionTo` supplied to `Client.v2()` callbacks", + )) +} diff --git a/src/agent-client-protocol/src/session/v2.rs b/src/agent-client-protocol/src/session/v2.rs new file mode 100644 index 0000000..576d101 --- /dev/null +++ b/src/agent-client-protocol/src/session/v2.rs @@ -0,0 +1,254 @@ +use std::path::Path; + +use crate::{Agent, SentRequest, V2ConnectionTo, role::HasPeer, schema::v2}; + +impl V2ConnectionTo +where + Counterpart: HasPeer, +{ + /// Build a draft protocol v2 `session/new` request. + pub fn build_session(&self, cwd: impl AsRef) -> V2SessionBuilder { + V2SessionBuilder::new(self, v2::NewSessionRequest::new(cwd.as_ref())) + } + + /// Build a draft protocol v2 session using the current working directory. + /// + /// Returns an error if the current directory cannot be determined. + pub fn build_session_cwd(&self) -> Result, crate::Error> { + let cwd = std::env::current_dir().map_err(|error| { + crate::Error::internal_error().data(format!("cannot get current directory: {error}")) + })?; + Ok(self.build_session(cwd)) + } + + /// Build a draft protocol v2 session from an existing `session/new` request. + pub fn build_session_from( + &self, + request: v2::NewSessionRequest, + ) -> V2SessionBuilder { + V2SessionBuilder::new(self, request) + } + + /// Resume a draft protocol v2 session. + /// + /// Use [`Self::resume_session_from`] to request history replay or set + /// other optional resume parameters. + pub fn resume_session( + &self, + session_id: impl Into, + cwd: impl AsRef, + ) -> SentRequest> { + self.resume_session_from(v2::ResumeSessionRequest::new(session_id, cwd.as_ref())) + } + + /// Resume a draft protocol v2 session from an existing request. + /// + /// Register typed session update handlers before connecting. When the + /// request asks for replay, the agent sends those updates before the + /// [`v2::ResumeSessionResponse`]. + pub fn resume_session_from( + &self, + request: v2::ResumeSessionRequest, + ) -> SentRequest> { + let session_id = request.session_id.clone(); + let session_connection = self.clone(); + + self.send_request_to(Agent, request).map(move |response| { + let session = V2Session { + session_id, + connection: session_connection, + }; + Ok(OpenedV2Session { session, response }) + }) + } +} + +/// Builder for a draft protocol v2 `session/new` request. +/// +/// Protocol v2 acknowledges `session/prompt` independently from inbound +/// session updates. Register typed [`v2::UpdateSessionNotification`] and +/// session request handlers on [`crate::Builder`] before connecting, then use +/// [`Self::start_session`] to create the command-only [`V2Session`] handle. +/// +/// Per-session MCP attachment and proxy-session helpers are currently available +/// only through the stable protocol v1 [`crate::SessionBuilder`]. +#[must_use = "call `start_session` to send the `session/new` request"] +#[derive(Debug)] +pub struct V2SessionBuilder +where + Counterpart: HasPeer, +{ + connection: V2ConnectionTo, + request: v2::NewSessionRequest, +} + +impl V2SessionBuilder +where + Counterpart: HasPeer, +{ + fn new(connection: &V2ConnectionTo, request: v2::NewSessionRequest) -> Self { + Self { + connection: connection.clone(), + request, + } + } + + /// Send `session/new` and return its independently consumable request. + /// + /// The successful result contains both a cloneable command handle and the + /// complete [`v2::NewSessionResponse`]. Consume the returned request with + /// [`SentRequest::block_task`], [`SentRequest::on_receiving_result`], or + /// another explicit [`SentRequest`] completion mode. + pub fn start_session( + self, + ) -> SentRequest> { + let Self { + connection, + request, + } = self; + let session_connection = connection.clone(); + + connection + .send_request_to(Agent, request) + .map(move |response| { + let session = V2Session { + session_id: response.session_id.clone(), + connection: session_connection, + }; + Ok(OpenedV2Session { session, response }) + }) + } +} + +/// A newly available protocol v2 session and its operation-specific response. +/// +/// Keeping the response separate from [`V2Session`] avoids treating +/// `session/new` setup data as mutable session state and lets each setup +/// operation, including `session/resume`, retain its own complete response +/// type. +#[derive(Debug)] +pub struct OpenedV2Session +where + Link: HasPeer, +{ + session: V2Session, + response: Response, +} + +impl OpenedV2Session +where + Link: HasPeer, +{ + /// Access the command handle for the opened session. + pub fn session(&self) -> &V2Session { + &self.session + } + + /// Access the complete response from the operation that opened the session. + pub fn response(&self) -> &Response { + &self.response + } + + /// Split this result into the command handle and complete setup response. + pub fn into_parts(self) -> (V2Session, Response) { + (self.session, self.response) + } + + /// Consume this result and return only the command handle. + pub fn into_session(self) -> V2Session { + self.session + } +} + +/// Cloneable command handle for a draft protocol v2 session. +/// +/// Inbound protocol traffic is intentionally not owned by this value. Receive +/// authoritative [`v2::UpdateSessionNotification`] values and interactive +/// requests such as [`v2::RequestPermissionRequest`] through typed handlers +/// installed on [`crate::Builder`]. +#[derive(Debug, Clone)] +pub struct V2Session +where + Link: HasPeer, +{ + session_id: v2::SessionId, + connection: V2ConnectionTo, +} + +impl V2Session +where + Link: HasPeer, +{ + /// Access the session ID. + pub fn session_id(&self) -> &v2::SessionId { + &self.session_id + } + + /// Access the underlying connection. + pub fn connection(&self) -> &V2ConnectionTo { + &self.connection + } + + /// Submit a text prompt and return its independent acceptance request. + /// + /// A successful response only acknowledges that the agent accepted the + /// prompt. The accepted user message, output, state changes, and completion + /// arrive independently through [`v2::UpdateSessionNotification`]. + pub fn send_prompt(&self, prompt: impl ToString) -> SentRequest { + self.send_prompt_blocks(vec![prompt.to_string().into()]) + } + + /// Submit arbitrary prompt content and return its acceptance request. + /// + /// The SDK does not track foreground state or gate prompt submission + /// locally. Wait for an `idle` state update before another ordinary prompt + /// unless using a separately defined admission mechanism. + pub fn send_prompt_blocks( + &self, + prompt: Vec, + ) -> SentRequest { + self.connection.send_request_to( + Agent, + v2::PromptRequest::new(self.session_id.clone(), prompt), + ) + } + + /// Ask the agent to cancel the session's current foreground work. + /// + /// This is independent from cancelling a prompt's [`SentRequest`]. + /// Cancellation completes when the agent reports an `idle` state update + /// with [`v2::StopReason::Cancelled`]. The client should immediately mark + /// unfinished tool calls for the active work as cancelled and remains + /// responsible for resolving every pending [`v2::RequestPermissionRequest`] + /// with the cancelled outcome. + pub fn cancel_active_work(&self) -> Result<(), crate::Error> { + self.connection.send_notification_to( + Agent, + v2::CancelSessionNotification::new(self.session_id.clone()), + ) + } + + /// Set a session configuration option. + /// + /// The response contains the full current option set. It is not cached on + /// this command handle. + pub fn set_config_option( + &self, + config_id: impl Into, + value: impl Into, + ) -> SentRequest { + self.connection.send_request_to( + Agent, + v2::SetSessionConfigOptionRequest::new(self.session_id.clone(), config_id, value), + ) + } + + /// Close the remote session and release its resources. + /// + /// Existing clones of this local command handle are not invalidated, but + /// the agent should reject subsequent commands for the closed session. + pub fn close(&self) -> SentRequest { + self.connection + .send_request_to(Agent, v2::CloseSessionRequest::new(self.session_id.clone())) + } +} diff --git a/src/agent-client-protocol/tests/protocol_v2.rs b/src/agent-client-protocol/tests/protocol_v2.rs index 6314278..022ffa3 100644 --- a/src/agent-client-protocol/tests/protocol_v2.rs +++ b/src/agent-client-protocol/tests/protocol_v2.rs @@ -11,9 +11,12 @@ use std::{ use agent_client_protocol::schema::{ProtocolVersion, v1, v2}; use agent_client_protocol::{ Agent, AgentProtocolRouter, Builder, ByteStreams, Client, ClientProtocolConnector, ConnectTo, - Error, JsonRpcMessage, JsonRpcRequest, JsonRpcResponse, NullHandler, RawJsonRpcMessage, Role, - TransportFrame, UntypedRole, + ConnectionContext, ConnectionTo, DynamicHandlerGuard, Error, HandleConnectionClose, + HandleDispatchFrom, JsonRpcMessage, JsonRpcRequest, JsonRpcResponse, NullHandler, + RawJsonRpcMessage, Role, RunWithConnectionTo, TransportFrame, UntypedRole, V2Builder, + V2ConnectionTo, }; +use agent_client_protocol_test::MockTransport; use agent_client_protocol_test::testy::Testy; use futures::StreamExt as _; use serde::{Deserialize, Serialize}; @@ -57,6 +60,25 @@ fn cwd() -> Result { std::env::current_dir().map_err(Error::into_internal_error) } +#[allow(dead_code)] +fn spawn_child_from_generic_context( + parent: &ConnectionTo, + builder: Builder< + R, + impl HandleDispatchFrom + 'static, + impl RunWithConnectionTo + 'static, + impl HandleConnectionClose + 'static, + Context, + >, + transport: impl ConnectTo + 'static, +) -> Result, Error> +where + R: Role, + Context: ConnectionContext, +{ + parent.spawn_connection_with_context(builder, transport) +} + fn v2_implementation() -> v2::Implementation { v2::Implementation::new("agent-client-protocol-test", env!("CARGO_PKG_VERSION")) } @@ -474,6 +496,24 @@ async fn assert_v2_client_rejected_by_v1_agent(agent: impl ConnectTo) -> .await } +#[tokio::test(flavor = "current_thread")] +async fn v2_context_preserves_protocol_neutral_connection_management() -> Result<(), Error> { + Client + .v2() + .connect_with(Agent.v2(), async |connection| { + let dynamic_handler: DynamicHandlerGuard = + connection.add_dynamic_handler(NullHandler)?; + let raw_child: ConnectionTo = + connection.spawn_connection::(Client.builder(), MockTransport)?; + let v2_child: V2ConnectionTo = + connection.spawn_connection(Client.v2(), MockTransport)?; + + drop((dynamic_handler, raw_child, v2_child)); + Ok(()) + }) + .await +} + #[tokio::test(flavor = "current_thread")] async fn non_acp_initialize_is_not_rewritten() -> Result<(), Error> { UntypedRole @@ -2140,7 +2180,7 @@ async fn protocol_router_routes_future_protocol_version_to_v2() -> Result<(), Er /// A v2 agent whose `session/new` handler only responds once the peer cancels /// the request via `$/cancel_request`. fn v2_agent_with_cancellable_new_session() --> Builder> { +-> V2Builder> { Agent .v2() .on_receive_request( diff --git a/src/agent-client-protocol/tests/session_v2.rs b/src/agent-client-protocol/tests/session_v2.rs new file mode 100644 index 0000000..9ce81ef --- /dev/null +++ b/src/agent-client-protocol/tests/session_v2.rs @@ -0,0 +1,1101 @@ +#![cfg(feature = "unstable_protocol_v2")] + +use std::{ + path::PathBuf, + sync::{Arc, Mutex}, + time::Duration, +}; + +use agent_client_protocol::{ + Agent, Client, ConnectionTo, Error, ErrorCode, Responder, RunWithConnectionTo, V2ConnectionTo, + schema::{ProtocolVersion, v2}, +}; +use futures::{ + StreamExt as _, + channel::{ + mpsc::{self, UnboundedReceiver}, + oneshot, + }, +}; + +const TIMEOUT: Duration = Duration::from_secs(10); + +fn cwd() -> Result { + std::env::current_dir().map_err(Error::into_internal_error) +} + +fn implementation() -> v2::Implementation { + v2::Implementation::new("session-v2-test", env!("CARGO_PKG_VERSION")) +} + +fn initialize_response(protocol_version: ProtocolVersion) -> v2::InitializeResponse { + v2::InitializeResponse::new(protocol_version, implementation()) + .capabilities(v2::AgentCapabilities::new().session(v2::SessionCapabilities::new())) +} + +struct V1SessionGuardRunner { + cwd: PathBuf, + completed: oneshot::Sender<()>, +} + +impl RunWithConnectionTo for V1SessionGuardRunner { + async fn run_with_connection_to(self, connection: ConnectionTo) -> Result<(), Error> { + let error = connection + .build_session(self.cwd) + .block_task() + .start_session() + .await + .expect_err("the v1 helper must reject the raw protocol v2 connection"); + assert_eq!(error.code, ErrorCode::InvalidRequest); + + let data = error + .data + .as_ref() + .and_then(serde_json::Value::as_str) + .unwrap_or_default(); + assert!(data.contains("V2ConnectionTo"), "{error:?}"); + + self.completed.send(()).map_err(|()| { + Error::internal_error().data("v1 session guard test receiver was dropped") + }) + } +} + +async fn next_update( + updates: &mut UnboundedReceiver, +) -> v2::UpdateSessionNotification { + updates + .next() + .await + .expect("the typed session update handler should remain active") +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn v2_prompt_acceptance_is_independent_from_session_updates() { + let session_id = v2::SessionId::new("v2-session"); + let config_option = v2::SessionConfigOption::boolean("thinking", "Thinking", true); + let meta = serde_json::Map::from_iter([( + "extension".to_owned(), + serde_json::json!({"preserved": true}), + )]); + let expected_new_response = v2::NewSessionResponse::new(session_id.clone()) + .config_options(vec![config_option]) + .meta(meta); + let agent_new_response = expected_new_response.clone(); + + let agent = Agent + .v2() + .on_receive_request( + async |request: v2::InitializeRequest, + responder: Responder, + _connection: V2ConnectionTo| { + assert_eq!(request.protocol_version, ProtocolVersion::V2); + responder.respond(initialize_response(request.protocol_version)) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async move |request: v2::NewSessionRequest, + responder: Responder, + _connection: V2ConnectionTo| { + assert!(AsRef::::as_ref(&request.cwd).is_absolute()); + responder.respond(agent_new_response.clone()) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async move |request: v2::PromptRequest, + responder: Responder, + connection: V2ConnectionTo| { + assert_eq!(request.session_id, v2::SessionId::new("v2-session")); + assert!(matches!( + request.prompt.as_slice(), + [v2::ContentBlock::Text(text)] if text.text == "hello" + )); + + // The update lane is independent and can be observed before + // this prompt's acceptance response. + connection.send_notification(v2::UpdateSessionNotification::new( + request.session_id.clone(), + v2::SessionUpdate::AgentMessage( + v2::AgentMessage::new(v2::MessageId::new("background")) + .content(vec!["unrelated".into()]), + ), + ))?; + responder.respond(v2::PromptResponse::new())?; + connection.send_notification(v2::UpdateSessionNotification::new( + request.session_id.clone(), + v2::SessionUpdate::UserMessage( + v2::UserMessage::new(v2::MessageId::new("accepted-user")) + .content(request.prompt), + ), + ))?; + connection.send_notification(v2::UpdateSessionNotification::new( + request.session_id.clone(), + v2::SessionUpdate::StateUpdate(v2::StateUpdate::Running( + v2::RunningStateUpdate::new(), + )), + ))?; + connection.send_notification(v2::UpdateSessionNotification::new( + request.session_id, + v2::SessionUpdate::StateUpdate(v2::StateUpdate::Idle( + v2::IdleStateUpdate::new().stop_reason(v2::StopReason::MaxTokens), + )), + )) + }, + agent_client_protocol::on_receive_request!(), + ); + + let (update_tx, mut update_rx) = mpsc::unbounded(); + let client = Client + .v2() + .on_receive_notification( + async move |update: v2::UpdateSessionNotification, + _connection: V2ConnectionTo| { + update_tx + .unbounded_send(update) + .map_err(Error::into_internal_error) + }, + agent_client_protocol::on_receive_notification!(), + ) + .connect_with(agent, async move |connection| { + let initialize = connection + .send_request(v2::InitializeRequest::new( + ProtocolVersion::V2, + implementation(), + )) + .block_task() + .await?; + assert_eq!(initialize.protocol_version, ProtocolVersion::V2); + + let opened = connection + .build_session(cwd()?) + .start_session() + .block_task() + .await?; + assert_eq!(opened.response(), &expected_new_response); + + let (session, response) = opened.into_parts(); + assert_eq!(response, expected_new_response); + assert_eq!(session.session_id(), &session_id); + + let acceptance = session.send_prompt("hello"); + + let background = next_update(&mut update_rx).await; + assert_eq!(background.session_id, session_id); + assert!(matches!( + background.update, + v2::SessionUpdate::AgentMessage(message) + if message.message_id == v2::MessageId::new("background") + )); + + assert_eq!(acceptance.block_task().await?, v2::PromptResponse::new()); + + assert!(matches!( + next_update(&mut update_rx).await.update, + v2::SessionUpdate::UserMessage(message) + if message.message_id == v2::MessageId::new("accepted-user") + )); + assert!(matches!( + next_update(&mut update_rx).await.update, + v2::SessionUpdate::StateUpdate(v2::StateUpdate::Running(_)) + )); + assert!(matches!( + next_update(&mut update_rx).await.update, + v2::SessionUpdate::StateUpdate(v2::StateUpdate::Idle(idle)) + if idle.stop_reason == Some(v2::StopReason::MaxTokens) + )); + Ok(()) + }); + + tokio::time::timeout(TIMEOUT, client) + .await + .expect("v2 session lifecycle timed out") + .expect("v2 session connection failed"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn v2_session_cancellation_completes_at_cancelled_idle() { + let agent = Agent + .v2() + .on_receive_request( + async |request: v2::InitializeRequest, + responder: Responder, + _connection: V2ConnectionTo| { + responder.respond(initialize_response(request.protocol_version)) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async |_request: v2::NewSessionRequest, + responder: Responder, + _connection: V2ConnectionTo| { + responder.respond(v2::NewSessionResponse::new(v2::SessionId::new( + "cancel-session", + ))) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async |request: v2::PromptRequest, + responder: Responder, + connection: V2ConnectionTo| { + responder.respond(v2::PromptResponse::new())?; + connection.send_notification(v2::UpdateSessionNotification::new( + request.session_id.clone(), + v2::SessionUpdate::StateUpdate(v2::StateUpdate::Running( + v2::RunningStateUpdate::new(), + )), + ))?; + connection.send_notification(v2::UpdateSessionNotification::new( + request.session_id, + v2::SessionUpdate::AgentMessageChunk(v2::ContentChunk::new( + "before".into(), + v2::MessageId::new("cancel-answer"), + )), + )) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_notification( + async |cancel: v2::CancelSessionNotification, connection: V2ConnectionTo| { + connection.send_notification(v2::UpdateSessionNotification::new( + cancel.session_id.clone(), + v2::SessionUpdate::AgentMessageChunk(v2::ContentChunk::new( + " after".into(), + v2::MessageId::new("cancel-answer"), + )), + ))?; + connection.send_notification(v2::UpdateSessionNotification::new( + cancel.session_id, + v2::SessionUpdate::StateUpdate(v2::StateUpdate::Idle( + v2::IdleStateUpdate::new().stop_reason(v2::StopReason::Cancelled), + )), + )) + }, + agent_client_protocol::on_receive_notification!(), + ); + + let (update_tx, mut update_rx) = mpsc::unbounded(); + let client = Client + .v2() + .on_receive_notification( + async move |update: v2::UpdateSessionNotification, + _connection: V2ConnectionTo| { + update_tx + .unbounded_send(update) + .map_err(Error::into_internal_error) + }, + agent_client_protocol::on_receive_notification!(), + ) + .connect_with(agent, async move |connection| { + connection + .send_request(v2::InitializeRequest::new( + ProtocolVersion::V2, + implementation(), + )) + .block_task() + .await?; + + let opened = connection + .build_session(cwd()?) + .start_session() + .block_task() + .await?; + let (session, _) = opened.into_parts(); + session.send_prompt("cancel me").block_task().await?; + + assert!(matches!( + next_update(&mut update_rx).await.update, + v2::SessionUpdate::StateUpdate(v2::StateUpdate::Running(_)) + )); + assert!(matches!( + next_update(&mut update_rx).await.update, + v2::SessionUpdate::AgentMessageChunk(_) + )); + + session.cancel_active_work()?; + assert!(matches!( + next_update(&mut update_rx).await.update, + v2::SessionUpdate::AgentMessageChunk(_) + )); + assert!(matches!( + next_update(&mut update_rx).await.update, + v2::SessionUpdate::StateUpdate(v2::StateUpdate::Idle(idle)) + if idle.stop_reason == Some(v2::StopReason::Cancelled) + )); + Ok(()) + }); + + tokio::time::timeout(TIMEOUT, client) + .await + .expect("v2 cancellation lifecycle timed out") + .expect("v2 cancellation connection failed"); +} + +#[tokio::test(flavor = "current_thread")] +async fn cloned_v2_session_handles_do_not_own_prompt_state() { + let agent = Agent + .v2() + .on_receive_request( + async |request: v2::InitializeRequest, + responder: Responder, + _connection: V2ConnectionTo| { + responder.respond(initialize_response(request.protocol_version)) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async |_request: v2::NewSessionRequest, + responder: Responder, + _connection: V2ConnectionTo| { + responder.respond(v2::NewSessionResponse::new(v2::SessionId::new( + "parallel-prompts", + ))) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async |request: v2::PromptRequest, + responder: Responder, + connection: V2ConnectionTo| { + assert!(matches!( + request.prompt.as_slice(), + [v2::ContentBlock::Text(text)] + if text.text == "first" || text.text == "second" + )); + responder.respond(v2::PromptResponse::new())?; + connection.send_notification(v2::UpdateSessionNotification::new( + request.session_id, + v2::SessionUpdate::StateUpdate(v2::StateUpdate::Idle( + v2::IdleStateUpdate::new().stop_reason(v2::StopReason::EndTurn), + )), + )) + }, + agent_client_protocol::on_receive_request!(), + ); + + let (update_tx, mut update_rx) = mpsc::unbounded(); + let client = Client + .v2() + .on_receive_notification( + async move |update: v2::UpdateSessionNotification, + _connection: V2ConnectionTo| { + update_tx + .unbounded_send(update) + .map_err(Error::into_internal_error) + }, + agent_client_protocol::on_receive_notification!(), + ) + .connect_with(agent, async move |connection| { + connection + .send_request(v2::InitializeRequest::new( + ProtocolVersion::V2, + implementation(), + )) + .block_task() + .await?; + + let session = connection + .build_session(cwd()?) + .start_session() + .block_task() + .await? + .into_session(); + let other_task = session.clone(); + + assert_eq!( + session.send_prompt("first").block_task().await?, + v2::PromptResponse::new() + ); + assert!(matches!( + next_update(&mut update_rx).await.update, + v2::SessionUpdate::StateUpdate(v2::StateUpdate::Idle(_)) + )); + + assert_eq!( + other_task.send_prompt("second").block_task().await?, + v2::PromptResponse::new() + ); + assert!(matches!( + next_update(&mut update_rx).await.update, + v2::SessionUpdate::StateUpdate(v2::StateUpdate::Idle(_)) + )); + Ok(()) + }); + + tokio::time::timeout(TIMEOUT, client) + .await + .expect("cloned v2 session handle test timed out") + .expect("cloned v2 session handle test failed"); +} + +#[tokio::test(flavor = "current_thread")] +async fn v2_prompt_error_does_not_stop_session_updates() { + let agent = Agent + .v2() + .on_receive_request( + async |request: v2::InitializeRequest, + responder: Responder, + _connection: V2ConnectionTo| { + responder.respond(initialize_response(request.protocol_version)) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async |_request: v2::NewSessionRequest, + responder: Responder, + _connection: V2ConnectionTo| { + responder.respond(v2::NewSessionResponse::new(v2::SessionId::new( + "rejected-prompt", + ))) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async |request: v2::PromptRequest, + responder: Responder, + connection: V2ConnectionTo| { + responder.respond_with_error(Error::invalid_params().data("prompt rejected"))?; + connection.send_notification(v2::UpdateSessionNotification::new( + request.session_id, + v2::SessionUpdate::AgentMessage( + v2::AgentMessage::new(v2::MessageId::new("unrelated-after-error")) + .content(vec!["background".into()]), + ), + )) + }, + agent_client_protocol::on_receive_request!(), + ); + + let (update_tx, mut update_rx) = mpsc::unbounded(); + let client = Client + .v2() + .on_receive_notification( + async move |update: v2::UpdateSessionNotification, + _connection: V2ConnectionTo| { + update_tx + .unbounded_send(update) + .map_err(Error::into_internal_error) + }, + agent_client_protocol::on_receive_notification!(), + ) + .connect_with(agent, async move |connection| { + connection + .send_request(v2::InitializeRequest::new( + ProtocolVersion::V2, + implementation(), + )) + .block_task() + .await?; + + let opened = connection + .build_session(cwd()?) + .start_session() + .block_task() + .await?; + let (session, _) = opened.into_parts(); + + let error = session + .send_prompt("reject this") + .block_task() + .await + .expect_err("the prompt error must reach its acceptance request"); + assert_eq!( + error, + Error::invalid_params().data("prompt rejected"), + "the helper must preserve the peer's prompt error" + ); + + assert!(matches!( + next_update(&mut update_rx).await.update, + v2::SessionUpdate::AgentMessage(message) + if message.message_id == v2::MessageId::new("unrelated-after-error") + )); + Ok(()) + }); + + tokio::time::timeout(TIMEOUT, client) + .await + .expect("v2 prompt rejection test timed out") + .expect("v2 prompt rejection stopped typed session updates"); +} + +#[tokio::test(flavor = "current_thread")] +async fn v2_permission_requests_use_a_separate_typed_handler() { + let agent = Agent + .v2() + .on_receive_request( + async |request: v2::InitializeRequest, + responder: Responder, + _connection: V2ConnectionTo| { + responder.respond(initialize_response(request.protocol_version)) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async |_request: v2::NewSessionRequest, + responder: Responder, + _connection: V2ConnectionTo| { + responder.respond(v2::NewSessionResponse::new(v2::SessionId::new( + "permission-session", + ))) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async |request: v2::PromptRequest, + responder: Responder, + connection: V2ConnectionTo| { + responder.respond(v2::PromptResponse::new())?; + + let session_id = request.session_id; + connection + .send_request(v2::RequestPermissionRequest::new( + session_id.clone(), + "Continue?", + vec![v2::PermissionOption::new( + "reject", + "Reject", + v2::PermissionOptionKind::RejectOnce, + )], + )) + .on_receiving_result({ + let connection = connection.clone(); + async move |response| { + assert_eq!( + response?.outcome, + v2::RequestPermissionOutcome::Selected( + v2::SelectedPermissionOutcome::new("reject") + ) + ); + connection.send_notification(v2::UpdateSessionNotification::new( + session_id, + v2::SessionUpdate::StateUpdate(v2::StateUpdate::Idle( + v2::IdleStateUpdate::new().stop_reason(v2::StopReason::Refusal), + )), + )) + } + }) + }, + agent_client_protocol::on_receive_request!(), + ); + + let (update_tx, mut update_rx) = mpsc::unbounded(); + let (permission_tx, mut permission_rx) = mpsc::unbounded::<( + v2::RequestPermissionRequest, + Responder, + )>(); + let client = Client + .v2() + .on_receive_notification( + async move |update: v2::UpdateSessionNotification, + _connection: V2ConnectionTo| { + update_tx + .unbounded_send(update) + .map_err(Error::into_internal_error) + }, + agent_client_protocol::on_receive_notification!(), + ) + .on_receive_request( + async move |request: v2::RequestPermissionRequest, + responder: Responder, + _connection: V2ConnectionTo| { + permission_tx + .unbounded_send((request, responder)) + .map_err(Error::into_internal_error) + }, + agent_client_protocol::on_receive_request!(), + ) + .connect_with(agent, async move |connection| { + connection + .send_request(v2::InitializeRequest::new( + ProtocolVersion::V2, + implementation(), + )) + .block_task() + .await?; + + let opened = connection + .build_session(cwd()?) + .start_session() + .block_task() + .await?; + let (session, _) = opened.into_parts(); + session.send_prompt("ask").block_task().await?; + + let (request, responder) = permission_rx + .next() + .await + .expect("permission handler should forward the request"); + assert_eq!(request.session_id, session.session_id().clone()); + assert_eq!(request.title, "Continue?"); + responder.respond(v2::RequestPermissionResponse::new( + v2::RequestPermissionOutcome::Selected(v2::SelectedPermissionOutcome::new( + "reject", + )), + ))?; + + assert!(matches!( + next_update(&mut update_rx).await.update, + v2::SessionUpdate::StateUpdate(v2::StateUpdate::Idle(idle)) + if idle.stop_reason == Some(v2::StopReason::Refusal) + )); + Ok(()) + }); + + tokio::time::timeout(TIMEOUT, client) + .await + .expect("v2 typed permission test timed out") + .expect("v2 typed permission test failed"); +} + +#[tokio::test(flavor = "current_thread")] +async fn unhandled_v2_session_messages_are_not_deferred() { + let (permission_result_tx, mut permission_result_rx) = + mpsc::unbounded::>(); + let agent = Agent + .v2() + .on_receive_request( + async |request: v2::InitializeRequest, + responder: Responder, + _connection: V2ConnectionTo| { + responder.respond(initialize_response(request.protocol_version)) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async |_request: v2::NewSessionRequest, + responder: Responder, + _connection: V2ConnectionTo| { + responder.respond(v2::NewSessionResponse::new(v2::SessionId::new( + "unhandled-session", + ))) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async move |request: v2::PromptRequest, + responder: Responder, + connection: V2ConnectionTo| { + responder.respond(v2::PromptResponse::new())?; + + // A v2 client without typed handlers should ignore an + // unhandled notification and reject an unhandled request, + // rather than retaining either for a dynamic v1 session route. + connection.send_notification(v2::UpdateSessionNotification::new( + request.session_id.clone(), + v2::SessionUpdate::StateUpdate(v2::StateUpdate::Running( + v2::RunningStateUpdate::new(), + )), + ))?; + + let permission_result_tx = permission_result_tx.clone(); + connection + .send_request(v2::RequestPermissionRequest::new( + request.session_id, + "Continue?", + vec![v2::PermissionOption::new( + "reject", + "Reject", + v2::PermissionOptionKind::RejectOnce, + )], + )) + .on_receiving_result(async move |result| { + permission_result_tx + .unbounded_send(result) + .map_err(Error::into_internal_error) + }) + }, + agent_client_protocol::on_receive_request!(), + ); + + let client = Client.v2().connect_with(agent, async move |connection| { + connection + .send_request(v2::InitializeRequest::new( + ProtocolVersion::V2, + implementation(), + )) + .block_task() + .await?; + + let session = connection + .build_session(cwd()?) + .start_session() + .block_task() + .await? + .into_session(); + session.send_prompt("ask").block_task().await?; + + let error = permission_result_rx + .next() + .await + .expect("the unhandled request should receive a response") + .expect_err("the client should reject an unhandled permission request"); + assert_eq!(error.code, ErrorCode::MethodNotFound); + Ok(()) + }); + + tokio::time::timeout(TIMEOUT, client) + .await + .expect("unhandled v2 session message test timed out") + .expect("unhandled v2 session message test failed"); +} + +#[tokio::test(flavor = "current_thread")] +async fn v2_resume_replay_is_handled_before_the_response() { + let session_id = v2::SessionId::new("resumed-session"); + let config_option = v2::SessionConfigOption::boolean("thinking", "Thinking", false); + let expected_response = v2::ResumeSessionResponse::new().config_options(vec![config_option]); + let agent_response = expected_response.clone(); + let agent_session_id = session_id.clone(); + + let agent = Agent + .v2() + .on_receive_request( + async |request: v2::InitializeRequest, + responder: Responder, + _connection: V2ConnectionTo| { + responder.respond(initialize_response(request.protocol_version)) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async move |request: v2::ResumeSessionRequest, + responder: Responder, + connection: V2ConnectionTo| { + assert_eq!(request.session_id, agent_session_id); + assert!(matches!( + request.replay_from, + Some(v2::ReplayFrom::Start(_)) + )); + + connection.send_notification(v2::UpdateSessionNotification::new( + request.session_id.clone(), + v2::SessionUpdate::UserMessage(v2::UserMessage::new(v2::MessageId::new( + "replayed-user", + ))), + ))?; + connection.send_notification(v2::UpdateSessionNotification::new( + request.session_id, + v2::SessionUpdate::AgentMessage(v2::AgentMessage::new(v2::MessageId::new( + "replayed-agent", + ))), + ))?; + responder.respond(agent_response.clone()) + }, + agent_client_protocol::on_receive_request!(), + ); + + let applied_replay = Arc::new(Mutex::new(Vec::new())); + let applied_replay_handler = applied_replay.clone(); + let client = Client + .v2() + .on_receive_notification( + async move |update: v2::UpdateSessionNotification, + _connection: V2ConnectionTo| { + let message_id = match update.update { + v2::SessionUpdate::UserMessage(message) => message.message_id, + v2::SessionUpdate::AgentMessage(message) => message.message_id, + other => panic!("unexpected replay update: {other:?}"), + }; + applied_replay_handler + .lock() + .expect("replay projection lock should not be poisoned") + .push(message_id); + Ok(()) + }, + agent_client_protocol::on_receive_notification!(), + ) + .connect_with(agent, async move |connection| { + connection + .send_request(v2::InitializeRequest::new( + ProtocolVersion::V2, + implementation(), + )) + .block_task() + .await?; + + let request = v2::ResumeSessionRequest::new(session_id.clone(), cwd()?) + .replay_from(v2::ReplayFrom::from(v2::ReplayFromStart::new())); + let opened = connection.resume_session_from(request).block_task().await?; + + assert_eq!(opened.session().session_id(), &session_id); + assert_eq!(opened.response(), &expected_response); + assert_eq!( + *applied_replay + .lock() + .expect("replay projection lock should not be poisoned"), + vec![ + v2::MessageId::new("replayed-user"), + v2::MessageId::new("replayed-agent"), + ], + "typed replay handlers must finish before the resume response is observed" + ); + Ok(()) + }); + + tokio::time::timeout(TIMEOUT, client) + .await + .expect("v2 resume replay test timed out") + .expect("v2 resume replay test failed"); +} + +#[tokio::test(flavor = "current_thread")] +async fn v2_session_commands_cover_configuration_and_close() { + let new_session_config_option = v2::SessionConfigOption::boolean("thinking", "Thinking", false); + let updated_config_option = v2::SessionConfigOption::boolean("thinking", "Thinking", true); + let set_response = v2::SetSessionConfigOptionResponse::new(vec![updated_config_option]); + let agent_set_response = set_response.clone(); + + let agent = Agent + .v2() + .on_receive_request( + async |request: v2::InitializeRequest, + responder: Responder, + _connection: V2ConnectionTo| { + responder.respond(initialize_response(request.protocol_version)) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async move |_request: v2::NewSessionRequest, + responder: Responder, + _connection: V2ConnectionTo| { + responder.respond( + v2::NewSessionResponse::new(v2::SessionId::new("command-session")) + .config_options(vec![new_session_config_option.clone()]), + ) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async move |request: v2::SetSessionConfigOptionRequest, + responder: Responder, + _connection: V2ConnectionTo| { + assert_eq!(request.session_id, v2::SessionId::new("command-session")); + assert_eq!(request.config_id, v2::SessionConfigId::new("thinking")); + assert_eq!(request.value, v2::SessionConfigOptionValue::from(true)); + responder.respond(agent_set_response.clone()) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async |request: v2::CloseSessionRequest, + responder: Responder, + _connection: V2ConnectionTo| { + assert_eq!(request.session_id, v2::SessionId::new("command-session")); + responder.respond(v2::CloseSessionResponse::new()) + }, + agent_client_protocol::on_receive_request!(), + ); + + let client = Client.v2().connect_with(agent, async move |connection| { + connection + .send_request(v2::InitializeRequest::new( + ProtocolVersion::V2, + implementation(), + )) + .block_task() + .await?; + + let opened = connection + .build_session(cwd()?) + .start_session() + .block_task() + .await?; + let (session, _) = opened.into_parts(); + + assert_eq!( + session + .set_config_option("thinking", true) + .block_task() + .await?, + set_response + ); + assert_eq!( + session.close().block_task().await?, + v2::CloseSessionResponse::new() + ); + Ok(()) + }); + + tokio::time::timeout(TIMEOUT, client) + .await + .expect("v2 session command test timed out") + .expect("v2 session command test failed"); +} + +#[tokio::test(flavor = "current_thread")] +async fn dropping_v2_session_does_not_unregister_update_handling() { + let session_id = v2::SessionId::new("dropped-handle"); + let agent_session_id = session_id.clone(); + + let agent = Agent + .v2() + .on_receive_request( + async |request: v2::InitializeRequest, + responder: Responder, + _connection: V2ConnectionTo| { + responder.respond(initialize_response(request.protocol_version)) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async move |_request: v2::NewSessionRequest, + responder: Responder, + _connection: V2ConnectionTo| { + responder.respond(v2::NewSessionResponse::new(agent_session_id.clone())) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async move |_request: v2::ListSessionsRequest, + responder: Responder, + connection: V2ConnectionTo| { + connection.send_notification(v2::UpdateSessionNotification::new( + v2::SessionId::new("dropped-handle"), + v2::SessionUpdate::AgentMessage( + v2::AgentMessage::new(v2::MessageId::new("background")) + .content(vec!["still routed".into()]), + ), + ))?; + responder.respond(v2::ListSessionsResponse::new(Vec::new())) + }, + agent_client_protocol::on_receive_request!(), + ); + + let (update_tx, mut update_rx) = mpsc::unbounded(); + let client = Client + .v2() + .on_receive_notification( + async move |update: v2::UpdateSessionNotification, + _connection: V2ConnectionTo| { + update_tx + .unbounded_send(update) + .map_err(Error::into_internal_error) + }, + agent_client_protocol::on_receive_notification!(), + ) + .connect_with(agent, async move |connection| { + connection + .send_request(v2::InitializeRequest::new( + ProtocolVersion::V2, + implementation(), + )) + .block_task() + .await?; + + let opened = connection + .build_session(cwd()?) + .start_session() + .block_task() + .await?; + let (session, _) = opened.into_parts(); + assert_eq!(session.session_id(), &session_id); + drop(session); + + let sessions = connection + .send_request(v2::ListSessionsRequest::new()) + .block_task() + .await?; + assert!(sessions.sessions.is_empty()); + + let update = next_update(&mut update_rx).await; + assert_eq!(update.session_id, session_id); + assert!(matches!( + update.update, + v2::SessionUpdate::AgentMessage(message) + if message.message_id == v2::MessageId::new("background") + )); + Ok(()) + }); + + tokio::time::timeout(TIMEOUT, client) + .await + .expect("v2 dropped-session update handling test timed out") + .expect("v2 dropped-session update handling test failed"); +} + +#[tokio::test(flavor = "current_thread")] +async fn v2_session_new_error_is_preserved_without_closing_connection() { + let agent = Agent + .v2() + .on_receive_request( + async |request: v2::InitializeRequest, + responder: Responder, + _connection: V2ConnectionTo| { + responder.respond(initialize_response(request.protocol_version)) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async |_request: v2::NewSessionRequest, + responder: Responder, + _connection: V2ConnectionTo| { + responder.respond_with_error(Error::invalid_params().data("session rejected")) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async |_request: v2::ListSessionsRequest, + responder: Responder, + _connection: V2ConnectionTo| { + responder.respond(v2::ListSessionsResponse::new(Vec::new())) + }, + agent_client_protocol::on_receive_request!(), + ); + + let client = Client.v2().connect_with(agent, async move |connection| { + connection + .send_request(v2::InitializeRequest::new( + ProtocolVersion::V2, + implementation(), + )) + .block_task() + .await?; + + let error = connection + .build_session(cwd()?) + .start_session() + .block_task() + .await + .expect_err("the session/new error must reach the caller"); + assert_eq!( + error, + Error::invalid_params().data("session rejected"), + "the helper must preserve the peer's JSON-RPC error" + ); + + let sessions = connection + .send_request(v2::ListSessionsRequest::new()) + .block_task() + .await?; + assert!(sessions.sessions.is_empty()); + Ok(()) + }); + + tokio::time::timeout(TIMEOUT, client) + .await + .expect("v2 session/new error test timed out") + .expect("v2 session/new error closed the connection"); +} + +#[tokio::test(flavor = "current_thread")] +async fn low_level_v2_runner_rejects_the_v1_session_helper() { + let (completed_tx, completed_rx) = oneshot::channel(); + let client = Client + .v2() + .with_runner(V1SessionGuardRunner { + cwd: cwd().expect("current directory should be available"), + completed: completed_tx, + }) + .connect_with(Agent.v2(), async move |_connection| { + completed_rx.await.map_err(Error::into_internal_error) + }); + + tokio::time::timeout(TIMEOUT, client) + .await + .expect("low-level v1 session guard test timed out") + .expect("low-level v1 session guard connection failed"); +}