rmcp_openapi/
tool_registry.rs

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