agent_client_protocol/session/

v2.rs

use std::{future::Future, path::Path};

use futures::{
    channel::oneshot,
    future::{self, Either},
};

use crate::{
    Agent, Client, ConnectionTo as RawConnectionTo, DynamicHandlerGuard, JsonRpcRequest, Responder,
    SentRequest, V2,
    jsonrpc::run::{NullRun, RunWithConnectionTo},
    role::{HasPeer, acp::ProxySessionMessages},
    v2::ConnectionTo,
    v2::schema as v2,
};

#[cfg(feature = "unstable_mcp_over_acp")]
use crate::{jsonrpc::run::ChainRun, mcp_server::McpServer};

use super::{
    SessionCleanup, close_session_registrations, drive_session_cleanup,
    run_attached_session_runner, session_cleanup, session_runner_connection, wait_session_cleanup,
};

async fn run_pending_session_setup<Counterpart, Run>(
    connection: RawConnectionTo<Counterpart>,
    run: Run,
    started_tx: oneshot::Sender<Result<(), crate::Error>>,
    promotion_rx: oneshot::Receiver<()>,
    cleanup: SessionCleanup,
) -> Result<(), crate::Error>
where
    Counterpart: HasPeer<Agent>,
    Run: RunWithConnectionTo<Counterpart>,
{
    let (connection, error_scope) = session_runner_connection(connection, &cleanup);
    let mut run = Box::pin(run.run_with_connection_to(connection));
    let first_poll =
        future::poll_fn(|cx| std::task::Poll::Ready(std::future::Future::poll(run.as_mut(), cx)))
            .await;
    // A composed runner may have observed an immediate error but remain Pending
    // while its sibling drives owned cleanup. It still failed before publication.
    if let Some(error) = error_scope.error() {
        close_session_registrations(&cleanup);
        if first_poll.is_pending() {
            drop(drive_session_cleanup(run, cleanup).await);
        } else {
            wait_session_cleanup(cleanup).await;
        }
        drop(started_tx.send(Err(error)));
        return Ok(());
    }
    let readiness = match &first_poll {
        std::task::Poll::Ready(result) => result.clone(),
        std::task::Poll::Pending => Ok(()),
    };
    drop(started_tx.send(readiness));

    match first_poll {
        std::task::Poll::Ready(Ok(())) => {
            if promotion_rx.await.is_err() {
                close_session_registrations(&cleanup);
                wait_session_cleanup(cleanup).await;
            }
            Ok(())
        }
        // The request has not been published yet, so report an immediate
        // startup failure through its readiness result without failing the
        // whole connection.
        std::task::Poll::Ready(Err(_)) => {
            close_session_registrations(&cleanup);
            wait_session_cleanup(cleanup).await;
            Ok(())
        }
        std::task::Poll::Pending => match future::select(run, promotion_rx).await {
            Either::Left((result, promotion_rx)) => {
                // A Pending first poll releases session setup for publication.
                // From that point onward the agent may already be using an
                // attachment, so runner failures are connection-fatal just as
                // they are for other connection runners.
                if result.is_err() || promotion_rx.await.is_err() {
                    close_session_registrations(&cleanup);
                    wait_session_cleanup(cleanup).await;
                }
                result
            }
            Either::Right((Ok(()), run)) => {
                run_attached_session_runner(run, cleanup, error_scope).await
            }
            Either::Right((Err(_), run)) => {
                close_session_registrations(&cleanup);
                run_attached_session_runner(run, cleanup, error_scope).await
            }
        },
    }
}

fn send_session_setup<Counterpart, Request, Run>(
    connection: ConnectionTo<Counterpart>,
    request: Request,
    dynamic_handler_registrations: Vec<DynamicHandlerGuard<Counterpart>>,
    run: Run,
    ordered: bool,
) -> SentRequest<Request::Response, V2>
where
    Counterpart: HasPeer<Agent>,
    Request: JsonRpcRequest + 'static,
    Request::Response: 'static,
    Run: RunWithConnectionTo<Counterpart> + 'static,
{
    let raw_connection = connection.raw_connection().clone();
    if dynamic_handler_registrations.is_empty() {
        drop(run);
        if ordered {
            raw_connection.send_ordered_request_to(Agent, request)
        } else {
            connection.send_request_to(Agent, request)
        }
    } else {
        let handlers_ready = raw_connection.dynamic_handler_barrier();
        let (runner_started_tx, runner_started_rx) = oneshot::channel();
        let (promotion_tx, promotion_rx) = oneshot::channel();
        let cleanup = session_cleanup(&dynamic_handler_registrations);
        let runner_started = match raw_connection.spawn(run_pending_session_setup(
            raw_connection.clone(),
            run,
            runner_started_tx,
            promotion_rx,
            cleanup,
        )) {
            Ok(()) => Either::Left(async move {
                runner_started_rx.await.map_err(|error| {
                    crate::util::internal_error(format!(
                        "session setup runner stopped before its initial poll: {error}"
                    ))
                })?
            }),
            Err(error) => Either::Right(future::ready(Err(error))),
        };
        let readiness = async move {
            future::try_join(handlers_ready, runner_started).await?;
            Ok(())
        };
        let response_hook = move |_response: &Request::Response| {
            promotion_tx.send(()).map_err(|()| {
                crate::util::internal_error("session setup runner stopped before setup completed")
            })?;
            dynamic_handler_registrations
                .into_iter()
                .for_each(DynamicHandlerGuard::detach);
            Ok(())
        };

        if ordered {
            raw_connection.send_ordered_request_to_with_response_hook_after(
                Agent,
                request,
                readiness,
                response_hook,
            )
        } else {
            raw_connection.send_request_to_with_response_hook_after(
                Agent,
                request,
                readiness,
                response_hook,
            )
        }
    }
}

impl<Counterpart> ConnectionTo<Counterpart>
where
    Counterpart: HasPeer<Agent>,
{
    /// Build a draft protocol v2 `session/new` request.
    ///
    /// `cwd` must be an absolute path: ACP requires one, and the SDK sends it
    /// without checking.
    pub fn build_session(&self, cwd: impl AsRef<Path>) -> SessionBuilder<Counterpart> {
        SessionBuilder::new(self, v2::NewSessionRequest::new(cwd.as_ref()))
    }

    /// Build a draft protocol v2 session using the current working directory.
    ///
    /// Returns an error if the current directory cannot be determined.
    pub fn build_session_cwd(&self) -> Result<SessionBuilder<Counterpart>, crate::Error> {
        let cwd = std::env::current_dir().map_err(|error| {
            crate::Error::internal_error().data(format!("cannot get current directory: {error}"))
        })?;
        Ok(self.build_session(cwd))
    }

    /// Build a draft protocol v2 session from an existing `session/new` request.
    pub fn build_session_from(
        &self,
        request: v2::NewSessionRequest,
    ) -> SessionBuilder<Counterpart> {
        SessionBuilder::new(self, request)
    }

    /// Build an unstable draft protocol v2 `session/fork` request.
    ///
    /// This helper is available with the `unstable_session_fork` feature. Call
    /// [`ForkSessionBuilder::start_session`] to publish the request and
    /// obtain a command handle for the newly created fork.
    ///
    /// `cwd` must be an absolute path: ACP requires one, and the SDK sends it
    /// without checking.
    #[cfg(feature = "unstable_session_fork")]
    pub fn fork_session(
        &self,
        session_id: impl Into<v2::SessionId>,
        cwd: impl AsRef<Path>,
    ) -> ForkSessionBuilder<Counterpart> {
        self.fork_session_from(v2::ForkSessionRequest::new(session_id, cwd.as_ref()))
    }

    /// Build an unstable draft protocol v2 `session/fork` request from an
    /// existing request.
    ///
    /// This helper is available with the `unstable_session_fork` feature. Call
    /// [`ForkSessionBuilder::start_session`] to publish the request.
    #[cfg(feature = "unstable_session_fork")]
    pub fn fork_session_from(
        &self,
        request: v2::ForkSessionRequest,
    ) -> ForkSessionBuilder<Counterpart> {
        ForkSessionBuilder::new(self, request)
    }

    /// Build a draft protocol v2 `session/resume` request.
    ///
    /// Use [`Self::resume_session_from`] to request history replay or set
    /// other optional resume parameters. Call
    /// [`ResumeSessionBuilder::start_session`] to publish the request.
    ///
    /// `cwd` must be an absolute path: ACP requires one, and the SDK sends it
    /// without checking.
    pub fn resume_session(
        &self,
        session_id: impl Into<v2::SessionId>,
        cwd: impl AsRef<Path>,
    ) -> ResumeSessionBuilder<Counterpart> {
        self.resume_session_from(v2::ResumeSessionRequest::new(session_id, cwd.as_ref()))
    }

    /// Build a draft protocol v2 `session/resume` request from an existing request.
    ///
    /// Register typed session update handlers before connecting. When the
    /// request asks for replay, the agent sends those updates before the
    /// [`v2::schema::ResumeSessionResponse`](v2::ResumeSessionResponse). Call
    /// [`ResumeSessionBuilder::start_session`] to publish the request.
    pub fn resume_session_from(
        &self,
        request: v2::ResumeSessionRequest,
    ) -> ResumeSessionBuilder<Counterpart> {
        ResumeSessionBuilder::new(self, request)
    }
}

/// Builder for a draft protocol v2 `session/new` request.
///
/// Protocol v2 acknowledges `session/prompt` independently from inbound
/// session updates. Register typed [`v2::schema::UpdateSessionNotification`](v2::UpdateSessionNotification) and
/// session request handlers on [`v2::Builder`](crate::v2::Builder) before connecting, then use
/// [`Self::start_session`] to create the command-only [`ActiveSession`] handle or
/// `on_proxy_session_start` to forward setup through a proxy.
///
/// With both the `unstable_protocol_v2` and `unstable_mcp_over_acp` features,
/// `with_mcp_server` attaches an MCP server to the new session.
#[must_use = "call `start_session` or `on_proxy_session_start` to send the `session/new` request"]
#[derive(Debug)]
pub struct SessionBuilder<Counterpart, Run = NullRun>
where
    Counterpart: HasPeer<Agent>,
    Run: RunWithConnectionTo<Counterpart>,
{
    connection: ConnectionTo<Counterpart>,
    request: v2::NewSessionRequest,
    dynamic_handler_registrations: Vec<DynamicHandlerGuard<Counterpart>>,
    run: Run,
}

impl<Counterpart> SessionBuilder<Counterpart, NullRun>
where
    Counterpart: HasPeer<Agent>,
{
    fn new(connection: &ConnectionTo<Counterpart>, request: v2::NewSessionRequest) -> Self {
        Self {
            connection: connection.clone(),
            request,
            dynamic_handler_registrations: Vec::new(),
            run: NullRun,
        }
    }
}

impl<Counterpart, Run> SessionBuilder<Counterpart, Run>
where
    Counterpart: HasPeer<Agent>,
    Run: RunWithConnectionTo<Counterpart>,
{
    /// Attach an MCP server to this new protocol v2 session.
    ///
    /// This method is available when both `unstable_protocol_v2` and
    /// `unstable_mcp_over_acp` are enabled. MCP routes are installed and their
    /// runner tasks receive an initial poll before `session/new` is published,
    /// allowing the agent to connect while handling session setup. A
    /// successful attachment remains active for the lifetime of the connection.
    ///
    /// The declaration is an `acp` transport (`McpServer::Acp`). The agent must
    /// advertise `session.mcp.acp`, or a compatibility proxy such as
    /// `McpOverAcpPolyfill` must sit in front of it; otherwise a conforming
    /// agent rejects the setup request with `-32602`. This method does not
    /// check the agent's capabilities.
    #[cfg(feature = "unstable_mcp_over_acp")]
    pub fn with_mcp_server<McpRun>(
        mut self,
        mcp_server: McpServer<Counterpart, McpRun>,
    ) -> Result<SessionBuilder<Counterpart, ChainRun<Run, McpRun>>, crate::Error>
    where
        McpRun: RunWithConnectionTo<Counterpart>,
    {
        let (handler, mcp_run) = mcp_server.into_v2_handler_and_runner();
        self.dynamic_handler_registrations
            .push(handler.into_dynamic_handler(&mut self.request.mcp_servers, &self.connection)?);
        Ok(SessionBuilder {
            connection: self.connection,
            request: self.request,
            dynamic_handler_registrations: self.dynamic_handler_registrations,
            run: ChainRun::new(self.run, mcp_run),
        })
    }

    fn send_new_session(self, ordered: bool) -> SentRequest<v2::NewSessionResponse, V2>
    where
        Run: 'static,
    {
        let Self {
            connection,
            request,
            dynamic_handler_registrations,
            run,
        } = self;
        send_session_setup(
            connection,
            request,
            dynamic_handler_registrations,
            run,
            ordered,
        )
    }

    /// Send `session/new` and return its independently consumable request.
    ///
    /// The successful result contains both a cloneable command handle and the
    /// complete [`v2::schema::NewSessionResponse`](v2::NewSessionResponse). Consume the returned request with
    /// [`SentRequest::block_task`], [`SentRequest::on_receiving_result`], or
    /// another explicit [`SentRequest`] completion mode.
    ///
    /// Attached MCP routes are installed and their runner tasks begin
    /// executing before the request is published. A valid success response
    /// promotes them to the connection lifetime, independently from how this
    /// request handle is consumed. Setup errors clean up the pending
    /// attachment.
    pub fn start_session(self) -> SentRequest<SessionSetup<Counterpart, v2::NewSessionResponse>, V2>
    where
        Run: 'static,
    {
        let session_connection = self.connection.clone();
        self.send_new_session(false).map(move |response| {
            let session = ActiveSession {
                session_id: response.session_id.clone(),
                connection: session_connection,
            };
            Ok(SessionSetup { session, response })
        })
    }

    /// Start a protocol v2 session through a proxy and forward its response.
    ///
    /// The downstream request is ordered and inherits cancellation from the
    /// upstream request. On success, this helper installs session routing before
    /// later inbound traffic is processed, forwards the complete response, and
    /// spawns `op` with an [`SessionSetup`] containing the command-only
    /// session handle plus the complete setup response. Inbound updates and
    /// interactive requests remain independent connection traffic.
    ///
    /// The callback runs outside the ordered response barrier, so it may wait
    /// for later connection traffic without deadlocking the dispatch loop.
    pub fn on_proxy_session_start<F, Fut>(
        self,
        responder: Responder<v2::NewSessionResponse, V2>,
        op: F,
    ) -> Result<(), crate::Error>
    where
        Counterpart: HasPeer<Client>,
        Run: 'static,
        F: FnOnce(SessionSetup<Counterpart, v2::NewSessionResponse>) -> Fut + Send + 'static,
        Fut: Future<Output = Result<(), crate::Error>> + Send,
    {
        let session_connection = self.connection.clone();
        self.send_new_session(true)
            .forward_cancellation_from(responder.cancellation())
            .on_receiving_ok_result(responder, async move |response, responder| {
                let session_id = response.session_id.clone();
                let raw_connection = session_connection.raw_connection();
                let route = match raw_connection.add_dynamic_handler(ProxySessionMessages::new(
                    crate::v1::schema::SessionId::from(session_id.0.clone()),
                )) {
                    Ok(route) => route,
                    Err(error) => return responder.respond_with_error(error),
                };

                let setup = SessionSetup {
                    session: ActiveSession {
                        session_id,
                        connection: session_connection.clone(),
                    },
                    response: response.clone(),
                };
                responder.respond(response)?;
                route.detach();
                raw_connection.spawn(async move { op(setup).await })
            })
    }
}

/// Builder for an unstable draft protocol v2 `session/fork` request.
///
/// A successful fork creates a new independent session whose ID comes from the
/// [`v2::schema::ForkSessionResponse`](v2::ForkSessionResponse). Register typed
/// [`v2::schema::UpdateSessionNotification`](v2::UpdateSessionNotification) and interactive request handlers on
/// [`v2::Builder`](crate::v2::Builder) before connecting, then use [`Self::start_session`] to
/// obtain a command handle for the fork or `on_proxy_session_start` to forward
/// setup through a proxy.
///
/// This type is available with the `unstable_session_fork` feature. With
/// `unstable_mcp_over_acp` as well, `with_mcp_server` attaches an MCP server to
/// the forked session.
#[cfg(feature = "unstable_session_fork")]
#[must_use = "call `start_session` or `on_proxy_session_start` to send the `session/fork` request"]
#[derive(Debug)]
pub struct ForkSessionBuilder<Counterpart, Run = NullRun>
where
    Counterpart: HasPeer<Agent>,
    Run: RunWithConnectionTo<Counterpart>,
{
    connection: ConnectionTo<Counterpart>,
    request: v2::ForkSessionRequest,
    dynamic_handler_registrations: Vec<DynamicHandlerGuard<Counterpart>>,
    run: Run,
}

#[cfg(feature = "unstable_session_fork")]
impl<Counterpart> ForkSessionBuilder<Counterpart, NullRun>
where
    Counterpart: HasPeer<Agent>,
{
    fn new(connection: &ConnectionTo<Counterpart>, request: v2::ForkSessionRequest) -> Self {
        Self {
            connection: connection.clone(),
            request,
            dynamic_handler_registrations: Vec::new(),
            run: NullRun,
        }
    }
}

#[cfg(feature = "unstable_session_fork")]
impl<Counterpart, Run> ForkSessionBuilder<Counterpart, Run>
where
    Counterpart: HasPeer<Agent>,
    Run: RunWithConnectionTo<Counterpart>,
{
    /// Attach an MCP server to this forked protocol v2 session.
    ///
    /// This method is available when `unstable_mcp_over_acp` is enabled in
    /// addition to `unstable_protocol_v2` and `unstable_session_fork`. MCP
    /// routes are installed and their runner tasks receive an initial poll
    /// before `session/fork` is published, allowing the agent to connect while
    /// handling session setup. A successful attachment remains active for the
    /// lifetime of the connection.
    ///
    /// The declaration is an `acp` transport (`McpServer::Acp`). The agent must
    /// advertise `session.mcp.acp`, or a compatibility proxy such as
    /// `McpOverAcpPolyfill` must sit in front of it; otherwise a conforming
    /// agent rejects the setup request with `-32602`. This method does not
    /// check the agent's capabilities.
    #[cfg(feature = "unstable_mcp_over_acp")]
    pub fn with_mcp_server<McpRun>(
        mut self,
        mcp_server: McpServer<Counterpart, McpRun>,
    ) -> Result<ForkSessionBuilder<Counterpart, ChainRun<Run, McpRun>>, crate::Error>
    where
        McpRun: RunWithConnectionTo<Counterpart>,
    {
        let (handler, mcp_run) = mcp_server.into_v2_handler_and_runner();
        self.dynamic_handler_registrations
            .push(handler.into_dynamic_handler(&mut self.request.mcp_servers, &self.connection)?);
        Ok(ForkSessionBuilder {
            connection: self.connection,
            request: self.request,
            dynamic_handler_registrations: self.dynamic_handler_registrations,
            run: ChainRun::new(self.run, mcp_run),
        })
    }

    fn send_fork_session(self, ordered: bool) -> SentRequest<v2::ForkSessionResponse, V2>
    where
        Run: 'static,
    {
        let Self {
            connection,
            request,
            dynamic_handler_registrations,
            run,
        } = self;
        send_session_setup(
            connection,
            request,
            dynamic_handler_registrations,
            run,
            ordered,
        )
    }

    /// Send `session/fork` and return its independently consumable request.
    ///
    /// The successful result contains both a cloneable command handle for the
    /// newly created fork and the complete [`v2::schema::ForkSessionResponse`](v2::ForkSessionResponse). Consume
    /// the returned request with [`SentRequest::block_task`],
    /// [`SentRequest::on_receiving_result`], or another explicit [`SentRequest`]
    /// completion mode.
    ///
    /// Attached MCP routes are installed and their runner tasks begin
    /// executing before the request is published. A valid success response
    /// promotes them to the connection lifetime independently from how this
    /// request handle is consumed; setup errors clean up the pending
    /// attachment.
    pub fn start_session(
        self,
    ) -> SentRequest<SessionSetup<Counterpart, v2::ForkSessionResponse>, V2>
    where
        Run: 'static,
    {
        let session_connection = self.connection.clone();
        self.send_fork_session(false).map(move |response| {
            let session = ActiveSession {
                session_id: response.session_id.clone(),
                connection: session_connection,
            };
            Ok(SessionSetup { session, response })
        })
    }

    /// Fork a protocol v2 session through a proxy and forward its response.
    ///
    /// The downstream request is ordered and inherits cancellation from the
    /// upstream request. On success, this helper obtains the new session ID
    /// from the response, installs session routing before later inbound traffic
    /// is processed, forwards the complete response, and spawns `op` with an
    /// [`SessionSetup`] containing the fork's command handle and response.
    /// Inbound updates and interactive requests remain independent connection
    /// traffic.
    ///
    /// The callback runs outside the ordered response barrier, so it may wait
    /// for later connection traffic without deadlocking the dispatch loop.
    pub fn on_proxy_session_start<F, Fut>(
        self,
        responder: Responder<v2::ForkSessionResponse, V2>,
        op: F,
    ) -> Result<(), crate::Error>
    where
        Counterpart: HasPeer<Client>,
        Run: 'static,
        F: FnOnce(SessionSetup<Counterpart, v2::ForkSessionResponse>) -> Fut + Send + 'static,
        Fut: Future<Output = Result<(), crate::Error>> + Send,
    {
        let session_connection = self.connection.clone();
        self.send_fork_session(true)
            .forward_cancellation_from(responder.cancellation())
            .on_receiving_ok_result(responder, async move |response, responder| {
                let session_id = response.session_id.clone();
                let raw_connection = session_connection.raw_connection();
                let route = match raw_connection.add_dynamic_handler(ProxySessionMessages::new(
                    crate::v1::schema::SessionId::from(session_id.0.clone()),
                )) {
                    Ok(route) => route,
                    Err(error) => return responder.respond_with_error(error),
                };

                let setup = SessionSetup {
                    session: ActiveSession {
                        session_id,
                        connection: session_connection.clone(),
                    },
                    response: response.clone(),
                };
                responder.respond(response)?;
                route.detach();
                raw_connection.spawn(async move { op(setup).await })
            })
    }
}

/// Builder for a draft protocol v2 `session/resume` request.
///
/// Replay updates arrive before the resume response. Direct clients must
/// register typed [`v2::schema::UpdateSessionNotification`](v2::UpdateSessionNotification) and interactive request
/// handlers on [`v2::Builder`](crate::v2::Builder) before connecting. Proxies should use
/// [`Self::on_proxy_session_start`], which makes downstream session routing
/// ready before publishing the resume request.
///
/// With both the `unstable_protocol_v2` and `unstable_mcp_over_acp` features,
/// `with_mcp_server` attaches an MCP server to the resumed session and makes it
/// ready before the agent can send replay or its response.
#[must_use = "call `start_session` or `on_proxy_session_start` to send the `session/resume` request"]
#[derive(Debug)]
pub struct ResumeSessionBuilder<Counterpart, Run = NullRun>
where
    Counterpart: HasPeer<Agent>,
    Run: RunWithConnectionTo<Counterpart>,
{
    connection: ConnectionTo<Counterpart>,
    request: v2::ResumeSessionRequest,
    dynamic_handler_registrations: Vec<DynamicHandlerGuard<Counterpart>>,
    run: Run,
}

impl<Counterpart> ResumeSessionBuilder<Counterpart, NullRun>
where
    Counterpart: HasPeer<Agent>,
{
    fn new(connection: &ConnectionTo<Counterpart>, request: v2::ResumeSessionRequest) -> Self {
        Self {
            connection: connection.clone(),
            request,
            dynamic_handler_registrations: Vec::new(),
            run: NullRun,
        }
    }
}

impl<Counterpart, Run> ResumeSessionBuilder<Counterpart, Run>
where
    Counterpart: HasPeer<Agent>,
    Run: RunWithConnectionTo<Counterpart>,
{
    /// Attach an MCP server to this resumed protocol v2 session.
    ///
    /// This method is available when both `unstable_protocol_v2` and
    /// `unstable_mcp_over_acp` are enabled. MCP routes are installed and their
    /// runner tasks receive an initial poll before `session/resume` is
    /// published, allowing the agent to use the server during replay and
    /// session setup. A successful attachment remains active for the lifetime
    /// of the connection.
    ///
    /// The declaration is an `acp` transport (`McpServer::Acp`). The agent must
    /// advertise `session.mcp.acp`, or a compatibility proxy such as
    /// `McpOverAcpPolyfill` must sit in front of it; otherwise a conforming
    /// agent rejects the setup request with `-32602`. This method does not
    /// check the agent's capabilities.
    #[cfg(feature = "unstable_mcp_over_acp")]
    pub fn with_mcp_server<McpRun>(
        mut self,
        mcp_server: McpServer<Counterpart, McpRun>,
    ) -> Result<ResumeSessionBuilder<Counterpart, ChainRun<Run, McpRun>>, crate::Error>
    where
        McpRun: RunWithConnectionTo<Counterpart>,
    {
        let (handler, mcp_run) = mcp_server.into_v2_handler_and_runner();
        self.dynamic_handler_registrations
            .push(handler.into_dynamic_handler(&mut self.request.mcp_servers, &self.connection)?);
        Ok(ResumeSessionBuilder {
            connection: self.connection,
            request: self.request,
            dynamic_handler_registrations: self.dynamic_handler_registrations,
            run: ChainRun::new(self.run, mcp_run),
        })
    }

    fn send_resume_session(self, ordered: bool) -> SentRequest<v2::ResumeSessionResponse, V2>
    where
        Run: 'static,
    {
        let Self {
            connection,
            request,
            dynamic_handler_registrations,
            run,
        } = self;
        send_session_setup(
            connection,
            request,
            dynamic_handler_registrations,
            run,
            ordered,
        )
    }

    /// Send `session/resume` and return its independently consumable request.
    ///
    /// The successful result contains both a cloneable command handle and the
    /// complete [`v2::schema::ResumeSessionResponse`](v2::ResumeSessionResponse). Consume the returned request
    /// with [`SentRequest::block_task`], [`SentRequest::on_receiving_result`],
    /// or another explicit [`SentRequest`] completion mode.
    ///
    /// Replay is delivered through the typed connection handlers before the
    /// response. Attached MCP routes and runner tasks are ready before the
    /// request is published. A valid success response promotes them to the
    /// connection lifetime independently from how this request handle is
    /// consumed; setup errors clean up the pending attachment.
    pub fn start_session(
        self,
    ) -> SentRequest<SessionSetup<Counterpart, v2::ResumeSessionResponse>, V2>
    where
        Run: 'static,
    {
        let session_id = self.request.session_id.clone();
        let session_connection = self.connection.clone();
        self.send_resume_session(false).map(move |response| {
            let session = ActiveSession {
                session_id,
                connection: session_connection,
            };
            Ok(SessionSetup { session, response })
        })
    }

    /// Resume a protocol v2 session through a proxy and forward its response.
    ///
    /// The session ID is known before the request, so this helper installs and
    /// acknowledges downstream session routing before publishing the ordered
    /// `session/resume` request. Replay can therefore be forwarded before the
    /// response as required by the protocol. The downstream request inherits
    /// cancellation from the upstream request.
    ///
    /// On success, the helper forwards the complete response and spawns `op`
    /// with an [`SessionSetup`] containing the command-only session handle
    /// plus that response. The callback runs outside the ordered response
    /// barrier, so it may wait for later connection traffic without
    /// deadlocking the dispatch loop.
    pub fn on_proxy_session_start<F, Fut>(
        mut self,
        responder: Responder<v2::ResumeSessionResponse, V2>,
        op: F,
    ) -> Result<(), crate::Error>
    where
        Counterpart: HasPeer<Client>,
        Run: 'static,
        F: FnOnce(SessionSetup<Counterpart, v2::ResumeSessionResponse>) -> Fut + Send + 'static,
        Fut: Future<Output = Result<(), crate::Error>> + Send,
    {
        let session_id = self.request.session_id.clone();
        let session_connection = self.connection.clone();
        self.dynamic_handler_registrations.push(
            session_connection
                .raw_connection()
                .add_dynamic_handler(ProxySessionMessages::new(
                    crate::v1::schema::SessionId::from(session_id.0.clone()),
                ))?,
        );

        self.send_resume_session(true)
            .forward_cancellation_from(responder.cancellation())
            .on_receiving_ok_result(responder, async move |response, responder| {
                let setup = SessionSetup {
                    session: ActiveSession {
                        session_id,
                        connection: session_connection.clone(),
                    },
                    response: response.clone(),
                };
                responder.respond(response)?;
                session_connection
                    .raw_connection()
                    .spawn(async move { op(setup).await })
            })
    }
}

/// A protocol v2 session and the response of the request that set it up.
///
/// Every setup method returns one: [`SessionBuilder::start_session`]
/// (`session/new`), [`ResumeSessionBuilder::start_session`]
/// (`session/resume`) and, with `unstable_session_fork`,
/// `ForkSessionBuilder::start_session` (`session/fork`), and their
/// `on_proxy_session_start` callbacks receive one too. The response is the
/// complete, operation-specific setup response; it is kept next to the
/// session rather than inside [`ActiveSession`], so the setup response lives
/// in the same place for every setup method.
///
/// New fields may be added, so destructure with `..`:
///
/// ```ignore
/// let SessionSetup { session, response, .. } = setup;
/// ```
#[non_exhaustive]
#[derive(Debug)]
pub struct SessionSetup<Link, Response>
where
    Link: HasPeer<Agent>,
{
    /// The command handle for the session.
    pub session: ActiveSession<Link>,
    /// The complete response of the request that set the session up.
    pub response: Response,
}

/// Cloneable command handle for a draft protocol v2 session.
///
/// Inbound protocol traffic is intentionally not owned by this value. Receive
/// authoritative [`v2::schema::UpdateSessionNotification`](v2::UpdateSessionNotification) values and interactive
/// requests such as [`v2::schema::RequestPermissionRequest`](v2::RequestPermissionRequest) through typed handlers
/// installed on [`v2::Builder`](crate::v2::Builder).
#[derive(Debug, Clone)]
pub struct ActiveSession<Link>
where
    Link: HasPeer<Agent>,
{
    session_id: v2::SessionId,
    connection: ConnectionTo<Link>,
}

impl<Link> ActiveSession<Link>
where
    Link: HasPeer<Agent>,
{
    /// Access the session ID.
    pub fn session_id(&self) -> &v2::SessionId {
        &self.session_id
    }

    /// Access the underlying connection.
    pub fn connection(&self) -> &ConnectionTo<Link> {
        &self.connection
    }

    /// Submit a text prompt and return its independent acceptance request.
    ///
    /// A successful response only acknowledges that the agent accepted the
    /// prompt. The accepted user message, output, state changes, and completion
    /// arrive independently through [`v2::schema::UpdateSessionNotification`](v2::UpdateSessionNotification).
    pub fn send_prompt(&self, prompt: impl ToString) -> SentRequest<v2::PromptResponse, V2> {
        self.send_prompt_blocks(vec![prompt.to_string().into()])
    }

    /// Submit arbitrary prompt content and return its acceptance request.
    ///
    /// `prompt` must not be empty; an agent may reject `[]` with `-32602`.
    ///
    /// The SDK does not track foreground state or gate prompt submission
    /// locally. A prompt may be sent in any state, including while foreground
    /// work is active: the agent either inserts it right away or holds it and
    /// inserts held prompts in received order, and it must not reject a prompt
    /// only because work is active. A held prompt's response stays pending
    /// until it is inserted. Cancel the returned [`SentRequest`] to withdraw a
    /// held prompt; [`Self::cancel_active_work`] and [`Self::close`] also
    /// answer prompts that were not yet inserted with `-32800`.
    pub fn send_prompt_blocks(
        &self,
        prompt: Vec<v2::ContentBlock>,
    ) -> SentRequest<v2::PromptResponse, V2> {
        self.connection.send_request_to(
            Agent,
            v2::PromptRequest::new(self.session_id.clone(), prompt),
        )
    }

    /// Ask the agent to cancel the session's current foreground work.
    ///
    /// This is independent from cancelling a prompt's [`SentRequest`].
    /// `session/cancel` covers the session's foreground work and prompts
    /// received before it that the agent has not yet inserted; those prompt
    /// requests fail with `-32800` (request cancelled). A prompt sent after
    /// this notification is not cancelled by it. A prompt whose user message was already
    /// inserted still succeeds with its `messageId`. Work in the latest
    /// `background` list is not covered and keeps running.
    ///
    /// If foreground work was active, cancellation completes when the agent
    /// reports an `idle` state update with [`v2::schema::StopReason::Cancelled`](v2::StopReason::Cancelled); no
    /// completion update is required when nothing was running. `cancelled` is only
    /// reported for work the client cancelled, through this notification or
    /// by closing the session. The client should immediately mark unfinished
    /// tool calls for the active work as cancelled.
    ///
    /// The agent stops waiting for pending
    /// [`v2::schema::RequestPermissionRequest`](v2::RequestPermissionRequest)s
    /// of the cancelled work and should cancel each of them with
    /// `$/cancel_request`. Answer a cancelled permission request with the
    /// `cancelled` outcome. A client may also answer a request it knows belongs
    /// to the cancelled work that way without waiting for `$/cancel_request`.
    ///
    /// See <https://agentclientprotocol.com/protocol/v2/prompt-lifecycle#cancellation>.
    pub fn cancel_active_work(&self) -> Result<(), crate::Error> {
        self.connection.send_notification_to(
            Agent,
            v2::CancelSessionNotification::new(self.session_id.clone()),
        )
    }

    /// Set a session configuration option.
    ///
    /// Only send a `config_id` the agent advertised for this session, with a
    /// value from that option. The agent rejects an unknown `config_id` or an
    /// invalid value with `-32602` and leaves the configuration unchanged.
    ///
    /// The response contains the complete option set as of the response and
    /// may report a different `currentValue` than requested; `currentValue`
    /// is always the value the agent is using. Apply responses and
    /// `config_option_update` notifications in the order received, each
    /// replacing the whole list. A duplicate `config_option_update`
    /// notification for the same change may arrive before the response. The
    /// response is not cached on this command handle.
    ///
    /// See <https://agentclientprotocol.com/protocol/v2/session-config-options>.
    pub fn set_config_option(
        &self,
        config_id: impl Into<v2::SessionConfigId>,
        value: impl Into<v2::SessionConfigOptionValue>,
    ) -> SentRequest<v2::SetSessionConfigOptionResponse, V2> {
        self.connection.send_request_to(
            Agent,
            v2::SetSessionConfigOptionRequest::new(self.session_id.clone(), config_id, value),
        )
    }

    /// Close the remote session and release its resources.
    ///
    /// The agent cancels the session's ongoing work as if `session/cancel` had
    /// been sent. When that cancels active foreground work, the agent sends the
    /// work's final updates and an `idle` state update with stop reason
    /// `cancelled` before it responds to `session/close`. A pending prompt
    /// whose user message closing kept from being inserted fails with `-32800`
    /// (request cancelled); one already inserted still succeeds with its
    /// message ID. Prompt responses need not precede the close response.
    /// Before it responds, the agent cancels every pending permission request
    /// and elicitation in the session with `$/cancel_request`.
    ///
    /// Existing clones of this local command handle are not invalidated, but
    /// the session is no longer active: the agent rejects later commands for it
    /// with `-32002` (resource not found).
    ///
    /// Closing is idempotent. An agent should answer a close of an unknown or
    /// already-closed session with success, and an older agent may answer
    /// `-32002` instead, so this helper resolves a `-32002` response as a
    /// successful, empty [`v2::schema::CloseSessionResponse`](v2::CloseSessionResponse). Other errors are
    /// returned unchanged.
    ///
    /// See <https://agentclientprotocol.com/protocol/v2/session-setup#closing-active-sessions>.
    pub fn close(&self) -> SentRequest<v2::CloseSessionResponse, V2> {
        self.connection
            .send_request_to(Agent, v2::CloseSessionRequest::new(self.session_id.clone()))
            .recover_error(|error| {
                if error.code == v2::ErrorCode::ResourceNotFound {
                    Ok(v2::CloseSessionResponse::new())
                } else {
                    Err(error)
                }
            })
    }
}