Skip to main content

claude_code_sdk_rust/internal/
sdk_mcp.rs

1use std::collections::HashMap;
2
3use crate::mcp::{MCPContent, SimpleMCPServer};
4use serde_json::Value;
5
6pub fn answer_mcp_message(
7    servers: &HashMap<String, SimpleMCPServer>,
8    server_name: &str,
9    message: &Value,
10) -> Value {
11    let Some(server) = servers.get(server_name) else {
12        return jsonrpc_error(
13            message.get("id").cloned(),
14            -32601,
15            &format!("Server '{server_name}' not found"),
16        );
17    };
18
19    let method = message.get("method").and_then(|v| v.as_str()).unwrap_or("");
20    match method {
21        "initialize" => serde_json::json!({
22            "jsonrpc": "2.0",
23            "id": message.get("id").cloned().unwrap_or(Value::Null),
24            "result": {
25                "protocolVersion": "2024-11-05",
26                "capabilities": {"tools": {}},
27                "serverInfo": {
28                    "name": server.name(),
29                    "version": server.version(),
30                },
31            },
32        }),
33        "tools/list" => serde_json::json!({
34            "jsonrpc": "2.0",
35            "id": message.get("id").cloned().unwrap_or(Value::Null),
36            "result": {
37                "tools": server
38                    .list_tools()
39                    .into_iter()
40                    .map(tool_to_wire)
41                    .collect::<Vec<_>>(),
42            },
43        }),
44        "tools/call" => call_tool(server, message),
45        "notifications/initialized" => serde_json::json!({
46            "jsonrpc": "2.0",
47            "result": {},
48        }),
49        _ => jsonrpc_error(
50            message.get("id").cloned(),
51            -32601,
52            &format!("Method '{method}' not found"),
53        ),
54    }
55}
56
57fn call_tool(server: &SimpleMCPServer, message: &Value) -> Value {
58    let id = message.get("id").cloned().unwrap_or(Value::Null);
59    let params = message.get("params").and_then(|v| v.as_object());
60    let name = params
61        .and_then(|params| params.get("name"))
62        .and_then(|v| v.as_str())
63        .unwrap_or("");
64    let arguments = params
65        .and_then(|params| params.get("arguments"))
66        .cloned()
67        .unwrap_or_else(|| serde_json::json!({}));
68
69    match server.call_tool(name, arguments) {
70        Ok(content) => serde_json::json!({
71            "jsonrpc": "2.0",
72            "id": id,
73            "result": {
74                "content": content.into_iter().map(content_to_wire).collect::<Vec<_>>(),
75            },
76        }),
77        Err(error) => serde_json::json!({
78            "jsonrpc": "2.0",
79            "id": id,
80            "result": {
81                "content": [{"type": "text", "text": error}],
82                "isError": true,
83            },
84        }),
85    }
86}
87
88fn tool_to_wire(tool: &crate::mcp::MCPTool) -> Value {
89    let mut value = serde_json::json!({
90        "name": tool.name,
91        "description": tool.description,
92        "inputSchema": tool.input_schema,
93    });
94    if let Some(annotations) = &tool.annotations {
95        value["annotations"] = serde_json::json!({
96            "title": annotations.title,
97            "readOnlyHint": annotations.read_only_hint,
98            "destructiveHint": annotations.destructive_hint,
99            "idempotentHint": annotations.idempotent_hint,
100            "openWorldHint": annotations.open_world_hint,
101        });
102        if let Some(max_size) = annotations.max_result_size_chars {
103            value["_meta"] = serde_json::json!({
104                "anthropic/maxResultSizeChars": max_size,
105            });
106        }
107    }
108    value
109}
110
111fn content_to_wire(content: MCPContent) -> Value {
112    match content {
113        MCPContent::Text { text } => serde_json::json!({"type": "text", "text": text}),
114        MCPContent::Image { data, mime_type } => {
115            serde_json::json!({"type": "image", "data": data, "mimeType": mime_type})
116        }
117        MCPContent::Audio { data, mime_type } => {
118            serde_json::json!({"type": "audio", "data": data, "mimeType": mime_type})
119        }
120        MCPContent::ResourceLink {
121            uri,
122            name,
123            description,
124            mime_type,
125        } => serde_json::json!({
126            "type": "resource_link",
127            "uri": uri,
128            "name": name,
129            "description": description,
130            "mimeType": mime_type,
131        }),
132        MCPContent::Resource {
133            uri,
134            mime_type,
135            text,
136        } => serde_json::json!({
137            "type": "resource",
138            "resource": {
139                "uri": uri,
140                "mimeType": mime_type,
141                "text": text,
142            },
143        }),
144    }
145}
146
147fn jsonrpc_error(id: Option<Value>, code: i32, message: &str) -> Value {
148    serde_json::json!({
149        "jsonrpc": "2.0",
150        "id": id.unwrap_or(Value::Null),
151        "error": {
152            "code": code,
153            "message": message,
154        },
155    })
156}
157
158#[cfg(test)]
159mod tests {
160    use super::*;
161    use crate::mcp::{MCPContent, MCPTool};
162
163    fn server() -> SimpleMCPServer {
164        let mut server = SimpleMCPServer::new("greeter");
165        server.register_tool(
166            MCPTool {
167                name: "greet".to_string(),
168                description: "Greet someone".to_string(),
169                input_schema: serde_json::json!({"type": "object"}),
170                annotations: None,
171            },
172            |input| {
173                let name = input
174                    .get("name")
175                    .and_then(|v| v.as_str())
176                    .unwrap_or("there");
177                Ok(vec![MCPContent::Text {
178                    text: format!("Hi {name}"),
179                }])
180            },
181        );
182        server
183    }
184
185    #[test]
186    fn answers_tools_list() {
187        let servers = HashMap::from([("greeter".to_string(), server())]);
188        let response = answer_mcp_message(
189            &servers,
190            "greeter",
191            &serde_json::json!({"jsonrpc": "2.0", "id": 1, "method": "tools/list"}),
192        );
193
194        assert_eq!(response["result"]["tools"][0]["name"], "greet");
195        assert_eq!(
196            response["result"]["tools"][0]["inputSchema"]["type"],
197            "object"
198        );
199    }
200
201    #[test]
202    fn answers_tools_call() {
203        let servers = HashMap::from([("greeter".to_string(), server())]);
204        let response = answer_mcp_message(
205            &servers,
206            "greeter",
207            &serde_json::json!({
208                "jsonrpc": "2.0",
209                "id": 2,
210                "method": "tools/call",
211                "params": {"name": "greet", "arguments": {"name": "Ada"}}
212            }),
213        );
214
215        assert_eq!(response["result"]["content"][0]["text"], "Hi Ada");
216    }
217}