polyc-tools 2026.8.3

//! Shared MCP streamable-HTTP serving helper.
//!
//! [`serve`] runs any [`ServerHandler`] over the MCP streamable-HTTP transport
//! until Ctrl-C / SIGTERM. The standalone MCP binaries in this workspace (e.g.
//! the conversation server) each build their own handler and hand it here, so
//! the transport wiring lives in one place. It also mounts the standard
//! `grpc.health.v1.Health` Connect health service alongside the MCP mount, so
//! a connector built on `serve()` gets a real `grpc:` readiness/liveness probe
//! for free instead of falling back to a bare `tcpSocket` check.
//!
//! Sessionless only (`#2272`): every server built on [`build_router`] serves
//! SEP-2567 stateless — `server/discover` plus self-contained per-request
//! `_meta`, never `initialize`/`initialized` and never `Mcp-Session-Id`. There
//! is no legacy session path left to opt into; a legacy client's `initialize`
//! is answered as an ordinary one-shot stateless request (no session id comes
//! back), which a legacy client cannot use, by design — this workspace has no
//! legacy clients to keep working (greenfield, modern-only both directions).

use std::{net::SocketAddr, sync::Arc};

use connectrpc::Router;
use rmcp::{
    ServerHandler,
    transport::streamable_http_server::{
        StreamableHttpServerConfig, StreamableHttpService, session::never::NeverSessionManager,
    },
};
use tokio_util::sync::CancellationToken;

/// Build the axum router `serve` binds.
///
/// Nests the MCP streamable-HTTP service under `path`, plus the standard
/// `grpc.health.v1.Health` Connect health service as the fallback. Split out
/// so tests can drive it without binding a socket — `pub` so the crate's own
/// integration tests (`tests/mcp_roundtrip.rs`) can dial the exact production
/// router rather than a hand-rolled stand-in.
pub fn build_router<H>(path: &str, handler: H) -> axum::Router
where
    H: ServerHandler + Clone + Send + Sync + 'static,
{
    let mut config = StreamableHttpServerConfig::default();
    // Discovery is the only dial mode: no legacy `initialize`/`initialized`
    // handshake, no `Mcp-Session-Id`, no session to mint. `NeverSessionManager`
    // below backs this up structurally — there is no session store left to
    // mint from even if a request tried to negotiate one.
    config.legacy_session_mode = false;
    config.sse_keep_alive = Some(std::time::Duration::from_secs(15));
    // These servers run behind a trusted in-cluster edge / k8s Service that owns
    // routing and the public boundary, so rmcp's browser-oriented DNS-rebinding
    // Host/Origin guard doesn't apply at this layer — turn it off. Set
    // `with_allowed_hosts([..])` instead if a server is fronted directly.
    let config = config.disable_allowed_hosts().disable_allowed_origins();
    let service = StreamableHttpService::new(
        move || Ok(handler.clone()),
        Arc::new(NeverSessionManager::default()),
        config,
    );

    // Standard gRPC health service (`grpc.health.v1.Health`) on the SAME port
    // that serves MCP, mirroring the control-plane/harness wiring
    // (`connectrpc_health::install_static`). There is no named Connect
    // service behind this connector — only the MCP streamable-HTTP mount —
    // so an empty service-name list is correct: it seeds just the
    // whole-process `""` health entry at `Serving`, which is what a
    // Kubernetes `grpc:` readiness/liveness probe checks by default. The
    // returned checker is dropped; the registered service keeps its own
    // `Arc`, so health stays `Serving` for the life of the router.
    let (health, _checker) = connectrpc_health::install_static(Router::new(), [] as [&str; 0]);
    axum::Router::new()
        .nest_service(path, service)
        .fallback_service(health.into_axum_service())
}

/// Serve an MCP [`ServerHandler`] over the streamable-HTTP transport on `bind`
/// under `path`, until Ctrl-C / SIGTERM. Shared by the standalone MCP binaries
/// so each only supplies a pre-built handler.
///
/// The handler is built once by the caller; the transport's per-session
/// service factory just clones it. Cloning is cheap because the handlers in
/// this crate are `Arc`-shaped, so the `ToolRouter` + schema JSON are built
/// once rather than per client connection.
///
/// # Errors
/// Returns an error if the listener cannot be bound or the signal wait fails.
pub async fn serve<H>(bind: SocketAddr, path: &str, handler: H) -> anyhow::Result<()>
where
    H: ServerHandler + Clone + Send + Sync + 'static,
{
    let router = build_router(path, handler);
    let listener = tokio::net::TcpListener::bind(bind).await?;
    let local = listener.local_addr()?;
    tracing::info!(addr = %local, path, "MCP server listening");
    println!("MCP server listening on http://{local}{path}");

    let ct = CancellationToken::new();
    let server_ct = ct.clone();
    let server = tokio::spawn(async move {
        let _ = axum::serve(listener, router)
            .with_graceful_shutdown(async move { server_ct.cancelled_owned().await })
            .await;
    });

    tokio::signal::ctrl_c().await?;
    tracing::info!("shutting down");
    ct.cancel();
    let _ = server.await;
    Ok(())
}

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

    use axum::body::Body;
    use axum::http::{Request, StatusCode, header};
    use rmcp::model::{Implementation, InitializeResult, ServerCapabilities, ServerInfo};
    use tower::ServiceExt; // for `oneshot`

    use super::*;

    /// A handler with no tools — the health mount is what's under test, not
    /// the MCP surface.
    #[derive(Clone)]
    struct NoopServer;

    impl ServerHandler for NoopServer {
        fn get_info(&self) -> ServerInfo {
            InitializeResult::new(ServerCapabilities::builder().build())
                .with_server_info(Implementation::new("noop", "0.0.0"))
        }
    }

    #[tokio::test]
    async fn health_check_reports_serving() {
        let router = build_router("/mcp", NoopServer);
        let resp = router
            .oneshot(
                Request::post("/grpc.health.v1.Health/Check")
                    .header(header::CONTENT_TYPE, "application/json")
                    .body(Body::from("{}"))
                    .unwrap(),
            )
            .await
            .unwrap();
        assert_eq!(resp.status(), StatusCode::OK);
        let body = axum::body::to_bytes(resp.into_body(), usize::MAX)
            .await
            .unwrap();
        let value: serde_json::Value = serde_json::from_slice(&body).unwrap();
        assert_eq!(value["status"], "SERVING");
    }
}