llm_kernel/mcp/
transport.rs

1//! JSON-RPC 2.0 transport and dispatch for MCP.
2//!
3//! Reads JSON-RPC requests from stdin (stdio transport) and routes them
4//! to the appropriate [`McpServer`] handler. Writes responses to stdout.
5
6use std::io::{self, BufRead, Write};
7
8use crate::mcp::server::McpServer;
9
10/// JSON-RPC 2.0 dispatcher for MCP stdio transport.
11pub struct JsonRpcDispatcher<'a> {
12    server: &'a McpServer,
13}
14
15impl<'a> JsonRpcDispatcher<'a> {
16    /// Create a new dispatcher wrapping an MCP server.
17    pub fn new(server: &'a McpServer) -> Self {
18        Self { server }
19    }
20
21    /// Dispatch a request, checking the `Authorization` header if auth is configured.
22    ///
23    /// Returns a JSON-RPC `-32001` error response when auth fails.
24    /// Pass `None` for `auth_header` when no header is present (stdio transport).
25    pub fn dispatch_authenticated(
26        &self,
27        request: &str,
28        auth_header: Option<&str>,
29    ) -> Option<String> {
30        let provided = auth_header.unwrap_or("");
31        if !self.server.check_auth(provided) {
32            // Echo back the request id (string or number) when we can parse one,
33            // else a null id per JSON-RPC.
34            let id = serde_json::from_str::<serde_json::Value>(request.trim())
35                .ok()
36                .and_then(|v| v.get("id").cloned())
37                .unwrap_or(serde_json::Value::Null);
38            return Some(self.error_response(id, -32001, "Unauthorized"));
39        }
40        self.dispatch(request)
41    }
42
43    /// Dispatch a JSON-RPC request (single or batch) and return the response.
44    pub fn dispatch(&self, request: &str) -> Option<String> {
45        let trimmed = request.trim();
46        if trimmed.starts_with('[') {
47            // Batch request
48            let reqs: Vec<serde_json::Value> = match serde_json::from_str(trimmed) {
49                Ok(v) => v,
50                Err(e) => {
51                    return Some(self.error_response(
52                        serde_json::Value::Null,
53                        -32700,
54                        &format!("Parse error: {e}"),
55                    ));
56                }
57            };
58            let responses: Vec<String> = reqs
59                .iter()
60                .filter_map(|req| self.dispatch_single(req))
61                .collect();
62            if responses.is_empty() {
63                None
64            } else {
65                Some(format!("[{}]", responses.join(",")))
66            }
67        } else {
68            let req: serde_json::Value = match serde_json::from_str(trimmed) {
69                Ok(v) => v,
70                Err(e) => {
71                    return Some(self.error_response(
72                        serde_json::Value::Null,
73                        -32700,
74                        &format!("Parse error: {e}"),
75                    ));
76                }
77            };
78            self.dispatch_single(&req)
79        }
80    }
81
82    /// Dispatch a single pre-parsed JSON-RPC request.
83    fn dispatch_single(&self, req: &serde_json::Value) -> Option<String> {
84        // Notifications (the `id` member is absent) don't get responses. This is
85        // distinct from a null id, which is a request that must be answered.
86        req.get("id")?;
87        // Preserve the id verbatim (JSON-RPC ids may be a string or a number).
88        let id = req.get("id").cloned().unwrap_or(serde_json::Value::Null);
89        let method = req.get("method").and_then(|v| v.as_str()).unwrap_or("");
90
91        let result = match method {
92            "initialize" => {
93                let requested = req
94                    .get("params")
95                    .and_then(|p| p.get("protocolVersion"))
96                    .and_then(|v| v.as_str());
97                Ok(self.server.initialize_response(requested))
98            }
99            "ping" => Ok(serde_json::json!({})),
100            "tools/list" => Ok(serde_json::json!({
101                "tools": self.server.tools()
102            })),
103            "resources/list" => Ok(serde_json::json!({
104                "resources": self.server.resources()
105            })),
106            "resources/templates/list" => Ok(serde_json::json!({
107                "resourceTemplates": []
108            })),
109            "prompts/list" => Ok(serde_json::json!({
110                "prompts": self.server.prompts()
111            })),
112            "prompts/get" => self.handle_prompt_get(req),
113            "tools/call" => return Some(self.handle_tool_call(&id, req)),
114            "resources/read" => self.handle_resource_read(req),
115            _ => Err((-32601, format!("Method not found: {method}"))),
116        };
117
118        match result {
119            Ok(value) => Some(self.success_response(id, value)),
120            Err((code, message)) => Some(self.error_response(id, code, &message)),
121        }
122    }
123
124    /// Run the stdio transport loop: read lines from stdin, dispatch, write to stdout.
125    pub fn run_stdio(&self) -> io::Result<()> {
126        let stdin = io::stdin();
127        let mut stdout = io::stdout().lock();
128
129        for line in stdin.lock().lines() {
130            let line = line?;
131            let trimmed = line.trim();
132            if trimmed.is_empty() {
133                continue;
134            }
135            if let Some(response) = self.dispatch(trimmed) {
136                writeln!(stdout, "{response}")?;
137                stdout.flush()?;
138            }
139        }
140        Ok(())
141    }
142
143    /// Handle `tools/call`. Returns a full JSON-RPC response string.
144    ///
145    /// An **unknown tool** is a protocol error (`-32602`, invalid params). A
146    /// tool that runs and **fails** is reported in-band as a successful result
147    /// with `isError: true`, per the MCP spec — so the model sees the error and
148    /// can adapt rather than the whole request failing at the transport layer.
149    fn handle_tool_call(&self, id: &serde_json::Value, req: &serde_json::Value) -> String {
150        let tool_name = req
151            .get("params")
152            .and_then(|p| p.get("name"))
153            .and_then(|n| n.as_str())
154            .unwrap_or("");
155
156        let params = req
157            .get("params")
158            .and_then(|p| p.get("arguments"))
159            .cloned()
160            .unwrap_or(serde_json::json!(null));
161
162        if !self.server.has_tool(tool_name) {
163            return self.error_response(id.clone(), -32602, &format!("Unknown tool: {tool_name}"));
164        }
165
166        match self.server.call_tool(tool_name, params) {
167            Ok(result) => self.success_response(
168                id.clone(),
169                serde_json::json!({
170                    "content": [{ "type": "text", "text": result.to_string() }],
171                    "isError": false
172                }),
173            ),
174            Err(e) => self.success_response(
175                id.clone(),
176                serde_json::json!({
177                    "content": [{ "type": "text", "text": e.to_string() }],
178                    "isError": true
179                }),
180            ),
181        }
182    }
183
184    /// Handle `prompts/get`: render a registered prompt with the given
185    /// arguments. An unknown prompt is an invalid-params error.
186    fn handle_prompt_get(
187        &self,
188        req: &serde_json::Value,
189    ) -> std::result::Result<serde_json::Value, (i32, String)> {
190        let name = req
191            .get("params")
192            .and_then(|p| p.get("name"))
193            .and_then(|n| n.as_str())
194            .unwrap_or("");
195        let args = req
196            .get("params")
197            .and_then(|p| p.get("arguments"))
198            .cloned()
199            .unwrap_or(serde_json::json!({}));
200        self.server
201            .get_prompt(name, args)
202            .map_err(|e| (-32602, e.to_string()))
203    }
204
205    fn handle_resource_read(
206        &self,
207        req: &serde_json::Value,
208    ) -> std::result::Result<serde_json::Value, (i32, String)> {
209        let uri = req
210            .get("params")
211            .and_then(|p| p.get("uri"))
212            .and_then(|u| u.as_str())
213            .unwrap_or("");
214        self.server
215            .read_resource(uri, serde_json::json!({}))
216            .map(|content| {
217                serde_json::json!({
218                    "contents": [{
219                        "uri": uri,
220                        "text": content.to_string()
221                    }]
222                })
223            })
224            .map_err(|e| (-32603, e.to_string()))
225    }
226
227    fn success_response(&self, id: serde_json::Value, result: serde_json::Value) -> String {
228        serde_json::to_string(&serde_json::json!({
229            "jsonrpc": "2.0",
230            "id": id,
231            "result": result,
232        }))
233        .unwrap_or_default()
234    }
235
236    fn error_response(&self, id: serde_json::Value, code: i32, message: &str) -> String {
237        serde_json::to_string(&serde_json::json!({
238            "jsonrpc": "2.0",
239            "id": id,
240            "error": {
241                "code": code,
242                "message": message,
243            }
244        }))
245        .unwrap_or_default()
246    }
247}
248
249#[cfg(test)]
250mod tests {
251    use super::*;
252    use crate::mcp::schema::ToolDescription;
253
254    fn test_server() -> McpServer {
255        let mut server = McpServer::new("test-server", "0.1.0");
256        server.register_tool(ToolDescription {
257            name: "echo".into(),
258            description: "Echo input".into(),
259            input_schema: serde_json::json!({"type": "object"}),
260        });
261        server.set_handler("echo", Ok);
262        server
263    }
264
265    #[test]
266    fn dispatch_initialize() {
267        let server = test_server();
268        let dispatcher = JsonRpcDispatcher::new(&server);
269        let resp = dispatcher
270            .dispatch(r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}"#)
271            .unwrap();
272        let parsed: serde_json::Value = serde_json::from_str(&resp).unwrap();
273        assert_eq!(parsed["result"]["serverInfo"]["name"], "test-server");
274    }
275
276    #[test]
277    fn dispatch_tools_list() {
278        let server = test_server();
279        let dispatcher = JsonRpcDispatcher::new(&server);
280        let resp = dispatcher
281            .dispatch(r#"{"jsonrpc":"2.0","id":2,"method":"tools/list"}"#)
282            .unwrap();
283        let parsed: serde_json::Value = serde_json::from_str(&resp).unwrap();
284        assert_eq!(parsed["result"]["tools"].as_array().unwrap().len(), 1);
285    }
286
287    #[test]
288    fn dispatch_tools_call() {
289        let server = test_server();
290        let dispatcher = JsonRpcDispatcher::new(&server);
291        let req = r#"{"jsonrpc":"2.0","id":3,"method":"tools/call","params":{"name":"echo","arguments":{"msg":"hello"}}}"#;
292        let resp = dispatcher.dispatch(req).unwrap();
293        let parsed: serde_json::Value = serde_json::from_str(&resp).unwrap();
294        let text = parsed["result"]["content"][0]["text"].as_str().unwrap();
295        assert!(text.contains("hello"));
296    }
297
298    #[test]
299    fn dispatch_unknown_method() {
300        let server = test_server();
301        let dispatcher = JsonRpcDispatcher::new(&server);
302        let resp = dispatcher
303            .dispatch(r#"{"jsonrpc":"2.0","id":4,"method":"nonexistent"}"#)
304            .unwrap();
305        let parsed: serde_json::Value = serde_json::from_str(&resp).unwrap();
306        assert_eq!(parsed["error"]["code"], -32601);
307    }
308
309    #[test]
310    fn dispatch_invalid_json() {
311        let server = test_server();
312        let dispatcher = JsonRpcDispatcher::new(&server);
313        let resp = dispatcher.dispatch("not json").unwrap();
314        let parsed: serde_json::Value = serde_json::from_str(&resp).unwrap();
315        assert_eq!(parsed["error"]["code"], -32700);
316    }
317
318    #[test]
319    fn dispatch_unknown_tool_is_invalid_params() {
320        // An unknown tool is a protocol error (-32602), not an in-band failure.
321        let server = test_server();
322        let dispatcher = JsonRpcDispatcher::new(&server);
323        let req = r#"{"jsonrpc":"2.0","id":5,"method":"tools/call","params":{"name":"missing","arguments":{}}}"#;
324        let resp = dispatcher.dispatch(req).unwrap();
325        let parsed: serde_json::Value = serde_json::from_str(&resp).unwrap();
326        assert_eq!(parsed["error"]["code"], -32602);
327    }
328
329    #[test]
330    fn dispatch_ping_returns_empty_result() {
331        let server = test_server();
332        let dispatcher = JsonRpcDispatcher::new(&server);
333        let resp = dispatcher
334            .dispatch(r#"{"jsonrpc":"2.0","id":9,"method":"ping"}"#)
335            .unwrap();
336        let parsed: serde_json::Value = serde_json::from_str(&resp).unwrap();
337        assert_eq!(parsed["id"], 9);
338        assert!(parsed["result"].is_object());
339        assert_eq!(parsed["result"].as_object().unwrap().len(), 0);
340    }
341
342    #[test]
343    fn dispatch_preserves_string_id() {
344        let server = test_server();
345        let dispatcher = JsonRpcDispatcher::new(&server);
346        let resp = dispatcher
347            .dispatch(r#"{"jsonrpc":"2.0","id":"req-abc","method":"tools/list"}"#)
348            .unwrap();
349        let parsed: serde_json::Value = serde_json::from_str(&resp).unwrap();
350        assert_eq!(parsed["id"], "req-abc");
351    }
352
353    #[test]
354    fn initialize_echoes_client_protocol_version() {
355        let server = test_server();
356        let dispatcher = JsonRpcDispatcher::new(&server);
357        let resp = dispatcher
358            .dispatch(
359                r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2024-11-05"}}"#,
360            )
361            .unwrap();
362        let parsed: serde_json::Value = serde_json::from_str(&resp).unwrap();
363        assert_eq!(parsed["result"]["protocolVersion"], "2024-11-05");
364    }
365
366    #[test]
367    fn tool_execution_error_reported_in_band() {
368        // A registered tool whose handler fails → result with isError: true,
369        // NOT a JSON-RPC error object.
370        let mut server = McpServer::new("t", "1.0");
371        server.register_tool(ToolDescription {
372            name: "boom".into(),
373            description: "always fails".into(),
374            input_schema: serde_json::json!({"type": "object"}),
375        });
376        server.set_handler("boom", |_| {
377            Err(crate::error::KernelError::Config("kaboom".into()))
378        });
379        let dispatcher = JsonRpcDispatcher::new(&server);
380        let req = r#"{"jsonrpc":"2.0","id":3,"method":"tools/call","params":{"name":"boom","arguments":{}}}"#;
381        let resp = dispatcher.dispatch(req).unwrap();
382        let parsed: serde_json::Value = serde_json::from_str(&resp).unwrap();
383        assert!(
384            parsed.get("error").is_none(),
385            "should not be a protocol error"
386        );
387        assert_eq!(parsed["result"]["isError"], true);
388        assert!(
389            parsed["result"]["content"][0]["text"]
390                .as_str()
391                .unwrap()
392                .contains("kaboom")
393        );
394    }
395
396    #[test]
397    fn dispatch_prompts_list_and_get() {
398        let mut server = McpServer::new("t", "1.0");
399        server.register_prompt(crate::mcp::schema::PromptDescription {
400            name: "greet".into(),
401            description: Some("Greet".into()),
402            arguments: Vec::new(),
403        });
404        server.set_prompt_handler("greet", |_| {
405            Ok(serde_json::json!({
406                "messages": [{ "role": "user", "content": { "type": "text", "text": "hi" } }]
407            }))
408        });
409        let dispatcher = JsonRpcDispatcher::new(&server);
410
411        let list = dispatcher
412            .dispatch(r#"{"jsonrpc":"2.0","id":1,"method":"prompts/list"}"#)
413            .unwrap();
414        let list: serde_json::Value = serde_json::from_str(&list).unwrap();
415        assert_eq!(list["result"]["prompts"][0]["name"], "greet");
416
417        let got = dispatcher
418            .dispatch(r#"{"jsonrpc":"2.0","id":2,"method":"prompts/get","params":{"name":"greet","arguments":{}}}"#)
419            .unwrap();
420        let got: serde_json::Value = serde_json::from_str(&got).unwrap();
421        assert_eq!(got["result"]["messages"][0]["content"]["text"], "hi");
422    }
423
424    #[test]
425    fn dispatch_resource_templates_list_is_empty() {
426        let server = test_server();
427        let dispatcher = JsonRpcDispatcher::new(&server);
428        let resp = dispatcher
429            .dispatch(r#"{"jsonrpc":"2.0","id":1,"method":"resources/templates/list"}"#)
430            .unwrap();
431        let parsed: serde_json::Value = serde_json::from_str(&resp).unwrap();
432        assert!(parsed["result"]["resourceTemplates"].is_array());
433    }
434
435    #[test]
436    fn notification_without_id_gets_no_response() {
437        let server = test_server();
438        let dispatcher = JsonRpcDispatcher::new(&server);
439        // `notifications/initialized` is a notification (no id) → no response.
440        assert!(
441            dispatcher
442                .dispatch(r#"{"jsonrpc":"2.0","method":"notifications/initialized"}"#)
443                .is_none()
444        );
445    }
446
447    #[test]
448    fn dispatch_batch_request() {
449        let server = test_server();
450        let dispatcher = JsonRpcDispatcher::new(&server);
451        let batch = r#"[
452            {"jsonrpc":"2.0","id":1,"method":"initialize","params":{}},
453            {"jsonrpc":"2.0","id":2,"method":"tools/list"}
454        ]"#;
455        let resp = dispatcher.dispatch(batch).unwrap();
456        let parsed: serde_json::Value = serde_json::from_str(&resp).unwrap();
457        let arr = parsed.as_array().unwrap();
458        assert_eq!(arr.len(), 2);
459        // First response: initialize
460        assert_eq!(arr[0]["result"]["serverInfo"]["name"], "test-server");
461        // Second response: tools/list
462        assert_eq!(arr[1]["result"]["tools"].as_array().unwrap().len(), 1);
463    }
464
465    #[test]
466    fn dispatch_batch_with_error() {
467        let server = test_server();
468        let dispatcher = JsonRpcDispatcher::new(&server);
469        let batch = r#"[
470            {"jsonrpc":"2.0","id":1,"method":"initialize","params":{}},
471            {"jsonrpc":"2.0","id":2,"method":"nonexistent"}
472        ]"#;
473        let resp = dispatcher.dispatch(batch).unwrap();
474        let parsed: serde_json::Value = serde_json::from_str(&resp).unwrap();
475        let arr = parsed.as_array().unwrap();
476        assert_eq!(arr[0]["result"]["serverInfo"]["name"], "test-server");
477        assert_eq!(arr[1]["error"]["code"], -32601);
478    }
479
480    #[test]
481    fn dispatch_authenticated_passes_with_no_auth_configured() {
482        let server = test_server();
483        let dispatcher = JsonRpcDispatcher::new(&server);
484        let req = r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}"#;
485        let resp = dispatcher.dispatch_authenticated(req, None).unwrap();
486        let parsed: serde_json::Value = serde_json::from_str(&resp).unwrap();
487        assert!(parsed["result"]["serverInfo"].is_object());
488    }
489
490    #[test]
491    fn dispatch_authenticated_rejects_missing_token() {
492        let server = McpServer::new("secured", "1.0").with_bearer_auth("secret-token");
493        let dispatcher = JsonRpcDispatcher::new(&server);
494        let req = r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}"#;
495        let resp = dispatcher.dispatch_authenticated(req, None).unwrap();
496        let parsed: serde_json::Value = serde_json::from_str(&resp).unwrap();
497        assert_eq!(parsed["error"]["code"], -32001);
498    }
499
500    #[test]
501    fn dispatch_authenticated_rejects_wrong_token() {
502        let server = McpServer::new("secured", "1.0").with_bearer_auth("correct-token");
503        let dispatcher = JsonRpcDispatcher::new(&server);
504        let req = r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}"#;
505        let resp = dispatcher
506            .dispatch_authenticated(req, Some("Bearer wrong-token"))
507            .unwrap();
508        let parsed: serde_json::Value = serde_json::from_str(&resp).unwrap();
509        assert_eq!(parsed["error"]["code"], -32001);
510    }
511
512    #[test]
513    fn dispatch_authenticated_passes_with_correct_token() {
514        let server = McpServer::new("secured", "1.0").with_bearer_auth("correct-token");
515        let dispatcher = JsonRpcDispatcher::new(&server);
516        let req = r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}"#;
517        let resp = dispatcher
518            .dispatch_authenticated(req, Some("Bearer correct-token"))
519            .unwrap();
520        let parsed: serde_json::Value = serde_json::from_str(&resp).unwrap();
521        assert!(parsed["result"]["serverInfo"].is_object());
522    }
523}
llm_kernel/mcp/transport.rs

llm_kernel/mcp/
transport.rs
