Skip to main content

llm_kernel/mcp/
http.rs

1//! HTTP/SSE remote transport for MCP.
2//!
3//! Exposes an [`McpServer`] over HTTP: a JSON-RPC endpoint (`POST /mcp`) and an
4//! SSE endpoint (`POST /mcp/sse`) that streams the response as a server-sent
5//! event. Both reuse the server's `Authorization` (Bearer) check, so a server
6//! secured for stdio is secured identically over HTTP.
7//!
8//! The transport holds the server behind an `Arc` (shared across request
9//! tasks) and dispatches `tools/call` via [`McpServer::call_tool_async`], so
10//! async handlers work transparently over HTTP.
11//!
12//! Requires the `mcp-http` feature (axum + tokio).
13
14use std::convert::Infallible;
15use std::net::SocketAddr;
16use std::sync::Arc;
17
18use axum::Json;
19use axum::extract::State;
20use axum::http::{HeaderMap, StatusCode};
21use axum::response::IntoResponse;
22use axum::response::sse::{Event, KeepAlive, Sse};
23use axum::routing::post;
24use serde_json::Value;
25use tokio_stream::wrappers::UnboundedReceiverStream;
26
27use crate::mcp::McpServer;
28
29/// Shared MCP server state for the HTTP transport.
30#[derive(Clone)]
31pub struct HttpTransport {
32    server: Arc<McpServer>,
33}
34
35impl HttpTransport {
36    /// Wrap a shared MCP server for HTTP serving.
37    pub fn new(server: Arc<McpServer>) -> Self {
38        Self { server }
39    }
40
41    /// Build the axum router with JSON-RPC and SSE routes.
42    pub fn router(&self) -> axum::Router {
43        axum::Router::new()
44            .route("/mcp", post(rpc_handler))
45            .route("/mcp/sse", post(sse_handler))
46            .with_state(self.clone())
47    }
48}
49
50/// Run the MCP HTTP transport on `addr` until the server is stopped.
51pub async fn serve(server: Arc<McpServer>, addr: SocketAddr) -> std::io::Result<()> {
52    let transport = HttpTransport::new(server);
53    let listener = tokio::net::TcpListener::bind(addr).await?;
54    axum::serve(listener, transport.router()).await?;
55    Ok(())
56}
57
58/// JSON-RPC code for "method not found".
59const ERR_METHOD_NOT_FOUND: i32 = -32601;
60/// JSON-RPC code for invalid params (unknown tool / prompt / resource).
61const ERR_INVALID_PARAMS: i32 = -32602;
62/// JSON-RPC code for a tool-execution / internal error.
63const ERR_INTERNAL: i32 = -32603;
64/// JSON-RPC code for unauthorized access.
65const ERR_UNAUTHORIZED: i32 = -32001;
66
67/// Dispatch a single JSON-RPC request against the server (async path).
68///
69/// `tools/call` is awaited via [`McpServer::call_tool_async`]; `initialize`,
70/// `ping`, `tools/list`, `resources/list`, `resources/templates/list`,
71/// `prompts/list`, `prompts/get`, and `resources/read` are handled
72/// synchronously. Notifications (no `id`) return `None`.
73async fn dispatch_async(server: &McpServer, req: &Value) -> Option<Value> {
74    // Notifications (no id) get no response.
75    let id = req.get("id")?.clone();
76    let method = req.get("method").and_then(|v| v.as_str()).unwrap_or("");
77
78    let result: Result<Value, (i32, String)> = match method {
79        "initialize" => {
80            let requested = req
81                .pointer("/params/protocolVersion")
82                .and_then(|v| v.as_str());
83            Ok(server.initialize_response(requested))
84        }
85        "ping" => Ok(serde_json::json!({})),
86        "tools/list" => Ok(serde_json::json!({ "tools": server.tools() })),
87        "resources/list" => Ok(serde_json::json!({ "resources": server.resources() })),
88        "resources/templates/list" => Ok(serde_json::json!({ "resourceTemplates": [] })),
89        "prompts/list" => Ok(serde_json::json!({ "prompts": server.prompts() })),
90        "prompts/get" => {
91            let name = req
92                .pointer("/params/name")
93                .and_then(|v| v.as_str())
94                .unwrap_or("");
95            let args = req
96                .pointer("/params/arguments")
97                .cloned()
98                .unwrap_or(serde_json::json!({}));
99            server
100                .get_prompt(name, args)
101                .map_err(|e| (ERR_INVALID_PARAMS, e.to_string()))
102        }
103        "resources/read" => {
104            let uri = req
105                .pointer("/params/uri")
106                .and_then(|v| v.as_str())
107                .unwrap_or("");
108            server
109                .read_resource(uri, serde_json::json!({}))
110                .map(|content| {
111                    serde_json::json!({
112                        "contents": [{ "uri": uri, "text": content.to_string() }]
113                    })
114                })
115                .map_err(|e| (ERR_INTERNAL, e.to_string()))
116        }
117        "tools/call" => {
118            let name = req
119                .pointer("/params/name")
120                .and_then(|v| v.as_str())
121                .unwrap_or("");
122            let params = req
123                .pointer("/params/arguments")
124                .cloned()
125                .unwrap_or(serde_json::json!(null));
126            if !server.has_tool(name) {
127                Err((ERR_INVALID_PARAMS, format!("Unknown tool: {name}")))
128            } else if let Err(e) = server.validate_tool_args(name, &params) {
129                Err((ERR_INVALID_PARAMS, e))
130            } else {
131                // Execution failures are reported in-band with isError: true.
132                match server.call_tool_async(name, params).await {
133                    Ok(r) => Ok(serde_json::json!({
134                        "content": [{ "type": "text", "text": r.to_string() }],
135                        "isError": false
136                    })),
137                    Err(e) => Ok(serde_json::json!({
138                        "content": [{ "type": "text", "text": e.to_string() }],
139                        "isError": true
140                    })),
141                }
142            }
143        }
144        _ => Err((ERR_METHOD_NOT_FOUND, format!("Method not found: {method}"))),
145    };
146
147    Some(match result {
148        Ok(value) => serde_json::json!({ "jsonrpc": "2.0", "id": id, "result": value }),
149        Err((code, message)) => serde_json::json!({
150            "jsonrpc": "2.0", "id": id,
151            "error": { "code": code, "message": message }
152        }),
153    })
154}
155
156/// Extract and validate the `Authorization` header. Returns `true` if the
157/// request may proceed.
158fn authorized(server: &McpServer, headers: &HeaderMap) -> bool {
159    let auth = headers
160        .get("authorization")
161        .and_then(|v| v.to_str().ok())
162        .unwrap_or("");
163    server.check_auth(auth)
164}
165
166/// Reject cross-origin browser requests (MCP spec: servers MUST validate
167/// `Origin` to prevent DNS rebinding). A page on any website can POST to a
168/// loopback MCP server; without this, that page executes tools.
169///
170/// Non-browser clients send no `Origin` and are unaffected.
171fn origin_allowed(headers: &HeaderMap) -> bool {
172    let Some(origin) = headers.get("origin").and_then(|v| v.to_str().ok()) else {
173        return true; // no Origin — not a browser-initiated request
174    };
175    if origin == "null" {
176        return false;
177    }
178    // Only loopback origins may drive a local MCP server. Parse the host
179    // bracket-aware: an IPv6 origin is `http://[::1]:3000`, where a naive
180    // `split(':').next()` yields "[" and rejects a legitimate loopback.
181    origin
182        .split_once("://")
183        .map(|(_, host_port)| {
184            if let Some(rest) = host_port.strip_prefix('[') {
185                rest.split(']').next().unwrap_or("") // "::1" (no brackets)
186            } else {
187                host_port.split(':').next().unwrap_or("")
188            }
189        })
190        .is_some_and(|host| host == "localhost" || host == "127.0.0.1" || host == "::1")
191}
192
193fn forbidden_response(id: Option<Value>) -> Json<Value> {
194    Json(serde_json::json!({
195        "jsonrpc": "2.0",
196        "id": id,
197        "error": { "code": ERR_UNAUTHORIZED, "message": "Forbidden origin" }
198    }))
199}
200
201fn unauthorized_response(id: Option<Value>) -> Json<Value> {
202    Json(serde_json::json!({
203        "jsonrpc": "2.0",
204        "id": id,
205        "error": { "code": ERR_UNAUTHORIZED, "message": "Unauthorized" }
206    }))
207}
208
209/// Dispatch a single request or a JSON-RPC batch (array). Returns `None` only
210/// when nothing needs answering (all notifications).
211async fn dispatch_any(server: &McpServer, req: &Value) -> Option<Value> {
212    let Some(batch) = req.as_array() else {
213        return dispatch_async(server, req).await;
214    };
215    let mut out = Vec::with_capacity(batch.len());
216    for item in batch {
217        if let Some(resp) = dispatch_async(server, item).await {
218            out.push(resp);
219        }
220    }
221    if out.is_empty() {
222        None
223    } else {
224        Some(Value::Array(out))
225    }
226}
227
228async fn rpc_handler(
229    State(state): State<HttpTransport>,
230    headers: HeaderMap,
231    Json(req): Json<Value>,
232) -> impl IntoResponse {
233    let id = req.get("id").cloned();
234    if !origin_allowed(&headers) {
235        return (StatusCode::FORBIDDEN, forbidden_response(id));
236    }
237    if !authorized(&state.server, &headers) {
238        return (StatusCode::UNAUTHORIZED, unauthorized_response(id));
239    }
240    match dispatch_any(&state.server, &req).await {
241        Some(resp) => (StatusCode::OK, Json(resp)),
242        // Notification — acknowledge with 204 No Content.
243        None => (StatusCode::NO_CONTENT, Json(serde_json::Value::Null)),
244    }
245}
246
247async fn sse_handler(
248    State(state): State<HttpTransport>,
249    headers: HeaderMap,
250    Json(req): Json<Value>,
251) -> impl IntoResponse {
252    let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
253    let server = state.server.clone();
254
255    // Produce the response for this request, then stream it as one SSE event.
256    tokio::spawn(async move {
257        let event = if !origin_allowed(&headers) {
258            Event::default().event("error").data(
259                serde_json::to_string(&forbidden_response(req.get("id").cloned()).0)
260                    .unwrap_or_default(),
261            )
262        } else if !authorized(&server, &headers) {
263            Event::default().event("error").data(
264                serde_json::to_string(&unauthorized_response(req.get("id").cloned()).0)
265                    .unwrap_or_default(),
266            )
267        } else if let Some(resp) = dispatch_any(&server, &req).await {
268            let data = serde_json::to_string(&resp).unwrap_or_default();
269            Event::default().event("message").data(data)
270        } else {
271            // Notification — no response event.
272            Event::default().event("noop")
273        };
274        let _ = tx.send(Ok::<_, Infallible>(event));
275    });
276
277    Sse::new(UnboundedReceiverStream::new(rx)).keep_alive(KeepAlive::default())
278}
279
280#[cfg(test)]
281mod tests {
282    use super::*;
283    use crate::mcp::schema::{ResourceDescription, ToolDescription};
284
285    fn server_with_echo() -> McpServer {
286        let mut server = McpServer::new("http-test", "1.0.0");
287        server.register_tool(ToolDescription {
288            name: "echo".into(),
289            description: "Echo".into(),
290            input_schema: serde_json::json!({"type": "object"}),
291        });
292        server.set_async_handler("echo", |params| async move { Ok(params) });
293        server
294    }
295
296    #[tokio::test]
297    async fn dispatch_initialize() {
298        let server = server_with_echo();
299        let req = serde_json::json!({"jsonrpc":"2.0","id":1,"method":"initialize","params":{}});
300        let resp = dispatch_async(&server, &req).await.unwrap();
301        assert_eq!(resp["result"]["serverInfo"]["name"], "http-test");
302    }
303
304    #[tokio::test]
305    async fn dispatch_tools_call_async() {
306        let server = server_with_echo();
307        let req = serde_json::json!({
308            "jsonrpc": "2.0", "id": 2, "method": "tools/call",
309            "params": { "name": "echo", "arguments": { "msg": "hello" } }
310        });
311        let resp = dispatch_async(&server, &req).await.unwrap();
312        let text = resp["result"]["content"][0]["text"].as_str().unwrap();
313        assert!(text.contains("hello"));
314    }
315
316    #[tokio::test]
317    async fn dispatch_unknown_method() {
318        let server = server_with_echo();
319        let req = serde_json::json!({"jsonrpc":"2.0","id":3,"method":"nope"});
320        let resp = dispatch_async(&server, &req).await.unwrap();
321        assert_eq!(resp["error"]["code"], ERR_METHOD_NOT_FOUND);
322    }
323
324    /// AC2: HTTP dispatch also serves `resources/read`, not just tools.
325    #[tokio::test]
326    async fn dispatch_resources_read() {
327        let mut server = McpServer::new("http-test", "1.0.0");
328        server.register_resource(ResourceDescription {
329            uri: "docs://x".into(),
330            name: "X".into(),
331            description: None,
332            mime_type: None,
333        });
334        server.set_resource_handler("docs://x", |_| Ok(serde_json::json!("# body")));
335        let req = serde_json::json!({
336            "jsonrpc": "2.0", "id": 4, "method": "resources/read",
337            "params": { "uri": "docs://x" }
338        });
339        let resp = dispatch_async(&server, &req).await.unwrap();
340        let text = resp["result"]["contents"][0]["text"].as_str().unwrap();
341        assert!(text.contains("body"));
342    }
343
344    #[test]
345    fn origin_validation_blocks_cross_site_browsers() {
346        let mut h = HeaderMap::new();
347        assert!(origin_allowed(&h), "no Origin (non-browser client) passes");
348        h.insert("origin", "http://localhost:3000".parse().unwrap());
349        assert!(origin_allowed(&h));
350        h.insert("origin", "http://127.0.0.1:8080".parse().unwrap());
351        assert!(origin_allowed(&h));
352        h.insert("origin", "https://evil.example.com".parse().unwrap());
353        assert!(!origin_allowed(&h), "DNS-rebinding origin must be rejected");
354        h.insert("origin", "null".parse().unwrap());
355        assert!(!origin_allowed(&h));
356        // Suffix trickery must not pass.
357        h.insert("origin", "https://localhost.evil.com".parse().unwrap());
358        assert!(!origin_allowed(&h));
359        // IPv6 loopback — a naive `split(':')` would see "[" and reject it.
360        h.insert("origin", "http://[::1]:3000".parse().unwrap());
361        assert!(origin_allowed(&h), "IPv6 loopback must pass: {h:?}");
362        h.insert("origin", "http://[::1]".parse().unwrap());
363        assert!(origin_allowed(&h), "IPv6 loopback (no port) must pass");
364        h.insert("origin", "http://[fe80::1]:3000".parse().unwrap());
365        assert!(!origin_allowed(&h), "non-loopback IPv6 must be rejected");
366    }
367
368    #[tokio::test]
369    async fn batch_requests_get_a_batch_response() {
370        let server = server_with_echo();
371        let batch = serde_json::json!([
372            {"jsonrpc":"2.0","id":1,"method":"ping"},
373            {"jsonrpc":"2.0","id":2,"method":"tools/call",
374             "params":{"name":"echo","arguments":{"v":1}}}
375        ]);
376        let resp = dispatch_any(&server, &batch)
377            .await
378            .expect("batch must not be dropped");
379        let arr = resp.as_array().expect("array response");
380        assert_eq!(arr.len(), 2);
381        assert_eq!(arr[0]["id"], 1);
382        assert_eq!(arr[1]["result"]["isError"], false);
383    }
384
385    /// AC2: a full HTTP round-trip — bind an ephemeral port, POST a tools/call,
386    /// and read the JSON-RPC response off the wire.
387    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
388    async fn http_round_trip_calls_tool() {
389        let server = Arc::new(server_with_echo());
390        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
391        let addr = listener.local_addr().unwrap();
392        // Hand the listener to axum in a background task.
393        let transport = HttpTransport::new(server);
394        tokio::spawn(async move {
395            let _ = axum::serve(listener, transport.router()).await;
396        });
397
398        let body = serde_json::to_string(&serde_json::json!({
399            "jsonrpc": "2.0", "id": 9, "method": "tools/call",
400            "params": { "name": "echo", "arguments": { "v": 42 } }
401        }))
402        .unwrap();
403        let req = format!(
404            "POST /mcp HTTP/1.1\r\nHost: localhost\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
405            body.len(),
406            body
407        );
408
409        let mut stream = tokio::net::TcpStream::connect(addr).await.unwrap();
410        use tokio::io::{AsyncReadExt, AsyncWriteExt};
411        stream.write_all(req.as_bytes()).await.unwrap();
412        let mut buf = Vec::new();
413        stream.read_to_end(&mut buf).await.unwrap();
414        let response = String::from_utf8_lossy(&buf);
415        assert!(response.contains("200 OK"), "response: {response}");
416        // The tool result is JSON-encoded inside the `text` field, so its quotes
417        // are escaped on the wire — assert on the unescaped value + content shape.
418        assert!(response.contains("\"content\""), "response: {response}");
419        assert!(response.contains("\\\"v\\\":42"), "response: {response}");
420    }
421}