emperor_mcp/serve.rs
1//!
2//! Streamable-HTTP server bootstrap.
3//!
4//! [`serve`] wires together:
5//! * the rmcp `StreamableHttpService` mounted at `/mcp`,
6//! * the [`super::headers::forward_headers_middleware`] that captures
7//! per-user credential headers for forwarding,
8//! * the [`super::credentials::credential_rejection_middleware`] that turns a tool's
9//! "upstream rejected the credentials" into a `401`/`403` response,
10//! * a `/health` endpoint,
11//! * graceful shutdown on Ctrl-C / cancellation.
12//!
13
14use std::net::SocketAddr;
15
16use rmcp::handler::server::ServerHandler;
17use rmcp::transport::StreamableHttpServerConfig;
18use rmcp::transport::streamable_http_server::session::local::LocalSessionManager;
19use rmcp::transport::streamable_http_server::tower::StreamableHttpService;
20use tokio_util::sync::CancellationToken;
21
22use crate::credentials::credential_rejection_middleware;
23use crate::headers::{ForwardConfig, forward_headers_middleware};
24
25const LEGACY_SESSION_MODE: bool = false;
26
27/// Options for [`serve`].
28pub struct ServeOptions {
29 /// Address to bind, e.g. `0.0.0.0:8000`.
30 pub bind: SocketAddr,
31 /// Which inbound headers to forward to the upstream service per request.
32 pub forward: ForwardConfig,
33 /// In-process `Host` validation (DNS-rebinding protection). Empty = off, for deployments
34 /// behind a validating ingress (Emperor Profile E5).
35 pub allowed_hosts: Vec<String>,
36 /// In-process `Origin` validation, entries with a scheme (`https://app.example`). Empty = off.
37 pub allowed_origins: Vec<String>,
38}
39
40impl ServeOptions {
41 /// Build options from a bind address and the common "forward only
42 /// Authorization" header policy.
43 pub fn new(bind: SocketAddr) -> Self {
44 Self {
45 bind,
46 forward: ForwardConfig::authorization_only(),
47 allowed_hosts: Vec::new(),
48 allowed_origins: Vec::new(),
49 }
50 }
51
52 /// Override the forwarded-header allow-list (e.g. for MCPs that use
53 /// `X-*` credential headers instead of / in addition to Authorization).
54 pub fn with_forward(mut self, forward: ForwardConfig) -> Self {
55 self.forward = forward;
56 self
57 }
58
59 /// Validate the `Host` header in-process (e.g. `["localhost", "mcp.example.com"]`), for
60 /// deployments without a validating ingress in front.
61 pub fn with_allowed_hosts<I, S>(mut self, hosts: I) -> Self
62 where
63 I: IntoIterator<Item = S>,
64 S: Into<String>,
65 {
66 self.allowed_hosts = hosts.into_iter().map(Into::into).collect();
67 self
68 }
69
70 /// Validate the `Origin` header in-process (entries include the scheme). Requests without
71 /// an `Origin` header (non-browser clients) still pass.
72 pub fn with_allowed_origins<I, S>(mut self, origins: I) -> Self
73 where
74 I: IntoIterator<Item = S>,
75 S: Into<String>,
76 {
77 self.allowed_origins = origins.into_iter().map(Into::into).collect();
78 self
79 }
80}
81
82/// Run a first-party MCP server over streamable-HTTP until shutdown.
83///
84/// `make_handler` is called once per session to build a fresh handler
85/// instance (rmcp's session model). The handler must implement
86/// [`ServerHandler`] (typically via `#[tool_router]` + `#[tool]`).
87///
88/// Per-user credential headers named in `options.forward` are captured by
89/// middleware and made available to tool handlers via
90/// [`crate::headers::ForwardedHeaders::from_request_context`]. A tool that gets `401`/`403`
91/// from the upstream returns [`crate::credentials::credentials_rejected`], and the MCP response
92/// carries that status.
93pub async fn serve<H, F>(options: ServeOptions, make_handler: F) -> anyhow::Result<()>
94where
95 H: ServerHandler + Send + 'static,
96 F: Fn() -> H + Send + Sync + 'static,
97{
98 let ct = CancellationToken::new();
99
100 let mut config = StreamableHttpServerConfig::default()
101 // Stateless: no `initialize`/`Mcp-Session-Id` handshake. Each tool call
102 // is a self-contained POST so per-request credential forwarding works
103 // and tool discovery (unauthenticated `tools/list`) still succeeds.
104 .with_legacy_session_mode(LEGACY_SESSION_MODE)
105 // Return plain JSON (no SSE framing) for simple request/response tools,
106 // which is what the gateway's sessionless path expects.
107 .with_json_response(!LEGACY_SESSION_MODE)
108 .with_cancellation_token(ct.child_token());
109 // Host/Origin validation: off by default (a validating ingress does it, Emperor Profile E5),
110 // in-process when the options name hosts or origins.
111 config = if options.allowed_hosts.is_empty() {
112 config.disable_allowed_hosts()
113 } else {
114 config.with_allowed_hosts(options.allowed_hosts.clone())
115 };
116 config = if options.allowed_origins.is_empty() {
117 config.disable_allowed_origins()
118 } else {
119 config.with_allowed_origins(options.allowed_origins.clone())
120 };
121
122 let service = StreamableHttpService::new(
123 move || Ok(make_handler()),
124 LocalSessionManager::default().into(),
125 config,
126 );
127
128 let app = axum::Router::new()
129 .route("/health", axum::routing::get(health))
130 .nest_service("/mcp", service)
131 .layer(axum::middleware::from_fn(credential_rejection_middleware))
132 .layer(axum::middleware::from_fn_with_state(
133 options.forward.clone(),
134 forward_headers_middleware,
135 ));
136
137 let listener = tokio::net::TcpListener::bind(options.bind).await?;
138 tracing::info!(bind = %options.bind, "MCP streamable-HTTP server listening on /mcp");
139
140 let shutdown = {
141 let ct = ct.clone();
142 async move {
143 let _ = tokio::signal::ctrl_c().await;
144 tracing::info!("shutdown signal received");
145 ct.cancel();
146 }
147 };
148
149 axum::serve(listener, app)
150 .with_graceful_shutdown(shutdown)
151 .await?;
152 Ok(())
153}
154
155/// Minimal liveness endpoint.
156async fn health() -> &'static str {
157 "OK"
158}