Skip to main content

forge_sdk/mcp/
server.rs

1//! The in-process MCP server.
2//!
3//! Holds a registry of [`Tool`] implementations and dispatches JSON-RPC
4//! requests to them. Transport is OUT of scope here - this type is a pure
5//! request/response router. The `mcp::orchestration` module wires `dispatch`
6//! calls to the `mcp_message` control-request arm of `Client::handle_control`.
7
8use 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/// What the client's requested version resolved to. Kept distinct from the
20/// answered string because an unsupported version and a missing one are
21/// different events: `protocolVersion` is required, so the latter is a broken
22/// client rather than a negotiation outcome.
23#[derive(Debug, PartialEq, Eq)]
24enum Negotiated<'a> {
25    /// A version we speak.
26    Agreed(&'static str),
27    /// A version string we do not speak.
28    Unsupported(&'a str),
29    /// Absent, or present but not a string.
30    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
45/// Echo the client's requested protocol version when we speak it, else answer
46/// with our latest - the spec's SHOULD for an unsupported request.
47fn 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/// A fully-constructed MCP server. Clone is cheap (just bumps Arcs).
62#[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    /// Whether any registered tool's name starts with `prefix`, so a
83    /// host can introspect which tool group it registered.
84    pub fn has_tool_prefix(&self, prefix: &str) -> bool {
85        self.tools.keys().any(|name| name.starts_with(prefix))
86    }
87
88    /// The names this server carries, in the order `tools/list` answers
89    /// them, so a host can record the surface it hands a child.
90    pub fn tool_names(&self) -> Vec<&str> {
91        self.tools.keys().map(String::as_str).collect()
92    }
93
94    /// Every registered tool's name and input schema, so a test can pin the
95    /// surface a family advertises. Test-only.
96    #[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    /// Dispatch one JSON-RPC request to the appropriate handler.
102    ///
103    /// Returns `None` for JSON-RPC notifications (no `id`) - notifications
104    /// must not elicit a response per the spec. Returns `Some(response)` for
105    /// every id-bearing request, success or error.
106    ///
107    /// This method is pure (modulo tool side-effects): no I/O, no
108    /// timeouts, no mutation of `self`. Thread-safe; cheap to call
109    /// concurrently.
110    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
180/// Builder for [`McpServer`].
181pub 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    /// Start a new builder.
201    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    /// The text the client delivers to the model once at session start:
211    /// how to use this server, in the protocol's own `instructions` field.
212    pub fn instructions(mut self, text: &str) -> Self {
213        self.instructions = Some(text.to_owned());
214        self
215    }
216
217    /// Register a tool. A duplicate name replaces the earlier registration -
218    /// five modules chain onto one builder, so a collision would otherwise
219    /// drop a tool with nothing to show for it.
220    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    /// Finalise into a runnable server.
229    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    /// A server with none omits the field, so a client that never learned
265    /// it sees exactly the old shape.
266    #[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    /// Both answer with our latest on the wire, so only the classifier can
280    /// tell them apart - an unsupported version is a negotiation outcome, a
281    /// missing one is a broken client.
282    #[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(&params)), 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    /// **The record of a surface is the surface itself.** A spawn logs
337    /// `tool_names()` as the list it handed its child, so a name the record
338    /// drops - or invents - would answer a later "tool not found" wrongly.
339    /// The sorting is pinned too: it is the order `tools/list` answers in.
340    #[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}