rmcp_openapi/
tool_registry.rs

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