rmcp_openapi/
tool_registry.rs

1use crate::error::Error;
2use crate::spec::{Filters, Spec};
3use crate::tool::ToolMetadata;
4use std::collections::HashMap;
5
6/// Registry for managing dynamically generated MCP tools from `OpenAPI` operations
7#[derive(Debug, Clone)]
8pub struct ToolRegistry {
9    /// Map of tool name to tool metadata
10    tools: HashMap<String, ToolMetadata>,
11    /// Map of tool name to `OpenAPI` operation for runtime lookup
12    operations: HashMap<String, (oas3::spec::Operation, String, String)>,
13    /// Source spec for reference
14    spec: Option<Spec>,
15}
16
17impl ToolRegistry {
18    /// Create a new empty tool registry
19    #[must_use]
20    pub fn new() -> Self {
21        Self {
22            tools: HashMap::new(),
23            operations: HashMap::new(),
24            spec: None,
25        }
26    }
27
28    /// Register tools from a specification
29    ///
30    /// # Errors
31    ///
32    /// Returns an error if any tool fails to be generated or registered
33    pub fn register_from_spec(
34        &mut self,
35        spec: Spec,
36        filters: Option<&Filters>,
37        skip_tool_descriptions: bool,
38        skip_parameter_descriptions: bool,
39    ) -> Result<usize, Error> {
40        // Clear existing tools
41        self.clear();
42
43        // Convert operations to tool metadata
44        let tools_metadata =
45            spec.to_tool_metadata(filters, skip_tool_descriptions, skip_parameter_descriptions)?;
46        let mut registered_count = 0;
47
48        // Register each tool
49        for tool in tools_metadata {
50            // Find corresponding operation
51            if let Some((operation, method, path)) = spec.get_operation(&tool.name) {
52                self.register_tool(tool, (operation.clone(), method, path))?;
53                registered_count += 1;
54            }
55        }
56
57        // Store the spec
58        self.spec = Some(spec);
59
60        Ok(registered_count)
61    }
62
63    /// Register a single tool with its corresponding operation
64    ///
65    /// # Errors
66    ///
67    /// Returns an error if the tool metadata is invalid or the tool name already exists
68    pub fn register_tool(
69        &mut self,
70        tool: ToolMetadata,
71        operation: (oas3::spec::Operation, String, String),
72    ) -> Result<(), Error> {
73        let tool_name = tool.name.clone();
74
75        // Validate tool metadata
76        self.validate_tool(&tool)?;
77
78        // Store tool metadata and operation
79        self.tools.insert(tool_name.clone(), tool);
80        self.operations.insert(tool_name, operation);
81
82        Ok(())
83    }
84
85    /// Validate tool metadata
86    fn validate_tool(&self, tool: &ToolMetadata) -> Result<(), Error> {
87        if tool.name.is_empty() {
88            return Err(Error::ToolGeneration(
89                "Tool name cannot be empty".to_string(),
90            ));
91        }
92
93        if tool.method.is_empty() {
94            return Err(Error::ToolGeneration(
95                "Tool method cannot be empty".to_string(),
96            ));
97        }
98
99        if tool.path.is_empty() {
100            return Err(Error::ToolGeneration(
101                "Tool path cannot be empty".to_string(),
102            ));
103        }
104
105        // Validate that the tool name is unique
106        if self.tools.contains_key(&tool.name) {
107            return Err(Error::ToolGeneration(format!(
108                "Tool '{}' already exists",
109                tool.name
110            )));
111        }
112
113        Ok(())
114    }
115
116    /// Get tool metadata by name
117    #[must_use]
118    pub fn get_tool(&self, name: &str) -> Option<&ToolMetadata> {
119        self.tools.get(name)
120    }
121
122    /// Get operation by tool name
123    #[must_use]
124    pub fn get_operation(
125        &self,
126        tool_name: &str,
127    ) -> Option<&(oas3::spec::Operation, String, String)> {
128        self.operations.get(tool_name)
129    }
130
131    /// Get all tool names
132    #[must_use]
133    pub fn get_tool_names(&self) -> Vec<String> {
134        self.tools.keys().cloned().collect()
135    }
136
137    /// Get all tools
138    #[must_use]
139    pub fn get_all_tools(&self) -> Vec<&ToolMetadata> {
140        self.tools.values().collect()
141    }
142
143    /// Get number of registered tools
144    #[must_use]
145    pub fn tool_count(&self) -> usize {
146        self.tools.len()
147    }
148
149    /// Check if a tool exists
150    #[must_use]
151    pub fn has_tool(&self, name: &str) -> bool {
152        self.tools.contains_key(name)
153    }
154
155    /// Remove a tool by name
156    pub fn remove_tool(&mut self, name: &str) -> Option<ToolMetadata> {
157        self.operations.remove(name);
158        self.tools.remove(name)
159    }
160
161    /// Clear all tools
162    pub fn clear(&mut self) {
163        self.tools.clear();
164        self.operations.clear();
165        self.spec = None;
166    }
167
168    /// Get the source `OpenAPI` spec
169    #[must_use]
170    pub fn get_spec(&self) -> Option<&Spec> {
171        self.spec.as_ref()
172    }
173
174    /// Get registry statistics
175    #[must_use]
176    pub fn get_stats(&self) -> ToolRegistryStats {
177        let mut method_counts = HashMap::new();
178        let mut path_counts = HashMap::new();
179
180        for tool in self.tools.values() {
181            *method_counts.entry(tool.method.clone()).or_insert(0) += 1;
182            *path_counts.entry(tool.path.clone()).or_insert(0) += 1;
183        }
184
185        ToolRegistryStats {
186            total_tools: self.tools.len(),
187            method_distribution: method_counts,
188            unique_paths: path_counts.len(),
189            has_spec: self.spec.is_some(),
190        }
191    }
192
193    /// Validate all tools in the registry
194    ///
195    /// # Errors
196    ///
197    /// Returns an error if any tool is missing its operation, has invalid metadata, or there are orphaned operations
198    pub fn validate_registry(&self) -> Result<(), Error> {
199        for tool in self.tools.values() {
200            // Check if corresponding operation exists
201            if !self.operations.contains_key(&tool.name) {
202                return Err(Error::ToolGeneration(format!(
203                    "Missing operation for tool '{}'",
204                    tool.name
205                )));
206            }
207
208            // Validate tool metadata schema
209            Self::validate_tool_metadata(&tool.name, tool)?;
210        }
211
212        // Check for orphaned operations
213        for operation_name in self.operations.keys() {
214            if !self.tools.contains_key(operation_name) {
215                return Err(Error::ToolGeneration(format!(
216                    "Orphaned operation '{operation_name}'"
217                )));
218            }
219        }
220
221        Ok(())
222    }
223
224    /// Validate a single tool's metadata
225    fn validate_tool_metadata(tool_name: &str, tool_metadata: &ToolMetadata) -> Result<(), Error> {
226        // Check that the tool has valid parameters schema
227        if !tool_metadata.parameters.is_object() {
228            return Err(Error::Validation(format!(
229                "Tool '{tool_name}' has invalid parameters schema - must be an object"
230            )));
231        }
232
233        let schema_obj = tool_metadata.parameters.as_object().unwrap();
234
235        // Check for required properties field
236        if let Some(properties) = schema_obj.get("properties") {
237            if !properties.is_object() {
238                return Err(Error::Validation(format!(
239                    "Tool '{tool_name}' properties field must be an object"
240                )));
241            }
242        } else {
243            return Err(Error::Validation(format!(
244                "Tool '{tool_name}' is missing properties field in parameters schema"
245            )));
246        }
247
248        // Validate required field if present
249        if let Some(required) = schema_obj.get("required")
250            && !required.is_array()
251        {
252            return Err(Error::Validation(format!(
253                "Tool '{tool_name}' required field must be an array"
254            )));
255        }
256
257        // Check HTTP method is valid
258        let valid_methods = ["GET", "POST", "PUT", "DELETE", "PATCH", "HEAD", "OPTIONS"];
259        if !valid_methods.contains(&tool_metadata.method.to_uppercase().as_str()) {
260            return Err(Error::Validation(format!(
261                "Tool '{}' has invalid HTTP method: {}",
262                tool_name, tool_metadata.method
263            )));
264        }
265
266        // Check path is not empty
267        if tool_metadata.path.is_empty() {
268            return Err(Error::Validation(format!(
269                "Tool '{tool_name}' has empty path"
270            )));
271        }
272
273        Ok(())
274    }
275}
276
277impl Default for ToolRegistry {
278    fn default() -> Self {
279        Self::new()
280    }
281}
282
283/// Statistics about the tool registry
284#[derive(Debug, Clone)]
285pub struct ToolRegistryStats {
286    pub total_tools: usize,
287    pub method_distribution: HashMap<String, usize>,
288    pub unique_paths: usize,
289    pub has_spec: bool,
290}
291
292impl ToolRegistryStats {
293    /// Get a summary string of the registry stats
294    #[must_use]
295    pub fn summary(&self) -> String {
296        let methods: Vec<String> = self
297            .method_distribution
298            .iter()
299            .map(|(method, count)| format!("{}: {}", method.to_uppercase(), count))
300            .collect();
301
302        format!(
303            "Tools: {}, Methods: [{}], Paths: {}, Spec: {}",
304            self.total_tools,
305            methods.join(", "),
306            self.unique_paths,
307            if self.has_spec { "loaded" } else { "none" }
308        )
309    }
310}