rmcp_openapi/
tool_registry.rs

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