agent_client_protocol/

session.rs

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

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

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

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

#[cfg(feature = "unstable_protocol_v2")]
mod v2;
#[cfg(feature = "unstable_protocol_v2")]
pub use 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: ConnectionTo<R>,
    cleanup: &SessionCleanup,
) -> (ConnectionTo<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 SessionBlockState for Blocking {}

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

/// Trait for marker types that indicate blocking vs non-blocking API.
/// See [`SessionBuilder::block_task`].
pub trait SessionBlockState: 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
    /// `V2ConnectionTo` with its own v2 `build_session` helper.
    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`.
    pub fn load_session(
        &self,
        session_id: impl Into<SessionId>,
        cwd: impl AsRef<Path>,
    ) -> RestoreSessionBuilder<Counterpart, LoadSessionRequest> {
        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,
    ) -> RestoreSessionBuilder<Counterpart, LoadSessionRequest> {
        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`.
    pub fn resume_session(
        &self,
        session_id: impl Into<SessionId>,
        cwd: impl AsRef<Path>,
    ) -> RestoreSessionBuilder<Counterpart, ResumeSessionRequest> {
        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,
    ) -> RestoreSessionBuilder<Counterpart, ResumeSessionRequest> {
        RestoreSessionBuilder::new(self, request)
    }

    /// Given a session response received from the agent,
    /// attach a handler to process messages related to this session
    /// and let you access them.
    ///
    /// 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<ActiveSession<'runner, Counterpart>, crate::Error> {
        let NewSessionResponse {
            session_id,
            modes,
            config_options,
            meta,
            ..
        } = response;

        let prepared = self.prepare_session_routing(&session_id)?;
        Ok(prepared.into_active_session(
            self.clone(),
            session_id,
            modes,
            config_options,
            meta,
            mcp_handler_registrations,
        ))
    }

    /// 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,
        modes: Option<SessionModeState>,
        config_options: Option<Vec<SessionConfigOption>>,
        meta: Option<Meta>,
        mcp_handler_registrations: Vec<DynamicHandlerGuard<Counterpart>>,
    ) -> ActiveSession<'runner, Counterpart> {
        ActiveSession {
            session_id,
            modes,
            config_options,
            meta,
            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;
    fn response_modes(response: &Self::Response) -> Option<SessionModeState>;
    fn response_config_options(response: &Self::Response) -> Option<Vec<SessionConfigOption>>;
    fn response_meta(response: &Self::Response) -> Option<Meta>;
}

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

    fn response_modes(response: &Self::Response) -> Option<SessionModeState> {
        response.modes.clone()
    }

    fn response_config_options(response: &Self::Response) -> Option<Vec<SessionConfigOption>> {
        response.config_options.clone()
    }

    fn response_meta(response: &Self::Response) -> Option<Meta> {
        response.meta.clone()
    }
}

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

    fn response_modes(response: &Self::Response) -> Option<SessionModeState> {
        response.modes.clone()
    }

    fn response_config_options(response: &Self::Response) -> Option<Vec<SessionConfigOption>> {
        response.config_options.clone()
    }

    fn response_meta(response: &Self::Response) -> Option<Meta> {
        response.meta.clone()
    }
}

/// Stable protocol v1 builder for `session/load` or `session/resume`.
///
/// 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>,
}

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,
) -> RestoredSession<'static, Counterpart, Request::Response>
where
    Counterpart: HasPeer<Agent>,
    Request: RestoreRequest,
{
    let session = prepared.into_active_session(
        connection,
        session_id,
        Request::response_modes(&response),
        Request::response_config_options(&response),
        Request::response_meta(&response),
        Vec::new(),
    );

    RestoredSession { 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(RestoredSession<'static, Counterpart, Request::Response>) -> Fut + Send + 'static,
    Fut: Future<Output = Result<(), crate::Error>> + Send,
{
    ensure_v1_session_protocol(&builder.connection)?;

    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.dynamic_handler_barrier();

    connection
        .send_ordered_request_to_after(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<RestoredSession<'static, Counterpart, Request::Response>, crate::Error>
where
    Counterpart: HasPeer<Agent>,
    Request: RestoreRequest,
{
    ensure_v1_session_protocol(&builder.connection)?;

    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.dynamic_handler_barrier();
    let session_connection = connection.clone();

    connection
        .send_ordered_request_to_after(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(RestoredSession<'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(RestoredSession<'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`](RestoreSessionBuilder::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<RestoredSession<'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`](RestoreSessionBuilder::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<RestoredSession<'static, Counterpart, ResumeSessionResponse>, crate::Error> {
        start_restored_session(self).await
    }
}

/// A restored stable-v1 session and the exact operation response that opened
/// it.
///
/// The session ID comes from the load or resume request because stable-v1
/// restore responses do not repeat it. Keeping the response separate preserves
/// every operation-specific field without reconstructing it from session
/// state.
pub struct RestoredSession<'runner, Link, Response>
where
    Link: HasPeer<Agent>,
{
    session: ActiveSession<'runner, Link>,
    response: Response,
}

impl<'runner, Link, Response> RestoredSession<'runner, Link, Response>
where
    Link: HasPeer<Agent>,
{
    /// Access the active session.
    pub fn session(&self) -> &ActiveSession<'runner, Link> {
        &self.session
    }

    /// Mutably access the active session, for example to consume replay.
    pub fn session_mut(&mut self) -> &mut ActiveSession<'runner, Link> {
        &mut self.session
    }

    /// Access the complete load or resume response.
    pub fn response(&self) -> &Response {
        &self.response
    }

    /// Split the restored value into its active session and exact response.
    pub fn into_parts(self) -> (ActiveSession<'runner, Link>, Response) {
        (self.session, self.response)
    }

    /// Consume this value and return only the active session.
    pub fn into_session(self) -> ActiveSession<'runner, Link> {
        self.session
    }
}

impl<Link, Response> std::fmt::Debug for RestoredSession<'_, Link, Response>
where
    Link: HasPeer<Agent>,
    Response: std::fmt::Debug,
{
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter
            .debug_struct("RestoredSession")
            .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.
    #[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)?);
        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 an `ActiveSession<'static, _>` 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::McpServerExt;
    /// # async fn example(transport: impl ConnectTo<Client>) -> Result<(), agent_client_protocol::Error> {
    /// # Client.builder().connect_with(transport, async |cx| {
    /// # let mcp = McpServer::<Agent, _>::builder("tools").build();
    /// cx.build_session_cwd()?
    ///     .with_mcp_server(mcp)?
    ///     .on_session_start(async |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(ActiveSession<'static, Counterpart>) -> Fut + Send + 'static,
        Fut: Future<Output = Result<(), crate::Error>> + Send,
    {
        ensure_v1_session_protocol(&self.connection)?;

        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.clone(), &cleanup);
        connection.spawn(run_attached_session_runner(
            run.run_with_connection_to(runner_connection),
            cleanup,
            error_scope,
        ))?;

        connection
            .send_ordered_request_to(Agent, request)
            .on_receiving_result({
                let connection = connection.clone();
                async move |result| {
                    let response = result?;

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

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

    /// Spawn a proxy session and run a closure with the session ID.
    ///
    /// 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 the `SessionId` once the session is established. 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::schema::v1::NewSessionRequest;
    /// # use agent_client_protocol::mcp_server::McpServer;
    /// # use agent_client_protocol_rmcp::McpServerExt;
    /// # async fn example(transport: impl ConnectTo<Proxy>) -> Result<(), agent_client_protocol::Error> {
    /// Proxy.builder()
    ///     .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 |session_id| {
    ///                 // 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>,
        op: F,
    ) -> Result<(), crate::Error>
    where
        F: FnOnce(SessionId) -> Fut + Send + 'static,
        Fut: Future<Output = Result<(), crate::Error>> + Send,
        Counterpart: HasPeer<Client>,
        R: 'static,
    {
        ensure_v1_session_protocol(&self.connection)?;

        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.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.send_ordered_request_to(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)?;

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

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

                connection.spawn(async move { op(session_id).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 `ActiveSession` passed to `op` has 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(
            ActiveSession<'runner, Counterpart>,
        ) -> 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.clone(), &cleanup);
        run_session_scope(
            run.run_with_connection_to(runner_connection),
            async move {
                ensure_v1_session_protocol(&connection)?;
                let response = connection
                    .send_request_to(Agent, request)
                    .block_task()
                    .await?;
                let active_session =
                    connection.attach_session(response, dynamic_handler_registrations)?;
                op(active_session).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 an `ActiveSession<'static, _>` because the session's runners are spawned into
    /// background tasks that live for the connection lifetime.
    ///
    /// Requires calling [`block_task`](Self::block_task) first.
    pub async fn start_session(self) -> Result<ActiveSession<'static, Counterpart>, crate::Error>
    where
        R: 'static,
    {
        ensure_v1_session_protocol(&self.connection)?;

        let Self {
            connection,
            request,
            dynamic_handler_registrations,
            run,
            block_state: _,
        } = self;

        let (active_session_tx, active_session_rx) = oneshot::channel();

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

        connection.clone().spawn(async move {
            let response = connection
                .send_request_to(Agent, request)
                .block_task()
                .await?;

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

            active_session_tx
                .send(active_session)
                .map_err(|_| crate::Error::internal_error())?;

            Ok(())
        })?;

        active_session_rx
            .await
            .map_err(|_| crate::Error::internal_error())
    }

    /// 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.
    ///
    /// Requires calling [`block_task`](Self::block_task) first.
    pub async fn start_session_proxy(
        self,
        responder: Responder<NewSessionResponse>,
    ) -> Result<SessionId, crate::Error>
    where
        Counterpart: HasPeer<Client>,
        R: 'static,
    {
        let active_session = self.start_session().await?;
        let session_id = active_session.session_id().clone();
        responder.respond(active_session.response())?;
        active_session.proxy_remaining_messages()?;
        Ok(session_id)
    }
}

/// Stable protocol v1 active session that lets you send prompts and receive updates.
///
/// 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>,
    modes: Option<SessionModeState>,
    config_options: Option<Vec<SessionConfigOption>>,
    meta: Option<serde_json::Map<String, serde_json::Value>>,
    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 ()>,
}

/// Incoming stable protocol v1 message from the agent.
#[non_exhaustive]
#[derive(Debug)]
#[allow(
    clippy::large_enum_variant,
    reason = "Dispatch messages vastly outnumber StopReason; boxing would add a heap allocation"
)]
pub enum SessionMessage {
    /// Periodic updates with new content, tool requests, etc.
    /// Use [`MatchDispatch`] to match on the message type.
    SessionMessage(Dispatch),

    /// When a prompt completes, the stop reason.
    StopReason(StopReason),
}

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

    /// Access modes available in this session.
    pub fn modes(&self) -> Option<&SessionModeState> {
        self.modes.as_ref()
    }

    /// Access the initial session configuration options returned by the agent.
    pub fn config_options(&self) -> Option<&[SessionConfigOption]> {
        self.config_options.as_deref()
    }

    /// Access meta data from session response.
    pub fn meta(&self) -> Option<&serde_json::Map<String, serde_json::Value>> {
        self.meta.as_ref()
    }

    /// Build a `NewSessionResponse` from the session information.
    ///
    /// Useful when you need to forward the session response to a client
    /// after doing some processing.
    pub fn response(&self) -> NewSessionResponse {
        NewSessionResponse::new(self.session_id.clone())
            .modes(self.modes.clone())
            .config_options(self.config_options.clone())
            .meta(self.meta.clone())
    }

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

    /// Send a prompt to the agent. You can then read messages sent in response.
    pub fn send_prompt(&mut self, prompt: impl ToString) -> Result<(), crate::Error> {
        let update_tx = self.update_tx.clone();
        self.connection
            .send_ordered_request_to(
                Agent,
                PromptRequest::new(self.session_id.clone(), vec![prompt.to_string().into()]),
            )
            .on_receiving_result(async move |result| {
                let PromptResponse { stop_reason, .. } = result?;

                update_tx
                    .unbounded_send(SessionMessage::StopReason(stop_reason))
                    .map_err(crate::util::internal_error)?;

                Ok(())
            })
    }

    /// Read an update from the agent in response to the 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 all updates until the end of the turn and create a string.
    /// Ignores non-text updates.
    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!(?update, "read_to_string update");
            match update {
                SessionMessage::SessionMessage(dispatch) => MatchDispatch::new(dispatch)
                    .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::StopReason(_stop_reason) => 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** - Drop the registration 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.
    ///
    /// This sequence prevents the race condition where messages could be delivered
    /// out of order or lost during the transition.
    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,
            // These fields are not needed for proxying
            modes: _,
            config_options: _,
            meta: _,
            _runner,
        } = self;

        // Step 1: Drop the session handler registration.
        // This unregisters the handler that was routing messages to update_rx.
        // After this point, no new messages will be added to the channel.
        drop(session_handler_registration);

        // 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::SessionMessage(dispatch) => {
                    // Forward the message to the client
                    connection.send_proxied_message_to(Client, dispatch)?;
                }
                SessionMessage::StopReason(_) => {
                    // StopReason is internal bookkeeping, 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();

        // 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: ConnectionTo<Counterpart>,
    ) -> Result<Handled<Dispatch>, crate::Error> {
        // If this is a message for our session, grab it.
        tracing::trace!(
            ?message,
            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::SessionMessage(message))
                            .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)
    }
}

#[cfg(not(feature = "unstable_protocol_v2"))]
#[allow(
    clippy::unnecessary_wraps,
    reason = "signature matches the feature-enabled protocol guard"
)]
fn ensure_v1_session_protocol<Counterpart: Role>(
    _connection: &ConnectionTo<Counterpart>,
) -> Result<(), crate::Error> {
    Ok(())
}

#[cfg(feature = "unstable_protocol_v2")]
fn ensure_v1_session_protocol<Counterpart: Role>(
    connection: &ConnectionTo<Counterpart>,
) -> Result<(), crate::Error> {
    if connection.acp_protocol_version() != Some(crate::schema::ProtocolVersion::V2) {
        return Ok(());
    }

    Err(crate::Error::invalid_request().data(
        "stable session builders use ACP protocol v1 types, but this is a protocol v2 connection; \
         use the `V2ConnectionTo` supplied to `Client.v2()` callbacks",
    ))
}