agent-client-protocol 5.0.0

//! MCP server construction, direct serving, and optional ACP attachment.

use std::{marker::PhantomData, sync::Arc};

use futures::{StreamExt, channel::mpsc};

use crate::{
    ConnectTo, Dispatch, DynConnectTo, Role,
    jsonrpc::run::{NullRun, RunWithConnectionTo},
    mcp_server::{McpConnectionContext, McpConnectionTo, McpServerConnect},
    role,
};

#[cfg(feature = "unstable_mcp_over_acp")]
use super::McpService;
#[cfg(feature = "unstable_mcp_over_acp")]
use uuid::Uuid;

#[cfg(feature = "unstable_mcp_over_acp")]
use crate::{
    Agent, Client, ConnectionTo, HandleDispatchFrom, Handled,
    jsonrpc::DynamicHandlerGuard,
    mcp_server::active_session::{McpActiveSession, V1McpProtocol},
    util::MatchDispatchFrom,
    v1::schema::{
        LoadSessionRequest, McpServer as SchemaMcpServer, McpServerAcp, McpServerAcpId,
        NewSessionRequest, ResumeSessionRequest,
    },
};

#[cfg(feature = "unstable_mcp_over_acp")]
use crate::{JsonRpcMessage, UntypedMessage};

#[cfg(feature = "unstable_mcp_over_acp")]
use crate::role::HasPeer;

#[cfg(all(feature = "unstable_mcp_over_acp", feature = "unstable_session_fork"))]
use crate::v1::schema::ForkSessionRequest;

/// Append an MCP server declaration to the `mcpServers` array of a raw
/// session setup request, leaving every other field as received.
#[cfg(feature = "unstable_mcp_over_acp")]
fn append_raw_declaration(
    request: &mut UntypedMessage,
    declaration: &impl serde::Serialize,
) -> Result<(), crate::Error> {
    let serde_json::Value::Object(params) = &mut request.params else {
        return Err(
            crate::Error::invalid_params().data("session setup parameters must be an object")
        );
    };
    let declaration =
        serde_json::to_value(declaration).map_err(crate::Error::into_internal_error)?;
    match params.get_mut("mcpServers") {
        Some(serde_json::Value::Array(servers)) => servers.push(declaration),
        Some(value) => *value = serde_json::Value::Array(vec![declaration]),
        None => {
            params.insert(
                "mcpServers".to_string(),
                serde_json::Value::Array(vec![declaration]),
            );
        }
    }
    Ok(())
}

/// A runtime-agnostic MCP server.
///
/// `McpServer` wraps an [`McpServerConnect`](`super::McpServerConnect`)
/// implementation. A server whose counterpart is [`role::mcp::Client`] can be
/// connected directly as a standalone MCP component. With the
/// `unstable_mcp_over_acp` feature, servers can instead be attached to ACP
/// session setup through `Builder::with_mcp_server` or
/// `SessionBuilder::with_mcp_server`. When `unstable_protocol_v2` is also
/// enabled, `Proxy.v2().with_mcp_server` attaches a server globally to draft
/// v2 setup requests. `v2::SessionBuilder::with_mcp_server` attaches one to a
/// single new v2 session, while `v2::ResumeSessionBuilder::with_mcp_server`
/// attaches one to a single resumed session. With `unstable_session_fork`,
/// `v2::ForkSessionBuilder::with_mcp_server` attaches one to a single forked
/// session.
///
/// # Creating an MCP Server
///
/// The `agent-client-protocol-rmcp` crate provides builder APIs for MCP tools
/// backed by the `rmcp` crate.
///
/// Or implement [`McpServerConnect`](`super::McpServerConnect`) for custom server behavior:
///
/// ```rust,ignore
/// let server = McpServer::new(MyCustomServerConnect, NullRun);
/// ```
pub struct McpServer<Counterpart: Role, Run = NullRun> {
    /// The host role that is serving up this MCP server
    phantom: PhantomData<Counterpart>,

    /// The "connect" instance
    connect: Arc<dyn McpServerConnect<Counterpart>>,

    #[cfg(feature = "unstable_mcp_over_acp")]
    service: Option<Arc<dyn McpService<Counterpart>>>,

    /// The runner is a task that should be run alongside the message handler.
    /// Some futures direct messages back through channels to this future which actually
    /// handles responding to the client.
    ///
    /// Some connector implementations use this to run support tasks alongside
    /// the message handler. It has no separate readiness protocol: communication
    /// primitives needed by `connect` must already exist when the server is
    /// constructed.
    runner: Run,
}

impl<Counterpart: Role + std::fmt::Debug, Run: std::fmt::Debug> std::fmt::Debug
    for McpServer<Counterpart, Run>
{
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("McpServer")
            .field("phantom", &self.phantom)
            .field("runner", &self.runner)
            .finish_non_exhaustive()
    }
}

impl<Counterpart: Role, Run> McpServer<Counterpart, Run>
where
    Run: RunWithConnectionTo<Counterpart>,
{
    /// Create an MCP server from something that implements the [`McpServerConnect`](`super::McpServerConnect`) trait.
    ///
    /// # See also
    ///
    /// See `agent-client-protocol-rmcp` to construct MCP servers from Rust code
    /// with `rmcp`.
    pub fn new(c: impl McpServerConnect<Counterpart>, runner: Run) -> Self {
        McpServer {
            phantom: PhantomData,
            connect: Arc::new(c),
            #[cfg(feature = "unstable_mcp_over_acp")]
            service: None,
            runner,
        }
    }

    /// Construct a reusable service for native MCP-over-ACP only.
    #[cfg(feature = "unstable_mcp_over_acp")]
    pub fn new_service(
        name: impl Into<String>,
        service: impl McpService<Counterpart>,
        runner: Run,
    ) -> Self {
        Self::new_service_with_standalone(service, NoStandalone { name: name.into() }, runner)
    }

    /// Construct a reusable native service and a separate standalone factory.
    ///
    /// Native requests do not create standalone sessions or perform hidden
    /// initialization/discovery exchanges.
    #[cfg(feature = "unstable_mcp_over_acp")]
    pub fn new_service_with_standalone(
        service: impl McpService<Counterpart>,
        standalone: impl McpServerConnect<Counterpart>,
        runner: Run,
    ) -> Self {
        Self {
            phantom: PhantomData,
            connect: Arc::new(standalone),
            service: Some(Arc::new(service)),
            runner,
        }
    }

    /// Split this MCP server into the message handler and a future that must be run while the handler is active.
    #[cfg(feature = "unstable_mcp_over_acp")]
    pub(crate) fn into_handler_and_runner(self) -> (McpSessionHandler<Counterpart>, Run)
    where
        Counterpart: HasPeer<Agent>,
    {
        let Self {
            phantom: _,
            connect,
            service,
            runner,
        } = self;
        let server_id = McpServerAcpId::new(format!("mcp-server:{}", Uuid::new_v4()));
        (
            McpSessionHandler::new_with_service(server_id, connect, service),
            runner,
        )
    }

    /// Split this MCP server into a protocol v2 session handler and its runner.
    #[cfg(all(feature = "unstable_mcp_over_acp", feature = "unstable_protocol_v2"))]
    pub(crate) fn into_v2_handler_and_runner(self) -> (V2McpSessionHandler<Counterpart>, Run)
    where
        Counterpart: HasPeer<Agent>,
    {
        let Self {
            phantom: _,
            connect,
            service,
            runner,
        } = self;
        let server_id = McpServerAcpId::new(format!("mcp-server:{}", Uuid::new_v4()));
        (
            V2McpSessionHandler::new_with_service(server_id, connect, service),
            runner,
        )
    }
}

#[cfg(feature = "unstable_mcp_over_acp")]
struct NoStandalone {
    name: String,
}

#[cfg(feature = "unstable_mcp_over_acp")]
impl<Counterpart: Role> McpServerConnect<Counterpart> for NoStandalone {
    fn name(&self) -> String {
        self.name.clone()
    }

    fn connect(&self, _context: McpConnectionTo<Counterpart>) -> DynConnectTo<role::mcp::Client> {
        struct Unavailable;
        impl ConnectTo<role::mcp::Client> for Unavailable {
            fn connect_to(
                self,
                _client: impl ConnectTo<role::mcp::Server>,
            ) -> impl std::future::Future<Output = Result<(), crate::Error>> + Send {
                std::future::ready(Err(crate::Error::method_not_found()
                    .data("this MCP service has no standalone transport adapter")))
            }
        }
        DynConnectTo::new(Unavailable)
    }
}

/// Message handler created from a [`McpServer`].
#[cfg(feature = "unstable_mcp_over_acp")]
pub(crate) struct McpSessionHandler<Counterpart: Role>
where
    Counterpart: HasPeer<Agent>,
{
    server_id: McpServerAcpId,
    connect: Arc<dyn McpServerConnect<Counterpart>>,
    active_session: McpActiveSession<Counterpart, V1McpProtocol>,
}

#[cfg(feature = "unstable_mcp_over_acp")]
impl<Counterpart: Role> McpSessionHandler<Counterpart>
where
    Counterpart: HasPeer<Agent>,
{
    fn new_with_service(
        server_id: McpServerAcpId,
        connect: Arc<dyn McpServerConnect<Counterpart>>,
        service: Option<Arc<dyn McpService<Counterpart>>>,
    ) -> Self {
        Self {
            active_session: McpActiveSession::new_with_service(
                server_id.clone(),
                connect.clone(),
                service,
            ),
            server_id,
            connect,
        }
    }

    fn declaration(&self) -> SchemaMcpServer {
        SchemaMcpServer::Acp(McpServerAcp::new(
            self.connect.name(),
            self.server_id.clone(),
        ))
    }

    /// Append this MCP server's native ACP declaration to a session setup request.
    fn append_declaration(&self, mcp_servers: &mut Vec<SchemaMcpServer>) {
        mcp_servers.push(self.declaration());
    }

    /// Whether `request` is a valid v1 session setup request this server
    /// attaches to. A setup request with malformed parameters is an error.
    fn validate_session_setup(request: &UntypedMessage) -> Result<bool, crate::Error> {
        match request.method() {
            "session/new" => {
                NewSessionRequest::parse_message(request.method(), request.params())?;
                Ok(true)
            }
            "session/load" => {
                LoadSessionRequest::parse_message(request.method(), request.params())?;
                Ok(true)
            }
            "session/resume" => {
                ResumeSessionRequest::parse_message(request.method(), request.params())?;
                Ok(true)
            }
            #[cfg(feature = "unstable_session_fork")]
            "session/fork" => {
                ForkSessionRequest::parse_message(request.method(), request.params())?;
                Ok(true)
            }
            _ => Ok(false),
        }
    }
}

/// Protocol v2 session adapter for the shared native MCP-over-ACP runtime.
#[cfg(all(feature = "unstable_mcp_over_acp", feature = "unstable_protocol_v2"))]
pub(crate) struct V2McpSessionHandler<Counterpart: Role>
where
    Counterpart: HasPeer<Agent>,
{
    server_id: McpServerAcpId,
    connect: Arc<dyn McpServerConnect<Counterpart>>,
    active_session: McpActiveSession<Counterpart, crate::mcp_server::active_session::V2McpProtocol>,
}

#[cfg(all(feature = "unstable_mcp_over_acp", feature = "unstable_protocol_v2"))]
impl<Counterpart: Role> V2McpSessionHandler<Counterpart>
where
    Counterpart: HasPeer<Agent>,
{
    #[cfg(test)]
    fn new(server_id: McpServerAcpId, connect: Arc<dyn McpServerConnect<Counterpart>>) -> Self {
        Self::new_with_service(server_id, connect, None)
    }

    fn new_with_service(
        server_id: McpServerAcpId,
        connect: Arc<dyn McpServerConnect<Counterpart>>,
        service: Option<Arc<dyn McpService<Counterpart>>>,
    ) -> Self {
        Self {
            active_session: McpActiveSession::new_with_service(
                server_id.clone(),
                connect.clone(),
                service,
            ),
            server_id,
            connect,
        }
    }

    fn declaration(&self) -> crate::v2::schema::McpServer {
        crate::v2::schema::McpServer::Acp(crate::v2::schema::McpServerAcp::new(
            self.connect.name(),
            crate::v2::schema::McpServerAcpId::from(self.server_id.0.clone()),
        ))
    }

    fn append_declaration(&self, mcp_servers: &mut Vec<crate::v2::schema::McpServer>) {
        mcp_servers.push(self.declaration());
    }

    fn validate_session_setup(request: &UntypedMessage) -> Result<bool, crate::Error> {
        match request.method() {
            "session/new" => {
                crate::v2::schema::NewSessionRequest::parse_message(
                    request.method(),
                    request.params(),
                )?;
                Ok(true)
            }
            "session/resume" => {
                crate::v2::schema::ResumeSessionRequest::parse_message(
                    request.method(),
                    request.params(),
                )?;
                Ok(true)
            }
            #[cfg(feature = "unstable_session_fork")]
            "session/fork" => {
                crate::v2::schema::ForkSessionRequest::parse_message(
                    request.method(),
                    request.params(),
                )?;
                Ok(true)
            }
            _ => Ok(false),
        }
    }

    /// Append this server's declaration to a raw v2 session setup request.
    ///
    /// MCP server names must be unique within `mcpServers`. If the client
    /// already declared a server with this server's name, this fails with
    /// `-32602` and leaves the request unchanged.
    fn append_declaration_to_raw(&self, request: &mut UntypedMessage) -> Result<(), crate::Error> {
        let serde_json::Value::Object(params) = &mut request.params else {
            return Err(
                crate::Error::invalid_params().data("session setup parameters must be an object")
            );
        };
        if let Some(serde_json::Value::Array(servers)) = params.get("mcpServers") {
            let name = self.connect.name();
            if servers.iter().any(|server| {
                server.get("name").and_then(serde_json::Value::as_str) == Some(name.as_str())
            }) {
                return Err(crate::Error::invalid_params().data(format!(
                    "MCP server `{name}` is configured by the proxy, so mcpServers cannot \
                     declare another server with that name; names must be unique within mcpServers"
                )));
            }
        }
        append_raw_declaration(request, &self.declaration())
    }

    /// Attach this server to a draft protocol v2 session setup request.
    ///
    /// MCP server names must be unique within a setup request's
    /// `mcpServers`, and agents reject duplicates with `-32602`. A server
    /// whose name is already declared fails here with `-32602` instead of
    /// sending a request the agent must reject.
    pub fn into_dynamic_handler(
        self,
        mcp_servers: &mut Vec<crate::v2::schema::McpServer>,
        cx: &crate::v2::ConnectionTo<Counterpart>,
    ) -> Result<DynamicHandlerGuard<Counterpart>, crate::Error> {
        let name = self.connect.name();
        if mcp_servers
            .iter()
            .any(|server| v2_mcp_server_name(server) == Some(name.as_str()))
        {
            return Err(crate::Error::invalid_params().data(format!(
                "MCP server name `{name}` is already declared; names must be unique within mcpServers"
            )));
        }
        self.append_declaration(mcp_servers);
        let cleanup = self.active_session.cleanup();
        Ok(cx
            .add_dynamic_handler(self.active_session)?
            .with_cleanup(cleanup))
    }
}

/// The `name` of a v2 MCP server declaration, if it has one.
#[cfg(all(feature = "unstable_mcp_over_acp", feature = "unstable_protocol_v2"))]
fn v2_mcp_server_name(server: &crate::v2::schema::McpServer) -> Option<&str> {
    use crate::v2::schema::McpServer;
    match server {
        McpServer::Http(server) => Some(&server.name),
        McpServer::Acp(server) => Some(&server.name),
        McpServer::Stdio(server) => Some(&server.name),
        McpServer::Other(server) => server
            .fields
            .get("name")
            .and_then(serde_json::Value::as_str),
        _ => None,
    }
}

#[cfg(all(feature = "unstable_mcp_over_acp", feature = "unstable_protocol_v2"))]
impl<Counterpart: Role> HandleDispatchFrom<Counterpart> for V2McpSessionHandler<Counterpart>
where
    Counterpart: HasPeer<Client> + HasPeer<Agent>,
{
    async fn handle_dispatch_from(
        &mut self,
        message: Dispatch,
        cx: ConnectionTo<Counterpart>,
    ) -> Result<Handled<Dispatch>, crate::Error> {
        MatchDispatchFrom::new(message, &cx)
            .if_request_from(Client, async |mut request: UntypedMessage, responder| {
                if !Self::validate_session_setup(&request)? {
                    return Ok(Handled::No {
                        message: (request, responder),
                        retry: false,
                    });
                }

                if let Err(error) = self.append_declaration_to_raw(&mut request) {
                    // Answer the client directly rather than forwarding a
                    // request the agent must reject.
                    responder.respond_with_error(error)?;
                    return Ok(Handled::Yes);
                }
                Ok(Handled::No {
                    message: (request, responder),
                    retry: false,
                })
            })
            .await
            .otherwise_delegate(&mut self.active_session)
            .await
    }

    fn describe_chain(&self) -> impl std::fmt::Debug {
        format!("V2McpServer({})", self.connect.name())
    }
}

#[cfg(feature = "unstable_mcp_over_acp")]
impl<Counterpart: Role> McpSessionHandler<Counterpart>
where
    Counterpart: HasPeer<Agent>,
{
    /// Attach this server to the new session, spawning off a dynamic handler that will
    /// manage requests coming from this session.
    ///
    /// # Return value
    ///
    /// Returns a [`DynamicHandlerGuard`] for the handler that intercepts messages
    /// related to this MCP server. Once the value is dropped, the MCP server messages
    /// 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(
        self,
        request: &mut NewSessionRequest,
        cx: &ConnectionTo<Counterpart>,
    ) -> Result<DynamicHandlerGuard<Counterpart>, crate::Error>
    where
        Counterpart: HasPeer<Agent>,
    {
        self.append_declaration(&mut request.mcp_servers);
        let cleanup = self.active_session.cleanup();
        Ok(cx
            .add_dynamic_handler(self.active_session)?
            .with_cleanup(cleanup))
    }
}

#[cfg(feature = "unstable_mcp_over_acp")]
impl<Counterpart: Role> HandleDispatchFrom<Counterpart> for McpSessionHandler<Counterpart>
where
    Counterpart: HasPeer<Client> + HasPeer<Agent>,
{
    async fn handle_dispatch_from(
        &mut self,
        message: Dispatch,
        cx: ConnectionTo<Counterpart>,
    ) -> Result<Handled<Dispatch>, crate::Error> {
        // Edit the raw request so fields the v1 schema types do not model
        // reach the agent unchanged.
        MatchDispatchFrom::new(message, &cx)
            .if_request_from(Client, async |mut request: UntypedMessage, responder| {
                if Self::validate_session_setup(&request)? {
                    append_raw_declaration(&mut request, &self.declaration())?;
                }
                Ok(Handled::No {
                    message: (request, responder),
                    retry: false,
                })
            })
            .await
            .otherwise_delegate(&mut self.active_session)
            .await
    }

    fn describe_chain(&self) -> impl std::fmt::Debug {
        format!("McpServer({})", self.connect.name())
    }
}

impl<Run> ConnectTo<role::mcp::Client> for McpServer<role::mcp::Client, Run>
where
    Run: RunWithConnectionTo<role::mcp::Client> + 'static,
{
    async fn connect_to(
        self,
        client: impl ConnectTo<role::mcp::Server>,
    ) -> Result<(), crate::Error> {
        let Self {
            connect,
            runner,
            phantom: _,
            #[cfg(feature = "unstable_mcp_over_acp")]
                service: _,
        } = self;

        let (tx, mut rx) = mpsc::unbounded();

        role::mcp::Server
            .builder()
            .with_runner(runner)
            .on_receive_dispatch(
                async |message_from_client: Dispatch, _cx| {
                    tx.unbounded_send(message_from_client)
                        .map_err(|_| crate::util::internal_error("nobody listening to mcp server"))
                },
                crate::on_receive_dispatch!(),
            )
            .with_spawned(async move |connection_to_client| {
                let spawned_server: DynConnectTo<role::mcp::Client> =
                    connect.connect(McpConnectionTo {
                        context: McpConnectionContext::Standalone,
                        connection: connection_to_client.clone(),
                        #[cfg(feature = "unstable_mcp_over_acp")]
                        cleanup: None,
                    });

                role::mcp::Client
                    .builder()
                    .on_receive_dispatch(
                        async |message_from_server: Dispatch, _| {
                            // when we receive a message from the server, fwd to the client
                            connection_to_client.send_proxied_message(message_from_server)
                        },
                        crate::on_receive_dispatch!(),
                    )
                    .connect_with(spawned_server, async |connection_to_server| {
                        while let Some(message_from_client) = rx.next().await {
                            connection_to_server.send_proxied_message(message_from_client)?;
                        }
                        Ok(())
                    })
                    .await
            })
            .connect_to(client)
            .await
    }
}

#[cfg(all(
    test,
    feature = "unstable_mcp_over_acp",
    feature = "unstable_protocol_v2"
))]
mod tests {
    use std::{collections::BTreeMap, path::PathBuf, sync::Arc};

    use serde::Serialize;
    use serde_json::{Value, json};

    use super::V2McpSessionHandler;
    use crate::{
        Conductor, DynConnectTo, Error, UntypedMessage,
        mcp_server::{McpConnectionTo, McpServerConnect},
        role,
        v1::schema::McpServerAcpId,
        v2::schema::{self as v2, McpServer},
    };

    struct UnusedMcpConnect;

    impl McpServerConnect<Conductor> for UnusedMcpConnect {
        fn name(&self) -> String {
            "global-v2-server".to_owned()
        }

        fn connect(&self, _context: McpConnectionTo<Conductor>) -> DynConnectTo<role::mcp::Client> {
            panic!("declaration tests must not connect to the MCP server")
        }
    }

    fn handler() -> V2McpSessionHandler<Conductor> {
        V2McpSessionHandler::new(
            McpServerAcpId::new("global-v2-server-id"),
            Arc::new(UnusedMcpConnect),
        )
    }

    fn existing_server() -> McpServer {
        McpServer::Other(v2::OtherMcpServer::new(
            "_future_transport",
            BTreeMap::from([("futureOption".to_owned(), json!({ "nested": true }))]),
        ))
    }

    fn meta() -> v2::Meta {
        v2::Meta::from_iter([("preserved".to_owned(), json!({ "nested": true }))])
    }

    fn assert_raw_append_preserves_params(
        handler: &V2McpSessionHandler<Conductor>,
        method: &str,
        params: impl Serialize,
    ) -> Result<v2::McpServerAcpId, Error> {
        let mut params = serde_json::to_value(params)?;
        let Value::Object(params_object) = &mut params else {
            panic!("session setup params should serialize as an object");
        };
        params_object.insert(
            "_futureSessionField".to_owned(),
            json!({ "must": ["remain", "untouched"] }),
        );

        let mut expected = params.clone();
        expected
            .get_mut("mcpServers")
            .and_then(Value::as_array_mut)
            .expect("test request should contain mcpServers")
            .push(serde_json::to_value(handler.declaration())?);

        let mut request = UntypedMessage::new(method, params)?;
        assert!(V2McpSessionHandler::<Conductor>::validate_session_setup(
            &request
        )?);
        handler.append_declaration_to_raw(&mut request)?;

        assert_eq!(
            request.params, expected,
            "global attachment must only append its declaration"
        );

        let appended = request
            .params
            .get("mcpServers")
            .and_then(Value::as_array)
            .and_then(|servers| servers.last())
            .cloned()
            .expect("global declaration should be appended");
        match serde_json::from_value::<McpServer>(appended)? {
            McpServer::Acp(server) => {
                assert_eq!(server.name, "global-v2-server");
                Ok(server.server_id)
            }
            server => panic!("expected an ACP server declaration, got {server:?}"),
        }
    }

    #[test]
    fn v2_global_mcp_declaration_preserves_all_session_setup_params() -> Result<(), Error> {
        let handler = handler();
        let cwd = PathBuf::from("/tmp/global-v2-mcp");
        let additional_directory = PathBuf::from("/tmp/global-v2-mcp-additional");
        let session_id = v2::SessionId::new("session-to-resume");
        let existing_server = existing_server();

        let new_server_id = assert_raw_append_preserves_params(
            &handler,
            "session/new",
            v2::NewSessionRequest::new(cwd.clone())
                .additional_directories([additional_directory.clone()])
                .mcp_servers(vec![existing_server.clone()])
                .meta(meta()),
        )?;

        let resume_server_id = assert_raw_append_preserves_params(
            &handler,
            "session/resume",
            v2::ResumeSessionRequest::new(session_id.clone(), cwd.clone())
                .additional_directories([additional_directory.clone()])
                .mcp_servers(vec![existing_server.clone()])
                .replay_from(v2::ReplayFrom::Start(
                    v2::ReplayFromStart::new().meta(meta()),
                ))
                .meta(meta()),
        )?;
        assert_eq!(resume_server_id, new_server_id);

        #[cfg(feature = "unstable_session_fork")]
        {
            let fork_server_id = assert_raw_append_preserves_params(
                &handler,
                "session/fork",
                v2::ForkSessionRequest::new(session_id, cwd)
                    .additional_directories([additional_directory])
                    .mcp_servers(vec![existing_server])
                    .meta(meta()),
            )?;
            assert_eq!(fork_server_id, new_server_id);
        }

        Ok(())
    }

    fn assert_raw_append_rejects_name_collision(
        handler: &V2McpSessionHandler<Conductor>,
        method: &str,
        params: impl Serialize,
    ) -> Result<(), Error> {
        let params = serde_json::to_value(params)?;
        let mut request = UntypedMessage::new(method, params.clone())?;
        assert!(V2McpSessionHandler::<Conductor>::validate_session_setup(
            &request
        )?);

        let error = handler
            .append_declaration_to_raw(&mut request)
            .expect_err("a colliding MCP server name must be rejected");
        assert_eq!(error.code, crate::ErrorCode::INVALID_PARAMS);
        let message = error
            .data
            .as_ref()
            .and_then(Value::as_str)
            .expect("the error should explain the collision");
        assert!(
            message.contains("`global-v2-server`") && message.contains("configured by the proxy"),
            "unexpected error message: {message}"
        );
        assert_eq!(
            request.params, params,
            "a rejected request must not be modified"
        );
        Ok(())
    }

    #[test]
    fn v2_global_mcp_declaration_rejects_client_server_with_same_name() -> Result<(), Error> {
        let handler = handler();
        let cwd = PathBuf::from("/tmp/global-v2-mcp");
        let session_id = v2::SessionId::new("session-to-resume");
        let servers = vec![
            existing_server(),
            McpServer::Http(v2::McpServerHttp::new(
                "global-v2-server",
                "https://example.com/mcp",
            )),
        ];

        assert_raw_append_rejects_name_collision(
            &handler,
            "session/new",
            v2::NewSessionRequest::new(cwd.clone()).mcp_servers(servers.clone()),
        )?;
        assert_raw_append_rejects_name_collision(
            &handler,
            "session/resume",
            v2::ResumeSessionRequest::new(session_id.clone(), cwd.clone())
                .mcp_servers(servers.clone()),
        )?;
        #[cfg(feature = "unstable_session_fork")]
        assert_raw_append_rejects_name_collision(
            &handler,
            "session/fork",
            v2::ForkSessionRequest::new(session_id, cwd).mcp_servers(servers),
        )?;

        Ok(())
    }

    #[test]
    fn v2_global_mcp_declaration_accepts_client_servers_with_other_names() -> Result<(), Error> {
        let handler = handler();
        let server_id = assert_raw_append_preserves_params(
            &handler,
            "session/new",
            v2::NewSessionRequest::new(PathBuf::from("/tmp/global-v2-mcp")).mcp_servers(vec![
                McpServer::Http(v2::McpServerHttp::new(
                    "client-server",
                    "https://example.com/mcp",
                )),
            ]),
        )?;
        assert_eq!(server_id.to_string(), "global-v2-server-id");
        Ok(())
    }

    #[test]
    fn v2_global_mcp_handler_ignores_non_setup_requests() -> Result<(), Error> {
        let request = UntypedMessage::new(
            "session/prompt",
            json!({
                "sessionId": "session-to-prompt",
                "prompt": []
            }),
        )?;

        assert!(!V2McpSessionHandler::<Conductor>::validate_session_setup(
            &request
        )?);
        assert_eq!(request.method(), "session/prompt");
        Ok(())
    }
}

#[cfg(all(test, feature = "unstable_mcp_over_acp"))]
mod v1_tests {
    use std::{path::PathBuf, sync::Arc};

    use serde::Serialize;
    use serde_json::{Value, json};

    use super::{McpSessionHandler, append_raw_declaration};
    use crate::{
        Conductor, DynConnectTo, Error, UntypedMessage,
        mcp_server::{McpConnectionTo, McpServerConnect},
        role,
        v1::schema::{
            LoadSessionRequest, McpServer, McpServerAcpId, McpServerHttp, NewSessionRequest,
            ResumeSessionRequest, SessionId,
        },
    };

    struct UnusedMcpConnect;

    impl McpServerConnect<Conductor> for UnusedMcpConnect {
        fn name(&self) -> String {
            "global-v1-server".to_owned()
        }

        fn connect(&self, _context: McpConnectionTo<Conductor>) -> DynConnectTo<role::mcp::Client> {
            panic!("declaration tests must not connect to the MCP server")
        }
    }

    fn handler() -> McpSessionHandler<Conductor> {
        McpSessionHandler::new_with_service(
            McpServerAcpId::new("global-v1-server-id"),
            Arc::new(UnusedMcpConnect),
            None,
        )
    }

    fn assert_raw_append_preserves_params(
        handler: &McpSessionHandler<Conductor>,
        method: &str,
        params: impl Serialize,
    ) -> Result<(), Error> {
        let mut params = serde_json::to_value(params)?;
        let Value::Object(params_object) = &mut params else {
            panic!("session setup params should serialize as an object");
        };
        // Fields the v1 schema types do not model must reach the agent.
        params_object.insert("futureField".to_owned(), json!("kept"));
        params_object.insert(
            "_futureSessionField".to_owned(),
            json!({ "must": ["remain", "untouched"] }),
        );

        let mut expected = params.clone();
        expected
            .get_mut("mcpServers")
            .and_then(Value::as_array_mut)
            .expect("test request should contain mcpServers")
            .push(serde_json::to_value(handler.declaration())?);

        let mut request = UntypedMessage::new(method, params)?;
        assert!(McpSessionHandler::<Conductor>::validate_session_setup(
            &request
        )?);
        append_raw_declaration(&mut request, &handler.declaration())?;

        assert_eq!(
            request.params, expected,
            "global attachment must only append its declaration"
        );
        Ok(())
    }

    #[test]
    fn v1_global_mcp_declaration_preserves_all_session_setup_params() -> Result<(), Error> {
        let handler = handler();
        let cwd = PathBuf::from("/tmp/global-v1-mcp");
        let session_id = SessionId::new("session-to-restore");
        let existing = vec![McpServer::Http(McpServerHttp::new(
            "client-server",
            "https://example.com/mcp",
        ))];

        assert_raw_append_preserves_params(
            &handler,
            "session/new",
            NewSessionRequest::new(cwd.clone()).mcp_servers(existing.clone()),
        )?;
        assert_raw_append_preserves_params(
            &handler,
            "session/load",
            LoadSessionRequest::new(session_id.clone(), cwd.clone()).mcp_servers(existing.clone()),
        )?;
        assert_raw_append_preserves_params(
            &handler,
            "session/resume",
            ResumeSessionRequest::new(session_id, cwd).mcp_servers(existing),
        )?;
        Ok(())
    }

    #[test]
    fn v1_global_mcp_declaration_ignores_other_requests() -> Result<(), Error> {
        let request = UntypedMessage::new("session/prompt", json!({ "sessionId": "s" }))?;
        assert!(!McpSessionHandler::<Conductor>::validate_session_setup(
            &request
        )?);
        Ok(())
    }
}