From 1139156961d940fbb9e9c607a9637d9fd4335f9c Mon Sep 17 00:00:00 2001 From: Ben Brandt Date: Thu, 9 Jul 2026 13:36:13 +0200 Subject: [PATCH 1/6] feat(unstable-v2): Add protocol router for agents --- md/protocol-v2.md | 59 +++- src/agent-client-protocol/src/lib.rs | 2 + src/agent-client-protocol/src/role/acp.rs | 271 +++++++++++++++ .../tests/protocol_v2.rs | 316 +++++++++++++++++- 4 files changed, 642 insertions(+), 6 deletions(-) diff --git a/md/protocol-v2.md b/md/protocol-v2.md index 7772e47..3beef37 100644 --- a/md/protocol-v2.md +++ b/md/protocol-v2.md @@ -78,6 +78,59 @@ The SDK handles the `initialize` negotiation at the JSON-RPC boundary: - After initialization, the local API version and negotiated wire version must match. The SDK does not convert traffic between v1 and v2. -That means v1 and v2 implementations need separate handlers today. A future SDK -API can route traffic to a registered implementation for the negotiated protocol -version, but `Agent.v2()` and `Client.v2()` are currently v2-only. +That means v1 and v2 implementations still need separate handlers. +`Agent.v2()` and `Client.v2()` are v2-only. While protocol v2 stabilizes, the +`unstable_protocol_v2` crate feature also exposes `Agent.protocol_router()`, +which can route each connection to the configured implementation compatible +with the client's requested protocol version. + +Agents can add protocol implementations independently, which makes it easy for +applications built with v2 support to control v2 rollout with a runtime feature +flag: + +```rust +use agent_client_protocol::schema::{v1, v2}; +use agent_client_protocol::{Agent, ConnectTo}; + +# fn implementation() -> v2::Implementation { +# v2::Implementation::new("example", "0.1.0") +# } +# async fn serve(client_transport: impl agent_client_protocol::ConnectTo) -> agent_client_protocol::Result<()> { +# let enable_protocol_v2 = true; +let v1_agent = Agent.builder().on_receive_request( + async |initialize: v1::InitializeRequest, responder, _cx| { + responder.respond(v1::InitializeResponse::new(initialize.protocol_version)) + }, + agent_client_protocol::on_receive_request!(), +); + +let agent = Agent.protocol_router().with_v1(v1_agent); + +let agent = if enable_protocol_v2 { + let v2_agent = Agent.v2().on_receive_request( + async |initialize: v2::InitializeRequest, responder, _cx| { + responder.respond(v2::InitializeResponse::new( + initialize.protocol_version, + implementation(), + )) + }, + agent_client_protocol::on_receive_request!(), + ); + + agent.with_v2(v2_agent) +} else { + agent +}; + +agent + .connect_to(client_transport) + .await?; +# Ok(()) +# } +``` + +The protocol router reads the initial `initialize` request, selects the +highest configured protocol version that is compatible with the requested +version, and then hands the connection to that implementation. If only v2 is +configured, v1 clients are rejected without changing the fluent API. The router +does not convert messages between v1 and v2 after routing. diff --git a/src/agent-client-protocol/src/lib.rs b/src/agent-client-protocol/src/lib.rs index 6cb2c7f..f8d8999 100644 --- a/src/agent-client-protocol/src/lib.rs +++ b/src/agent-client-protocol/src/lib.rs @@ -103,6 +103,8 @@ pub use jsonrpc::{ }; pub use jsonrpc::{RequestCancellation, is_cancel_request_notification}; +#[cfg(feature = "unstable_protocol_v2")] +pub use role::acp::AgentProtocolRouter; pub use role::{ Role, RoleId, UntypedRole, acp::{Agent, Client, Conductor, Proxy}, diff --git a/src/agent-client-protocol/src/role/acp.rs b/src/agent-client-protocol/src/role/acp.rs index d5919f8..39abbe6 100644 --- a/src/agent-client-protocol/src/role/acp.rs +++ b/src/agent-client-protocol/src/role/acp.rs @@ -1,10 +1,21 @@ use std::{fmt::Debug, hash::Hash}; +#[cfg(feature = "unstable_protocol_v2")] +use futures::{StreamExt as _, future}; + +#[cfg(feature = "unstable_protocol_v2")] +use crate::DynConnectTo; use crate::jsonrpc::{Builder, handlers::NullHandler, run::NullRun}; use crate::role::{HasPeer, RemoteStyle}; +#[cfg(feature = "unstable_protocol_v2")] +use crate::schema::ProtocolVersion; +#[cfg(feature = "unstable_protocol_v2")] +use crate::schema::v1::RequestId; use crate::schema::v1::{InitializeRequest, NewSessionRequest, NewSessionResponse, SessionId}; use crate::schema::{InitializeProxyRequest, METHOD_INITIALIZE_PROXY}; use crate::util::MatchDispatchFrom; +#[cfg(feature = "unstable_protocol_v2")] +use crate::{Channel, RawJsonRpcMessage, RawJsonRpcParams}; use crate::{ConnectTo, ConnectionTo, Dispatch, HandleDispatchFrom, Handled, Role, RoleId}; /// The client role - typically an IDE or CLI that controls an agent. @@ -137,6 +148,266 @@ impl Agent { pub fn v2(self) -> Builder { self.builder().v2_agent() } + + /// Create a router that chooses between configured protocol implementations. + /// + /// Add implementations with [`AgentProtocolRouter::with_v1`] and + /// [`AgentProtocolRouter::with_v2`]. + /// The resulting router reads the initial + /// `initialize` request, selects the highest configured implementation + /// compatible with the client's requested protocol version, then forwards + /// the connection to that implementation. It does not convert traffic + /// between protocol versions after routing. + /// + /// Requires the `unstable_protocol_v2` crate feature while protocol v2 + /// stabilizes. + #[cfg(feature = "unstable_protocol_v2")] + #[must_use] + pub fn protocol_router(self) -> AgentProtocolRouter { + AgentProtocolRouter::new() + } +} + +/// Agent component that routes each connection to a configured protocol implementation. +/// +/// Use [`Agent::protocol_router`] to start the builder, then add each supported +/// protocol version independently. The selected implementation owns the +/// connection after the initial `initialize` negotiation. +#[cfg(feature = "unstable_protocol_v2")] +#[derive(Debug, Default)] +pub struct AgentProtocolRouter { + v1: Option>, + v2: Option>, +} + +#[cfg(feature = "unstable_protocol_v2")] +impl AgentProtocolRouter { + /// Create an empty agent protocol router. + #[must_use] + pub fn new() -> Self { + Self::default() + } + + /// Return this router with an ACP v1 implementation configured. + #[must_use] + pub fn with_v1(mut self, agent: impl ConnectTo) -> Self { + self.v1 = Some(DynConnectTo::new(agent)); + self + } + + /// Return this router with an ACP v2 implementation configured. + #[must_use] + pub fn with_v2(mut self, agent: impl ConnectTo) -> Self { + self.v2 = Some(DynConnectTo::new(agent)); + self + } +} + +#[cfg(feature = "unstable_protocol_v2")] +impl ConnectTo for AgentProtocolRouter { + async fn connect_to(self, client: impl ConnectTo) -> Result<(), crate::Error> { + let (mut client_channel, client_future) = client.into_channel_and_future(); + let (first_message, client_future): ( + Result, + crate::BoxFuture<'static, Result<(), crate::Error>>, + ) = match future::select(Box::pin(client_channel.rx.next()), client_future).await { + future::Either::Left((Some(first_message), client_future)) => { + (first_message, client_future) + } + future::Either::Left((None, client_future)) => return client_future.await, + future::Either::Right((result, first_message)) => { + result?; + drop(first_message); + let Some(first_message) = client_channel.rx.next().await else { + return Ok(()); + }; + (first_message, Box::pin(future::ready(Ok(())))) + } + }; + + let mut first_message = first_message?; + let supported = SupportedAgentProtocols { + v1: self.v1.is_some(), + v2: self.v2.is_some(), + }; + let selected = match select_agent_protocol(&mut first_message, supported) { + Ok(selected) => selected, + Err(error) => { + send_initialize_error(&client_channel.tx, &first_message, error.clone())?; + return Err(error); + } + }; + let Some(agent) = selected.take_agent(self) else { + let error = selected.unsupported_error(supported); + send_initialize_error(&client_channel.tx, &first_message, error.clone())?; + return Err(error); + }; + + let (router_channel, agent_channel) = Channel::duplex(); + let Channel { + rx: from_agent, + tx: to_agent, + } = router_channel; + + to_agent + .unbounded_send(Ok(first_message)) + .map_err(crate::util::internal_error)?; + + let agent_future = Box::pin(agent.connect_to(agent_channel)); + + let ((), (), (), ()) = futures::try_join!( + client_future, + agent_future, + Channel { + rx: client_channel.rx, + tx: to_agent, + } + .copy(), + Channel { + rx: from_agent, + tx: client_channel.tx, + } + .copy(), + )?; + + Ok(()) + } +} + +#[cfg(feature = "unstable_protocol_v2")] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum AgentProtocol { + V1, + V2, +} + +#[cfg(feature = "unstable_protocol_v2")] +impl AgentProtocol { + fn version(self) -> ProtocolVersion { + match self { + Self::V1 => ProtocolVersion::V1, + Self::V2 => ProtocolVersion::V2, + } + } + + fn take_agent(self, agent: AgentProtocolRouter) -> Option> { + match self { + Self::V1 => agent.v1, + Self::V2 => agent.v2, + } + } + + fn name(self) -> &'static str { + match self { + Self::V1 => "1", + Self::V2 => "2", + } + } + + fn unsupported_error(self, supported: SupportedAgentProtocols) -> crate::Error { + crate::Error::invalid_request().data(format!( + "ACP protocol version {} is not configured; this endpoint supports {}", + self.name(), + supported.description() + )) + } +} + +#[cfg(feature = "unstable_protocol_v2")] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct SupportedAgentProtocols { + v1: bool, + v2: bool, +} + +#[cfg(feature = "unstable_protocol_v2")] +impl SupportedAgentProtocols { + fn highest_compatible(self, requested: ProtocolVersion) -> Option { + if self.v2 && requested >= ProtocolVersion::V2 { + return Some(AgentProtocol::V2); + } + + if self.v1 && requested >= ProtocolVersion::V1 { + return Some(AgentProtocol::V1); + } + + None + } + + fn description(self) -> String { + match (self.v1, self.v2) { + (true, true) => "ACP protocol versions 1 and 2".into(), + (true, false) => "ACP protocol version 1".into(), + (false, true) => "ACP protocol version 2".into(), + (false, false) => "no ACP protocol versions".into(), + } + } +} + +#[cfg(feature = "unstable_protocol_v2")] +fn select_agent_protocol( + message: &mut RawJsonRpcMessage, + supported: SupportedAgentProtocols, +) -> Result { + let RawJsonRpcMessage::Request(request) = message else { + return Err( + crate::Error::invalid_request().data("first ACP message must be an initialize request") + ); + }; + + if request.method.as_ref() != "initialize" { + return Err(crate::Error::invalid_request().data("first ACP request must be initialize")); + } + + let Some(RawJsonRpcParams::Object(params)) = &mut request.params else { + return Err(invalid_initialize_protocol_version()); + }; + let Some(protocol_version) = params.get("protocolVersion") else { + return Err(invalid_initialize_protocol_version()); + }; + + let requested = serde_json::from_value::(protocol_version.clone()) + .map_err(|_| invalid_initialize_protocol_version())?; + let selected = highest_compatible_agent_protocol(requested, supported)?; + params.insert( + "protocolVersion".into(), + serde_json::to_value(selected.version()).map_err(crate::Error::into_internal_error)?, + ); + + Ok(selected) +} + +#[cfg(feature = "unstable_protocol_v2")] +fn highest_compatible_agent_protocol( + requested: ProtocolVersion, + supported: SupportedAgentProtocols, +) -> Result { + supported.highest_compatible(requested).ok_or_else(|| { + crate::Error::invalid_request().data(format!( + "unsupported ACP protocol version {requested}; this endpoint supports {}", + supported.description() + )) + }) +} + +#[cfg(feature = "unstable_protocol_v2")] +fn invalid_initialize_protocol_version() -> crate::Error { + crate::Error::invalid_params() + .data("initialize.protocolVersion must be a valid ACP protocol version") +} + +#[cfg(feature = "unstable_protocol_v2")] +fn send_initialize_error( + tx: &futures::channel::mpsc::UnboundedSender>, + message: &RawJsonRpcMessage, + error: crate::Error, +) -> Result<(), crate::Error> { + let id = match message { + RawJsonRpcMessage::Request(request) => request.id.clone(), + RawJsonRpcMessage::Notification(_) | RawJsonRpcMessage::Response(_) => RequestId::Null, + }; + tx.unbounded_send(Ok(RawJsonRpcMessage::response(id, Err(error)))) + .map_err(crate::util::internal_error) } impl HasPeer for Agent { diff --git a/src/agent-client-protocol/tests/protocol_v2.rs b/src/agent-client-protocol/tests/protocol_v2.rs index ac660b8..6c61d7e 100644 --- a/src/agent-client-protocol/tests/protocol_v2.rs +++ b/src/agent-client-protocol/tests/protocol_v2.rs @@ -4,8 +4,8 @@ use std::path::PathBuf; use agent_client_protocol::schema::{ProtocolVersion, v1, v2}; use agent_client_protocol::{ - Agent, Builder, Client, ConnectTo, Error, JsonRpcMessage, JsonRpcRequest, JsonRpcResponse, - NullHandler, RawJsonRpcMessage, Role, UntypedRole, + Agent, AgentProtocolRouter, Builder, Client, ConnectTo, Error, JsonRpcMessage, JsonRpcRequest, + JsonRpcResponse, NullHandler, RawJsonRpcMessage, Role, UntypedRole, }; use agent_client_protocol_test::testy::Testy; use futures::StreamExt as _; @@ -72,11 +72,38 @@ fn v2_initialize_response_with_session( .capabilities(v2::AgentCapabilities::new().session(v2::SessionCapabilities::new())) } -#[cfg(feature = "unstable_mcp_over_acp")] fn json_value(value: impl Serialize) -> Result { serde_json::to_value(value).map_err(Error::into_internal_error) } +fn runtime_flag_protocol_router(enable_protocol_v2: bool) -> AgentProtocolRouter { + let v1_agent = Agent.builder().on_receive_request( + async |initialize: v1::InitializeRequest, responder, _cx| { + assert_eq!(initialize.protocol_version, ProtocolVersion::V1); + responder.respond(v1::InitializeResponse::new(initialize.protocol_version)) + }, + agent_client_protocol::on_receive_request!(), + ); + + let agent = Agent.protocol_router().with_v1(v1_agent); + + if enable_protocol_v2 { + let v2_agent = Agent.v2().on_receive_request( + async |initialize: v2::InitializeRequest, responder, _cx| { + assert_eq!(initialize.protocol_version, ProtocolVersion::V2); + responder.respond(v2_initialize_response_with_session( + initialize.protocol_version, + )) + }, + agent_client_protocol::on_receive_request!(), + ); + + agent.with_v2(v2_agent) + } else { + agent + } +} + async fn assert_malformed_initialize_rejected(params: Map) -> Result<(), Error> { let agent = Agent.v2().on_receive_request( async |_initialize: v2::InitializeRequest, responder, _cx| { @@ -596,6 +623,289 @@ async fn v2_client_and_agent_negotiate_v2() -> Result<(), Error> { .await } +#[tokio::test(flavor = "current_thread")] +async fn protocol_router_routes_v1_client_to_v1_implementation() -> Result<(), Error> { + let agent = Agent + .protocol_router() + .with_v1( + Agent + .builder() + .on_receive_request( + async |initialize: v1::InitializeRequest, responder, _cx| { + assert_eq!(initialize.protocol_version, ProtocolVersion::V1); + responder.respond(v1::InitializeResponse::new(initialize.protocol_version)) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async |request: v1::NewSessionRequest, responder, _cx| { + assert!(request.cwd.is_absolute()); + responder.respond(v1::NewSessionResponse::new(v1::SessionId::new( + "v1-protocol-router-session", + ))) + }, + agent_client_protocol::on_receive_request!(), + ), + ) + .with_v2(Agent.v2().on_receive_request( + async |_initialize: v2::InitializeRequest, responder, _cx| { + responder.respond_with_internal_error("v2 implementation should not run") + }, + agent_client_protocol::on_receive_request!(), + )); + + Client + .builder() + .connect_with(agent, async |cx| { + let initialize = cx + .send_request(v1_initialize_request(ProtocolVersion::V1)) + .block_task() + .await?; + assert_eq!(initialize.protocol_version, ProtocolVersion::V1); + + let session = cx + .send_request(v1::NewSessionRequest::new(cwd()?)) + .block_task() + .await?; + assert_eq!(session.session_id.0.as_ref(), "v1-protocol-router-session"); + Ok(()) + }) + .await +} + +#[tokio::test(flavor = "current_thread")] +async fn protocol_router_routes_v2_client_to_v2_implementation() -> Result<(), Error> { + let agent = Agent + .protocol_router() + .with_v1(Agent.builder().on_receive_request( + async |_initialize: v1::InitializeRequest, responder, _cx| { + responder.respond_with_internal_error("v1 implementation should not run") + }, + agent_client_protocol::on_receive_request!(), + )) + .with_v2( + Agent + .v2() + .on_receive_request( + async |initialize: v2::InitializeRequest, responder, _cx| { + assert_eq!(initialize.protocol_version, ProtocolVersion::V2); + responder.respond(v2_initialize_response_with_session( + initialize.protocol_version, + )) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async |request: v2::NewSessionRequest, responder, _cx| { + assert!(request.cwd.is_absolute()); + responder.respond(v2::NewSessionResponse::new(v2::SessionId::new( + "v2-protocol-router-session", + ))) + }, + agent_client_protocol::on_receive_request!(), + ), + ); + + Client + .v2() + .connect_with(agent, async |cx| { + let initialize = cx + .send_request(v2_initialize_request(ProtocolVersion::V1)) + .block_task() + .await?; + assert_eq!(initialize.protocol_version, ProtocolVersion::V2); + + let session = cx + .send_request(v2::NewSessionRequest::new(cwd()?)) + .block_task() + .await?; + assert_eq!(session.session_id.0.as_ref(), "v2-protocol-router-session"); + Ok(()) + }) + .await +} + +#[tokio::test(flavor = "current_thread")] +async fn protocol_router_can_route_only_v1() -> Result<(), Error> { + let agent = Agent + .protocol_router() + .with_v1(Agent.builder().on_receive_request( + async |initialize: v1::InitializeRequest, responder, _cx| { + assert_eq!(initialize.protocol_version, ProtocolVersion::V1); + responder.respond(v1::InitializeResponse::new(initialize.protocol_version)) + }, + agent_client_protocol::on_receive_request!(), + )); + + Client + .builder() + .connect_with(agent, async |cx| { + let initialize = cx + .send_request(v1_initialize_request(ProtocolVersion::V1)) + .block_task() + .await?; + assert_eq!(initialize.protocol_version, ProtocolVersion::V1); + Ok(()) + }) + .await +} + +#[tokio::test(flavor = "current_thread")] +async fn protocol_router_can_route_only_v2() -> Result<(), Error> { + let agent = Agent + .protocol_router() + .with_v2(Agent.v2().on_receive_request( + async |initialize: v2::InitializeRequest, responder, _cx| { + assert_eq!(initialize.protocol_version, ProtocolVersion::V2); + responder.respond(v2_initialize_response_with_session( + initialize.protocol_version, + )) + }, + agent_client_protocol::on_receive_request!(), + )); + + Client + .v2() + .connect_with(agent, async |cx| { + let initialize = cx + .send_request(v2_initialize_request(ProtocolVersion::V1)) + .block_task() + .await?; + assert_eq!(initialize.protocol_version, ProtocolVersion::V2); + Ok(()) + }) + .await +} + +#[tokio::test(flavor = "current_thread")] +async fn protocol_router_supports_runtime_v2_registration_flag() -> Result<(), Error> { + Client + .v2() + .connect_with(runtime_flag_protocol_router(false), async |cx| { + let error = cx + .send_request(v2_initialize_request(ProtocolVersion::V2)) + .block_task() + .await + .expect_err("runtime-disabled v2 should route to v1 and fail v2 negotiation"); + let data = error + .data + .as_ref() + .and_then(|data| data.as_str()) + .unwrap_or_default(); + assert!(data.contains("peer negotiated 1"), "{error:?}"); + Ok(()) + }) + .await?; + + Client + .v2() + .connect_with(runtime_flag_protocol_router(true), async |cx| { + let initialize = cx + .send_request(v2_initialize_request(ProtocolVersion::V1)) + .block_task() + .await?; + assert_eq!(initialize.protocol_version, ProtocolVersion::V2); + Ok(()) + }) + .await +} + +#[tokio::test(flavor = "current_thread")] +async fn protocol_router_v2_only_rejects_v1_client() -> Result<(), Error> { + let agent = Agent + .protocol_router() + .with_v2(Agent.v2().on_receive_request( + async |_initialize: v2::InitializeRequest, responder, _cx| { + responder.respond_with_internal_error("v2 implementation should not run") + }, + agent_client_protocol::on_receive_request!(), + )); + + let (mut channel, agent_future) = ConnectTo::::into_channel_and_future(agent); + let agent_task = tokio::spawn(agent_future); + + channel + .tx + .unbounded_send(Ok(RawJsonRpcMessage::request( + "initialize".into(), + json_value(v1_initialize_request(ProtocolVersion::V1))?, + v1::RequestId::Number(1), + )?)) + .map_err(Error::into_internal_error)?; + + while let Some(message) = channel.rx.next().await { + let message = message?; + let RawJsonRpcMessage::Response(v1::Response::Error { error, .. }) = message else { + continue; + }; + let data = error + .data + .as_ref() + .and_then(|data| data.as_str()) + .unwrap_or_default(); + assert!( + data.contains("supports ACP protocol version 2"), + "{error:?}" + ); + agent_task.abort(); + return Ok(()); + } + + agent_task.abort(); + Err(agent_client_protocol::util::internal_error( + "protocol router did not reject v1 initialize", + )) +} + +#[tokio::test(flavor = "current_thread")] +async fn protocol_router_routes_future_protocol_version_to_v2() -> Result<(), Error> { + let agent = Agent + .protocol_router() + .with_v1(Agent.builder().on_receive_request( + async |_initialize: v1::InitializeRequest, responder, _cx| { + responder.respond_with_internal_error("v1 implementation should not run") + }, + agent_client_protocol::on_receive_request!(), + )) + .with_v2(Agent.v2().on_receive_request( + async |initialize: v2::InitializeRequest, responder, _cx| { + assert_eq!(initialize.protocol_version, ProtocolVersion::V2); + responder.respond(v2_initialize_response_with_session( + initialize.protocol_version, + )) + }, + agent_client_protocol::on_receive_request!(), + )); + + let (mut channel, agent_future) = ConnectTo::::into_channel_and_future(agent); + let agent_task = tokio::spawn(agent_future); + + channel + .tx + .unbounded_send(Ok(RawJsonRpcMessage::request( + "initialize".into(), + json_value(v2_initialize_request(ProtocolVersion::from(3_u16)))?, + v1::RequestId::Number(1), + )?)) + .map_err(Error::into_internal_error)?; + + while let Some(message) = channel.rx.next().await { + let message = message?; + let RawJsonRpcMessage::Response(v1::Response::Result { result, .. }) = message else { + continue; + }; + let initialize = v2::InitializeResponse::from_value("initialize", result)?; + assert_eq!(initialize.protocol_version, ProtocolVersion::V2); + agent_task.abort(); + return Ok(()); + } + + agent_task.abort(); + Err(agent_client_protocol::util::internal_error( + "protocol router did not respond to initialize", + )) +} + /// A v2 agent whose `session/new` handler only responds once the peer cancels /// the request via `$/cancel_request`. fn v2_agent_with_cancellable_new_session() From e2255965bbd5b78b9cf1ad3e4d9d19e20860f524 Mon Sep 17 00:00:00 2001 From: Ben Brandt Date: Thu, 9 Jul 2026 14:22:14 +0200 Subject: [PATCH 2/6] feat(unstable-v2): Add client protocol router --- md/protocol-v2.md | 40 +- src/agent-client-protocol/src/lib.rs | 2 + src/agent-client-protocol/src/role/acp.rs | 404 +++++++++++++++--- .../tests/protocol_v2.rs | 209 ++++++++- 4 files changed, 599 insertions(+), 56 deletions(-) diff --git a/md/protocol-v2.md b/md/protocol-v2.md index 3beef37..25af4af 100644 --- a/md/protocol-v2.md +++ b/md/protocol-v2.md @@ -80,9 +80,8 @@ The SDK handles the `initialize` negotiation at the JSON-RPC boundary: That means v1 and v2 implementations still need separate handlers. `Agent.v2()` and `Client.v2()` are v2-only. While protocol v2 stabilizes, the -`unstable_protocol_v2` crate feature also exposes `Agent.protocol_router()`, -which can route each connection to the configured implementation compatible -with the client's requested protocol version. +`unstable_protocol_v2` crate feature also exposes `Agent.protocol_router()` and +`Client.protocol_router()` for composing version-specific implementations. Agents can add protocol implementations independently, which makes it easy for applications built with v2 support to control v2 rollout with a runtime feature @@ -134,3 +133,38 @@ highest configured protocol version that is compatible with the requested version, and then hands the connection to that implementation. If only v2 is configured, v1 clients are rejected without changing the fluent API. The router does not convert messages between v1 and v2 after routing. + +Clients use the same fluent shape: + +```rust +use agent_client_protocol::{Client, ConnectTo}; + +# fn v1_client() -> impl agent_client_protocol::ConnectTo { +# Client.builder() +# } +# fn v2_client() -> impl agent_client_protocol::ConnectTo { +# Client.v2() +# } +# async fn run(agent_transport: impl agent_client_protocol::ConnectTo) -> agent_client_protocol::Result<()> { +# let enable_protocol_v2 = true; +let client = Client.protocol_router().with_v1(v1_client()); + +let client = if enable_protocol_v2 { + client.with_v2(v2_client()) +} else { + client +}; + +client + .connect_to(agent_transport) + .await?; +# Ok(()) +# } +``` + +The client router starts the highest configured implementation. If v2 +initialization negotiates v1 and a v1 implementation is configured, it switches +to the local v1 client implementation and forwards the original initialize +response. It does not send a second initialize request on the same connection. +If v2 initialization is rejected, that error is surfaced to the v2 client +implementation. diff --git a/src/agent-client-protocol/src/lib.rs b/src/agent-client-protocol/src/lib.rs index f8d8999..710401e 100644 --- a/src/agent-client-protocol/src/lib.rs +++ b/src/agent-client-protocol/src/lib.rs @@ -105,6 +105,8 @@ pub use jsonrpc::{RequestCancellation, is_cancel_request_notification}; #[cfg(feature = "unstable_protocol_v2")] pub use role::acp::AgentProtocolRouter; +#[cfg(feature = "unstable_protocol_v2")] +pub use role::acp::ClientProtocolRouter; pub use role::{ Role, RoleId, UntypedRole, acp::{Agent, Client, Conductor, Proxy}, diff --git a/src/agent-client-protocol/src/role/acp.rs b/src/agent-client-protocol/src/role/acp.rs index 39abbe6..9d09ed9 100644 --- a/src/agent-client-protocol/src/role/acp.rs +++ b/src/agent-client-protocol/src/role/acp.rs @@ -9,9 +9,9 @@ use crate::jsonrpc::{Builder, handlers::NullHandler, run::NullRun}; use crate::role::{HasPeer, RemoteStyle}; #[cfg(feature = "unstable_protocol_v2")] use crate::schema::ProtocolVersion; -#[cfg(feature = "unstable_protocol_v2")] -use crate::schema::v1::RequestId; use crate::schema::v1::{InitializeRequest, NewSessionRequest, NewSessionResponse, SessionId}; +#[cfg(feature = "unstable_protocol_v2")] +use crate::schema::v1::{RequestId, Response as RpcResponse}; use crate::schema::{InitializeProxyRequest, METHOD_INITIALIZE_PROXY}; use crate::util::MatchDispatchFrom; #[cfg(feature = "unstable_protocol_v2")] @@ -69,6 +69,23 @@ impl Client { self.builder().v2_client() } + /// Create a router that chooses between configured protocol implementations. + /// + /// Add implementations with [`ClientProtocolRouter::with_v1`] and + /// [`ClientProtocolRouter::with_v2`]. The resulting router starts the + /// highest configured protocol implementation, observes the initial + /// `initialize` response, and falls back to a configured v1 implementation + /// when v2 initialization negotiates v1. It does not convert traffic + /// between protocol versions after routing. + /// + /// Requires the `unstable_protocol_v2` crate feature while protocol v2 + /// stabilizes. + #[cfg(feature = "unstable_protocol_v2")] + #[must_use] + pub fn protocol_router(self) -> ClientProtocolRouter { + ClientProtocolRouter::new() + } + /// Connect to `agent` and run `main_fn` with the [`ConnectionTo`]. /// Returns the result of `main_fn` (or an error if something goes wrong). /// @@ -82,6 +99,81 @@ impl Client { } } +/// Client component that routes each connection to a configured protocol implementation. +/// +/// Use [`Client::protocol_router`] to start the builder, then add each supported +/// protocol version independently. The selected implementation owns the +/// connection after the initial `initialize` negotiation. +#[cfg(feature = "unstable_protocol_v2")] +#[derive(Debug, Default)] +pub struct ClientProtocolRouter { + v1: Option>, + v2: Option>, +} + +#[cfg(feature = "unstable_protocol_v2")] +impl ClientProtocolRouter { + /// Create an empty client protocol router. + #[must_use] + pub fn new() -> Self { + Self::default() + } + + /// Return this router with an ACP v1 implementation configured. + #[must_use] + pub fn with_v1(mut self, client: impl ConnectTo) -> Self { + self.v1 = Some(DynConnectTo::new(client)); + self + } + + /// Return this router with an ACP v2 implementation configured. + #[must_use] + pub fn with_v2(mut self, client: impl ConnectTo) -> Self { + self.v2 = Some(DynConnectTo::new(client)); + self + } +} + +#[cfg(feature = "unstable_protocol_v2")] +impl ConnectTo for ClientProtocolRouter { + async fn connect_to(mut self, agent: impl ConnectTo) -> Result<(), crate::Error> { + let supported = SupportedClientProtocols { + v1: self.v1.is_some(), + v2: self.v2.is_some(), + }; + let Some(selected) = supported.highest_configured() else { + return Err(crate::Error::invalid_request() + .data("client protocol router has no configured ACP protocol implementations")); + }; + + let agent = RunningProtocolPeer::new(agent); + let client = selected + .take_client(&mut self) + .expect("selected protocol is configured"); + let (client, initialize) = start_client_protocol(selected, client).await?; + let (client, agent, initialize_response) = + send_initialize_and_receive(client, agent, initialize).await?; + + if selected == ClientProtocol::V2 + && let Some(client) = self.v1.take() + && initialize_response_negotiated_v1(&initialize_response) + { + let (fallback_client, fallback_initialize) = + start_client_protocol(ClientProtocol::V1, client).await?; + + let fallback_response = initialize_response.with_id( + initialize_request_id(&fallback_initialize) + .expect("validated initialize request has an id"), + ); + fallback_client.send(fallback_response)?; + return pipe_protocol_peers_until_done(fallback_client, agent).await; + } + + client.send(initialize_response.into_message())?; + pipe_protocol_peers_until_done(client, agent).await + } +} + impl HasPeer for Client { fn remote_style(&self, _peer: Client) -> RemoteStyle { RemoteStyle::Counterpart @@ -206,26 +298,10 @@ impl AgentProtocolRouter { #[cfg(feature = "unstable_protocol_v2")] impl ConnectTo for AgentProtocolRouter { async fn connect_to(self, client: impl ConnectTo) -> Result<(), crate::Error> { - let (mut client_channel, client_future) = client.into_channel_and_future(); - let (first_message, client_future): ( - Result, - crate::BoxFuture<'static, Result<(), crate::Error>>, - ) = match future::select(Box::pin(client_channel.rx.next()), client_future).await { - future::Either::Left((Some(first_message), client_future)) => { - (first_message, client_future) - } - future::Either::Left((None, client_future)) => return client_future.await, - future::Either::Right((result, first_message)) => { - result?; - drop(first_message); - let Some(first_message) = client_channel.rx.next().await else { - return Ok(()); - }; - (first_message, Box::pin(future::ready(Ok(())))) - } + let client = RunningProtocolPeer::new(client); + let Some((mut first_message, client)) = client.next_message().await? else { + return Ok(()); }; - - let mut first_message = first_message?; let supported = SupportedAgentProtocols { v1: self.v1.is_some(), v2: self.v2.is_some(), @@ -233,44 +309,19 @@ impl ConnectTo for AgentProtocolRouter { let selected = match select_agent_protocol(&mut first_message, supported) { Ok(selected) => selected, Err(error) => { - send_initialize_error(&client_channel.tx, &first_message, error.clone())?; + send_initialize_error(&client.tx, &first_message, error.clone())?; return Err(error); } }; let Some(agent) = selected.take_agent(self) else { let error = selected.unsupported_error(supported); - send_initialize_error(&client_channel.tx, &first_message, error.clone())?; + send_initialize_error(&client.tx, &first_message, error.clone())?; return Err(error); }; - let (router_channel, agent_channel) = Channel::duplex(); - let Channel { - rx: from_agent, - tx: to_agent, - } = router_channel; - - to_agent - .unbounded_send(Ok(first_message)) - .map_err(crate::util::internal_error)?; - - let agent_future = Box::pin(agent.connect_to(agent_channel)); - - let ((), (), (), ()) = futures::try_join!( - client_future, - agent_future, - Channel { - rx: client_channel.rx, - tx: to_agent, - } - .copy(), - Channel { - rx: from_agent, - tx: client_channel.tx, - } - .copy(), - )?; - - Ok(()) + let agent = RunningProtocolPeer::new(agent); + agent.send(first_message)?; + pipe_protocol_peers_until_closed(client, agent).await } } @@ -410,6 +461,257 @@ fn send_initialize_error( .map_err(crate::util::internal_error) } +#[cfg(feature = "unstable_protocol_v2")] +struct RunningProtocolPeer { + rx: futures::channel::mpsc::UnboundedReceiver>, + tx: futures::channel::mpsc::UnboundedSender>, + future: crate::BoxFuture<'static, Result<(), crate::Error>>, +} + +#[cfg(feature = "unstable_protocol_v2")] +impl RunningProtocolPeer { + fn new(component: impl ConnectTo) -> Self { + let (Channel { rx, tx }, future) = component.into_channel_and_future(); + Self { rx, tx, future } + } + + async fn next_message(self) -> Result, crate::Error> { + let Self { mut rx, tx, future } = self; + match future::select(Box::pin(rx.next()), future).await { + future::Either::Left((Some(message), future)) => { + Ok(Some((message?, Self { rx, tx, future }))) + } + future::Either::Left((None, future)) => { + future.await?; + Ok(None) + } + future::Either::Right((result, next_message)) => { + result?; + drop(next_message); + let Some(message) = rx.next().await else { + return Ok(None); + }; + Ok(Some(( + message?, + Self { + rx, + tx, + future: Box::pin(future::ready(Ok(()))), + }, + ))) + } + } + } + + fn send(&self, message: RawJsonRpcMessage) -> Result<(), crate::Error> { + self.tx + .unbounded_send(Ok(message)) + .map_err(crate::util::internal_error) + } +} + +#[cfg(feature = "unstable_protocol_v2")] +async fn pipe_protocol_peers_until_closed( + left: RunningProtocolPeer, + right: RunningProtocolPeer, +) -> Result<(), crate::Error> { + let ((), (), (), ()) = futures::try_join!( + left.future, + right.future, + Channel { + rx: left.rx, + tx: right.tx, + } + .copy(), + Channel { + rx: right.rx, + tx: left.tx, + } + .copy(), + )?; + + Ok(()) +} + +#[cfg(feature = "unstable_protocol_v2")] +async fn pipe_protocol_peers_until_done( + left: RunningProtocolPeer, + right: RunningProtocolPeer, +) -> Result<(), crate::Error> { + let bridge = Box::pin(async move { + let ((), ()) = futures::try_join!( + Channel { + rx: left.rx, + tx: right.tx, + } + .copy(), + Channel { + rx: right.rx, + tx: left.tx, + } + .copy(), + )?; + Ok(()) + }); + + match future::select(left.future, future::select(right.future, bridge)).await { + future::Either::Left((result, _)) + | future::Either::Right(( + future::Either::Left((result, _)) | future::Either::Right((result, _)), + _, + )) => result, + } +} + +#[cfg(feature = "unstable_protocol_v2")] +#[derive(Debug)] +struct InitializeResponse { + id: RequestId, + result: Result, +} + +#[cfg(feature = "unstable_protocol_v2")] +impl InitializeResponse { + fn from_message(message: RawJsonRpcMessage) -> Result { + match message { + RawJsonRpcMessage::Response(RpcResponse::Result { id, result }) => Ok(Self { + id, + result: Ok(result), + }), + RawJsonRpcMessage::Response(RpcResponse::Error { id, error }) => Ok(Self { + id, + result: Err(error), + }), + message => Err(crate::Error::invalid_request().data(format!( + "first ACP response must be an initialize response, got {message:?}", + ))), + } + } + + fn into_message(self) -> RawJsonRpcMessage { + RawJsonRpcMessage::response(self.id, self.result) + } + + fn with_id(self, id: RequestId) -> RawJsonRpcMessage { + RawJsonRpcMessage::response(id, self.result) + } + + fn protocol_version(&self) -> Option { + serde_json::from_value(self.result.as_ref().ok()?.get("protocolVersion")?.clone()).ok() + } +} + +#[cfg(feature = "unstable_protocol_v2")] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum ClientProtocol { + V1, + V2, +} + +#[cfg(feature = "unstable_protocol_v2")] +impl ClientProtocol { + fn take_client(self, router: &mut ClientProtocolRouter) -> Option> { + match self { + Self::V1 => router.v1.take(), + Self::V2 => router.v2.take(), + } + } + + fn name(self) -> &'static str { + match self { + Self::V1 => "1", + Self::V2 => "2", + } + } +} + +#[cfg(feature = "unstable_protocol_v2")] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct SupportedClientProtocols { + v1: bool, + v2: bool, +} + +#[cfg(feature = "unstable_protocol_v2")] +impl SupportedClientProtocols { + fn highest_configured(self) -> Option { + if self.v2 { + return Some(ClientProtocol::V2); + } + + if self.v1 { + return Some(ClientProtocol::V1); + } + + None + } +} + +#[cfg(feature = "unstable_protocol_v2")] +async fn start_client_protocol( + protocol: ClientProtocol, + client: DynConnectTo, +) -> Result<(RunningProtocolPeer, RawJsonRpcMessage), crate::Error> { + let client = RunningProtocolPeer::new(client); + let Some((initialize, client)) = client.next_message().await? else { + return Err(crate::Error::invalid_request().data(format!( + "ACP protocol version {} client implementation ended before initialize", + protocol.name() + ))); + }; + ensure_client_initialize_request(protocol, &initialize)?; + Ok((client, initialize)) +} + +#[cfg(feature = "unstable_protocol_v2")] +async fn send_initialize_and_receive( + client: RunningProtocolPeer, + agent: RunningProtocolPeer, + initialize: RawJsonRpcMessage, +) -> Result<(RunningProtocolPeer, RunningProtocolPeer, InitializeResponse), crate::Error> { + agent.send(initialize)?; + let Some((response, agent)) = agent.next_message().await? else { + return Err(crate::Error::internal_error().data("agent closed before initialize response")); + }; + let response = InitializeResponse::from_message(response)?; + Ok((client, agent, response)) +} + +#[cfg(feature = "unstable_protocol_v2")] +fn ensure_client_initialize_request( + protocol: ClientProtocol, + message: &RawJsonRpcMessage, +) -> Result<(), crate::Error> { + let RawJsonRpcMessage::Request(request) = message else { + return Err(crate::Error::invalid_request().data(format!( + "ACP protocol version {} client implementation must send initialize first", + protocol.name() + ))); + }; + + if request.method.as_ref() != "initialize" { + return Err(crate::Error::invalid_request().data(format!( + "ACP protocol version {} client implementation must send initialize first", + protocol.name() + ))); + } + + Ok(()) +} + +#[cfg(feature = "unstable_protocol_v2")] +fn initialize_request_id(message: &RawJsonRpcMessage) -> Option { + let RawJsonRpcMessage::Request(request) = message else { + return None; + }; + Some(request.id.clone()) +} + +#[cfg(feature = "unstable_protocol_v2")] +fn initialize_response_negotiated_v1(response: &InitializeResponse) -> bool { + response.protocol_version() == Some(ProtocolVersion::V1) +} + impl HasPeer for Agent { fn remote_style(&self, _peer: Agent) -> RemoteStyle { RemoteStyle::Counterpart diff --git a/src/agent-client-protocol/tests/protocol_v2.rs b/src/agent-client-protocol/tests/protocol_v2.rs index 6c61d7e..28dd578 100644 --- a/src/agent-client-protocol/tests/protocol_v2.rs +++ b/src/agent-client-protocol/tests/protocol_v2.rs @@ -4,8 +4,9 @@ use std::path::PathBuf; use agent_client_protocol::schema::{ProtocolVersion, v1, v2}; use agent_client_protocol::{ - Agent, AgentProtocolRouter, Builder, Client, ConnectTo, Error, JsonRpcMessage, JsonRpcRequest, - JsonRpcResponse, NullHandler, RawJsonRpcMessage, Role, UntypedRole, + Agent, AgentProtocolRouter, Builder, Client, ClientProtocolRouter, ConnectTo, Error, + JsonRpcMessage, JsonRpcRequest, JsonRpcResponse, NullHandler, RawJsonRpcMessage, Role, + UntypedRole, }; use agent_client_protocol_test::testy::Testy; use futures::StreamExt as _; @@ -104,6 +105,163 @@ fn runtime_flag_protocol_router(enable_protocol_v2: bool) -> AgentProtocolRouter } } +fn runtime_flag_client_protocol_router(enable_protocol_v2: bool) -> ClientProtocolRouter { + let client = Client + .protocol_router() + .with_v1(InitializingV1Client::new("v1-client-router-session")); + + if enable_protocol_v2 { + client.with_v2(InitializingV2Client::new("v2-client-router-session")) + } else { + client + } +} + +struct InitializingV1Client { + expected_session_id: &'static str, +} + +impl InitializingV1Client { + fn new(expected_session_id: &'static str) -> Self { + Self { + expected_session_id, + } + } +} + +impl ConnectTo for InitializingV1Client { + async fn connect_to(self, agent: impl ConnectTo) -> Result<(), Error> { + let expected_session_id = self.expected_session_id; + Client + .builder() + .connect_with(agent, async move |cx| { + let initialize = cx + .send_request(v1_initialize_request(ProtocolVersion::V1)) + .block_task() + .await?; + assert_eq!(initialize.protocol_version, ProtocolVersion::V1); + + let session = cx + .send_request(v1::NewSessionRequest::new(cwd()?)) + .block_task() + .await?; + assert_eq!(session.session_id.0.as_ref(), expected_session_id); + Ok(()) + }) + .await + } +} + +struct RejectingV1Client; + +impl ConnectTo for RejectingV1Client { + async fn connect_to(self, _agent: impl ConnectTo) -> Result<(), Error> { + Err(Error::internal_error().data("v1 client fallback should not run")) + } +} + +struct ExpectingV2InitializeErrorClient; + +impl ConnectTo for ExpectingV2InitializeErrorClient { + async fn connect_to(self, agent: impl ConnectTo) -> Result<(), Error> { + Client + .v2() + .connect_with(agent, async |cx| { + let error = cx + .send_request(v2_initialize_request(ProtocolVersion::V1)) + .block_task() + .await + .expect_err("v2 initialize rejection should be surfaced"); + let data = error + .data + .as_ref() + .and_then(|data| data.as_str()) + .unwrap_or_default(); + assert!( + data.contains("only supports ACP protocol version 1"), + "{error:?}" + ); + Ok(()) + }) + .await + } +} + +struct InitializingV2Client { + expected_session_id: &'static str, +} + +impl InitializingV2Client { + fn new(expected_session_id: &'static str) -> Self { + Self { + expected_session_id, + } + } +} + +impl ConnectTo for InitializingV2Client { + async fn connect_to(self, agent: impl ConnectTo) -> Result<(), Error> { + let expected_session_id = self.expected_session_id; + Client + .v2() + .connect_with(agent, async move |cx| { + let initialize = cx + .send_request(v2_initialize_request(ProtocolVersion::V1)) + .block_task() + .await?; + assert_eq!(initialize.protocol_version, ProtocolVersion::V2); + + let session = cx + .send_request(v2::NewSessionRequest::new(cwd()?)) + .block_task() + .await?; + assert_eq!(session.session_id.0.as_ref(), expected_session_id); + Ok(()) + }) + .await + } +} + +fn v1_agent_with_session(session_id: &'static str) -> impl ConnectTo { + Agent + .builder() + .on_receive_request( + async |initialize: v1::InitializeRequest, responder, _cx| { + assert_eq!(initialize.protocol_version, ProtocolVersion::V1); + responder.respond(v1::InitializeResponse::new(initialize.protocol_version)) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async move |request: v1::NewSessionRequest, responder, _cx| { + assert!(request.cwd.is_absolute()); + responder.respond(v1::NewSessionResponse::new(v1::SessionId::new(session_id))) + }, + agent_client_protocol::on_receive_request!(), + ) +} + +fn v2_agent_with_session(session_id: &'static str) -> impl ConnectTo { + Agent + .v2() + .on_receive_request( + async |initialize: v2::InitializeRequest, responder, _cx| { + assert_eq!(initialize.protocol_version, ProtocolVersion::V2); + responder.respond(v2_initialize_response_with_session( + initialize.protocol_version, + )) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async move |request: v2::NewSessionRequest, responder, _cx| { + assert!(request.cwd.is_absolute()); + responder.respond(v2::NewSessionResponse::new(v2::SessionId::new(session_id))) + }, + agent_client_protocol::on_receive_request!(), + ) +} + async fn assert_malformed_initialize_rejected(params: Map) -> Result<(), Error> { let agent = Agent.v2().on_receive_request( async |_initialize: v2::InitializeRequest, responder, _cx| { @@ -623,6 +781,53 @@ async fn v2_client_and_agent_negotiate_v2() -> Result<(), Error> { .await } +#[tokio::test(flavor = "current_thread")] +async fn client_protocol_router_routes_to_v2_client_for_v2_agent() -> Result<(), Error> { + Client + .protocol_router() + .with_v1(InitializingV1Client::new("v1-client-router-session")) + .with_v2(InitializingV2Client::new("v2-client-router-session")) + .connect_to(v2_agent_with_session("v2-client-router-session")) + .await +} + +#[tokio::test(flavor = "current_thread")] +async fn client_protocol_router_does_not_retry_after_v2_initialize_rejection() -> Result<(), Error> +{ + Client + .protocol_router() + .with_v1(RejectingV1Client) + .with_v2(ExpectingV2InitializeErrorClient) + .connect_to(v1_agent_with_session("v1-client-router-session")) + .await +} + +#[tokio::test(flavor = "current_thread")] +async fn client_protocol_router_falls_back_to_v1_when_agent_router_negotiates_v1() +-> Result<(), Error> { + let agent = Agent + .protocol_router() + .with_v1(v1_agent_with_session("v1-client-router-session")); + + Client + .protocol_router() + .with_v1(InitializingV1Client::new("v1-client-router-session")) + .with_v2(InitializingV2Client::new("v2-client-router-session")) + .connect_to(agent) + .await +} + +#[tokio::test(flavor = "current_thread")] +async fn client_protocol_router_supports_runtime_v2_registration_flag() -> Result<(), Error> { + runtime_flag_client_protocol_router(false) + .connect_to(v1_agent_with_session("v1-client-router-session")) + .await?; + + runtime_flag_client_protocol_router(true) + .connect_to(v2_agent_with_session("v2-client-router-session")) + .await +} + #[tokio::test(flavor = "current_thread")] async fn protocol_router_routes_v1_client_to_v1_implementation() -> Result<(), Error> { let agent = Agent From 83bb835dbe549d6faa163319b5e532d8677c7a75 Mon Sep 17 00:00:00 2001 From: Ben Brandt Date: Thu, 9 Jul 2026 14:53:57 +0200 Subject: [PATCH 3/6] fix(acp): Flush protocol router rejections --- src/agent-client-protocol/src/role/acp.rs | 27 +++++- .../tests/protocol_v2.rs | 83 ++++++++++++++++++- 2 files changed, 104 insertions(+), 6 deletions(-) diff --git a/src/agent-client-protocol/src/role/acp.rs b/src/agent-client-protocol/src/role/acp.rs index 9d09ed9..784af19 100644 --- a/src/agent-client-protocol/src/role/acp.rs +++ b/src/agent-client-protocol/src/role/acp.rs @@ -309,14 +309,12 @@ impl ConnectTo for AgentProtocolRouter { let selected = match select_agent_protocol(&mut first_message, supported) { Ok(selected) => selected, Err(error) => { - send_initialize_error(&client.tx, &first_message, error.clone())?; - return Err(error); + return reject_initialize(client, &first_message, error).await; } }; let Some(agent) = selected.take_agent(self) else { let error = selected.unsupported_error(supported); - send_initialize_error(&client.tx, &first_message, error.clone())?; - return Err(error); + return reject_initialize(client, &first_message, error).await; }; let agent = RunningProtocolPeer::new(agent); @@ -461,6 +459,27 @@ fn send_initialize_error( .map_err(crate::util::internal_error) } +#[cfg(feature = "unstable_protocol_v2")] +async fn reject_initialize( + client: RunningProtocolPeer, + message: &RawJsonRpcMessage, + error: crate::Error, +) -> Result<(), crate::Error> { + let RunningProtocolPeer { mut rx, tx, future } = client; + send_initialize_error(&tx, message, error)?; + drop(tx); + + let drain_incoming = async move { + while let Some(message) = rx.next().await { + message?; + } + Ok(()) + }; + + let ((), ()) = futures::try_join!(future, drain_incoming)?; + Ok(()) +} + #[cfg(feature = "unstable_protocol_v2")] struct RunningProtocolPeer { rx: futures::channel::mpsc::UnboundedReceiver>, diff --git a/src/agent-client-protocol/tests/protocol_v2.rs b/src/agent-client-protocol/tests/protocol_v2.rs index 28dd578..c959e0b 100644 --- a/src/agent-client-protocol/tests/protocol_v2.rs +++ b/src/agent-client-protocol/tests/protocol_v2.rs @@ -4,14 +4,15 @@ use std::path::PathBuf; use agent_client_protocol::schema::{ProtocolVersion, v1, v2}; use agent_client_protocol::{ - Agent, AgentProtocolRouter, Builder, Client, ClientProtocolRouter, ConnectTo, Error, - JsonRpcMessage, JsonRpcRequest, JsonRpcResponse, NullHandler, RawJsonRpcMessage, Role, + Agent, AgentProtocolRouter, Builder, ByteStreams, Client, ClientProtocolRouter, ConnectTo, + Error, JsonRpcMessage, JsonRpcRequest, JsonRpcResponse, NullHandler, RawJsonRpcMessage, Role, UntypedRole, }; use agent_client_protocol_test::testy::Testy; use futures::StreamExt as _; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; +use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt}; #[derive(Debug, Clone, Serialize, Deserialize, JsonRpcRequest)] #[request(method = "initialize", response = ForeignInitializeResponse)] @@ -1062,6 +1063,84 @@ async fn protocol_router_v2_only_rejects_v1_client() -> Result<(), Error> { )) } +#[tokio::test(flavor = "current_thread")] +async fn protocol_router_rejection_is_initialize_request_error() -> Result<(), Error> { + let agent = Agent + .protocol_router() + .with_v2(Agent.v2().on_receive_request( + async |_initialize: v2::InitializeRequest, responder, _cx| { + responder.respond_with_internal_error("v2 implementation should not run") + }, + agent_client_protocol::on_receive_request!(), + )); + + Client + .builder() + .connect_with(agent, async |cx| { + let error = cx + .send_request(v1_initialize_request(ProtocolVersion::V1)) + .block_task() + .await + .expect_err("v1 initialize should be rejected by the v2-only router"); + let data = error + .data + .as_ref() + .and_then(|data| data.as_str()) + .unwrap_or_default(); + assert!( + data.contains("supports ACP protocol version 2"), + "{error:?}" + ); + Ok(()) + }) + .await +} + +#[tokio::test(flavor = "current_thread")] +async fn protocol_router_rejection_flushes_over_byte_streams() -> Result<(), Error> { + let agent = Agent + .protocol_router() + .with_v2(Agent.v2().on_receive_request( + async |_initialize: v2::InitializeRequest, responder, _cx| { + responder.respond_with_internal_error("v2 implementation should not run") + }, + agent_client_protocol::on_receive_request!(), + )); + + let (client_writer, server_reader) = tokio::io::duplex(1024); + let (server_writer, client_reader) = tokio::io::duplex(1024); + let server_transport = ByteStreams::new(server_writer.compat_write(), server_reader.compat()); + let client_transport = ByteStreams::new(client_writer.compat_write(), client_reader.compat()); + + let agent_task = tokio::spawn(agent.connect_to(server_transport)); + + Client + .builder() + .connect_with(client_transport, async |cx| { + let error = cx + .send_request(v1_initialize_request(ProtocolVersion::V1)) + .block_task() + .await + .expect_err("v1 initialize should be rejected by the v2-only router"); + let data = error + .data + .as_ref() + .and_then(|data| data.as_str()) + .unwrap_or_default(); + assert!( + data.contains("supports ACP protocol version 2"), + "{error:?}" + ); + Ok(()) + }) + .await?; + + agent_task + .await + .map_err(agent_client_protocol::util::internal_error)??; + Ok(()) +} + #[tokio::test(flavor = "current_thread")] async fn protocol_router_routes_future_protocol_version_to_v2() -> Result<(), Error> { let agent = Agent From bb68c795be5289221cece2cffc6071c1a1cbca96 Mon Sep 17 00:00:00 2001 From: Ben Brandt Date: Thu, 9 Jul 2026 15:13:59 +0200 Subject: [PATCH 4/6] fix(acp): Rewrite initialize params on protocol downgrade --- src/agent-client-protocol/src/role/acp.rs | 76 +++++++++++++++---- .../tests/protocol_v2.rs | 46 +++++++++++ 2 files changed, 109 insertions(+), 13 deletions(-) diff --git a/src/agent-client-protocol/src/role/acp.rs b/src/agent-client-protocol/src/role/acp.rs index 784af19..0a59b6c 100644 --- a/src/agent-client-protocol/src/role/acp.rs +++ b/src/agent-client-protocol/src/role/acp.rs @@ -2,17 +2,19 @@ use std::{fmt::Debug, hash::Hash}; #[cfg(feature = "unstable_protocol_v2")] use futures::{StreamExt as _, future}; +#[cfg(feature = "unstable_protocol_v2")] +use serde::{Serialize, de::DeserializeOwned}; #[cfg(feature = "unstable_protocol_v2")] use crate::DynConnectTo; use crate::jsonrpc::{Builder, handlers::NullHandler, run::NullRun}; use crate::role::{HasPeer, RemoteStyle}; -#[cfg(feature = "unstable_protocol_v2")] -use crate::schema::ProtocolVersion; use crate::schema::v1::{InitializeRequest, NewSessionRequest, NewSessionResponse, SessionId}; #[cfg(feature = "unstable_protocol_v2")] use crate::schema::v1::{RequestId, Response as RpcResponse}; use crate::schema::{InitializeProxyRequest, METHOD_INITIALIZE_PROXY}; +#[cfg(feature = "unstable_protocol_v2")] +use crate::schema::{ProtocolVersion, v2}; use crate::util::MatchDispatchFrom; #[cfg(feature = "unstable_protocol_v2")] use crate::{Channel, RawJsonRpcMessage, RawJsonRpcParams}; @@ -332,13 +334,6 @@ enum AgentProtocol { #[cfg(feature = "unstable_protocol_v2")] impl AgentProtocol { - fn version(self) -> ProtocolVersion { - match self { - Self::V1 => ProtocolVersion::V1, - Self::V2 => ProtocolVersion::V2, - } - } - fn take_agent(self, agent: AgentProtocolRouter) -> Option> { match self { Self::V1 => agent.v1, @@ -418,14 +413,64 @@ fn select_agent_protocol( let requested = serde_json::from_value::(protocol_version.clone()) .map_err(|_| invalid_initialize_protocol_version())?; let selected = highest_compatible_agent_protocol(requested, supported)?; - params.insert( - "protocolVersion".into(), - serde_json::to_value(selected.version()).map_err(crate::Error::into_internal_error)?, - ); + rewrite_initialize_params(params, requested, selected)?; Ok(selected) } +#[cfg(feature = "unstable_protocol_v2")] +fn rewrite_initialize_params( + params: &mut serde_json::Map, + requested: ProtocolVersion, + selected: AgentProtocol, +) -> Result<(), crate::Error> { + match selected { + AgentProtocol::V1 => { + let mut initialize = if requested >= ProtocolVersion::V2 { + v2::conversion::v2_to_v1(parse_initialize_params::(params)?) + .map_err(invalid_initialize_params)? + } else { + parse_initialize_params::(params)? + }; + initialize.protocol_version = ProtocolVersion::V1; + replace_initialize_params(params, initialize) + } + AgentProtocol::V2 => { + let mut initialize = if requested >= ProtocolVersion::V2 { + parse_initialize_params::(params)? + } else { + v2::conversion::v1_to_v2(parse_initialize_params::(params)?) + .map_err(invalid_initialize_params)? + }; + initialize.protocol_version = ProtocolVersion::V2; + replace_initialize_params(params, initialize) + } + } +} + +#[cfg(feature = "unstable_protocol_v2")] +fn parse_initialize_params( + params: &serde_json::Map, +) -> Result { + serde_json::from_value(serde_json::Value::Object(params.clone())) + .map_err(invalid_initialize_params) +} + +#[cfg(feature = "unstable_protocol_v2")] +fn replace_initialize_params( + params: &mut serde_json::Map, + initialize: impl Serialize, +) -> Result<(), crate::Error> { + let value = serde_json::to_value(initialize).map_err(crate::Error::into_internal_error)?; + let serde_json::Value::Object(object) = value else { + return Err(crate::util::internal_error( + "initialize params did not serialize to an object", + )); + }; + *params = object; + Ok(()) +} + #[cfg(feature = "unstable_protocol_v2")] fn highest_compatible_agent_protocol( requested: ProtocolVersion, @@ -445,6 +490,11 @@ fn invalid_initialize_protocol_version() -> crate::Error { .data("initialize.protocolVersion must be a valid ACP protocol version") } +#[cfg(feature = "unstable_protocol_v2")] +fn invalid_initialize_params(error: impl ToString) -> crate::Error { + crate::Error::invalid_params().data(format!("invalid initialize params: {}", error.to_string())) +} + #[cfg(feature = "unstable_protocol_v2")] fn send_initialize_error( tx: &futures::channel::mpsc::UnboundedSender>, diff --git a/src/agent-client-protocol/tests/protocol_v2.rs b/src/agent-client-protocol/tests/protocol_v2.rs index c959e0b..af6f4c4 100644 --- a/src/agent-client-protocol/tests/protocol_v2.rs +++ b/src/agent-client-protocol/tests/protocol_v2.rs @@ -1016,6 +1016,52 @@ async fn protocol_router_supports_runtime_v2_registration_flag() -> Result<(), E .await } +#[tokio::test(flavor = "current_thread")] +async fn protocol_router_downgrades_v2_initialize_metadata_to_v1() -> Result<(), Error> { + let agent = Agent + .protocol_router() + .with_v1(Agent.builder().on_receive_request( + async |initialize: v1::InitializeRequest, responder, _cx| { + assert_eq!(initialize.protocol_version, ProtocolVersion::V1); + let client_info = initialize + .client_info + .as_ref() + .expect("v2 info should become v1 clientInfo"); + assert_eq!(client_info.name, "v2-metadata-client"); + assert_eq!(client_info.title.as_deref(), Some("V2 Metadata Client")); + assert!( + initialize.client_capabilities.session.is_some(), + "{:?}", + initialize.client_capabilities + ); + responder.respond(v1::InitializeResponse::new(initialize.protocol_version)) + }, + agent_client_protocol::on_receive_request!(), + )); + + Client + .v2() + .connect_with(agent, async |cx| { + let request = v2::InitializeRequest::new( + ProtocolVersion::V2, + v2::Implementation::new("v2-metadata-client", "9.9.9").title("V2 Metadata Client"), + ); + let error = cx + .send_request(request) + .block_task() + .await + .expect_err("v2 client should reject the downgraded v1 initialize response"); + let data = error + .data + .as_ref() + .and_then(|data| data.as_str()) + .unwrap_or_default(); + assert!(data.contains("peer negotiated 1"), "{error:?}"); + Ok(()) + }) + .await +} + #[tokio::test(flavor = "current_thread")] async fn protocol_router_v2_only_rejects_v1_client() -> Result<(), Error> { let agent = Agent From e14218b8cd66af289ef1fd2971e7ae75fc9498f5 Mon Sep 17 00:00:00 2001 From: Ben Brandt Date: Thu, 9 Jul 2026 16:56:29 +0200 Subject: [PATCH 5/6] feat(acp): Rework client protocol connector --- src/agent-client-protocol/src/lib.rs | 2 +- src/agent-client-protocol/src/role/acp.rs | 192 +++++++++++------- .../tests/protocol_v2.rs | 145 ++++++++++--- 3 files changed, 235 insertions(+), 104 deletions(-) diff --git a/src/agent-client-protocol/src/lib.rs b/src/agent-client-protocol/src/lib.rs index 710401e..7d9453a 100644 --- a/src/agent-client-protocol/src/lib.rs +++ b/src/agent-client-protocol/src/lib.rs @@ -106,7 +106,7 @@ pub use jsonrpc::{RequestCancellation, is_cancel_request_notification}; #[cfg(feature = "unstable_protocol_v2")] pub use role::acp::AgentProtocolRouter; #[cfg(feature = "unstable_protocol_v2")] -pub use role::acp::ClientProtocolRouter; +pub use role::acp::ClientProtocolConnector; pub use role::{ Role, RoleId, UntypedRole, acp::{Agent, Client, Conductor, Proxy}, diff --git a/src/agent-client-protocol/src/role/acp.rs b/src/agent-client-protocol/src/role/acp.rs index 0a59b6c..002a026 100644 --- a/src/agent-client-protocol/src/role/acp.rs +++ b/src/agent-client-protocol/src/role/acp.rs @@ -71,21 +71,22 @@ impl Client { self.builder().v2_client() } - /// Create a router that chooses between configured protocol implementations. + /// Create a connector that chooses between configured protocol implementations. /// - /// Add implementations with [`ClientProtocolRouter::with_v1`] and - /// [`ClientProtocolRouter::with_v2`]. The resulting router starts the - /// highest configured protocol implementation, observes the initial - /// `initialize` response, and falls back to a configured v1 implementation - /// when v2 initialization negotiates v1. It does not convert traffic - /// between protocol versions after routing. + /// Add implementation factories with [`ClientProtocolConnector::with_v1`] + /// and [`ClientProtocolConnector::with_v2`]. The resulting connector starts + /// the highest configured protocol implementation. If a v2 implementation + /// successfully negotiates v1 and a v1 implementation is configured, the + /// connector opens a fresh agent connection and restarts with the v1 + /// implementation so the agent observes the v1 implementation's own + /// `initialize` metadata and capabilities. /// /// Requires the `unstable_protocol_v2` crate feature while protocol v2 /// stabilizes. #[cfg(feature = "unstable_protocol_v2")] #[must_use] - pub fn protocol_router(self) -> ClientProtocolRouter { - ClientProtocolRouter::new() + pub fn protocol_connector(self) -> ClientProtocolConnector { + ClientProtocolConnector::new() } /// Connect to `agent` and run `main_fn` with the [`ConnectionTo`]. @@ -101,78 +102,123 @@ impl Client { } } -/// Client component that routes each connection to a configured protocol implementation. +/// Client connector that opens an agent connection with a configured protocol implementation. /// -/// Use [`Client::protocol_router`] to start the builder, then add each supported -/// protocol version independently. The selected implementation owns the -/// connection after the initial `initialize` negotiation. +/// Use [`Client::protocol_connector`] to start the builder, then add each +/// supported protocol version independently. Implementations and the agent +/// connection are provided as factories because fallback from v2 to v1 requires +/// a fresh connection initialized by the v1 implementation. #[cfg(feature = "unstable_protocol_v2")] #[derive(Debug, Default)] -pub struct ClientProtocolRouter { - v1: Option>, - v2: Option>, +pub struct ClientProtocolConnector { + v1: Option>, + v2: Option>, } #[cfg(feature = "unstable_protocol_v2")] -impl ClientProtocolRouter { - /// Create an empty client protocol router. +impl ClientProtocolConnector { + /// Create an empty client protocol connector. #[must_use] pub fn new() -> Self { Self::default() } - /// Return this router with an ACP v1 implementation configured. + /// Return this connector with an ACP v1 implementation factory configured. #[must_use] - pub fn with_v1(mut self, client: impl ConnectTo) -> Self { - self.v1 = Some(DynConnectTo::new(client)); + pub fn with_v1(mut self, client: impl FnMut() -> C + Send + 'static) -> Self + where + C: ConnectTo, + { + self.v1 = Some(DynConnectToFactory::new(client)); self } - /// Return this router with an ACP v2 implementation configured. + /// Return this connector with an ACP v2 implementation factory configured. #[must_use] - pub fn with_v2(mut self, client: impl ConnectTo) -> Self { - self.v2 = Some(DynConnectTo::new(client)); + pub fn with_v2(mut self, client: impl FnMut() -> C + Send + 'static) -> Self + where + C: ConnectTo, + { + self.v2 = Some(DynConnectToFactory::new(client)); self } -} -#[cfg(feature = "unstable_protocol_v2")] -impl ConnectTo for ClientProtocolRouter { - async fn connect_to(mut self, agent: impl ConnectTo) -> Result<(), crate::Error> { + /// Connect to an agent produced by `agent` using the highest configured + /// compatible protocol implementation. + pub async fn connect_to( + mut self, + mut agent: impl FnMut() -> C + Send + 'static, + ) -> Result<(), crate::Error> + where + C: ConnectTo, + { let supported = SupportedClientProtocols { v1: self.v1.is_some(), v2: self.v2.is_some(), }; let Some(selected) = supported.highest_configured() else { return Err(crate::Error::invalid_request() - .data("client protocol router has no configured ACP protocol implementations")); + .data("client protocol connector has no configured ACP protocol implementations")); }; - let agent = RunningProtocolPeer::new(agent); - let client = selected - .take_client(&mut self) - .expect("selected protocol is configured"); - let (client, initialize) = start_client_protocol(selected, client).await?; - let (client, agent, initialize_response) = - send_initialize_and_receive(client, agent, initialize).await?; - - if selected == ClientProtocol::V2 - && let Some(client) = self.v1.take() - && initialize_response_negotiated_v1(&initialize_response) - { - let (fallback_client, fallback_initialize) = - start_client_protocol(ClientProtocol::V1, client).await?; - - let fallback_response = initialize_response.with_id( - initialize_request_id(&fallback_initialize) - .expect("validated initialize request has an id"), - ); - fallback_client.send(fallback_response)?; - return pipe_protocol_peers_until_done(fallback_client, agent).await; + match selected { + ClientProtocol::V1 => { + let client = self + .v1 + .as_mut() + .expect("selected protocol is configured") + .create(); + connect_client_protocol(ClientProtocol::V1, client, agent()).await + } + ClientProtocol::V2 => { + let client = self + .v2 + .as_mut() + .expect("selected protocol is configured") + .create(); + let (client, agent_connection, initialize_response) = + initialize_client_protocol(ClientProtocol::V2, client, agent()).await?; + + if initialize_response_negotiated_v1(&initialize_response) + && let Some(v1) = self.v1.as_mut() + { + drop((client, agent_connection, initialize_response)); + return connect_client_protocol(ClientProtocol::V1, v1.create(), agent()).await; + } + + client.send(initialize_response.into_message())?; + pipe_protocol_peers_until_done(client, agent_connection).await + } + } + } +} + +#[cfg(feature = "unstable_protocol_v2")] +struct DynConnectToFactory { + inner: Box DynConnectTo + Send>, +} + +#[cfg(feature = "unstable_protocol_v2")] +impl DynConnectToFactory { + fn new(mut factory: impl FnMut() -> C + Send + 'static) -> Self + where + C: ConnectTo, + { + Self { + inner: Box::new(move || DynConnectTo::new(factory())), } + } + + fn create(&mut self) -> DynConnectTo { + (self.inner)() + } +} - client.send(initialize_response.into_message())?; - pipe_protocol_peers_until_done(client, agent).await +#[cfg(feature = "unstable_protocol_v2")] +impl Debug for DynConnectToFactory { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("DynConnectToFactory") + .finish_non_exhaustive() } } @@ -661,10 +707,6 @@ impl InitializeResponse { RawJsonRpcMessage::response(self.id, self.result) } - fn with_id(self, id: RequestId) -> RawJsonRpcMessage { - RawJsonRpcMessage::response(id, self.result) - } - fn protocol_version(&self) -> Option { serde_json::from_value(self.result.as_ref().ok()?.get("protocolVersion")?.clone()).ok() } @@ -679,13 +721,6 @@ enum ClientProtocol { #[cfg(feature = "unstable_protocol_v2")] impl ClientProtocol { - fn take_client(self, router: &mut ClientProtocolRouter) -> Option> { - match self { - Self::V1 => router.v1.take(), - Self::V2 => router.v2.take(), - } - } - fn name(self) -> &'static str { match self { Self::V1 => "1", @@ -746,6 +781,29 @@ async fn send_initialize_and_receive( Ok((client, agent, response)) } +#[cfg(feature = "unstable_protocol_v2")] +async fn initialize_client_protocol( + protocol: ClientProtocol, + client: DynConnectTo, + agent: impl ConnectTo, +) -> Result<(RunningProtocolPeer, RunningProtocolPeer, InitializeResponse), crate::Error> { + let agent = RunningProtocolPeer::new(agent); + let (client, initialize) = start_client_protocol(protocol, client).await?; + send_initialize_and_receive(client, agent, initialize).await +} + +#[cfg(feature = "unstable_protocol_v2")] +async fn connect_client_protocol( + protocol: ClientProtocol, + client: DynConnectTo, + agent: impl ConnectTo, +) -> Result<(), crate::Error> { + let (client, agent, initialize_response) = + initialize_client_protocol(protocol, client, agent).await?; + client.send(initialize_response.into_message())?; + pipe_protocol_peers_until_done(client, agent).await +} + #[cfg(feature = "unstable_protocol_v2")] fn ensure_client_initialize_request( protocol: ClientProtocol, @@ -768,14 +826,6 @@ fn ensure_client_initialize_request( Ok(()) } -#[cfg(feature = "unstable_protocol_v2")] -fn initialize_request_id(message: &RawJsonRpcMessage) -> Option { - let RawJsonRpcMessage::Request(request) = message else { - return None; - }; - Some(request.id.clone()) -} - #[cfg(feature = "unstable_protocol_v2")] fn initialize_response_negotiated_v1(response: &InitializeResponse) -> bool { response.protocol_version() == Some(ProtocolVersion::V1) diff --git a/src/agent-client-protocol/tests/protocol_v2.rs b/src/agent-client-protocol/tests/protocol_v2.rs index af6f4c4..219835f 100644 --- a/src/agent-client-protocol/tests/protocol_v2.rs +++ b/src/agent-client-protocol/tests/protocol_v2.rs @@ -1,10 +1,16 @@ #![cfg(feature = "unstable_protocol_v2")] -use std::path::PathBuf; +use std::{ + path::PathBuf, + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, +}; use agent_client_protocol::schema::{ProtocolVersion, v1, v2}; use agent_client_protocol::{ - Agent, AgentProtocolRouter, Builder, ByteStreams, Client, ClientProtocolRouter, ConnectTo, + Agent, AgentProtocolRouter, Builder, ByteStreams, Client, ClientProtocolConnector, ConnectTo, Error, JsonRpcMessage, JsonRpcRequest, JsonRpcResponse, NullHandler, RawJsonRpcMessage, Role, UntypedRole, }; @@ -106,13 +112,13 @@ fn runtime_flag_protocol_router(enable_protocol_v2: bool) -> AgentProtocolRouter } } -fn runtime_flag_client_protocol_router(enable_protocol_v2: bool) -> ClientProtocolRouter { +fn runtime_flag_client_protocol_connector(enable_protocol_v2: bool) -> ClientProtocolConnector { let client = Client - .protocol_router() - .with_v1(InitializingV1Client::new("v1-client-router-session")); + .protocol_connector() + .with_v1(|| InitializingV1Client::new("v1-client-connector-session")); if enable_protocol_v2 { - client.with_v2(InitializingV2Client::new("v2-client-router-session")) + client.with_v2(|| InitializingV2Client::new("v2-client-connector-session")) } else { client } @@ -120,12 +126,24 @@ fn runtime_flag_client_protocol_router(enable_protocol_v2: bool) -> ClientProtoc struct InitializingV1Client { expected_session_id: &'static str, + implementation_name: &'static str, } impl InitializingV1Client { fn new(expected_session_id: &'static str) -> Self { Self { expected_session_id, + implementation_name: "agent-client-protocol-test", + } + } + + fn with_implementation_name( + expected_session_id: &'static str, + implementation_name: &'static str, + ) -> Self { + Self { + expected_session_id, + implementation_name, } } } @@ -133,11 +151,14 @@ impl InitializingV1Client { impl ConnectTo for InitializingV1Client { async fn connect_to(self, agent: impl ConnectTo) -> Result<(), Error> { let expected_session_id = self.expected_session_id; + let implementation_name = self.implementation_name; Client .builder() .connect_with(agent, async move |cx| { let initialize = cx - .send_request(v1_initialize_request(ProtocolVersion::V1)) + .send_request(v1::InitializeRequest::new(ProtocolVersion::V1).client_info( + v1::Implementation::new(implementation_name, env!("CARGO_PKG_VERSION")), + )) .block_task() .await?; assert_eq!(initialize.protocol_version, ProtocolVersion::V1); @@ -783,49 +804,109 @@ async fn v2_client_and_agent_negotiate_v2() -> Result<(), Error> { } #[tokio::test(flavor = "current_thread")] -async fn client_protocol_router_routes_to_v2_client_for_v2_agent() -> Result<(), Error> { +async fn client_protocol_connector_routes_to_v2_client_for_v2_agent() -> Result<(), Error> { Client - .protocol_router() - .with_v1(InitializingV1Client::new("v1-client-router-session")) - .with_v2(InitializingV2Client::new("v2-client-router-session")) - .connect_to(v2_agent_with_session("v2-client-router-session")) + .protocol_connector() + .with_v1(|| InitializingV1Client::new("v1-client-connector-session")) + .with_v2(|| InitializingV2Client::new("v2-client-connector-session")) + .connect_to(|| v2_agent_with_session("v2-client-connector-session")) .await } #[tokio::test(flavor = "current_thread")] -async fn client_protocol_router_does_not_retry_after_v2_initialize_rejection() -> Result<(), Error> -{ +async fn client_protocol_connector_does_not_retry_after_v2_initialize_rejection() +-> Result<(), Error> { Client - .protocol_router() - .with_v1(RejectingV1Client) - .with_v2(ExpectingV2InitializeErrorClient) - .connect_to(v1_agent_with_session("v1-client-router-session")) + .protocol_connector() + .with_v1(|| RejectingV1Client) + .with_v2(|| ExpectingV2InitializeErrorClient) + .connect_to(|| v1_agent_with_session("v1-client-connector-session")) .await } #[tokio::test(flavor = "current_thread")] -async fn client_protocol_router_falls_back_to_v1_when_agent_router_negotiates_v1() +async fn client_protocol_connector_falls_back_to_v1_when_agent_router_negotiates_v1() -> Result<(), Error> { - let agent = Agent - .protocol_router() - .with_v1(v1_agent_with_session("v1-client-router-session")); - Client - .protocol_router() - .with_v1(InitializingV1Client::new("v1-client-router-session")) - .with_v2(InitializingV2Client::new("v2-client-router-session")) - .connect_to(agent) + .protocol_connector() + .with_v1(|| InitializingV1Client::new("v1-client-connector-session")) + .with_v2(|| InitializingV2Client::new("v2-client-connector-session")) + .connect_to(|| { + Agent + .protocol_router() + .with_v1(v1_agent_with_session("v1-client-connector-session")) + }) .await } #[tokio::test(flavor = "current_thread")] -async fn client_protocol_router_supports_runtime_v2_registration_flag() -> Result<(), Error> { - runtime_flag_client_protocol_router(false) - .connect_to(v1_agent_with_session("v1-client-router-session")) +async fn client_protocol_connector_reconnects_before_v1_fallback() -> Result<(), Error> { + let connections = Arc::new(AtomicUsize::new(0)); + let agent_connections = Arc::clone(&connections); + + Client + .protocol_connector() + .with_v1(|| { + InitializingV1Client::with_implementation_name( + "v1-reconnected-session", + "v1-reconnected-client", + ) + }) + .with_v2(|| InitializingV2Client::new("v2-client-should-not-continue")) + .connect_to(move || { + let connection_number = agent_connections.fetch_add(1, Ordering::SeqCst) + 1; + let initialize_connection_number = connection_number; + let session_connection_number = connection_number; + + Agent.protocol_router().with_v1( + Agent + .builder() + .on_receive_request( + async move |initialize: v1::InitializeRequest, responder, _cx| { + assert_eq!(initialize.protocol_version, ProtocolVersion::V1); + let info = initialize + .client_info + .as_ref() + .expect("initialize should include client info"); + + if initialize_connection_number == 1 { + assert_eq!(&*info.name, "agent-client-protocol-test"); + } else { + assert_eq!(initialize_connection_number, 2); + assert_eq!(&*info.name, "v1-reconnected-client"); + } + + responder + .respond(v1::InitializeResponse::new(initialize.protocol_version)) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async move |request: v1::NewSessionRequest, responder, _cx| { + assert_eq!(session_connection_number, 2); + assert!(request.cwd.is_absolute()); + responder.respond(v1::NewSessionResponse::new(v1::SessionId::new( + "v1-reconnected-session", + ))) + }, + agent_client_protocol::on_receive_request!(), + ), + ) + }) + .await?; + + assert_eq!(connections.load(Ordering::SeqCst), 2); + Ok(()) +} + +#[tokio::test(flavor = "current_thread")] +async fn client_protocol_connector_supports_runtime_v2_registration_flag() -> Result<(), Error> { + runtime_flag_client_protocol_connector(false) + .connect_to(|| v1_agent_with_session("v1-client-connector-session")) .await?; - runtime_flag_client_protocol_router(true) - .connect_to(v2_agent_with_session("v2-client-router-session")) + runtime_flag_client_protocol_connector(true) + .connect_to(|| v2_agent_with_session("v2-client-connector-session")) .await } From 6808830a07a3a298e7e71bf4a2c9f6be688bac9a Mon Sep 17 00:00:00 2001 From: Ben Brandt Date: Thu, 9 Jul 2026 18:05:57 +0200 Subject: [PATCH 6/6] fix(acp): reuse matching v1 fallback connections --- src/agent-client-protocol/src/role/acp.rs | 76 +++++++++++++++-- .../tests/protocol_v2.rs | 83 +++++++++++++++++-- 2 files changed, 147 insertions(+), 12 deletions(-) diff --git a/src/agent-client-protocol/src/role/acp.rs b/src/agent-client-protocol/src/role/acp.rs index 002a026..32249a3 100644 --- a/src/agent-client-protocol/src/role/acp.rs +++ b/src/agent-client-protocol/src/role/acp.rs @@ -77,9 +77,9 @@ impl Client { /// and [`ClientProtocolConnector::with_v2`]. The resulting connector starts /// the highest configured protocol implementation. If a v2 implementation /// successfully negotiates v1 and a v1 implementation is configured, the - /// connector opens a fresh agent connection and restarts with the v1 - /// implementation so the agent observes the v1 implementation's own - /// `initialize` metadata and capabilities. + /// connector reuses the connection only when the v1 implementation's + /// `initialize` metadata and capabilities match what the agent already saw; + /// otherwise it opens a fresh agent connection and restarts with v1. /// /// Requires the `unstable_protocol_v2` crate feature while protocol v2 /// stabilizes. @@ -106,8 +106,8 @@ impl Client { /// /// Use [`Client::protocol_connector`] to start the builder, then add each /// supported protocol version independently. Implementations and the agent -/// connection are provided as factories because fallback from v2 to v1 requires -/// a fresh connection initialized by the v1 implementation. +/// connection are provided as factories because fallback from v2 to v1 may +/// require a fresh connection initialized by the v1 implementation. #[cfg(feature = "unstable_protocol_v2")] #[derive(Debug, Default)] pub struct ClientProtocolConnector { @@ -176,12 +176,35 @@ impl ClientProtocolConnector { .as_mut() .expect("selected protocol is configured") .create(); + let agent_connection = RunningProtocolPeer::new(agent()); + let (client, initialize) = + start_client_protocol(ClientProtocol::V2, client).await?; + let v2_initialize_as_v1 = + normalize_initialize_params_for_agent_protocol(&initialize, AgentProtocol::V1)?; let (client, agent_connection, initialize_response) = - initialize_client_protocol(ClientProtocol::V2, client, agent()).await?; + send_initialize_and_receive(client, agent_connection, initialize).await?; if initialize_response_negotiated_v1(&initialize_response) && let Some(v1) = self.v1.as_mut() { + let fallback_client = v1.create(); + let (fallback_client, fallback_initialize) = + start_client_protocol(ClientProtocol::V1, fallback_client).await?; + let v1_initialize = normalize_initialize_params_for_agent_protocol( + &fallback_initialize, + AgentProtocol::V1, + )?; + + if v1_initialize == v2_initialize_as_v1 { + let fallback_response = initialize_response.with_id( + initialize_request_id(&fallback_initialize) + .expect("validated initialize request has an id"), + ); + fallback_client.send(fallback_response)?; + return pipe_protocol_peers_until_done(fallback_client, agent_connection) + .await; + } + drop((client, agent_connection, initialize_response)); return connect_client_protocol(ClientProtocol::V1, v1.create(), agent()).await; } @@ -464,6 +487,35 @@ fn select_agent_protocol( Ok(selected) } +#[cfg(feature = "unstable_protocol_v2")] +fn normalize_initialize_params_for_agent_protocol( + message: &RawJsonRpcMessage, + selected: AgentProtocol, +) -> Result, crate::Error> { + let RawJsonRpcMessage::Request(request) = message else { + return Err( + crate::Error::invalid_request().data("first ACP message must be an initialize request") + ); + }; + + if request.method.as_ref() != "initialize" { + return Err(crate::Error::invalid_request().data("first ACP request must be initialize")); + } + + let Some(RawJsonRpcParams::Object(params)) = &request.params else { + return Err(invalid_initialize_protocol_version()); + }; + let Some(protocol_version) = params.get("protocolVersion") else { + return Err(invalid_initialize_protocol_version()); + }; + + let requested = serde_json::from_value::(protocol_version.clone()) + .map_err(|_| invalid_initialize_protocol_version())?; + let mut params = params.clone(); + rewrite_initialize_params(&mut params, requested, selected)?; + Ok(params) +} + #[cfg(feature = "unstable_protocol_v2")] fn rewrite_initialize_params( params: &mut serde_json::Map, @@ -707,6 +759,10 @@ impl InitializeResponse { RawJsonRpcMessage::response(self.id, self.result) } + fn with_id(self, id: RequestId) -> RawJsonRpcMessage { + RawJsonRpcMessage::response(id, self.result) + } + fn protocol_version(&self) -> Option { serde_json::from_value(self.result.as_ref().ok()?.get("protocolVersion")?.clone()).ok() } @@ -826,6 +882,14 @@ fn ensure_client_initialize_request( Ok(()) } +#[cfg(feature = "unstable_protocol_v2")] +fn initialize_request_id(message: &RawJsonRpcMessage) -> Option { + let RawJsonRpcMessage::Request(request) = message else { + return None; + }; + Some(request.id.clone()) +} + #[cfg(feature = "unstable_protocol_v2")] fn initialize_response_negotiated_v1(response: &InitializeResponse) -> bool { response.protocol_version() == Some(ProtocolVersion::V1) diff --git a/src/agent-client-protocol/tests/protocol_v2.rs b/src/agent-client-protocol/tests/protocol_v2.rs index 219835f..9ef82bc 100644 --- a/src/agent-client-protocol/tests/protocol_v2.rs +++ b/src/agent-client-protocol/tests/protocol_v2.rs @@ -127,6 +127,7 @@ fn runtime_flag_client_protocol_connector(enable_protocol_v2: bool) -> ClientPro struct InitializingV1Client { expected_session_id: &'static str, implementation_name: &'static str, + client_capabilities: Option, } impl InitializingV1Client { @@ -134,6 +135,7 @@ impl InitializingV1Client { Self { expected_session_id, implementation_name: "agent-client-protocol-test", + client_capabilities: None, } } @@ -144,23 +146,32 @@ impl InitializingV1Client { Self { expected_session_id, implementation_name, + client_capabilities: None, } } + + fn with_client_capabilities(mut self, client_capabilities: v1::ClientCapabilities) -> Self { + self.client_capabilities = Some(client_capabilities); + self + } } impl ConnectTo for InitializingV1Client { async fn connect_to(self, agent: impl ConnectTo) -> Result<(), Error> { let expected_session_id = self.expected_session_id; let implementation_name = self.implementation_name; + let client_capabilities = self.client_capabilities; Client .builder() .connect_with(agent, async move |cx| { - let initialize = cx - .send_request(v1::InitializeRequest::new(ProtocolVersion::V1).client_info( - v1::Implementation::new(implementation_name, env!("CARGO_PKG_VERSION")), - )) - .block_task() - .await?; + let mut request = v1::InitializeRequest::new(ProtocolVersion::V1).client_info( + v1::Implementation::new(implementation_name, env!("CARGO_PKG_VERSION")), + ); + if let Some(client_capabilities) = client_capabilities { + request = request.client_capabilities(client_capabilities); + } + + let initialize = cx.send_request(request).block_task().await?; assert_eq!(initialize.protocol_version, ProtocolVersion::V1); let session = cx @@ -839,6 +850,66 @@ async fn client_protocol_connector_falls_back_to_v1_when_agent_router_negotiates .await } +#[tokio::test(flavor = "current_thread")] +async fn client_protocol_connector_reuses_matching_connection_before_v1_fallback() +-> Result<(), Error> { + let connections = Arc::new(AtomicUsize::new(0)); + let agent_connections = Arc::clone(&connections); + + Client + .protocol_connector() + .with_v1(|| { + InitializingV1Client::new("v1-reused-session").with_client_capabilities( + v1::ClientCapabilities::new().session( + v1::ClientSessionCapabilities::new().config_options( + v1::SessionConfigOptionsCapabilities::new() + .boolean(v1::BooleanConfigOptionCapabilities::new()), + ), + ), + ) + }) + .with_v2(|| InitializingV2Client::new("v2-client-should-not-continue")) + .connect_to(move || { + let connection_number = agent_connections.fetch_add(1, Ordering::SeqCst) + 1; + let initialize_connection_number = connection_number; + let session_connection_number = connection_number; + + Agent.protocol_router().with_v1( + Agent + .builder() + .on_receive_request( + async move |initialize: v1::InitializeRequest, responder, _cx| { + assert_eq!(initialize_connection_number, 1); + assert_eq!(initialize.protocol_version, ProtocolVersion::V1); + let info = initialize + .client_info + .as_ref() + .expect("initialize should include client info"); + assert_eq!(&*info.name, "agent-client-protocol-test"); + + responder + .respond(v1::InitializeResponse::new(initialize.protocol_version)) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async move |request: v1::NewSessionRequest, responder, _cx| { + assert_eq!(session_connection_number, 1); + assert!(request.cwd.is_absolute()); + responder.respond(v1::NewSessionResponse::new(v1::SessionId::new( + "v1-reused-session", + ))) + }, + agent_client_protocol::on_receive_request!(), + ), + ) + }) + .await?; + + assert_eq!(connections.load(Ordering::SeqCst), 1); + Ok(()) +} + #[tokio::test(flavor = "current_thread")] async fn client_protocol_connector_reconnects_before_v1_fallback() -> Result<(), Error> { let connections = Arc::new(AtomicUsize::new(0));