From e8bb2a256c99f4c8f466313320b933ed2ff3b491 Mon Sep 17 00:00:00 2001 From: Gurdas Nijor Date: Mon, 10 Aug 2026 16:39:22 -0700 Subject: [PATCH] feat(acp): attach MCP servers to restored sessions --- .../tests/mcp_server_handler_chain.rs | 89 +++++++++ src/agent-client-protocol-cookbook/src/lib.rs | 25 +++ src/agent-client-protocol/CHANGELOG.md | 3 + .../src/concepts/sessions.rs | 23 +++ .../src/mcp_server/server.rs | 28 ++- src/agent-client-protocol/src/session.rs | 169 +++++++++++++++++- 6 files changed, 331 insertions(+), 6 deletions(-) diff --git a/src/agent-client-protocol-conductor/tests/mcp_server_handler_chain.rs b/src/agent-client-protocol-conductor/tests/mcp_server_handler_chain.rs index e41bf32a..aec732e7 100644 --- a/src/agent-client-protocol-conductor/tests/mcp_server_handler_chain.rs +++ b/src/agent-client-protocol-conductor/tests/mcp_server_handler_chain.rs @@ -92,6 +92,55 @@ struct ProxyWithMcpAndHandler { config: Arc, } +struct RestoredSessionMcpProxy; + +fn restored_session_server() +-> McpServer> { + McpServer::builder("test-server") + .tool_fn_mut( + "echo", + "Echoes back the input", + async |params: EchoParams, _cx| { + Ok(EchoOutput { + result: format!("Echo: {}", params.message), + }) + }, + agent_client_protocol::tool_fn_mut!(), + ) + .build() +} + +impl ConnectTo for RestoredSessionMcpProxy { + async fn connect_to( + self, + client: impl ConnectTo, + ) -> Result<(), agent_client_protocol::Error> { + Proxy + .builder() + .name("restored-session-mcp") + .on_receive_request_from( + Client, + async |request: LoadSessionRequest, responder, cx| { + cx.build_restored_session_from(request) + .with_mcp_server(restored_session_server())? + .on_proxy_session_start(responder, async |_session_id| Ok(())) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request_from( + Client, + async |request: ResumeSessionRequest, responder, cx| { + cx.build_restored_session_from(request) + .with_mcp_server(restored_session_server())? + .on_proxy_session_start(responder, async |_session_id| Ok(())) + }, + agent_client_protocol::on_receive_request!(), + ) + .connect_to(client) + .await + } +} + impl ConnectTo for ProxyWithMcpAndHandler { async fn connect_to( self, @@ -324,3 +373,43 @@ async fn test_mcp_server_injected_into_all_session_setup_requests() Ok(()) } + +/// Restored sessions can each receive an independently scoped MCP server. +#[tokio::test] +async fn test_per_session_mcp_server_attached_to_load_and_resume() +-> Result<(), agent_client_protocol::Error> { + let setup_requests = Arc::new(SetupRequests::default()); + let proxy = DynConnectTo::::new(RestoredSessionMcpProxy); + let agent = DynConnectTo::::new(SimpleAgent { + setup_requests: setup_requests.clone(), + }); + + run_test(vec![proxy], agent, async |connection_to_editor| { + recv(connection_to_editor.send_request(InitializeRequest::new(ProtocolVersion::V1))) + .await?; + + let cwd = PathBuf::from("/tmp"); + recv(connection_to_editor.send_request(LoadSessionRequest::new( + SessionId::new("loaded-session"), + cwd.clone(), + ))) + .await?; + recv(connection_to_editor.send_request(ResumeSessionRequest::new( + SessionId::new("resumed-session"), + cwd, + ))) + .await?; + + Ok(()) + }) + .await?; + + let server_ids = setup_requests.server_ids.lock().unwrap(); + assert_eq!(server_ids.len(), 2); + assert_ne!( + server_ids[0], server_ids[1], + "each restored session should advertise its own MCP server ID" + ); + + Ok(()) +} diff --git a/src/agent-client-protocol-cookbook/src/lib.rs b/src/agent-client-protocol-cookbook/src/lib.rs index 8eb0fe22..e5fc116f 100644 --- a/src/agent-client-protocol-cookbook/src/lib.rs +++ b/src/agent-client-protocol-cookbook/src/lib.rs @@ -655,6 +655,31 @@ pub mod per_session_mcp_server { //! ID-keyed state, preinstall a gate or placeholder that later handlers //! await, then populate it from the callback. //! + //! # Restoring stable v1 sessions + //! + //! `session/load` and `session/resume` can attach a new per-session server + //! through the corresponding restore builder: + //! + //! ```rust,ignore + //! Proxy.builder() + //! .on_receive_request_from( + //! Client, + //! async |request: LoadSessionRequest, responder, connection| { + //! connection + //! .build_restored_session_from(request) + //! .with_mcp_server(build_workspace_server())? + //! .on_proxy_session_start(responder, async |session_id| { + //! tracing::info!(%session_id, "Session restored"); + //! Ok(()) + //! }) + //! }, + //! agent_client_protocol::on_receive_request!(), + //! ); + //! ``` + //! + //! Use the same pattern with `ResumeSessionRequest`. Each invocation gets + //! its own MCP server ID and SDK-managed session route. + //! //! # Draft v2 pattern //! //! `Proxy.v2()` exposes the same non-blocking setup shape with v2 schema diff --git a/src/agent-client-protocol/CHANGELOG.md b/src/agent-client-protocol/CHANGELOG.md index 5ec93e8e..88117915 100644 --- a/src/agent-client-protocol/CHANGELOG.md +++ b/src/agent-client-protocol/CHANGELOG.md @@ -4,6 +4,9 @@ ### Added +- *(unstable)* Add `ConnectionTo::build_restored_session_from` for attaching + per-session MCP servers while proxying stable-v1 `session/load` and + `session/resume` requests. - *(unstable)* Expose programmatic tool-call names through the `unstable_tool_call_name` feature. - *(unstable)* Expose v1 and draft-v2 plan operations through the diff --git a/src/agent-client-protocol/src/concepts/sessions.rs b/src/agent-client-protocol/src/concepts/sessions.rs index 54782e51..8cc0a793 100644 --- a/src/agent-client-protocol/src/concepts/sessions.rs +++ b/src/agent-client-protocol/src/concepts/sessions.rs @@ -110,6 +110,29 @@ //! //! See the cookbook for detailed MCP server examples. //! +//! A proxy restoring a v1 session can attach a fresh MCP server to that +//! session with `build_restored_session_from`. The same builder accepts +//! `LoadSessionRequest` and `ResumeSessionRequest`: +//! +//! ```rust,ignore +//! Proxy.builder().on_receive_request_from( +//! Client, +//! async |request: LoadSessionRequest, responder, cx| { +//! cx.build_restored_session_from(request) +//! .with_mcp_server(session_tools())? +//! .on_proxy_session_start(responder, async |session_id| { +//! track_session(session_id); +//! Ok(()) +//! }) +//! }, +//! agent_client_protocol::on_receive_request!(), +//! ); +//! ``` +//! +//! The SDK installs session routing and the MCP handler before forwarding the +//! restore request. Successful handlers remain active for the connection +//! lifetime; failed restore requests remove the pending handlers. +//! //! # Non-Blocking Session Start //! //! If you're inside an `on_receive_*` callback and need to start a session, diff --git a/src/agent-client-protocol/src/mcp_server/server.rs b/src/agent-client-protocol/src/mcp_server/server.rs index 36e2d4d7..7fb8ce36 100644 --- a/src/agent-client-protocol/src/mcp_server/server.rs +++ b/src/agent-client-protocol/src/mcp_server/server.rs @@ -11,6 +11,27 @@ use crate::{ role, }; +#[cfg(feature = "unstable_mcp_over_acp")] +pub(crate) trait McpSessionSetupRequest { + fn mcp_servers_mut(&mut self) -> &mut Vec; +} + +#[cfg(feature = "unstable_mcp_over_acp")] +macro_rules! impl_mcp_session_setup_request { + ($($request:ty),+ $(,)?) => { + $( + impl McpSessionSetupRequest for $request { + fn mcp_servers_mut(&mut self) -> &mut Vec { + &mut self.mcp_servers + } + } + )+ + }; +} + +#[cfg(feature = "unstable_mcp_over_acp")] +impl_mcp_session_setup_request!(NewSessionRequest, LoadSessionRequest, ResumeSessionRequest); + #[cfg(feature = "unstable_mcp_over_acp")] use uuid::Uuid; @@ -313,15 +334,16 @@ where /// will no longer be received, so you need to keep this value alive as long as the session /// is in use. You can also invoke [`DynamicHandlerGuard::detach`] if you /// want to keep the handler registered for the life of the connection. - pub fn into_dynamic_handler( + pub fn into_dynamic_handler( self, - request: &mut NewSessionRequest, + request: &mut Request, cx: &ConnectionTo, ) -> Result, crate::Error> where Counterpart: HasPeer, + Request: McpSessionSetupRequest, { - self.append_declaration(&mut request.mcp_servers); + self.append_declaration(request.mcp_servers_mut()); cx.add_dynamic_handler(self.active_session) } } diff --git a/src/agent-client-protocol/src/session.rs b/src/agent-client-protocol/src/session.rs index 47200c07..cd061a41 100644 --- a/src/agent-client-protocol/src/session.rs +++ b/src/agent-client-protocol/src/session.rs @@ -10,9 +10,9 @@ use crate::{ }, role::{HasPeer, acp::ProxySessionMessages}, schema::v1::{ - ContentBlock, ContentChunk, NewSessionRequest, NewSessionResponse, PromptRequest, - PromptResponse, SessionId, SessionModeState, SessionNotification, SessionUpdate, - StopReason, + ContentBlock, ContentChunk, LoadSessionRequest, NewSessionRequest, NewSessionResponse, + PromptRequest, PromptResponse, ResumeSessionRequest, SessionId, SessionModeState, + SessionNotification, SessionUpdate, StopReason, }, util::{MatchDispatch, MatchDispatchFrom, run_until}, }; @@ -75,6 +75,23 @@ where SessionBuilder::new(self, request) } + /// Stable protocol v1 builder for restoring an existing session through a proxy. + /// + /// This accepts either a [`LoadSessionRequest`] or [`ResumeSessionRequest`]. + /// Use [`RestoredSessionBuilder::with_mcp_server`] to attach a server only to + /// this restored session, then + /// [`RestoredSessionBuilder::on_proxy_session_start`] to forward setup and + /// install session routing. + pub fn build_restored_session_from( + &self, + request: Request, + ) -> RestoredSessionBuilder + where + Request: RestoredSessionRequest, + { + RestoredSessionBuilder::new(self, request) + } + /// Given a session response received from the agent, /// attach a handler to process messages related to this session /// and let you access them. @@ -115,6 +132,152 @@ where } } +/// A stable protocol v1 request that restores an existing session. +/// +/// This trait is implemented by [`LoadSessionRequest`] and +/// [`ResumeSessionRequest`]. +pub trait RestoredSessionRequest: crate::JsonRpcRequest + Send + 'static { + /// The session being restored. + fn session_id(&self) -> &SessionId; +} + +impl RestoredSessionRequest for LoadSessionRequest { + fn session_id(&self) -> &SessionId { + &self.session_id + } +} + +impl RestoredSessionRequest for ResumeSessionRequest { + fn session_id(&self) -> &SessionId { + &self.session_id + } +} + +/// Stable protocol v1 builder for restoring a session through a proxy. +/// +/// The builder attaches per-session MCP servers and installs SDK-owned session +/// routing while forwarding either `session/load` or `session/resume`. +#[must_use = "call `on_proxy_session_start` to restore the session"] +#[derive(Debug)] +pub struct RestoredSessionBuilder +where + Counterpart: HasPeer, + Request: RestoredSessionRequest, + Run: RunWithConnectionTo, +{ + connection: ConnectionTo, + request: Request, + dynamic_handler_registrations: Vec>, + run: Run, +} + +impl RestoredSessionBuilder +where + Counterpart: HasPeer, + Request: RestoredSessionRequest, +{ + fn new(connection: &ConnectionTo, request: Request) -> Self { + Self { + connection: connection.clone(), + request, + dynamic_handler_registrations: Vec::new(), + run: NullRun, + } + } +} + +impl RestoredSessionBuilder +where + Counterpart: HasPeer, + Request: RestoredSessionRequest, + R: RunWithConnectionTo, +{ + /// Restore a proxy session and run a closure with its session ID. + /// + /// Session routing and attached MCP handlers are installed before the + /// setup request is sent. On success, they remain active for the lifetime + /// of the connection and the complete setup response is forwarded to the + /// client. On failure, the pending handlers are removed. + pub fn on_proxy_session_start( + self, + responder: Responder, + op: F, + ) -> Result<(), crate::Error> + where + Counterpart: HasPeer, + R: 'static, + F: FnOnce(SessionId) -> Fut + Send + 'static, + Fut: Future> + Send, + { + ensure_v1_session_protocol(&self.connection)?; + + let Self { + connection, + request, + dynamic_handler_registrations, + run, + } = self; + let session_id = request.session_id().clone(); + let session_route = + connection.add_dynamic_handler(ProxySessionMessages::new(session_id.clone()))?; + + connection + .send_ordered_request_to(Agent, request) + .forward_cancellation_from(responder.cancellation()) + .on_receiving_ok_result(responder, { + let connection = connection.clone(); + async move |response, responder| { + responder.respond(response)?; + connection.spawn(run.run_with_connection_to(connection.clone()))?; + dynamic_handler_registrations + .into_iter() + .for_each(DynamicHandlerGuard::detach); + session_route.detach(); + connection.spawn(async move { op(session_id).await }) + } + }) + } +} + +#[cfg(feature = "unstable_mcp_over_acp")] +macro_rules! impl_restored_session_mcp { + ($($request:ty),+ $(,)?) => { + $( + impl RestoredSessionBuilder + where + Counterpart: HasPeer, + R: RunWithConnectionTo, + { + /// Attach an MCP server only to this restored session. + pub fn with_mcp_server( + mut self, + mcp_server: McpServer, + ) -> Result< + RestoredSessionBuilder>, + crate::Error, + > + where + McpRun: RunWithConnectionTo, + { + let (handler, mcp_run) = mcp_server.into_handler_and_runner(); + self.dynamic_handler_registrations.push( + handler.into_dynamic_handler(&mut self.request, &self.connection)?, + ); + Ok(RestoredSessionBuilder { + connection: self.connection, + request: self.request, + dynamic_handler_registrations: self.dynamic_handler_registrations, + run: ChainRun::new(self.run, mcp_run), + }) + } + } + )+ + }; +} + +#[cfg(feature = "unstable_mcp_over_acp")] +impl_restored_session_mcp!(LoadSessionRequest, ResumeSessionRequest); + /// Stable protocol v1 session builder for a new session request. /// Allows you to add MCP servers or set other details for this session. ///