From 384a13946841df9fb8a0a9c8ca2e59292b0fe318 Mon Sep 17 00:00:00 2001 From: Ben Brandt Date: Thu, 20 Aug 2026 12:05:59 -0700 Subject: [PATCH 1/3] feat(unstable-v2): Add runnable v2 quickstart examples --- Cargo.lock | 1 + README.md | 14 +- md/SUMMARY.md | 1 + md/mcp-bridge.md | 12 +- md/migration_v2.0.md | 9 +- md/protocol-v2-quickstart.md | 93 ++ md/protocol-v2.md | 156 ++-- .../CHANGELOG.md | 6 + src/agent-client-protocol-cookbook/Cargo.toml | 3 +- src/agent-client-protocol-cookbook/src/lib.rs | 173 +++- src/agent-client-protocol/CHANGELOG.md | 47 +- src/agent-client-protocol/Cargo.toml | 8 + src/agent-client-protocol/README.md | 28 +- .../examples/simple_agent_v2.rs | 472 ++++++++++ .../examples/v2_one_shot_client.rs | 262 ++++++ .../src/concepts/proxies.rs | 15 +- .../src/concepts/sessions.rs | 33 +- src/agent-client-protocol/src/jsonrpc.rs | 14 +- .../src/jsonrpc/incoming_actor.rs | 2 +- .../src/jsonrpc/outgoing_actor.rs | 8 +- .../src/jsonrpc/protocol_compat.rs | 781 +++++++++++++++- src/agent-client-protocol/src/lib.rs | 8 +- .../src/mcp_server/mod.rs | 6 +- .../src/mcp_server/server.rs | 4 +- src/agent-client-protocol/src/session/v2.rs | 197 +++++ .../tests/protocol_v2.rs | 834 ++++++++++++++---- .../tests/session_ordering.rs | 185 +++- src/agent-client-protocol/tests/session_v2.rs | 74 ++ .../tests/session_v2_mcp.rs | 204 ++++- 29 files changed, 3261 insertions(+), 389 deletions(-) create mode 100644 md/protocol-v2-quickstart.md create mode 100644 src/agent-client-protocol/examples/simple_agent_v2.rs create mode 100644 src/agent-client-protocol/examples/v2_one_shot_client.rs diff --git a/Cargo.lock b/Cargo.lock index a898f20e..5cca7202 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -65,6 +65,7 @@ version = "2.0.0" dependencies = [ "agent-client-protocol", "agent-client-protocol-rmcp", + "futures", "rmcp", "schemars 1.2.2", "serde", diff --git a/README.md b/README.md index d657f23a..f62a2572 100644 --- a/README.md +++ b/README.md @@ -29,9 +29,10 @@ when both `unstable_protocol_v2` and `unstable_mcp_over_acp` are enabled: supported session setup request, while `V2SessionBuilder::with_mcp_server(...)` attaches a server to one new session and `V2ResumeSessionBuilder::with_mcp_server(...)` attaches one while resuming. -Successful v2 attachments remain active for the connection lifetime, and both -builders expose `on_proxy_session_start` to forward proxied setup without -coupling later session events to that response. +With `unstable_session_fork`, `V2ForkSessionBuilder::with_mcp_server(...)` +attaches one while forking. Successful v2 attachments remain active for the +connection lifetime, and all three builders expose `on_proxy_session_start` to +forward proxied setup without coupling later session events to that response. **Proxy orchestration** @@ -50,8 +51,9 @@ coupling later session events to that response. - **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). + session usage are covered in [Protocol V2](./md/protocol-v2.md). To run a + complete v2 agent/client pair first, use the + [Runnable V2 Quickstart](./md/protocol-v2-quickstart.md). `Client.builder()`, `Agent.builder()`, and `Proxy.builder()` remain stable-v1 entry points; their `.v2()` counterparts select the draft-v2 API. With @@ -66,7 +68,7 @@ infrastructure can use `without_acp_version_guard`. - [Clients](https://agentclientprotocol.com/overview/clients) - Official Libraries - **Kotlin**: [`acp-kotlin`](https://github.com/agentclientprotocol/kotlin-sdk) – supports JVM, other targets are in progress - - **Rust**: [`agent-client-protocol`](https://crates.io/crates/agent-client-protocol) - See the [agent](./src/agent-client-protocol/examples/simple_agent.rs) and [client](./src/agent-client-protocol/examples/yolo_one_shot_client.rs) examples + - **Rust**: [`agent-client-protocol`](https://crates.io/crates/agent-client-protocol) - See the stable-v1 [agent](./src/agent-client-protocol/examples/simple_agent.rs) and [client](./src/agent-client-protocol/examples/yolo_one_shot_client.rs), or the draft-v2 [agent](./src/agent-client-protocol/examples/simple_agent_v2.rs) and [client](./src/agent-client-protocol/examples/v2_one_shot_client.rs) examples - **TypeScript**: [`@agentclientprotocol/sdk`](https://www.npmjs.com/package/@agentclientprotocol/sdk) - See [examples/](https://github.com/agentclientprotocol/typescript-sdk/tree/main/src/examples) - [Community Libraries](https://agentclientprotocol.com/libraries/community) diff --git a/md/SUMMARY.md b/md/SUMMARY.md index 350c0652..76274074 100644 --- a/md/SUMMARY.md +++ b/md/SUMMARY.md @@ -9,6 +9,7 @@ - [Request Cancellation](./request-cancellation.md) - [Configurable LLM Providers](./llm-providers.md) - [Protocol V2](./protocol-v2.md) +- [Runnable V2 Quickstart](./protocol-v2-quickstart.md) # Transports diff --git a/md/mcp-bridge.md b/md/mcp-bridge.md index 413e70f8..7caddd40 100644 --- a/md/mcp-bridge.md +++ b/md/mcp-bridge.md @@ -33,8 +33,10 @@ initialization, capability, session setup, and `mcp/*` wire types. It does not change the core attachment API. `Proxy.v2().with_mcp_server(...)` provides connection-global attachment. `V2SessionBuilder::with_mcp_server(...)` and `V2ResumeSessionBuilder::with_mcp_server(...)` provide per-session attachment -for new and resumed sessions respectively. The polyfill adapts their native -declarations when the final agent supports only HTTP MCP. +for new and resumed sessions respectively. With `unstable_session_fork`, +`V2ForkSessionBuilder::with_mcp_server(...)` provides per-session fork +attachment. The polyfill adapts their native declarations when the final agent +supports only HTTP MCP. ## Placement @@ -60,8 +62,10 @@ setup requests as `McpServer::Acp`; callers do not need to construct a transport placeholder themselves. In v2, `Proxy.v2().with_mcp_server(...)` provides connection-global attachment. `V2SessionBuilder::with_mcp_server(...)` and `V2ResumeSessionBuilder::with_mcp_server(...)` provide per-session attachment -for new and resumed sessions respectively. The polyfill translates those -native declarations at the final compatibility boundary. +for new and resumed sessions respectively. With `unstable_session_fork`, +`V2ForkSessionBuilder::with_mcp_server(...)` provides per-session fork +attachment. The polyfill translates those native declarations at the final +compatibility boundary. During initialization, the polyfill forwards the request to its successor. When the successor advertises HTTP MCP support, the polyfill advertises native ACP diff --git a/md/migration_v2.0.md b/md/migration_v2.0.md index 19fbc39f..d0938970 100644 --- a/md/migration_v2.0.md +++ b/md/migration_v2.0.md @@ -275,10 +275,11 @@ high-level server in the same place; the emitted declaration and wire methods ch builder attachment advertises the same server ID on `session/new`, `session/load`, `session/resume`, and feature-gated `session/fork`. Stable v1 per-session attachment remains specific to `session/new`; draft v2 additionally supports per-session resume attachment through -`V2ResumeSessionBuilder::with_mcp_server`. Do not construct an HTTP server with an `acp:` URL. If -the final agent accepts HTTP but not native ACP MCP servers, insert `McpOverAcpPolyfill` -immediately before it. The polyfill now consumes native `McpServer::Acp` declarations and adapts -only its final-agent-facing side. +`V2ResumeSessionBuilder::with_mcp_server` and feature-gated fork attachment through +`V2ForkSessionBuilder::with_mcp_server`. Do not construct an HTTP server with an `acp:` URL. If the +final agent accepts HTTP but not native ACP MCP servers, insert `McpOverAcpPolyfill` immediately +before it. The polyfill now consumes native `McpServer::Acp` declarations and adapts only its +final-agent-facing side. The polyfill's public `BridgeMode` enum and `McpOverAcpPolyfill::stdio` were removed because the required conductor `mcp` helper subcommand no longer exists. The polyfill has one supported mode; diff --git a/md/protocol-v2-quickstart.md b/md/protocol-v2-quickstart.md new file mode 100644 index 00000000..a97da360 --- /dev/null +++ b/md/protocol-v2-quickstart.md @@ -0,0 +1,93 @@ +# Runnable Protocol V2 Quickstart + +The core crate includes a small ACP v2 agent and client that run together over +stdio. Both are compiled examples behind the `unstable_protocol_v2` feature: + +- [`simple_agent_v2.rs`](https://github.com/agentclientprotocol/rust-sdk/blob/main/src/agent-client-protocol/examples/simple_agent_v2.rs) + implements initialization and the complete baseline session lifecycle. +- [`v2_one_shot_client.rs`](https://github.com/agentclientprotocol/rust-sdk/blob/main/src/agent-client-protocol/examples/v2_one_shot_client.rs) + initializes the agent, creates a session, sends one prompt, renders text + output, waits for the matching session to become idle, and closes it. + +## Run the pair + +Build both examples from the repository root: + +```bash +cargo build -p agent-client-protocol \ + --features unstable_protocol_v2 \ + --examples +``` + +Then point the client at the agent executable: + +```bash +./target/debug/examples/v2_one_shot_client \ + --command ./target/debug/examples/simple_agent_v2 \ + "Hello from ACP v2" +``` + +The result shows the two independent parts of a v2 prompt: + +```text +Prompt accepted; waiting for session output and completion... +Echo: Hello from ACP v2 +Session is idle: Some(EndTurn) +``` + +The agent writes only JSON-RPC to stdout because ACP uses stdout as the wire. +Write logs and diagnostics to stderr when extending it. + +## Client lifecycle + +The client installs its `session/update` and `session/request_permission` +handlers before it opens a session. Permission requests are part of the +baseline client surface and have no capability marker; this non-interactive +example cancels them explicitly. It then follows this sequence: + +1. Send `initialize` and verify that the agent advertised session support. +2. Send `session/new` and retain the returned command handle and session ID. +3. Send `session/prompt` and await its response. This only confirms acceptance. +4. Ignore queued updates for that new session until its foreground state becomes + `running`. The running update may already be queued when prompt acceptance + arrives. +5. Project subsequent message updates, and treat the next matching + `state_update` with `idle` as completion of foreground work. An idle update + queued before running is only the session's earlier ready state. Background + updates may still arrive afterward. +6. Send `session/close` when the client no longer needs the active session. + +Real clients normally maintain one shared update projection for every session. +Do not install a temporary handler after sending a prompt: updates can arrive +before the prompt response and are not scoped to a prompt or turn ID. +Within that projection, message chunks append by `messageId`; a later message +snapshot with concrete content replaces the accumulated chunks, `null` clears +them, and omitted content preserves them. Rendering chunks and then rendering a +snapshot again would duplicate output. + +## Agent lifecycle + +Advertising `AgentCapabilities::session(SessionCapabilities::new())` commits an +agent to the baseline session surface. The example handles: + +- `session/new` +- `session/list` +- `session/resume` +- `session/close` +- `session/prompt` +- `session/cancel` +- `session/update` notifications sent to the client + +The prompt handler validates and marks the session busy, responds to +`session/prompt` immediately, and moves the actual work into a spawned task so +the connection can continue dispatching cancellation and other traffic. That +task sends the accepted user message, a running update, output, and finally an +idle update with a stop reason. + +The example keeps history in memory and supports replay from the start before +the `session/resume` response. A production agent should replace this with +durable session storage, define its supported replay cursors, and make resource +cleanup and cancellation robust across process failure. + +For the connection APIs, proxy routing, and compatibility details surrounding +these examples, continue with [Protocol V2](./protocol-v2.md). diff --git a/md/protocol-v2.md b/md/protocol-v2.md index feaf9c0f..951727af 100644 --- a/md/protocol-v2.md +++ b/md/protocol-v2.md @@ -10,6 +10,11 @@ agent-client-protocol = { version = "...", features = ["unstable_protocol_v2"] } This feature is separate from the broad `unstable` feature because protocol v2 is a versioning experiment, not just an unstable method family. +To start from working code, build and run the companion agent and client in the +[Runnable Protocol V2 Quickstart](./protocol-v2-quickstart.md). The examples +exercise prompt acceptance, independent session updates, and the terminal idle +state over a real stdio connection. + ## JSON-RPC batches Batch framing is a shared JSON-RPC transport feature, not a v2-only protocol @@ -74,26 +79,32 @@ Agent When v2 mode is enabled, application code should use types from `agent_client_protocol::schema::v2`. The flat `agent_client_protocol::schema::*` exports remain the stable v1 schema. This will likely change as v2 gets closer -to release. +to release. The preceding agent fragment demonstrates version negotiation only; +an agent that advertises session support must also implement the complete +baseline session surface shown in the runnable quickstart. ## 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. The `resume_session*` helpers return a -`V2ResumeSessionBuilder`; call `start_session` to publish the request and obtain -an `OpenedV2Session` containing the command handle and complete -`ResumeSessionResponse`. +the v2 `build_session*` and `resume_session*` helpers, plus feature-gated +`fork_session*` helpers when `unstable_session_fork` is enabled. The shared +names describe the same lifecycle operations while the connection type selects +their schema and return types at compile time. Resume and fork return +`V2ResumeSessionBuilder` and `V2ForkSessionBuilder`; call `start_session` to +publish the request and obtain an `OpenedV2Session` containing the command +handle and complete operation-specific response. 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<_>`. +remains schema-agnostic. These generic APIs do not infer a protocol version from +the Rust payload type: callers on a v2 connection must use `schema::v2` types, +or deliberately send extension or untyped messages. The connection guard +enforces negotiation and initialization lifecycle, not Rust-type provenance. +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<_>`, @@ -109,62 +120,14 @@ 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. +A complete client installs update and interactive-request handlers before +connecting. After `session/prompt` is accepted, it must keep the connection +alive and consume updates until the matching session reaches idle. The +[`v2_one_shot_client`](https://github.com/agentclientprotocol/rust-sdk/blob/main/src/agent-client-protocol/examples/v2_one_shot_client.rs) +example demonstrates the full sequence, while the compiled cookbook +`v2_one_shot_prompt` recipe shows how to embed it in an application. Permission +handlers should transfer the request and responder to application-owned work +rather than waiting for user input inside the dispatch callback. V2 deliberately separates prompt submission from session observation: @@ -180,6 +143,10 @@ V2 deliberately separates prompt submission from session observation: `V2ResumeSessionBuilder`. Its `start_session` method publishes `session/resume` and returns an `OpenedV2Session` containing the complete `ResumeSessionResponse` without reconstructing it. +- With `unstable_session_fork`, `V2ConnectionTo::fork_session` and + `fork_session_from` return a `V2ForkSessionBuilder`. Its `start_session` + publishes `session/fork`, preserves the complete `ForkSessionResponse`, and + uses that response's newly allocated session ID for the command handle. - `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` @@ -207,25 +174,27 @@ V2 deliberately separates prompt submission from session observation: `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 before calling `start_session` on a -`V2ResumeSessionBuilder`: 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. +Install connection handlers before `session/new`, `session/resume`, and +feature-gated `session/fork` requests. This is especially important before +calling `start_session` on a `V2ResumeSessionBuilder`: 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. For a session configured with `V2SessionBuilder::with_mcp_server` or -`V2ResumeSessionBuilder::with_mcp_server`, the SDK installs the MCP routes and -initially polls their runner tasks before publishing `session/new` or -`session/resume`, so the agent can connect to those servers during setup or -resume replay. Runners may continue asynchronous initialization; custom -connectors must be able to queue connections and messages once constructed. A -successful setup promotes the attachment to the connection lifetime; a setup -failure, including an error response after cancellation, cleans up the pending +`V2ResumeSessionBuilder::with_mcp_server`, or feature-gated +`V2ForkSessionBuilder::with_mcp_server`, the SDK installs the MCP routes and +initially polls their runner tasks before publishing the corresponding setup +request, so the agent can connect to those servers during setup or resume +replay. Runners may continue asynchronous initialization; custom connectors +must be able to queue connections and messages once constructed. A successful +setup promotes the attachment to the connection lifetime; a setup failure, +including an error response after cancellation, cleans up the pending attachment. This attachment requires both `unstable_protocol_v2` and -`unstable_mcp_over_acp`. +`unstable_mcp_over_acp`; fork additionally requires +`unstable_session_fork`. A v2 proxy can instead attach one server globally with `Proxy.v2().with_mcp_server(...)`. The proxy reuses one connection-scoped @@ -234,8 +203,9 @@ feature-gated `session/fork` requests. It modifies only the `mcpServers` field, preserving unrelated setup fields and extensions for downstream handlers. `V2SessionBuilder::on_proxy_session_start` and -`V2ResumeSessionBuilder::on_proxy_session_start` are the non-blocking setup -helpers for a v2 proxy: +`V2ResumeSessionBuilder::on_proxy_session_start`, plus +`V2ForkSessionBuilder::on_proxy_session_start` when enabled, are the +non-blocking setup helpers for a v2 proxy: ```rust,ignore use agent_client_protocol::schema::v2; @@ -258,18 +228,19 @@ Proxy ); ``` -Both helpers forward request cancellation, send an ordered downstream setup +These helpers forward request cancellation, send an ordered downstream setup request, and forward the complete operation-specific response without -reconstruction. For `session/new`, routing is installed when its response makes -the new session ID available and before later inbound traffic is dispatched. -For `session/resume`, the ID is already known, so routing and any per-session +reconstruction. For `session/new` and `session/fork`, routing is installed when +the response makes the new session ID available and before later inbound +traffic is dispatched. Fork routing uses the response's new ID rather than the +source session ID. For `session/resume`, the ID is already known, so routing and any per-session MCP attachment are ready before the downstream request is published. Replay updates can therefore be forwarded upstream before the complete `ResumeSessionResponse`, as required by the protocol. A failed or cancelled downstream response drops pending routing and MCP attachment; successful setup keeps them for the connection lifetime. A cancellation signal itself remains advisory: it is forwarded downstream while the helper awaits that response. -The helper then spawns the callback outside the ordering barrier with an +Each helper then spawns the callback outside the ordering barrier with an `OpenedV2Session` containing the command-only session handle and complete setup response. Updates and interactive requests remain independent connection traffic and should still be handled by typed callbacks on `Proxy.v2()`. @@ -347,6 +318,13 @@ as described above. The SDK handles the `initialize` negotiation at the JSON-RPC boundary: +- Native `Client.v2()` and `Agent.v2()` connections reject ordinary protocol + traffic until the initialization response completes; `$/cancel_request` + remains available while initialization is in progress. The client is the + initializer and the agent is the responder; attempts in the opposite + direction are rejected. An initialization error leaves the connection + uninitialized so the client can retry, while a second initialization after a + successful handshake is rejected. - A v2 client advertises protocol v2 as its latest supported version. - A v2 client requires a v2 agent. If the agent responds with v1, the `initialize` request resolves with an error and the caller must explicitly diff --git a/src/agent-client-protocol-cookbook/CHANGELOG.md b/src/agent-client-protocol-cookbook/CHANGELOG.md index 3f5eb618..bdab0842 100644 --- a/src/agent-client-protocol-cookbook/CHANGELOG.md +++ b/src/agent-client-protocol-cookbook/CHANGELOG.md @@ -7,6 +7,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added + +- Add a compiled draft-v2 one-shot prompt recipe that handles initialization, + permissions, message patch projection, prompt acceptance, idle completion, + and session close. + ## [2.0.0](https://github.com/agentclientprotocol/rust-sdk/compare/agent-client-protocol-cookbook-v1.3.0...agent-client-protocol-cookbook-v2.0.0) - 2026-07-23 ### Changed diff --git a/src/agent-client-protocol-cookbook/Cargo.toml b/src/agent-client-protocol-cookbook/Cargo.toml index 741d3152..7c5893fd 100644 --- a/src/agent-client-protocol-cookbook/Cargo.toml +++ b/src/agent-client-protocol-cookbook/Cargo.toml @@ -14,8 +14,9 @@ categories = ["development-tools"] [dependencies] [dev-dependencies] -agent-client-protocol = { workspace = true, features = ["unstable_mcp_over_acp"] } +agent-client-protocol = { workspace = true, features = ["unstable_mcp_over_acp", "unstable_protocol_v2"] } agent-client-protocol-rmcp.workspace = true +futures.workspace = true rmcp.workspace = true schemars.workspace = true serde.workspace = true diff --git a/src/agent-client-protocol-cookbook/src/lib.rs b/src/agent-client-protocol-cookbook/src/lib.rs index b33cdf31..41866e45 100644 --- a/src/agent-client-protocol-cookbook/src/lib.rs +++ b/src/agent-client-protocol-cookbook/src/lib.rs @@ -15,6 +15,7 @@ //! [`Client.builder()`](agent_client_protocol::Client) to build connections. //! //! - [`one_shot_prompt`] - Send a single prompt and get a response (simplest pattern) +//! - [`v2_one_shot_prompt`] - Send a draft-v2 prompt and wait for the independent idle update //! - [`connecting_as_client`] - More details on connection setup and permission handling //! //! # Building Proxies @@ -120,6 +121,152 @@ pub mod one_shot_prompt { //! [`RequestPermissionRequest`]: agent_client_protocol::schema::v1::RequestPermissionRequest } +pub mod v2_one_shot_prompt { + //! Pattern: One prompt with the draft protocol v2 lifecycle. + //! + //! A successful v2 `session/prompt` response acknowledges acceptance; it + //! does not contain output and does not mean the work is complete. Install + //! update handlers before connecting, then consume matching updates until + //! the new session reports `running` and subsequently reports `idle`. + //! + //! This module is compiled with the cookbook's v2 feature coverage. For a + //! runnable CLI pair, see the SDK's `simple_agent_v2` and + //! `v2_one_shot_client` examples. + //! + //! # Example + //! + //! ``` + //! use std::collections::HashMap; + //! use agent_client_protocol::{Agent, Client, ConnectTo, Error, Responder, V2ConnectionTo}; + //! use agent_client_protocol::schema::{MaybeUndefined, ProtocolVersion, v2}; + //! use futures::{StreamExt, channel::mpsc}; + //! + //! #[derive(Default)] + //! struct AgentTextProjection { + //! order: Vec, + //! messages: HashMap>, + //! } + //! + //! impl AgentTextProjection { + //! fn apply(&mut self, update: v2::SessionUpdate) { + //! match update { + //! v2::SessionUpdate::AgentMessageChunk(chunk) => { + //! self.message_content(chunk.message_id).push(chunk.content); + //! } + //! v2::SessionUpdate::AgentMessage(message) => { + //! let content = self.message_content(message.message_id); + //! match message.content { + //! // Snapshots patch chunks accumulated for the same message ID. + //! MaybeUndefined::Undefined => {} + //! MaybeUndefined::Null => content.clear(), + //! MaybeUndefined::Value(replacement) => *content = replacement, + //! } + //! } + //! _ => {} + //! } + //! } + //! + //! fn message_content( + //! &mut self, + //! message_id: v2::MessageId, + //! ) -> &mut Vec { + //! if !self.messages.contains_key(&message_id) { + //! self.order.push(message_id.clone()); + //! } + //! self.messages.entry(message_id).or_default() + //! } + //! + //! fn text(&self) -> String { + //! self.order + //! .iter() + //! .filter_map(|message_id| self.messages.get(message_id)) + //! .flatten() + //! .filter_map(|content| match content { + //! v2::ContentBlock::Text(text) => Some(text.text.as_str()), + //! _ => None, + //! }) + //! .collect() + //! } + //! } + //! + //! async fn ask_agent( + //! transport: impl ConnectTo + 'static, + //! prompt: &str, + //! ) -> Result { + //! let (update_tx, mut update_rx) = mpsc::unbounded(); + //! + //! 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| { + //! // This non-interactive recipe rejects permission requests. + //! responder.respond(v2::RequestPermissionResponse::new( + //! v2::RequestPermissionOutcome::Cancelled, + //! )) + //! }, + //! agent_client_protocol::on_receive_request!(), + //! ) + //! .connect_with(transport, async move |connection| { + //! let initialized = connection + //! .send_request(v2::InitializeRequest::new( + //! ProtocolVersion::V2, + //! v2::Implementation::new("example-client", "0.1.0"), + //! )) + //! .block_task() + //! .await?; + //! if initialized.capabilities.session.is_none() { + //! return Err(Error::invalid_params() + //! .data("agent did not advertise session support")); + //! } + //! + //! let session = connection + //! .build_session_cwd()? + //! .start_session() + //! .block_task() + //! .await? + //! .into_session(); + //! + //! // This response only means that the prompt was accepted. + //! session.send_prompt(prompt).block_task().await?; + //! + //! let mut projection = AgentTextProjection::default(); + //! let mut observed_running = false; + //! while let Some(notification) = update_rx.next().await { + //! if ¬ification.session_id != session.session_id() { + //! continue; + //! } + //! match notification.update { + //! v2::SessionUpdate::StateUpdate(v2::StateUpdate::Running(_)) => { + //! observed_running = true; + //! } + //! v2::SessionUpdate::StateUpdate(v2::StateUpdate::Idle(_)) + //! if observed_running => + //! { + //! session.close().block_task().await?; + //! return Ok(projection.text()); + //! } + //! update if observed_running => projection.apply(update), + //! _ => {} + //! } + //! } + //! Err(Error::internal_error() + //! .data("agent disconnected before the prompt ran to completion")) + //! }) + //! .await + //! } + //! ``` +} + pub mod connecting_as_client { //! Pattern: Connecting as a client. //! @@ -660,10 +807,11 @@ pub mod per_session_mcp_server { //! //! `Proxy.v2()` exposes the same non-blocking setup shape with v2 schema //! types. `V2SessionBuilder` handles `session/new`, while - //! `V2ResumeSessionBuilder` handles `session/resume`. Their - //! `on_proxy_session_start` callbacks receive an `OpenedV2Session`, not - //! just a session ID, so they retain both the command-only handle and the - //! complete operation-specific response: + //! `V2ResumeSessionBuilder` handles `session/resume`. With + //! `unstable_session_fork`, `V2ForkSessionBuilder` handles `session/fork`. + //! Their `on_proxy_session_start` callbacks receive an `OpenedV2Session`, + //! not just a session ID, so they retain both the command-only handle and + //! the complete operation-specific response: //! //! ```rust,ignore //! use agent_client_protocol::schema::v2; @@ -713,12 +861,17 @@ pub mod per_session_mcp_server { //! ); //! ``` //! - //! Both helpers forward upstream cancellation and the complete setup - //! response. New-session routing is installed before later inbound - //! traffic; resume routing and MCP readiness are established before the - //! request is published so replay can precede its response. The callback - //! runs outside the ordering barrier. V2 updates and interactive requests - //! remain independent connection traffic. + //! Fork uses the same terminal helper after + //! `connection.fork_session_from(request)`. Its returned handle and route + //! use the newly allocated ID from the complete `ForkSessionResponse`, not + //! the source session ID. + //! + //! All helpers forward upstream cancellation and the complete setup + //! response. New-session and fork routing are installed before later + //! inbound traffic; resume routing and MCP readiness are established + //! before the request is published so replay can precede its response. The + //! callback runs outside the ordering barrier. V2 updates and interactive + //! requests remain independent connection traffic. //! //! # Stable v1 alternative: spawning `start_session_proxy` //! diff --git a/src/agent-client-protocol/CHANGELOG.md b/src/agent-client-protocol/CHANGELOG.md index 8d613e03..aff8de80 100644 --- a/src/agent-client-protocol/CHANGELOG.md +++ b/src/agent-client-protocol/CHANGELOG.md @@ -4,6 +4,11 @@ ### Added +- *(unstable-v2)* Add runnable draft-v2 agent and one-shot client examples. The + agent implements the complete baseline session lifecycle; the client handles + permissions, projects chunk and snapshot updates by message ID, and waits for + the matching idle state. Add a compiled cookbook recipe and mdbook quickstart + for the same lifecycle. - *(unstable-v2)* Add `Proxy::protocol_router` and `ProxyProtocolRouter` to compose strict v1 and v2 proxy implementations behind one connection. Routing requires the exact version selected by the conductor and preserves @@ -13,23 +18,25 @@ - *(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. The resume helpers return a - `V2ResumeSessionBuilder`, so `session/resume` is not published until - `start_session` or `on_proxy_session_start` is called. Session updates and - interactive requests remain on typed connection handlers. With - `unstable_mcp_over_acp`, both session builders support native per-session MCP - attachment while preserving independent response consumption. MCP routes - are installed and runners begin executing before the setup request is - published; successful attachments remain active for the connection lifetime - while setup failures, including cancellation error responses, clean up - pending state. The builders' `on_proxy_session_start` helpers forward the - complete setup response and cancellation, install routing before later - inbound traffic, and then spawn user work with the `OpenedV2Session`. Resume - routing is ready before request publication so replay can precede the + independent prompt-acceptance requests, create, resume, and feature-gated + fork setup responses, cloneable command handles, configuration, close, and + session-wide cancellation. `Client.v2()` and `Agent.v2()` callbacks receive + a `V2ConnectionTo` whose unversioned `build_session*`, `resume_session*`, and + feature-gated `fork_session*` methods expose only the v2 lifecycle. Resume + and fork return `V2ResumeSessionBuilder` and `V2ForkSessionBuilder`, so their + requests are not published until `start_session` or + `on_proxy_session_start` is called. Session updates and interactive requests + remain on typed connection handlers. With `unstable_mcp_over_acp`, all + available session builders support native per-session MCP attachment while + preserving independent response consumption. MCP routes are installed and + runners begin executing before the setup request is published; successful + attachments remain active for the connection lifetime while setup failures, + including cancellation error responses, clean up pending state. The + builders' `on_proxy_session_start` helpers forward the complete setup + response and cancellation, install routing before later inbound traffic, and + then spawn user work with the `OpenedV2Session`. New and fork builders use + the response session ID for the installed route and returned command handle; + resume routing is ready before request publication so replay can precede the response as required by the protocol. - *(unstable-v2)* Add `Proxy::v2()` as the draft-v2-only proxy builder while keeping `Proxy::builder()` on stable v1. With `unstable_mcp_over_acp`, @@ -59,6 +66,12 @@ ### Fixed +- *(unstable-v2)* Require native v2 client and agent connections to complete + their single initialization handshake before sending or accepting other + protocol traffic. Reject initialization in the wrong direction and reject + reinitialization after a successful handshake while allowing retries after an + initialization error. Protocol-level request cancellation remains available + while initialization is in progress. - *(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 diff --git a/src/agent-client-protocol/Cargo.toml b/src/agent-client-protocol/Cargo.toml index 491cc4da..51352ed0 100644 --- a/src/agent-client-protocol/Cargo.toml +++ b/src/agent-client-protocol/Cargo.toml @@ -15,6 +15,14 @@ categories = ["development-tools"] all-features = true rustdoc-args = ["--cfg", "docsrs"] +[[example]] +name = "simple_agent_v2" +required-features = ["unstable_protocol_v2"] + +[[example]] +name = "v2_one_shot_client" +required-features = ["unstable_protocol_v2"] + [features] default = [] diff --git a/src/agent-client-protocol/README.md b/src/agent-client-protocol/README.md index faca129f..d146e5d9 100644 --- a/src/agent-client-protocol/README.md +++ b/src/agent-client-protocol/README.md @@ -51,6 +51,26 @@ connection handlers. See [Protocol V2](https://agentclientprotocol.github.io/rus implementations as one component; custom raw routing infrastructure can use `Proxy.builder().without_acp_version_guard()`. +### Runnable draft-v2 pair + +The `simple_agent_v2` example implements the complete baseline session +lifecycle, and `v2_one_shot_client` demonstrates the split prompt lifecycle: +the prompt response acknowledges acceptance, while output and completion arrive +through `session/update` notifications. + +```bash +cargo build -p agent-client-protocol \ + --features unstable_protocol_v2 \ + --examples + +./target/debug/examples/v2_one_shot_client \ + --command ./target/debug/examples/simple_agent_v2 \ + "Hello from ACP v2" +``` + +See the [Runnable Protocol V2 Quickstart](https://agentclientprotocol.github.io/rust-sdk/protocol-v2-quickstart.html) +for the lifecycle invariants to preserve when adapting these examples. + ## MCP Server Attachment The runtime-agnostic `mcp_server` module can build and directly serve standalone @@ -64,8 +84,10 @@ protocol v2 supports both scopes when both unstable features are enabled: `Proxy.v2().with_mcp_server(...)` injects a global server into supported setup requests, `V2SessionBuilder::with_mcp_server(...)` attaches one to a single `session/new`, and `V2ResumeSessionBuilder::with_mcp_server(...)` attaches one -to a single `session/resume`. Successful attachments remain active for the -connection lifetime. A v2 proxy can forward either setup operation with the +to a single `session/resume`. With `unstable_session_fork`, +`V2ForkSessionBuilder::with_mcp_server(...)` attaches one to a single +`session/fork`. Successful attachments remain active for the connection +lifetime. A v2 proxy can forward any of these setup operations with the builder's `on_proxy_session_start`; updates and interactive requests remain independent connection traffic. @@ -74,7 +96,7 @@ independent connection traffic. See the [crate documentation](https://docs.rs/agent-client-protocol) for: - **[Cookbook](https://docs.rs/agent-client-protocol-cookbook)** — Patterns for building clients, proxies, and agents -- **[Examples](https://github.com/agentclientprotocol/rust-sdk/tree/main/src/agent-client-protocol/examples)** — Working code you can run +- **[Examples](https://github.com/agentclientprotocol/rust-sdk/tree/main/src/agent-client-protocol/examples)** — Runnable stable-v1 and draft-v2 clients and agents ## Related Crates diff --git a/src/agent-client-protocol/examples/simple_agent_v2.rs b/src/agent-client-protocol/examples/simple_agent_v2.rs new file mode 100644 index 00000000..09ac3a92 --- /dev/null +++ b/src/agent-client-protocol/examples/simple_agent_v2.rs @@ -0,0 +1,472 @@ +//! A small, complete ACP v2 agent that echoes prompts over stdio. +//! +//! The agent implements the baseline v2 session lifecycle: new, list, resume, +//! close, prompt, cancel, and update. It keeps session history in memory so a +//! client can request replay from the beginning when resuming a session. +//! +//! Run it with the companion `v2_one_shot_client` example. An ACP agent owns +//! stdout for JSON-RPC, so diagnostics belong on stderr. + +use std::{ + collections::HashMap, + sync::{Arc, Mutex}, +}; + +use agent_client_protocol::schema::v2; +use agent_client_protocol::{Agent, Client, Error, Responder, Result, Stdio, V2ConnectionTo}; +use tokio::sync::Notify; + +#[derive(Clone, Debug, Default)] +struct EchoAgent { + state: Arc>, + state_changed: Arc, +} + +#[derive(Debug, Default)] +struct AgentState { + sessions: HashMap, + next_session_id: u64, + next_message_id: u64, +} + +#[derive(Clone, Debug)] +struct Session { + cwd: v2::AbsolutePath, + additional_directories: Vec, + active: bool, + foreground_work: bool, + cancelled: bool, + history: Vec, +} + +impl EchoAgent { + fn create_session(&self, request: v2::NewSessionRequest) -> v2::SessionId { + let mut state = self.state.lock().expect("session state lock poisoned"); + state.next_session_id += 1; + let session_id = v2::SessionId::new(format!("echo-session-{}", state.next_session_id)); + state.sessions.insert( + session_id.clone(), + Session { + cwd: request.cwd, + additional_directories: request.additional_directories, + active: true, + foreground_work: false, + cancelled: false, + history: Vec::new(), + }, + ); + session_id + } + + fn list_sessions(&self, request: &v2::ListSessionsRequest) -> Vec { + let state = self.state.lock().expect("session state lock poisoned"); + let mut sessions = state + .sessions + .iter() + .filter(|(_, session)| request.cwd.as_ref().is_none_or(|cwd| cwd == &session.cwd)) + .map(|(session_id, session)| { + v2::SessionInfo::new(session_id.clone(), session.cwd.clone()) + .additional_directories(session.additional_directories.clone()) + .title("Echo agent session") + }) + .collect::>(); + sessions.sort_by(|left, right| { + left.session_id + .to_string() + .cmp(&right.session_id.to_string()) + }); + sessions + } + + fn resume_session(&self, request: &v2::ResumeSessionRequest) -> Result> { + let mut state = self.state.lock().expect("session state lock poisoned"); + let session = state + .sessions + .get_mut(&request.session_id) + .ok_or_else(|| invalid_params(format!("unknown session `{}`", request.session_id)))?; + if session.foreground_work { + return Err(invalid_params(format!( + "session `{}` still has foreground work", + request.session_id + ))); + } + if session.cwd != request.cwd { + return Err(invalid_params(format!( + "session `{}` has a different working directory", + request.session_id + ))); + } + + let history = match &request.replay_from { + None => Vec::new(), + Some(v2::ReplayFrom::Start(_)) => session.history.clone(), + Some(_) => return Err(invalid_params("unsupported replay cursor")), + }; + session + .additional_directories + .clone_from(&request.additional_directories); + session.active = true; + session.cancelled = false; + Ok(history) + } + + fn begin_prompt(&self, session_id: &v2::SessionId) -> Result<()> { + let mut state = self.state.lock().expect("session state lock poisoned"); + let session = state + .sessions + .get_mut(session_id) + .ok_or_else(|| invalid_params(format!("unknown session `{session_id}`")))?; + if !session.active { + return Err(invalid_params(format!("closed session `{session_id}`"))); + } + if session.foreground_work { + return Err(invalid_params(format!( + "session `{session_id}` already has foreground work" + ))); + } + session.foreground_work = true; + session.cancelled = false; + Ok(()) + } + + fn finish_prompt( + &self, + session_id: &v2::SessionId, + connection: &V2ConnectionTo, + ) -> Result<()> { + let mut state = self.state.lock().expect("session state lock poisoned"); + let session = state + .sessions + .get_mut(session_id) + .ok_or_else(|| invalid_params(format!("unknown session `{session_id}`")))?; + let stop_reason = if session.cancelled || !session.active { + v2::StopReason::Cancelled + } else { + v2::StopReason::EndTurn + }; + send_update( + connection, + session_id, + v2::SessionUpdate::StateUpdate(v2::StateUpdate::Idle( + v2::IdleStateUpdate::new().stop_reason(stop_reason), + )), + )?; + session.foreground_work = false; + session.cancelled = false; + drop(state); + self.state_changed.notify_waiters(); + Ok(()) + } + + fn abandon_prompt(&self, session_id: &v2::SessionId) { + if let Some(session) = self + .state + .lock() + .expect("session state lock poisoned") + .sessions + .get_mut(session_id) + { + session.foreground_work = false; + session.cancelled = false; + } + self.state_changed.notify_waiters(); + } + + fn next_message_id(&self, kind: &str) -> v2::MessageId { + let mut state = self.state.lock().expect("session state lock poisoned"); + state.next_message_id += 1; + v2::MessageId::new(format!("{kind}-{}", state.next_message_id)) + } + + fn record_history(&self, session_id: &v2::SessionId, update: v2::SessionUpdate) { + if let Some(session) = self + .state + .lock() + .expect("session state lock poisoned") + .sessions + .get_mut(session_id) + { + session.history.push(update); + } + } + + fn is_cancelled(&self, session_id: &v2::SessionId) -> bool { + self.state + .lock() + .expect("session state lock poisoned") + .sessions + .get(session_id) + .is_none_or(|session| session.cancelled || !session.active) + } + + fn cancel(&self, session_id: &v2::SessionId) { + if let Some(session) = self + .state + .lock() + .expect("session state lock poisoned") + .sessions + .get_mut(session_id) + && session.foreground_work + { + session.cancelled = true; + } + self.state_changed.notify_waiters(); + } + + fn close(&self, session_id: &v2::SessionId) -> Result { + let mut state = self.state.lock().expect("session state lock poisoned"); + let session = state + .sessions + .get_mut(session_id) + .ok_or_else(|| invalid_params(format!("unknown session `{session_id}`")))?; + session.active = false; + if session.foreground_work { + session.cancelled = true; + } + let foreground_work = session.foreground_work; + drop(state); + self.state_changed.notify_waiters(); + Ok(foreground_work) + } + + fn has_foreground_work(&self, session_id: &v2::SessionId) -> bool { + self.state + .lock() + .expect("session state lock poisoned") + .sessions + .get(session_id) + .is_some_and(|session| session.foreground_work) + } + + async fn wait_for_foreground_work(&self, session_id: &v2::SessionId) { + let notified = self.state_changed.notified(); + tokio::pin!(notified); + loop { + notified.as_mut().enable(); + if !self.has_foreground_work(session_id) { + return; + } + notified.as_mut().await; + notified.set(self.state_changed.notified()); + } + } + + async fn process_prompt( + &self, + request: v2::PromptRequest, + connection: V2ConnectionTo, + ) -> Result<()> { + let session_id = request.session_id.clone(); + let result = self.process_prompt_inner(request, &connection).await; + match result { + Ok(()) => self.finish_prompt(&session_id, &connection), + Err(error) => { + self.abandon_prompt(&session_id); + Err(error) + } + } + } + + async fn process_prompt_inner( + &self, + request: v2::PromptRequest, + connection: &V2ConnectionTo, + ) -> Result<()> { + let session_id = request.session_id; + let user_message = v2::SessionUpdate::UserMessage( + v2::UserMessage::new(self.next_message_id("user-message")) + .content(request.prompt.clone()), + ); + send_update(connection, &session_id, user_message.clone())?; + self.record_history(&session_id, user_message); + + send_update( + connection, + &session_id, + v2::SessionUpdate::StateUpdate(v2::StateUpdate::Running(v2::RunningStateUpdate::new())), + )?; + + // Prompt acceptance is independent from this work. Yield once so the + // client can observe the response before output starts arriving. + tokio::task::yield_now().await; + if self.is_cancelled(&session_id) { + return Ok(()); + } + + let prompt = request + .prompt + .iter() + .filter_map(|block| match block { + v2::ContentBlock::Text(text) => Some(text.text.as_str()), + _ => None, + }) + .collect::>() + .join(" "); + let response = if prompt.is_empty() { + "I received a prompt without text content.".to_string() + } else { + format!("Echo: {prompt}") + }; + let message_id = self.next_message_id("agent-message"); + send_update( + connection, + &session_id, + v2::SessionUpdate::AgentMessageChunk(v2::ContentChunk::new( + response.clone().into(), + message_id.clone(), + )), + )?; + // A complete snapshot replaces the content accumulated from chunks + // with the same message ID; clients must not render it a second time. + let agent_message = v2::SessionUpdate::AgentMessage( + v2::AgentMessage::new(message_id).content(vec![response.into()]), + ); + send_update(connection, &session_id, agent_message.clone())?; + self.record_history(&session_id, agent_message); + Ok(()) + } +} + +#[tokio::main] +async fn main() -> Result<()> { + let agent = EchoAgent::default(); + + Agent + .v2() + .name("simple-agent-v2") + .on_receive_request( + async |request: v2::InitializeRequest, + responder: Responder, + _connection: V2ConnectionTo| { + responder.respond( + v2::InitializeResponse::new( + request.protocol_version, + v2::Implementation::new("simple-agent-v2", env!("CARGO_PKG_VERSION")), + ) + .capabilities( + v2::AgentCapabilities::new().session(v2::SessionCapabilities::new()), + ), + ) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + { + let agent = agent.clone(); + async move |request: v2::NewSessionRequest, + responder: Responder, + connection: V2ConnectionTo| { + let session_id = agent.create_session(request); + responder.respond(v2::NewSessionResponse::new(session_id.clone()))?; + // A client can already have this ready-state update queued + // when a later prompt response arrives. Prompt completion + // must wait for running and the subsequent idle instead. + send_update( + &connection, + &session_id, + v2::SessionUpdate::StateUpdate(v2::StateUpdate::Idle( + v2::IdleStateUpdate::new(), + )), + ) + } + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + { + let agent = agent.clone(); + async move |request: v2::ListSessionsRequest, + responder: Responder, + _connection: V2ConnectionTo| { + responder.respond(v2::ListSessionsResponse::new(agent.list_sessions(&request))) + } + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + { + let agent = agent.clone(); + async move |request: v2::ResumeSessionRequest, + responder: Responder, + connection: V2ConnectionTo| { + let history = match agent.resume_session(&request) { + Ok(history) => history, + Err(error) => return responder.respond_with_error(error), + }; + for update in history { + send_update(&connection, &request.session_id, update)?; + } + responder.respond(v2::ResumeSessionResponse::new()) + } + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + { + let agent = agent.clone(); + async move |request: v2::CloseSessionRequest, + responder: Responder, + connection: V2ConnectionTo| { + if !agent.close(&request.session_id)? { + return responder.respond(v2::CloseSessionResponse::new()); + } + let waiting_agent = agent.clone(); + connection.spawn(async move { + waiting_agent + .wait_for_foreground_work(&request.session_id) + .await; + responder.respond(v2::CloseSessionResponse::new()) + }) + } + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + { + let agent = agent.clone(); + async move |request: v2::PromptRequest, + responder: Responder, + connection: V2ConnectionTo| { + agent.begin_prompt(&request.session_id)?; + responder.respond(v2::PromptResponse::new())?; + + let prompt_connection = connection.clone(); + let session_id = request.session_id.clone(); + if let Err(error) = connection.spawn({ + let agent = agent.clone(); + async move { agent.process_prompt(request, prompt_connection).await } + }) { + agent.abandon_prompt(&session_id); + return Err(error); + } + Ok(()) + } + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_notification( + async move |notification: v2::CancelSessionNotification, + _connection: V2ConnectionTo| { + agent.cancel(¬ification.session_id); + Ok(()) + }, + agent_client_protocol::on_receive_notification!(), + ) + .connect_to(Stdio::new()) + .await +} + +fn send_update( + connection: &V2ConnectionTo, + session_id: &v2::SessionId, + update: v2::SessionUpdate, +) -> Result<()> { + connection.send_notification(v2::UpdateSessionNotification::new( + session_id.clone(), + update, + )) +} + +fn invalid_params(message: impl ToString) -> Error { + Error::invalid_params().data(message.to_string()) +} diff --git a/src/agent-client-protocol/examples/v2_one_shot_client.rs b/src/agent-client-protocol/examples/v2_one_shot_client.rs new file mode 100644 index 00000000..cc10009a --- /dev/null +++ b/src/agent-client-protocol/examples/v2_one_shot_client.rs @@ -0,0 +1,262 @@ +//! A one-shot ACP v2 client that waits for the agent's idle update. +//! +//! Unlike ACP v1, a successful v2 `session/prompt` response only means the +//! prompt was accepted. Output and completion arrive independently through +//! `session/update`, so this example keeps receiving updates until the same +//! session reports that its foreground work is idle. +//! +//! ```text +//! cargo run -p agent-client-protocol --features unstable_protocol_v2 \ +//! --example v2_one_shot_client -- \ +//! --command ./target/debug/examples/simple_agent_v2 \ +//! "What should a v2 client wait for?" +//! ``` + +use std::{collections::HashMap, str::FromStr}; + +use agent_client_protocol::schema::{MaybeUndefined, ProtocolVersion, v2}; +use agent_client_protocol::{AcpAgent, Agent, Client, Error, Responder, V2ConnectionTo}; +use clap::Parser; +use tokio::sync::mpsc::{UnboundedReceiver, unbounded_channel}; + +#[derive(Parser)] +#[command(name = "v2-one-shot-client")] +#[command(about = "Send one prompt to an ACP v2 agent and wait for idle")] +struct Cli { + /// Command used to start the agent. + #[arg(short, long)] + command: String, + + /// Text to send to the agent. + prompt: String, +} + +#[tokio::main] +async fn main() -> Result<(), Box> { + let cli = Cli::parse(); + let agent = AcpAgent::from_str(&cli.command)?; + let (update_tx, mut update_rx) = unbounded_channel(); + + Client + .v2() + .name("v2-one-shot-client") + .on_receive_notification( + async move |notification: v2::UpdateSessionNotification, + _connection: V2ConnectionTo| { + update_tx + .send(notification) + .map_err(Error::into_internal_error) + }, + agent_client_protocol::on_receive_notification!(), + ) + .on_receive_request( + async move |request: v2::RequestPermissionRequest, + responder: Responder, + _connection: V2ConnectionTo| { + eprintln!( + "Agent requested permission for session {}; cancelling in this non-interactive example", + request.session_id + ); + responder.respond(v2::RequestPermissionResponse::new( + v2::RequestPermissionOutcome::Cancelled, + )) + }, + agent_client_protocol::on_receive_request!(), + ) + .connect_with(agent, async move |connection| { + let initialize = connection + .send_request(v2::InitializeRequest::new( + ProtocolVersion::V2, + v2::Implementation::new("v2-one-shot-client", env!("CARGO_PKG_VERSION")), + )) + .block_task() + .await?; + if initialize.capabilities.session.is_none() { + return Err(Error::invalid_params() + .data("agent did not advertise the v2 session capability")); + } + + let opened = connection + .build_session_cwd()? + .start_session() + .block_task() + .await?; + let session = opened.into_session(); + + session.send_prompt(&cli.prompt).block_task().await?; + eprintln!("Prompt accepted; waiting for session output and completion..."); + + let (output, stop_reason) = + wait_until_idle(&mut update_rx, session.session_id()).await?; + println!("{output}"); + eprintln!("Session is idle: {stop_reason:?}"); + + session.close().block_task().await?; + Ok(()) + }) + .await?; + + Ok(()) +} + +async fn wait_until_idle( + updates: &mut UnboundedReceiver, + session_id: &v2::SessionId, +) -> Result<(String, Option), Error> { + let mut projection = AgentTextProjection::default(); + let mut observed_running = false; + loop { + let notification = updates.recv().await.ok_or_else(|| { + Error::internal_error().data("agent disconnected before the prompt ran to completion") + })?; + if ¬ification.session_id != session_id { + continue; + } + + match notification.update { + v2::SessionUpdate::StateUpdate(v2::StateUpdate::Running(_)) => { + observed_running = true; + } + v2::SessionUpdate::StateUpdate(v2::StateUpdate::Idle(idle)) if observed_running => { + return Ok((projection.text(), idle.stop_reason)); + } + update if observed_running => projection.apply(update), + _ => {} + } + } +} + +/// Minimal projection of v2 agent messages. +/// +/// Chunks append to a message, while an `agent_message` snapshot patches the +/// accumulated content for the same `messageId`: a value replaces it, null +/// clears it, and an omitted field preserves it. +#[derive(Default)] +struct AgentTextProjection { + order: Vec, + messages: HashMap>, +} + +impl AgentTextProjection { + fn apply(&mut self, update: v2::SessionUpdate) { + match update { + v2::SessionUpdate::AgentMessageChunk(chunk) => { + self.message_content(chunk.message_id).push(chunk.content); + } + v2::SessionUpdate::AgentMessage(message) => { + let content = self.message_content(message.message_id); + match message.content { + MaybeUndefined::Undefined => {} + MaybeUndefined::Null => content.clear(), + MaybeUndefined::Value(replacement) => *content = replacement, + } + } + _ => {} + } + } + + fn message_content(&mut self, message_id: v2::MessageId) -> &mut Vec { + if !self.messages.contains_key(&message_id) { + self.order.push(message_id.clone()); + } + self.messages.entry(message_id).or_default() + } + + fn text(&self) -> String { + self.order + .iter() + .filter_map(|message_id| self.messages.get(message_id)) + .flatten() + .filter_map(|content| match content { + v2::ContentBlock::Text(text) => Some(text.text.as_str()), + _ => None, + }) + .collect() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn agent_message_snapshots_patch_accumulated_chunks() { + let message_id = v2::MessageId::new("message-1"); + let mut projection = AgentTextProjection::default(); + + projection.apply(v2::SessionUpdate::AgentMessageChunk(v2::ContentChunk::new( + "hel".into(), + message_id.clone(), + ))); + projection.apply(v2::SessionUpdate::AgentMessageChunk(v2::ContentChunk::new( + "lo".into(), + message_id.clone(), + ))); + assert_eq!(projection.text(), "hello"); + + projection.apply(v2::SessionUpdate::AgentMessage( + v2::AgentMessage::new(message_id.clone()).content(vec!["replacement".into()]), + )); + assert_eq!(projection.text(), "replacement"); + + projection.apply(v2::SessionUpdate::AgentMessage(v2::AgentMessage::new( + message_id.clone(), + ))); + assert_eq!(projection.text(), "replacement"); + + projection.apply(v2::SessionUpdate::AgentMessage( + v2::AgentMessage::new(message_id.clone()).content(MaybeUndefined::Null), + )); + assert_eq!(projection.text(), ""); + + projection.apply(v2::SessionUpdate::AgentMessageChunk(v2::ContentChunk::new( + "again".into(), + message_id, + ))); + assert_eq!(projection.text(), "again"); + } + + #[tokio::test] + async fn prompt_completion_ignores_updates_until_running() { + let session_id = v2::SessionId::new("session-1"); + let (update_tx, mut update_rx) = unbounded_channel(); + let notification = |update| v2::UpdateSessionNotification::new(session_id.clone(), update); + + update_tx + .send(notification(v2::SessionUpdate::AgentMessageChunk( + v2::ContentChunk::new("stale".into(), "stale-message"), + ))) + .unwrap(); + update_tx + .send(notification(v2::SessionUpdate::StateUpdate( + v2::StateUpdate::Idle(v2::IdleStateUpdate::new()), + ))) + .unwrap(); + update_tx + .send(notification(v2::SessionUpdate::StateUpdate( + v2::StateUpdate::Running(v2::RunningStateUpdate::new()), + ))) + .unwrap(); + update_tx + .send(notification(v2::SessionUpdate::AgentMessageChunk( + v2::ContentChunk::new("current".into(), "current-message"), + ))) + .unwrap(); + update_tx + .send(notification(v2::SessionUpdate::AgentMessage( + v2::AgentMessage::new("current-message").content(vec!["current".into()]), + ))) + .unwrap(); + update_tx + .send(notification(v2::SessionUpdate::StateUpdate( + v2::StateUpdate::Idle( + v2::IdleStateUpdate::new().stop_reason(v2::StopReason::EndTurn), + ), + ))) + .unwrap(); + + let (text, stop_reason) = wait_until_idle(&mut update_rx, &session_id).await.unwrap(); + assert_eq!(text, "current"); + assert_eq!(stop_reason, Some(v2::StopReason::EndTurn)); + } +} diff --git a/src/agent-client-protocol/src/concepts/proxies.rs b/src/agent-client-protocol/src/concepts/proxies.rs index 42990e0e..b2d1d032 100644 --- a/src/agent-client-protocol/src/concepts/proxies.rs +++ b/src/agent-client-protocol/src/concepts/proxies.rs @@ -165,12 +165,15 @@ //! //! For `schema::v2::ResumeSessionRequest`, use //! `cx.resume_session_from(request)` and the resulting -//! `V2ResumeSessionBuilder` in the same shape. Resume routing and any -//! per-session MCP attachment are ready before the downstream request is -//! published, allowing replay to precede the complete response. Both setup -//! helpers forward that operation's response before spawning the callback. -//! Later updates and interactive requests remain independent traffic handled -//! by typed connection callbacks. +//! `V2ResumeSessionBuilder` in the same shape. With +//! `unstable_session_fork`, use `cx.fork_session_from(request)` and +//! `V2ForkSessionBuilder` for `ForkSessionRequest`; the returned session and +//! installed route use the new ID from `ForkSessionResponse`, not the source +//! session ID. Resume routing and any per-session MCP attachment are ready +//! before the downstream request is published, allowing replay to precede the +//! complete response. All setup helpers forward that operation's response +//! before spawning the callback. Later updates and interactive requests remain +//! independent traffic handled by typed connection callbacks. //! //! # The Conductor //! diff --git a/src/agent-client-protocol/src/concepts/sessions.rs b/src/agent-client-protocol/src/concepts/sessions.rs index ec95b422..ccc359fd 100644 --- a/src/agent-client-protocol/src/concepts/sessions.rs +++ b/src/agent-client-protocol/src/concepts/sessions.rs @@ -8,8 +8,9 @@ //! `ActiveSession`. With the `unstable_protocol_v2` feature, callbacks created //! through `Client.v2()` receive `V2ConnectionTo` and its `build_session*`, //! `V2SessionBuilder`, `resume_session*`, `V2ResumeSessionBuilder`, and -//! command-only `V2Session` APIs. The v2 resume helpers return a builder and do -//! not publish `session/resume` until `start_session` or +//! command-only `V2Session` APIs. With `unstable_session_fork`, it also exposes +//! `fork_session*` and `V2ForkSessionBuilder`. The v2 resume and fork helpers +//! return builders and do not publish their requests until `start_session` or //! `on_proxy_session_start` is called. V2 prompt responses acknowledge //! acceptance independently; receive session-wide updates and interactive //! requests through typed connection handlers. @@ -90,12 +91,14 @@ //! MCP attachment requires the `unstable_mcp_over_acp` feature. Standalone MCP //! servers remain available without it. Draft protocol v2 per-session //! attachment uses `V2SessionBuilder::with_mcp_server` for new sessions or -//! `V2ResumeSessionBuilder::with_mcp_server` for resumed sessions and -//! additionally requires `unstable_protocol_v2`. The SDK installs the routes -//! and initially polls the runners before publishing the setup request, so the -//! agent can use them during setup or resume replay. Successful attachments -//! remain active for the connection lifetime; setup failures, including an -//! error response after cancellation, clean up the pending attachment. +//! `V2ResumeSessionBuilder::with_mcp_server` for resumed sessions. With +//! `unstable_session_fork`, `V2ForkSessionBuilder::with_mcp_server` provides the +//! same attachment for forked sessions. These APIs additionally require +//! `unstable_protocol_v2`. The SDK installs the routes and initially polls the +//! runners before publishing the setup request, so the agent can use them +//! during setup or resume replay. Successful attachments remain active for the +//! connection lifetime; setup failures, including an error response after +//! cancellation, clean up the pending attachment. //! //! ```ignore //! # use agent_client_protocol::{Client, Agent, ConnectTo}; @@ -149,10 +152,11 @@ //! for details. //! //! For a draft v2 proxy, use `V2SessionBuilder::on_proxy_session_start` or -//! `V2ResumeSessionBuilder::on_proxy_session_start` instead. Each forwards the -//! complete operation-specific response and then spawns the callback with an -//! `OpenedV2Session`, so the callback keeps both the command-only session handle -//! and that exact response: +//! `V2ResumeSessionBuilder::on_proxy_session_start` instead. The feature-gated +//! `V2ForkSessionBuilder` exposes the same helper. Each forwards the complete +//! operation-specific response and then spawns the callback with an +//! `OpenedV2Session`, so the callback keeps both the command-only session +//! handle and that exact response: //! //! ```rust,ignore //! Proxy.v2() @@ -170,7 +174,10 @@ //! ); //! ``` //! -//! For `session/resume`, the builder installs and acknowledges session routing +//! For `session/new` and feature-gated `session/fork`, the builder installs +//! routing with the newly allocated response session ID before later inbound +//! traffic is dispatched. For `session/resume`, the builder installs and +//! acknowledges session routing //! before publishing the downstream request, allowing replay updates to be //! forwarded before the resume response. The downstream request inherits //! upstream cancellation. An unsuccessful downstream response drops pending diff --git a/src/agent-client-protocol/src/jsonrpc.rs b/src/agent-client-protocol/src/jsonrpc.rs index c7c46095..b7913610 100644 --- a/src/agent-client-protocol/src/jsonrpc.rs +++ b/src/agent-client-protocol/src/jsonrpc.rs @@ -1164,8 +1164,11 @@ impl< /// /// This is intended for protocol-routing infrastructure that has already /// selected v2 but still needs protocol-neutral [`ConnectionTo`] values in - /// its callbacks. Most clients should use [`Client::v2`](crate::Client::v2), - /// which also exposes the version-typed [`V2ConnectionTo`] API. + /// its callbacks. The guarded child must still send and receive the + /// `initialize` round trip; a router that consumes initialization itself + /// must use [`Builder::without_acp_version_guard`] for the selected child. + /// Most clients should use [`Client::v2`](crate::Client::v2), which also + /// exposes the version-typed [`V2ConnectionTo`] API. pub fn with_v2_protocol_guard(mut self) -> Self { self.protocol_mode = ProtocolMode::v2_client(); self @@ -1183,8 +1186,11 @@ impl< /// /// This is intended for protocol-routing infrastructure that has already /// selected v2 but still needs protocol-neutral [`ConnectionTo`] values in - /// its callbacks. Most agents should use [`Agent::v2`](crate::Agent::v2), - /// which also exposes the version-typed [`V2ConnectionTo`] API. + /// its callbacks. The guarded child must still receive and answer the + /// `initialize` request; a router that consumes initialization itself must + /// use [`Builder::without_acp_version_guard`] for the selected child. Most + /// agents should use [`Agent::v2`](crate::Agent::v2), which also exposes the + /// version-typed [`V2ConnectionTo`] API. pub fn with_v2_protocol_guard(mut self) -> Self { self.protocol_mode = ProtocolMode::v2_agent(); self diff --git a/src/agent-client-protocol/src/jsonrpc/incoming_actor.rs b/src/agent-client-protocol/src/jsonrpc/incoming_actor.rs index 55d0aec0..b65fccab 100644 --- a/src/agent-client-protocol/src/jsonrpc/incoming_actor.rs +++ b/src/agent-client-protocol/src/jsonrpc/incoming_actor.rs @@ -441,9 +441,9 @@ fn dispatch_from_message( .expect("well-formed JSON"); if let Some(id) = id { - let message = protocol_compat.incoming_message(message)?; let response_destination = response_destination.expect("incoming requests always have a response destination"); + let message = protocol_compat.incoming_request(&id, message)?; Ok(vec![Dispatch::Request( message, Responder::new( diff --git a/src/agent-client-protocol/src/jsonrpc/outgoing_actor.rs b/src/agent-client-protocol/src/jsonrpc/outgoing_actor.rs index b5114939..effe164c 100644 --- a/src/agent-client-protocol/src/jsonrpc/outgoing_actor.rs +++ b/src/agent-client-protocol/src/jsonrpc/outgoing_actor.rs @@ -70,12 +70,14 @@ pub(super) async fn outgoing_protocol_actor( %method, "Completing abandoned JSON-RPC batch request with Internal Error" ); - let fallback = RawJsonRpcMessage::response( - id, + let fallback = protocol_compat.outgoing_response_to( + &id, + &method, Err(crate::Error::internal_error().data(format!( "request handler dropped its responder for `{method}`" ))), ); + let fallback = RawJsonRpcMessage::response(id, fallback); if let Some(frame) = destination.abandon(fallback) { transport_tx .unbounded_send(frame) @@ -178,7 +180,7 @@ pub(super) async fn outgoing_protocol_actor( method, response, destination, - } => match protocol_compat.outgoing_response(&method, response) { + } => match protocol_compat.outgoing_response_to(&id, &method, response) { Ok(value) => { tracing::debug!(?id, "Sending success response"); (RawJsonRpcMessage::response(id, Ok(value)), destination) diff --git a/src/agent-client-protocol/src/jsonrpc/protocol_compat.rs b/src/agent-client-protocol/src/jsonrpc/protocol_compat.rs index 0d68b80a..536a74c7 100644 --- a/src/agent-client-protocol/src/jsonrpc/protocol_compat.rs +++ b/src/agent-client-protocol/src/jsonrpc/protocol_compat.rs @@ -1,6 +1,7 @@ #[cfg(not(feature = "unstable_protocol_v2"))] mod imp { #![allow(clippy::unused_self, clippy::unnecessary_wraps)] + use crate::schema::v1::RequestId; use crate::{UntypedMessage, role::RemoteStyle}; #[derive(Clone, Copy, Debug, Default)] @@ -43,6 +44,14 @@ mod imp { Ok(message) } + pub(crate) fn incoming_request( + &self, + _id: &RequestId, + message: UntypedMessage, + ) -> Result { + self.incoming_message(message) + } + pub(crate) fn outgoing_message( &self, message: UntypedMessage, @@ -80,6 +89,15 @@ mod imp { ) -> Result { result } + + pub(crate) fn outgoing_response_to( + &self, + _id: &RequestId, + method: &str, + result: Result, + ) -> Result { + self.outgoing_response(method, result) + } } } @@ -87,7 +105,7 @@ mod imp { mod imp { use std::sync::{Arc, Mutex}; - use crate::schema::ProtocolVersion; + use crate::schema::{ProtocolVersion, v1::RequestId}; use crate::{UntypedMessage, role::RemoteStyle}; #[derive(Clone, Copy, Debug)] @@ -100,6 +118,7 @@ mod imp { pub(crate) struct AcpProtocolMode { api: ProtocolVersionKind, initialize_surface: InitializeSurface, + initialization_role: InitializationRole, } #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -108,6 +127,17 @@ mod imp { Proxy, } + #[derive(Clone, Copy, Debug, PartialEq, Eq)] + enum InitializationRole { + /// Preserve the v1 and proxy initialization behavior. A v2 proxy has + /// both an incoming predecessor initialization and an outgoing + /// successor initialization, so the peer lifecycle does not apply to + /// it. + Unchecked, + Initiator, + Responder, + } + impl AcpProtocolMode { fn is_incoming_initialize_request(self, method: &str) -> bool { match self.initialize_surface { @@ -140,6 +170,7 @@ mod imp { Self::Acp(AcpProtocolMode { api: ProtocolVersionKind::V1, initialize_surface: InitializeSurface::Peer, + initialization_role: InitializationRole::Unchecked, }) } @@ -147,6 +178,7 @@ mod imp { Self::Acp(AcpProtocolMode { api: ProtocolVersionKind::V1, initialize_surface: InitializeSurface::Peer, + initialization_role: InitializationRole::Unchecked, }) } @@ -154,6 +186,7 @@ mod imp { Self::Acp(AcpProtocolMode { api: ProtocolVersionKind::V1, initialize_surface: InitializeSurface::Proxy, + initialization_role: InitializationRole::Unchecked, }) } @@ -161,6 +194,7 @@ mod imp { Self::Acp(AcpProtocolMode { api: ProtocolVersionKind::V2, initialize_surface: InitializeSurface::Peer, + initialization_role: InitializationRole::Responder, }) } @@ -168,6 +202,7 @@ mod imp { Self::Acp(AcpProtocolMode { api: ProtocolVersionKind::V2, initialize_surface: InitializeSurface::Peer, + initialization_role: InitializationRole::Initiator, }) } @@ -175,6 +210,7 @@ mod imp { Self::Acp(AcpProtocolMode { api: ProtocolVersionKind::V2, initialize_surface: InitializeSurface::Proxy, + initialization_role: InitializationRole::Unchecked, }) } @@ -193,6 +229,11 @@ mod imp { "cannot merge standard ACP and proxy ACP builders; \ handler chains share one initialization surface", ); + assert_eq!( + this.initialization_role, other.initialization_role, + "cannot merge ACP builders with different initialization roles; \ + handler chains share one connection lifecycle", + ); Self::Acp(this) } } @@ -216,6 +257,16 @@ mod imp { struct ProtocolState { negotiated: ProtocolVersionKind, pending_initialize: Option, + incoming_initialize_id: Option, + initialization: InitializationState, + } + + #[derive(Clone, Copy, Debug, PartialEq, Eq)] + enum InitializationState { + Unchecked, + Uninitialized, + Initializing, + Ready, } #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)] @@ -250,26 +301,43 @@ mod imp { ProtocolMode::Acp(mode) => Some(mode), }; let negotiated = mode.map_or(ProtocolVersionKind::V1, |mode| mode.api); + let initialization = match mode.map(|mode| mode.initialization_role) { + Some(InitializationRole::Initiator | InitializationRole::Responder) => { + InitializationState::Uninitialized + } + Some(InitializationRole::Unchecked) | None => InitializationState::Unchecked, + }; Self { mode, state: Arc::new(Mutex::new(ProtocolState { negotiated, pending_initialize: None, + incoming_initialize_id: None, + initialization, })), } } + #[cfg(test)] pub(crate) fn incoming_message( &self, message: UntypedMessage, + ) -> Result { + self.incoming_request(&RequestId::Null, message) + } + + pub(crate) fn incoming_request( + &self, + id: &RequestId, + message: UntypedMessage, ) -> Result { let Some(mode) = self.mode else { return Ok(message); }; if mode.is_incoming_initialize_request(message.method()) { - return self.incoming_initialize_request(mode, message); + return self.incoming_initialize_request(mode, id, message); } if mode.initialize_surface == InitializeSurface::Proxy && (message.method() == "initialize" || successor_encloses_initialize(&message)) @@ -277,6 +345,7 @@ mod imp { return Err(invalid_proxy_initialize_direction()); } + self.ensure_initialized(message.method())?; ensure_matching_protocol_version( message.method(), self.active_wire_version(), @@ -298,9 +367,11 @@ mod imp { outgoing_initialize_params(mode, remote_style, &mut message)? { set_protocol_version(params, mode.api)?; - self.set_pending_initialize(mode.api); + validate_native_v2_initialize_request(mode, params)?; + self.begin_outgoing_initialize(mode)?; mode.api } else { + self.ensure_initialized(message.method())?; self.active_wire_version() }; @@ -316,6 +387,7 @@ mod imp { return Ok(vec![message]); }; + self.ensure_notification_allowed(message.method())?; ensure_matching_protocol_version( message.method(), self.active_wire_version(), @@ -332,6 +404,7 @@ mod imp { return Ok(vec![message]); }; + self.ensure_notification_allowed(message.method())?; ensure_matching_protocol_version( message.method(), mode.api, @@ -354,38 +427,70 @@ mod imp { } let value = result?; + self.ensure_initialized(method)?; ensure_matching_protocol_version(method, self.active_wire_version(), mode.api)?; Ok(value) } + #[cfg(test)] pub(crate) fn outgoing_response( &self, method: &str, result: Result, + ) -> Result { + self.outgoing_response_to(&RequestId::Null, method, result) + } + + pub(crate) fn outgoing_response_to( + &self, + id: &RequestId, + method: &str, + result: Result, ) -> Result { let Some(mode) = self.mode else { return result; }; - // Always drain any pending initialize state so a failed initialize - // doesn't leak negotiation state to a subsequent request. - let pending_initialize = if mode.is_outgoing_initialize_response(method) { - self.take_pending_initialize() - } else { - None - }; - - let mut value = result?; + if mode.is_outgoing_initialize_response(method) { + match mode.initialization_role { + InitializationRole::Initiator => { + return result.and_then(|_| { + Err(unexpected_initialize_response(mode.initialization_role)) + }); + } + InitializationRole::Responder if !self.is_pending_incoming_initialize(id) => { + return result.and_then(|_| { + Err(unexpected_initialize_response(mode.initialization_role)) + }); + } + InitializationRole::Unchecked | InitializationRole::Responder => {} + } + let mut value = match result { + Ok(value) => value, + Err(error) => { + self.fail_initialize(); + return Err(error); + } + }; + let negotiated = self.pending_initialize().or_else(|| { + (mode.initialization_role == InitializationRole::Unchecked).then_some(mode.api) + }); + let negotiated = negotiated + .ok_or_else(|| unexpected_initialize_response(mode.initialization_role))?; + if let Err(error) = ensure_matching_protocol_version(method, mode.api, negotiated) + .and_then(|()| set_protocol_version(&mut value, negotiated)) + .and_then(|()| validate_native_v2_initialize_response(mode, &value)) + { + self.fail_initialize(); + return Err(error); + } + self.complete_initialize(negotiated)?; + return Ok(value); + } - let wire_version = if mode.is_outgoing_initialize_response(method) { - let negotiated = pending_initialize.unwrap_or(mode.api); - ensure_matching_protocol_version(method, mode.api, negotiated)?; - set_protocol_version(&mut value, negotiated)?; - self.set_negotiated(negotiated); - negotiated - } else { - self.active_wire_version() - }; + let value = result?; + self.ensure_initialized(method)?; + let wire_version = self.active_wire_version(); ensure_matching_protocol_version(method, mode.api, wire_version)?; Ok(value) @@ -394,6 +499,7 @@ mod imp { fn incoming_initialize_request( &self, mode: AcpProtocolMode, + id: &RequestId, mut message: UntypedMessage, ) -> Result { let requested = required_protocol_version_from_value(message.params())?; @@ -403,8 +509,9 @@ mod imp { return Err(unsupported_protocol_version(requested, mode.api)); } - self.set_pending_initialize(mode.api); set_protocol_version(&mut message.params, mode.api)?; + validate_native_v2_initialize_request(mode, message.params())?; + self.begin_incoming_initialize(mode, id)?; Ok(message) } @@ -413,48 +520,186 @@ mod imp { mode: AcpProtocolMode, result: Result, ) -> Result { - let _pending_initialize = self.take_pending_initialize(); - let mut value = result?; - let response_version = required_protocol_version_from_value(&value)?; - let wire_version = ProtocolVersionKind::from_protocol_version(response_version) - .ok_or_else(|| unsupported_protocol_version(response_version, mode.api))?; - if wire_version != mode.api { - return Err(required_protocol_version(mode.api, wire_version)); + let mut value = match result { + Ok(value) => value, + Err(error) => { + self.fail_initialize(); + return Err(error); + } + }; + let response = (|| { + let pending = self.pending_initialize().or_else(|| { + (mode.initialization_role == InitializationRole::Unchecked).then_some(mode.api) + }); + let pending = pending + .ok_or_else(|| unexpected_initialize_response(mode.initialization_role))?; + let response_version = required_protocol_version_from_value(&value)?; + let wire_version = ProtocolVersionKind::from_protocol_version(response_version) + .ok_or_else(|| unsupported_protocol_version(response_version, mode.api))?; + if wire_version != mode.api { + return Err(required_protocol_version(mode.api, wire_version)); + } + ensure_matching_protocol_version("initialize", pending, wire_version)?; + set_protocol_version(&mut value, wire_version)?; + validate_native_v2_initialize_response(mode, &value)?; + Ok(wire_version) + })(); + + match response { + Ok(wire_version) => { + self.complete_initialize(wire_version)?; + Ok(value) + } + Err(error) => { + self.fail_initialize(); + Err(error) + } } - self.set_negotiated(wire_version); + } - set_protocol_version(&mut value, wire_version)?; - Ok(value) + fn begin_incoming_initialize( + &self, + mode: AcpProtocolMode, + id: &RequestId, + ) -> Result<(), crate::Error> { + match mode.initialization_role { + InitializationRole::Initiator => return Err(invalid_initialize_direction()), + InitializationRole::Unchecked | InitializationRole::Responder => {} + } + self.begin_initialize(mode.api, Some(id)) } - fn active_wire_version(&self) -> ProtocolVersionKind { - let state = self + fn begin_outgoing_initialize(&self, mode: AcpProtocolMode) -> Result<(), crate::Error> { + match mode.initialization_role { + InitializationRole::Responder => return Err(invalid_initialize_direction()), + InitializationRole::Unchecked | InitializationRole::Initiator => {} + } + self.begin_initialize(mode.api, None) + } + + fn begin_initialize( + &self, + requested: ProtocolVersionKind, + incoming_id: Option<&RequestId>, + ) -> Result<(), crate::Error> { + let mut state = self .state .lock() .expect("protocol compatibility state mutex poisoned"); - state.pending_initialize.unwrap_or(state.negotiated) + match state.initialization { + InitializationState::Unchecked => { + state.pending_initialize = Some(requested); + Ok(()) + } + InitializationState::Uninitialized => { + state.initialization = InitializationState::Initializing; + state.pending_initialize = Some(requested); + state.incoming_initialize_id = incoming_id.cloned(); + Ok(()) + } + InitializationState::Initializing => Err(crate::Error::invalid_request() + .data("ACP initialization is already in progress on this connection")), + InitializationState::Ready => Err(crate::Error::invalid_request().data( + "ACP connections may only be initialized once; reconnect to initialize again", + )), + } } - fn set_negotiated(&self, negotiated: ProtocolVersionKind) { - self.state + fn ensure_initialized(&self, method: &str) -> Result<(), crate::Error> { + let state = self + .state + .lock() + .expect("protocol compatibility state mutex poisoned") + .initialization; + match state { + InitializationState::Unchecked | InitializationState::Ready => Ok(()), + InitializationState::Uninitialized | InitializationState::Initializing => { + Err(crate::Error::invalid_request().data(format!( + "ACP initialization must complete before `{method}` can be used", + ))) + } + } + } + + fn ensure_notification_allowed(&self, method: &str) -> Result<(), crate::Error> { + let initialization = self + .state .lock() .expect("protocol compatibility state mutex poisoned") - .negotiated = negotiated; + .initialization; + if initialization == InitializationState::Initializing && method == "$/cancel_request" { + return Ok(()); + } + self.ensure_initialized(method) + } + + fn complete_initialize(&self, negotiated: ProtocolVersionKind) -> Result<(), crate::Error> { + let mut state = self + .state + .lock() + .expect("protocol compatibility state mutex poisoned"); + if state.pending_initialize.is_none() + && state.initialization != InitializationState::Unchecked + { + return Err(unexpected_initialize_response( + self.mode + .expect("protocol initialization requires an ACP mode") + .initialization_role, + )); + } + if !matches!( + state.initialization, + InitializationState::Unchecked | InitializationState::Initializing + ) { + return Err(unexpected_initialize_response( + self.mode + .expect("protocol initialization requires an ACP mode") + .initialization_role, + )); + } + state.pending_initialize = None; + state.incoming_initialize_id = None; + state.negotiated = negotiated; + if state.initialization == InitializationState::Initializing { + state.initialization = InitializationState::Ready; + } + Ok(()) + } + + fn fail_initialize(&self) { + let mut state = self + .state + .lock() + .expect("protocol compatibility state mutex poisoned"); + state.pending_initialize = None; + state.incoming_initialize_id = None; + if state.initialization == InitializationState::Initializing { + state.initialization = InitializationState::Uninitialized; + } } - fn set_pending_initialize(&self, negotiated: ProtocolVersionKind) { + fn is_pending_incoming_initialize(&self, id: &RequestId) -> bool { self.state .lock() .expect("protocol compatibility state mutex poisoned") - .pending_initialize = Some(negotiated); + .incoming_initialize_id + .as_ref() + == Some(id) + } + + fn active_wire_version(&self) -> ProtocolVersionKind { + let state = self + .state + .lock() + .expect("protocol compatibility state mutex poisoned"); + state.pending_initialize.unwrap_or(state.negotiated) } - fn take_pending_initialize(&self) -> Option { + fn pending_initialize(&self) -> Option { self.state .lock() .expect("protocol compatibility state mutex poisoned") .pending_initialize - .take() } } @@ -524,6 +769,20 @@ mod imp { .data("initialize.protocolVersion must be a valid ACP protocol version") } + fn invalid_initialize_direction() -> crate::Error { + crate::Error::invalid_request() + .data("ACP clients send `initialize` requests and ACP agents respond to them") + } + + fn unexpected_initialize_response(role: InitializationRole) -> crate::Error { + let detail = match role { + InitializationRole::Initiator => "before an initialize request is pending", + InitializationRole::Responder => "before an initialize request was received", + InitializationRole::Unchecked => "without a pending initialize request", + }; + crate::Error::invalid_request().data(format!("received an initialize response {detail}")) + } + fn invalid_proxy_initialize_direction() -> crate::Error { crate::Error::invalid_request().data( "proxy initialization must arrive as `_proxy/initialize`; outgoing `initialize` must target the successor so the connection can apply `_proxy/successor`", @@ -551,6 +810,34 @@ mod imp { Ok(()) } + fn validate_native_v2_initialize_request( + mode: AcpProtocolMode, + value: &serde_json::Value, + ) -> Result<(), crate::Error> { + if mode.initialization_role == InitializationRole::Unchecked { + return Ok(()); + } + ::parse_message( + "initialize", + value, + )?; + Ok(()) + } + + fn validate_native_v2_initialize_response( + mode: AcpProtocolMode, + value: &serde_json::Value, + ) -> Result<(), crate::Error> { + if mode.initialization_role == InitializationRole::Unchecked { + return Ok(()); + } + ::from_value( + "initialize", + value.clone(), + )?; + Ok(()) + } + fn ensure_matching_protocol_version( method: &str, from: ProtocolVersionKind, @@ -609,6 +896,14 @@ mod imp { .pending_initialize } + fn initialization_state(compat: &ProtocolCompat) -> InitializationState { + compat + .state + .lock() + .expect("protocol compatibility state mutex poisoned") + .initialization + } + fn v2_implementation() -> v2::Implementation { v2::Implementation::new("protocol-compat-test", env!("CARGO_PKG_VERSION")) } @@ -626,6 +921,10 @@ mod imp { { let compat = ProtocolCompat::new(ProtocolMode::v2_agent()); assert_eq!(compat.active_wire_version(), ProtocolVersionKind::V2); + assert_eq!( + initialization_state(&compat), + InitializationState::Uninitialized + ); compat.incoming_message(UntypedMessage::new( "initialize", @@ -634,6 +933,10 @@ mod imp { assert_eq!(negotiated(&compat), ProtocolVersionKind::V2); assert_eq!(compat.active_wire_version(), ProtocolVersionKind::V2); + assert_eq!( + initialization_state(&compat), + InitializationState::Initializing + ); compat.outgoing_response( "initialize", @@ -644,6 +947,7 @@ mod imp { assert_eq!(negotiated(&compat), ProtocolVersionKind::V2); assert_eq!(compat.active_wire_version(), ProtocolVersionKind::V2); + assert_eq!(initialization_state(&compat), InitializationState::Ready); Ok(()) } @@ -652,6 +956,10 @@ mod imp { { let compat = ProtocolCompat::new(ProtocolMode::v2_client()); assert_eq!(compat.active_wire_version(), ProtocolVersionKind::V2); + assert_eq!( + initialization_state(&compat), + InitializationState::Uninitialized + ); compat.outgoing_message( UntypedMessage::new("initialize", v2_initialize_request(ProtocolVersion::V1))?, @@ -660,6 +968,10 @@ mod imp { assert_eq!(negotiated(&compat), ProtocolVersionKind::V2); assert_eq!(compat.active_wire_version(), ProtocolVersionKind::V2); + assert_eq!( + initialization_state(&compat), + InitializationState::Initializing + ); compat.incoming_response( "initialize", @@ -670,6 +982,7 @@ mod imp { assert_eq!(negotiated(&compat), ProtocolVersionKind::V2); assert_eq!(compat.active_wire_version(), ProtocolVersionKind::V2); + assert_eq!(initialization_state(&compat), InitializationState::Ready); Ok(()) } @@ -695,6 +1008,18 @@ mod imp { assert!(result.is_err()); assert_eq!(negotiated(&compat), ProtocolVersionKind::V2); assert_eq!(compat.active_wire_version(), ProtocolVersionKind::V2); + assert_eq!( + initialization_state(&compat), + InitializationState::Uninitialized + ); + compat.outgoing_message( + UntypedMessage::new("initialize", v2_initialize_request(ProtocolVersion::V2))?, + RemoteStyle::Counterpart, + )?; + assert_eq!( + initialization_state(&compat), + InitializationState::Initializing + ); Ok(()) } @@ -721,11 +1046,383 @@ mod imp { assert!(data.contains("protocolVersion"), "{error:?}"); assert_eq!(negotiated(&compat), ProtocolVersionKind::V2); assert_eq!(compat.active_wire_version(), ProtocolVersionKind::V2); + assert_eq!( + initialization_state(&compat), + InitializationState::Uninitialized + ); } Ok(()) } + #[test] + fn malformed_v2_initialize_requests_do_not_begin_a_handshake() -> Result<(), crate::Error> { + let malformed = serde_json::json!({ "protocolVersion": ProtocolVersion::V2 }); + + let client = ProtocolCompat::new(ProtocolMode::v2_client()); + client + .outgoing_message( + UntypedMessage::new("initialize", malformed.clone())?, + RemoteStyle::Counterpart, + ) + .expect_err("native v2 clients must send the complete initialize request shape"); + assert_eq!( + initialization_state(&client), + InitializationState::Uninitialized + ); + client.outgoing_message( + UntypedMessage::new("initialize", v2_initialize_request(ProtocolVersion::V2))?, + RemoteStyle::Counterpart, + )?; + assert_eq!( + initialization_state(&client), + InitializationState::Initializing + ); + + let agent = ProtocolCompat::new(ProtocolMode::v2_agent()); + agent + .incoming_message(UntypedMessage::new("initialize", malformed)?) + .expect_err("native v2 agents must receive the complete initialize request shape"); + assert_eq!( + initialization_state(&agent), + InitializationState::Uninitialized + ); + agent.incoming_message(UntypedMessage::new( + "initialize", + v2_initialize_request(ProtocolVersion::V2), + )?)?; + assert_eq!( + initialization_state(&agent), + InitializationState::Initializing + ); + Ok(()) + } + + #[test] + fn malformed_v2_initialize_success_leaves_client_uninitialized_for_retry() + -> Result<(), crate::Error> { + let compat = ProtocolCompat::new(ProtocolMode::v2_client()); + compat.outgoing_message( + UntypedMessage::new("initialize", v2_initialize_request(ProtocolVersion::V2))?, + RemoteStyle::Counterpart, + )?; + + compat + .incoming_response( + "initialize", + Ok(serde_json::json!({ "protocolVersion": ProtocolVersion::V2 })), + ) + .expect_err("initialize success must contain the complete v2 response shape"); + assert_eq!( + initialization_state(&compat), + InitializationState::Uninitialized + ); + assert_eq!(pending_initialize(&compat), None); + + compat.outgoing_message( + UntypedMessage::new("initialize", v2_initialize_request(ProtocolVersion::V2))?, + RemoteStyle::Counterpart, + )?; + assert_eq!( + initialization_state(&compat), + InitializationState::Initializing + ); + Ok(()) + } + + #[test] + fn malformed_v2_initialize_success_leaves_agent_uninitialized_for_retry() + -> Result<(), crate::Error> { + let compat = ProtocolCompat::new(ProtocolMode::v2_agent()); + compat.incoming_message(UntypedMessage::new( + "initialize", + v2_initialize_request(ProtocolVersion::V2), + )?)?; + + compat + .outgoing_response( + "initialize", + Ok(serde_json::json!({ "protocolVersion": ProtocolVersion::V2 })), + ) + .expect_err("initialize success must contain the complete v2 response shape"); + assert_eq!( + initialization_state(&compat), + InitializationState::Uninitialized + ); + assert_eq!(pending_initialize(&compat), None); + + compat.incoming_message(UntypedMessage::new( + "initialize", + v2_initialize_request(ProtocolVersion::V2), + )?)?; + assert_eq!( + initialization_state(&compat), + InitializationState::Initializing + ); + Ok(()) + } + + #[test] + fn v2_peer_traffic_requires_completed_initialization() -> Result<(), crate::Error> { + for compat in [ + ProtocolCompat::new(ProtocolMode::v2_agent()), + ProtocolCompat::new(ProtocolMode::v2_client()), + ] { + for error in [ + compat + .incoming_message(UntypedMessage::new( + "session/new", + serde_json::json!({}), + )?) + .expect_err("incoming requests must wait for initialization"), + compat + .outgoing_message( + UntypedMessage::new("session/new", serde_json::json!({}))?, + RemoteStyle::Counterpart, + ) + .expect_err("outgoing requests must wait for initialization"), + compat + .incoming_notification(UntypedMessage::new( + "session/update", + serde_json::json!({}), + )?) + .expect_err("incoming notifications must wait for initialization"), + compat + .outgoing_notification(UntypedMessage::new( + "session/update", + serde_json::json!({}), + )?) + .expect_err("outgoing notifications must wait for initialization"), + compat + .incoming_response("session/new", Ok(serde_json::json!({}))) + .expect_err("incoming responses must wait for initialization"), + compat + .outgoing_response("session/new", Ok(serde_json::json!({}))) + .expect_err("outgoing responses must wait for initialization"), + ] { + let data = error + .data + .as_ref() + .and_then(|data| data.as_str()) + .unwrap_or_default(); + assert!(data.contains("initialization must complete"), "{error:?}"); + } + } + Ok(()) + } + + #[test] + fn v2_initialization_allows_only_protocol_cancellation_notifications() + -> Result<(), crate::Error> { + let client = ProtocolCompat::new(ProtocolMode::v2_client()); + client.outgoing_message( + UntypedMessage::new("initialize", v2_initialize_request(ProtocolVersion::V2))?, + RemoteStyle::Counterpart, + )?; + client.outgoing_notification(UntypedMessage::new( + "$/cancel_request", + serde_json::json!({ "requestId": 1 }), + )?)?; + client + .outgoing_notification(UntypedMessage::new( + "session/update", + serde_json::json!({}), + )?) + .expect_err("ordinary notifications must wait for initialization"); + + let agent = ProtocolCompat::new(ProtocolMode::v2_agent()); + agent.incoming_message(UntypedMessage::new( + "initialize", + v2_initialize_request(ProtocolVersion::V2), + )?)?; + agent.incoming_notification(UntypedMessage::new( + "$/cancel_request", + serde_json::json!({ "requestId": 1 }), + )?)?; + agent + .incoming_notification(UntypedMessage::new( + "session/update", + serde_json::json!({}), + )?) + .expect_err("ordinary notifications must wait for initialization"); + Ok(()) + } + + #[test] + fn v2_peer_initialization_has_one_direction_and_one_successful_round_trip() + -> Result<(), crate::Error> { + let agent = ProtocolCompat::new(ProtocolMode::v2_agent()); + let error = agent + .outgoing_message( + UntypedMessage::new("initialize", v2_initialize_request(ProtocolVersion::V2))?, + RemoteStyle::Counterpart, + ) + .expect_err("agents must not initiate initialization"); + assert!( + error + .data + .as_ref() + .and_then(|data| data.as_str()) + .is_some_and(|data| data.contains("clients send `initialize`")), + "{error:?}" + ); + assert_eq!( + initialization_state(&agent), + InitializationState::Uninitialized + ); + + agent.incoming_message(UntypedMessage::new( + "initialize", + v2_initialize_request(ProtocolVersion::V2), + )?)?; + agent.outgoing_response( + "initialize", + Ok(serde_json::to_value(v2_initialize_response( + ProtocolVersion::V2, + ))?), + )?; + let error = agent + .incoming_message(UntypedMessage::new( + "initialize", + v2_initialize_request(ProtocolVersion::V2), + )?) + .expect_err("agents must reject reinitialization"); + assert!( + error + .data + .as_ref() + .and_then(|data| data.as_str()) + .is_some_and(|data| data.contains("only be initialized once")), + "{error:?}" + ); + assert_eq!(initialization_state(&agent), InitializationState::Ready); + + let client = ProtocolCompat::new(ProtocolMode::v2_client()); + let error = client + .incoming_message(UntypedMessage::new( + "initialize", + v2_initialize_request(ProtocolVersion::V2), + )?) + .expect_err("clients must not receive initialization requests"); + assert!( + error + .data + .as_ref() + .and_then(|data| data.as_str()) + .is_some_and(|data| data.contains("clients send `initialize`")), + "{error:?}" + ); + assert_eq!( + initialization_state(&client), + InitializationState::Uninitialized + ); + + client.outgoing_message( + UntypedMessage::new("initialize", v2_initialize_request(ProtocolVersion::V2))?, + RemoteStyle::Counterpart, + )?; + client.incoming_response( + "initialize", + Ok(serde_json::to_value(v2_initialize_response( + ProtocolVersion::V2, + ))?), + )?; + let error = client + .outgoing_message( + UntypedMessage::new("initialize", v2_initialize_request(ProtocolVersion::V2))?, + RemoteStyle::Counterpart, + ) + .expect_err("clients must reject reinitialization"); + assert!( + error + .data + .as_ref() + .and_then(|data| data.as_str()) + .is_some_and(|data| data.contains("only be initialized once")), + "{error:?}" + ); + assert_eq!(initialization_state(&client), InitializationState::Ready); + Ok(()) + } + + #[test] + fn rejected_concurrent_initialize_does_not_clear_the_active_handshake() + -> Result<(), crate::Error> { + let compat = ProtocolCompat::new(ProtocolMode::v2_agent()); + let accepted_id = RequestId::Number(1); + let rejected_id = RequestId::Number(2); + + compat.incoming_request( + &accepted_id, + UntypedMessage::new("initialize", v2_initialize_request(ProtocolVersion::V2))?, + )?; + let duplicate_error = compat + .incoming_request( + &rejected_id, + UntypedMessage::new("initialize", v2_initialize_request(ProtocolVersion::V2))?, + ) + .expect_err("a second initialize request must be rejected while one is active"); + compat + .outgoing_response_to(&rejected_id, "initialize", Err(duplicate_error)) + .expect_err("the rejected initialize receives its own error response"); + + assert_eq!( + initialization_state(&compat), + InitializationState::Initializing + ); + assert_eq!(pending_initialize(&compat), Some(ProtocolVersionKind::V2)); + + compat.outgoing_response_to( + &accepted_id, + "initialize", + Ok(serde_json::to_value(v2_initialize_response( + ProtocolVersion::V2, + ))?), + )?; + assert_eq!(initialization_state(&compat), InitializationState::Ready); + Ok(()) + } + + #[test] + fn rejected_wrong_direction_initialize_does_not_clear_client_handshake() + -> Result<(), crate::Error> { + let compat = ProtocolCompat::new(ProtocolMode::v2_client()); + compat.outgoing_message( + UntypedMessage::new("initialize", v2_initialize_request(ProtocolVersion::V2))?, + RemoteStyle::Counterpart, + )?; + + let wrong_direction_id = RequestId::Number(2); + let wrong_direction_error = compat + .incoming_request( + &wrong_direction_id, + UntypedMessage::new("initialize", v2_initialize_request(ProtocolVersion::V2))?, + ) + .expect_err("agents must not send initialize requests to clients"); + compat + .outgoing_response_to( + &wrong_direction_id, + "initialize", + Err(wrong_direction_error), + ) + .expect_err("the wrong-direction initialize receives its own error response"); + + assert_eq!( + initialization_state(&compat), + InitializationState::Initializing + ); + assert_eq!(pending_initialize(&compat), Some(ProtocolVersionKind::V2)); + + compat.incoming_response( + "initialize", + Ok(serde_json::to_value(v2_initialize_response( + ProtocolVersion::V2, + ))?), + )?; + assert_eq!(initialization_state(&compat), InitializationState::Ready); + Ok(()) + } + #[test] fn incoming_initialize_request_rejects_unsupported_protocol_version() -> Result<(), crate::Error> { diff --git a/src/agent-client-protocol/src/lib.rs b/src/agent-client-protocol/src/lib.rs index 409ee556..4f71dbde 100644 --- a/src/agent-client-protocol/src/lib.rs +++ b/src/agent-client-protocol/src/lib.rs @@ -24,9 +24,11 @@ //! inbound traffic are independent. With both v2 and MCP-over-ACP features, //! `Proxy.v2()` supports global MCP attachment, while `V2SessionBuilder` and //! `V2ResumeSessionBuilder` support per-session attachment plus non-blocking -//! proxy setup. Per-session MCP routes and runners are ready before a setup -//! request is published, as is proxy session routing for resume replay; -//! successful attachments remain active for the connection lifetime. +//! proxy setup. With `unstable_session_fork`, `V2ForkSessionBuilder` provides +//! the same shape for forked sessions and uses the response's new session ID. +//! Per-session MCP routes and runners are ready before a setup request is +//! published, as is proxy session routing for resume replay; successful +//! attachments remain active for the connection lifetime. //! //! Here's a minimal example that initializes a v1 connection, creates a //! session, and sends a prompt: diff --git a/src/agent-client-protocol/src/mcp_server/mod.rs b/src/agent-client-protocol/src/mcp_server/mod.rs index 6eea588c..d33605ea 100644 --- a/src/agent-client-protocol/src/mcp_server/mod.rs +++ b/src/agent-client-protocol/src/mcp_server/mod.rs @@ -8,8 +8,10 @@ //! attachment and per-session attachment. V2 uses //! `Proxy.v2().with_mcp_server(...)` or //! `V2SessionBuilder::with_mcp_server(...)` for new sessions and -//! `V2ResumeSessionBuilder::with_mcp_server(...)` for resumed sessions, and -//! additionally requires `unstable_protocol_v2`. +//! `V2ResumeSessionBuilder::with_mcp_server(...)` for resumed sessions. With +//! `unstable_session_fork`, `V2ForkSessionBuilder::with_mcp_server(...)` +//! attaches a server to a forked session. V2 attachment additionally requires +//! `unstable_protocol_v2`. //! //! ## Building MCP servers with tools //! diff --git a/src/agent-client-protocol/src/mcp_server/server.rs b/src/agent-client-protocol/src/mcp_server/server.rs index b0f97d23..cd12d10f 100644 --- a/src/agent-client-protocol/src/mcp_server/server.rs +++ b/src/agent-client-protocol/src/mcp_server/server.rs @@ -46,7 +46,9 @@ use crate::schema::v1::ForkSessionRequest; /// enabled, `Proxy.v2().with_mcp_server` attaches a server globally to draft /// v2 setup requests. `V2SessionBuilder::with_mcp_server` attaches one to a /// single new v2 session, while `V2ResumeSessionBuilder::with_mcp_server` -/// attaches one to a single resumed session. +/// attaches one to a single resumed session. With `unstable_session_fork`, +/// `V2ForkSessionBuilder::with_mcp_server` attaches one to a single forked +/// session. /// /// # Creating an MCP Server /// diff --git a/src/agent-client-protocol/src/session/v2.rs b/src/agent-client-protocol/src/session/v2.rs index 82816ead..dea34aa5 100644 --- a/src/agent-client-protocol/src/session/v2.rs +++ b/src/agent-client-protocol/src/session/v2.rs @@ -160,6 +160,33 @@ where V2SessionBuilder::new(self, request) } + /// Build an unstable draft protocol v2 `session/fork` request. + /// + /// This helper is available with the `unstable_session_fork` feature. Call + /// [`V2ForkSessionBuilder::start_session`] to publish the request and + /// obtain a command handle for the newly created fork. + #[cfg(feature = "unstable_session_fork")] + pub fn fork_session( + &self, + session_id: impl Into, + cwd: impl AsRef, + ) -> V2ForkSessionBuilder { + self.fork_session_from(v2::ForkSessionRequest::new(session_id, cwd.as_ref())) + } + + /// Build an unstable draft protocol v2 `session/fork` request from an + /// existing request. + /// + /// This helper is available with the `unstable_session_fork` feature. Call + /// [`V2ForkSessionBuilder::start_session`] to publish the request. + #[cfg(feature = "unstable_session_fork")] + pub fn fork_session_from( + &self, + request: v2::ForkSessionRequest, + ) -> V2ForkSessionBuilder { + V2ForkSessionBuilder::new(self, request) + } + /// Build a draft protocol v2 `session/resume` request. /// /// Use [`Self::resume_session_from`] to request history replay or set @@ -349,6 +376,176 @@ where } } +/// Builder for an unstable draft protocol v2 `session/fork` request. +/// +/// A successful fork creates a new independent session whose ID comes from the +/// [`v2::ForkSessionResponse`]. Register typed +/// [`v2::UpdateSessionNotification`] and interactive request handlers on +/// [`crate::Builder`] before connecting, then use [`Self::start_session`] to +/// obtain a command handle for the fork or `on_proxy_session_start` to forward +/// setup through a proxy. +/// +/// This type is available with the `unstable_session_fork` feature. With +/// `unstable_mcp_over_acp` as well, `with_mcp_server` attaches an MCP server to +/// the forked session. +#[cfg(feature = "unstable_session_fork")] +#[must_use = "call `start_session` or `on_proxy_session_start` to send the `session/fork` request"] +#[derive(Debug)] +pub struct V2ForkSessionBuilder +where + Counterpart: HasPeer, + Run: RunWithConnectionTo, +{ + connection: V2ConnectionTo, + request: v2::ForkSessionRequest, + dynamic_handler_registrations: Vec>, + run: Run, +} + +#[cfg(feature = "unstable_session_fork")] +impl V2ForkSessionBuilder +where + Counterpart: HasPeer, +{ + fn new(connection: &V2ConnectionTo, request: v2::ForkSessionRequest) -> Self { + Self { + connection: connection.clone(), + request, + dynamic_handler_registrations: Vec::new(), + run: NullRun, + } + } +} + +#[cfg(feature = "unstable_session_fork")] +impl V2ForkSessionBuilder +where + Counterpart: HasPeer, + Run: RunWithConnectionTo, +{ + /// Attach an MCP server to this forked protocol v2 session. + /// + /// This method is available when `unstable_mcp_over_acp` is enabled in + /// addition to `unstable_protocol_v2` and `unstable_session_fork`. MCP + /// routes are installed and their runner tasks receive an initial poll + /// before `session/fork` is published, allowing the agent to connect while + /// handling session setup. A successful attachment remains active for the + /// lifetime of the connection. + #[cfg(feature = "unstable_mcp_over_acp")] + pub fn with_mcp_server( + mut self, + mcp_server: McpServer, + ) -> Result>, crate::Error> + where + McpRun: RunWithConnectionTo, + { + let (handler, mcp_run) = mcp_server.into_v2_handler_and_runner(); + self.dynamic_handler_registrations + .push(handler.into_dynamic_handler(&mut self.request.mcp_servers, &self.connection)?); + Ok(V2ForkSessionBuilder { + connection: self.connection, + request: self.request, + dynamic_handler_registrations: self.dynamic_handler_registrations, + run: ChainRun::new(self.run, mcp_run), + }) + } + + fn send_fork_session(self, ordered: bool) -> SentRequest + where + Run: 'static, + { + let Self { + connection, + request, + dynamic_handler_registrations, + run, + } = self; + send_session_setup( + connection, + request, + dynamic_handler_registrations, + run, + ordered, + ) + } + + /// Send `session/fork` and return its independently consumable request. + /// + /// The successful result contains both a cloneable command handle for the + /// newly created fork and the complete [`v2::ForkSessionResponse`]. Consume + /// the returned request with [`SentRequest::block_task`], + /// [`SentRequest::on_receiving_result`], or another explicit [`SentRequest`] + /// completion mode. + /// + /// Attached MCP routes are installed and their runner tasks begin + /// executing before the request is published. A valid success response + /// promotes them to the connection lifetime independently from how this + /// request handle is consumed; setup errors clean up the pending + /// attachment. + pub fn start_session(self) -> SentRequest> + where + Run: 'static, + { + let session_connection = self.connection.clone(); + self.send_fork_session(false).map(move |response| { + let session = V2Session { + session_id: response.session_id.clone(), + connection: session_connection, + }; + Ok(OpenedV2Session { session, response }) + }) + } + + /// Fork a protocol v2 session through a proxy and forward its response. + /// + /// The downstream request is ordered and inherits cancellation from the + /// upstream request. On success, this helper obtains the new session ID + /// from the response, installs session routing before later inbound traffic + /// is processed, forwards the complete response, and spawns `op` with an + /// [`OpenedV2Session`] containing the fork's command handle and response. + /// Inbound updates and interactive requests remain independent connection + /// traffic. + /// + /// The callback runs outside the ordered response barrier, so it may wait + /// for later connection traffic without deadlocking the dispatch loop. + pub fn on_proxy_session_start( + self, + responder: Responder, + op: F, + ) -> Result<(), crate::Error> + where + Counterpart: HasPeer, + Run: 'static, + F: FnOnce(OpenedV2Session) -> Fut + Send + 'static, + Fut: Future> + Send, + { + let session_connection = self.connection.clone(); + self.send_fork_session(true) + .forward_cancellation_from(responder.cancellation()) + .on_receiving_ok_result(responder, async move |response, responder| { + let session_id = response.session_id.clone(); + let raw_connection = session_connection.raw_connection(); + let route = match raw_connection.add_dynamic_handler(ProxySessionMessages::new( + crate::schema::v1::SessionId::from(session_id.clone()), + )) { + Ok(route) => route, + Err(error) => return responder.respond_with_error(error), + }; + + let opened = OpenedV2Session { + session: V2Session { + session_id, + connection: session_connection.clone(), + }, + response: response.clone(), + }; + responder.respond(response)?; + route.detach(); + raw_connection.spawn(async move { op(opened).await }) + }) + } +} + /// Builder for a draft protocol v2 `session/resume` request. /// /// Replay updates arrive before the resume response. Direct clients must diff --git a/src/agent-client-protocol/tests/protocol_v2.rs b/src/agent-client-protocol/tests/protocol_v2.rs index 3dc93dfb..b4d0b939 100644 --- a/src/agent-client-protocol/tests/protocol_v2.rs +++ b/src/agent-client-protocol/tests/protocol_v2.rs @@ -12,9 +12,9 @@ use agent_client_protocol::schema::{ProtocolVersion, SuccessorMessage, v1, v2}; use agent_client_protocol::{ Agent, AgentProtocolRouter, Builder, ByteStreams, Client, ClientProtocolConnector, Conductor, ConnectTo, ConnectionContext, ConnectionTo, DynamicHandlerGuard, Error, HandleConnectionClose, - HandleDispatchFrom, JsonRpcMessage, JsonRpcRequest, JsonRpcResponse, NullHandler, Proxy, - RawJsonRpcMessage, Role, RunWithConnectionTo, TransportFrame, UntypedMessage, UntypedRole, - V2Builder, V2ConnectionTo, + HandleDispatchFrom, JsonRpcMessage, JsonRpcNotification, JsonRpcRequest, JsonRpcResponse, + NullHandler, Proxy, RawJsonRpcMessage, Role, RunWithConnectionTo, TransportFrame, + UntypedMessage, UntypedRole, V2Builder, V2ConnectionTo, }; use agent_client_protocol_test::MockTransport; use agent_client_protocol_test::testy::Testy; @@ -36,6 +36,13 @@ struct ForeignInitializeResponse { protocol_version: String, } +#[derive(Debug, Clone, Serialize, Deserialize, JsonRpcRequest)] +#[request(method = "initialize", response = Value)] +struct RawInitializeRequest { + #[serde(flatten)] + params: Map, +} + struct ForeignPeer; impl ConnectTo for ForeignPeer { @@ -669,87 +676,364 @@ fn v2_extension_enum_parsing_preserves_method_prefix() -> Result<(), Error> { Ok(()) } +fn assert_v2_client_request_mapping( + method: &str, + request: Req, + response: Value, + request_variant: impl FnOnce(v2::ClientRequest) -> bool, + response_variant: impl FnOnce(v2::AgentResponse) -> bool, +) -> Result<(), Error> +where + Req: JsonRpcRequest + Serialize, + Req::Response: JsonRpcResponse, +{ + let params = json_value(request)?; + let request = Req::parse_message(method, ¶ms)?; + assert_eq!(request.method(), method); + assert_eq!(request.to_untyped_message()?.method(), method); + let request = v2::ClientRequest::parse_message(method, ¶ms)?; + assert_eq!(request.method(), method); + assert_eq!(request.to_untyped_message()?.method(), method); + assert!(request_variant(request)); + + ::from_value(method, response.clone())?; + assert!(response_variant(v2::AgentResponse::from_value( + method, response + )?)); + Ok(()) +} + +fn assert_v2_agent_request_mapping( + method: &str, + request: Req, + response: Value, + request_variant: impl FnOnce(v2::AgentRequest) -> bool, + response_variant: impl FnOnce(v2::ClientResponse) -> bool, +) -> Result<(), Error> +where + Req: JsonRpcRequest + Serialize, + Req::Response: JsonRpcResponse, +{ + let params = json_value(request)?; + let request = Req::parse_message(method, ¶ms)?; + assert_eq!(request.method(), method); + assert_eq!(request.to_untyped_message()?.method(), method); + let request = v2::AgentRequest::parse_message(method, ¶ms)?; + assert_eq!(request.method(), method); + assert_eq!(request.to_untyped_message()?.method(), method); + assert!(request_variant(request)); + + ::from_value(method, response.clone())?; + assert!(response_variant(v2::ClientResponse::from_value( + method, response + )?)); + Ok(()) +} + +fn assert_v2_client_notification_mapping( + method: &str, + notification: Notif, + notification_variant: impl FnOnce(v2::ClientNotification) -> bool, +) -> Result<(), Error> +where + Notif: JsonRpcNotification + Serialize, +{ + let params = json_value(notification)?; + let notification = Notif::parse_message(method, ¶ms)?; + assert_eq!(notification.method(), method); + assert_eq!(notification.to_untyped_message()?.method(), method); + let notification = v2::ClientNotification::parse_message(method, ¶ms)?; + assert_eq!(notification.method(), method); + assert_eq!(notification.to_untyped_message()?.method(), method); + assert!(notification_variant(notification)); + Ok(()) +} + +fn assert_v2_agent_notification_mapping( + method: &str, + notification: Notif, + notification_variant: impl FnOnce(v2::AgentNotification) -> bool, +) -> Result<(), Error> +where + Notif: JsonRpcNotification + Serialize, +{ + let params = json_value(notification)?; + let notification = Notif::parse_message(method, ¶ms)?; + assert_eq!(notification.method(), method); + assert_eq!(notification.to_untyped_message()?.method(), method); + let notification = v2::AgentNotification::parse_message(method, ¶ms)?; + assert_eq!(notification.method(), method); + assert_eq!(notification.to_untyped_message()?.method(), method); + assert!(notification_variant(notification)); + Ok(()) +} + #[test] -fn v2_schema_1_4_method_names_are_jsonrpc_mapped() -> Result<(), Error> { - fn assert_request() {} - fn assert_notification() {} - - assert_request::(); - assert_request::(); - assert_notification::(); - assert_notification::(); - assert_notification::(); - - let login_params = serde_json::json!({ "methodId": "browser" }); - let login = v2::LoginAuthRequest::parse_message("auth/login", &login_params)?; - assert_eq!(login.method(), "auth/login"); - let client_request = v2::ClientRequest::parse_message("auth/login", &login_params)?; - assert!(matches!( - client_request, - v2::ClientRequest::LoginAuthRequest(_) - )); - let login_response = v2::AgentResponse::from_value("auth/login", serde_json::json!({}))?; - assert!(matches!( - login_response, - v2::AgentResponse::LoginAuthResponse(_) - )); +fn sdk_supported_v2_method_surface_is_jsonrpc_mapped() -> Result<(), Error> { + macro_rules! assert_client_request { + ($request:ident, $response:ident, $method:literal, $request_value:expr, $response_value:expr) => { + assert_v2_client_request_mapping::( + $method, + $request_value, + json_value($response_value)?, + |request| matches!(request, v2::ClientRequest::$request(_)), + |response| matches!(response, v2::AgentResponse::$response(_)), + )?; + }; + } - let logout = v2::LogoutAuthRequest::parse_message("auth/logout", &serde_json::json!({}))?; - assert_eq!(logout.method(), "auth/logout"); - let client_request = v2::ClientRequest::parse_message("auth/logout", &serde_json::json!({}))?; - assert!(matches!( - client_request, - v2::ClientRequest::LogoutAuthRequest(_) - )); - let logout_response = v2::AgentResponse::from_value("auth/logout", serde_json::json!({}))?; - assert!(matches!( - logout_response, - v2::AgentResponse::LogoutAuthResponse(_) - )); + macro_rules! assert_agent_request { + ($request:ident, $response:ident, $method:literal, $request_value:expr, $response_value:expr) => { + assert_v2_agent_request_mapping::( + $method, + $request_value, + json_value($response_value)?, + |request| matches!(request, v2::AgentRequest::$request(_)), + |response| matches!(response, v2::ClientResponse::$response(_)), + )?; + }; + } - let cancel_params = serde_json::json!({ "requestId": "req-1" }); + assert_client_request!( + InitializeRequest, + InitializeResponse, + "initialize", + v2_initialize_request(ProtocolVersion::V2), + v2::InitializeResponse::new(ProtocolVersion::V2, v2_implementation()) + ); + assert_client_request!( + LoginAuthRequest, + LoginAuthResponse, + "auth/login", + v2::LoginAuthRequest::new("browser"), + v2::LoginAuthResponse::new() + ); + assert_client_request!( + LogoutAuthRequest, + LogoutAuthResponse, + "auth/logout", + v2::LogoutAuthRequest::new(), + v2::LogoutAuthResponse::new() + ); + assert_client_request!( + NewSessionRequest, + NewSessionResponse, + "session/new", + v2::NewSessionRequest::new(cwd()?), + v2::NewSessionResponse::new("new-session") + ); + assert_client_request!( + ListSessionsRequest, + ListSessionsResponse, + "session/list", + v2::ListSessionsRequest::new(), + v2::ListSessionsResponse::new(Vec::new()) + ); + assert_client_request!( + DeleteSessionRequest, + DeleteSessionResponse, + "session/delete", + v2::DeleteSessionRequest::new("session-1"), + v2::DeleteSessionResponse::new() + ); + assert_client_request!( + ResumeSessionRequest, + ResumeSessionResponse, + "session/resume", + v2::ResumeSessionRequest::new("session-1", cwd()?), + v2::ResumeSessionResponse::new() + ); + assert_client_request!( + CloseSessionRequest, + CloseSessionResponse, + "session/close", + v2::CloseSessionRequest::new("session-1"), + v2::CloseSessionResponse::new() + ); + assert_client_request!( + SetSessionConfigOptionRequest, + SetSessionConfigOptionResponse, + "session/set_config_option", + v2::SetSessionConfigOptionRequest::new("session-1", "model", "model-1"), + v2::SetSessionConfigOptionResponse::new(Vec::new()) + ); + assert_client_request!( + PromptRequest, + PromptResponse, + "session/prompt", + v2::PromptRequest::new("session-1", Vec::new()), + v2::PromptResponse::new() + ); + + #[cfg(feature = "unstable_session_fork")] + assert_client_request!( + ForkSessionRequest, + ForkSessionResponse, + "session/fork", + v2::ForkSessionRequest::new("session-1", cwd()?), + v2::ForkSessionResponse::new("forked-session") + ); + + #[cfg(feature = "unstable_llm_providers")] + { + assert_client_request!( + ListProvidersRequest, + ListProvidersResponse, + "providers/list", + v2::ListProvidersRequest::new(), + v2::ListProvidersResponse::new(Vec::new()) + ); + assert_client_request!( + SetProviderRequest, + SetProviderResponse, + "providers/set", + v2::SetProviderRequest::new( + "provider-1", + v2::LlmProtocol::OpenAi, + "https://example.com" + ), + v2::SetProviderResponse::new() + ); + assert_client_request!( + DisableProviderRequest, + DisableProviderResponse, + "providers/disable", + v2::DisableProviderRequest::new("provider-1"), + v2::DisableProviderResponse::new() + ); + } + + assert_v2_client_notification_mapping( + "session/cancel", + v2::CancelSessionNotification::new("session-1"), + |notification| { + matches!( + notification, + v2::ClientNotification::CancelSessionNotification(_) + ) + }, + )?; + + assert_agent_request!( + RequestPermissionRequest, + RequestPermissionResponse, + "session/request_permission", + v2::RequestPermissionRequest::new("session-1", "Run command?", Vec::new()), + v2::RequestPermissionResponse::new(v2::RequestPermissionOutcome::Cancelled) + ); + + let update = v2::UpdateSessionNotification::new( + "session-1", + v2::SessionUpdate::StateUpdate(v2::StateUpdate::Running(v2::RunningStateUpdate::new())), + ); + assert_v2_agent_notification_mapping("session/update", update, |notification| { + matches!( + notification, + v2::AgentNotification::UpdateSessionNotification(_) + ) + })?; + + #[cfg(feature = "unstable_elicitation")] + { + assert_agent_request!( + CreateElicitationRequest, + CreateElicitationResponse, + "elicitation/create", + v2::CreateElicitationRequest::new( + v2::ElicitationFormMode::new( + v2::ElicitationSessionScope::new("session-1"), + v2::ElicitationSchema::new(), + ), + "Choose a value", + ), + v2::CreateElicitationResponse::new(v2::ElicitationAction::Decline) + ); + assert_v2_agent_notification_mapping( + "elicitation/complete", + v2::CompleteElicitationNotification::new("elicitation-1"), + |notification| { + matches!( + notification, + v2::AgentNotification::CompleteElicitationNotification(_) + ) + }, + )?; + } + + #[cfg(feature = "unstable_mcp_over_acp")] + { + fn message_response() -> Result { + serde_json::from_value(serde_json::json!({ "tools": [] })) + .map_err(Error::into_internal_error) + } + + assert_client_request!( + MessageMcpRequest, + MessageMcpResponse, + "mcp/message", + v2::MessageMcpRequest::new("connection-1", "tools/list"), + message_response()? + ); + assert_v2_client_notification_mapping( + "mcp/message", + v2::MessageMcpNotification::new("connection-1", "notifications/tools/list"), + |notification| { + matches!( + notification, + v2::ClientNotification::MessageMcpNotification(_) + ) + }, + )?; + + assert_agent_request!( + ConnectMcpRequest, + ConnectMcpResponse, + "mcp/connect", + v2::ConnectMcpRequest::new("server-1"), + v2::ConnectMcpResponse::new("connection-1") + ); + assert_agent_request!( + MessageMcpRequest, + MessageMcpResponse, + "mcp/message", + v2::MessageMcpRequest::new("connection-1", "tools/list"), + message_response()? + ); + assert_agent_request!( + DisconnectMcpRequest, + DisconnectMcpResponse, + "mcp/disconnect", + v2::DisconnectMcpRequest::new("connection-1"), + v2::DisconnectMcpResponse::new() + ); + assert_v2_agent_notification_mapping( + "mcp/message", + v2::MessageMcpNotification::new("connection-1", "notifications/tools/list"), + |notification| { + matches!( + notification, + v2::AgentNotification::MessageMcpNotification(_) + ) + }, + )?; + } + + let cancel_params = json_value(v2::CancelRequestNotification::new(String::from( + "request-1", + )))?; let cancel = v2::CancelRequestNotification::parse_message("$/cancel_request", &cancel_params)?; assert_eq!(cancel.method(), "$/cancel_request"); - let protocol_notification = - v2::ProtocolLevelNotification::parse_message("$/cancel_request", &cancel_params)?; assert!(matches!( - protocol_notification, + v2::ProtocolLevelNotification::parse_message("$/cancel_request", &cancel_params)?, v2::ProtocolLevelNotification::CancelRequestNotification(_) )); - let session_cancel_params = serde_json::json!({ "sessionId": "session-1" }); - let session_cancel = - v2::CancelSessionNotification::parse_message("session/cancel", &session_cancel_params)?; - assert_eq!(session_cancel.method(), "session/cancel"); - let client_notification = - v2::ClientNotification::parse_message("session/cancel", &session_cancel_params)?; - assert!(matches!( - client_notification, - v2::ClientNotification::CancelSessionNotification(_) - )); - - let update_params = serde_json::json!({ - "sessionId": "session-1", - "update": { "sessionUpdate": "_custom" } - }); - let update = v2::UpdateSessionNotification::parse_message("session/update", &update_params)?; - assert_eq!(update.method(), "session/update"); - let agent_notification = - v2::AgentNotification::parse_message("session/update", &update_params)?; - assert!(matches!( - agent_notification, - v2::AgentNotification::UpdateSessionNotification(_) - )); - Ok(()) } #[cfg(feature = "unstable_mcp_over_acp")] #[test] -fn mcp_over_acp_variants_are_jsonrpc_mapped() -> Result<(), Error> { - fn assert_request() {} - fn assert_notification() {} - +fn mcp_over_acp_v1_variants_are_jsonrpc_mapped() -> Result<(), Error> { macro_rules! assert_message_mapping { ($ty:ty, $method:literal, $params:expr, $pattern:pat) => {{ let message = <$ty as JsonRpcMessage>::parse_message($method, &$params)?; @@ -766,11 +1050,6 @@ fn mcp_over_acp_variants_are_jsonrpc_mapped() -> Result<(), Error> { }}; } - assert_request::(); - assert_request::(); - assert_request::(); - assert_notification::(); - assert_message_mapping!( v1::ClientRequest, "mcp/message", @@ -838,101 +1117,6 @@ fn mcp_over_acp_variants_are_jsonrpc_mapped() -> Result<(), Error> { v1::AgentNotification::MessageMcpNotification(_) ); - assert_message_mapping!( - v2::MessageMcpRequest, - "mcp/message", - json_value(v2::MessageMcpRequest::new("conn-1", "tools/list"))?, - v2::MessageMcpRequest { .. } - ); - assert_message_mapping!( - v2::MessageMcpNotification, - "mcp/message", - json_value(v2::MessageMcpNotification::new( - "conn-1", - "notifications/tools/list" - ))?, - v2::MessageMcpNotification { .. } - ); - assert_message_mapping!( - v2::ConnectMcpRequest, - "mcp/connect", - json_value(v2::ConnectMcpRequest::new("server-1"))?, - v2::ConnectMcpRequest { .. } - ); - assert_message_mapping!( - v2::DisconnectMcpRequest, - "mcp/disconnect", - json_value(v2::DisconnectMcpRequest::new("conn-1"))?, - v2::DisconnectMcpRequest { .. } - ); - - assert_message_mapping!( - v2::ClientRequest, - "mcp/message", - json_value(v2::MessageMcpRequest::new("conn-1", "tools/list"))?, - v2::ClientRequest::MessageMcpRequest(_) - ); - assert_response_mapping!( - v2::AgentResponse, - "mcp/message", - serde_json::json!({ "tools": [] }), - v2::AgentResponse::MessageMcpResponse(_) - ); - assert_message_mapping!( - v2::ClientNotification, - "mcp/message", - json_value(v2::MessageMcpNotification::new( - "conn-1", - "notifications/tools/list" - ))?, - v2::ClientNotification::MessageMcpNotification(_) - ); - assert_message_mapping!( - v2::AgentRequest, - "mcp/connect", - json_value(v2::ConnectMcpRequest::new("server-1"))?, - v2::AgentRequest::ConnectMcpRequest(_) - ); - assert_message_mapping!( - v2::AgentRequest, - "mcp/message", - json_value(v2::MessageMcpRequest::new("conn-1", "tools/list"))?, - v2::AgentRequest::MessageMcpRequest(_) - ); - assert_message_mapping!( - v2::AgentRequest, - "mcp/disconnect", - json_value(v2::DisconnectMcpRequest::new("conn-1"))?, - v2::AgentRequest::DisconnectMcpRequest(_) - ); - assert_response_mapping!( - v2::ClientResponse, - "mcp/connect", - json_value(v2::ConnectMcpResponse::new("conn-1"))?, - v2::ClientResponse::ConnectMcpResponse(_) - ); - assert_response_mapping!( - v2::ClientResponse, - "mcp/message", - serde_json::json!({ "tools": [] }), - v2::ClientResponse::MessageMcpResponse(_) - ); - assert_response_mapping!( - v2::ClientResponse, - "mcp/disconnect", - serde_json::json!({}), - v2::ClientResponse::DisconnectMcpResponse(_) - ); - assert_message_mapping!( - v2::AgentNotification, - "mcp/message", - json_value(v2::MessageMcpNotification::new( - "conn-1", - "notifications/tools/list" - ))?, - v2::AgentNotification::MessageMcpNotification(_) - ); - Ok(()) } @@ -1002,6 +1186,316 @@ async fn v2_client_and_agent_negotiate_v2() -> Result<(), Error> { .await } +#[tokio::test(flavor = "current_thread")] +async fn v2_client_does_not_send_session_requests_before_initialization() -> Result<(), Error> { + let handler_ran = Arc::new(AtomicBool::new(false)); + let handler_flag = Arc::clone(&handler_ran); + let agent = Agent + .builder() + .without_acp_version_guard() + .on_receive_request( + async move |_request: v2::NewSessionRequest, responder, _cx| { + handler_flag.store(true, Ordering::SeqCst); + responder.respond(v2::NewSessionResponse::new(v2::SessionId::new( + "unexpected-session", + ))) + }, + agent_client_protocol::on_receive_request!(), + ); + + Client + .v2() + .connect_with(agent, async |cx| { + let error = cx + .send_request(v2::NewSessionRequest::new(cwd()?)) + .block_task() + .await + .expect_err("v2 clients must initialize before session requests"); + let data = error + .data + .as_ref() + .and_then(|data| data.as_str()) + .unwrap_or_default(); + assert!(data.contains("initialization must complete"), "{error:?}"); + Ok(()) + }) + .await?; + + assert!(!handler_ran.load(Ordering::SeqCst)); + Ok(()) +} + +#[tokio::test(flavor = "current_thread")] +async fn v2_agent_rejects_session_requests_before_initialization() -> Result<(), Error> { + let handler_ran = Arc::new(AtomicBool::new(false)); + let handler_flag = Arc::clone(&handler_ran); + let agent = Agent.v2().on_receive_request( + async move |_request: v2::NewSessionRequest, responder, _cx| { + handler_flag.store(true, Ordering::SeqCst); + responder.respond(v2::NewSessionResponse::new(v2::SessionId::new( + "unexpected-session", + ))) + }, + agent_client_protocol::on_receive_request!(), + ); + + Client + .builder() + .without_acp_version_guard() + .connect_with(agent, async |cx| { + let error = cx + .send_request(v2::NewSessionRequest::new(cwd()?)) + .block_task() + .await + .expect_err("v2 agents must reject session requests before initialization"); + let data = error + .data + .as_ref() + .and_then(|data| data.as_str()) + .unwrap_or_default(); + assert!(data.contains("initialization must complete"), "{error:?}"); + Ok(()) + }) + .await?; + + assert!(!handler_ran.load(Ordering::SeqCst)); + Ok(()) +} + +#[tokio::test(flavor = "current_thread")] +async fn v2_agent_rejects_reinitialization_without_losing_ready_state() -> Result<(), Error> { + let initialize_count = Arc::new(AtomicUsize::new(0)); + let initialize_counter = Arc::clone(&initialize_count); + let agent = Agent + .v2() + .on_receive_request( + async move |initialize: v2::InitializeRequest, responder, _cx| { + initialize_counter.fetch_add(1, Ordering::SeqCst); + responder.respond(v2_initialize_response_with_session( + initialize.protocol_version, + )) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async |_request: v2::NewSessionRequest, responder, _cx| { + responder.respond(v2::NewSessionResponse::new(v2::SessionId::new( + "ready-session", + ))) + }, + agent_client_protocol::on_receive_request!(), + ); + + Client + .builder() + .without_acp_version_guard() + .connect_with(agent, async |cx| { + cx.send_request(v2_initialize_request(ProtocolVersion::V2)) + .block_task() + .await?; + + let error = cx + .send_request(v2_initialize_request(ProtocolVersion::V2)) + .block_task() + .await + .expect_err("v2 agents must reject reinitialization"); + let data = error + .data + .as_ref() + .and_then(|data| data.as_str()) + .unwrap_or_default(); + assert!(data.contains("only be initialized once"), "{error:?}"); + + let session = cx + .send_request(v2::NewSessionRequest::new(cwd()?)) + .block_task() + .await?; + assert_eq!(session.session_id.0.as_ref(), "ready-session"); + Ok(()) + }) + .await?; + + assert_eq!(initialize_count.load(Ordering::SeqCst), 1); + Ok(()) +} + +#[tokio::test(flavor = "current_thread")] +async fn v2_agent_can_retry_after_batched_initialize_responder_is_dropped() -> Result<(), Error> { + use tokio::io::{AsyncWriteExt as _, BufReader}; + + let initialize_count = Arc::new(AtomicUsize::new(0)); + let initialize_counter = Arc::clone(&initialize_count); + let agent = Agent.v2().on_receive_request( + async move |initialize: v2::InitializeRequest, responder, _cx| { + if initialize_counter.fetch_add(1, Ordering::SeqCst) == 0 { + drop(responder); + Ok(()) + } else { + responder.respond(v2_initialize_response_with_session( + initialize.protocol_version, + )) + } + }, + agent_client_protocol::on_receive_request!(), + ); + + let (mut client_writer, server_reader) = tokio::io::duplex(4096); + let (server_writer, client_reader) = tokio::io::duplex(4096); + let server_transport = ByteStreams::new(server_writer.compat_write(), server_reader.compat()); + let agent_task = tokio::spawn(agent.connect_to(server_transport)); + let mut client_reader = BufReader::new(client_reader); + + write_wire_json( + &mut client_writer, + &serde_json::json!([{ + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": json_value(v2_initialize_request(ProtocolVersion::V2))?, + }]), + ) + .await?; + + let abandoned = read_wire_json(&mut client_reader).await?; + let abandoned = abandoned + .as_array() + .and_then(|responses| responses.first()) + .ok_or_else(|| Error::internal_error().data("expected initialize error batch"))?; + assert_eq!(abandoned["id"], 1); + assert_eq!(abandoned["error"]["code"], -32603); + assert!( + abandoned["error"]["data"] + .as_str() + .is_some_and(|data| data.contains("dropped its responder")), + "{abandoned:?}" + ); + + write_wire_json( + &mut client_writer, + &serde_json::json!({ + "jsonrpc": "2.0", + "id": 2, + "method": "initialize", + "params": json_value(v2_initialize_request(ProtocolVersion::V2))?, + }), + ) + .await?; + + let retry = read_wire_json(&mut client_reader).await?; + assert_eq!(retry["id"], 2); + assert_eq!(retry["result"]["protocolVersion"], 2); + assert_eq!(initialize_count.load(Ordering::SeqCst), 2); + + client_writer + .shutdown() + .await + .map_err(Error::into_internal_error)?; + agent_task + .await + .map_err(agent_client_protocol::util::internal_error)??; + Ok(()) +} + +#[tokio::test(flavor = "current_thread")] +async fn v2_client_can_retry_after_malformed_initialize_success() -> Result<(), Error> { + let initialize_count = Arc::new(AtomicUsize::new(0)); + let initialize_counter = Arc::clone(&initialize_count); + let agent = Agent + .builder() + .without_acp_version_guard() + .on_receive_request( + async move |_initialize: RawInitializeRequest, responder, _cx| { + if initialize_counter.fetch_add(1, Ordering::SeqCst) == 0 { + responder.respond(serde_json::json!({ + "protocolVersion": ProtocolVersion::V2, + })) + } else { + responder.respond(json_value(v2_initialize_response_with_session( + ProtocolVersion::V2, + ))?) + } + }, + agent_client_protocol::on_receive_request!(), + ); + + Client + .v2() + .connect_with(agent, async |cx| { + let error = cx + .send_request(v2_initialize_request(ProtocolVersion::V2)) + .block_task() + .await + .expect_err("a malformed v2 initialize success must fail typed decoding"); + let data = error + .data + .as_ref() + .map(Value::to_string) + .unwrap_or_default(); + assert!(data.contains("info"), "{error:?}"); + + let retry = cx + .send_request(v2_initialize_request(ProtocolVersion::V2)) + .block_task() + .await?; + assert_eq!(retry.protocol_version, ProtocolVersion::V2); + assert!(retry.capabilities.session.is_some()); + Ok(()) + }) + .await?; + + assert_eq!(initialize_count.load(Ordering::SeqCst), 2); + Ok(()) +} + +#[tokio::test(flavor = "current_thread")] +async fn v2_agent_can_retry_after_malformed_initialize_success() -> Result<(), Error> { + let initialize_count = Arc::new(AtomicUsize::new(0)); + let initialize_counter = Arc::clone(&initialize_count); + let agent = Agent.v2().on_receive_request( + async move |_initialize: RawInitializeRequest, responder, _cx| { + if initialize_counter.fetch_add(1, Ordering::SeqCst) == 0 { + responder.respond(serde_json::json!({ + "protocolVersion": ProtocolVersion::V2, + })) + } else { + responder.respond(json_value(v2_initialize_response_with_session( + ProtocolVersion::V2, + ))?) + } + }, + agent_client_protocol::on_receive_request!(), + ); + + Client + .builder() + .without_acp_version_guard() + .connect_with(agent, async |cx| { + let error = cx + .send_request(v2_initialize_request(ProtocolVersion::V2)) + .block_task() + .await + .expect_err("a malformed v2 initialize success must become a wire error"); + let data = error + .data + .as_ref() + .map(Value::to_string) + .unwrap_or_default(); + assert!(data.contains("info"), "{error:?}"); + + let retry = cx + .send_request(v2_initialize_request(ProtocolVersion::V2)) + .block_task() + .await?; + assert_eq!(retry.protocol_version, ProtocolVersion::V2); + assert!(retry.capabilities.session.is_some()); + Ok(()) + }) + .await?; + + assert_eq!(initialize_count.load(Ordering::SeqCst), 2); + Ok(()) +} + #[tokio::test(flavor = "current_thread")] async fn client_protocol_connector_routes_to_v2_client_for_v2_agent() -> Result<(), Error> { Client diff --git a/src/agent-client-protocol/tests/session_ordering.rs b/src/agent-client-protocol/tests/session_ordering.rs index b62fcda8..53805f92 100644 --- a/src/agent-client-protocol/tests/session_ordering.rs +++ b/src/agent-client-protocol/tests/session_ordering.rs @@ -8,10 +8,10 @@ use agent_client_protocol::{ PromptResponse, SessionId, SessionNotification, SessionUpdate, StopReason, TextContent, }, }; -use futures::{ - StreamExt as _, - channel::{mpsc, oneshot}, -}; +use futures::{StreamExt as _, channel::oneshot}; + +#[cfg(feature = "unstable_protocol_v2")] +use futures::channel::mpsc; #[cfg(feature = "unstable_protocol_v2")] use agent_client_protocol::{ @@ -103,6 +103,15 @@ mod callback_future_lifetimes { |_opened| LifetimeTaggedFuture(PhantomData) } + #[cfg(all(feature = "unstable_protocol_v2", feature = "unstable_session_fork"))] + fn v2_proxy_fork_callback<'a>() -> impl FnOnce( + agent_client_protocol::OpenedV2Session, + ) -> LifetimeTaggedFuture<'a> + + Send + + 'static { + |_opened| LifetimeTaggedFuture(PhantomData) + } + fn on_session_start_accepts_non_static_callback_future<'a>( connection: &ConnectionTo, _scope: &'a str, @@ -146,6 +155,18 @@ mod callback_future_lifetimes { .resume_session_from(request) .on_proxy_session_start(responder, v2_proxy_resume_callback::<'a>()) } + + #[cfg(all(feature = "unstable_protocol_v2", feature = "unstable_session_fork"))] + fn v2_on_proxy_fork_accepts_non_static_callback_future<'a>( + connection: &V2ConnectionTo, + request: v2::ForkSessionRequest, + responder: Responder, + _scope: &'a str, + ) -> Result<(), agent_client_protocol::Error> { + connection + .fork_session_from(request) + .on_proxy_session_start(responder, v2_proxy_fork_callback::<'a>()) + } } #[tokio::test(flavor = "current_thread")] @@ -424,6 +445,162 @@ async fn v2_proxy_session_start_installs_routing_before_later_batch_entry() { .expect("v2 proxy session connection failed"); } +#[cfg(all(feature = "unstable_protocol_v2", feature = "unstable_session_fork"))] +#[tokio::test(flavor = "current_thread")] +async fn v2_proxy_fork_installs_response_id_routing_before_later_batch_entry() { + let source_session_id = v2::SessionId::new("same-batch-v2-source"); + let forked_session_id = v2::SessionId::new("same-batch-v2-fork"); + let setup_response = + v2::ForkSessionResponse::new(forked_session_id.clone()).config_options(vec![ + v2::SessionConfigOption::boolean("thinking", "Thinking", true), + ]); + let callback_response = setup_response.clone(); + let callback_session_id = forked_session_id.clone(); + let notification_session_id = forked_session_id.clone(); + let expected_source_session_id = source_session_id.clone(); + let (transport, mut peer) = Channel::duplex(); + let (callback_tx, mut callback_rx) = mpsc::unbounded(); + let (peer_done_tx, peer_done_rx) = oneshot::channel(); + + let proxy = Proxy + .v2() + .on_receive_request_from( + Client, + async |request: v2::InitializeProxyRequest, responder, _connection| { + responder.respond(v2::InitializeResponse::new( + request.initialize.protocol_version, + v2::Implementation::new("same-batch-fork-proxy", env!("CARGO_PKG_VERSION")), + )) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request_from( + Client, + async move |request: v2::ForkSessionRequest, + responder, + connection: V2ConnectionTo| { + let callback_response = callback_response.clone(); + let callback_session_id = callback_session_id.clone(); + let callback_tx = callback_tx.clone(); + connection + .fork_session_from(request) + .on_proxy_session_start(responder, move |opened| async move { + assert_eq!(opened.session().session_id(), &callback_session_id); + assert_eq!(opened.response(), &callback_response); + callback_tx.unbounded_send(()).map_err(|_| { + agent_client_protocol::Error::internal_error() + .data("v2 fork callback receiver was dropped") + }) + }) + }, + agent_client_protocol::on_receive_request!(), + ) + .connect_with(transport, async move |_connection| { + callback_rx.next().await.ok_or_else(|| { + agent_client_protocol::Error::internal_error().data("v2 fork callback did not run") + })?; + peer_done_rx.await.map_err(|_| { + agent_client_protocol::Error::internal_error().data("raw peer stopped early") + }) + }); + + let peer = async move { + initialize_raw_v2_proxy(&mut peer, "same-batch-fork-client").await?; + + let upstream_id = agent_client_protocol::schema::v1::RequestId::Number(2); + peer.tx + .unbounded_send(TransportFrame::Single(RawJsonRpcMessage::request( + "session/fork".to_owned(), + serde_json::to_value(v2::ForkSessionRequest::new( + source_session_id, + "/same-batch-v2-fork", + )) + .expect("fork request should serialize"), + upstream_id.clone(), + )?)) + .expect("proxy should accept session/fork"); + + let Some(TransportFrame::Single(RawJsonRpcMessage::Request(forwarded))) = + peer.rx.next().await + else { + panic!("expected a forwarded session/fork request"); + }; + let successor = SuccessorMessage::::parse_message( + forwarded.method.as_ref(), + &forwarded.params, + )?; + assert_eq!(successor.message.session_id, expected_source_session_id); + assert_eq!( + successor.message.cwd, + v2::AbsolutePath::new("/same-batch-v2-fork") + ); + + let response = RawJsonRpcMessage::response( + forwarded.id, + Ok(serde_json::to_value(setup_response).expect("fork response should serialize")), + ); + let update = SuccessorMessage { + message: v2::UpdateSessionNotification::new( + notification_session_id, + v2::SessionUpdate::StateUpdate(v2::StateUpdate::Running( + v2::RunningStateUpdate::new(), + )), + ), + meta: None, + } + .to_untyped_message()?; + let (method, params) = update.into_parts(); + let notification = RawJsonRpcMessage::notification(method, params)?; + let batch = TransportBatch::from_messages([response, notification]) + .expect("test response batch should be non-empty"); + peer.tx + .unbounded_send(TransportFrame::Batch(batch)) + .expect("proxy should accept the response batch"); + + let mut saw_response = false; + let mut saw_update = false; + for _ in 0..2 { + let Some(TransportFrame::Single(message)) = peer.rx.next().await else { + panic!("expected a forwarded fork response and update"); + }; + match message { + RawJsonRpcMessage::Response( + agent_client_protocol::schema::v1::Response::Result { id, result }, + ) => { + assert_eq!(id, upstream_id); + let response = v2::ForkSessionResponse::from_value("session/fork", result)?; + assert_eq!(response.session_id, forked_session_id); + assert_eq!(response.config_options.len(), 1); + saw_response = true; + } + RawJsonRpcMessage::Notification(notification) => { + let update = v2::UpdateSessionNotification::parse_message( + notification.method.as_ref(), + ¬ification.params, + )?; + assert_eq!(update.session_id, forked_session_id); + assert!(matches!( + update.update, + v2::SessionUpdate::StateUpdate(v2::StateUpdate::Running(_)) + )); + saw_update = true; + } + message => panic!("unexpected proxy output: {message:?}"), + } + } + assert!(saw_response); + assert!(saw_update); + peer_done_tx + .send(()) + .map_err(|()| agent_client_protocol::Error::internal_error()) + }; + + tokio::time::timeout(TIMEOUT, async { futures::try_join!(proxy, peer) }) + .await + .expect("same-batch v2 fork update was not routed") + .expect("v2 proxy fork connection failed"); +} + #[cfg(feature = "unstable_protocol_v2")] #[tokio::test(flavor = "current_thread")] async fn v2_proxy_resume_forwards_replay_before_same_batch_response() { diff --git a/src/agent-client-protocol/tests/session_v2.rs b/src/agent-client-protocol/tests/session_v2.rs index 0dad7aa3..87943dad 100644 --- a/src/agent-client-protocol/tests/session_v2.rs +++ b/src/agent-client-protocol/tests/session_v2.rs @@ -214,6 +214,80 @@ async fn v2_prompt_acceptance_is_independent_from_session_updates() { .expect("v2 session connection failed"); } +#[cfg(feature = "unstable_session_fork")] +#[tokio::test(flavor = "current_thread")] +async fn v2_fork_builder_uses_the_forked_response_id_and_preserves_the_response() { + let source_session_id = v2::SessionId::new("source-session"); + let forked_session_id = v2::SessionId::new("forked-session"); + let fork_cwd = cwd().expect("test cwd should be available"); + let expected_cwd = v2::AbsolutePath::new(fork_cwd.clone()); + let config_option = v2::SessionConfigOption::boolean("thinking", "Thinking", true); + let response_meta = serde_json::Map::from_iter([( + "extension".to_owned(), + serde_json::json!({"preserved": true}), + )]); + let expected_response = v2::ForkSessionResponse::new(forked_session_id.clone()) + .config_options(vec![config_option]) + .meta(response_meta); + let agent_response = expected_response.clone(); + let agent_source_session_id = source_session_id.clone(); + + let agent = Agent + .v2() + .on_receive_request( + async |request: v2::InitializeRequest, + responder: Responder, + _connection: V2ConnectionTo| { + responder.respond( + v2::InitializeResponse::new(request.protocol_version, implementation()) + .capabilities(v2::AgentCapabilities::new().session( + v2::SessionCapabilities::new().fork(v2::SessionForkCapabilities::new()), + )), + ) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async move |request: v2::ForkSessionRequest, + responder: Responder, + _connection: V2ConnectionTo| { + assert_eq!(request.session_id, agent_source_session_id); + assert_eq!(request.cwd, expected_cwd); + responder.respond(agent_response.clone()) + }, + 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 + .fork_session(source_session_id.clone(), fork_cwd) + .start_session() + .block_task() + .await?; + assert_eq!(opened.session().session_id(), &forked_session_id); + assert_ne!(opened.session().session_id(), &source_session_id); + assert_eq!(opened.response(), &expected_response); + + let (session, response) = opened.into_parts(); + assert_eq!(session.session_id(), &forked_session_id); + assert_eq!(response, expected_response); + Ok(()) + }); + + tokio::time::timeout(TIMEOUT, client) + .await + .expect("v2 fork builder test timed out") + .expect("v2 fork builder test failed"); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn v2_session_cancellation_completes_at_cancelled_idle() { let agent = Agent diff --git a/src/agent-client-protocol/tests/session_v2_mcp.rs b/src/agent-client-protocol/tests/session_v2_mcp.rs index 8ee17180..b1c20083 100644 --- a/src/agent-client-protocol/tests/session_v2_mcp.rs +++ b/src/agent-client-protocol/tests/session_v2_mcp.rs @@ -38,12 +38,13 @@ fn implementation() -> v2::Implementation { } fn initialize_response(protocol_version: ProtocolVersion) -> v2::InitializeResponse { - v2::InitializeResponse::new(protocol_version, implementation()).capabilities( - v2::AgentCapabilities::new().session( - v2::SessionCapabilities::new() - .mcp(v2::McpCapabilities::new().acp(v2::McpAcpCapabilities::new())), - ), - ) + let session_capabilities = v2::SessionCapabilities::new() + .mcp(v2::McpCapabilities::new().acp(v2::McpAcpCapabilities::new())); + #[cfg(feature = "unstable_session_fork")] + let session_capabilities = session_capabilities.fork(v2::SessionForkCapabilities::new()); + + v2::InitializeResponse::new(protocol_version, implementation()) + .capabilities(v2::AgentCapabilities::new().session(session_capabilities)) } fn object(value: Value) -> Map { @@ -436,6 +437,197 @@ async fn v2_session_mcp_attachment_is_ready_during_setup_and_lives_for_connectio .expect("v2 MCP attachment test timed out") } +#[cfg(feature = "unstable_session_fork")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn v2_fork_mcp_attachment_preserves_request_and_lives_for_connection() -> Result<(), Error> { + let (server_id_tx, mut server_id_rx) = mpsc::unbounded(); + let (round_trip_trigger_tx, mut round_trip_trigger_rx) = mpsc::unbounded(); + let (round_trip_tx, mut round_trip_rx) = mpsc::unbounded(); + let first_round_trip_tx = round_trip_tx.clone(); + let source_session_id = v2::SessionId::new("v2-source-mcp-session"); + let expected_source_session_id = source_session_id.clone(); + let forked_session_id = v2::SessionId::new("v2-forked-mcp-session"); + let expected_forked_session_id = forked_session_id.clone(); + let fork_cwd = cwd()?; + let expected_fork_cwd = v2::AbsolutePath::new(fork_cwd.clone()); + let additional_directory = fork_cwd.join("additional"); + let expected_additional_directory = v2::AbsolutePath::new(additional_directory.clone()); + let existing_mcp_server = v2::McpServer::Other(v2::OtherMcpServer::new( + "_test_transport", + BTreeMap::from([("extension".to_owned(), json!({"preserved": true}))]), + )); + let expected_existing_mcp_server = existing_mcp_server.clone(); + let setup_meta = Map::from_iter([("setup".to_owned(), json!({"preserved": true}))]); + let expected_setup_meta = setup_meta.clone(); + let expected_response = v2::ForkSessionResponse::new(forked_session_id.clone()) + .config_options(vec![v2::SessionConfigOption::boolean( + "thinking", "Thinking", true, + )]) + .meta(Map::from_iter([( + "response".to_owned(), + json!({"preserved": true}), + )])); + let agent_response = expected_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::ForkSessionRequest, + responder: Responder, + connection: V2ConnectionTo| { + assert_eq!(request.session_id, expected_source_session_id); + assert_eq!(request.cwd, expected_fork_cwd); + assert_eq!( + request.additional_directories, + vec![expected_additional_directory.clone()] + ); + assert_eq!(request.meta.as_ref(), Some(&expected_setup_meta)); + let server = match request.mcp_servers.as_slice() { + [existing, v2::McpServer::Acp(server)] + if existing == &expected_existing_mcp_server => + { + server + } + servers => { + panic!("expected the existing declaration followed by ACP, got {servers:?}") + } + }; + assert_eq!(server.name, "v2-echo"); + let server_id = server.server_id.clone(); + server_id_tx + .unbounded_send(server_id.clone()) + .map_err(Error::into_internal_error)?; + let round_trip_connection = connection.clone(); + let first_round_trip_tx = first_round_trip_tx.clone(); + let agent_response = agent_response.clone(); + connection.spawn(async move { + match run_mcp_round_trip(&round_trip_connection, &server_id, 1).await { + Ok(round_trip) => { + first_round_trip_tx + .unbounded_send(Ok(round_trip)) + .map_err(Error::into_internal_error)?; + responder.respond(agent_response) + } + Err(error) => { + first_round_trip_tx + .unbounded_send(Err(error.clone())) + .map_err(Error::into_internal_error)?; + responder.respond_with_error(error) + } + } + }) + }, + agent_client_protocol::on_receive_request!(), + ) + .with_spawned(move |connection: V2ConnectionTo| async move { + let server_id = server_id_rx.next().await.ok_or_else(|| { + Error::internal_error().data("session/fork did not advertise an MCP server") + })?; + let mut sequence = 1; + + while round_trip_trigger_rx.next().await.is_some() { + sequence += 1; + let result = run_mcp_round_trip(&connection, &server_id, sequence).await; + let failed = result.is_err(); + round_trip_tx + .unbounded_send(result) + .map_err(Error::into_internal_error)?; + if failed { + break; + } + } + Ok(()) + }); + + let test = async move { + let (context_tx, mut context_rx) = mpsc::unbounded(); + let (notice_tx, mut notice_rx) = mpsc::unbounded(); + let (connector_dropped_tx, connector_dropped_rx) = oneshot::channel(); + let (runner_started_tx, runner_started_rx) = oneshot::channel(); + let (runner_dropped_tx, runner_dropped_rx) = oneshot::channel(); + let runner_started = Arc::new(AtomicBool::new(false)); + + Client + .v2() + .connect_with(agent, async move |connection| { + connection + .send_request(v2::InitializeRequest::new( + ProtocolVersion::V2, + implementation(), + )) + .block_task() + .await?; + + let mcp_server = McpServer::::new( + EchoMcpConnect { + context_tx, + notice_tx, + runner_started: runner_started.clone(), + dropped_tx: Mutex::new(Some(connector_dropped_tx)), + }, + ProbeRunner { + started: runner_started.clone(), + started_tx: Some(runner_started_tx), + dropped_tx: Some(runner_dropped_tx), + }, + ); + let pending_session = connection + .fork_session_from( + v2::ForkSessionRequest::new(source_session_id.clone(), fork_cwd) + .additional_directories([additional_directory]) + .mcp_servers(vec![existing_mcp_server]) + .meta(setup_meta), + ) + .with_mcp_server(mcp_server)? + .start_session(); + + runner_started_rx + .await + .map_err(Error::into_internal_error)?; + assert!( + runner_started.load(Ordering::Acquire), + "the MCP runner must be first-polled before session/fork is published" + ); + + assert_round_trip(1, &mut round_trip_rx, &mut context_rx, &mut notice_rx).await?; + + let opened = pending_session.block_task().await?; + assert_eq!(opened.session().session_id(), &expected_forked_session_id); + assert_ne!(opened.session().session_id(), &source_session_id); + assert_eq!(opened.response(), &expected_response); + let session = opened.into_session(); + let remaining_session = session.clone(); + drop(session); + drop(remaining_session); + + round_trip_trigger_tx + .unbounded_send(()) + .map_err(Error::into_internal_error)?; + assert_round_trip(2, &mut round_trip_rx, &mut context_rx, &mut notice_rx).await?; + + Ok(()) + }) + .await?; + + connector_dropped_rx + .await + .map_err(Error::into_internal_error)?; + runner_dropped_rx.await.map_err(Error::into_internal_error) + }; + + tokio::time::timeout(TIMEOUT, test) + .await + .expect("v2 fork MCP attachment test timed out") +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn v2_resume_mcp_attachment_preserves_request_and_lives_for_connection() -> Result<(), Error> { From 7b79654fdd046c1da791f972debbf2f5d4895719 Mon Sep 17 00:00:00 2001 From: Ben Brandt Date: Thu, 20 Aug 2026 12:23:29 -0700 Subject: [PATCH 2/3] fix test --- src/agent-client-protocol/examples/v2_one_shot_client.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/agent-client-protocol/examples/v2_one_shot_client.rs b/src/agent-client-protocol/examples/v2_one_shot_client.rs index cc10009a..51926681 100644 --- a/src/agent-client-protocol/examples/v2_one_shot_client.rs +++ b/src/agent-client-protocol/examples/v2_one_shot_client.rs @@ -185,14 +185,14 @@ mod tests { let mut projection = AgentTextProjection::default(); projection.apply(v2::SessionUpdate::AgentMessageChunk(v2::ContentChunk::new( - "hel".into(), + "hello ".into(), message_id.clone(), ))); projection.apply(v2::SessionUpdate::AgentMessageChunk(v2::ContentChunk::new( - "lo".into(), + "world".into(), message_id.clone(), ))); - assert_eq!(projection.text(), "hello"); + assert_eq!(projection.text(), "hello world"); projection.apply(v2::SessionUpdate::AgentMessage( v2::AgentMessage::new(message_id.clone()).content(vec!["replacement".into()]), From fa5685ae557e5cbe3ce40c5e30b8ccaea63a5ef4 Mon Sep 17 00:00:00 2001 From: Ben Brandt Date: Thu, 20 Aug 2026 12:54:49 -0700 Subject: [PATCH 3/3] clippy --- .../src/conductor.rs | 4 +- .../src/debug_logger.rs | 2 +- src/agent-client-protocol-http/src/client.rs | 13 +++-- .../src/websocket_server.rs | 2 +- .../examples/with_mcp_server.rs | 1 + src/agent-client-protocol-rmcp/src/builder.rs | 8 +-- .../src/bin/mcp_echo_server.rs | 1 + src/agent-client-protocol-test/src/lib.rs | 11 +++- src/agent-client-protocol/src/jsonrpc.rs | 6 +-- .../src/jsonrpc/close.rs | 6 +-- .../src/jsonrpc/handlers.rs | 9 ++-- src/agent-client-protocol/src/jsonrpc/run.rs | 6 +-- src/agent-client-protocol/src/role.rs | 8 +-- src/agent-client-protocol/src/role/acp.rs | 18 +++---- src/agent-client-protocol/src/role/mcp.rs | 18 ++++--- .../tests/jsonrpc_advanced.rs | 46 +++++++++------- .../tests/jsonrpc_batch.rs | 17 +++--- .../tests/jsonrpc_request_cancellation.rs | 54 ++++++++++--------- .../tests/jsonrpc_transport_close.rs | 10 ++-- .../tests/protocol_v2.rs | 10 +++- .../tests/session_v2_mcp.rs | 10 +++- 21 files changed, 154 insertions(+), 106 deletions(-) diff --git a/src/agent-client-protocol-conductor/src/conductor.rs b/src/agent-client-protocol-conductor/src/conductor.rs index 67f55575..c15cfdc7 100644 --- a/src/agent-client-protocol-conductor/src/conductor.rs +++ b/src/agent-client-protocol-conductor/src/conductor.rs @@ -908,7 +908,7 @@ where // The feature-off implementation awaits typed dispatch matchers; the v2 // implementation is intentionally raw and completes synchronously. - #[allow(clippy::unused_async)] + #[allow(unknown_lints, clippy::unused_async, clippy::unused_async_trait_impl)] async fn forward_message_from_client_to_proxy( &mut self, target_component_index: usize, @@ -988,7 +988,7 @@ where /// running as a proxy). // The feature-off implementation awaits typed dispatch matchers; the v2 // implementation is intentionally raw and completes synchronously. - #[allow(clippy::unused_async)] + #[allow(unknown_lints, clippy::unused_async, clippy::unused_async_trait_impl)] async fn forward_message_to_agent( &mut self, _client_connection: ConnectionTo, diff --git a/src/agent-client-protocol-conductor/src/debug_logger.rs b/src/agent-client-protocol-conductor/src/debug_logger.rs index e465ff9b..108b7b27 100644 --- a/src/agent-client-protocol-conductor/src/debug_logger.rs +++ b/src/agent-client-protocol-conductor/src/debug_logger.rs @@ -159,7 +159,7 @@ impl Write for DebugLogWriter { fn flush(&mut self) -> std::io::Result<()> { if !self.buffer.is_empty() { - let line = self.buffer.drain(..).collect::>(); + let line = std::mem::take(&mut self.buffer); let line_str = String::from_utf8_lossy(&line); self.logger.write_tracing_log(&line_str); } diff --git a/src/agent-client-protocol-http/src/client.rs b/src/agent-client-protocol-http/src/client.rs index 0661b61e..1c899872 100644 --- a/src/agent-client-protocol-http/src/client.rs +++ b/src/agent-client-protocol-http/src/client.rs @@ -1473,10 +1473,15 @@ mod tests { } impl WsSink for RecordingWsSink { - async fn send(&mut self, message: WsMessage) -> Result<(), String> { - self.0 - .unbounded_send(message) - .map_err(|error| error.to_string()) + fn send( + &mut self, + message: WsMessage, + ) -> impl std::future::Future> + Send { + std::future::ready( + self.0 + .unbounded_send(message) + .map_err(|error| error.to_string()), + ) } } diff --git a/src/agent-client-protocol-http/src/websocket_server.rs b/src/agent-client-protocol-http/src/websocket_server.rs index 4ed7fac9..d051f4d0 100644 --- a/src/agent-client-protocol-http/src/websocket_server.rs +++ b/src/agent-client-protocol-http/src/websocket_server.rs @@ -240,7 +240,7 @@ mod tests { }; use async_tungstenite::{tokio::connect_async, tungstenite::Message as ClientWsMessage}; use axum::{Router, extract::WebSocketUpgrade, routing::get}; - use futures::{StreamExt as _, future::BoxFuture}; + use futures::future::BoxFuture; use serde_json::json; use tokio::{ net::TcpListener, diff --git a/src/agent-client-protocol-rmcp/examples/with_mcp_server.rs b/src/agent-client-protocol-rmcp/examples/with_mcp_server.rs index 9479a1bb..36b5d7e5 100644 --- a/src/agent-client-protocol-rmcp/examples/with_mcp_server.rs +++ b/src/agent-client-protocol-rmcp/examples/with_mcp_server.rs @@ -63,6 +63,7 @@ impl ExampleMcpServer { } } +#[allow(unknown_lints, clippy::unused_async_trait_impl)] #[tool_handler] impl ServerHandler for ExampleMcpServer { fn get_info(&self) -> ServerInfo { diff --git a/src/agent-client-protocol-rmcp/src/builder.rs b/src/agent-client-protocol-rmcp/src/builder.rs index 82ff2d03..a41da1ee 100644 --- a/src/agent-client-protocol-rmcp/src/builder.rs +++ b/src/agent-client-protocol-rmcp/src/builder.rs @@ -1,6 +1,6 @@ //! MCP server builder for creating MCP servers. -use std::{marker::PhantomData, pin::pin, sync::Arc}; +use std::{future::Future, marker::PhantomData, pin::pin, sync::Arc}; use futures::future::{BoxFuture, Either}; use futures_concurrency::future::TryJoin; @@ -345,18 +345,18 @@ impl ServerHandler for McpServerConnection { } } - async fn list_tools( + fn list_tools( &self, _request: Option, _context: rmcp::service::RequestContext, - ) -> Result { + ) -> impl Future> + Send { // Return only enabled tools let tools: Vec<_> = self .data .enabled_tools() .map(|tool| make_tool_model(tool.metadata())) .collect(); - Ok(ListToolsResult::with_all_items(tools)) + std::future::ready(Ok(ListToolsResult::with_all_items(tools))) } fn get_info(&self) -> rmcp::model::ServerInfo { diff --git a/src/agent-client-protocol-test/src/bin/mcp_echo_server.rs b/src/agent-client-protocol-test/src/bin/mcp_echo_server.rs index 02cc404e..73462aeb 100644 --- a/src/agent-client-protocol-test/src/bin/mcp_echo_server.rs +++ b/src/agent-client-protocol-test/src/bin/mcp_echo_server.rs @@ -48,6 +48,7 @@ impl EchoServer { } } +#[allow(unknown_lints, clippy::unused_async_trait_impl)] #[tool_handler] impl ServerHandler for EchoServer { fn get_info(&self) -> ServerInfo { diff --git a/src/agent-client-protocol-test/src/lib.rs b/src/agent-client-protocol-test/src/lib.rs index 166c7970..39655e5a 100644 --- a/src/agent-client-protocol-test/src/lib.rs +++ b/src/agent-client-protocol-test/src/lib.rs @@ -1,3 +1,5 @@ +use std::future::Future; + use agent_client_protocol::*; use serde::{Deserialize, Serialize}; @@ -11,8 +13,13 @@ pub mod testy; pub struct MockTransport; impl ConnectTo for MockTransport { - async fn connect_to(self, _client: impl ConnectTo) -> Result<(), Error> { - panic!("MockTransport should never be used in running code - it's only for doctests") + fn connect_to( + self, + _client: impl ConnectTo, + ) -> impl Future> + Send { + std::future::poll_fn(|_| { + panic!("MockTransport should never be used in running code - it's only for doctests") + }) } } diff --git a/src/agent-client-protocol/src/jsonrpc.rs b/src/agent-client-protocol/src/jsonrpc.rs index b7913610..9cff5a6d 100644 --- a/src/agent-client-protocol/src/jsonrpc.rs +++ b/src/agent-client-protocol/src/jsonrpc.rs @@ -6816,12 +6816,12 @@ mod tests { struct ClaimingDynamicHandler; impl HandleDispatchFrom for ClaimingDynamicHandler { - async fn handle_dispatch_from( + fn handle_dispatch_from( &mut self, _message: Dispatch, _connection: ConnectionTo, - ) -> Result, crate::Error> { - Ok(Handled::Yes) + ) -> impl Future, crate::Error>> + Send { + future::ready(Ok(Handled::Yes)) } fn describe_chain(&self) -> impl Debug { diff --git a/src/agent-client-protocol/src/jsonrpc/close.rs b/src/agent-client-protocol/src/jsonrpc/close.rs index dc614776..8c56220b 100644 --- a/src/agent-client-protocol/src/jsonrpc/close.rs +++ b/src/agent-client-protocol/src/jsonrpc/close.rs @@ -29,11 +29,11 @@ pub trait HandleConnectionClose: Send { pub struct NullClose; impl HandleConnectionClose for NullClose { - async fn handle_connection_close( + fn handle_connection_close( self, _connection: ConnectionTo, - ) -> Result<(), crate::Error> { - Ok(()) + ) -> impl Future> + Send { + std::future::ready(Ok(())) } } diff --git a/src/agent-client-protocol/src/jsonrpc/handlers.rs b/src/agent-client-protocol/src/jsonrpc/handlers.rs index efc904b7..30bb4fbe 100644 --- a/src/agent-client-protocol/src/jsonrpc/handlers.rs +++ b/src/agent-client-protocol/src/jsonrpc/handlers.rs @@ -7,6 +7,7 @@ use crate::role::{HasPeer, Role, handle_incoming_dispatch}; use crate::{ConnectionTo, Dispatch, JsonRpcNotification, JsonRpcRequest, UntypedMessage}; // Types re-exported from crate root use super::Responder; +use std::future::Future; use std::marker::PhantomData; use std::ops::AsyncFnMut; @@ -25,15 +26,15 @@ impl HandleDispatchFrom for NullHandler { "(null)" } - async fn handle_dispatch_from( + fn handle_dispatch_from( &mut self, message: Dispatch, _cx: ConnectionTo, - ) -> Result, crate::Error> { - Ok(Handled::No { + ) -> impl Future, crate::Error>> + Send { + std::future::ready(Ok(Handled::No { message, retry: false, - }) + })) } } diff --git a/src/agent-client-protocol/src/jsonrpc/run.rs b/src/agent-client-protocol/src/jsonrpc/run.rs index e8adb790..571bf61a 100644 --- a/src/agent-client-protocol/src/jsonrpc/run.rs +++ b/src/agent-client-protocol/src/jsonrpc/run.rs @@ -33,11 +33,11 @@ pub trait RunWithConnectionTo: Send { pub struct NullRun; impl RunWithConnectionTo for NullRun { - async fn run_with_connection_to( + fn run_with_connection_to( self, _cx: ConnectionTo, - ) -> Result<(), crate::Error> { - Ok(()) + ) -> impl Future> + Send { + std::future::ready(Ok(())) } } diff --git a/src/agent-client-protocol/src/role.rs b/src/agent-client-protocol/src/role.rs index 5c537403..0b8e1057 100644 --- a/src/agent-client-protocol/src/role.rs +++ b/src/agent-client-protocol/src/role.rs @@ -284,15 +284,15 @@ impl Role for UntypedRole { RoleId::from_singleton(self) } - async fn default_handle_dispatch_from( + fn default_handle_dispatch_from( &self, message: Dispatch, _connection: ConnectionTo, - ) -> Result, crate::Error> { - Ok(Handled::No { + ) -> impl Future, crate::Error>> + Send { + std::future::ready(Ok(Handled::No { message, retry: false, - }) + })) } fn counterpart(&self) -> Self::Counterpart { diff --git a/src/agent-client-protocol/src/role/acp.rs b/src/agent-client-protocol/src/role/acp.rs index 754eb57b..e2c6f6c5 100644 --- a/src/agent-client-protocol/src/role/acp.rs +++ b/src/agent-client-protocol/src/role/acp.rs @@ -1,4 +1,4 @@ -use std::{fmt::Debug, hash::Hash}; +use std::{fmt::Debug, future::Future, hash::Hash}; #[cfg(feature = "unstable_protocol_v2")] use futures::{StreamExt as _, future}; @@ -49,15 +49,15 @@ impl Role for Client { Builder::new(self).v1_client() } - async fn default_handle_dispatch_from( + fn default_handle_dispatch_from( &self, message: Dispatch, _connection: ConnectionTo, - ) -> Result, crate::Error> { - Ok(Handled::No { + ) -> impl Future, crate::Error>> + Send { + std::future::ready(Ok(Handled::No { message, retry: false, - }) + })) } fn role_id(&self) -> RoleId { @@ -1303,15 +1303,15 @@ pub struct Proxy; impl Role for Proxy { type Counterpart = Conductor; - async fn default_handle_dispatch_from( + fn default_handle_dispatch_from( &self, message: crate::Dispatch, _connection: crate::ConnectionTo, - ) -> Result, crate::Error> { - Ok(Handled::No { + ) -> impl Future, crate::Error>> + Send { + std::future::ready(Ok(Handled::No { message, retry: false, - }) + })) } fn role_id(&self) -> RoleId { diff --git a/src/agent-client-protocol/src/role/mcp.rs b/src/agent-client-protocol/src/role/mcp.rs index 508c0f40..f33a35e5 100644 --- a/src/agent-client-protocol/src/role/mcp.rs +++ b/src/agent-client-protocol/src/role/mcp.rs @@ -3,6 +3,8 @@ //! These roles are used for MCP connections, which are separate from ACP but //! use the same underlying connection infrastructure. +use std::future::Future; + use crate::{ Handled, RoleId, jsonrpc::{Builder, handlers::NullHandler, run::NullRun}, @@ -24,15 +26,15 @@ impl Role for Client { Server } - async fn default_handle_dispatch_from( + fn default_handle_dispatch_from( &self, message: crate::Dispatch, _connection: crate::ConnectionTo, - ) -> Result, crate::Error> { - Ok(Handled::No { + ) -> impl Future, crate::Error>> + Send { + std::future::ready(Ok(Handled::No { message, retry: false, - }) + })) } } @@ -64,15 +66,15 @@ impl Role for Server { Client } - async fn default_handle_dispatch_from( + fn default_handle_dispatch_from( &self, message: crate::Dispatch, _connection: crate::ConnectionTo, - ) -> Result, crate::Error> { - Ok(Handled::No { + ) -> impl Future, crate::Error>> + Send { + std::future::ready(Ok(Handled::No { message, retry: false, - }) + })) } } diff --git a/src/agent-client-protocol/tests/jsonrpc_advanced.rs b/src/agent-client-protocol/tests/jsonrpc_advanced.rs index e065441a..1ed60d53 100644 --- a/src/agent-client-protocol/tests/jsonrpc_advanced.rs +++ b/src/agent-client-protocol/tests/jsonrpc_advanced.rs @@ -13,7 +13,11 @@ use agent_client_protocol::{ use futures::channel::{mpsc, oneshot}; use futures::{AsyncRead, AsyncWrite, StreamExt as _}; use serde::{Deserialize, Serialize}; -use std::{marker::PhantomData, time::Duration}; +use std::{ + future::{Future, ready}, + marker::PhantomData, + time::Duration, +}; use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt}; /// Test helper to block and wait for a JSON-RPC response. @@ -215,28 +219,30 @@ struct AfterResponseCollector { } impl HandleDispatchFrom for AfterResponseCollector { - async fn handle_dispatch_from( + fn handle_dispatch_from( &mut self, message: Dispatch, _connection: ConnectionTo, - ) -> Result, agent_client_protocol::Error> { - if let Dispatch::Notification(notification) = &message - && AfterResponseNotification::matches_method(¬ification.method) - { - let notification = AfterResponseNotification::parse_message( - ¬ification.method, - ¬ification.params, - )?; - self.notification_tx - .unbounded_send(notification.value) - .map_err(agent_client_protocol::Error::into_internal_error)?; - return Ok(Handled::Yes); - } - - Ok(Handled::No { - message, - retry: false, - }) + ) -> impl Future, agent_client_protocol::Error>> + Send { + ready((|| { + if let Dispatch::Notification(notification) = &message + && AfterResponseNotification::matches_method(¬ification.method) + { + let notification = AfterResponseNotification::parse_message( + ¬ification.method, + ¬ification.params, + )?; + self.notification_tx + .unbounded_send(notification.value) + .map_err(agent_client_protocol::Error::into_internal_error)?; + return Ok(Handled::Yes); + } + + Ok(Handled::No { + message, + retry: false, + }) + })()) } fn describe_chain(&self) -> impl std::fmt::Debug { diff --git a/src/agent-client-protocol/tests/jsonrpc_batch.rs b/src/agent-client-protocol/tests/jsonrpc_batch.rs index 8ebc9c74..97df192e 100644 --- a/src/agent-client-protocol/tests/jsonrpc_batch.rs +++ b/src/agent-client-protocol/tests/jsonrpc_batch.rs @@ -5,11 +5,14 @@ //! one consolidated response array, while requests and notifications initiated //! by the SDK remain individual messages. -use std::sync::{ - Arc, - atomic::{AtomicUsize, Ordering}, +use std::{ + future::{Future, ready}, + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, + time::Duration, }; -use std::time::Duration; use agent_client_protocol::{ Agent, ByteStreams, Channel, ConnectTo, ConnectionTo, Dispatch, Error, HandleDispatchFrom, @@ -128,12 +131,12 @@ impl JsonRpcNotification for TestNotification {} struct RetryErrorHandler; impl HandleDispatchFrom for RetryErrorHandler { - async fn handle_dispatch_from( + fn handle_dispatch_from( &mut self, _message: Dispatch, _connection: ConnectionTo, - ) -> Result, Error> { - Err(Error::internal_error().data("retry handler error won")) + ) -> impl Future, Error>> + Send { + ready(Err(Error::internal_error().data("retry handler error won"))) } fn describe_chain(&self) -> impl std::fmt::Debug { diff --git a/src/agent-client-protocol/tests/jsonrpc_request_cancellation.rs b/src/agent-client-protocol/tests/jsonrpc_request_cancellation.rs index 3e2b50d4..3be9a7c2 100644 --- a/src/agent-client-protocol/tests/jsonrpc_request_cancellation.rs +++ b/src/agent-client-protocol/tests/jsonrpc_request_cancellation.rs @@ -10,7 +10,10 @@ //! - Test handlers report observed cancellations through in-process channels, //! which the test awaits (with a timeout) instead of sleeping. -use std::sync::{Arc, Mutex}; +use std::{ + future::{Future, ready}, + sync::{Arc, Mutex}, +}; use agent_client_protocol::{ Channel, ConnectionTo, Dispatch, HandleDispatchFrom, Handled, JsonRpcMessage, JsonRpcRequest, @@ -151,15 +154,15 @@ impl Role for WrappedHost { RoleId::from_singleton(self) } - async fn default_handle_dispatch_from( + fn default_handle_dispatch_from( &self, message: Dispatch, _connection: ConnectionTo, - ) -> Result, agent_client_protocol::Error> { - Ok(Handled::No { + ) -> impl Future, agent_client_protocol::Error>> + Send { + ready(Ok(Handled::No { message, retry: false, - }) + })) } fn counterpart(&self) -> Self::Counterpart { @@ -174,15 +177,15 @@ impl Role for WrappedCounterpart { RoleId::from_singleton(self) } - async fn default_handle_dispatch_from( + fn default_handle_dispatch_from( &self, message: Dispatch, _connection: ConnectionTo, - ) -> Result, agent_client_protocol::Error> { - Ok(Handled::No { + ) -> impl Future, agent_client_protocol::Error>> + Send { + ready(Ok(Handled::No { message, retry: false, - }) + })) } fn counterpart(&self) -> Self::Counterpart { @@ -197,15 +200,15 @@ impl Role for WrappedSuccessor { RoleId::from_singleton(self) } - async fn default_handle_dispatch_from( + fn default_handle_dispatch_from( &self, message: Dispatch, _connection: ConnectionTo, - ) -> Result, agent_client_protocol::Error> { - Ok(Handled::No { + ) -> impl Future, agent_client_protocol::Error>> + Send { + ready(Ok(Handled::No { message, retry: false, - }) + })) } fn counterpart(&self) -> Self::Counterpart { @@ -220,15 +223,15 @@ impl Role for WrappedSuccessorCounterpart { RoleId::from_singleton(self) } - async fn default_handle_dispatch_from( + fn default_handle_dispatch_from( &self, message: Dispatch, _connection: ConnectionTo, - ) -> Result, agent_client_protocol::Error> { - Ok(Handled::No { + ) -> impl Future, agent_client_protocol::Error>> + Send { + ready(Ok(Handled::No { message, retry: false, - }) + })) } fn counterpart(&self) -> Self::Counterpart { @@ -2040,26 +2043,29 @@ struct CancelCollector { } impl HandleDispatchFrom for CancelCollector { - async fn handle_dispatch_from( + fn handle_dispatch_from( &mut self, message: Dispatch, _connection: ConnectionTo, - ) -> Result, agent_client_protocol::Error> { + ) -> impl Future, agent_client_protocol::Error>> + Send { if let Dispatch::Notification(notification) = &message && CancelRequestNotification::matches_method(¬ification.method) { - let cancel = CancelRequestNotification::parse_message( + let cancel = match CancelRequestNotification::parse_message( ¬ification.method, ¬ification.params, - )?; + ) { + Ok(cancel) => cancel, + Err(error) => return ready(Err(error)), + }; self.tx.unbounded_send(cancel.request_id).unwrap(); - return Ok(Handled::Yes); + return ready(Ok(Handled::Yes)); } - Ok(Handled::No { + ready(Ok(Handled::No { message, retry: false, - }) + })) } fn describe_chain(&self) -> impl std::fmt::Debug { diff --git a/src/agent-client-protocol/tests/jsonrpc_transport_close.rs b/src/agent-client-protocol/tests/jsonrpc_transport_close.rs index 0dd8fdf4..4643af1c 100644 --- a/src/agent-client-protocol/tests/jsonrpc_transport_close.rs +++ b/src/agent-client-protocol/tests/jsonrpc_transport_close.rs @@ -1,7 +1,8 @@ //! Regression tests for incoming transport closure. use std::{ - future, io, + future::{self, Future}, + io, panic::{RefUnwindSafe, UnwindSafe}, sync::{ Arc, Mutex, @@ -160,8 +161,11 @@ where } impl ConnectTo for ImmediateClient { - async fn connect_to(self, _client: impl ConnectTo) -> Result<(), Error> { - Ok(()) + fn connect_to( + self, + _client: impl ConnectTo, + ) -> impl Future> + Send { + future::ready(Ok(())) } } diff --git a/src/agent-client-protocol/tests/protocol_v2.rs b/src/agent-client-protocol/tests/protocol_v2.rs index b4d0b939..1eb24668 100644 --- a/src/agent-client-protocol/tests/protocol_v2.rs +++ b/src/agent-client-protocol/tests/protocol_v2.rs @@ -1,6 +1,7 @@ #![cfg(feature = "unstable_protocol_v2")] use std::{ + future::{Future, ready}, path::PathBuf, sync::{ Arc, @@ -300,8 +301,13 @@ impl ConnectTo for InitializingV1Client { struct RejectingV1Client; impl ConnectTo for RejectingV1Client { - async fn connect_to(self, _agent: impl ConnectTo) -> Result<(), Error> { - Err(Error::internal_error().data("v1 client fallback should not run")) + fn connect_to( + self, + _agent: impl ConnectTo, + ) -> impl Future> + Send { + ready(Err( + Error::internal_error().data("v1 client fallback should not run") + )) } } diff --git a/src/agent-client-protocol/tests/session_v2_mcp.rs b/src/agent-client-protocol/tests/session_v2_mcp.rs index b1c20083..904e7f52 100644 --- a/src/agent-client-protocol/tests/session_v2_mcp.rs +++ b/src/agent-client-protocol/tests/session_v2_mcp.rs @@ -2,6 +2,7 @@ use std::{ collections::BTreeMap, + future::{Future, ready}, path::PathBuf, sync::{ Arc, Mutex, @@ -852,8 +853,13 @@ impl McpServerConnect for DropTrackedMcpConnect { struct ImmediateErrorRunner; impl RunWithConnectionTo for ImmediateErrorRunner { - async fn run_with_connection_to(self, _connection: ConnectionTo) -> Result<(), Error> { - Err(Error::internal_error().data("runner failed before publication")) + fn run_with_connection_to( + self, + _connection: ConnectionTo, + ) -> impl Future> + Send { + ready(Err( + Error::internal_error().data("runner failed before publication") + )) } }