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