Skip to main content

codoseo_web/routes/
mcp.rs

1//! `/mcp`: the cloud MCP server (streamable HTTP, stateless, JSON responses) in front of
2//! [`CloudMcp`]. Authentication is the API's: `Authorization: Bearer <key>` and nothing else
3//! (a session cookie never counts). One endpoint serves both tiers: a request with a key gets the
4//! keyed tools, a request without one gets the no-key tools on the cloud and a 401 when
5//! self-hosted. A malformed, unknown or revoked key is a 401 JSON error, never "no key".
6//!
7//! A small middleware resolves the caller once per HTTP request and puts it in the request's
8//! extensions; rmcp copies the request parts into every tool context, where the handler reads
9//! it. `initialize`, `tools/list` and pings are free; each keyed `tools/call` costs one call
10//! (see [`AgentBackend`]). The route is exempt from the `Origin` check (no cookie reaches it),
11//! and in the cloud rmcp itself refuses a `Host` that isn't this app's (DNS rebinding).
12
13use std::net::SocketAddr;
14use std::sync::Arc;
15
16use axum::Router;
17use axum::extract::{ConnectInfo, Request, State};
18use axum::http::{HeaderValue, header};
19use axum::middleware::{self, Next};
20use axum::response::{IntoResponse, Response};
21use codoseo_mcp::cloud::{AnonBackend, Caller, CloudMcp};
22use rmcp::transport::streamable_http_server::session::never::NeverSessionManager;
23use rmcp::transport::streamable_http_server::{StreamableHttpServerConfig, StreamableHttpService};
24
25use crate::agent::anon::AnonCaller;
26use crate::agent::auth::{self, ApiCaller};
27use crate::agent::mcp::AgentBackend;
28use crate::config::{Config, Mode};
29use crate::metrics::{self, Surface, Tier};
30use crate::state::AppState;
31
32/// Where the MCP server lives; the `Origin` check exempts exactly this path.
33pub const PATH: &str = "/mcp";
34
35/// The caller as the middleware resolves it for [`AgentBackend`].
36type McpCaller = Caller<ApiCaller, AnonCaller>;
37
38/// The most a request body may hold: a tool call is a few hundred bytes.
39const MAX_BODY_BYTES: usize = 64 * 1024;
40
41pub fn routes(state: &AppState) -> Router<AppState> {
42    router_for(state, CloudMcp::new(AgentBackend::new(state.clone())))
43}
44
45/// `/mcp` in front of `handler` (a test serves one with a faulty backend).
46pub fn router_for<B>(state: &AppState, handler: CloudMcp<B>) -> Router<AppState>
47where
48    B: AnonBackend<Keyed = ApiCaller, Anon = AnonCaller> + Clone,
49{
50    let config = StreamableHttpServerConfig::default()
51        // No sessions: every request stands alone, so any instance can answer it.
52        .with_legacy_session_mode(false)
53        .with_json_response(true)
54        .with_max_request_body_bytes(MAX_BODY_BYTES)
55        // Open calls end when the server starts shutting down.
56        .with_cancellation_token(state.shutdown.clone());
57    let config = match state.config.mode {
58        Mode::Cloud => config.with_allowed_hosts(allowed_hosts(&state.config)),
59        // The Host check guards a local server against DNS rebinding from a browser, which
60        // needs a cookie or ambient network position to matter. This endpoint takes a Bearer
61        // key only, and a self-hosted address is whatever the operator makes it (a LAN name, a
62        // proxy that rewrites Host), so there is nothing to allow-list.
63        Mode::SelfHost => config.disable_allowed_hosts(),
64    };
65    let service = StreamableHttpService::new(
66        move || Ok(handler.clone()),
67        Arc::new(NeverSessionManager::default()),
68        config,
69    );
70    Router::new()
71        .route_service(PATH, service)
72        .layer(middleware::from_fn_with_state(
73            state.clone(),
74            resolve_caller,
75        ))
76}
77
78/// The cloud's `Host` values rmcp accepts: this app's own (with its port when `BASE_URL` has one) and
79/// the loopback names, which a local client or a test server uses.
80fn allowed_hosts(config: &Config) -> Vec<String> {
81    let mut hosts: Vec<String> = ["localhost", "127.0.0.1", "::1"].map(str::to_owned).into();
82    if let Some(host) = config.base_url.host_str() {
83        hosts.push(match config.base_url.port() {
84            Some(port) => format!("{host}:{port}"),
85            None => host.to_owned(),
86        });
87    }
88    hosts
89}
90
91/// Works out who is calling before rmcp sees the request. A bad key, or no key on a self-hosted
92/// server, ends here as the API's JSON 401.
93async fn resolve_caller(State(state): State<AppState>, mut req: Request, next: Next) -> Response {
94    let headers = req.headers();
95    let caller: McpCaller = match auth::bearer_key(headers) {
96        Ok(None) if state.config.mode == Mode::Cloud => {
97            let peer = req
98                .extensions()
99                .get::<ConnectInfo<SocketAddr>>()
100                .map(|c| c.0.ip());
101            Caller::Anon(AnonCaller::from_request(&state, headers, peer))
102        }
103        // A key, a bad key, or no key where one is required: `authenticate` says which.
104        _ => match auth::authenticate(&state, headers).await {
105            Ok(keyed) => Caller::Keyed(keyed),
106            Err(e) => {
107                metrics::api_request(Surface::Mcp, Tier::Key, e.code());
108                return e.into_response();
109            }
110        },
111    };
112    req.extensions_mut().insert(caller);
113    // The request's parts travel into the tool context; keep the key out of any debug print.
114    if let Some(value) = req.headers_mut().get_mut(header::AUTHORIZATION) {
115        value.set_sensitive(true);
116    }
117    let mut res = next.run(req).await;
118    res.headers_mut()
119        .insert(header::CACHE_CONTROL, HeaderValue::from_static("no-store"));
120    res
121}