polyc-tools 2026.8.3

//! Process-local reuse of live MCP connector sessions, keyed by
//! (connector, principal).
//!
//! Cold-dialing every connector on every turn wastes the `initialize` handshake
//! a live session already paid for. A [`ConnectionPool`] caches the live session
//! so a later turn of the SAME conversation reuses it. The key carries the
//! principal (the conversation) as well as the connector, so a session is NEVER
//! shared across principals — a stateful MCP server could otherwise leak one
//! conversation's server-side context into another (invariant 6 of `#582`). The
//! connector's bearer token is namespace-scoped, so URL + token alone is not
//! principal-distinct; the principal must be part of the key.
//!
//! # What is cached, and what is not
//!
//! Only the live SESSION is cached — never a connector's advertised tool
//! specs. Catalog freshness is not this pool's concern and deliberately does
//! not ride its eviction clock: a connector's `tools/list` response carries
//! its own SEP-2549 `ttlMs`/`cacheScope` hint, and rmcp's per-peer response
//! cache is what honours it (configured in `mcp_client`'s `dial_service`).
//! Every server in this surface stamps `ttlMs: 0`, so in practice a pooled
//! session re-lists on every turn and no catalog outlives the turn that
//! fetched it.
//!
//! `#2275` briefly cached a live-listed catalog alongside its session entry,
//! pinning it for the session's whole lifetime. That bought one `tools/list`
//! round trip per reused session and cost INV-14: a connector could
//! re-declare a tool destructive and the conversation would keep gating it
//! under the old, more permissive declaration until the session was evicted.
//! The session's eviction clock (`IDLE_TTL`, a rotated bearer, a dead peer)
//! is not a freshness signal about the remote's catalog, so that cache is
//! gone. See `docs/specifications/invariants-mcp.md`'s INV-14.
//!
//! # Rotation and dead peers
//!
//! Each entry remembers the bearer it dialed with. A descriptor arriving with a
//! different bearer for the same key (the acquire path sees the stored header no
//! longer matches) evicts the stale entry and re-dials with the new token —
//! token rotation (`#395`) awareness without a health check. A call that fails
//! with the transport kind after its retry budget, or is auth-rejected, evicts
//! the entry so the next call re-dials fresh.
//!
//! # Idle reclamation
//!
//! An entry unused for longer than `IDLE_TTL` is dropped opportunistically on
//! the next acquire (under the same lock, no extra await).
//! In the shared-Service production mode one pool serves many conversations, so
//! a departed conversation's session would otherwise pin a remote's server-side
//! session and a local transport task for the pod's lifetime. Re-dialing is
//! cheap (one handshake), so reclaiming beats hoarding.
//!
//! # Concurrency
//!
//! At most one entry is kept per key, but concurrent misses for the same key
//! each dial and the LAST write wins — the simpler of the two (single-flight vs.
//! last-write-wins) options. A superseded session's `Arc` drops and cancels its
//! own transport; any in-flight call on it still completes against the clone the
//! caller already holds. The lock is a plain [`std::sync::Mutex`] never held
//! across the dial `await`, so acquiring is cancellation-safe.

use std::{
    collections::HashMap,
    sync::{Arc, Mutex},
    time::{Duration, Instant},
};

use polyc_crypto::sensitive::Sensitive;
use rmcp::service::{RoleClient, RunningService};

use crate::mcp_client::{McpClientError, dial_service};

/// How long a cached session may sit idle before the next [`ConnectionPool::acquire`]
/// reclaims it.
///
/// A session is cheap to re-dial — a single `initialize` round-trip — so a
/// long-idle one earns nothing while it pins the remote's server-side session
/// and a local transport task. In the shared-Service production mode one pool
/// serves many conversations, so without a ceiling a departed conversation's
/// session lingers for the pod's whole lifetime. Fifteen minutes sits well above
/// the gap between turns of an active conversation (a live chat keeps its warm
/// session) yet reclaims an abandoned one promptly.
pub(crate) const IDLE_TTL: Duration = Duration::from_mins(15);

/// The identity a cached session is keyed by: the connector's canonical resource
/// and the principal (conversation) it was dialed for.
///
/// The principal is load-bearing: a session is never shared across principals,
/// so two conversations reaching the same connector hold two distinct sessions.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub(crate) struct ConnectionKey {
    /// The connector's canonical resource indicator (RFC 8707), shared by every
    /// turn that dials the same endpoint.
    connector: String,
    /// The principal a session is scoped to — the conversation routing key.
    principal: String,
}

impl ConnectionKey {
    /// Key a connector session by its canonical `connector` resource and the
    /// `principal` (conversation) it serves.
    pub(crate) fn new(connector: impl Into<String>, principal: impl Into<String>) -> Self {
        Self {
            connector: connector.into(),
            principal: principal.into(),
        }
    }
}

/// One cached live session plus the bearer it was dialed with (so a rotated
/// token forces a re-dial) and when it was last used (so an idle one is
/// reclaimed).
struct CachedSession {
    /// The live MCP client, shared with every source that reuses it.
    service: Arc<RunningService<RoleClient, ()>>,
    /// The `Authorization` header presented at dial (`None` for an anonymous
    /// dial). A later acquire carrying a different header means the token
    /// rotated, so the entry is evicted and re-dialed. Compared by value rather
    /// than by a hash: the raw token already lives in this session's live
    /// transport in the same address space, so hashing the copy kept here would
    /// add ceremony without shrinking the trust boundary.
    auth_header: Option<String>,
    /// The caller identity header presented at dial (`None` for a turn with no
    /// attributed persona). A later acquire carrying a different caller under
    /// the SAME (connector, principal) key — a group conversation where a
    /// different participant's turn draws next — must never reuse this
    /// session: the header rides the transport's fixed request config, so a
    /// stale value would silently misattribute a later caller's calls to the
    /// caller who first dialed. Mismatched entries are evicted and re-dialed,
    /// exactly like a rotated `auth_header`.
    caller: Option<String>,
    /// When this session was last acquired, refreshed on every reuse. Drives the
    /// [`IDLE_TTL`] sweep in [`ConnectionPool::acquire`].
    last_used: Instant,
}

/// A process-local cache of live MCP connector sessions, keyed by
/// (connector, principal). Cheap to clone — clones share the one backing store.
///
/// Owned by the harness process (one instance shared across the turns it serves)
/// and passed into the per-turn tool composition, never a global static.
#[derive(Clone)]
pub struct ConnectionPool {
    /// The backing store, shared by every clone.
    inner: Arc<Mutex<HashMap<ConnectionKey, CachedSession>>>,
    /// How long an entry may sit idle before an acquire reclaims it. Defaults to
    /// [`IDLE_TTL`]; a test may shrink it via [`Self::with_idle_ttl`].
    idle_ttl: Duration,
}

impl Default for ConnectionPool {
    fn default() -> Self {
        Self {
            inner: Arc::default(),
            idle_ttl: IDLE_TTL,
        }
    }
}

impl std::fmt::Debug for ConnectionPool {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let live = self.inner.lock().map_or(0, |m| m.len());
        f.debug_struct("ConnectionPool")
            .field("live", &live)
            .field("idle_ttl", &self.idle_ttl)
            .finish()
    }
}

impl ConnectionPool {
    /// An empty pool with the default `IDLE_TTL`.
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// An empty pool whose idle sweep uses `idle_ttl` instead of [`IDLE_TTL`], so
    /// a test can force reclamation without waiting the production window.
    #[cfg(test)]
    pub(crate) fn with_idle_ttl(idle_ttl: Duration) -> Self {
        Self {
            inner: Arc::default(),
            idle_ttl,
        }
    }

    /// Reuse the live session cached under `key`, or dial a fresh one and cache
    /// it. A cached session dialed with a different bearer than `auth_header`
    /// (a rotated token) is evicted and re-dialed.
    ///
    /// The connect `timeout`, when set, bounds the whole handshake so a slow or
    /// dead peer cannot hang the caller; the pool's `Mutex` is never held across
    /// that `await`. A real dial logs `connected dynamic tool source` so a warm
    /// (reused) acquire is silent — the connect metric counts real dials only.
    ///
    /// # Errors
    ///
    /// Returns the classified [`McpClientError`] from [`dial_service`] on a
    /// handshake failure, or [`McpClientError::Timeout`] if the dial exceeds
    /// `timeout`.
    pub(crate) async fn acquire(
        &self,
        key: &ConnectionKey,
        uri: &Arc<str>,
        auth_header: Option<&str>,
        caller: Option<&str>,
        label: &str,
        timeout: Option<Duration>,
    ) -> Result<Arc<RunningService<RoleClient, ()>>, McpClientError> {
        // Fast path: a live entry dialed with the same bearer AND caller is
        // reused as-is. This same locked pass first reclaims any idle entries.
        if let Some(service) = self.reuse(key, auth_header, caller) {
            return Ok(service);
        }
        // Miss (or a rotated token / changed caller): dial fresh with no lock
        // held across the await. The connect budget bounds a slow or dead
        // handshake.
        let dial = dial_service(
            Arc::clone(uri),
            auth_header.map(str::to_owned),
            caller.map(str::to_owned),
            None,
        );
        let service = match timeout {
            Some(budget) => tokio::time::timeout(budget, dial)
                .await
                .map_err(|_elapsed| McpClientError::Timeout(budget))??,
            None => dial.await?,
        };
        let service = Arc::new(service);
        tracing::info!(label = %label, uri = %uri, "connected dynamic tool source");
        // Last-write-wins: a concurrent dial for the same key is overwritten and
        // its (now unreferenced-by-the-pool) session drops, cancelling its own
        // transport. The clone returned below keeps THIS caller's call alive.
        self.inner
            .lock()
            .expect("connection pool mutex poisoned")
            .insert(
                key.clone(),
                CachedSession {
                    service: Arc::clone(&service),
                    auth_header: auth_header.map(str::to_owned),
                    caller: caller.map(str::to_owned),
                    last_used: Instant::now(),
                },
            );
        Ok(service)
    }

    /// The live session cached under `key` when its stored bearer AND caller
    /// match `auth_header`/`caller`, cloning the shared `Arc` and refreshing
    /// the entry's last-used stamp. A mismatch (rotated token, or a different
    /// caller drawing the same principal-scoped key — a group conversation's
    /// next turn) or a miss returns `None`, so [`Self::acquire`] re-dials
    /// rather than handing a later caller a session whose fixed transport
    /// config still carries an earlier caller's identity header.
    ///
    /// The same locked pass first reclaims every entry idle past [`Self::idle_ttl`]
    /// — an opportunistic sweep with no extra lock or await.
    fn reuse(
        &self,
        key: &ConnectionKey,
        auth_header: Option<&str>,
        caller: Option<&str>,
    ) -> Option<Arc<RunningService<RoleClient, ()>>> {
        let mut map = self.inner.lock().expect("connection pool mutex poisoned");
        let now = Instant::now();
        // Idle sweep: drop any session unused past the TTL, cancelling its
        // transport if the pool held its last reference.
        map.retain(|_, entry| {
            if now.duration_since(entry.last_used) < self.idle_ttl {
                true
            } else {
                shutdown_if_orphaned(&entry.service);
                false
            }
        });
        let reused = map
            .get_mut(key)
            .filter(|entry| {
                entry.auth_header.as_deref() == auth_header && entry.caller.as_deref() == caller
            })
            .map(|entry| {
                entry.last_used = now;
                Arc::clone(&entry.service)
            });
        drop(map);
        reused
    }

    /// Drop the cached session under `key` so the next [`Self::acquire`] re-dials
    /// a fresh one. Called after a transport-dead or auth-rejected call; a no-op
    /// when nothing is cached.
    ///
    /// If the pool held the last reference the session's transport is cancelled;
    /// otherwise another caller still using it owns its shutdown.
    pub(crate) fn evict(&self, key: &ConnectionKey) {
        let evicted = self
            .inner
            .lock()
            .expect("connection pool mutex poisoned")
            .remove(key);
        if let Some(session) = evicted {
            shutdown_if_orphaned(&session.service);
        }
    }
}

/// Cancel `service`'s transport when the pool held its last reference, so a
/// dropped-from-the-pool session's background tasks exit promptly. A no-op while
/// any caller still holds a clone — that caller owns its shutdown.
fn shutdown_if_orphaned(service: &Arc<RunningService<RoleClient, ()>>) {
    if Arc::strong_count(service) == 1 {
        service.cancellation_token().cancel();
    }
}

/// How an [`McpToolSource`](crate::McpToolSource) obtains its live MCP session
/// for an `execute`.
///
/// The advertised specs are identical either way; only session ownership
/// differs, so every other method on the source is oblivious to it.
pub(crate) enum SessionHandle {
    /// A session the source dialed and OWNS: its transport is cancelled when the
    /// last clone drops. The eager `McpToolSource::connect` path.
    Direct(Arc<RunningService<RoleClient, ()>>),
    /// A session drawn on demand from a shared [`ConnectionPool`], keyed by
    /// (connector, principal). The pool owns the session's lifecycle: a dead or
    /// auth-rejected peer is evicted so the next call re-dials. The
    /// `McpToolSource::pooled` path.
    Pooled(PooledSession),
}

/// The state a [`SessionHandle::Pooled`] source needs to draw (dial-or-reuse)
/// and evict its session through the shared [`ConnectionPool`].
pub(crate) struct PooledSession {
    /// The shared pool this source draws from (cheap clone — shared store).
    pub(crate) pool: ConnectionPool,
    /// This source's (connector, principal) cache key.
    pub(crate) key: ConnectionKey,
    /// The connector endpoint dialed on a cache miss.
    pub(crate) uri: Arc<str>,
    /// The audience-checked bearer header presented at dial — already released
    /// from its `AudienceBoundToken` for this exact `uri` at compose time (a
    /// mismatch failed closed before advertising). The pool compares it to
    /// detect a rotated token. `None` dials anonymously. Held as
    /// [`Sensitive`] for this source's whole lifetime — exposed only at the
    /// point [`SessionHandle::acquire`] hands it to the pool.
    pub(crate) auth_header: Option<Sensitive<String>>,
    /// The turn's caller identity presented at dial (`None` for no attributed
    /// persona). The pool compares it on reuse exactly like `auth_header` —
    /// see [`CachedSession::caller`].
    pub(crate) caller: Option<String>,
    /// Connect budget bounding a lazy (first-`execute`) dial, so a dead peer
    /// discovered only at call time still cannot hang the turn.
    pub(crate) connect_timeout: Option<Duration>,
    /// Connector label for the pool's `connected dynamic tool source` log.
    pub(crate) label: String,
}

impl SessionHandle {
    /// The live session for a call: a [`Self::Direct`] source hands back its own
    /// client; a [`Self::Pooled`] source reuses (or dials) through the pool.
    pub(crate) async fn acquire(
        &self,
    ) -> Result<Arc<RunningService<RoleClient, ()>>, McpClientError> {
        match self {
            Self::Direct(service) => Ok(Arc::clone(service)),
            Self::Pooled(p) => {
                p.pool
                    .acquire(
                        &p.key,
                        &p.uri,
                        p.auth_header.as_ref().map(|h| h.expose().as_str()),
                        p.caller.as_deref(),
                        &p.label,
                        p.connect_timeout,
                    )
                    .await
            }
        }
    }

    /// Drop a pooled source's cached session so the next call re-dials (dead
    /// peer / rotated credentials). A no-op for a [`Self::Direct`] source, which
    /// owns its one session for its lifetime.
    pub(crate) fn evict(&self) {
        if let Self::Pooled(p) = self {
            p.pool.evict(&p.key);
        }
    }

    /// Cancel a [`Self::Direct`] source's transport. A [`Self::Pooled`] source
    /// leaves shutdown to the pool, which may still be lending the session to
    /// another turn of the same conversation.
    pub(crate) fn shutdown(&self) {
        if let Self::Direct(service) = self {
            service.cancellation_token().cancel();
        }
    }
}

#[cfg(test)]
mod tests {
    #![allow(clippy::pedantic, clippy::nursery, missing_docs)]

    use std::{borrow::Cow, net::SocketAddr};

    use rmcp::{
        ErrorData as McpError, ServerHandler,
        handler::server::{
            router::tool::ToolRouter,
            tool::{ToolCallContext, ToolRoute},
        },
        model::{
            CallToolRequestParams, CallToolResponse, CallToolResult, Implementation,
            InitializeResult, ListToolsResult, PaginatedRequestParams, ServerCapabilities, Tool,
        },
        service::{RequestContext, RoleServer},
        transport::streamable_http_server::{
            StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager,
        },
    };
    use serde_json::json;
    use tokio_util::sync::CancellationToken;

    use super::*;

    /// Minimal MCP server fixture: a single read-only `echo` tool, enough for a
    /// real `initialize` + `list_tools` handshake so the pool caches a live
    /// session.
    #[derive(Clone)]
    struct EchoServer {
        router: Arc<ToolRouter<Self>>,
    }

    impl EchoServer {
        fn new() -> Self {
            let mut router: ToolRouter<Self> = ToolRouter::new();
            let schema = json!({ "type": "object", "properties": {} });
            let echo = Tool::new(
                Cow::Borrowed("echo"),
                Cow::Borrowed("Echo."),
                schema.as_object().cloned().unwrap_or_default(),
            );
            router.add_route(ToolRoute::new_dyn(echo, |_ctx: ToolCallContext<Self>| {
                Box::pin(
                    async move { Ok(CallToolResult::structured(json!({ "echo": true })).into()) },
                )
            }));
            Self {
                router: Arc::new(router),
            }
        }
    }

    impl std::fmt::Debug for EchoServer {
        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            f.debug_struct("EchoServer").finish_non_exhaustive()
        }
    }

    impl ServerHandler for EchoServer {
        fn get_info(&self) -> rmcp::model::ServerInfo {
            InitializeResult::new(ServerCapabilities::builder().enable_tools().build())
                .with_server_info(Implementation::new("echo", env!("CARGO_PKG_VERSION")))
        }

        fn list_tools(
            &self,
            _request: Option<PaginatedRequestParams>,
            _context: RequestContext<RoleServer>,
        ) -> impl Future<Output = Result<ListToolsResult, McpError>> + Send + '_ {
            let tools = self.router.list_all();
            async move { Ok(ListToolsResult::with_all_items(tools)) }
        }

        fn call_tool(
            &self,
            request: CallToolRequestParams,
            context: RequestContext<RoleServer>,
        ) -> impl Future<Output = Result<CallToolResponse, McpError>> + Send + '_ {
            let router = self.router.clone();
            async move {
                router
                    .call(ToolCallContext::new(self, request, context))
                    .await
            }
        }
    }

    async fn spawn_server() -> (SocketAddr, CancellationToken, tokio::task::JoinHandle<()>) {
        let mut config = StreamableHttpServerConfig::default();
        config.legacy_session_mode = true;
        config.sse_keep_alive = None;
        let config = config.disable_allowed_hosts().disable_allowed_origins();
        let service = StreamableHttpService::new(
            || Ok(EchoServer::new()),
            Arc::new(LocalSessionManager::default()),
            config,
        );
        let router = axum::Router::new().nest_service("/mcp", service);
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        let ct = CancellationToken::new();
        let server_ct = ct.clone();
        let handle = tokio::spawn(async move {
            let _ = axum::serve(listener, router)
                .with_graceful_shutdown(async move { server_ct.cancelled_owned().await })
                .await;
        });
        tokio::time::sleep(Duration::from_millis(50)).await;
        (addr, ct, handle)
    }

    /// An entry idle past the pool's TTL is swept on the next acquire, forcing a
    /// re-dial — the returned session is a FRESH `Arc`, not the reclaimed one.
    /// The generous-TTL control in the same test proves it is the sweep, not a
    /// per-acquire re-dial, that supplies the new session.
    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn idle_session_is_swept_on_next_acquire() {
        let (addr, ct, handle) = spawn_server().await;
        let uri: Arc<str> = Arc::from(format!("http://{addr}/mcp").as_str());
        let key = ConnectionKey::new("http://conn/mcp", "conv-1");

        // Control: with a generous TTL the second acquire reuses the SAME session.
        let warm = ConnectionPool::new();
        let a = warm
            .acquire(&key, &uri, None, None, "notes", None)
            .await
            .unwrap();
        let b = warm
            .acquire(&key, &uri, None, None, "notes", None)
            .await
            .unwrap();
        assert!(
            Arc::ptr_eq(&a, &b),
            "a still-fresh session must be reused, not re-dialed"
        );

        // A sub-millisecond TTL: the first session is idle by the next acquire,
        // which sweeps it and re-dials a fresh one.
        let pool = ConnectionPool::with_idle_ttl(Duration::from_millis(1));
        let first = pool
            .acquire(&key, &uri, None, None, "notes", None)
            .await
            .unwrap();
        tokio::time::sleep(Duration::from_millis(20)).await;
        let second = pool
            .acquire(&key, &uri, None, None, "notes", None)
            .await
            .unwrap();
        assert!(
            !Arc::ptr_eq(&first, &second),
            "an entry idle past the TTL must be swept and re-dialed, not reused"
        );

        drop((a, b, first, second));
        ct.cancel();
        let _ = tokio::time::timeout(Duration::from_secs(5), handle).await;
    }
}