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 {
129                // Execution failures are reported in-band with isError: true.
130                match server.call_tool_async(name, params).await {
131                    Ok(r) => Ok(serde_json::json!({
132                        "content": [{ "type": "text", "text": r.to_string() }],
133                        "isError": false
134                    })),
135                    Err(e) => Ok(serde_json::json!({
136                        "content": [{ "type": "text", "text": e.to_string() }],
137                        "isError": true
138                    })),
139                }
140            }
141        }
142        _ => Err((ERR_METHOD_NOT_FOUND, format!("Method not found: {method}"))),
143    };
144
145    Some(match result {
146        Ok(value) => serde_json::json!({ "jsonrpc": "2.0", "id": id, "result": value }),
147        Err((code, message)) => serde_json::json!({
148            "jsonrpc": "2.0", "id": id,
149            "error": { "code": code, "message": message }
150        }),
151    })
152}
153
154/// Extract and validate the `Authorization` header. Returns `true` if the
155/// request may proceed.
156fn authorized(server: &McpServer, headers: &HeaderMap) -> bool {
157    let auth = headers
158        .get("authorization")
159        .and_then(|v| v.to_str().ok())
160        .unwrap_or("");
161    server.check_auth(auth)
162}
163
164fn unauthorized_response(id: Option<Value>) -> Json<Value> {
165    Json(serde_json::json!({
166        "jsonrpc": "2.0",
167        "id": id,
168        "error": { "code": ERR_UNAUTHORIZED, "message": "Unauthorized" }
169    }))
170}
171
172async fn rpc_handler(
173    State(state): State<HttpTransport>,
174    headers: HeaderMap,
175    Json(req): Json<Value>,
176) -> impl IntoResponse {
177    let id = req.get("id").cloned();
178    if !authorized(&state.server, &headers) {
179        return (StatusCode::UNAUTHORIZED, unauthorized_response(id));
180    }
181    match dispatch_async(&state.server, &req).await {
182        Some(resp) => (StatusCode::OK, Json(resp)),
183        // Notification — acknowledge with 204 No Content.
184        None => (StatusCode::NO_CONTENT, Json(serde_json::Value::Null)),
185    }
186}
187
188async fn sse_handler(
189    State(state): State<HttpTransport>,
190    headers: HeaderMap,
191    Json(req): Json<Value>,
192) -> impl IntoResponse {
193    let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
194    let server = state.server.clone();
195
196    // Produce the response for this request, then stream it as one SSE event.
197    tokio::spawn(async move {
198        let event = if !authorized(&server, &headers) {
199            Event::default().event("error").data(
200                serde_json::to_string(&unauthorized_response(req.get("id").cloned()).0)
201                    .unwrap_or_default(),
202            )
203        } else if let Some(resp) = dispatch_async(&server, &req).await {
204            let data = serde_json::to_string(&resp).unwrap_or_default();
205            Event::default().event("message").data(data)
206        } else {
207            // Notification — no response event.
208            Event::default().event("noop")
209        };
210        let _ = tx.send(Ok::<_, Infallible>(event));
211    });
212
213    Sse::new(UnboundedReceiverStream::new(rx)).keep_alive(KeepAlive::default())
214}
215
216#[cfg(test)]
217mod tests {
218    use super::*;
219    use crate::mcp::schema::{ResourceDescription, ToolDescription};
220
221    fn server_with_echo() -> McpServer {
222        let mut server = McpServer::new("http-test", "1.0.0");
223        server.register_tool(ToolDescription {
224            name: "echo".into(),
225            description: "Echo".into(),
226            input_schema: serde_json::json!({"type": "object"}),
227        });
228        server.set_async_handler("echo", |params| async move { Ok(params) });
229        server
230    }
231
232    #[tokio::test]
233    async fn dispatch_initialize() {
234        let server = server_with_echo();
235        let req = serde_json::json!({"jsonrpc":"2.0","id":1,"method":"initialize","params":{}});
236        let resp = dispatch_async(&server, &req).await.unwrap();
237        assert_eq!(resp["result"]["serverInfo"]["name"], "http-test");
238    }
239
240    #[tokio::test]
241    async fn dispatch_tools_call_async() {
242        let server = server_with_echo();
243        let req = serde_json::json!({
244            "jsonrpc": "2.0", "id": 2, "method": "tools/call",
245            "params": { "name": "echo", "arguments": { "msg": "hello" } }
246        });
247        let resp = dispatch_async(&server, &req).await.unwrap();
248        let text = resp["result"]["content"][0]["text"].as_str().unwrap();
249        assert!(text.contains("hello"));
250    }
251
252    #[tokio::test]
253    async fn dispatch_unknown_method() {
254        let server = server_with_echo();
255        let req = serde_json::json!({"jsonrpc":"2.0","id":3,"method":"nope"});
256        let resp = dispatch_async(&server, &req).await.unwrap();
257        assert_eq!(resp["error"]["code"], ERR_METHOD_NOT_FOUND);
258    }
259
260    /// AC2: HTTP dispatch also serves `resources/read`, not just tools.
261    #[tokio::test]
262    async fn dispatch_resources_read() {
263        let mut server = McpServer::new("http-test", "1.0.0");
264        server.register_resource(ResourceDescription {
265            uri: "docs://x".into(),
266            name: "X".into(),
267            description: None,
268            mime_type: None,
269        });
270        server.set_resource_handler("docs://x", |_| Ok(serde_json::json!("# body")));
271        let req = serde_json::json!({
272            "jsonrpc": "2.0", "id": 4, "method": "resources/read",
273            "params": { "uri": "docs://x" }
274        });
275        let resp = dispatch_async(&server, &req).await.unwrap();
276        let text = resp["result"]["contents"][0]["text"].as_str().unwrap();
277        assert!(text.contains("body"));
278    }
279
280    /// AC2: a full HTTP round-trip — bind an ephemeral port, POST a tools/call,
281    /// and read the JSON-RPC response off the wire.
282    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
283    async fn http_round_trip_calls_tool() {
284        let server = Arc::new(server_with_echo());
285        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
286        let addr = listener.local_addr().unwrap();
287        // Hand the listener to axum in a background task.
288        let transport = HttpTransport::new(server);
289        tokio::spawn(async move {
290            let _ = axum::serve(listener, transport.router()).await;
291        });
292
293        let body = serde_json::to_string(&serde_json::json!({
294            "jsonrpc": "2.0", "id": 9, "method": "tools/call",
295            "params": { "name": "echo", "arguments": { "v": 42 } }
296        }))
297        .unwrap();
298        let req = format!(
299            "POST /mcp HTTP/1.1\r\nHost: localhost\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
300            body.len(),
301            body
302        );
303
304        let mut stream = tokio::net::TcpStream::connect(addr).await.unwrap();
305        use tokio::io::{AsyncReadExt, AsyncWriteExt};
306        stream.write_all(req.as_bytes()).await.unwrap();
307        let mut buf = Vec::new();
308        stream.read_to_end(&mut buf).await.unwrap();
309        let response = String::from_utf8_lossy(&buf);
310        assert!(response.contains("200 OK"), "response: {response}");
311        // The tool result is JSON-encoded inside the `text` field, so its quotes
312        // are escaped on the wire — assert on the unescaped value + content shape.
313        assert!(response.contains("\"content\""), "response: {response}");
314        assert!(response.contains("\\\"v\\\":42"), "response: {response}");
315    }
316}
llm_kernel/mcp/http.rs

llm_kernel/mcp/
http.rs
