agent_client_protocol/

session.rs

use std::{future::Future, marker::PhantomData, path::Path, sync::Arc};

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

use crate::{
    Agent, Client, ConnectionTo as RawConnectionTo, Dispatch, HandleDispatchFrom, Handled,
    JsonRpcRequest, Responder, Role,
    jsonrpc::{
        DynamicHandlerCleanup, DynamicHandlerGuard,
        run::{NullRun, RunWithConnectionTo, RunnerErrorScope},
    },
    role::{HasPeer, acp::ProxySessionMessages},
    util::{MatchDispatch, MatchDispatchFrom},
    v1::ConnectionTo,
    v1::schema::{
        ContentBlock, ContentChunk, LoadSessionRequest, LoadSessionResponse, NewSessionRequest,
        NewSessionResponse, PromptRequest, PromptResponse, ResumeSessionRequest,
        ResumeSessionResponse, SessionId, SessionNotification, SessionUpdate,
    },
};

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

#[cfg(feature = "unstable_protocol_v2")]
pub(crate) mod v2;

type SessionCleanup = Vec<Arc<dyn DynamicHandlerCleanup>>;

fn session_cleanup<R: Role>(guards: &[DynamicHandlerGuard<R>]) -> SessionCleanup {
    guards
        .iter()
        .filter_map(DynamicHandlerGuard::cleanup)
        .collect()
}

fn close_session_registrations(cleanup: &SessionCleanup) {
    for registration in cleanup {
        registration.close();
    }
}

async fn wait_session_cleanup(cleanup: SessionCleanup) {
    future::join_all(cleanup.iter().map(|registration| registration.wait())).await;
}

fn session_runner_connection<R: Role>(
    connection: RawConnectionTo<R>,
    cleanup: &SessionCleanup,
) -> (RawConnectionTo<R>, RunnerErrorScope) {
    let closing = cleanup.clone();
    let scope = RunnerErrorScope::new(
        move || close_session_registrations(&closing),
        wait_session_cleanup(cleanup.clone()),
    );
    (connection.with_runner_error_scope(scope.clone()), scope)
}

/// Keep the actual (possibly borrowed) runner alive until only these
/// registrations finish. Never seal connection-wide protected admission here.
async fn drive_session_cleanup(
    run: impl Future<Output = Result<(), crate::Error>>,
    cleanup: SessionCleanup,
) -> Result<(), crate::Error> {
    match future::select(
        Box::pin(wait_session_cleanup(cleanup.clone())),
        Box::pin(run),
    )
    .await
    {
        Either::Left(((), _run)) => Ok(()),
        Either::Right((result, waiting)) => {
            if result.is_err() {
                close_session_registrations(&cleanup);
            }
            waiting.await;
            result
        }
    }
}

async fn run_attached_session_runner(
    run: impl Future<Output = Result<(), crate::Error>>,
    cleanup: SessionCleanup,
    error_scope: RunnerErrorScope,
) -> Result<(), crate::Error> {
    let result = if cleanup.is_empty() {
        run.await
    } else {
        drive_session_cleanup(run, cleanup).await
    };
    // Local cleanup can win the race against the retained chain's final poll.
    // Once attached, an observed runner failure must still reach the task actor.
    error_scope.error().map_or(result, Err)
}

async fn run_session_scope<T>(
    run: impl Future<Output = Result<(), crate::Error>>,
    op: impl Future<Output = Result<T, crate::Error>>,
    cleanup: SessionCleanup,
    error_scope: RunnerErrorScope,
) -> Result<T, crate::Error> {
    let result = match future::select(Box::pin(run), Box::pin(op)).await {
        Either::Left((run_result, op)) => {
            let result = match run_result {
                Ok(()) => op.await,
                Err(error) => {
                    drop(op);
                    Err(error)
                }
            };
            close_session_registrations(&cleanup);
            wait_session_cleanup(cleanup).await;
            result
        }
        Either::Right((result, run)) => {
            close_session_registrations(&cleanup);
            let runner_result = drive_session_cleanup(run, cleanup).await;
            // Foreground failure remains authoritative over cleanup failures.
            result.and_then(|value| runner_result.map(|()| value))
        }
    };
    // The chain may have observed an error while still Pending for cleanup.
    // Successful foreground completion must not hide that recorded failure.
    // An explicit foreground error remains authoritative.
    result.and_then(|value| error_scope.error().map_or(Ok(value), Err))
}

/// Marker type indicating the session builder will block the current task.
#[derive(Debug)]
pub struct Blocking;
impl block_state::Sealed for Blocking {}
impl SessionBlockState for Blocking {}

/// Marker type indicating the session builder will not block the current task.
#[derive(Debug)]
pub struct NonBlocking;
impl block_state::Sealed for NonBlocking {}
impl SessionBlockState for NonBlocking {}

mod block_state {
    pub trait Sealed {}
}

/// Trait for marker types that indicate blocking vs non-blocking API.
/// See [`SessionBuilder::block_task`].
///
/// This trait is sealed: [`Blocking`] and [`NonBlocking`] are its only
/// implementations.
pub trait SessionBlockState: block_state::Sealed + Send + 'static + Sync + std::fmt::Debug {}

impl<Counterpart: Role> ConnectionTo<Counterpart>
where
    Counterpart: HasPeer<Agent>,
{
    /// Stable protocol v1 session builder for a new session request.
    ///
    /// With `unstable_protocol_v2`, a `Client.v2()` callback receives a
    /// `v2::ConnectionTo` with its own v2 `build_session` helper.
    ///
    /// `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, NullRun> {
        SessionBuilder::new(self, NewSessionRequest::new(cwd.as_ref()))
    }

    /// Stable protocol v1 session builder using the current working directory.
    ///
    /// This is a convenience wrapper around [`build_session`](Self::build_session)
    /// that uses [`std::env::current_dir`] to get the working directory.
    ///
    /// Returns an error if the current directory cannot be determined.
    pub fn build_session_cwd(&self) -> Result<SessionBuilder<Counterpart, NullRun>, crate::Error> {
        let cwd = std::env::current_dir().map_err(|e| {
            crate::Error::internal_error().data(format!("cannot get current directory: {e}"))
        })?;
        Ok(self.build_session(cwd))
    }

    /// Stable protocol v1 session builder starting from an existing request.
    ///
    /// Use this when you've intercepted a `session.new` request and want to
    /// modify it (e.g., inject MCP servers) before forwarding.
    pub fn build_session_from(
        &self,
        request: NewSessionRequest,
    ) -> SessionBuilder<Counterpart, NullRun> {
        SessionBuilder::new(self, request)
    }

    /// Stable protocol v1 session builder that loads an existing session.
    ///
    /// The returned builder installs session routing before publishing
    /// `session/load`, so replay notifications sent before the response are
    /// available through the restored [`ActiveSession`].
    ///
    /// Call this only when the initialization response advertises
    /// `agentCapabilities.loadSession`.
    ///
    /// `cwd` must be an absolute path: ACP requires one, and the SDK sends it
    /// without checking.
    pub fn load_session(
        &self,
        session_id: impl Into<SessionId>,
        cwd: impl AsRef<Path>,
    ) -> LoadSessionBuilder<Counterpart> {
        self.load_session_from(LoadSessionRequest::new(session_id, cwd.as_ref()))
    }

    /// Stable protocol v1 session builder from an existing `session/load`
    /// request.
    ///
    /// Use this to send a typed request assembled or intercepted elsewhere
    /// without rebuilding it.
    pub fn load_session_from(
        &self,
        request: LoadSessionRequest,
    ) -> LoadSessionBuilder<Counterpart> {
        RestoreSessionBuilder::new(self, request)
    }

    /// Stable protocol v1 session builder that resumes an existing session.
    ///
    /// This is the `session/resume` counterpart of
    /// [`load_session`](Self::load_session), but continues without replaying
    /// conversation history. Call this only when the initialization response
    /// advertises `agentCapabilities.sessionCapabilities.resume`.
    ///
    /// `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<SessionId>,
        cwd: impl AsRef<Path>,
    ) -> ResumeSessionBuilder<Counterpart> {
        self.resume_session_from(ResumeSessionRequest::new(session_id, cwd.as_ref()))
    }

    /// Stable protocol v1 session builder from an existing `session/resume`
    /// request.
    ///
    /// Use this to send a typed request assembled or intercepted elsewhere
    /// without rebuilding it.
    pub fn resume_session_from(
        &self,
        request: ResumeSessionRequest,
    ) -> ResumeSessionBuilder<Counterpart> {
        RestoreSessionBuilder::new(self, request)
    }

    /// Given a session response received from the agent,
    /// attach a handler to process messages related to this session
    /// and return the session together with the response.
    ///
    /// Normally you would not use this method directly but would
    /// instead use [`Self::build_session`] and then [`SessionBuilder::start_session`].
    ///
    /// The vector `dynamic_handler_registrations` contains any dynamic
    /// handle registrations associated with this session (e.g., from MCP servers).
    /// You can simply pass `Default::default()` if not applicable.
    pub(crate) fn attach_session<'runner>(
        &self,
        response: NewSessionResponse,
        mcp_handler_registrations: Vec<DynamicHandlerGuard<Counterpart>>,
    ) -> Result<SessionSetup<'runner, Counterpart, NewSessionResponse>, crate::Error> {
        let prepared = self.prepare_session_routing(&response.session_id)?;
        let session = prepared.into_active_session(
            self.clone(),
            response.session_id.clone(),
            mcp_handler_registrations,
        );
        Ok(SessionSetup { session, response })
    }

    /// Install the update channel and handler for `session_id`.
    ///
    /// Restore requests call this before request publication. Dropping the
    /// returned value deactivates and removes the route.
    fn prepare_session_routing(
        &self,
        session_id: &SessionId,
    ) -> Result<PreparedSession<Counterpart>, crate::Error> {
        let (update_tx, update_rx) = mpsc::unbounded();
        let handler = ActiveSessionHandler::new(session_id.clone(), update_tx.clone());
        let session_handler_registration = self.add_dynamic_handler(handler)?;

        Ok(PreparedSession {
            update_rx,
            update_tx,
            session_handler_registration,
        })
    }
}

/// Session-routing state installed before a restore request is published.
struct PreparedSession<Counterpart: Role>
where
    Counterpart: HasPeer<Agent>,
{
    update_rx: mpsc::UnboundedReceiver<SessionMessage>,
    update_tx: mpsc::UnboundedSender<SessionMessage>,
    session_handler_registration: DynamicHandlerGuard<Counterpart>,
}

impl<Counterpart> PreparedSession<Counterpart>
where
    Counterpart: HasPeer<Agent>,
{
    fn into_active_session<'runner>(
        self,
        connection: ConnectionTo<Counterpart>,
        session_id: SessionId,
        mcp_handler_registrations: Vec<DynamicHandlerGuard<Counterpart>>,
    ) -> ActiveSession<'runner, Counterpart> {
        ActiveSession {
            session_id,
            update_rx: self.update_rx,
            update_tx: self.update_tx,
            connection,
            session_handler_registration: self.session_handler_registration,
            mcp_handler_registrations,
            _runner: PhantomData,
        }
    }
}

/// Internal behavior shared by the two stable restore operations.
trait RestoreRequest: JsonRpcRequest {
    fn session_id(&self) -> &SessionId;
}

impl RestoreRequest for LoadSessionRequest {
    fn session_id(&self) -> &SessionId {
        &self.session_id
    }
}

impl RestoreRequest for ResumeSessionRequest {
    fn session_id(&self) -> &SessionId {
        &self.session_id
    }
}

/// Stable protocol v1 builder for `session/load` or `session/resume`.
///
/// This type is not exported. Callers name it through the
/// [`LoadSessionBuilder`] and [`ResumeSessionBuilder`] aliases.
///
/// Use [`ConnectionTo::load_session`] or [`ConnectionTo::resume_session`] to
/// construct this builder. Use the matching `_from` method to send an existing
/// typed request without rebuilding it.
///
/// The `BlockState` parameter mirrors [`SessionBuilder`]:
/// - [`NonBlocking`] exposes `on_session_start` on each concrete operation.
/// - [`Blocking`], selected with [`Self::block_task`], exposes
///   `start_session`.
///
/// Session routing is acknowledged before the request can reach the peer.
/// Dropping a pending blocking start removes that routing and applies the
/// standard [`SentRequest`](crate::SentRequest) drop-time cancellation
/// behavior. Error responses remove the route before later entries in the same
/// transport frame are dispatched.
#[must_use = "use `start_session` or `on_session_start` to restore the session"]
#[derive(Debug)]
pub struct RestoreSessionBuilder<Counterpart, Request, BlockState = NonBlocking>
where
    Counterpart: HasPeer<Agent>,
    BlockState: SessionBlockState,
{
    connection: ConnectionTo<Counterpart>,
    request: Request,
    block_state: PhantomData<BlockState>,
}

/// Stable protocol v1 builder for a `session/load` request, returned by
/// [`ConnectionTo::load_session`] and [`ConnectionTo::load_session_from`].
///
/// The `BlockState` parameter mirrors [`SessionBuilder`]:
/// - [`NonBlocking`] exposes `on_session_start`.
/// - [`Blocking`], selected with `block_task`, exposes `start_session`.
///
/// Session routing is installed before the request can reach the peer, so
/// replay notifications sent before the response are kept on the restored
/// [`ActiveSession`].
///
/// Unlike [`SessionBuilder`], this builder has no `with_mcp_server` and no
/// proxy helpers (`on_proxy_session_start`, `start_session_proxy`). To add an
/// MCP server to load requests a proxy forwards, use `with_mcp_server` on the
/// proxy's [`v1::Builder`](crate::v1::Builder).
pub type LoadSessionBuilder<Counterpart, BlockState = NonBlocking> =
    RestoreSessionBuilder<Counterpart, LoadSessionRequest, BlockState>;

/// Stable protocol v1 builder for a `session/resume` request, returned by
/// [`ConnectionTo::resume_session`] and [`ConnectionTo::resume_session_from`].
///
/// The `BlockState` parameter mirrors [`SessionBuilder`]:
/// - [`NonBlocking`] exposes `on_session_start`.
/// - [`Blocking`], selected with `block_task`, exposes `start_session`.
///
/// Unlike [`SessionBuilder`], this builder has no `with_mcp_server` and no
/// proxy helpers (`on_proxy_session_start`, `start_session_proxy`). To add an
/// MCP server to resume requests a proxy forwards, use `with_mcp_server` on
/// the proxy's [`v1::Builder`](crate::v1::Builder).
pub type ResumeSessionBuilder<Counterpart, BlockState = NonBlocking> =
    RestoreSessionBuilder<Counterpart, ResumeSessionRequest, BlockState>;

impl<Counterpart, Request> RestoreSessionBuilder<Counterpart, Request, NonBlocking>
where
    Counterpart: HasPeer<Agent>,
{
    fn new(connection: &ConnectionTo<Counterpart>, request: Request) -> Self {
        Self {
            connection: connection.clone(),
            request,
            block_state: PhantomData,
        }
    }

    /// Mark this restore builder as able to block the current task.
    ///
    /// Do not use the resulting blocking methods inside a message handler.
    pub fn block_task(self) -> RestoreSessionBuilder<Counterpart, Request, Blocking> {
        RestoreSessionBuilder {
            connection: self.connection,
            request: self.request,
            block_state: PhantomData,
        }
    }
}

fn restored_session<Counterpart, Request>(
    connection: ConnectionTo<Counterpart>,
    session_id: SessionId,
    prepared: PreparedSession<Counterpart>,
    response: Request::Response,
) -> SessionSetup<'static, Counterpart, Request::Response>
where
    Counterpart: HasPeer<Agent>,
    Request: RestoreRequest,
{
    let session = prepared.into_active_session(connection, session_id, Vec::new());

    SessionSetup { session, response }
}

fn on_restore_session_start<Counterpart, Request, F, Fut>(
    builder: RestoreSessionBuilder<Counterpart, Request>,
    op: F,
) -> Result<(), crate::Error>
where
    Counterpart: HasPeer<Agent>,
    Request: RestoreRequest,
    F: FnOnce(SessionSetup<'static, Counterpart, Request::Response>) -> Fut + Send + 'static,
    Fut: Future<Output = Result<(), crate::Error>> + Send,
{
    let RestoreSessionBuilder {
        connection,
        request,
        block_state: _,
    } = builder;
    let session_id = request.session_id().clone();
    let prepared = connection.prepare_session_routing(&session_id)?;
    let routing_ready = connection.raw_connection().dynamic_handler_barrier();

    connection
        .raw_connection()
        .send_ordered_request_to_after::<Raw, _, _, _>(Agent, request, routing_ready)
        .on_receiving_result({
            let connection = connection.clone();
            async move |result| {
                let response = result?;
                let restored = restored_session::<_, Request>(
                    connection.clone(),
                    session_id,
                    prepared,
                    response,
                );
                connection.spawn(async move { op(restored).await })
            }
        })
}

async fn start_restored_session<Counterpart, Request>(
    builder: RestoreSessionBuilder<Counterpart, Request, Blocking>,
) -> Result<SessionSetup<'static, Counterpart, Request::Response>, crate::Error>
where
    Counterpart: HasPeer<Agent>,
    Request: RestoreRequest,
{
    let RestoreSessionBuilder {
        connection,
        request,
        block_state: _,
    } = builder;
    let session_id = request.session_id().clone();
    let prepared = connection.prepare_session_routing(&session_id)?;
    let routing_ready = connection.raw_connection().dynamic_handler_barrier();
    let session_connection = connection.clone();

    connection
        .raw_connection()
        .send_ordered_request_to_after::<Raw, _, _, _>(Agent, request, routing_ready)
        .block_task_with_ordered_result(move |result| {
            let response = result?;
            Ok(restored_session::<_, Request>(
                session_connection,
                session_id,
                prepared,
                response,
            ))
        })
        .await
}

impl<Counterpart> RestoreSessionBuilder<Counterpart, LoadSessionRequest>
where
    Counterpart: HasPeer<Agent>,
{
    /// Restore with `session/load` in the background and run `op` once its
    /// exact response and active session are available.
    ///
    /// This returns immediately and is safe to call from a message handler.
    /// Replay notifications can arrive before the response and are retained by
    /// the returned session.
    pub fn on_session_start<F, Fut>(self, op: F) -> Result<(), crate::Error>
    where
        F: FnOnce(SessionSetup<'static, Counterpart, LoadSessionResponse>) -> Fut + Send + 'static,
        Fut: Future<Output = Result<(), crate::Error>> + Send,
    {
        on_restore_session_start(self, op)
    }
}

impl<Counterpart> RestoreSessionBuilder<Counterpart, ResumeSessionRequest>
where
    Counterpart: HasPeer<Agent>,
{
    /// Restore with `session/resume` in the background and run `op` once its
    /// exact response and active session are available.
    ///
    /// This returns immediately and is safe to call from a message handler.
    /// The returned session receives subsequent session traffic.
    pub fn on_session_start<F, Fut>(self, op: F) -> Result<(), crate::Error>
    where
        F: FnOnce(SessionSetup<'static, Counterpart, ResumeSessionResponse>) -> Fut
            + Send
            + 'static,
        Fut: Future<Output = Result<(), crate::Error>> + Send,
    {
        on_restore_session_start(self, op)
    }
}

impl<Counterpart> RestoreSessionBuilder<Counterpart, LoadSessionRequest, Blocking>
where
    Counterpart: HasPeer<Agent>,
{
    /// Publish `session/load`, wait on the current task, and return an
    /// [`ActiveSession`] together with the exact [`LoadSessionResponse`].
    ///
    /// Requires [`block_task`](LoadSessionBuilder::block_task). Dropping
    /// this future while it is pending cancels the request and removes the
    /// provisional session route.
    pub async fn start_session(
        self,
    ) -> Result<SessionSetup<'static, Counterpart, LoadSessionResponse>, crate::Error> {
        start_restored_session(self).await
    }
}

impl<Counterpart> RestoreSessionBuilder<Counterpart, ResumeSessionRequest, Blocking>
where
    Counterpart: HasPeer<Agent>,
{
    /// Publish `session/resume`, wait on the current task, and return an
    /// [`ActiveSession`] together with the exact [`ResumeSessionResponse`].
    ///
    /// Requires [`block_task`](ResumeSessionBuilder::block_task). Dropping
    /// this future while it is pending cancels the request and removes the
    /// provisional session route.
    pub async fn start_session(
        self,
    ) -> Result<SessionSetup<'static, Counterpart, ResumeSessionResponse>, crate::Error> {
        start_restored_session(self).await
    }
}

/// A stable protocol v1 session and the response of the request that set it
/// up.
///
/// Every setup method returns one: [`SessionBuilder::start_session`]
/// (`session/new`), and `start_session` on [`LoadSessionBuilder`]
/// (`session/load`) and [`ResumeSessionBuilder`] (`session/resume`). The `on_session_start` and
/// [`run_until`](SessionBuilder::run_until) 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. For load and
/// resume, the session ID comes from the request, because stable v1 restore
/// responses do not repeat it.
///
/// New fields may be added, so destructure with `..`:
///
/// ```ignore
/// let SessionSetup { mut session, response, .. } = builder.start_session().await?;
/// ```
#[non_exhaustive]
pub struct SessionSetup<'runner, Link, Response>
where
    Link: HasPeer<Agent>,
{
    /// The session, with its update stream.
    pub session: ActiveSession<'runner, Link>,
    /// The complete response of the request that set the session up.
    pub response: Response,
}

impl<Link, Response> std::fmt::Debug for SessionSetup<'_, Link, Response>
where
    Link: HasPeer<Agent>,
    Response: std::fmt::Debug,
{
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter
            .debug_struct("SessionSetup")
            .field("session_id", self.session.session_id())
            .field("response", &self.response)
            .finish()
    }
}

/// Stable protocol v1 session builder for a new session request.
/// Allows you to add MCP servers or set other details for this session.
///
/// The `BlockState` type parameter tracks whether blocking methods are available:
/// - `NonBlocking` (default): Only [`on_session_start`](Self::on_session_start) is available
/// - `Blocking` (after calling [`block_task`](Self::block_task)):
///   [`run_until`](Self::run_until) and [`start_session`](Self::start_session) become available
#[must_use = "use `start_session`, `run_until`, or `on_session_start` to start the session"]
#[derive(Debug)]
pub struct SessionBuilder<
    Counterpart,
    Run: RunWithConnectionTo<Counterpart> = NullRun,
    BlockState: SessionBlockState = NonBlocking,
> where
    Counterpart: HasPeer<Agent>,
{
    connection: ConnectionTo<Counterpart>,
    request: NewSessionRequest,
    dynamic_handler_registrations: Vec<DynamicHandlerGuard<Counterpart>>,
    run: Run,
    block_state: PhantomData<BlockState>,
}

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

impl<Counterpart, R, BlockState> SessionBuilder<Counterpart, R, BlockState>
where
    Counterpart: HasPeer<Agent>,
    R: RunWithConnectionTo<Counterpart>,
    BlockState: SessionBlockState,
{
    /// Attach an MCP server to this new session.
    ///
    /// The declaration is an `acp` transport (`McpServer::Acp`). The agent must
    /// advertise `agentCapabilities.mcpCapabilities.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<R, McpRun>, BlockState>, crate::Error>
    where
        McpRun: RunWithConnectionTo<Counterpart>,
    {
        let (handler, mcp_run) = mcp_server.into_handler_and_runner();
        self.dynamic_handler_registrations.push(
            handler.into_dynamic_handler(&mut self.request, self.connection.raw_connection())?,
        );
        Ok(SessionBuilder {
            connection: self.connection,
            request: self.request,
            dynamic_handler_registrations: self.dynamic_handler_registrations,
            run: ChainRun::new(self.run, mcp_run),
            block_state: self.block_state,
        })
    }

    /// Spawn a task that runs the provided closure once the session starts.
    ///
    /// Unlike [`start_session`](Self::start_session), this method returns immediately
    /// without blocking the current task. The session handshake and closure execution
    /// happen in a spawned background task.
    ///
    /// The closure receives a [`SessionSetup`] with an
    /// `ActiveSession<'static, _>` and the [`NewSessionResponse`], and runs in
    /// a spawned task. If it returns an error, the error propagates to the
    /// connection's task handling.
    ///
    /// # Example
    ///
    /// ```ignore
    /// # use agent_client_protocol::{Client, Agent, ConnectTo};
    /// # use agent_client_protocol::mcp_server::McpServer;
    /// # use agent_client_protocol_rmcp::rmcp_3::McpServerExt;
    /// # async fn example(transport: impl ConnectTo<Client>) -> Result<(), agent_client_protocol::Error> {
    /// # Client.v1().connect_with(transport, async |cx| {
    /// # let mcp = McpServer::<Agent, _>::builder("tools").build();
    /// cx.build_session_cwd()?
    ///     .with_mcp_server(mcp)?
    ///     .on_session_start(async |SessionSetup { mut session, .. }| {
    ///         // Do something with the session
    ///         session.send_prompt("Hello")?;
    ///         let response = session.read_to_string().await?;
    ///         Ok(())
    ///     })?;
    /// // Returns immediately, session runs in background
    /// # Ok(())
    /// # }).await?;
    /// # Ok(())
    /// # }
    /// ```
    ///
    /// # Ordering
    ///
    /// Session runners are scheduled and routing setup is installed before the
    /// dispatch loop processes the next message when the session response is
    /// routed during its original dispatch. No user callback code runs under
    /// that ordering guarantee: the callback is invoked in a spawned task, so
    /// it may wait for later session traffic without deadlocking the connection.
    /// A response interceptor that retains the response and routes it later
    /// cannot retroactively order session setup before messages the dispatch
    /// loop has already processed.
    pub fn on_session_start<F, Fut>(self, op: F) -> Result<(), crate::Error>
    where
        R: 'static,
        F: FnOnce(SessionSetup<'static, Counterpart, NewSessionResponse>) -> Fut + Send + 'static,
        Fut: Future<Output = Result<(), crate::Error>> + Send,
    {
        let Self {
            connection,
            request,
            dynamic_handler_registrations,
            run,
            block_state: _,
        } = self;

        let cleanup = session_cleanup(&dynamic_handler_registrations);
        let (runner_connection, error_scope) =
            session_runner_connection(connection.raw_connection().clone(), &cleanup);
        connection.spawn(run_attached_session_runner(
            run.run_with_connection_to(runner_connection),
            cleanup,
            error_scope,
        ))?;

        connection
            .raw_connection()
            .send_ordered_request_to::<Raw, _, _>(Agent, request)
            .on_receiving_result({
                let connection = connection.clone();
                async move |result| {
                    let response = result?;

                    let setup =
                        connection.attach_session(response, dynamic_handler_registrations)?;

                    connection.spawn(async move { op(setup).await })
                }
            })
    }

    /// Spawn a proxy session and run a closure with the `session/new` response.
    ///
    /// A **proxy session** starts the session with the agent and then automatically
    /// proxies all session updates (prompts, tool calls, etc.) from the agent back
    /// to the client. You don't need to handle any messages yourself - the proxy
    /// takes care of forwarding everything. This is useful when you want to inject
    /// and/or filter prompts coming from the client but otherwise not be involved
    /// in the session.
    ///
    /// Unlike [`start_session_proxy`](Self::start_session_proxy), this method returns
    /// immediately without blocking the current task. The session handshake, client
    /// response, and proxy setup all happen in a spawned background task.
    ///
    /// The closure receives a clone of the complete [`NewSessionResponse`], which was also
    /// forwarded to the client, once the session is established. Read the session ID from its
    /// `session_id` field. Use it for logging
    /// or eventual tracking; it runs concurrently with later connection traffic. Register
    /// ID-independent state that later handlers must observe before calling this helper. For
    /// ID-keyed bookkeeping, install a gate or placeholder first, make later handlers await it,
    /// and populate it from the closure.
    ///
    /// # Example
    ///
    /// ```ignore
    /// # use agent_client_protocol::{Proxy, Client, Conductor, ConnectTo};
    /// # use agent_client_protocol::v1::schema::NewSessionRequest;
    /// # use agent_client_protocol::mcp_server::McpServer;
    /// # use agent_client_protocol_rmcp::rmcp_3::McpServerExt;
    /// # async fn example(transport: impl ConnectTo<Proxy>) -> Result<(), agent_client_protocol::Error> {
    /// Proxy.v1()
    ///     .on_receive_request_from(Client, async |request: NewSessionRequest, responder, cx| {
    ///         let mcp = McpServer::<Conductor, _>::builder("tools").build();
    ///         cx.build_session_from(request)
    ///             .with_mcp_server(mcp)?
    ///             .on_proxy_session_start(responder, async |response| {
    ///                 // Session started
    ///                 Ok(())
    ///             })
    ///     }, agent_client_protocol::on_receive_request!())
    ///     .connect_to(transport)
    ///     .await?;
    /// # Ok(())
    /// # }
    /// ```
    ///
    /// # Ordering
    ///
    /// The client response is queued, proxy routing is installed, and session runners are
    /// scheduled before the dispatch loop processes the next message when the session response
    /// is routed during its original dispatch. This is a local ordering guarantee, not a
    /// guarantee that the response reaches the client before later wire traffic. No user callback
    /// code runs under the barrier: the callback is invoked in a spawned task, so it may wait for
    /// later connection traffic. A response interceptor that retains the response and routes it
    /// later cannot retroactively order this setup before messages the loop already processed.
    pub fn on_proxy_session_start<F, Fut>(
        self,
        responder: Responder<NewSessionResponse, crate::V1>,
        op: F,
    ) -> Result<(), crate::Error>
    where
        F: FnOnce(NewSessionResponse) -> Fut + Send + 'static,
        Fut: Future<Output = Result<(), crate::Error>> + Send,
        Counterpart: HasPeer<Client>,
        R: 'static,
    {
        let Self {
            connection,
            request,
            dynamic_handler_registrations,
            run,
            block_state: _,
        } = self;

        let cleanup = session_cleanup(&dynamic_handler_registrations);
        let (runner_connection, error_scope) =
            session_runner_connection(connection.raw_connection().clone(), &cleanup);
        connection.spawn(run_attached_session_runner(
            run.run_with_connection_to(runner_connection),
            cleanup,
            error_scope,
        ))?;

        // Send the "new session" request to the agent.
        let sent = connection
            .raw_connection()
            .send_ordered_request_to::<Raw, _, _>(Agent, request);
        let sent = sent.forward_cancellation_from(responder.cancellation());

        sent.on_receiving_ok_result(responder, {
            let connection = connection.clone();
            async move |response, responder| {
                // Extract the session-id from the response and forward
                // the response back to the client
                let session_id = response.session_id.clone();
                responder.respond(response.clone())?;

                // Install a dynamic handler to proxy messages from this session
                connection
                    .add_dynamic_handler(ProxySessionMessages::new(session_id))?
                    .detach();

                // Keep dynamic handlers live for the connection.
                dynamic_handler_registrations
                    .into_iter()
                    .for_each(DynamicHandlerGuard::detach);

                connection.spawn(async move { op(response).await })
            }
        })
    }
}

impl<Counterpart, R> SessionBuilder<Counterpart, R, NonBlocking>
where
    Counterpart: HasPeer<Agent>,
    R: RunWithConnectionTo<Counterpart>,
{
    /// Mark this session builder as being able to block the current task.
    ///
    /// After calling this, you can use [`run_until`](Self::run_until) or
    /// [`start_session`](Self::start_session) which block the current task.
    ///
    /// This should not be used from inside a message handler like
    /// [`Builder::on_receive_request`](`crate::Builder::on_receive_request`) or [`HandleDispatchFrom`]
    /// implementations.
    pub fn block_task(self) -> SessionBuilder<Counterpart, R, Blocking> {
        SessionBuilder {
            connection: self.connection,
            request: self.request,
            dynamic_handler_registrations: self.dynamic_handler_registrations,
            run: self.run,
            block_state: PhantomData,
        }
    }
}

impl<Counterpart, R> SessionBuilder<Counterpart, R, Blocking>
where
    Counterpart: HasPeer<Agent>,
    R: RunWithConnectionTo<Counterpart>,
{
    /// Run this session synchronously. The current task will be blocked
    /// and `op` will be executed with the active session information.
    /// This is useful when you have MCP servers that are borrowed from your local
    /// stack frame.
    ///
    /// The [`SessionSetup`] passed to `op` holds an `ActiveSession` with a
    /// non-`'static` lifetime, which
    /// prevents calling [`ActiveSession::proxy_remaining_messages`] (since the
    /// session's background runners would terminate when `op` returns).
    ///
    /// Requires calling [`block_task`](Self::block_task) first.
    pub async fn run_until<T>(
        self,
        op: impl for<'runner> AsyncFnOnce(
            SessionSetup<'runner, Counterpart, NewSessionResponse>,
        ) -> Result<T, crate::Error>,
    ) -> Result<T, crate::Error> {
        let Self {
            connection,
            request,
            dynamic_handler_registrations,
            run,
            block_state: _,
        } = self;

        let cleanup = session_cleanup(&dynamic_handler_registrations);
        let (runner_connection, error_scope) =
            session_runner_connection(connection.raw_connection().clone(), &cleanup);
        run_session_scope(
            run.run_with_connection_to(runner_connection),
            async move {
                let response = connection
                    .send_request_to(Agent, request)
                    .block_task()
                    .await?;
                let setup = connection.attach_session(response, dynamic_handler_registrations)?;
                op(setup).await
            },
            cleanup,
            error_scope,
        )
        .await
    }

    /// Send the request to create the session and return a handle.
    /// This is an alternative to [`Self::run_until`] that avoids rightward
    /// drift but at the cost of requiring MCP servers that are `Send` and
    /// don't access data from the surrounding scope.
    ///
    /// Returns a [`SessionSetup`] with the `ActiveSession<'static, _>` and the
    /// complete [`NewSessionResponse`]. The session is `'static` because its
    /// runners are spawned into background tasks that live for the connection
    /// lifetime.
    ///
    /// An error response to `session/new` is returned to the caller and leaves
    /// the connection running. The session's dynamic handlers are removed, so
    /// its runners stop once their cleanup finishes. Session routing is
    /// installed before later inbound messages are dispatched.
    ///
    /// Requires calling [`block_task`](Self::block_task) first.
    pub async fn start_session(
        self,
    ) -> Result<SessionSetup<'static, Counterpart, NewSessionResponse>, crate::Error>
    where
        R: 'static,
    {
        let Self {
            connection,
            request,
            dynamic_handler_registrations,
            run,
            block_state: _,
        } = self;

        let cleanup = session_cleanup(&dynamic_handler_registrations);
        let (runner_connection, error_scope) =
            session_runner_connection(connection.raw_connection().clone(), &cleanup);
        connection.spawn(run_attached_session_runner(
            run.run_with_connection_to(runner_connection),
            cleanup,
            error_scope,
        ))?;

        let session_connection = connection.clone();
        connection
            .raw_connection()
            .send_ordered_request_to::<Raw, _, _>(Agent, request)
            .block_task_with_ordered_result(move |result| {
                let response = result?;
                session_connection.attach_session(response, dynamic_handler_registrations)
            })
            .await
    }

    /// Start a proxy session that forwards all messages between client and agent.
    ///
    /// A **proxy session** starts the session with the agent and then automatically
    /// proxies all session updates (prompts, tool calls, etc.) from the agent back
    /// to the client. You don't need to handle any messages yourself - the proxy
    /// takes care of forwarding everything. This is useful when you want to inject
    /// and/or filter prompts coming from the client but otherwise not be involved
    /// in the session.
    ///
    /// This is a convenience method that combines [`start_session`](Self::start_session),
    /// responding to the client, and [`ActiveSession::proxy_remaining_messages`].
    ///
    /// For more control (e.g., to send some messages before proxying), use
    /// [`start_session`](Self::start_session) instead and call
    /// [`proxy_remaining_messages`](ActiveSession::proxy_remaining_messages) manually.
    ///
    /// Returns a clone of the complete [`NewSessionResponse`] forwarded to the
    /// client.
    ///
    /// Requires calling [`block_task`](Self::block_task) first.
    pub async fn start_session_proxy(
        self,
        responder: Responder<NewSessionResponse, crate::V1>,
    ) -> Result<NewSessionResponse, crate::Error>
    where
        Counterpart: HasPeer<Client>,
        R: 'static,
    {
        let SessionSetup {
            session, response, ..
        } = self.start_session().await?;
        responder.respond(response.clone())?;
        session.proxy_remaining_messages()?;
        Ok(response)
    }
}

/// Stable protocol v1 active session that lets you send prompts and receive updates.
///
/// The session owns a stream of [`SessionMessage`]s: the requests and
/// notifications the agent sends for this session, and the responses to
/// prompts sent with [`send_prompt`](Self::send_prompt) or
/// [`send_prompt_blocks`](Self::send_prompt_blocks).
///
/// This handle covers prompting and reading the stream only. It has no
/// helpers for `session/cancel`, `session/set_mode`,
/// `session/set_config_option`, `session/close` or `session/fork`. Send those
/// messages through [`connection`](Self::connection) with
/// [`session_id`](Self::session_id), for example
/// `session.connection().send_notification_to(Agent, CancelNotification::new(session.session_id().clone()))`.
///
/// The `'runner` lifetime represents the span during which session support runners
/// (such as MCP servers) are active. When created via [`SessionBuilder::start_session`],
/// this is `'static` because the runners are spawned into background tasks.
/// When created via [`SessionBuilder::run_until`], this is tied to the
/// closure scope, preventing [`Self::proxy_remaining_messages`] from being called
/// (since the runners would stop when the closure returns).
#[derive(Debug)]
pub struct ActiveSession<'runner, Link>
where
    Link: HasPeer<Agent>,
{
    session_id: SessionId,
    update_rx: mpsc::UnboundedReceiver<SessionMessage>,
    update_tx: mpsc::UnboundedSender<SessionMessage>,
    connection: ConnectionTo<Link>,

    /// Registration for the handler that routes session messages to `update_rx`.
    /// This is separate from MCP handlers so it can be dropped independently
    /// when switching to proxy mode.
    session_handler_registration: DynamicHandlerGuard<Link>,

    /// Registrations for MCP server handlers.
    /// These will be dropped once the active-session struct is dropped
    /// which will cause them to be deregistered.
    mcp_handler_registrations: Vec<DynamicHandlerGuard<Link>>,

    /// Phantom lifetime representing the session-runner lifetime.
    _runner: PhantomData<&'runner ()>,
}

/// An item on a stable protocol v1 session's update stream, returned by
/// [`ActiveSession::read_update`].
#[non_exhaustive]
#[derive(Debug)]
#[allow(
    clippy::large_enum_variant,
    reason = "messages vastly outnumber prompt responses; boxing would add a heap allocation"
)]
pub enum SessionMessage {
    /// A request or notification the agent sent for this session, such as a
    /// `session/update` notification or a `session/request_permission`
    /// request.
    ///
    /// Request handles are typed for protocol v1. To match on the message
    /// type with [`MatchDispatch`], convert it with
    /// [`Dispatch::into_raw`](crate::Dispatch::into_raw) first.
    Message(crate::v1::Dispatch),

    /// The response to a `session/prompt` request sent with
    /// [`ActiveSession::send_prompt`] or
    /// [`ActiveSession::send_prompt_blocks`], which ends that prompt turn.
    ///
    /// An error response from the agent arrives here as `Err`. It does not
    /// close the connection or affect other sessions.
    PromptResponse(Result<PromptResponse, crate::Error>),
}

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

    /// Access the underlying connection context used to communicate with the agent.
    pub fn connection(&self) -> &ConnectionTo<Link> {
        &self.connection
    }

    /// Send a text prompt to the agent.
    ///
    /// Read the agent's output with [`read_update`](Self::read_update) or
    /// [`read_to_string`](Self::read_to_string). The `session/prompt`
    /// response, or the error the agent answered with, arrives on the same
    /// stream as [`SessionMessage::PromptResponse`].
    pub fn send_prompt(&mut self, prompt: impl ToString) -> Result<(), crate::Error> {
        self.send_prompt_blocks(vec![prompt.to_string().into()])
    }

    /// Send a prompt made of arbitrary content blocks to the agent.
    ///
    /// Content other than text and resource links needs the matching
    /// `promptCapabilities` entry from the agent's initialize response. Like
    /// [`send_prompt`](Self::send_prompt), the
    /// response arrives on the update stream as
    /// [`SessionMessage::PromptResponse`].
    pub fn send_prompt_blocks(&mut self, prompt: Vec<ContentBlock>) -> Result<(), crate::Error> {
        let update_tx = self.update_tx.clone();
        self.connection
            .raw_connection()
            .send_ordered_request_to::<Raw, _, _>(
                Agent,
                PromptRequest::new(self.session_id.clone(), prompt),
            )
            .on_receiving_result(async move |result| {
                // A closed channel means the session was dropped or handed to
                // `proxy_remaining_messages`, so nobody reads the response.
                drop(update_tx.unbounded_send(SessionMessage::PromptResponse(result)));
                Ok(())
            })
    }

    /// Read the next item on the session's update stream: a message the agent
    /// sent for this session, or the response to a prompt.
    pub async fn read_update(&mut self) -> Result<SessionMessage, crate::Error> {
        use futures::StreamExt;
        let message =
            self.update_rx.next().await.ok_or_else(|| {
                crate::util::internal_error("session channel closed unexpectedly")
            })?;

        Ok(message)
    }

    /// Read updates until the next prompt response and return the text of the
    /// agent's message chunks.
    ///
    /// Ignores non-text updates. Requests from the agent that reach the
    /// session, such as `session/request_permission`, are answered with a
    /// method not found error (`-32601`). Use
    /// [`read_update`](Self::read_update) instead to answer them yourself. If
    /// the agent answers the prompt with an error, this returns that error.
    pub async fn read_to_string(&mut self) -> Result<String, crate::Error> {
        let mut output = String::new();
        loop {
            let update = self.read_update().await?;
            tracing::trace!("read_to_string update");
            match update {
                SessionMessage::Message(dispatch) => MatchDispatch::new(dispatch.into_raw())
                    .if_notification(async |notif: SessionNotification| match notif.update {
                        SessionUpdate::AgentMessageChunk(ContentChunk {
                            content: ContentBlock::Text(text),
                            ..
                        }) => {
                            output.push_str(&text.text);
                            Ok(())
                        }
                        _ => Ok(()),
                    })
                    .await
                    .otherwise_ignore()?,
                SessionMessage::PromptResponse(response) => {
                    response?;
                    break;
                }
            }
        }
        Ok(output)
    }
}

impl<Link> ActiveSession<'static, Link>
where
    Link: HasPeer<Agent>,
{
    /// Proxy all remaining messages for this session between client and agent.
    ///
    /// Use this when you want to inject MCP servers into a session but don't need
    /// to actively interact with it after setup. The session messages will be proxied
    /// between client and agent automatically.
    ///
    /// This consumes the `ActiveSession` since you're giving up active control.
    ///
    /// This method is only available on `ActiveSession<'static, _>` (from
    /// [`SessionBuilder::start_session`]) because it requires the session's runners to outlive
    /// the method call.
    ///
    /// # Message Ordering Guarantees
    ///
    /// This method ensures proper handoff from active session mode to proxy mode
    /// without losing or reordering messages:
    ///
    /// 1. **Stop the session handler** - Deactivate the handler that routes messages
    ///    to `update_rx`. After this, no new messages will be queued.
    /// 2. **Close the channel** - Drop `update_tx` so we can detect when the channel
    ///    is fully drained.
    /// 3. **Drain queued messages** - Forward any messages that were already queued
    ///    in `update_rx` to the client, preserving order.
    /// 4. **Install proxy handler** - Now that all queued messages are forwarded,
    ///    install the proxy handler to handle future messages.
    /// 5. **Remove the session handler** - Drop its registration.
    ///
    /// This sequence prevents the race condition where messages could be delivered
    /// out of order or lost during the transition. Messages that arrive between
    /// steps 1 and 4 wait for the proxy handler. The deactivated handler stays
    /// registered until the proxy handler is installed, so a request in that
    /// window waits too instead of being rejected.
    pub fn proxy_remaining_messages(self) -> Result<(), crate::Error>
    where
        Link: HasPeer<Client>,
    {
        // Destructure self to get ownership of all fields
        let ActiveSession {
            session_id,
            mut update_rx,
            update_tx,
            connection,
            session_handler_registration,
            mcp_handler_registrations,
            _runner,
        } = self;

        // Step 1: Deactivate the session handler.
        // After this point, no new messages will be added to the channel.
        // Its removal is queued only after the proxy handler (step 5): while a
        // deactivated handler is registered, the incoming actor parks session
        // requests for the replacement instead of rejecting them.
        session_handler_registration.deactivate();

        // Step 2: Drop the sender side of the channel.
        // This allows us to detect when the channel is fully drained
        // (recv will return None when empty and sender is dropped).
        drop(update_tx);

        // Step 3: Drain any messages that were already queued and forward to client.
        // These messages arrived before we dropped the handler but haven't been
        // consumed yet. We must forward them to maintain message ordering.
        while let Ok(message) = update_rx.try_recv() {
            match message {
                SessionMessage::Message(dispatch) => {
                    // Forward the message to the client
                    connection.send_proxied_message_to(Client, dispatch)?;
                }
                SessionMessage::PromptResponse(_) => {
                    // Responses to the proxy's own prompts are not forwarded
                }
            }
        }

        // Step 4: Install the proxy handler for future messages.
        // Now that all queued messages have been forwarded, the proxy handler
        // can take over. Any new messages will go directly through the proxy.
        connection
            .add_dynamic_handler(ProxySessionMessages::new(session_id))?
            .detach();

        // Step 5: Remove the deactivated session handler.
        drop(session_handler_registration);

        // Keep MCP server handlers alive for the lifetime of the proxy
        for registration in mcp_handler_registrations {
            registration.detach();
        }

        Ok(())
    }
}

struct ActiveSessionHandler {
    session_id: SessionId,
    update_tx: mpsc::UnboundedSender<SessionMessage>,
}

impl ActiveSessionHandler {
    pub fn new(session_id: SessionId, update_tx: mpsc::UnboundedSender<SessionMessage>) -> Self {
        Self {
            session_id,
            update_tx,
        }
    }
}

impl<Counterpart: Role> HandleDispatchFrom<Counterpart> for ActiveSessionHandler
where
    Counterpart: HasPeer<Agent>,
{
    async fn handle_dispatch_from(
        &mut self,
        message: Dispatch,
        cx: RawConnectionTo<Counterpart>,
    ) -> Result<Handled<Dispatch>, crate::Error> {
        // If this is a message for our session, grab it.
        tracing::trace!(
            method = message.method(),
            handler_session_id = ?self.session_id,
            "ActiveSessionHandler::handle_dispatch"
        );
        MatchDispatchFrom::new(message, &cx)
            .if_dispatch_from(Agent, async |message| {
                if let Some(session_id) = message.get_session_id()? {
                    tracing::trace!(
                        message_session_id = ?session_id,
                        handler_session_id = ?self.session_id,
                        "ActiveSessionHandler::handle_dispatch"
                    );
                    if session_id == self.session_id {
                        self.update_tx
                            .unbounded_send(SessionMessage::Message(message.into_protocol()?))
                            .map_err(crate::util::internal_error)?;
                        return Ok(Handled::Yes);
                    }
                }

                // Otherwise, pass it through.
                Ok(Handled::No {
                    message,
                    retry: false,
                })
            })
            .await
            .done()
    }

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