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/// Cap on a single stdio JSON-RPC line. `BufRead::lines` grows without bound,
11/// so a peer that never sends `\n` can exhaust memory — fatal under a
12/// `panic = "abort"` release profile.
13const MAX_LINE_BYTES: usize = 16 * 1024 * 1024;
14
15/// JSON-RPC 2.0 dispatcher for MCP stdio transport.
16pub struct JsonRpcDispatcher<'a> {
17    server: &'a McpServer,
18}
19
20impl<'a> JsonRpcDispatcher<'a> {
21    /// Create a new dispatcher wrapping an MCP server.
22    pub fn new(server: &'a McpServer) -> Self {
23        Self { server }
24    }
25
26    /// Dispatch a request, checking the `Authorization` header if auth is configured.
27    ///
28    /// Returns a JSON-RPC `-32001` error response when auth fails.
29    /// Pass `None` for `auth_header` when no header is present (stdio transport).
30    pub fn dispatch_authenticated(
31        &self,
32        request: &str,
33        auth_header: Option<&str>,
34    ) -> Option<String> {
35        let provided = auth_header.unwrap_or("");
36        if !self.server.check_auth(provided) {
37            // Echo back the request id (string or number) when we can parse one,
38            // else a null id per JSON-RPC.
39            let id = serde_json::from_str::<serde_json::Value>(request.trim())
40                .ok()
41                .and_then(|v| v.get("id").cloned())
42                .unwrap_or(serde_json::Value::Null);
43            return Some(self.error_response(id, -32001, "Unauthorized"));
44        }
45        self.dispatch(request)
46    }
47
48    /// Dispatch a JSON-RPC request (single or batch) and return the response.
49    pub fn dispatch(&self, request: &str) -> Option<String> {
50        let trimmed = request.trim();
51        if trimmed.starts_with('[') {
52            // Batch request
53            let reqs: Vec<serde_json::Value> = match serde_json::from_str(trimmed) {
54                Ok(v) => v,
55                Err(e) => {
56                    return Some(self.error_response(
57                        serde_json::Value::Null,
58                        -32700,
59                        &format!("Parse error: {e}"),
60                    ));
61                }
62            };
63            let responses: Vec<String> = reqs
64                .iter()
65                .filter_map(|req| self.dispatch_single(req))
66                .collect();
67            if responses.is_empty() {
68                None
69            } else {
70                Some(format!("[{}]", responses.join(",")))
71            }
72        } else {
73            let req: serde_json::Value = match serde_json::from_str(trimmed) {
74                Ok(v) => v,
75                Err(e) => {
76                    return Some(self.error_response(
77                        serde_json::Value::Null,
78                        -32700,
79                        &format!("Parse error: {e}"),
80                    ));
81                }
82            };
83            self.dispatch_single(&req)
84        }
85    }
86
87    /// Dispatch a single pre-parsed JSON-RPC request.
88    fn dispatch_single(&self, req: &serde_json::Value) -> Option<String> {
89        // A non-object is never a valid request — answering with silence
90        // (the notification path) leaves the client waiting forever.
91        if !req.is_object() {
92            return Some(self.error_response(
93                serde_json::Value::Null,
94                -32600,
95                "Invalid Request: expected a JSON object",
96            ));
97        }
98        // Notifications (the `id` member is absent) don't get responses. This is
99        // distinct from a null id, which is a request that must be answered.
100        req.get("id")?;
101        // Preserve the id verbatim (JSON-RPC ids may be a string or a number).
102        let id = req.get("id").cloned().unwrap_or(serde_json::Value::Null);
103        let method = req.get("method").and_then(|v| v.as_str()).unwrap_or("");
104
105        let result = match method {
106            "initialize" => {
107                let requested = req
108                    .get("params")
109                    .and_then(|p| p.get("protocolVersion"))
110                    .and_then(|v| v.as_str());
111                Ok(self.server.initialize_response(requested))
112            }
113            "ping" => Ok(serde_json::json!({})),
114            "tools/list" => Ok(serde_json::json!({
115                "tools": self.server.tools()
116            })),
117            "resources/list" => Ok(serde_json::json!({
118                "resources": self.server.resources()
119            })),
120            "resources/templates/list" => Ok(serde_json::json!({
121                "resourceTemplates": []
122            })),
123            "prompts/list" => Ok(serde_json::json!({
124                "prompts": self.server.prompts()
125            })),
126            "prompts/get" => self.handle_prompt_get(req),
127            "tools/call" => return Some(self.handle_tool_call(&id, req)),
128            "resources/read" => self.handle_resource_read(req),
129            _ => Err((-32601, format!("Method not found: {method}"))),
130        };
131
132        match result {
133            Ok(value) => Some(self.success_response(id, value)),
134            Err((code, message)) => Some(self.error_response(id, code, &message)),
135        }
136    }
137
138    /// Dispatch a JSON-RPC request, awaiting async tool handlers.
139    ///
140    /// Identical to [`JsonRpcDispatcher::dispatch`] except that `tools/call`
141    /// resolves through [`McpServer::call_tool_async`], so tools registered
142    /// with `set_async_handler` work. The synchronous [`Self::dispatch`]
143    /// cannot invoke them — prefer this one whenever the server has any
144    /// async handler.
145    pub async fn dispatch_async(&self, request: &str) -> Option<String> {
146        let trimmed = request.trim();
147        if trimmed.starts_with('[') {
148            let reqs: Vec<serde_json::Value> = match serde_json::from_str(trimmed) {
149                Ok(v) => v,
150                Err(e) => {
151                    return Some(self.error_response(
152                        serde_json::Value::Null,
153                        -32700,
154                        &format!("Parse error: {e}"),
155                    ));
156                }
157            };
158            let mut responses: Vec<String> = Vec::with_capacity(reqs.len());
159            for req in &reqs {
160                if let Some(r) = self.dispatch_single_async(req).await {
161                    responses.push(r);
162                }
163            }
164            if responses.is_empty() {
165                None
166            } else {
167                Some(format!("[{}]", responses.join(",")))
168            }
169        } else {
170            let req: serde_json::Value = match serde_json::from_str(trimmed) {
171                Ok(v) => v,
172                Err(e) => {
173                    return Some(self.error_response(
174                        serde_json::Value::Null,
175                        -32700,
176                        &format!("Parse error: {e}"),
177                    ));
178                }
179            };
180            self.dispatch_single_async(&req).await
181        }
182    }
183
184    /// [`Self::dispatch_authenticated`] with async tool handler support.
185    pub async fn dispatch_authenticated_async(
186        &self,
187        request: &str,
188        auth_header: Option<&str>,
189    ) -> Option<String> {
190        if !self.server.check_auth(auth_header.unwrap_or("")) {
191            let id = serde_json::from_str::<serde_json::Value>(request.trim())
192                .ok()
193                .and_then(|v| v.get("id").cloned())
194                .unwrap_or(serde_json::Value::Null);
195            return Some(self.error_response(id, -32001, "Unauthorized"));
196        }
197        self.dispatch_async(request).await
198    }
199
200    async fn dispatch_single_async(&self, req: &serde_json::Value) -> Option<String> {
201        // Only `tools/call` can reach an async handler; everything else is
202        // pure metadata and shares the synchronous path.
203        if req.get("method").and_then(|v| v.as_str()) == Some("tools/call") {
204            req.get("id")?;
205            let id = req.get("id").cloned().unwrap_or(serde_json::Value::Null);
206            return Some(self.handle_tool_call_async(&id, req).await);
207        }
208        self.dispatch_single(req)
209    }
210
211    /// Run the stdio transport loop: read lines from stdin, dispatch, write to stdout.
212    ///
213    /// Tools registered with `set_async_handler` are NOT callable from this
214    /// loop — use [`Self::run_stdio_async`] when the server has any.
215    ///
216    /// # Errors
217    ///
218    /// Fails immediately (`InvalidInput`) when the server has any async-only
219    /// tool: every call to it would otherwise return `isError` at runtime and
220    /// look like a handler bug.
221    pub fn run_stdio(&self) -> io::Result<()> {
222        if let Some(name) = self.server.async_only_tools().first() {
223            return Err(io::Error::new(
224                io::ErrorKind::InvalidInput,
225                format!(
226                    "tool '{name}' has only an async handler; use run_stdio_async instead of run_stdio"
227                ),
228            ));
229        }
230        let stdin = io::stdin();
231        let mut stdout = io::stdout().lock();
232
233        for line in stdin.lock().lines() {
234            let line = line?;
235            if line.len() > MAX_LINE_BYTES {
236                return Err(io::Error::new(
237                    io::ErrorKind::InvalidData,
238                    format!("JSON-RPC line exceeds {MAX_LINE_BYTES} bytes"),
239                ));
240            }
241            let trimmed = line.trim();
242            if trimmed.is_empty() {
243                continue;
244            }
245            if let Some(response) = self.dispatch(trimmed) {
246                writeln!(stdout, "{response}")?;
247                stdout.flush()?;
248            }
249        }
250        Ok(())
251    }
252
253    /// [`Self::run_stdio`] for servers with async tool handlers.
254    ///
255    /// Reads stdin with blocking I/O between awaits — an MCP stdio server is
256    /// a dedicated process, so occupying the calling task is intended. Drive
257    /// it from a runtime's blocking-friendly context (e.g. a dedicated task).
258    pub async fn run_stdio_async(&self) -> io::Result<()> {
259        let stdin = io::stdin();
260        let mut stdout = io::stdout().lock();
261
262        for line in stdin.lock().lines() {
263            let line = line?;
264            if line.len() > MAX_LINE_BYTES {
265                return Err(io::Error::new(
266                    io::ErrorKind::InvalidData,
267                    format!("JSON-RPC line exceeds {MAX_LINE_BYTES} bytes"),
268                ));
269            }
270            let trimmed = line.trim();
271            if trimmed.is_empty() {
272                continue;
273            }
274            if let Some(response) = self.dispatch_async(trimmed).await {
275                writeln!(stdout, "{response}")?;
276                stdout.flush()?;
277            }
278        }
279        Ok(())
280    }
281
282    /// [`Self::handle_tool_call`] resolving through `call_tool_async`.
283    async fn handle_tool_call_async(
284        &self,
285        id: &serde_json::Value,
286        req: &serde_json::Value,
287    ) -> String {
288        let tool_name = req
289            .get("params")
290            .and_then(|p| p.get("name"))
291            .and_then(|n| n.as_str())
292            .unwrap_or("");
293        let params = req
294            .get("params")
295            .and_then(|p| p.get("arguments"))
296            .cloned()
297            .unwrap_or(serde_json::json!(null));
298
299        if !self.server.has_tool(tool_name) {
300            return self.error_response(id.clone(), -32602, &format!("Unknown tool: {tool_name}"));
301        }
302        if let Err(e) = self.server.validate_tool_args(tool_name, &params) {
303            return self.error_response(id.clone(), -32602, &e);
304        }
305
306        match self.server.call_tool_async(tool_name, params).await {
307            Ok(result) => self.success_response(
308                id.clone(),
309                serde_json::json!({
310                    "content": [{ "type": "text", "text": result.to_string() }],
311                    "isError": false
312                }),
313            ),
314            Err(e) => self.success_response(
315                id.clone(),
316                serde_json::json!({
317                    "content": [{ "type": "text", "text": e.to_string() }],
318                    "isError": true
319                }),
320            ),
321        }
322    }
323
324    /// Handle `tools/call`. Returns a full JSON-RPC response string.
325    ///
326    /// An **unknown tool** is a protocol error (`-32602`, invalid params). A
327    /// tool that runs and **fails** is reported in-band as a successful result
328    /// with `isError: true`, per the MCP spec — so the model sees the error and
329    /// can adapt rather than the whole request failing at the transport layer.
330    fn handle_tool_call(&self, id: &serde_json::Value, req: &serde_json::Value) -> String {
331        let tool_name = req
332            .get("params")
333            .and_then(|p| p.get("name"))
334            .and_then(|n| n.as_str())
335            .unwrap_or("");
336
337        let params = req
338            .get("params")
339            .and_then(|p| p.get("arguments"))
340            .cloned()
341            .unwrap_or(serde_json::json!(null));
342
343        if !self.server.has_tool(tool_name) {
344            return self.error_response(id.clone(), -32602, &format!("Unknown tool: {tool_name}"));
345        }
346        if let Err(e) = self.server.validate_tool_args(tool_name, &params) {
347            return self.error_response(id.clone(), -32602, &e);
348        }
349
350        match self.server.call_tool(tool_name, params) {
351            Ok(result) => self.success_response(
352                id.clone(),
353                serde_json::json!({
354                    "content": [{ "type": "text", "text": result.to_string() }],
355                    "isError": false
356                }),
357            ),
358            Err(e) => self.success_response(
359                id.clone(),
360                serde_json::json!({
361                    "content": [{ "type": "text", "text": e.to_string() }],
362                    "isError": true
363                }),
364            ),
365        }
366    }
367
368    /// Handle `prompts/get`: render a registered prompt with the given
369    /// arguments. An unknown prompt is an invalid-params error.
370    fn handle_prompt_get(
371        &self,
372        req: &serde_json::Value,
373    ) -> std::result::Result<serde_json::Value, (i32, String)> {
374        let name = req
375            .get("params")
376            .and_then(|p| p.get("name"))
377            .and_then(|n| n.as_str())
378            .unwrap_or("");
379        let args = req
380            .get("params")
381            .and_then(|p| p.get("arguments"))
382            .cloned()
383            .unwrap_or(serde_json::json!({}));
384        self.server
385            .get_prompt(name, args)
386            .map_err(|e| (-32602, e.to_string()))
387    }
388
389    fn handle_resource_read(
390        &self,
391        req: &serde_json::Value,
392    ) -> std::result::Result<serde_json::Value, (i32, String)> {
393        let uri = req
394            .get("params")
395            .and_then(|p| p.get("uri"))
396            .and_then(|u| u.as_str())
397            .unwrap_or("");
398        self.server
399            .read_resource(uri, serde_json::json!({}))
400            .map(|content| {
401                serde_json::json!({
402                    "contents": [{
403                        "uri": uri,
404                        "text": content.to_string()
405                    }]
406                })
407            })
408            .map_err(|e| (-32603, e.to_string()))
409    }
410
411    fn success_response(&self, id: serde_json::Value, result: serde_json::Value) -> String {
412        serde_json::to_string(&serde_json::json!({
413            "jsonrpc": "2.0",
414            "id": id,
415            "result": result,
416        }))
417        .unwrap_or_default()
418    }
419
420    fn error_response(&self, id: serde_json::Value, code: i32, message: &str) -> String {
421        serde_json::to_string(&serde_json::json!({
422            "jsonrpc": "2.0",
423            "id": id,
424            "error": {
425                "code": code,
426                "message": message,
427            }
428        }))
429        .unwrap_or_default()
430    }
431}
432
433#[cfg(test)]
434mod tests {
435    use super::*;
436    use crate::mcp::schema::ToolDescription;
437
438    /// A server whose tool is registered ONLY via `set_async_handler` — the
439    /// shape real consumers use when their handlers await I/O.
440    fn async_only_server() -> McpServer {
441        let mut server = McpServer::new("async-server", "0.1.0");
442        server.register_tool(ToolDescription {
443            name: "search".into(),
444            description: "async search".into(),
445            input_schema: serde_json::json!({"type": "object"}),
446        });
447        server.set_async_handler("search", |p: serde_json::Value| async move { Ok(p) });
448        server
449    }
450
451    #[tokio::test]
452    async fn async_only_tool_runs_over_dispatch_async() {
453        let server = async_only_server();
454        let dispatcher = JsonRpcDispatcher::new(&server);
455        let req = r#"{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"search","arguments":{"q":"x"}}}"#;
456        let resp = dispatcher.dispatch_async(req).await.expect("response");
457        let parsed: serde_json::Value = serde_json::from_str(&resp).unwrap();
458        assert_eq!(parsed["result"]["isError"], false, "{parsed}");
459        assert!(
460            parsed["result"]["content"][0]["text"]
461                .as_str()
462                .unwrap()
463                .contains("\"q\""),
464            "{parsed}"
465        );
466    }
467
468    #[tokio::test]
469    async fn dispatch_async_still_serves_metadata_methods() {
470        let server = async_only_server();
471        let dispatcher = JsonRpcDispatcher::new(&server);
472        let resp = dispatcher
473            .dispatch_async(r#"{"jsonrpc":"2.0","id":7,"method":"tools/list"}"#)
474            .await
475            .expect("response");
476        let parsed: serde_json::Value = serde_json::from_str(&resp).unwrap();
477        assert_eq!(parsed["id"], 7);
478        assert_eq!(parsed["result"]["tools"][0]["name"], "search");
479    }
480
481    #[tokio::test]
482    async fn missing_required_argument_is_rejected_not_defaulted() {
483        let mut server = McpServer::new("s", "1.0");
484        server.register_tool(ToolDescription {
485            name: "search".into(),
486            description: "d".into(),
487            input_schema: serde_json::json!({
488                "type": "object",
489                "properties": {"query": {"type": "string"}},
490                "required": ["query"]
491            }),
492        });
493        server.set_async_handler("search", |p: serde_json::Value| async move { Ok(p) });
494        let dispatcher = JsonRpcDispatcher::new(&server);
495
496        // No arguments at all — the handler would see null and search for "".
497        let resp = dispatcher
498            .dispatch_async(
499                r#"{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"search"}}"#,
500            )
501            .await
502            .expect("response");
503        assert!(resp.contains("-32602"), "{resp}");
504        assert!(resp.contains("query"), "{resp}");
505
506        // Supplying it dispatches normally.
507        let ok = dispatcher
508            .dispatch_async(
509                r#"{"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"search","arguments":{"query":"x"}}}"#,
510            )
511            .await
512            .expect("response");
513        let parsed: serde_json::Value = serde_json::from_str(&ok).unwrap();
514        assert_eq!(parsed["result"]["isError"], false, "{ok}");
515    }
516
517    #[test]
518    fn non_object_request_gets_invalid_request_not_silence() {
519        let server = async_only_server();
520        let dispatcher = JsonRpcDispatcher::new(&server);
521        let resp = dispatcher
522            .dispatch("[1, 2]")
523            .expect("must answer, not hang");
524        assert!(resp.contains("-32600"), "{resp}");
525    }
526
527    #[test]
528    fn sync_call_of_async_only_tool_says_so() {
529        let server = async_only_server();
530        let err = server
531            .call_tool("search", serde_json::json!({}))
532            .unwrap_err()
533            .to_string();
534        assert!(err.contains("async"), "misleading error: {err}");
535    }
536
537    fn test_server() -> McpServer {
538        let mut server = McpServer::new("test-server", "0.1.0");
539        server.register_tool(ToolDescription {
540            name: "echo".into(),
541            description: "Echo input".into(),
542            input_schema: serde_json::json!({"type": "object"}),
543        });
544        server.set_handler("echo", Ok);
545        server
546    }
547
548    #[test]
549    fn dispatch_initialize() {
550        let server = test_server();
551        let dispatcher = JsonRpcDispatcher::new(&server);
552        let resp = dispatcher
553            .dispatch(r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}"#)
554            .unwrap();
555        let parsed: serde_json::Value = serde_json::from_str(&resp).unwrap();
556        assert_eq!(parsed["result"]["serverInfo"]["name"], "test-server");
557    }
558
559    #[test]
560    fn dispatch_tools_list() {
561        let server = test_server();
562        let dispatcher = JsonRpcDispatcher::new(&server);
563        let resp = dispatcher
564            .dispatch(r#"{"jsonrpc":"2.0","id":2,"method":"tools/list"}"#)
565            .unwrap();
566        let parsed: serde_json::Value = serde_json::from_str(&resp).unwrap();
567        assert_eq!(parsed["result"]["tools"].as_array().unwrap().len(), 1);
568    }
569
570    #[test]
571    fn dispatch_tools_call() {
572        let server = test_server();
573        let dispatcher = JsonRpcDispatcher::new(&server);
574        let req = r#"{"jsonrpc":"2.0","id":3,"method":"tools/call","params":{"name":"echo","arguments":{"msg":"hello"}}}"#;
575        let resp = dispatcher.dispatch(req).unwrap();
576        let parsed: serde_json::Value = serde_json::from_str(&resp).unwrap();
577        let text = parsed["result"]["content"][0]["text"].as_str().unwrap();
578        assert!(text.contains("hello"));
579    }
580
581    #[test]
582    fn dispatch_unknown_method() {
583        let server = test_server();
584        let dispatcher = JsonRpcDispatcher::new(&server);
585        let resp = dispatcher
586            .dispatch(r#"{"jsonrpc":"2.0","id":4,"method":"nonexistent"}"#)
587            .unwrap();
588        let parsed: serde_json::Value = serde_json::from_str(&resp).unwrap();
589        assert_eq!(parsed["error"]["code"], -32601);
590    }
591
592    #[test]
593    fn dispatch_invalid_json() {
594        let server = test_server();
595        let dispatcher = JsonRpcDispatcher::new(&server);
596        let resp = dispatcher.dispatch("not json").unwrap();
597        let parsed: serde_json::Value = serde_json::from_str(&resp).unwrap();
598        assert_eq!(parsed["error"]["code"], -32700);
599    }
600
601    #[test]
602    fn dispatch_unknown_tool_is_invalid_params() {
603        // An unknown tool is a protocol error (-32602), not an in-band failure.
604        let server = test_server();
605        let dispatcher = JsonRpcDispatcher::new(&server);
606        let req = r#"{"jsonrpc":"2.0","id":5,"method":"tools/call","params":{"name":"missing","arguments":{}}}"#;
607        let resp = dispatcher.dispatch(req).unwrap();
608        let parsed: serde_json::Value = serde_json::from_str(&resp).unwrap();
609        assert_eq!(parsed["error"]["code"], -32602);
610    }
611
612    #[test]
613    fn dispatch_ping_returns_empty_result() {
614        let server = test_server();
615        let dispatcher = JsonRpcDispatcher::new(&server);
616        let resp = dispatcher
617            .dispatch(r#"{"jsonrpc":"2.0","id":9,"method":"ping"}"#)
618            .unwrap();
619        let parsed: serde_json::Value = serde_json::from_str(&resp).unwrap();
620        assert_eq!(parsed["id"], 9);
621        assert!(parsed["result"].is_object());
622        assert_eq!(parsed["result"].as_object().unwrap().len(), 0);
623    }
624
625    #[test]
626    fn dispatch_preserves_string_id() {
627        let server = test_server();
628        let dispatcher = JsonRpcDispatcher::new(&server);
629        let resp = dispatcher
630            .dispatch(r#"{"jsonrpc":"2.0","id":"req-abc","method":"tools/list"}"#)
631            .unwrap();
632        let parsed: serde_json::Value = serde_json::from_str(&resp).unwrap();
633        assert_eq!(parsed["id"], "req-abc");
634    }
635
636    #[test]
637    fn initialize_echoes_client_protocol_version() {
638        let server = test_server();
639        let dispatcher = JsonRpcDispatcher::new(&server);
640        let resp = dispatcher
641            .dispatch(
642                r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2024-11-05"}}"#,
643            )
644            .unwrap();
645        let parsed: serde_json::Value = serde_json::from_str(&resp).unwrap();
646        assert_eq!(parsed["result"]["protocolVersion"], "2024-11-05");
647    }
648
649    #[test]
650    fn tool_execution_error_reported_in_band() {
651        // A registered tool whose handler fails → result with isError: true,
652        // NOT a JSON-RPC error object.
653        let mut server = McpServer::new("t", "1.0");
654        server.register_tool(ToolDescription {
655            name: "boom".into(),
656            description: "always fails".into(),
657            input_schema: serde_json::json!({"type": "object"}),
658        });
659        server.set_handler("boom", |_| {
660            Err(crate::error::KernelError::Config("kaboom".into()))
661        });
662        let dispatcher = JsonRpcDispatcher::new(&server);
663        let req = r#"{"jsonrpc":"2.0","id":3,"method":"tools/call","params":{"name":"boom","arguments":{}}}"#;
664        let resp = dispatcher.dispatch(req).unwrap();
665        let parsed: serde_json::Value = serde_json::from_str(&resp).unwrap();
666        assert!(
667            parsed.get("error").is_none(),
668            "should not be a protocol error"
669        );
670        assert_eq!(parsed["result"]["isError"], true);
671        assert!(
672            parsed["result"]["content"][0]["text"]
673                .as_str()
674                .unwrap()
675                .contains("kaboom")
676        );
677    }
678
679    #[test]
680    fn dispatch_prompts_list_and_get() {
681        let mut server = McpServer::new("t", "1.0");
682        server.register_prompt(crate::mcp::schema::PromptDescription {
683            name: "greet".into(),
684            description: Some("Greet".into()),
685            arguments: Vec::new(),
686        });
687        server.set_prompt_handler("greet", |_| {
688            Ok(serde_json::json!({
689                "messages": [{ "role": "user", "content": { "type": "text", "text": "hi" } }]
690            }))
691        });
692        let dispatcher = JsonRpcDispatcher::new(&server);
693
694        let list = dispatcher
695            .dispatch(r#"{"jsonrpc":"2.0","id":1,"method":"prompts/list"}"#)
696            .unwrap();
697        let list: serde_json::Value = serde_json::from_str(&list).unwrap();
698        assert_eq!(list["result"]["prompts"][0]["name"], "greet");
699
700        let got = dispatcher
701            .dispatch(r#"{"jsonrpc":"2.0","id":2,"method":"prompts/get","params":{"name":"greet","arguments":{}}}"#)
702            .unwrap();
703        let got: serde_json::Value = serde_json::from_str(&got).unwrap();
704        assert_eq!(got["result"]["messages"][0]["content"]["text"], "hi");
705    }
706
707    #[test]
708    fn dispatch_resource_templates_list_is_empty() {
709        let server = test_server();
710        let dispatcher = JsonRpcDispatcher::new(&server);
711        let resp = dispatcher
712            .dispatch(r#"{"jsonrpc":"2.0","id":1,"method":"resources/templates/list"}"#)
713            .unwrap();
714        let parsed: serde_json::Value = serde_json::from_str(&resp).unwrap();
715        assert!(parsed["result"]["resourceTemplates"].is_array());
716    }
717
718    #[test]
719    fn notification_without_id_gets_no_response() {
720        let server = test_server();
721        let dispatcher = JsonRpcDispatcher::new(&server);
722        // `notifications/initialized` is a notification (no id) → no response.
723        assert!(
724            dispatcher
725                .dispatch(r#"{"jsonrpc":"2.0","method":"notifications/initialized"}"#)
726                .is_none()
727        );
728    }
729
730    #[test]
731    fn dispatch_batch_request() {
732        let server = test_server();
733        let dispatcher = JsonRpcDispatcher::new(&server);
734        let batch = r#"[
735            {"jsonrpc":"2.0","id":1,"method":"initialize","params":{}},
736            {"jsonrpc":"2.0","id":2,"method":"tools/list"}
737        ]"#;
738        let resp = dispatcher.dispatch(batch).unwrap();
739        let parsed: serde_json::Value = serde_json::from_str(&resp).unwrap();
740        let arr = parsed.as_array().unwrap();
741        assert_eq!(arr.len(), 2);
742        // First response: initialize
743        assert_eq!(arr[0]["result"]["serverInfo"]["name"], "test-server");
744        // Second response: tools/list
745        assert_eq!(arr[1]["result"]["tools"].as_array().unwrap().len(), 1);
746    }
747
748    #[test]
749    fn dispatch_batch_with_error() {
750        let server = test_server();
751        let dispatcher = JsonRpcDispatcher::new(&server);
752        let batch = r#"[
753            {"jsonrpc":"2.0","id":1,"method":"initialize","params":{}},
754            {"jsonrpc":"2.0","id":2,"method":"nonexistent"}
755        ]"#;
756        let resp = dispatcher.dispatch(batch).unwrap();
757        let parsed: serde_json::Value = serde_json::from_str(&resp).unwrap();
758        let arr = parsed.as_array().unwrap();
759        assert_eq!(arr[0]["result"]["serverInfo"]["name"], "test-server");
760        assert_eq!(arr[1]["error"]["code"], -32601);
761    }
762
763    #[test]
764    fn dispatch_authenticated_passes_with_no_auth_configured() {
765        let server = test_server();
766        let dispatcher = JsonRpcDispatcher::new(&server);
767        let req = r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}"#;
768        let resp = dispatcher.dispatch_authenticated(req, None).unwrap();
769        let parsed: serde_json::Value = serde_json::from_str(&resp).unwrap();
770        assert!(parsed["result"]["serverInfo"].is_object());
771    }
772
773    #[test]
774    fn dispatch_authenticated_rejects_missing_token() {
775        let server = McpServer::new("secured", "1.0").with_bearer_auth("secret-token");
776        let dispatcher = JsonRpcDispatcher::new(&server);
777        let req = r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}"#;
778        let resp = dispatcher.dispatch_authenticated(req, None).unwrap();
779        let parsed: serde_json::Value = serde_json::from_str(&resp).unwrap();
780        assert_eq!(parsed["error"]["code"], -32001);
781    }
782
783    #[test]
784    fn dispatch_authenticated_rejects_wrong_token() {
785        let server = McpServer::new("secured", "1.0").with_bearer_auth("correct-token");
786        let dispatcher = JsonRpcDispatcher::new(&server);
787        let req = r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}"#;
788        let resp = dispatcher
789            .dispatch_authenticated(req, Some("Bearer wrong-token"))
790            .unwrap();
791        let parsed: serde_json::Value = serde_json::from_str(&resp).unwrap();
792        assert_eq!(parsed["error"]["code"], -32001);
793    }
794
795    #[test]
796    fn dispatch_authenticated_passes_with_correct_token() {
797        let server = McpServer::new("secured", "1.0").with_bearer_auth("correct-token");
798        let dispatcher = JsonRpcDispatcher::new(&server);
799        let req = r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}"#;
800        let resp = dispatcher
801            .dispatch_authenticated(req, Some("Bearer correct-token"))
802            .unwrap();
803        let parsed: serde_json::Value = serde_json::from_str(&resp).unwrap();
804        assert!(parsed["result"]["serverInfo"].is_object());
805    }
806}
llm_kernel/mcp/transport.rs

llm_kernel/mcp/
transport.rs
