diff --git a/Cargo.toml b/Cargo.toml index 457b90e5..867564c1 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -97,7 +97,7 @@ regex = "1.12.2" [workspace.lints.rust] future_incompatible = { level = "warn", priority = -1 } -let-underscore = "warn" +let_underscore = "warn" missing_debug_implementations = "warn" # missing_docs = "warn" nonstandard_style = { level = "warn", priority = -1 } diff --git a/README.md b/README.md index 8a6b0e54..05b33cdd 100644 --- a/README.md +++ b/README.md @@ -23,11 +23,14 @@ Native MCP-over-ACP support is currently opt-in through the core crate's `unstable_mcp_over_acp` feature. Standalone MCP servers need no ACP transport feature; the rmcp integration exposes a matching passthrough feature when those servers are attached to ACP. Stable protocol v1 supports per-session and global -proxy attachment. Per-session attachment through the draft `V2SessionBuilder` -is also available when both `unstable_protocol_v2` and -`unstable_mcp_over_acp` are enabled. Successful v2 attachments remain active -for the connection lifetime; global proxy attachment and proxy-session helpers -remain v1-only. +proxy attachment. Draft protocol v2 supports the same two attachment scopes +when both `unstable_protocol_v2` and `unstable_mcp_over_acp` are enabled: +`Proxy.v2().with_mcp_server(...)` injects a global declaration into each +supported session setup request, while +`V2SessionBuilder::with_mcp_server(...)` attaches a server to one new session. +Successful v2 attachments remain active for the connection lifetime, and +`V2SessionBuilder::on_proxy_session_start` forwards a proxied setup response +without coupling later session events to that response. **Proxy orchestration** @@ -49,6 +52,12 @@ remain v1-only. session usage are covered in [Protocol V2](./md/protocol-v2.md). +`Client.builder()`, `Agent.builder()`, and `Proxy.builder()` remain stable-v1 +entry points; their `.v2()` counterparts select the draft-v2 API. Raw proxy +routing infrastructure that selects and validates a version itself can use +`without_acp_version_guard`, but ordinary v2 proxy implementations should use +`Proxy.v2()`. + ## Integrations - [Protocol schema and documentation](https://agentclientprotocol.com/) diff --git a/md/protocol-v2.md b/md/protocol-v2.md index 88b520f9..28e85455 100644 --- a/md/protocol-v2.md +++ b/md/protocol-v2.md @@ -19,12 +19,12 @@ of requests or notifications. See [Transport Architecture: JSON-RPC Batch Behavior](./transport-architecture.md#json-rpc-batch-behavior) for the complete rules. -By default, `Client.builder()` and `Agent.builder()` continue to expose the -stable v1 API and advertise protocol v1. To use the v2 API for a connection, -construct the builder with `Client.v2()` or `Agent.v2()`. Fluent typed handlers, -spawned tasks, close callbacks, and `connect_with` receive -`V2ConnectionTo<_>`, so the protocol version is reflected in the high-level -Rust API as well as on the wire: +By default, `Client.builder()`, `Agent.builder()`, and `Proxy.builder()` +continue to expose the stable v1 API. To use the v2 API for a connection, +construct the builder with `Client.v2()`, `Agent.v2()`, or `Proxy.v2()`. +Fluent typed handlers, spawned tasks, close callbacks, and `connect_with` +receive `V2ConnectionTo<_>`, so the protocol version is reflected in the +high-level Rust API as well as on the wire: ```rust use agent_client_protocol::schema::{ProtocolVersion, v2}; @@ -215,14 +215,63 @@ 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; any setup failure cleans it up. This attachment requires both `unstable_protocol_v2` and -`unstable_mcp_over_acp`. Global proxy attachment and proxy-session helpers -remain v1-only. +`unstable_mcp_over_acp`. + +A v2 proxy can instead attach one server globally with +`Proxy.v2().with_mcp_server(...)`. The proxy reuses one connection-scoped +server ID and adds its declaration to v2 `session/new`, `session/resume`, and +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` is the non-blocking setup helper for +a v2 proxy: + +```rust,ignore +use agent_client_protocol::schema::v2; +use agent_client_protocol::{Client, Proxy}; + +Proxy + .v2() + .on_receive_request_from( + Client, + async |request: v2::NewSessionRequest, responder, cx| { + cx.build_session_from(request) + .with_mcp_server(session_server)? + .on_proxy_session_start(responder, async |opened| { + let (session, setup_response) = opened.into_parts(); + record_session(session.session_id(), setup_response); + Ok(()) + }) + }, + agent_client_protocol::on_receive_request!(), + ); +``` + +The helper forwards request cancellation, sends an ordered downstream +`session/new`, installs session routing before later inbound traffic is +dispatched, and forwards the complete `NewSessionResponse`. It 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()`. If an application wants stream ergonomics, it can fan typed updates out from the connection handler with an explicit buffering and subscriber policy. ## Conductor and proxy initialization +Proxy authors should make the version boundary explicit. `Proxy.builder()` is +the stable v1 builder, while `Proxy.v2()` is v2-only and requires +`_proxy/initialize` to select protocol v2. A proxy built for one version rejects +the other version instead of parsing it through a permissive schema. + +Raw routing infrastructure is the exception. If a component deliberately +selects and validates the version itself, it can use +`Proxy.builder().without_acp_version_guard()` and keep protocol-neutral +`ConnectionTo` callbacks. This disables the SDK's automatic version guard and +is not a substitute for selecting `Proxy.v2()` in an ordinary v2 proxy +implementation. + Enable `unstable_protocol_v2` on `agent-client-protocol-conductor` to carry a v2 connection through a conductor proxy chain. The conductor inspects the raw `protocolVersion` before parsing initialization, rewrites ordinary `initialize` diff --git a/src/agent-client-protocol-conductor/src/conductor.rs b/src/agent-client-protocol-conductor/src/conductor.rs index 36944099..67f55575 100644 --- a/src/agent-client-protocol-conductor/src/conductor.rs +++ b/src/agent-client-protocol-conductor/src/conductor.rs @@ -779,12 +779,16 @@ where // passes through messages but which can trigger the // tracing events. if self.trace_handle.is_some() && num_proxies == 0 { + let trace_proxy = Proxy.builder(); + #[cfg(feature = "unstable_protocol_v2")] + let trace_proxy = trace_proxy.without_acp_version_guard(); + self.connect_to_proxy( &client, 0, ComponentIndex::Client, ComponentIndex::Agent, - Proxy.builder(), + trace_proxy, )?; } else { // Spawn each proxy component diff --git a/src/agent-client-protocol-conductor/tests/initialization_sequence.rs b/src/agent-client-protocol-conductor/tests/initialization_sequence.rs index d6291aa3..49bed9aa 100644 --- a/src/agent-client-protocol-conductor/tests/initialization_sequence.rs +++ b/src/agent-client-protocol-conductor/tests/initialization_sequence.rs @@ -312,6 +312,18 @@ async fn run_bad_proxy_test( .await } +fn assert_initialize_proxy_rejection(error: &agent_client_protocol::Error) { + #[cfg(feature = "unstable_protocol_v2")] + let expected = "_proxy/initialize"; + #[cfg(not(feature = "unstable_protocol_v2"))] + let expected = "initialize/proxy"; + + assert!( + error.to_string().contains(expected), + "error should mention {expected}: {error:?}" + ); +} + #[tokio::test] async fn test_conductor_rejects_initialize_proxy_forwarded_to_agent() -> Result<(), agent_client_protocol::Error> { @@ -327,10 +339,7 @@ async fn test_conductor_rejects_initialize_proxy_forwarded_to_agent() .await; if let Err(err) = init_response { - assert!( - err.to_string().contains("initialize/proxy"), - "Error should mention initialize/proxy: {err:?}" - ); + assert_initialize_proxy_rejection(&err); } Ok::<(), agent_client_protocol::Error>(()) @@ -340,12 +349,7 @@ async fn test_conductor_rejects_initialize_proxy_forwarded_to_agent() match result { Ok(()) => panic!("Expected error when proxy forwards InitializeProxyRequest to agent"), - Err(err) => { - assert!( - err.to_string().contains("initialize/proxy"), - "Error should mention initialize/proxy: {err:?}" - ); - } + Err(err) => assert_initialize_proxy_rejection(&err), } Ok(()) @@ -370,10 +374,7 @@ async fn test_conductor_rejects_initialize_proxy_forwarded_to_proxy() // The error may come through recv() or bubble up through the test harness if let Err(err) = init_response { - assert!( - err.to_string().contains("initialize/proxy"), - "Error should mention initialize/proxy: {err:?}" - ); + assert_initialize_proxy_rejection(&err); } Ok::<(), agent_client_protocol::Error>(()) @@ -384,12 +385,7 @@ async fn test_conductor_rejects_initialize_proxy_forwarded_to_proxy() // The error might bubble up through run_test_with_components instead match result { Ok(()) => panic!("Expected error when proxy forwards InitializeProxyRequest to proxy"), - Err(err) => { - assert!( - err.to_string().contains("initialize/proxy"), - "Error should mention initialize/proxy: {err:?}" - ); - } + Err(err) => assert_initialize_proxy_rejection(&err), } Ok(()) diff --git a/src/agent-client-protocol-conductor/tests/initialization_v2.rs b/src/agent-client-protocol-conductor/tests/initialization_v2.rs index 28e9f855..3ec23cf1 100644 --- a/src/agent-client-protocol-conductor/tests/initialization_v2.rs +++ b/src/agent-client-protocol-conductor/tests/initialization_v2.rs @@ -203,13 +203,13 @@ impl ConnectTo for RecordingPassthroughProxy { let sequence = self.sequence; Proxy - .builder() + .v2() .name("v2-passthrough-proxy") .on_receive_request_from( Client, async move |request: v2::InitializeProxyRequest, responder, - cx: ConnectionTo| { + cx: V2ConnectionTo| { assert_eq!( sequence.fetch_add(1, Ordering::SeqCst), 0, @@ -226,7 +226,7 @@ impl ConnectTo for RecordingPassthroughProxy { Client, async move |request: v2::NewSessionRequest, responder, - cx: ConnectionTo| { + cx: V2ConnectionTo| { cx.send_request_to(Agent, request) .forward_response_to(responder) }, @@ -509,6 +509,41 @@ async fn v2_proxy_initialize_precedes_agent_initialize() -> Result<(), Error> { Ok(()) } +#[tokio::test] +async fn v2_tracing_without_user_proxies_uses_version_neutral_bridge() -> Result<(), Error> { + let request = initialize_request(); + let response = initialize_response(); + let sequence = Arc::new(AtomicUsize::new(0)); + let agent = recording_agent(request.clone(), response.clone(), Arc::clone(&sequence), 0); + let components = ProxiesAndAgent::new(agent); + let (trace_tx, _trace_rx) = mpsc::unbounded(); + let (editor_out, conductor_in) = duplex(4096); + let (conductor_out, editor_in) = duplex(4096); + let transport = ByteStreams::new(editor_out.compat_write(), editor_in.compat()); + + Client + .v2() + .name("v2-editor") + .with_spawned(|_cx| async move { + ConductorImpl::new_agent("v2-conductor", components) + .trace_to(trace_tx) + .run(ByteStreams::new( + conductor_out.compat_write(), + conductor_in.compat(), + )) + .await + }) + .connect_with(transport, async move |cx| { + let received = cx.send_request(request).block_task().await?; + assert_eq!(received, response); + Ok(()) + }) + .await?; + + assert_eq!(sequence.load(Ordering::SeqCst), 1); + Ok(()) +} + #[tokio::test] async fn v2_nested_proxy_instantiator_cannot_change_selected_version() -> Result<(), Error> { let request = initialize_request(); @@ -657,6 +692,316 @@ async fn v2_session_new_preserves_request_and_response_through_proxy() -> Result Ok(()) } +#[tokio::test] +async fn v2_proxy_session_helper_preserves_response_and_routes_later_updates() -> Result<(), Error> +{ + let initialize_request = initialize_request(); + let initialize_response = initialize_response(); + let session_id = v2::SessionId::new("v2-helper-session"); + let session_response = v2::NewSessionResponse::new(session_id.clone()) + .config_options(vec![v2::SessionConfigOption::boolean( + "thinking", "Thinking", true, + )]) + .meta(meta("response", "proxy-helper-response")); + let expected_callback_response = session_response.clone(); + let agent_initialize_response = initialize_response.clone(); + let agent_session_response = session_response.clone(); + let agent_session_id = session_id.clone(); + let agent = Agent + .v2() + .on_receive_request( + async move |_request: v2::InitializeRequest, responder, _cx| { + responder.respond(agent_initialize_response.clone()) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async move |_request: v2::NewSessionRequest, responder, cx| { + responder.respond(agent_session_response.clone())?; + cx.send_notification(v2::UpdateSessionNotification::new( + agent_session_id.clone(), + v2::SessionUpdate::StateUpdate(v2::StateUpdate::Running( + v2::RunningStateUpdate::new(), + )), + )) + }, + agent_client_protocol::on_receive_request!(), + ); + + let (callback_tx, mut callback_rx) = mpsc::unbounded(); + let proxy = Proxy.v2().on_receive_request_from( + Client, + async move |request: v2::NewSessionRequest, responder, cx: V2ConnectionTo| { + let callback_tx = callback_tx.clone(); + cx.build_session_from(request).on_proxy_session_start( + responder, + move |opened| async move { + callback_tx + .unbounded_send(( + opened.session().session_id().clone(), + opened.response().clone(), + )) + .map_err(Error::into_internal_error) + }, + ) + }, + agent_client_protocol::on_receive_request!(), + ); + + let (update_tx, mut update_rx) = mpsc::unbounded(); + let (editor_out, conductor_in) = duplex(4096); + let (conductor_out, editor_in) = duplex(4096); + let transport = ByteStreams::new(editor_out.compat_write(), editor_in.compat()); + Client + .v2() + .on_receive_notification( + async move |update: v2::UpdateSessionNotification, _cx| { + update_tx + .unbounded_send(update) + .map_err(Error::into_internal_error) + }, + agent_client_protocol::on_receive_notification!(), + ) + .with_spawned(|_cx| async move { + ConductorImpl::new_agent("v2-conductor", ProxiesAndAgent::new(agent).proxy(proxy)) + .run(ByteStreams::new( + conductor_out.compat_write(), + conductor_in.compat(), + )) + .await + }) + .connect_with(transport, async move |cx| { + cx.send_request(initialize_request).block_task().await?; + + let received = cx + .send_request(v2::NewSessionRequest::new("/v2-helper-session")) + .block_task() + .await?; + assert_eq!(received, session_response); + + let (callback_session_id, callback_response) = + tokio::time::timeout(std::time::Duration::from_secs(2), callback_rx.next()) + .await + .expect("proxy session callback should not hang") + .ok_or_else(|| Error::internal_error().data("proxy callback channel closed"))?; + assert_eq!(callback_session_id, session_id); + assert_eq!(callback_response, expected_callback_response); + + let update = tokio::time::timeout(std::time::Duration::from_secs(2), update_rx.next()) + .await + .expect("post-response session update should not hang") + .ok_or_else(|| Error::internal_error().data("session update channel closed"))?; + assert_eq!(update.session_id, session_id); + assert!(matches!( + update.update, + v2::SessionUpdate::StateUpdate(v2::StateUpdate::Running(_)) + )); + Ok(()) + }) + .await +} + +#[tokio::test] +async fn v2_proxy_session_helper_reissues_cancellation_for_the_downstream_hop() -> Result<(), Error> +{ + let (parked_id_tx, mut parked_id_rx) = mpsc::unbounded(); + let (cancel_tx, mut cancel_rx) = mpsc::unbounded(); + let agent = Agent + .v2() + .on_receive_request( + async |_request: v2::InitializeRequest, responder, _cx| { + responder.respond(initialize_response()) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async move |request: v2::NewSessionRequest, responder, cx| { + if AsRef::::as_ref(&request.cwd).ends_with("park-session") { + parked_id_tx + .unbounded_send(responder.id().clone()) + .map_err(Error::into_internal_error)?; + let cancellation = responder.cancellation(); + cx.spawn(async move { + let result = cancellation + .run_until_cancelled(std::future::pending::< + Result, + >()) + .await; + responder.respond_with_result(result) + })?; + return Ok(()); + } + + responder.respond(v2::NewSessionResponse::new(v2::SessionId::new( + "normal-session", + ))) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_notification( + async move |cancel: v1::CancelRequestNotification, _cx| { + cancel_tx + .unbounded_send(cancel.request_id) + .map_err(Error::into_internal_error) + }, + agent_client_protocol::on_receive_notification!(), + ); + let proxy = Proxy.v2().on_receive_request_from( + Client, + async |request: v2::NewSessionRequest, responder, cx: V2ConnectionTo| { + cx.build_session_from(request) + .on_proxy_session_start(responder, |_opened| async { Ok(()) }) + }, + agent_client_protocol::on_receive_request!(), + ); + + let (editor_out, conductor_in) = duplex(4096); + let (conductor_out, editor_in) = duplex(4096); + let transport = ByteStreams::new(editor_out.compat_write(), editor_in.compat()); + let client_request_id = tokio::time::timeout( + std::time::Duration::from_secs(10), + Client + .v2() + .with_spawned(|_cx| async move { + ConductorImpl::new_agent( + "v2-cancellation-conductor", + ProxiesAndAgent::new(agent).proxy(proxy), + ) + .run(ByteStreams::new( + conductor_out.compat_write(), + conductor_in.compat(), + )) + .await + }) + .connect_with(transport, async move |cx| { + cx.send_request(initialize_request()).block_task().await?; + + let pending = cx.send_request(v2::NewSessionRequest::new("/park-session")); + let client_request_id = pending.id().clone(); + pending.cancel()?; + let error = pending + .block_task() + .await + .expect_err("cancelled v2 session/new should fail"); + assert_eq!(i32::from(error.code), -32800); + + let response = cx + .send_request(v2::NewSessionRequest::new("/normal-session")) + .block_task() + .await?; + assert_eq!(response.session_id, v2::SessionId::new("normal-session")); + Ok(client_request_id) + }), + ) + .await + .expect("v2 proxy cancellation test timed out")?; + + let parked_id = tokio::time::timeout(std::time::Duration::from_secs(2), parked_id_rx.next()) + .await + .expect("agent should observe the forwarded request") + .ok_or_else(|| Error::internal_error().data("parked request channel closed"))?; + assert_ne!( + parked_id, client_request_id, + "each proxy hop must allocate its own request ID" + ); + let cancelled_id = tokio::time::timeout(std::time::Duration::from_secs(2), cancel_rx.next()) + .await + .expect("agent should observe the reissued cancellation") + .ok_or_else(|| Error::internal_error().data("cancellation channel closed"))?; + assert_eq!(cancelled_id, parked_id); + assert!( + cancel_rx.try_recv().is_err(), + "the downstream hop must receive exactly one cancellation" + ); + Ok(()) +} + +#[tokio::test] +async fn v2_proxy_session_helper_forwards_invalid_success_without_closing_connection() +-> Result<(), Error> { + let attempts = Arc::new(AtomicUsize::new(0)); + let agent_attempts = Arc::clone(&attempts); + let agent = Agent + .builder() + .without_acp_version_guard() + .on_receive_request( + async move |request: UntypedMessage, + responder: agent_client_protocol::Responder, + _cx| { + match request.method() { + "initialize" => responder.respond( + serde_json::to_value(initialize_response()) + .map_err(Error::into_internal_error)?, + ), + "session/new" if agent_attempts.fetch_add(1, Ordering::SeqCst) == 0 => { + responder.respond(serde_json::json!({ + "_futureResponseField": { + "preserved": false, + }, + })) + } + "session/new" => responder.respond(serde_json::json!({ + "sessionId": "recovered-session", + "_futureResponseField": { + "preserved": true, + }, + })), + method => responder.respond_with_error( + Error::method_not_found().data(format!("unexpected method `{method}`")), + ), + } + }, + agent_client_protocol::on_receive_request!(), + ); + + let callback_count = Arc::new(AtomicUsize::new(0)); + let proxy_callback_count = Arc::clone(&callback_count); + let proxy = Proxy.v2().on_receive_request_from( + Client, + async move |request: v2::NewSessionRequest, responder, cx: V2ConnectionTo| { + let callback_count = Arc::clone(&proxy_callback_count); + cx.build_session_from(request).on_proxy_session_start( + responder, + move |_opened| async move { + callback_count.fetch_add(1, Ordering::SeqCst); + Ok(()) + }, + ) + }, + agent_client_protocol::on_receive_request!(), + ); + + run_with_conductor(ProxiesAndAgent::new(agent).proxy(proxy), async move |cx| { + cx.send_request(initialize_request()).block_task().await?; + + let error = cx + .send_request(v2::NewSessionRequest::new("/malformed-session")) + .block_task() + .await + .expect_err("a malformed downstream success must be rejected"); + assert!( + error.to_string().contains("sessionId"), + "unexpected malformed response error: {error:?}" + ); + + let response = cx + .send_request(v2::NewSessionRequest::new("/recovered-session")) + .block_task() + .await?; + assert_eq!(response.session_id, v2::SessionId::new("recovered-session")); + Ok(()) + }) + .await?; + + assert_eq!(attempts.load(Ordering::SeqCst), 2); + assert_eq!( + callback_count.load(Ordering::SeqCst), + 1, + "the callback must run only for the valid setup" + ); + Ok(()) +} + #[tokio::test] async fn v2_nested_conductor_preserves_exact_version_unknown_fields() -> Result<(), Error> { let initialize = initialize_request(); diff --git a/src/agent-client-protocol-conductor/tests/mcp_over_acp_polyfill_v2.rs b/src/agent-client-protocol-conductor/tests/mcp_over_acp_polyfill_v2.rs index bf1548bd..284ee0ce 100644 --- a/src/agent-client-protocol-conductor/tests/mcp_over_acp_polyfill_v2.rs +++ b/src/agent-client-protocol-conductor/tests/mcp_over_acp_polyfill_v2.rs @@ -124,7 +124,7 @@ impl ConnectTo for NativeMcpProvider { let disconnect_count = Arc::clone(&self.disconnect_count); Proxy - .builder() + .v2() .name("native-v2-mcp-provider") .on_receive_request_from( Agent, diff --git a/src/agent-client-protocol-conductor/tests/mcp_server_handler_chain_v2.rs b/src/agent-client-protocol-conductor/tests/mcp_server_handler_chain_v2.rs new file mode 100644 index 00000000..75da7d16 --- /dev/null +++ b/src/agent-client-protocol-conductor/tests/mcp_server_handler_chain_v2.rs @@ -0,0 +1,376 @@ +#![cfg(feature = "unstable_protocol_v2")] + +use std::{ + collections::BTreeMap, + path::PathBuf, + sync::{ + Arc, Mutex, + atomic::{AtomicUsize, Ordering}, + }, +}; + +use agent_client_protocol::{ + Agent, ByteStreams, Client, Conductor, ConnectTo, DynConnectTo, Error, NullRun, Proxy, + Responder, V2ConnectionTo, + mcp_server::{McpConnectionTo, McpServer, McpServerConnect}, + role, + schema::{ProtocolVersion, v2}, +}; +use agent_client_protocol_conductor::{ConductorImpl, ProxiesAndAgent}; +use futures::{StreamExt as _, channel::mpsc}; +use serde_json::json; +use tokio::io::duplex; +use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt}; + +fn implementation(name: &str) -> v2::Implementation { + v2::Implementation::new(name, env!("CARGO_PKG_VERSION")) +} + +fn meta() -> v2::Meta { + v2::Meta::from_iter([("preserved".to_owned(), json!({ "nested": true }))]) +} + +fn existing_server() -> v2::McpServer { + v2::McpServer::Other(v2::OtherMcpServer::new( + "_future_transport", + BTreeMap::from([("futureOption".to_owned(), json!({ "nested": true }))]), + )) +} + +#[derive(Debug, PartialEq, Eq)] +struct ObservedMcpContext { + server_id: String, + connection_id: String, +} + +struct RecordingMcpConnect { + contexts: Arc>>, +} + +impl McpServerConnect for RecordingMcpConnect { + fn name(&self) -> String { + "global-v2-server".to_owned() + } + + fn connect(&self, context: McpConnectionTo) -> DynConnectTo { + self.contexts.lock().unwrap().push(ObservedMcpContext { + server_id: context + .server_id() + .expect("the global MCP server should be attached through ACP") + .to_string(), + connection_id: context + .connection_id() + .expect("an attached MCP connection should have an ID") + .to_string(), + }); + DynConnectTo::new(PendingMcpComponent) + } +} + +struct PendingMcpComponent; + +impl ConnectTo for PendingMcpComponent { + async fn connect_to(self, client: impl ConnectTo) -> Result<(), Error> { + role::mcp::Server + .builder() + .connect_with(client, async |_connection| { + std::future::pending::>().await + }) + .await + } +} + +struct GlobalMcpProxy { + setup_handler_calls: Arc, + mcp_contexts: Arc>>, +} + +impl ConnectTo for GlobalMcpProxy { + async fn connect_to(self, conductor: impl ConnectTo) -> Result<(), Error> { + let new_calls = self.setup_handler_calls.clone(); + let resume_calls = self.setup_handler_calls; + let mcp_server = McpServer::new( + RecordingMcpConnect { + contexts: self.mcp_contexts, + }, + NullRun, + ); + + Proxy + .v2() + .name("global-v2-mcp-proxy") + .with_mcp_server(mcp_server) + .on_receive_request_from( + Client, + async move |request: v2::NewSessionRequest, + responder: Responder, + connection: V2ConnectionTo| { + new_calls.fetch_add(1, Ordering::SeqCst); + connection + .send_request_to(Agent, request) + .forward_response_to(responder) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request_from( + Client, + async move |request: v2::ResumeSessionRequest, + responder: Responder, + connection: V2ConnectionTo| { + resume_calls.fetch_add(1, Ordering::SeqCst); + connection + .send_request_to(Agent, request) + .forward_response_to(responder) + }, + agent_client_protocol::on_receive_request!(), + ) + .connect_to(conductor) + .await + } +} + +#[derive(Default)] +struct ObservedSetup { + server_ids: Mutex>, +} + +impl ObservedSetup { + fn record( + &self, + servers: &[v2::McpServer], + expected_existing: &v2::McpServer, + ) -> Result { + let server = match servers { + [existing, v2::McpServer::Acp(server)] if existing == expected_existing => server, + servers => { + return Err(Error::internal_error() + .data(format!("unexpected MCP declarations: {servers:?}"))); + } + }; + if server.name != "global-v2-server" { + return Err(Error::internal_error().data(format!( + "unexpected global MCP server name: {}", + server.name + ))); + } + self.server_ids + .lock() + .unwrap() + .push(server.server_id.clone()); + Ok(server.server_id.clone()) + } +} + +struct RecordingAgent { + setup: Arc, + expected_existing: v2::McpServer, + cwd: PathBuf, + additional_directory: PathBuf, + round_trip_tx: mpsc::UnboundedSender>, +} + +impl ConnectTo for RecordingAgent { + async fn connect_to(self, client: impl ConnectTo) -> Result<(), Error> { + let new_setup = self.setup.clone(); + let resume_setup = self.setup; + let new_existing = self.expected_existing.clone(); + let resume_existing = self.expected_existing; + let new_cwd = self.cwd.clone(); + let resume_cwd = self.cwd; + let new_additional = self.additional_directory.clone(); + let resume_additional = self.additional_directory; + let round_trip_tx = self.round_trip_tx; + + Agent + .v2() + .name("recording-v2-agent") + .on_receive_request( + async |request: v2::InitializeRequest, + responder: Responder, + _connection: V2ConnectionTo| { + responder.respond(v2::InitializeResponse::new( + request.protocol_version, + implementation("recording-v2-agent"), + )) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async move |request: v2::NewSessionRequest, + responder: Responder, + connection: V2ConnectionTo| { + let new_setup = new_setup.clone(); + let new_existing = new_existing.clone(); + let new_cwd = new_cwd.clone(); + let new_additional = new_additional.clone(); + let round_trip_tx = round_trip_tx.clone(); + assert_eq!(request.cwd, v2::AbsolutePath::new(new_cwd)); + assert_eq!( + request.additional_directories.as_slice(), + [v2::AbsolutePath::new(new_additional)] + ); + assert_eq!(request.meta.as_ref(), Some(&meta())); + let server_id = new_setup.record(&request.mcp_servers, &new_existing)?; + let mcp_connection = connection.clone(); + connection.spawn(async move { + let result = async { + let connected = mcp_connection + .send_request(v2::ConnectMcpRequest::new(server_id)) + .block_task() + .await?; + mcp_connection + .send_request(v2::DisconnectMcpRequest::new( + connected.connection_id, + )) + .block_task() + .await?; + Ok(()) + } + .await; + round_trip_tx + .unbounded_send(result) + .map_err(Error::into_internal_error) + })?; + responder.respond(v2::NewSessionResponse::new(v2::SessionId::new( + "global-v2-session", + ))) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async move |request: v2::ResumeSessionRequest, + responder: Responder, + _connection: V2ConnectionTo| { + let resume_setup = resume_setup.clone(); + let resume_existing = resume_existing.clone(); + let resume_cwd = resume_cwd.clone(); + let resume_additional = resume_additional.clone(); + assert_eq!(request.cwd, v2::AbsolutePath::new(resume_cwd)); + assert_eq!( + request.additional_directories.as_slice(), + [v2::AbsolutePath::new(resume_additional)] + ); + assert_eq!( + request.replay_from, + Some(v2::ReplayFrom::Start( + v2::ReplayFromStart::new().meta(meta()) + )) + ); + assert_eq!(request.meta.as_ref(), Some(&meta())); + resume_setup.record(&request.mcp_servers, &resume_existing)?; + responder.respond(v2::ResumeSessionResponse::new()) + }, + agent_client_protocol::on_receive_request!(), + ) + .connect_to(client) + .await + } +} + +async fn run_with_conductor( + proxy: DynConnectTo, + agent: DynConnectTo, + editor_task: impl AsyncFnOnce(V2ConnectionTo) -> Result<(), Error>, +) -> Result<(), Error> { + let (editor_out, conductor_in) = duplex(4096); + let (conductor_out, editor_in) = duplex(4096); + let transport = ByteStreams::new(editor_out.compat_write(), editor_in.compat()); + + Client + .v2() + .name("v2-editor") + .with_spawned(|_connection| async move { + ConductorImpl::new_agent("v2-conductor", ProxiesAndAgent::new(agent).proxy(proxy)) + .run(ByteStreams::new( + conductor_out.compat_write(), + conductor_in.compat(), + )) + .await + }) + .connect_with(transport, editor_task) + .await +} + +#[tokio::test] +async fn v2_global_mcp_attachment_preserves_setup_and_continues_handler_chain() -> Result<(), Error> +{ + let setup_handler_calls = Arc::new(AtomicUsize::new(0)); + let observed_setup = Arc::new(ObservedSetup::default()); + let mcp_contexts = Arc::new(Mutex::new(Vec::new())); + let (round_trip_tx, mut round_trip_rx) = mpsc::unbounded(); + let cwd = PathBuf::from("/tmp/global-v2-mcp"); + let additional_directory = PathBuf::from("/tmp/global-v2-mcp-additional"); + let existing_server = existing_server(); + + let proxy = DynConnectTo::new(GlobalMcpProxy { + setup_handler_calls: setup_handler_calls.clone(), + mcp_contexts: mcp_contexts.clone(), + }); + let agent = DynConnectTo::new(RecordingAgent { + setup: observed_setup.clone(), + expected_existing: existing_server.clone(), + cwd: cwd.clone(), + additional_directory: additional_directory.clone(), + round_trip_tx, + }); + + run_with_conductor(proxy, agent, async move |connection| { + connection + .send_request(v2::InitializeRequest::new( + ProtocolVersion::V2, + implementation("v2-editor"), + )) + .block_task() + .await?; + + let new_session = connection + .send_request( + v2::NewSessionRequest::new(cwd.clone()) + .additional_directories([additional_directory.clone()]) + .mcp_servers(vec![existing_server.clone()]) + .meta(meta()), + ) + .block_task() + .await?; + + tokio::time::timeout(std::time::Duration::from_secs(2), round_trip_rx.next()) + .await + .expect("global MCP connect/disconnect round trip should not hang") + .ok_or_else(|| Error::internal_error().data("MCP round-trip channel closed"))??; + + connection + .send_request( + v2::ResumeSessionRequest::new(new_session.session_id, cwd) + .additional_directories([additional_directory]) + .mcp_servers(vec![existing_server]) + .replay_from(v2::ReplayFrom::Start( + v2::ReplayFromStart::new().meta(meta()), + )) + .meta(meta()), + ) + .block_task() + .await?; + Ok(()) + }) + .await?; + + assert_eq!( + setup_handler_calls.load(Ordering::SeqCst), + 2, + "both typed handlers after the global MCP handler should run" + ); + let server_ids = observed_setup.server_ids.lock().unwrap(); + assert_eq!(server_ids.len(), 2); + assert_eq!( + server_ids[0], server_ids[1], + "a global MCP server must advertise one stable server ID" + ); + let mcp_contexts = mcp_contexts.lock().unwrap(); + assert_eq!(mcp_contexts.len(), 1); + assert_eq!(mcp_contexts[0].server_id, server_ids[0].to_string()); + assert!( + !mcp_contexts[0].connection_id.is_empty(), + "the global MCP connection should receive a connection ID" + ); + Ok(()) +} diff --git a/src/agent-client-protocol-cookbook/src/lib.rs b/src/agent-client-protocol-cookbook/src/lib.rs index 3716be9e..8eb0fe22 100644 --- a/src/agent-client-protocol-cookbook/src/lib.rs +++ b/src/agent-client-protocol-cookbook/src/lib.rs @@ -20,8 +20,10 @@ //! # Building Proxies //! //! A proxy sits between client and agent, intercepting and optionally modifying -//! messages. The most common use case is adding MCP tools. Use [`Proxy.builder()`](agent_client_protocol::Proxy) -//! to build proxy connections. +//! messages. The most common use case is adding MCP tools. Use +//! [`Proxy.builder()`](agent_client_protocol::Proxy) for stable protocol v1 +//! proxy connections. With the core SDK's `unstable_protocol_v2` feature, use +//! `Proxy.v2()` for a draft-v2-only proxy. //! //! **Important:** Proxies don't run standalone—they need the [`agent-client-protocol-conductor`] to //! orchestrate the connection between client, proxies, and agent. See @@ -439,7 +441,8 @@ pub mod global_mcp_server { //! for all sessions. The server is added to the connection's handler chain and //! automatically injects itself into every supported session setup request. //! This pattern requires the core SDK's `unstable_mcp_over_acp` feature (or - //! the rmcp crate's matching passthrough feature). + //! the rmcp crate's matching passthrough feature). Draft v2 additionally + //! requires `unstable_protocol_v2`. //! //! # When to use //! @@ -490,6 +493,16 @@ pub mod global_mcp_server { //! let proxy = MyProxy { mcp_server }; //! ``` //! + //! The example uses stable protocol v1. A draft v2 proxy selects its API + //! before attaching the server: + //! + //! ```rust,ignore + //! Proxy.v2() + //! .with_mcp_server(mcp_server) + //! .connect_to(conductor) + //! .await?; + //! ``` + //! //! # Using rmcp //! //! If you have an existing [rmcp](https://docs.rs/rmcp) server implementation, @@ -551,9 +564,11 @@ pub mod global_mcp_server { //! handler. It: //! //! 1. Intercepts session setup requests and adds a schema-native - //! `McpServer::Acp` declaration with one connection-scoped server ID, - //! reused in each request's `mcp_servers` list (`session/new`, - //! `session/load`, `session/resume`, and feature-gated `session/fork`) + //! `McpServer::Acp` declaration with one connection-scoped server ID. + //! V1 injects it into `session/new`, `session/load`, `session/resume`, + //! and feature-gated `session/fork`; v2 injects it into + //! `session/new`, `session/resume`, and feature-gated `session/fork` + //! while preserving unrelated request fields //! 2. Passes the modified request through to the next handler //! 3. Handles `mcp/connect`, `mcp/message`, and `mcp/disconnect` for that server ID //! @@ -568,7 +583,8 @@ pub mod per_session_mcp_server { //! Use this pattern when each session needs its own MCP server instance //! with access to session-specific context like the working directory. //! It requires the core SDK's `unstable_mcp_over_acp` feature (or the rmcp - //! crate's matching passthrough feature). + //! crate's matching passthrough feature). Draft v2 additionally requires + //! `unstable_protocol_v2`. //! //! # When to use //! @@ -576,7 +592,7 @@ pub mod per_session_mcp_server { //! - You want eventual active-session tracking that does not need to precede later traffic //! - Tools need to customize behavior based on session parameters //! - //! # Basic pattern with `on_proxy_session_start` + //! # Stable v1 pattern with `on_proxy_session_start` //! //! The most common pattern intercepts [`NewSessionRequest`], extracts context, //! creates a per-session MCP server, and uses [`on_proxy_session_start`] to @@ -639,7 +655,46 @@ pub mod per_session_mcp_server { //! ID-keyed state, preinstall a gate or placeholder that later handlers //! await, then populate it from the callback. //! - //! # Alternative: spawning `start_session_proxy` + //! # Draft v2 pattern + //! + //! `Proxy.v2()` exposes the same non-blocking setup shape with v2 schema + //! types. Its `V2SessionBuilder::on_proxy_session_start` callback receives + //! an `OpenedV2Session`, not just a session ID, so it retains both the + //! command-only handle and the complete `NewSessionResponse`: + //! + //! ```rust,ignore + //! use agent_client_protocol::schema::v2; + //! + //! Proxy.v2() + //! .on_receive_request_from( + //! Client, + //! async |request: v2::NewSessionRequest, responder, connection| { + //! let workspace_path = request.cwd.clone(); + //! let mcp_server = build_workspace_server(workspace_path); + //! + //! connection + //! .build_session_from(request) + //! .with_mcp_server(mcp_server)? + //! .on_proxy_session_start(responder, async move |opened| { + //! let (session, setup_response) = opened.into_parts(); + //! tracing::info!( + //! session_id = %session.session_id(), + //! ?setup_response, + //! "Session started" + //! ); + //! Ok(()) + //! }) + //! }, + //! agent_client_protocol::on_receive_request!(), + //! ); + //! ``` + //! + //! The helper forwards upstream cancellation and the complete setup + //! response, and installs session routing before later inbound traffic. + //! The callback runs outside that ordering barrier. V2 updates and + //! interactive requests remain independent connection traffic. + //! + //! # Stable v1 alternative: spawning `start_session_proxy` //! //! If you need the linear [`start_session_proxy`] API, move it into a //! spawned task. Awaiting it directly in the request handler would block diff --git a/src/agent-client-protocol-polyfill/src/mcp_over_acp/mod.rs b/src/agent-client-protocol-polyfill/src/mcp_over_acp/mod.rs index 8b15edf6..7dc7ca97 100644 --- a/src/agent-client-protocol-polyfill/src/mcp_over_acp/mod.rs +++ b/src/agent-client-protocol-polyfill/src/mcp_over_acp/mod.rs @@ -129,8 +129,11 @@ impl ConnectTo for McpOverAcpPolyfill { ) -> Result<(), agent_client_protocol::Error> { let (bridge_tx, bridge_rx) = mpsc::channel(128); - Proxy - .builder() + let proxy = Proxy.builder(); + #[cfg(feature = "unstable_protocol_v2")] + let proxy = proxy.without_acp_version_guard(); + + proxy .name("mcp-over-acp-polyfill") .with_runner(BridgeRunner { bridge_tx: bridge_tx.clone(), diff --git a/src/agent-client-protocol/CHANGELOG.md b/src/agent-client-protocol/CHANGELOG.md index e5d07476..5ec93e8e 100644 --- a/src/agent-client-protocol/CHANGELOG.md +++ b/src/agent-client-protocol/CHANGELOG.md @@ -19,8 +19,17 @@ per-session MCP attachment while preserving independent response consumption. MCP routes are installed and runners begin executing before `session/new` is published; successful attachments remain active for the - connection lifetime. Global proxy attachment and proxy-session helpers - remain v1-only. + connection lifetime. `V2SessionBuilder::on_proxy_session_start` forwards the + complete setup response, preserves cancellation, installs session routing + before later inbound traffic, and then spawns user work with the + `OpenedV2Session`. +- *(unstable-v2)* Add `Proxy::v2()` as the draft-v2-only proxy builder while + keeping `Proxy::builder()` on stable v1. With `unstable_mcp_over_acp`, + `Proxy.v2().with_mcp_server(...)` injects a connection-scoped MCP server + declaration into v2 new, resume, and feature-gated fork setup requests while + preserving unrelated fields. Version-neutral routing infrastructure can + continue to use `Builder::without_acp_version_guard` when it owns raw + version selection and validation. - *(unstable-v2)* Expose protocol-neutral dynamic handler registration on `V2ConnectionTo`. - *(unstable-v2)* Add `ConnectionTo::spawn_connection_with_context` for raw diff --git a/src/agent-client-protocol/README.md b/src/agent-client-protocol/README.md index fd194056..69ea1a64 100644 --- a/src/agent-client-protocol/README.md +++ b/src/agent-client-protocol/README.md @@ -40,14 +40,15 @@ Client.builder() # } ``` -Draft protocol v2 is opt-in through `unstable_protocol_v2`. `Client.v2()` and -`Agent.v2()` callbacks receive a version-typed `V2ConnectionTo` with +Draft protocol v2 is opt-in through `unstable_protocol_v2`. `Client.v2()`, +`Agent.v2()`, and `Proxy.v2()` callbacks receive a version-typed +`V2ConnectionTo` with high-level, command-only session helpers because prompt acceptance and inbound traffic are independent. Session updates and interactive requests use typed connection handlers. See [Protocol V2](https://agentclientprotocol.github.io/rust-sdk/protocol-v2.html#high-level-v2-sessions). -Per-session MCP attachment through `V2SessionBuilder` is available with both -`unstable_protocol_v2` and `unstable_mcp_over_acp`. Global proxy attachment and -proxy-session helpers remain v1-only. +`Proxy.builder()` remains the stable v1 entry point; raw routing infrastructure +that selects and validates the version itself can use +`Proxy.builder().without_acp_version_guard()`. ## MCP Server Attachment @@ -58,9 +59,13 @@ Attached servers are advertised with native `McpServer::Acp` declarations and communicate through `mcp/connect`, `mcp/message`, and `mcp/disconnect`. Use `agent-client-protocol-polyfill` immediately before an HTTP-capable agent. Stable protocol v1 supports per-session and global proxy attachment. Draft -protocol v2 supports per-session `session/new` attachment when both unstable -features are enabled; successful attachments remain active for the connection -lifetime, and its global proxy path remains v1-only. +protocol v2 supports both scopes when both unstable features are enabled: +`Proxy.v2().with_mcp_server(...)` injects a global server into supported setup +requests, and `V2SessionBuilder::with_mcp_server(...)` attaches one to a single +`session/new`. Successful attachments remain active for the connection +lifetime. A v2 proxy can forward setup with +`V2SessionBuilder::on_proxy_session_start`; updates and interactive requests +remain independent connection traffic. ## Learning More diff --git a/src/agent-client-protocol/src/concepts/proxies.rs b/src/agent-client-protocol/src/concepts/proxies.rs index 677f37b3..8479afc7 100644 --- a/src/agent-client-protocol/src/concepts/proxies.rs +++ b/src/agent-client-protocol/src/concepts/proxies.rs @@ -14,10 +14,24 @@ //! Unlike simpler links, there's no default peer - you must always specify //! which direction you're communicating with. //! +//! # Choosing a Protocol Version +//! +//! `Proxy::builder` creates a stable protocol v1 proxy. With the +//! `unstable_protocol_v2` feature, `Proxy.v2()` creates a v2-only proxy whose +//! fluent callbacks receive `V2ConnectionTo`. The builder +//! validates `_proxy/initialize` and later traffic against the selected +//! version. +//! +//! Low-level routing infrastructure that deliberately selects and validates a +//! raw protocol version itself can use +//! `Proxy.builder().without_acp_version_guard()`. Disabling the guard is an +//! explicit version-neutral escape hatch, not the ordinary way to author a v2 +//! proxy. +//! //! # Default Forwarding //! //! By default, [`Proxy`] forwards all messages it doesn't handle. -//! This means a minimal proxy that does nothing is just: +//! This means a minimal stable v1 proxy that does nothing is just: //! //! ``` //! # use agent_client_protocol::{Proxy, Conductor, ConnectTo}; @@ -65,6 +79,7 @@ //! (available in all sessions) or per-session. //! //! These ACP attachment APIs require the `unstable_mcp_over_acp` feature. +//! Draft v2 attachment additionally requires `unstable_protocol_v2`. //! //! ## Global MCP Server //! @@ -82,6 +97,21 @@ //! # } //! ``` //! +//! For draft v2, select the v2 proxy builder before attaching the global +//! server: +//! +//! ```rust,ignore +//! Proxy.v2() +//! .with_mcp_server(my_mcp_server) +//! .connect_to(transport) +//! .await?; +//! ``` +//! +//! The v1 builder injects the declaration into new, load, resume, and +//! feature-gated fork requests. The v2 builder injects it into new, resume, +//! and feature-gated fork requests while preserving unrelated setup fields. +//! Both reuse one connection-scoped server ID. +//! //! ## Per-Session MCP Server //! //! ```ignore @@ -106,6 +136,33 @@ //! # } //! ``` //! +//! The corresponding v2 proxy uses `Proxy.v2()`, a +//! `schema::v2::NewSessionRequest`, and the same fluent session-builder shape. +//! Its `V2SessionBuilder::on_proxy_session_start` callback receives an +//! `OpenedV2Session` containing both the command-only v2 session handle and the +//! complete `NewSessionResponse`: +//! +//! ```rust,ignore +//! Proxy.v2() +//! .on_receive_request_from( +//! Client, +//! async |request: schema::v2::NewSessionRequest, responder, cx| { +//! cx.build_session_from(request) +//! .with_mcp_server(my_mcp_server)? +//! .on_proxy_session_start(responder, async |opened| { +//! let (session, response) = opened.into_parts(); +//! track_session(session.session_id(), response); +//! Ok(()) +//! }) +//! }, +//! agent_client_protocol::on_receive_request!(), +//! ); +//! ``` +//! +//! The setup helper installs session routing and forwards the complete response +//! before spawning the callback. Later updates and interactive requests remain +//! independent traffic handled by typed connection callbacks. +//! //! # The Conductor //! //! Proxies don't run standalone - they're orchestrated by a **conductor**. @@ -138,6 +195,8 @@ //! | Task | Approach | //! |------|----------| //! | Forward everything | Just `connect_to(transport)` | +//! | Author a v1 or v2 proxy | `Proxy.builder()` or `Proxy.v2()` | +//! | Route versions yourself | `without_acp_version_guard` on the raw proxy builder | //! | Intercept specific messages | `on_receive_*_from` with explicit peers | //! | Add global tools | `with_mcp_server` on builder | //! | Add per-session tools | `with_mcp_server` on session builder | diff --git a/src/agent-client-protocol/src/concepts/sessions.rs b/src/agent-client-protocol/src/concepts/sessions.rs index 16dad326..54782e51 100644 --- a/src/agent-client-protocol/src/concepts/sessions.rs +++ b/src/agent-client-protocol/src/concepts/sessions.rs @@ -141,6 +141,33 @@ //! order setup before messages already processed. See [Ordering](super::ordering) //! for details. //! +//! For a draft v2 proxy, use +//! `V2SessionBuilder::on_proxy_session_start` instead. It forwards the complete +//! `NewSessionResponse` and then spawns the callback with an +//! `OpenedV2Session`, so the callback keeps both the command-only session handle +//! and the operation-specific response: +//! +//! ```rust,ignore +//! Proxy.v2() +//! .on_receive_request_from( +//! Client, +//! async |request: schema::v2::NewSessionRequest, responder, cx| { +//! cx.build_session_from(request) +//! .on_proxy_session_start(responder, async |opened| { +//! let (session, setup_response) = opened.into_parts(); +//! track_session(session.session_id(), setup_response); +//! Ok(()) +//! }) +//! }, +//! agent_client_protocol::on_receive_request!(), +//! ); +//! ``` +//! +//! The downstream request inherits upstream cancellation. Session routing is +//! installed before later inbound traffic is dispatched, but user work runs +//! outside that ordering barrier. V2 session updates and interactive requests +//! remain independent traffic handled by typed connection callbacks. +//! //! # Next Steps //! //! - [Callbacks](super::callbacks) - Handle incoming requests diff --git a/src/agent-client-protocol/src/jsonrpc.rs b/src/agent-client-protocol/src/jsonrpc.rs index d946f0ca..a51795a2 100644 --- a/src/agent-client-protocol/src/jsonrpc.rs +++ b/src/agent-client-protocol/src/jsonrpc.rs @@ -50,7 +50,7 @@ use crate::jsonrpc::task_actor::{Task, TaskTx}; use crate::mcp_server::McpServer; use crate::role::HasPeer; use crate::role::Role; -use crate::{Agent, Client, ConnectTo, RoleId}; +use crate::{Agent, Client, ConnectTo, Proxy, RoleId}; /// One valid JSON-RPC message carried inside a [`TransportFrame`]. /// @@ -1042,6 +1042,8 @@ fn default_protocol_mode() -> ProtocolMode { ProtocolMode::v1_agent() } else if role == TypeId::of::() { ProtocolMode::v1_client() + } else if role == TypeId::of::() { + ProtocolMode::v1_proxy() } else { ProtocolMode::disabled() } @@ -1114,6 +1116,18 @@ impl< } } + pub(crate) fn v2_proxy(self) -> V2Builder { + Builder { + host: self.host, + name: self.name, + handler: self.handler, + runner: self.runner, + protocol_mode: ProtocolMode::v2_proxy(), + on_close: self.on_close, + context: PhantomData, + } + } + /// Disable all automatic ACP protocol-version tracking and validation. /// /// This is a low-level escape hatch for protocol-routing infrastructure @@ -1124,8 +1138,10 @@ impl< /// This method is deliberately available only on builders whose callbacks /// receive raw [`ConnectionTo`] values. Applications should normally use /// [`Client::builder`](crate::Client::builder), - /// [`Agent::builder`](crate::Agent::builder), [`Client::v2`](crate::Client::v2), - /// or [`Agent::v2`](crate::Agent::v2) instead. + /// [`Agent::builder`](crate::Agent::builder), + /// [`Proxy::builder`](crate::Proxy::builder), [`Client::v2`](crate::Client::v2), + /// [`Agent::v2`](crate::Agent::v2), or [`Proxy::v2`](crate::Proxy::v2) + /// instead. /// /// ```compile_fail /// # use agent_client_protocol::Client; @@ -1729,30 +1745,6 @@ impl< self.with_handler(handler) } - /// Add an MCP server to session setup requests proxied through this connection. - /// - /// The same native MCP server declaration is added to new, load, and resume - /// requests, plus fork requests when `unstable_session_fork` is enabled. - /// - /// Only applicable to proxies. - #[cfg(feature = "unstable_mcp_over_acp")] - pub fn with_mcp_server( - self, - mcp_server: McpServer>, - ) -> Builder< - Host, - impl HandleDispatchFrom, - impl RunWithConnectionTo, - Close, - Context, - > - where - Host::Counterpart: HasPeer + HasPeer, - { - let (handler, runner) = mcp_server.into_handler_and_runner(); - self.with_handler(handler).with_runner(runner) - } - /// Run in server mode with the provided transport. /// /// This drives the connection by continuously processing messages from the transport @@ -1999,6 +1991,74 @@ impl< } } +#[cfg(feature = "unstable_mcp_over_acp")] +impl< + Host: Role, + Handler: HandleDispatchFrom, + Runner: RunWithConnectionTo, + Close: HandleConnectionClose, +> Builder +{ + /// Add an MCP server to protocol v1 session setup requests proxied through + /// this connection. + /// + /// The same native MCP server declaration is added to new, load, and resume + /// requests, plus fork requests when `unstable_session_fork` is enabled. + /// + /// Only applicable to proxies. Use the same method on `V2Builder` to + /// attach the server to protocol v2 setup requests. + pub fn with_mcp_server( + self, + mcp_server: McpServer>, + ) -> Builder< + Host, + impl HandleDispatchFrom, + impl RunWithConnectionTo, + Close, + RawConnectionContext, + > + where + Host::Counterpart: HasPeer + HasPeer, + { + let (handler, runner) = mcp_server.into_handler_and_runner(); + self.with_handler(handler).with_runner(runner) + } +} + +#[cfg(all(feature = "unstable_mcp_over_acp", feature = "unstable_protocol_v2"))] +impl< + Host: Role, + Handler: HandleDispatchFrom, + Runner: RunWithConnectionTo, + Close: HandleConnectionClose, +> Builder +{ + /// Add an MCP server to protocol v2 session setup requests proxied through + /// this connection. + /// + /// The same native MCP server declaration is added to new and resume + /// requests, plus fork requests when `unstable_session_fork` is enabled. + /// Unrelated request fields are preserved exactly. + /// + /// Only applicable to proxies. + pub fn with_mcp_server( + self, + mcp_server: McpServer>, + ) -> Builder< + Host, + impl HandleDispatchFrom, + impl RunWithConnectionTo, + Close, + V2ConnectionContext, + > + where + Host::Counterpart: HasPeer + HasPeer, + { + let (handler, runner) = mcp_server.into_v2_handler_and_runner(); + self.with_handler(handler).with_runner(runner) + } +} + impl ConnectTo for Builder where R: Role, @@ -3044,10 +3104,12 @@ enum OutgoingMessage { /// the original method method: String, - /// the message to send; this may have a distinct method - /// depending on the peer + /// The logical message before peer-direction wrapping. untyped: UntypedMessage, + /// How to transform the logical message for its target peer. + remote_style: crate::role::RemoteStyle, + /// Optional prerequisite that must finish before the request becomes /// visible on the transport. readiness: Option, @@ -3932,6 +3994,35 @@ impl ConnectionTo { ) } + /// Send an ordered request with readiness and valid-success hooks. + #[cfg(all(feature = "unstable_protocol_v2", feature = "unstable_mcp_over_acp"))] + pub(crate) fn send_ordered_request_to_with_response_hook_after< + Peer: Role, + Req: JsonRpcRequest, + BeforeSend: Future> + Send + 'static, + >( + &self, + peer: Peer, + request: Req, + before_send: BeforeSend, + response_hook: impl FnOnce(&Req::Response) -> Result<(), crate::Error> + Send + 'static, + ) -> SentRequest + where + Counterpart: HasPeer, + { + let hook: ResponseRouteHook = Box::new(move |method, value| { + let response = Req::Response::from_value(method, value.clone())?; + response_hook(&response) + }); + self.send_request_to_with_options( + peer, + request, + true, + Some(RequestReadiness::new(before_send)), + Some(hook), + ) + } + /// Send a request whose callback must run before later inbound messages. /// /// The ordering marker is installed before the request enters the outgoing @@ -3988,7 +4079,7 @@ impl ConnectionTo { .map(move |json| ::from_value(&method, json)); } - match remote_style.transform_outgoing_message(request) { + match request.to_untyped_message() { Ok(untyped) => { // Register before enqueueing so incoming EOF can fail every // observable request before close callbacks begin. The @@ -4011,6 +4102,7 @@ impl ConnectionTo { id: id.clone(), method: method.clone(), untyped, + remote_style, readiness, }; @@ -6462,6 +6554,142 @@ mod tests { }); } + #[cfg(feature = "unstable_protocol_v2")] + #[test] + fn proxy_builders_select_exact_proxy_protocol_guards() -> Result<(), crate::Error> { + use crate::schema::ProtocolVersion; + + for (mode, selected, unsupported) in [ + ( + Proxy.builder().protocol_mode, + ProtocolVersion::V1, + ProtocolVersion::V2, + ), + ( + Proxy.v2().protocol_mode, + ProtocolVersion::V2, + ProtocolVersion::V1, + ), + ] { + assert_eq!(mode.api_protocol_version(), Some(selected)); + + let error = ProtocolCompat::new(mode) + .incoming_message(UntypedMessage::new( + "_proxy/initialize", + serde_json::json!({ "protocolVersion": unsupported }), + )?) + .expect_err("a proxy builder must reject the other protocol version"); + let data = error + .data + .as_ref() + .and_then(|data| data.as_str()) + .unwrap_or_default(); + assert!( + data.contains(&format!("only supports ACP protocol version {selected}")), + "{error:?}" + ); + } + + Ok(()) + } + + #[cfg(feature = "unstable_protocol_v2")] + #[test] + fn v2_proxy_rejects_explicitly_prewrapped_initialize_request() { + let (message_tx, message_rx) = mpsc::unbounded(); + let (task_tx, _task_rx) = mpsc::unbounded(); + let (dynamic_handler_tx, _dynamic_handler_rx) = mpsc::unbounded(); + let transport_completion: SharedTransportCompletion = + future::ready(Ok::<(), crate::Error>(())).boxed().shared(); + let pending_replies = PendingReplies::default(); + let connection = ConnectionTo::new( + crate::Conductor, + message_tx, + task_tx, + dynamic_handler_tx, + transport_completion, + pending_replies.registrar(), + ProtocolMode::v2_proxy(), + ); + + let request = crate::schema::SuccessorMessage { + message: UntypedMessage::new( + "initialize", + serde_json::json!({ "protocolVersion": crate::schema::ProtocolVersion::V1 }), + ) + .expect("test initialize request should serialize"), + meta: None, + }; + let sent = connection.send_request_to(Agent, request); + + let (transport_tx, mut transport_rx) = mpsc::unbounded(); + let mut actor = Box::pin(outgoing_actor::outgoing_protocol_actor( + message_rx, + pending_replies, + transport_tx, + ProtocolCompat::new(ProtocolMode::v2_proxy()), + )); + assert!( + actor.as_mut().now_or_never().is_none(), + "the outgoing actor should continue after rejecting the request" + ); + assert!( + transport_rx.next().now_or_never().is_none(), + "an explicitly prewrapped initialize must not reach the transport" + ); + + let error = futures::executor::block_on(sent.block_task()) + .expect_err("connection routing must own successor wrapping"); + let data = error + .data + .as_ref() + .and_then(|data| data.as_str()) + .unwrap_or_default(); + assert!(data.contains("logical `initialize`"), "{error:?}"); + assert!(data.contains("_proxy/successor"), "{error:?}"); + } + + #[cfg(feature = "unstable_protocol_v2")] + #[test] + fn v2_proxy_builder_exposes_typed_context_to_user_callbacks() { + fn assert_v2_context(_connection: &V2ConnectionTo) {} + + let _builder = Proxy + .v2() + .on_receive_request_from( + Client, + async |_request: UntypedMessage, _responder, connection| { + assert_v2_context(&connection); + Ok(()) + }, + crate::on_receive_request!(), + ) + .on_receive_notification_from( + Agent, + async |_notification: UntypedMessage, connection| { + assert_v2_context(&connection); + Ok(()) + }, + crate::on_receive_notification!(), + ) + .on_receive_dispatch_from( + Client, + async |_dispatch: Dispatch, connection| { + assert_v2_context(&connection); + Ok(()) + }, + crate::on_receive_dispatch!(), + ) + .with_spawned(async |connection| { + assert_v2_context(&connection); + Ok(()) + }) + .on_close(async |connection| { + assert_v2_context(&connection); + Ok(()) + }); + } + #[cfg(feature = "unstable_protocol_v2")] #[test] fn raw_connection_spawns_v2_builder_with_typed_child_callback() { diff --git a/src/agent-client-protocol/src/jsonrpc/outgoing_actor.rs b/src/agent-client-protocol/src/jsonrpc/outgoing_actor.rs index 351964ba..b5114939 100644 --- a/src/agent-client-protocol/src/jsonrpc/outgoing_actor.rs +++ b/src/agent-client-protocol/src/jsonrpc/outgoing_actor.rs @@ -87,6 +87,7 @@ pub(super) async fn outgoing_protocol_actor( id, method, untyped, + remote_style, readiness, } => { // Requests register their response destination synchronously @@ -116,7 +117,8 @@ pub(super) async fn outgoing_protocol_actor( } let request = match protocol_compat - .outgoing_message(untyped) + .outgoing_message(untyped, remote_style) + .and_then(|untyped| remote_style.transform_outgoing_message(untyped)) .and_then(|untyped| untyped.into_raw_jsonrpc_message(Some(id.clone()))) { Ok(request) => request, diff --git a/src/agent-client-protocol/src/jsonrpc/protocol_compat.rs b/src/agent-client-protocol/src/jsonrpc/protocol_compat.rs index b5c788f8..0d68b80a 100644 --- a/src/agent-client-protocol/src/jsonrpc/protocol_compat.rs +++ b/src/agent-client-protocol/src/jsonrpc/protocol_compat.rs @@ -1,7 +1,7 @@ #[cfg(not(feature = "unstable_protocol_v2"))] mod imp { #![allow(clippy::unused_self, clippy::unnecessary_wraps)] - use crate::UntypedMessage; + use crate::{UntypedMessage, role::RemoteStyle}; #[derive(Clone, Copy, Debug, Default)] pub(crate) struct ProtocolMode; @@ -19,6 +19,10 @@ mod imp { Self } + pub(crate) fn v1_proxy() -> Self { + Self + } + pub(crate) fn merge(self, _other: Self) -> Self { self } @@ -42,6 +46,7 @@ mod imp { pub(crate) fn outgoing_message( &self, message: UntypedMessage, + _remote_style: RemoteStyle, ) -> Result { Ok(message) } @@ -82,8 +87,8 @@ mod imp { mod imp { use std::sync::{Arc, Mutex}; - use crate::UntypedMessage; use crate::schema::ProtocolVersion; + use crate::{UntypedMessage, role::RemoteStyle}; #[derive(Clone, Copy, Debug)] pub(crate) enum ProtocolMode { @@ -91,9 +96,39 @@ mod imp { Acp(AcpProtocolMode), } - #[derive(Clone, Copy, Debug)] + #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub(crate) struct AcpProtocolMode { api: ProtocolVersionKind, + initialize_surface: InitializeSurface, + } + + #[derive(Clone, Copy, Debug, PartialEq, Eq)] + enum InitializeSurface { + Peer, + Proxy, + } + + impl AcpProtocolMode { + fn is_incoming_initialize_request(self, method: &str) -> bool { + match self.initialize_surface { + InitializeSurface::Peer => method == "initialize", + InitializeSurface::Proxy => method == "_proxy/initialize", + } + } + + fn is_incoming_initialize_response(method: &str) -> bool { + // Pending replies retain the logical method from before successor + // wrapping, so a proxy's downstream initialize response is also + // keyed by `initialize`. + method == "initialize" + } + + fn is_outgoing_initialize_response(self, method: &str) -> bool { + match self.initialize_surface { + InitializeSurface::Peer => method == "initialize", + InitializeSurface::Proxy => method == "_proxy/initialize", + } + } } impl ProtocolMode { @@ -104,24 +139,42 @@ mod imp { pub(crate) fn v1_agent() -> Self { Self::Acp(AcpProtocolMode { api: ProtocolVersionKind::V1, + initialize_surface: InitializeSurface::Peer, }) } pub(crate) fn v1_client() -> Self { Self::Acp(AcpProtocolMode { api: ProtocolVersionKind::V1, + initialize_surface: InitializeSurface::Peer, + }) + } + + pub(crate) fn v1_proxy() -> Self { + Self::Acp(AcpProtocolMode { + api: ProtocolVersionKind::V1, + initialize_surface: InitializeSurface::Proxy, }) } pub(crate) fn v2_agent() -> Self { Self::Acp(AcpProtocolMode { api: ProtocolVersionKind::V2, + initialize_surface: InitializeSurface::Peer, }) } pub(crate) fn v2_client() -> Self { Self::Acp(AcpProtocolMode { api: ProtocolVersionKind::V2, + initialize_surface: InitializeSurface::Peer, + }) + } + + pub(crate) fn v2_proxy() -> Self { + Self::Acp(AcpProtocolMode { + api: ProtocolVersionKind::V2, + initialize_surface: InitializeSurface::Proxy, }) } @@ -135,6 +188,11 @@ mod imp { "cannot merge ACP builders with different API protocol versions; \ handler chains share a single API surface", ); + assert_eq!( + this.initialize_surface, other.initialize_surface, + "cannot merge standard ACP and proxy ACP builders; \ + handler chains share one initialization surface", + ); Self::Acp(this) } } @@ -210,9 +268,14 @@ mod imp { return Ok(message); }; - if message.method() == "initialize" { + if mode.is_incoming_initialize_request(message.method()) { return self.incoming_initialize_request(mode, message); } + if mode.initialize_surface == InitializeSurface::Proxy + && (message.method() == "initialize" || successor_encloses_initialize(&message)) + { + return Err(invalid_proxy_initialize_direction()); + } ensure_matching_protocol_version( message.method(), @@ -225,13 +288,16 @@ mod imp { pub(crate) fn outgoing_message( &self, mut message: UntypedMessage, + remote_style: RemoteStyle, ) -> Result { let Some(mode) = self.mode else { return Ok(message); }; - let wire_version = if message.method() == "initialize" { - set_protocol_version(&mut message.params, mode.api)?; + let wire_version = if let Some(params) = + outgoing_initialize_params(mode, remote_style, &mut message)? + { + set_protocol_version(params, mode.api)?; self.set_pending_initialize(mode.api); mode.api } else { @@ -283,7 +349,7 @@ mod imp { return result; }; - if method == "initialize" { + if AcpProtocolMode::is_incoming_initialize_response(method) { return self.incoming_initialize_response(mode, result); } @@ -303,7 +369,7 @@ mod imp { // Always drain any pending initialize state so a failed initialize // doesn't leak negotiation state to a subsequent request. - let pending_initialize = if method == "initialize" { + let pending_initialize = if mode.is_outgoing_initialize_response(method) { self.take_pending_initialize() } else { None @@ -311,7 +377,7 @@ mod imp { let mut value = result?; - let wire_version = if method == "initialize" { + 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)?; @@ -402,22 +468,86 @@ mod imp { serde_json::from_value(version.clone()).map_err(|_| invalid_initialize_protocol_version()) } + fn outgoing_initialize_params( + mode: AcpProtocolMode, + remote_style: RemoteStyle, + message: &mut UntypedMessage, + ) -> Result, crate::Error> { + if successor_encloses_initialize(message) { + return Err(invalid_prewrapped_initialize()); + } + + match mode.initialize_surface { + InitializeSurface::Peer => { + Ok((message.method() == "initialize").then_some(&mut message.params)) + } + InitializeSurface::Proxy => { + if message.method() == "initialize" { + if remote_style != RemoteStyle::Successor { + return Err(invalid_proxy_initialize_direction()); + } + return Ok(Some(&mut message.params)); + } + if message.method() == "_proxy/initialize" { + return Err(invalid_proxy_initialize_direction()); + } + Ok(None) + } + } + } + + fn successor_encloses_initialize(message: &UntypedMessage) -> bool { + let mut method = message.method(); + let mut params = message.params(); + let mut wrapped = false; + + while method == "_proxy/successor" { + wrapped = true; + let Some(inner_method) = params.get("method").and_then(serde_json::Value::as_str) + else { + return false; + }; + method = inner_method; + if method == "_proxy/successor" { + let Some(inner_params) = params.get("params") else { + return false; + }; + params = inner_params; + } + } + + wrapped && matches!(method, "initialize" | "_proxy/initialize") + } + fn invalid_initialize_protocol_version() -> crate::Error { crate::Error::invalid_params() .data("initialize.protocolVersion must be a valid ACP protocol version") } + 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`", + ) + } + + fn invalid_prewrapped_initialize() -> crate::Error { + crate::Error::invalid_request().data( + "initialize requests must be sent as a logical `initialize` message; `_proxy/successor` wrapping is applied by connection routing", + ) + } + fn set_protocol_version( value: &mut serde_json::Value, version: ProtocolVersionKind, ) -> Result<(), crate::Error> { - if let serde_json::Value::Object(object) = value { - object.insert( - "protocolVersion".into(), - serde_json::to_value(version.as_protocol_version()) - .map_err(crate::Error::into_internal_error)?, - ); - } + let serde_json::Value::Object(object) = value else { + return Err(invalid_initialize_protocol_version()); + }; + object.insert( + "protocolVersion".into(), + serde_json::to_value(version.as_protocol_version()) + .map_err(crate::Error::into_internal_error)?, + ); Ok(()) } @@ -471,6 +601,14 @@ mod imp { .negotiated } + fn pending_initialize(compat: &ProtocolCompat) -> Option { + compat + .state + .lock() + .expect("protocol compatibility state mutex poisoned") + .pending_initialize + } + fn v2_implementation() -> v2::Implementation { v2::Implementation::new("protocol-compat-test", env!("CARGO_PKG_VERSION")) } @@ -515,10 +653,10 @@ mod imp { let compat = ProtocolCompat::new(ProtocolMode::v2_client()); assert_eq!(compat.active_wire_version(), ProtocolVersionKind::V2); - compat.outgoing_message(UntypedMessage::new( - "initialize", - v2_initialize_request(ProtocolVersion::V1), - )?)?; + compat.outgoing_message( + UntypedMessage::new("initialize", v2_initialize_request(ProtocolVersion::V1))?, + RemoteStyle::Counterpart, + )?; assert_eq!(negotiated(&compat), ProtocolVersionKind::V2); assert_eq!(compat.active_wire_version(), ProtocolVersionKind::V2); @@ -541,10 +679,10 @@ mod imp { let compat = ProtocolCompat::new(ProtocolMode::v2_client()); assert_eq!(compat.active_wire_version(), ProtocolVersionKind::V2); - compat.outgoing_message(UntypedMessage::new( - "initialize", - v2_initialize_request(ProtocolVersion::V1), - )?)?; + compat.outgoing_message( + UntypedMessage::new("initialize", v2_initialize_request(ProtocolVersion::V1))?, + RemoteStyle::Counterpart, + )?; assert_eq!(negotiated(&compat), ProtocolVersionKind::V2); assert_eq!(compat.active_wire_version(), ProtocolVersionKind::V2); @@ -567,10 +705,10 @@ mod imp { serde_json::json!({ "protocolVersion": 100_000 }), ] { let compat = ProtocolCompat::new(ProtocolMode::v2_client()); - compat.outgoing_message(UntypedMessage::new( - "initialize", - v2_initialize_request(ProtocolVersion::V1), - )?)?; + compat.outgoing_message( + UntypedMessage::new("initialize", v2_initialize_request(ProtocolVersion::V1))?, + RemoteStyle::Counterpart, + )?; let error = compat .incoming_response("initialize", Ok(value)) @@ -613,11 +751,288 @@ mod imp { Ok(()) } + #[test] + fn proxy_initialize_request_requires_the_selected_protocol_version() + -> Result<(), crate::Error> { + for (mode, selected, unsupported, selected_kind) in [ + ( + ProtocolMode::v1_proxy(), + ProtocolVersion::V1, + ProtocolVersion::V2, + ProtocolVersionKind::V1, + ), + ( + ProtocolMode::v2_proxy(), + ProtocolVersion::V2, + ProtocolVersion::V1, + ProtocolVersionKind::V2, + ), + ] { + let compat = ProtocolCompat::new(mode); + let error = compat + .incoming_message(UntypedMessage::new( + "_proxy/initialize", + serde_json::json!({ "protocolVersion": unsupported }), + )?) + .expect_err( + "proxy initialization must reject a protocol version outside its API", + ); + let data = error + .data + .as_ref() + .and_then(|data| data.as_str()) + .unwrap_or_default(); + assert!( + data.contains(&format!("only supports ACP protocol version {selected}")), + "{error:?}" + ); + assert_eq!(negotiated(&compat), selected_kind); + assert_eq!(pending_initialize(&compat), None); + assert_eq!(compat.active_wire_version(), selected_kind); + } + + Ok(()) + } + + #[test] + fn proxy_initialize_requests_reject_wrong_directions_and_wrapping() + -> Result<(), crate::Error> { + for mode in [ProtocolMode::v1_proxy(), ProtocolMode::v2_proxy()] { + let compat = ProtocolCompat::new(mode); + for error in [ + compat + .incoming_message(UntypedMessage::new( + "initialize", + serde_json::json!({ "protocolVersion": ProtocolVersion::V2 }), + )?) + .expect_err("proxy initialization must arrive through `_proxy/initialize`"), + compat + .incoming_message(UntypedMessage::new( + "_proxy/successor", + serde_json::json!({ + "method": "initialize", + "params": { "protocolVersion": ProtocolVersion::V2 } + }), + )?) + .expect_err("successor traffic must not carry initialization"), + compat + .incoming_message(UntypedMessage::new( + "_proxy/successor", + serde_json::json!({ + "method": "_proxy/successor", + "params": { + "method": "_proxy/initialize", + "params": { "protocolVersion": ProtocolVersion::V2 } + } + }), + )?) + .expect_err("nested successor traffic must not carry initialization"), + compat + .outgoing_message( + UntypedMessage::new( + "initialize", + serde_json::json!({ "protocolVersion": ProtocolVersion::V2 }), + )?, + RemoteStyle::Predecessor, + ) + .expect_err("downstream proxy initialization must use `_proxy/successor`"), + compat + .outgoing_message( + UntypedMessage::new( + "_proxy/initialize", + serde_json::json!({ "protocolVersion": ProtocolVersion::V2 }), + )?, + RemoteStyle::Successor, + ) + .expect_err("a proxy must not send `_proxy/initialize` downstream"), + ] { + let data = error + .data + .as_ref() + .and_then(|data| data.as_str()) + .unwrap_or_default(); + assert!(data.contains("_proxy/initialize"), "{error:?}"); + assert!(data.contains("_proxy/successor"), "{error:?}"); + } + assert_eq!(pending_initialize(&compat), None); + } + + Ok(()) + } + + #[test] + fn outgoing_initialize_rejects_non_object_params() -> Result<(), crate::Error> { + for (compat, message, remote_style) in [ + ( + ProtocolCompat::new(ProtocolMode::v2_client()), + UntypedMessage::new("initialize", serde_json::Value::Null)?, + RemoteStyle::Counterpart, + ), + ( + ProtocolCompat::new(ProtocolMode::v2_proxy()), + UntypedMessage::new("initialize", serde_json::Value::Null)?, + RemoteStyle::Successor, + ), + ] { + let error = compat + .outgoing_message(message, remote_style) + .expect_err("initialize params must be an object"); + let data = error + .data + .as_ref() + .and_then(|data| data.as_str()) + .unwrap_or_default(); + assert!(data.contains("protocolVersion"), "{error:?}"); + assert_eq!(pending_initialize(&compat), None); + } + + Ok(()) + } + + #[test] + fn outgoing_initialize_rejects_explicit_successor_wrapping() -> Result<(), crate::Error> { + for message in [ + UntypedMessage::new( + "_proxy/successor", + serde_json::json!({ + "method": "initialize", + "params": { "protocolVersion": ProtocolVersion::V1 } + }), + )?, + UntypedMessage::new( + "_proxy/successor", + serde_json::json!({ + "method": "_proxy/successor", + "params": { + "method": "initialize", + "params": { "protocolVersion": ProtocolVersion::V1 } + } + }), + )?, + ] { + let compat = ProtocolCompat::new(ProtocolMode::v2_proxy()); + let error = compat + .outgoing_message(message, RemoteStyle::Successor) + .expect_err("connection routing must own successor wrapping"); + let data = error + .data + .as_ref() + .and_then(|data| data.as_str()) + .unwrap_or_default(); + assert!(data.contains("logical `initialize`"), "{error:?}"); + assert!(data.contains("_proxy/successor"), "{error:?}"); + assert_eq!(pending_initialize(&compat), None); + } + + Ok(()) + } + + #[test] + fn proxy_initialize_round_trip_uses_the_selected_protocol_version() + -> Result<(), crate::Error> { + for (mode, selected, selected_kind) in [ + ( + ProtocolMode::v1_proxy(), + ProtocolVersion::V1, + ProtocolVersionKind::V1, + ), + ( + ProtocolMode::v2_proxy(), + ProtocolVersion::V2, + ProtocolVersionKind::V2, + ), + ] { + let compat = ProtocolCompat::new(mode); + let request = compat.incoming_message(UntypedMessage::new( + "_proxy/initialize", + serde_json::json!({ "protocolVersion": selected }), + )?)?; + assert_eq!( + required_protocol_version_from_value(request.params())?, + selected + ); + assert_eq!(pending_initialize(&compat), Some(selected_kind)); + + let response = compat.outgoing_response( + "_proxy/initialize", + Ok(serde_json::json!({ "protocolVersion": selected })), + )?; + assert_eq!(required_protocol_version_from_value(&response)?, selected); + assert_eq!(negotiated(&compat), selected_kind); + assert_eq!(pending_initialize(&compat), None); + assert_eq!(compat.active_wire_version(), selected_kind); + } + + Ok(()) + } + + #[test] + fn proxy_initialize_response_rejects_the_wrong_version_and_clears_pending_state() + -> Result<(), crate::Error> { + for (mode, selected, unsupported, selected_kind) in [ + ( + ProtocolMode::v1_proxy(), + ProtocolVersion::V1, + ProtocolVersion::V2, + ProtocolVersionKind::V1, + ), + ( + ProtocolMode::v2_proxy(), + ProtocolVersion::V2, + ProtocolVersion::V1, + ProtocolVersionKind::V2, + ), + ] { + let compat = ProtocolCompat::new(mode); + let request = compat.outgoing_message( + UntypedMessage::new( + "initialize", + serde_json::json!({ "protocolVersion": unsupported }), + )?, + RemoteStyle::Successor, + )?; + assert_eq!( + required_protocol_version_from_value(request.params())?, + selected + ); + assert_eq!(pending_initialize(&compat), Some(selected_kind)); + + let error = compat + .incoming_response( + "initialize", + Ok(serde_json::json!({ "protocolVersion": unsupported })), + ) + .expect_err("proxy initialization must reject a mismatched response version"); + let data = error + .data + .as_ref() + .and_then(|data| data.as_str()) + .unwrap_or_default(); + assert!( + data.contains(&format!( + "required ACP protocol version {selected} but peer negotiated {unsupported}" + )), + "{error:?}" + ); + assert_eq!(negotiated(&compat), selected_kind); + assert_eq!(pending_initialize(&compat), None); + assert_eq!(compat.active_wire_version(), selected_kind); + } + + Ok(()) + } + #[test] #[should_panic(expected = "cannot merge ACP builders with different API protocol versions")] fn merging_different_api_protocol_modes_panics() { let _ = ProtocolMode::v1_agent().merge(ProtocolMode::v2_agent()); } + + #[test] + #[should_panic(expected = "cannot merge standard ACP and proxy ACP builders")] + fn merging_standard_and_proxy_protocol_modes_panics() { + let _ = ProtocolMode::v1_agent().merge(ProtocolMode::v1_proxy()); + } } } diff --git a/src/agent-client-protocol/src/lib.rs b/src/agent-client-protocol/src/lib.rs index 7698eeb1..3cb8edd6 100644 --- a/src/agent-client-protocol/src/lib.rs +++ b/src/agent-client-protocol/src/lib.rs @@ -21,10 +21,10 @@ //! This example uses stable ACP protocol v1. The draft protocol v2 feature //! provides a command-only `V2Session` API and receives updates and interactive //! requests through typed connection handlers because prompt acceptance and -//! inbound traffic are independent. Per-session MCP attachment is available -//! with both v2 and MCP-over-ACP features and remains active for the connection -//! lifetime after successful setup; global proxy attachment and proxy-session -//! helpers remain v1-only. +//! inbound traffic are independent. With both v2 and MCP-over-ACP features, +//! `Proxy.v2()` supports global MCP attachment and `V2SessionBuilder` supports +//! per-session attachment plus non-blocking proxy setup. 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 fc5c3160..b6f92ffd 100644 --- a/src/agent-client-protocol/src/mcp_server/mod.rs +++ b/src/agent-client-protocol/src/mcp_server/mod.rs @@ -4,7 +4,10 @@ //! the core SDK to a particular MCP implementation or async runtime. With the //! `unstable_mcp_over_acp` feature, the same servers can be attached to ACP //! session setup requests through the `with_mcp_server` builder methods. -//! Draft protocol v2 per-session attachment additionally requires +//! Stable protocol v1 and draft protocol v2 both support global proxy +//! attachment and per-session attachment. V2 uses +//! `Proxy.v2().with_mcp_server(...)` or +//! `V2SessionBuilder::with_mcp_server(...)` and 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 fca151ab..36e2d4d7 100644 --- a/src/agent-client-protocol/src/mcp_server/server.rs +++ b/src/agent-client-protocol/src/mcp_server/server.rs @@ -26,6 +26,9 @@ use crate::{ util::MatchDispatchFrom, }; +#[cfg(all(feature = "unstable_mcp_over_acp", feature = "unstable_protocol_v2"))] +use crate::{JsonRpcMessage, UntypedMessage}; + #[cfg(feature = "unstable_mcp_over_acp")] use crate::role::HasPeer; @@ -39,9 +42,10 @@ use crate::schema::v1::ForkSessionRequest; /// connected directly as a standalone MCP component. With the /// `unstable_mcp_over_acp` feature, servers can instead be attached to ACP /// session setup through `Builder::with_mcp_server` or -/// `SessionBuilder::with_mcp_server`. Draft protocol v2 supports per-session -/// attachment through `V2SessionBuilder::with_mcp_server` when -/// `unstable_protocol_v2` is also enabled. +/// `SessionBuilder::with_mcp_server`. When `unstable_protocol_v2` is also +/// enabled, `Proxy.v2().with_mcp_server` attaches a server globally to draft +/// v2 setup requests and `V2SessionBuilder::with_mcp_server` attaches one to a +/// single new v2 session. /// /// # Creating an MCP Server /// @@ -188,22 +192,112 @@ where } } + fn declaration(&self) -> crate::schema::v2::McpServer { + crate::schema::v2::McpServer::Acp(crate::schema::v2::McpServerAcp::new( + self.connect.name(), + crate::schema::v2::McpServerAcpId::from(self.server_id.clone()), + )) + } + + fn append_declaration(&self, request: &mut crate::schema::v2::NewSessionRequest) { + request.mcp_servers.push(self.declaration()); + } + + fn validate_session_setup(request: &UntypedMessage) -> Result { + match request.method() { + "session/new" => { + crate::schema::v2::NewSessionRequest::parse_message( + request.method(), + request.params(), + )?; + Ok(true) + } + "session/resume" => { + crate::schema::v2::ResumeSessionRequest::parse_message( + request.method(), + request.params(), + )?; + Ok(true) + } + #[cfg(feature = "unstable_session_fork")] + "session/fork" => { + crate::schema::v2::ForkSessionRequest::parse_message( + request.method(), + request.params(), + )?; + Ok(true) + } + _ => Ok(false), + } + } + + fn append_declaration_to_raw(&self, request: &mut UntypedMessage) -> Result<(), crate::Error> { + let serde_json::Value::Object(params) = &mut request.params else { + return Err( + crate::Error::invalid_params().data("session setup parameters must be an object") + ); + }; + let declaration = + serde_json::to_value(self.declaration()).map_err(crate::Error::into_internal_error)?; + match params.get_mut("mcpServers") { + Some(serde_json::Value::Array(servers)) => servers.push(declaration), + Some(value) => *value = serde_json::Value::Array(vec![declaration]), + None => { + params.insert( + "mcpServers".to_string(), + serde_json::Value::Array(vec![declaration]), + ); + } + } + Ok(()) + } + /// Attach this server to a draft protocol v2 `session/new` request. pub fn into_dynamic_handler( self, request: &mut crate::schema::v2::NewSessionRequest, cx: &crate::V2ConnectionTo, ) -> Result, crate::Error> { - request.mcp_servers.push(crate::schema::v2::McpServer::Acp( - crate::schema::v2::McpServerAcp::new( - self.connect.name(), - crate::schema::v2::McpServerAcpId::new(self.server_id.0), - ), - )); + self.append_declaration(request); cx.add_dynamic_handler(self.active_session) } } +#[cfg(all(feature = "unstable_mcp_over_acp", feature = "unstable_protocol_v2"))] +impl HandleDispatchFrom for V2McpSessionHandler +where + Counterpart: HasPeer + HasPeer, +{ + async fn handle_dispatch_from( + &mut self, + message: Dispatch, + cx: ConnectionTo, + ) -> Result, crate::Error> { + MatchDispatchFrom::new(message, &cx) + .if_request_from(Client, async |mut request: UntypedMessage, responder| { + if !Self::validate_session_setup(&request)? { + return Ok(Handled::No { + message: (request, responder), + retry: false, + }); + } + + self.append_declaration_to_raw(&mut request)?; + Ok(Handled::No { + message: (request, responder), + retry: false, + }) + }) + .await + .otherwise_delegate(&mut self.active_session) + .await + } + + fn describe_chain(&self) -> impl std::fmt::Debug { + format!("V2McpServer({})", self.connect.name()) + } +} + #[cfg(feature = "unstable_mcp_over_acp")] impl McpSessionHandler where @@ -350,3 +444,167 @@ where .await } } + +#[cfg(all( + test, + feature = "unstable_mcp_over_acp", + feature = "unstable_protocol_v2" +))] +mod tests { + use std::{collections::BTreeMap, path::PathBuf, sync::Arc}; + + use serde::Serialize; + use serde_json::{Value, json}; + + use super::V2McpSessionHandler; + use crate::{ + Conductor, DynConnectTo, Error, UntypedMessage, + mcp_server::{McpConnectionTo, McpServerConnect}, + role, + schema::{ + v1::McpServerAcpId, + v2::{self, McpServer}, + }, + }; + + struct UnusedMcpConnect; + + impl McpServerConnect for UnusedMcpConnect { + fn name(&self) -> String { + "global-v2-server".to_owned() + } + + fn connect(&self, _context: McpConnectionTo) -> DynConnectTo { + panic!("declaration tests must not connect to the MCP server") + } + } + + fn handler() -> V2McpSessionHandler { + V2McpSessionHandler::new( + McpServerAcpId::new("global-v2-server-id"), + Arc::new(UnusedMcpConnect), + ) + } + + fn existing_server() -> McpServer { + McpServer::Other(v2::OtherMcpServer::new( + "_future_transport", + BTreeMap::from([("futureOption".to_owned(), json!({ "nested": true }))]), + )) + } + + fn meta() -> v2::Meta { + v2::Meta::from_iter([("preserved".to_owned(), json!({ "nested": true }))]) + } + + fn assert_raw_append_preserves_params( + handler: &V2McpSessionHandler, + method: &str, + params: impl Serialize, + ) -> Result { + let mut params = serde_json::to_value(params)?; + let Value::Object(params_object) = &mut params else { + panic!("session setup params should serialize as an object"); + }; + params_object.insert( + "_futureSessionField".to_owned(), + json!({ "must": ["remain", "untouched"] }), + ); + + let mut expected = params.clone(); + expected + .get_mut("mcpServers") + .and_then(Value::as_array_mut) + .expect("test request should contain mcpServers") + .push(serde_json::to_value(handler.declaration())?); + + let mut request = UntypedMessage::new(method, params)?; + assert!(V2McpSessionHandler::::validate_session_setup( + &request + )?); + handler.append_declaration_to_raw(&mut request)?; + + assert_eq!( + request.params, expected, + "global attachment must only append its declaration" + ); + + let appended = request + .params + .get("mcpServers") + .and_then(Value::as_array) + .and_then(|servers| servers.last()) + .cloned() + .expect("global declaration should be appended"); + match serde_json::from_value::(appended)? { + McpServer::Acp(server) => { + assert_eq!(server.name, "global-v2-server"); + Ok(server.server_id) + } + server => panic!("expected an ACP server declaration, got {server:?}"), + } + } + + #[test] + fn v2_global_mcp_declaration_preserves_all_session_setup_params() -> Result<(), Error> { + let handler = handler(); + let cwd = PathBuf::from("/tmp/global-v2-mcp"); + let additional_directory = PathBuf::from("/tmp/global-v2-mcp-additional"); + let session_id = v2::SessionId::new("session-to-resume"); + let existing_server = existing_server(); + + let new_server_id = assert_raw_append_preserves_params( + &handler, + "session/new", + v2::NewSessionRequest::new(cwd.clone()) + .additional_directories([additional_directory.clone()]) + .mcp_servers(vec![existing_server.clone()]) + .meta(meta()), + )?; + + let resume_server_id = assert_raw_append_preserves_params( + &handler, + "session/resume", + v2::ResumeSessionRequest::new(session_id.clone(), cwd.clone()) + .additional_directories([additional_directory.clone()]) + .mcp_servers(vec![existing_server.clone()]) + .replay_from(v2::ReplayFrom::Start( + v2::ReplayFromStart::new().meta(meta()), + )) + .meta(meta()), + )?; + assert_eq!(resume_server_id, new_server_id); + + #[cfg(feature = "unstable_session_fork")] + { + let fork_server_id = assert_raw_append_preserves_params( + &handler, + "session/fork", + v2::ForkSessionRequest::new(session_id, cwd) + .additional_directories([additional_directory]) + .mcp_servers(vec![existing_server]) + .meta(meta()), + )?; + assert_eq!(fork_server_id, new_server_id); + } + + Ok(()) + } + + #[test] + fn v2_global_mcp_handler_ignores_non_setup_requests() -> Result<(), Error> { + let request = UntypedMessage::new( + "session/prompt", + json!({ + "sessionId": "session-to-prompt", + "prompt": [] + }), + )?; + + assert!(!V2McpSessionHandler::::validate_session_setup( + &request + )?); + assert_eq!(request.method(), "session/prompt"); + Ok(()) + } +} diff --git a/src/agent-client-protocol/src/role/acp.rs b/src/agent-client-protocol/src/role/acp.rs index 24cef8a6..2466c0c8 100644 --- a/src/agent-client-protocol/src/role/acp.rs +++ b/src/agent-client-protocol/src/role/acp.rs @@ -1098,10 +1098,28 @@ impl Role for Proxy { } impl Proxy { - /// Create a connection builder for a proxy. + /// Create a stable protocol v1 connection builder for a proxy. + /// + /// Use `Proxy::v2` for a protocol-v2-only proxy with typed callbacks and + /// wire validation. Protocol-routing infrastructure that deliberately + /// selects a version itself can disable the guard with + /// `Builder::without_acp_version_guard`. pub fn builder(self) -> Builder { Builder::new(self) } + + /// Create a proxy builder that uses the ACP protocol v2 API. + /// + /// This builder requires `_proxy/initialize` to select protocol v2. + /// Fluent callbacks receive [`crate::V2ConnectionTo`], while + /// low-level custom handlers and runners retain the protocol-neutral + /// [`ConnectionTo`] interface. + /// + /// Requires the `unstable_protocol_v2` crate feature. + #[cfg(feature = "unstable_protocol_v2")] + pub fn v2(self) -> V2Builder { + self.builder().v2_proxy() + } } impl HasPeer for Proxy { diff --git a/src/agent-client-protocol/src/session/v2.rs b/src/agent-client-protocol/src/session/v2.rs index 5aa1ede5..d1f9291a 100644 --- a/src/agent-client-protocol/src/session/v2.rs +++ b/src/agent-client-protocol/src/session/v2.rs @@ -1,4 +1,4 @@ -use std::path::Path; +use std::{future::Future, path::Path}; #[cfg(feature = "unstable_mcp_over_acp")] use futures::{ @@ -7,9 +7,9 @@ use futures::{ }; use crate::{ - Agent, DynamicHandlerGuard, SentRequest, V2ConnectionTo, + Agent, Client, DynamicHandlerGuard, Responder, SentRequest, V2ConnectionTo, jsonrpc::run::{NullRun, RunWithConnectionTo}, - role::HasPeer, + role::{HasPeer, acp::ProxySessionMessages}, schema::v2, }; @@ -128,13 +128,12 @@ where /// Protocol v2 acknowledges `session/prompt` independently from inbound /// session updates. Register typed [`v2::UpdateSessionNotification`] and /// session request handlers on [`crate::Builder`] before connecting, then use -/// [`Self::start_session`] to create the command-only [`V2Session`] handle. +/// [`Self::start_session`] to create the command-only [`V2Session`] handle or +/// `on_proxy_session_start` to forward setup through a proxy. /// /// With both the `unstable_protocol_v2` and `unstable_mcp_over_acp` features, -/// [`Self::with_mcp_server`] attaches an MCP server to the new session. -/// Proxy-session helpers remain available only through the stable protocol v1 -/// [`crate::SessionBuilder`]. -#[must_use = "call `start_session` to send the `session/new` request"] +/// `with_mcp_server` attaches an MCP server to the new session. +#[must_use = "call `start_session` or `on_proxy_session_start` to send the `session/new` request"] #[derive(Debug)] pub struct V2SessionBuilder where @@ -192,19 +191,7 @@ where }) } - /// Send `session/new` and return its independently consumable request. - /// - /// The successful result contains both a cloneable command handle and the - /// complete [`v2::NewSessionResponse`]. Consume the returned request with - /// [`SentRequest::block_task`], [`SentRequest::on_receiving_result`], or - /// another explicit [`SentRequest`] completion mode. - /// - /// 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> + fn send_new_session(self, ordered: bool) -> SentRequest where Run: 'static, { @@ -214,42 +201,41 @@ where dynamic_handler_registrations, run, } = self; - let session_connection = connection.clone(); let raw_connection = connection.raw_connection().clone(); #[cfg(feature = "unstable_mcp_over_acp")] - let sent_request = if dynamic_handler_registrations.is_empty() { - drop(run); - raw_connection.send_request_to(Agent, request) - } else { - let handlers_ready = raw_connection.dynamic_handler_barrier(); - let (runner_started_tx, runner_started_rx) = oneshot::channel(); - let (promotion_tx, promotion_rx) = oneshot::channel(); - let runner_started = match raw_connection.spawn(run_pending_mcp_attachment( - raw_connection.clone(), - run, - runner_started_tx, - promotion_rx, - )) { - Ok(()) => Either::Left(async move { - runner_started_rx.await.map_err(|error| { - crate::util::internal_error(format!( - "MCP runner stopped before its initial poll: {error}" - )) - })? - }), - Err(error) => Either::Right(future::ready(Err(error))), - }; - let readiness = async move { - future::try_join(handlers_ready, runner_started).await?; - Ok(()) - }; - - raw_connection.send_request_to_with_response_hook_after( - Agent, - request, - readiness, - move |_response| { + { + if dynamic_handler_registrations.is_empty() { + drop(run); + if ordered { + raw_connection.send_ordered_request_to(Agent, request) + } else { + raw_connection.send_request_to(Agent, request) + } + } else { + let handlers_ready = raw_connection.dynamic_handler_barrier(); + let (runner_started_tx, runner_started_rx) = oneshot::channel(); + let (promotion_tx, promotion_rx) = oneshot::channel(); + let runner_started = match raw_connection.spawn(run_pending_mcp_attachment( + raw_connection.clone(), + run, + runner_started_tx, + promotion_rx, + )) { + Ok(()) => Either::Left(async move { + runner_started_rx.await.map_err(|error| { + crate::util::internal_error(format!( + "MCP runner stopped before its initial poll: {error}" + )) + })? + }), + Err(error) => Either::Right(future::ready(Err(error))), + }; + let readiness = async move { + future::try_join(handlers_ready, runner_started).await?; + Ok(()) + }; + let response_hook = move |_response: &v2::NewSessionResponse| { promotion_tx.send(()).map_err(|()| { crate::util::internal_error( "MCP runner stopped before session setup completed", @@ -259,18 +245,56 @@ where .into_iter() .for_each(DynamicHandlerGuard::detach); Ok(()) - }, - ) - }; + }; + + if ordered { + raw_connection.send_ordered_request_to_with_response_hook_after( + Agent, + request, + readiness, + response_hook, + ) + } else { + raw_connection.send_request_to_with_response_hook_after( + Agent, + request, + readiness, + response_hook, + ) + } + } + } #[cfg(not(feature = "unstable_mcp_over_acp"))] - let sent_request = { + { drop(dynamic_handler_registrations); drop(run); - raw_connection.send_request_to(Agent, request) - }; + if ordered { + raw_connection.send_ordered_request_to(Agent, request) + } else { + raw_connection.send_request_to(Agent, request) + } + } + } - sent_request.map(move |response| { + /// Send `session/new` and return its independently consumable request. + /// + /// The successful result contains both a cloneable command handle and the + /// complete [`v2::NewSessionResponse`]. Consume the returned request with + /// [`SentRequest::block_task`], [`SentRequest::on_receiving_result`], or + /// another explicit [`SentRequest`] completion mode. + /// + /// 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_new_session(false).map(move |response| { let session = V2Session { session_id: response.session_id.clone(), connection: session_connection, @@ -278,6 +302,54 @@ where Ok(OpenedV2Session { session, response }) }) } + + /// Start 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 installs session routing before + /// later inbound traffic is processed, forwards the complete response, and + /// spawns `op` with an [`OpenedV2Session`] containing the command-only + /// session handle plus the complete setup 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_new_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 }) + }) + } } /// A newly available protocol v2 session and its operation-specific response. diff --git a/src/agent-client-protocol/tests/session_ordering.rs b/src/agent-client-protocol/tests/session_ordering.rs index 9b749ab3..cf67f988 100644 --- a/src/agent-client-protocol/tests/session_ordering.rs +++ b/src/agent-client-protocol/tests/session_ordering.rs @@ -8,7 +8,16 @@ use agent_client_protocol::{ PromptResponse, SessionId, SessionNotification, SessionUpdate, StopReason, TextContent, }, }; -use futures::{StreamExt as _, channel::oneshot}; +use futures::{ + StreamExt as _, + channel::{mpsc, oneshot}, +}; + +#[cfg(feature = "unstable_protocol_v2")] +use agent_client_protocol::{ + JsonRpcMessage, JsonRpcResponse, Proxy, V2ConnectionTo, + schema::{ProtocolVersion, SuccessorMessage, v2}, +}; const TIMEOUT: Duration = Duration::from_secs(10); @@ -46,6 +55,15 @@ mod callback_future_lifetimes { |_session_id| LifetimeTaggedFuture(PhantomData) } + #[cfg(feature = "unstable_protocol_v2")] + fn v2_proxy_session_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, @@ -65,6 +83,18 @@ mod callback_future_lifetimes { .build_session_from(request) .on_proxy_session_start(responder, proxy_session_callback::<'a>()) } + + #[cfg(feature = "unstable_protocol_v2")] + fn v2_on_proxy_session_start_accepts_non_static_callback_future<'a>( + connection: &V2ConnectionTo, + request: v2::NewSessionRequest, + responder: Responder, + _scope: &'a str, + ) -> Result<(), agent_client_protocol::Error> { + connection + .build_session_from(request) + .on_proxy_session_start(responder, v2_proxy_session_callback::<'a>()) + } } #[tokio::test(flavor = "current_thread")] @@ -194,3 +224,172 @@ async fn on_session_start_installs_routing_before_later_batch_entry() { .expect("same-batch session update was not routed") .expect("session connection failed"); } + +#[cfg(feature = "unstable_protocol_v2")] +#[tokio::test(flavor = "current_thread")] +async fn v2_proxy_session_start_installs_routing_before_later_batch_entry() { + let session_id = v2::SessionId::new("same-batch-v2-session"); + let setup_response = v2::NewSessionResponse::new(session_id.clone()).config_options(vec![ + v2::SessionConfigOption::boolean("thinking", "Thinking", true), + ]); + let callback_response = setup_response.clone(); + let response_session_id = session_id.clone(); + let notification_session_id = 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-proxy", env!("CARGO_PKG_VERSION")), + )) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request_from( + Client, + async move |request: v2::NewSessionRequest, + responder, + connection: V2ConnectionTo| { + let callback_response = callback_response.clone(); + let response_session_id = response_session_id.clone(); + let callback_tx = callback_tx.clone(); + connection + .build_session_from(request) + .on_proxy_session_start(responder, move |opened| async move { + assert_eq!(opened.session().session_id(), &response_session_id); + assert_eq!(opened.response(), &callback_response); + callback_tx.unbounded_send(()).map_err(|_| { + agent_client_protocol::Error::internal_error() + .data("v2 proxy 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 proxy 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 { + let initialize_id = agent_client_protocol::schema::v1::RequestId::Number(1); + let initialize = v2::InitializeProxyRequest::new(v2::InitializeRequest::new( + ProtocolVersion::V2, + v2::Implementation::new("same-batch-client", env!("CARGO_PKG_VERSION")), + )); + peer.tx + .unbounded_send(TransportFrame::Single(RawJsonRpcMessage::request( + "_proxy/initialize".to_owned(), + serde_json::to_value(initialize).expect("initialize request should serialize"), + initialize_id.clone(), + )?)) + .expect("proxy should accept initialization"); + + let Some(TransportFrame::Single(RawJsonRpcMessage::Response( + agent_client_protocol::schema::v1::Response::Result { id, result }, + ))) = peer.rx.next().await + else { + panic!("expected the proxy initialize response"); + }; + assert_eq!(id, initialize_id); + let initialize_response = v2::InitializeResponse::from_value("_proxy/initialize", result)?; + assert_eq!(initialize_response.protocol_version, ProtocolVersion::V2); + + let upstream_id = agent_client_protocol::schema::v1::RequestId::Number(2); + peer.tx + .unbounded_send(TransportFrame::Single(RawJsonRpcMessage::request( + "session/new".to_owned(), + serde_json::to_value(v2::NewSessionRequest::new("/same-batch-v2-session")) + .expect("session request should serialize"), + upstream_id.clone(), + )?)) + .expect("proxy should accept session/new"); + + let Some(TransportFrame::Single(RawJsonRpcMessage::Request(forwarded))) = + peer.rx.next().await + else { + panic!("expected a forwarded session/new request"); + }; + let successor = SuccessorMessage::::parse_message( + forwarded.method.as_ref(), + &forwarded.params, + )?; + assert_eq!( + successor.message.cwd, + v2::AbsolutePath::new("/same-batch-v2-session") + ); + + let response = RawJsonRpcMessage::response( + forwarded.id, + Ok(serde_json::to_value(setup_response).expect("session 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 session response and update"); + }; + match message { + RawJsonRpcMessage::Response( + agent_client_protocol::schema::v1::Response::Result { id, result }, + ) => { + assert_eq!(id, upstream_id); + let response = v2::NewSessionResponse::from_value("session/new", result)?; + assert_eq!(response.session_id, session_id); + saw_response = true; + } + RawJsonRpcMessage::Notification(notification) => { + let update = v2::UpdateSessionNotification::parse_message( + notification.method.as_ref(), + ¬ification.params, + )?; + assert_eq!(update.session_id, 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 session update was not routed") + .expect("v2 proxy session connection failed"); +}