agent_client_protocol/mcp_server/
registry.rs

1//! Runtime-neutral MCP tool registration and dispatch.
2
3use std::{collections::HashSet, sync::Arc};
4
5use futures::future::BoxFuture;
6use rustc_hash::FxHashMap;
7use schemars::{JsonSchema, generate::SchemaSettings};
8use serde_json::{Map, Value};
9
10use crate::{Error, Role};
11
12use super::{McpConnectionTo, McpTool};
13
14/// JSON Schema object used to describe MCP tool inputs and outputs.
15pub type McpToolSchema = Map<String, Value>;
16
17/// Tracks which tools are enabled.
18///
19/// - `DenyList`: All tools enabled except those in the set (default)
20/// - `AllowList`: Only tools in the set are enabled
21#[derive(Clone, Debug)]
22enum EnabledTools {
23    /// All tools enabled except those in the deny set.
24    DenyList(HashSet<String>),
25    /// Only tools in the allow set are enabled.
26    AllowList(HashSet<String>),
27}
28
29impl Default for EnabledTools {
30    fn default() -> Self {
31        EnabledTools::DenyList(HashSet::new())
32    }
33}
34
35impl EnabledTools {
36    /// Check if a tool is enabled.
37    fn is_enabled(&self, name: &str) -> bool {
38        match self {
39            EnabledTools::DenyList(deny) => !deny.contains(name),
40            EnabledTools::AllowList(allow) => allow.contains(name),
41        }
42    }
43}
44
45/// [`McpToolRegistry::enable_tool`] or [`McpToolRegistry::disable_tool`] named
46/// a tool that is not registered.
47///
48/// This is a local configuration error, usually a typo, not a JSON-RPC error.
49/// Converting it into an [`Error`] with `?` produces an internal error
50/// (`-32603`) whose data names the tool.
51#[derive(Clone, Debug, PartialEq, Eq)]
52#[non_exhaustive]
53pub struct UnknownToolError {
54    name: String,
55}
56
57impl UnknownToolError {
58    fn new(name: &str) -> Self {
59        Self {
60            name: name.to_string(),
61        }
62    }
63
64    /// The tool name that is not registered.
65    #[must_use]
66    pub fn name(&self) -> &str {
67        &self.name
68    }
69}
70
71impl std::fmt::Display for UnknownToolError {
72    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
73        write!(f, "unknown MCP tool `{}`", self.name)
74    }
75}
76
77impl std::error::Error for UnknownToolError {}
78
79impl From<UnknownToolError> for Error {
80    fn from(error: UnknownToolError) -> Self {
81        Error::internal_error().data(error.to_string())
82    }
83}
84
85/// Runtime-neutral metadata for an MCP tool.
86#[derive(Clone, Debug)]
87pub struct McpToolMetadata {
88    name: String,
89    title: Option<String>,
90    description: String,
91    input_schema: Arc<McpToolSchema>,
92    output_schema: Option<Arc<McpToolSchema>>,
93}
94
95impl McpToolMetadata {
96    fn from_tool<R: Role, M: McpTool<R>>(tool: &M) -> Self {
97        Self {
98            name: tool.name(),
99            title: tool.title(),
100            description: tool.description(),
101            input_schema: schema_for_type::<M::Input>(),
102            output_schema: schema_for_output::<M::Output>(),
103        }
104    }
105
106    /// The tool name.
107    #[must_use]
108    pub fn name(&self) -> &str {
109        &self.name
110    }
111
112    /// A human-readable title for the tool.
113    #[must_use]
114    pub fn title(&self) -> Option<&str> {
115        self.title.as_deref()
116    }
117
118    /// A description of what the tool does.
119    #[must_use]
120    pub fn description(&self) -> &str {
121        &self.description
122    }
123
124    /// JSON Schema object defining the expected parameters for the tool.
125    #[must_use]
126    pub fn input_schema(&self) -> &Arc<McpToolSchema> {
127        &self.input_schema
128    }
129
130    /// Optional JSON Schema object defining the structure of the tool's output.
131    #[must_use]
132    pub fn output_schema(&self) -> Option<&Arc<McpToolSchema>> {
133        self.output_schema.as_ref()
134    }
135}
136
137/// A registered MCP tool that can be dispatched with erased JSON values.
138pub struct RegisteredMcpTool<Counterpart: Role> {
139    metadata: McpToolMetadata,
140    tool: Arc<dyn ErasedMcpTool<Counterpart>>,
141}
142
143impl<Counterpart: Role + std::fmt::Debug> std::fmt::Debug for RegisteredMcpTool<Counterpart> {
144    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
145        f.debug_struct("RegisteredMcpTool")
146            .field("metadata", &self.metadata)
147            .field("has_structured_output", &self.has_structured_output())
148            .finish_non_exhaustive()
149    }
150}
151
152impl<Counterpart: Role> RegisteredMcpTool<Counterpart> {
153    fn new(tool: impl McpTool<Counterpart> + 'static) -> Self {
154        let metadata = McpToolMetadata::from_tool(&tool);
155        Self {
156            metadata,
157            tool: make_erased_mcp_tool(tool),
158        }
159    }
160
161    /// Tool metadata.
162    #[must_use]
163    pub fn metadata(&self) -> &McpToolMetadata {
164        &self.metadata
165    }
166
167    /// The tool name.
168    #[must_use]
169    pub fn name(&self) -> &str {
170        self.metadata.name()
171    }
172
173    /// Whether the tool returns structured output.
174    #[must_use]
175    pub fn has_structured_output(&self) -> bool {
176        self.metadata.output_schema().is_some()
177    }
178
179    /// Invoke the registered tool using JSON input and output values.
180    ///
181    /// # Errors
182    ///
183    /// Returns an invalid params error (`-32602`) if `input` does not
184    /// deserialize into the tool's input type, an internal error (`-32603`) if
185    /// the output does not serialize, and otherwise the tool's own error.
186    pub fn call_tool(
187        &self,
188        input: Value,
189        connection: McpConnectionTo<Counterpart>,
190    ) -> BoxFuture<'_, Result<Value, Error>> {
191        self.tool.call_tool(input, connection)
192    }
193}
194
195/// Runtime-neutral registry for MCP tools.
196#[derive(Debug)]
197pub struct McpToolRegistry<Counterpart: Role> {
198    instructions: Option<String>,
199    tool_indices: FxHashMap<String, usize>,
200    tools: Vec<RegisteredMcpTool<Counterpart>>,
201    enabled_tools: EnabledTools,
202}
203
204impl<Counterpart: Role> Default for McpToolRegistry<Counterpart> {
205    fn default() -> Self {
206        Self {
207            instructions: None,
208            tool_indices: FxHashMap::default(),
209            tools: Vec::new(),
210            enabled_tools: EnabledTools::default(),
211        }
212    }
213}
214
215impl<Counterpart: Role> McpToolRegistry<Counterpart> {
216    /// Set the server instructions that are provided to the client.
217    pub fn set_instructions(&mut self, instructions: impl ToString) {
218        self.instructions = Some(instructions.to_string());
219    }
220
221    /// Server instructions provided to the client.
222    #[must_use]
223    pub fn instructions(&self) -> Option<&str> {
224        self.instructions.as_deref()
225    }
226
227    /// Register a tool.
228    pub fn register_tool(&mut self, tool: impl McpTool<Counterpart> + 'static) {
229        let registered_tool = RegisteredMcpTool::new(tool);
230        let name = registered_tool.name().to_string();
231
232        if let Some(&index) = self.tool_indices.get(&name) {
233            self.tools[index] = registered_tool;
234        } else {
235            self.tool_indices.insert(name, self.tools.len());
236            self.tools.push(registered_tool);
237        }
238    }
239
240    /// Return all registered tools in registration order.
241    pub fn tools(&self) -> impl Iterator<Item = &RegisteredMcpTool<Counterpart>> {
242        self.tools.iter()
243    }
244
245    /// Return enabled registered tools in registration order.
246    pub fn enabled_tools(&self) -> impl Iterator<Item = &RegisteredMcpTool<Counterpart>> {
247        self.tools
248            .iter()
249            .filter(|tool| self.enabled_tools.is_enabled(tool.name()))
250    }
251
252    /// Return a registered tool by name, even if it is disabled.
253    #[must_use]
254    pub fn tool(&self, name: &str) -> Option<&RegisteredMcpTool<Counterpart>> {
255        self.tool_indices
256            .get(name)
257            .and_then(|&index| self.tools.get(index))
258    }
259
260    /// Return an enabled tool by name.
261    #[must_use]
262    pub fn enabled_tool(&self, name: &str) -> Option<&RegisteredMcpTool<Counterpart>> {
263        self.tool(name)
264            .filter(|tool| self.enabled_tools.is_enabled(tool.name()))
265    }
266
267    /// Check whether a tool is registered.
268    #[must_use]
269    pub fn contains_tool(&self, name: &str) -> bool {
270        self.tool_indices.contains_key(name)
271    }
272
273    /// Disable all tools. After calling this, only tools explicitly enabled
274    /// with [`enable_tool`](Self::enable_tool) will be available.
275    pub fn disable_all_tools(&mut self) {
276        self.enabled_tools = EnabledTools::AllowList(HashSet::new());
277    }
278
279    /// Enable all tools. After calling this, all tools will be available
280    /// except those explicitly disabled with [`disable_tool`](Self::disable_tool).
281    pub fn enable_all_tools(&mut self) {
282        self.enabled_tools = EnabledTools::DenyList(HashSet::new());
283    }
284
285    /// Disable a specific tool by name.
286    ///
287    /// # Errors
288    ///
289    /// Returns [`UnknownToolError`] if no tool with this name is registered.
290    pub fn disable_tool(&mut self, name: &str) -> Result<(), UnknownToolError> {
291        if !self.contains_tool(name) {
292            return Err(UnknownToolError::new(name));
293        }
294        match &mut self.enabled_tools {
295            EnabledTools::DenyList(deny) => {
296                deny.insert(name.to_string());
297            }
298            EnabledTools::AllowList(allow) => {
299                allow.remove(name);
300            }
301        }
302        Ok(())
303    }
304
305    /// Enable a specific tool by name.
306    ///
307    /// # Errors
308    ///
309    /// Returns [`UnknownToolError`] if no tool with this name is registered.
310    pub fn enable_tool(&mut self, name: &str) -> Result<(), UnknownToolError> {
311        if !self.contains_tool(name) {
312            return Err(UnknownToolError::new(name));
313        }
314        match &mut self.enabled_tools {
315            EnabledTools::DenyList(deny) => {
316                deny.remove(name);
317            }
318            EnabledTools::AllowList(allow) => {
319                allow.insert(name.to_string());
320            }
321        }
322        Ok(())
323    }
324}
325
326/// Erased version of the MCP tool trait that is dyn-compatible.
327trait ErasedMcpTool<Counterpart: Role>: Send + Sync {
328    fn call_tool(
329        &self,
330        input: Value,
331        connection: McpConnectionTo<Counterpart>,
332    ) -> BoxFuture<'_, Result<Value, Error>>;
333}
334
335fn make_erased_mcp_tool<R, M>(tool: M) -> Arc<dyn ErasedMcpTool<R>>
336where
337    R: Role,
338    M: McpTool<R> + 'static,
339{
340    struct ErasedMcpToolImpl<M> {
341        tool: M,
342    }
343
344    impl<R, M> ErasedMcpTool<R> for ErasedMcpToolImpl<M>
345    where
346        R: Role,
347        M: McpTool<R>,
348    {
349        fn call_tool(
350            &self,
351            input: Value,
352            context: McpConnectionTo<R>,
353        ) -> BoxFuture<'_, Result<Value, Error>> {
354            Box::pin(async move {
355                let input = serde_json::from_value(input)
356                    .map_err(|error| Error::invalid_params().data(error.to_string()))?;
357                serde_json::to_value(self.tool.call_tool(input, context).await?)
358                    .map_err(|error| Error::internal_error().data(error.to_string()))
359            })
360        }
361    }
362
363    Arc::new(ErasedMcpToolImpl { tool })
364}
365
366fn schema_for_type<T: JsonSchema>() -> Arc<McpToolSchema> {
367    let settings = SchemaSettings::draft2020_12();
368    let generator = settings.into_generator();
369    let schema = generator.into_root_schema_for::<T>();
370    let object = serde_json::to_value(schema).expect("failed to serialize schema");
371    let Value::Object(object) = object else {
372        panic!("Schema serialization produced non-object value: expected JSON object");
373    };
374    Arc::new(object)
375}
376
377fn schema_for_output<T: JsonSchema>() -> Option<Arc<McpToolSchema>> {
378    let schema = schema_for_type::<T>();
379    match schema.get("type") {
380        Some(Value::String(t)) if t == "object" => Some(schema),
381        _ => None,
382    }
383}
384
385#[cfg(test)]
386mod tests {
387    use serde::Deserialize;
388
389    use super::*;
390    use crate::{Channel, mcp_server::McpConnectionContext, role::mcp};
391
392    #[derive(Deserialize, JsonSchema)]
393    struct EchoInput {
394        message: String,
395    }
396
397    struct Echo;
398
399    impl McpTool<mcp::Client> for Echo {
400        type Input = EchoInput;
401        type Output = String;
402
403        fn name(&self) -> String {
404            "echo".to_string()
405        }
406
407        fn description(&self) -> String {
408            "Echoes the message".to_string()
409        }
410
411        fn call_tool(
412            &self,
413            input: EchoInput,
414            _context: McpConnectionTo<mcp::Client>,
415        ) -> impl Future<Output = Result<String, Error>> + Send {
416            std::future::ready(Ok(input.message))
417        }
418    }
419
420    #[test]
421    fn call_tool_reports_bad_arguments_as_invalid_params() {
422        let (channel, _peer) = Channel::duplex();
423        futures::executor::block_on(mcp::Server.builder().connect_with(
424            channel,
425            async |connection| {
426                let context = McpConnectionTo {
427                    context: McpConnectionContext::Standalone,
428                    connection,
429                    #[cfg(feature = "unstable_mcp_over_acp")]
430                    cleanup: None,
431                };
432                let mut registry = McpToolRegistry::default();
433                registry.register_tool(Echo);
434                let tool = registry.tool("echo").expect("echo is registered");
435
436                let output = tool
437                    .call_tool(serde_json::json!({ "message": "hi" }), context.clone())
438                    .await?;
439                assert_eq!(output, Value::String("hi".to_string()));
440
441                let error = tool
442                    .call_tool(serde_json::json!({ "msg": "hi" }), context)
443                    .await
444                    .expect_err("mismatched arguments are rejected");
445                assert_eq!(error.code, Error::invalid_params().code);
446                Ok(())
447            },
448        ))
449        .unwrap();
450    }
451}
agent_client_protocol/mcp_server/registry.rs

agent_client_protocol/mcp_server/
registry.rs
