diff --git a/md/protocol-v2.md b/md/protocol-v2.md index 7772e47..25af4af 100644 --- a/md/protocol-v2.md +++ b/md/protocol-v2.md @@ -78,6 +78,93 @@ 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()` 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 +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. + +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 6cb2c7f..7d9453a 100644 --- a/src/agent-client-protocol/src/lib.rs +++ b/src/agent-client-protocol/src/lib.rs @@ -103,6 +103,10 @@ pub use jsonrpc::{ }; 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::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 d5919f8..32249a3 100644 --- a/src/agent-client-protocol/src/role/acp.rs +++ b/src/agent-client-protocol/src/role/acp.rs @@ -1,10 +1,23 @@ 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}; 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}; use crate::{ConnectTo, ConnectionTo, Dispatch, HandleDispatchFrom, Handled, Role, RoleId}; /// The client role - typically an IDE or CLI that controls an agent. @@ -58,6 +71,24 @@ impl Client { self.builder().v2_client() } + /// Create a connector that chooses between configured protocol implementations. + /// + /// 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 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. + #[cfg(feature = "unstable_protocol_v2")] + #[must_use] + pub fn protocol_connector(self) -> ClientProtocolConnector { + ClientProtocolConnector::new() + } + /// Connect to `agent` and run `main_fn` with the [`ConnectionTo`]. /// Returns the result of `main_fn` (or an error if something goes wrong). /// @@ -71,6 +102,149 @@ impl Client { } } +/// Client connector that opens an agent connection with a configured protocol implementation. +/// +/// 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 may +/// require a fresh connection initialized by the v1 implementation. +#[cfg(feature = "unstable_protocol_v2")] +#[derive(Debug, Default)] +pub struct ClientProtocolConnector { + v1: Option>, + v2: Option>, +} + +#[cfg(feature = "unstable_protocol_v2")] +impl ClientProtocolConnector { + /// Create an empty client protocol connector. + #[must_use] + pub fn new() -> Self { + Self::default() + } + + /// Return this connector with an ACP v1 implementation factory configured. + #[must_use] + 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 connector with an ACP v2 implementation factory configured. + #[must_use] + pub fn with_v2(mut self, client: impl FnMut() -> C + Send + 'static) -> Self + where + C: ConnectTo, + { + self.v2 = Some(DynConnectToFactory::new(client)); + self + } + + /// 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 connector has no configured ACP protocol implementations")); + }; + + 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 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) = + 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; + } + + 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)() + } +} + +#[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() + } +} + impl HasPeer for Client { fn remote_style(&self, _peer: Client) -> RemoteStyle { RemoteStyle::Counterpart @@ -137,6 +311,588 @@ 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 client = RunningProtocolPeer::new(client); + let Some((mut first_message, client)) = client.next_message().await? else { + return Ok(()); + }; + 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) => { + return reject_initialize(client, &first_message, error).await; + } + }; + let Some(agent) = selected.take_agent(self) else { + let error = selected.unsupported_error(supported); + return reject_initialize(client, &first_message, error).await; + }; + + let agent = RunningProtocolPeer::new(agent); + agent.send(first_message)?; + pipe_protocol_peers_until_closed(client, agent).await + } +} + +#[cfg(feature = "unstable_protocol_v2")] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum AgentProtocol { + V1, + V2, +} + +#[cfg(feature = "unstable_protocol_v2")] +impl AgentProtocol { + 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)?; + rewrite_initialize_params(params, requested, selected)?; + + 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, + 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, + 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 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>, + 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) +} + +#[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>, + 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 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")] +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, + 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 { diff --git a/src/agent-client-protocol/tests/protocol_v2.rs b/src/agent-client-protocol/tests/protocol_v2.rs index ac660b8..9ef82bc 100644 --- a/src/agent-client-protocol/tests/protocol_v2.rs +++ b/src/agent-client-protocol/tests/protocol_v2.rs @@ -1,16 +1,24 @@ #![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, Builder, Client, ConnectTo, Error, JsonRpcMessage, JsonRpcRequest, JsonRpcResponse, - NullHandler, RawJsonRpcMessage, Role, UntypedRole, + Agent, AgentProtocolRouter, Builder, ByteStreams, Client, ClientProtocolConnector, 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)] @@ -72,11 +80,221 @@ 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 + } +} + +fn runtime_flag_client_protocol_connector(enable_protocol_v2: bool) -> ClientProtocolConnector { + let client = Client + .protocol_connector() + .with_v1(|| InitializingV1Client::new("v1-client-connector-session")); + + if enable_protocol_v2 { + client.with_v2(|| InitializingV2Client::new("v2-client-connector-session")) + } else { + client + } +} + +struct InitializingV1Client { + expected_session_id: &'static str, + implementation_name: &'static str, + client_capabilities: Option, +} + +impl InitializingV1Client { + fn new(expected_session_id: &'static str) -> Self { + Self { + expected_session_id, + implementation_name: "agent-client-protocol-test", + client_capabilities: None, + } + } + + fn with_implementation_name( + expected_session_id: &'static str, + implementation_name: &'static str, + ) -> Self { + 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 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 + .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| { @@ -596,6 +814,580 @@ async fn v2_client_and_agent_negotiate_v2() -> Result<(), Error> { .await } +#[tokio::test(flavor = "current_thread")] +async fn client_protocol_connector_routes_to_v2_client_for_v2_agent() -> Result<(), Error> { + Client + .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_connector_does_not_retry_after_v2_initialize_rejection() +-> Result<(), Error> { + Client + .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_connector_falls_back_to_v1_when_agent_router_negotiates_v1() +-> Result<(), Error> { + Client + .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_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)); + 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_connector(true) + .connect_to(|| v2_agent_with_session("v2-client-connector-session")) + .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_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 + .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_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 + .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()