1use std::collections::BTreeMap;
9use std::sync::Arc;
10
11use tracing::{debug, warn};
12
13use crate::mcp::protocol::{
14 JsonRpcRequest, JsonRpcResponse, JsonRpcResult, LATEST_PROTOCOL_VERSION, McpError,
15 SUPPORTED_PROTOCOL_VERSIONS, ServerInfo, ToolDescription,
16};
17use crate::mcp::tool::{Tool, ToolInput, ToolOutput};
18
19#[derive(Debug, PartialEq, Eq)]
24enum Negotiated<'a> {
25 Agreed(&'static str),
27 Unsupported(&'a str),
29 Malformed,
31}
32
33fn classify_protocol_version(params: Option<&serde_json::Value>) -> Negotiated<'_> {
34 let Some(want) =
35 params.and_then(|p| p.get("protocolVersion")).and_then(serde_json::Value::as_str)
36 else {
37 return Negotiated::Malformed;
38 };
39 SUPPORTED_PROTOCOL_VERSIONS
40 .into_iter()
41 .find(|v| *v == want)
42 .map_or(Negotiated::Unsupported(want), Negotiated::Agreed)
43}
44
45fn negotiate_protocol_version(params: Option<&serde_json::Value>) -> &'static str {
48 match classify_protocol_version(params) {
49 Negotiated::Agreed(v) => v,
50 Negotiated::Unsupported(want) => {
51 debug!(requested = %want, "mcp server: unsupported protocol version, answering latest");
52 LATEST_PROTOCOL_VERSION
53 }
54 Negotiated::Malformed => {
55 warn!("mcp server: initialize has no string protocolVersion, answering latest");
56 LATEST_PROTOCOL_VERSION
57 }
58 }
59}
60
61#[derive(Clone)]
63pub struct McpServer {
64 name: String,
65 version: String,
66 instructions: Option<String>,
67 tools: BTreeMap<String, Arc<dyn Tool>>,
68}
69
70impl std::fmt::Debug for McpServer {
71 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
72 f.debug_struct("McpServer")
73 .field("name", &self.name)
74 .field("version", &self.version)
75 .field("instructions", &self.instructions.is_some())
76 .field("tools", &self.tools.keys().collect::<Vec<_>>())
77 .finish()
78 }
79}
80
81impl McpServer {
82 pub fn has_tool_prefix(&self, prefix: &str) -> bool {
85 self.tools.keys().any(|name| name.starts_with(prefix))
86 }
87
88 pub fn tool_names(&self) -> Vec<&str> {
91 self.tools.keys().map(String::as_str).collect()
92 }
93
94 #[cfg(any(test, feature = "testing"))]
97 pub fn tool_schemas(&self) -> Vec<(String, serde_json::Value)> {
98 self.tools.iter().map(|(name, tool)| (name.clone(), tool.input_schema())).collect()
99 }
100
101 pub async fn dispatch(&self, req: &JsonRpcRequest) -> Option<JsonRpcResponse> {
111 let Some(id) = req.id.clone() else {
112 debug!(method = %req.method, "mcp server: notification ignored");
113 return None;
114 };
115 let result: Result<JsonRpcResult, McpError> = match req.method.as_str() {
116 "initialize" => Ok(JsonRpcResult::Initialize {
117 protocol_version: negotiate_protocol_version(req.params.as_ref()).into(),
118 capabilities: serde_json::json!({"tools": {"listChanged": false}}),
119 server_info: ServerInfo { name: self.name.clone(), version: self.version.clone() },
120 instructions: self.instructions.clone(),
121 }),
122 "ping" => Ok(JsonRpcResult::Empty {}),
123 "tools/list" => {
124 let tools: Vec<ToolDescription> = self
125 .tools
126 .values()
127 .map(|t| ToolDescription {
128 name: t.name().into(),
129 description: t.description().into(),
130 input_schema: t.input_schema(),
131 })
132 .collect();
133 Ok(JsonRpcResult::ToolsList { tools })
134 }
135 "tools/call" => match self.call_tool(req.params.as_ref()).await {
136 Ok(output) => Ok(JsonRpcResult::ToolsCall {
137 content: output.to_mcp_content(),
138 is_error: output.is_error,
139 }),
140 Err(err) => Err(err),
141 },
142 other => Err(McpError {
143 code: -32601,
144 message: format!("method not found: {other}"),
145 data: None,
146 }),
147 };
148
149 Some(match result {
150 Ok(r) => JsonRpcResponse::success(id, r),
151 Err(e) => JsonRpcResponse::error(id, e),
152 })
153 }
154
155 async fn call_tool(&self, params: Option<&serde_json::Value>) -> Result<ToolOutput, McpError> {
156 let params = params.ok_or_else(|| McpError {
157 code: -32602,
158 message: "tools/call requires params".into(),
159 data: None,
160 })?;
161 let name = params.get("name").and_then(|v| v.as_str()).ok_or_else(|| McpError {
162 code: -32602,
163 message: "tools/call params.name missing".into(),
164 data: None,
165 })?;
166 let tool = self.tools.get(name).ok_or_else(|| McpError {
167 code: -32602,
168 message: format!("unknown tool: {name}"),
169 data: None,
170 })?;
171 let args = params
172 .get("arguments")
173 .cloned()
174 .unwrap_or(serde_json::Value::Object(serde_json::Map::new()));
175 let input = ToolInput { value: args };
176 Ok(tool.call(input).await)
177 }
178}
179
180pub struct McpServerBuilder {
182 name: String,
183 version: String,
184 instructions: Option<String>,
185 tools: BTreeMap<String, Arc<dyn Tool>>,
186}
187
188impl std::fmt::Debug for McpServerBuilder {
189 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
190 f.debug_struct("McpServerBuilder")
191 .field("name", &self.name)
192 .field("version", &self.version)
193 .field("instructions", &self.instructions.is_some())
194 .field("tools", &self.tools.keys().collect::<Vec<_>>())
195 .finish()
196 }
197}
198
199impl McpServerBuilder {
200 pub fn new(name: impl Into<String>, version: impl Into<String>) -> Self {
202 Self {
203 name: name.into(),
204 version: version.into(),
205 instructions: None,
206 tools: BTreeMap::new(),
207 }
208 }
209
210 pub fn instructions(mut self, text: &str) -> Self {
213 self.instructions = Some(text.to_owned());
214 self
215 }
216
217 pub fn tool<T: Tool + 'static>(mut self, tool: T) -> Self {
221 let name = tool.name().to_string();
222 if self.tools.insert(name.clone(), Arc::new(tool)).is_some() {
223 warn!(%name, server = %self.name, "mcp tool registered twice; keeping the later one");
224 }
225 self
226 }
227
228 pub fn build(self) -> McpServer {
230 McpServer {
231 name: self.name,
232 version: self.version,
233 instructions: self.instructions,
234 tools: self.tools,
235 }
236 }
237}
238
239#[cfg(test)]
240mod tests_initialize {
241 use super::{JsonRpcRequest, McpServerBuilder};
242 use serde_json::json;
243
244 async fn initialize_result(server: &super::McpServer) -> serde_json::Value {
245 let request: JsonRpcRequest = serde_json::from_value(json!({
246 "jsonrpc": "2.0",
247 "id": 1,
248 "method": "initialize",
249 "params": { "protocolVersion": "2025-06-18" },
250 }))
251 .expect("an initialize request");
252 let response = server.dispatch(&request).await.expect("an id-bearing request is answered");
253 serde_json::to_value(response.result.expect("initialize succeeds")).expect("serializes")
254 }
255
256 #[tokio::test]
257 async fn an_initialize_result_carries_the_servers_instructions_when_set() {
258 let server =
259 McpServerBuilder::new("forge", "1.2.3").instructions("claim before working").build();
260 let result = initialize_result(&server).await;
261 assert_eq!(result["instructions"], "claim before working");
262 }
263
264 #[tokio::test]
267 async fn no_instructions_leaves_the_field_out_of_the_json() {
268 let server = McpServerBuilder::new("forge", "1.2.3").build();
269 let result = initialize_result(&server).await;
270 assert!(result.get("instructions").is_none(), "absent, not null: {result}");
271 }
272}
273
274#[cfg(test)]
275mod tests_negotiation {
276 use super::{Negotiated, classify_protocol_version};
277 use serde_json::json;
278
279 #[test]
283 fn unsupported_and_malformed_are_distinct() {
284 let unsupported = json!({"protocolVersion": "1999-01-01"});
285 assert_eq!(
286 classify_protocol_version(Some(&unsupported)),
287 Negotiated::Unsupported("1999-01-01")
288 );
289
290 for malformed in [json!({"protocolVersion": 5}), json!({"capabilities": {}}), json!(null)] {
291 assert_eq!(
292 classify_protocol_version(Some(&malformed)),
293 Negotiated::Malformed,
294 "params: {malformed}"
295 );
296 }
297 assert_eq!(classify_protocol_version(None), Negotiated::Malformed);
298 }
299
300 #[test]
301 fn every_supported_version_is_agreed() {
302 for v in super::SUPPORTED_PROTOCOL_VERSIONS {
303 let params = json!({ "protocolVersion": v });
304 assert_eq!(classify_protocol_version(Some(¶ms)), Negotiated::Agreed(v));
305 }
306 }
307}
308
309#[cfg(test)]
310mod tests_tool_names {
311 use super::{JsonRpcRequest, McpServerBuilder};
312 use crate::mcp::tool::{Tool, ToolInput, ToolOutput};
313 use serde_json::{Value, json};
314
315 struct Named(&'static str);
316
317 #[async_trait::async_trait]
318 impl Tool for Named {
319 fn name(&self) -> &'static str {
320 self.0
321 }
322
323 fn description(&self) -> &'static str {
324 "a tool the record test registers"
325 }
326
327 fn input_schema(&self) -> Value {
328 json!({"type": "object", "properties": {}, "additionalProperties": false})
329 }
330
331 async fn call(&self, _input: ToolInput) -> ToolOutput {
332 ToolOutput::text("ok")
333 }
334 }
335
336 #[tokio::test]
341 async fn the_recorded_names_are_exactly_what_tools_list_answers() {
342 let server = McpServerBuilder::new("forge", "1.2.3")
343 .tool(Named("tasks__list"))
344 .tool(Named("browser_form_state"))
345 .tool(Named("agents__list"))
346 .build();
347
348 let request: JsonRpcRequest = serde_json::from_value(json!({
349 "jsonrpc": "2.0",
350 "id": 1,
351 "method": "tools/list",
352 }))
353 .expect("a tools/list request");
354 let response = server.dispatch(&request).await.expect("an id-bearing request is answered");
355 let result = serde_json::to_value(response.result.expect("tools/list succeeds"))
356 .expect("serializes");
357 let listed: Vec<String> = result["tools"]
358 .as_array()
359 .expect("a tools array")
360 .iter()
361 .map(|tool| tool["name"].as_str().expect("a tool name").to_owned())
362 .collect();
363
364 assert_eq!(
365 server.tool_names(),
366 listed,
367 "what a spawn records is exactly the set tools/list answers, in its order",
368 );
369 }
370}