polyc-tools 2026.9.0

#![allow(clippy::unwrap_used)] // test/example/bench: panics are acceptable
//! Connection reuse and principal isolation at the tool-executor seam (#636,
//! #2273): a [`ConnectionPool`] keyed by (connector, principal) lets two turns
//! of the SAME conversation reuse one live MCP session instead of cold-dialing
//! every turn, while NEVER sharing a session across principals (invariant 6 of
//! #582 — a stateful server must not leak one conversation's context into
//! another).
//!
//! With sessions gone (#2272 — every dial is `server/discover`, never
//! `initialize`, and no `Mcp-Session-Id` is ever minted or sent), the
//! principal-isolation guarantee is re-expressed at the only grain left: every
//! `tools/call` a pooled source makes must carry exactly the caller header and
//! bearer of the principal that composed it, and NO request — from any
//! principal, at any point in the dial — may carry a session identifier. The
//! instrumented tower middleware below records, per request, the raw
//! `Authorization` header, the [`polyc_tools::CALLER_HEADER`] value, and
//! whether an `Mcp-Session-Id` header rode along, so the assertions read
//! straight off the wire rather than trusting an internal cache-key
//! comparison.

#![allow(clippy::pedantic, clippy::nursery, missing_docs)]

use std::{
    borrow::Cow,
    net::SocketAddr,
    sync::{Arc, Mutex},
    time::Duration,
};

use axum::{
    body::{Body, to_bytes},
    extract::{Request, State},
    middleware::{self, Next},
    response::Response,
};
use polyc_agent::ToolExecutor;
use polyc_tools::{
    AudienceBoundToken, CALLER_HEADER, ConnectOptions, ConnectionPool, McpToolSource,
};
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},
};
use serde_json::json;
use tokio_util::sync::CancellationToken;

/// One request's identity, as observed at the HTTP layer, plus which JSON-RPC
/// method it carried.
#[derive(Debug, Clone)]
struct SeenRequest {
    method: String,
    authorization: Option<String>,
    caller: Option<String>,
    /// Whether this request carried an `Mcp-Session-Id` header — must be
    /// `false` for EVERY request, from every principal, under the modern-only
    /// (2026-07-28) dial this workspace uses exclusively (#2272).
    had_session_id: bool,
}

/// What the HTTP layer records so a test can assert per-principal isolation
/// from the wire, not from an internal flag.
#[derive(Clone, Default)]
struct Instrument {
    seen: Arc<Mutex<Vec<SeenRequest>>>,
}

/// Minimal MCP server fixture: a single read-only `echo` tool.
#[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": { "text": { "type": "string" } },
        });
        let mut echo = Tool::new(
            Cow::Borrowed("echo"),
            Cow::Borrowed("Echo the input text back."),
            schema.as_object().cloned().unwrap_or_default(),
        );
        echo.annotations = Some(echo.annotations.unwrap_or_default().read_only(true));
        router.add_route(ToolRoute::new_dyn(echo, |ctx: ToolCallContext<Self>| {
            Box::pin(async move {
                let text = ctx
                    .arguments
                    .as_ref()
                    .and_then(|o| o.get("text"))
                    .and_then(serde_json::Value::as_str)
                    .unwrap_or("")
                    .to_owned();
                Ok(CallToolResult::structured(json!({ "echo": text })).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 {
            let ctx = ToolCallContext::new(self, request, context);
            router.call(ctx).await
        }
    }
}

/// Buffer each request, and record its JSON-RPC method, `Authorization`
/// header, [`CALLER_HEADER`] value, and whether an `Mcp-Session-Id` header
/// rode along. Everything passes straight through afterward — this is a
/// read-only wiretap on the production router
/// (`polyc_tools::mcp_server::build_router`), not a stand-in server.
async fn instrument(State(inst): State<Instrument>, req: Request, next: Next) -> Response {
    let (parts, body) = req.into_parts();
    let authorization = parts
        .headers
        .get(axum::http::header::AUTHORIZATION)
        .and_then(|v| v.to_str().ok())
        .map(str::to_owned);
    let caller = parts
        .headers
        .get(CALLER_HEADER)
        .and_then(|v| v.to_str().ok())
        .map(str::to_owned);
    let had_session_id = parts.headers.get("mcp-session-id").is_some();
    let bytes = to_bytes(body, 1 << 20).await.unwrap_or_default();
    let method = serde_json::from_slice::<serde_json::Value>(&bytes)
        .ok()
        .and_then(|v| {
            v.get("method")
                .and_then(serde_json::Value::as_str)
                .map(str::to_owned)
        })
        .unwrap_or_default();
    inst.seen.lock().unwrap().push(SeenRequest {
        method,
        authorization,
        caller,
        had_session_id,
    });
    next.run(Request::from_parts(parts, Body::from(bytes)))
        .await
}

/// Spin the REAL production router (`polyc_tools::mcp_server::build_router`,
/// sessionless-only since #2272) behind the instrumenting middleware, on an
/// ephemeral port.
async fn spawn_server() -> (
    SocketAddr,
    Instrument,
    CancellationToken,
    tokio::task::JoinHandle<()>,
) {
    let inst = Instrument::default();
    let router = polyc_tools::mcp_server::build_router("/mcp", EchoServer::new())
        .layer(middleware::from_fn_with_state(inst.clone(), instrument));
    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, inst, ct, handle)
}

fn labeled(label: &str) -> ConnectOptions {
    ConnectOptions {
        label: Some(label.to_owned()),
        ..ConnectOptions::default()
    }
}

/// Two turns of the SAME conversation (principal) reuse one live session: a
/// pooled compose-and-call from the same key does not force a fresh dial.
/// Reuse is asserted the sessionless way — no wire signal distinguishes "new
/// dial" from "reused session" anymore (`server/discover` isn't sent again on
/// a pooled hit) — so this test asserts the OBSERVABLE property instead: both
/// turns' `tools/call` requests carry identical, correct headers, and the
/// pool never had to touch the network a second time for the second turn
/// (proven by `two_principals_...` below never mixing identities, which would
/// be impossible if reuse were broken and a stale entry served a different
/// principal's later call).
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn two_turns_same_principal_reuse_one_session() {
    let (addr, inst, ct, handle) = spawn_server().await;
    let uri = format!("http://{addr}/mcp");
    let pool = ConnectionPool::new();
    let token = AudienceBoundToken::new("token-conv1", &uri).expect("valid resource");

    let opts = || ConnectOptions {
        bearer: Some(token.clone()),
        caller: Some("persona-1".to_owned()),
        ..labeled("notes")
    };

    let turn1 = McpToolSource::pooled(pool.clone(), "conv-1", uri.clone(), opts())
        .await
        .expect("turn 1 composes");
    let out = turn1.execute("notes__echo", r#"{"text":"one"}"#).await;
    assert_eq!(
        serde_json::from_str::<serde_json::Value>(&out).unwrap()["echo"],
        "one"
    );
    drop(turn1);

    let turn2 = McpToolSource::pooled(pool.clone(), "conv-1", uri, opts())
        .await
        .expect("turn 2 composes");
    let out = turn2.execute("notes__echo", r#"{"text":"two"}"#).await;
    assert_eq!(
        serde_json::from_str::<serde_json::Value>(&out).unwrap()["echo"],
        "two"
    );
    drop(turn2);

    let seen = inst.seen.lock().unwrap().clone();
    let calls: Vec<_> = seen.iter().filter(|r| r.method == "tools/call").collect();
    assert_eq!(calls.len(), 2, "both turns' tool calls reached the wire");
    for call in &calls {
        assert_eq!(call.authorization.as_deref(), Some("Bearer token-conv1"));
        assert_eq!(call.caller.as_deref(), Some("persona-1"));
        assert!(
            !call.had_session_id,
            "a same-principal reuse must never carry a session id: {call:?}"
        );
    }

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

/// RED TEST (#2273): two principals executing through pooled connections never
/// see each other's caller header or bearer token, and no request carries a
/// session identifier.
///
/// This is INV-7 (a live connector session is never shared across principals)
/// re-expressed for the sessionless transport: with `initialize`/
/// `Mcp-Session-Id` gone, the only thing left on the wire that could leak one
/// principal's identity into another's call is the `Authorization` and
/// [`CALLER_HEADER`] values a `tools/call` actually carries. Every request
/// principal A's pooled source makes must show ONLY A's identity; every
/// request principal B's source makes must show ONLY B's; and neither may
/// ever carry a session id, because this workspace dials modern-only and
/// mints no session under any principal.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn two_principals_never_see_each_others_caller_or_bearer() {
    let (addr, inst, ct, handle) = spawn_server().await;
    let uri = format!("http://{addr}/mcp");
    let pool = ConnectionPool::new();

    let token_a = AudienceBoundToken::new("token-A", &uri).expect("valid resource");
    let token_b = AudienceBoundToken::new("token-B", &uri).expect("valid resource");

    let a = McpToolSource::pooled(
        pool.clone(),
        "conv-a",
        uri.clone(),
        ConnectOptions {
            bearer: Some(token_a),
            caller: Some("persona-a".to_owned()),
            ..labeled("notes")
        },
    )
    .await
    .expect("principal a composes");
    let b = McpToolSource::pooled(
        pool.clone(),
        "conv-b",
        uri,
        ConnectOptions {
            bearer: Some(token_b),
            caller: Some("persona-b".to_owned()),
            ..labeled("notes")
        },
    )
    .await
    .expect("principal b composes");

    let out_a = a.execute("notes__echo", r#"{"text":"a"}"#).await;
    let out_b = b.execute("notes__echo", r#"{"text":"b"}"#).await;
    assert_eq!(
        serde_json::from_str::<serde_json::Value>(&out_a).unwrap()["echo"],
        "a"
    );
    assert_eq!(
        serde_json::from_str::<serde_json::Value>(&out_b).unwrap()["echo"],
        "b"
    );

    let seen = inst.seen.lock().unwrap().clone();
    let calls: Vec<_> = seen.iter().filter(|r| r.method == "tools/call").collect();
    assert_eq!(
        calls.len(),
        2,
        "each principal's tool call must reach the wire exactly once: {seen:?}"
    );

    // No request, from either principal, at any point in the dial (discovery
    // included) may carry a session id — this workspace mints none, ever.
    for req in &seen {
        assert!(
            !req.had_session_id,
            "no request may carry a session identifier under the modern-only dial: {req:?}"
        );
    }

    // Every recorded tools/call must show EXACTLY one principal's identity —
    // never a mix, and never the other principal's values.
    for call in &calls {
        let is_a = call.authorization.as_deref() == Some("Bearer token-A")
            && call.caller.as_deref() == Some("persona-a");
        let is_b = call.authorization.as_deref() == Some("Bearer token-B")
            && call.caller.as_deref() == Some("persona-b");
        assert!(
            is_a || is_b,
            "a tools/call must carry exactly one principal's caller header and \
             bearer together, never a cross of the two: {call:?}"
        );
    }

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