rmcp_openapi/
tool_registry.rs

1use crate::error::OpenApiError;
2use crate::openapi_spec::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, (openapiv3::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    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, method, path)) = spec.get_operation(&tool.name) {
40                self.register_tool(tool, (operation.clone(), method, path))?;
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: (openapiv3::Operation, String, String),
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(
107        &self,
108        tool_name: &str,
109    ) -> Option<&(openapiv3::Operation, String, String)> {
110        self.operations.get(tool_name)
111    }
112
113    /// Get all tool names
114    pub fn get_tool_names(&self) -> Vec<String> {
115        self.tools.keys().cloned().collect()
116    }
117
118    /// Get all tools
119    pub fn get_all_tools(&self) -> Vec<&ToolMetadata> {
120        self.tools.values().collect()
121    }
122
123    /// Get number of registered tools
124    pub fn tool_count(&self) -> usize {
125        self.tools.len()
126    }
127
128    /// Check if a tool exists
129    pub fn has_tool(&self, name: &str) -> bool {
130        self.tools.contains_key(name)
131    }
132
133    /// Remove a tool by name
134    pub fn remove_tool(&mut self, name: &str) -> Option<ToolMetadata> {
135        self.operations.remove(name);
136        self.tools.remove(name)
137    }
138
139    /// Clear all tools
140    pub fn clear(&mut self) {
141        self.tools.clear();
142        self.operations.clear();
143        self.spec = None;
144    }
145
146    /// Get the source OpenAPI spec
147    pub fn get_spec(&self) -> Option<&OpenApiSpec> {
148        self.spec.as_ref()
149    }
150
151    /// Get registry statistics
152    pub fn get_stats(&self) -> ToolRegistryStats {
153        let mut method_counts = HashMap::new();
154        let mut path_counts = HashMap::new();
155
156        for tool in self.tools.values() {
157            *method_counts.entry(tool.method.clone()).or_insert(0) += 1;
158            *path_counts.entry(tool.path.clone()).or_insert(0) += 1;
159        }
160
161        ToolRegistryStats {
162            total_tools: self.tools.len(),
163            method_distribution: method_counts,
164            unique_paths: path_counts.len(),
165            has_spec: self.spec.is_some(),
166        }
167    }
168
169    /// Validate all tools in the registry
170    pub fn validate_registry(&self) -> Result<(), OpenApiError> {
171        for tool in self.tools.values() {
172            // Check if corresponding operation exists
173            if !self.operations.contains_key(&tool.name) {
174                return Err(OpenApiError::ToolGeneration(format!(
175                    "Missing operation for tool '{}'",
176                    tool.name
177                )));
178            }
179
180            // Validate tool metadata schema
181            self.validate_tool_metadata(&tool.name, tool)?;
182        }
183
184        // Check for orphaned operations
185        for operation_name in self.operations.keys() {
186            if !self.tools.contains_key(operation_name) {
187                return Err(OpenApiError::ToolGeneration(format!(
188                    "Orphaned operation '{operation_name}'"
189                )));
190            }
191        }
192
193        Ok(())
194    }
195
196    /// Validate a single tool's metadata
197    fn validate_tool_metadata(
198        &self,
199        tool_name: &str,
200        tool_metadata: &ToolMetadata,
201    ) -> Result<(), OpenApiError> {
202        // Check that the tool has valid parameters schema
203        if !tool_metadata.parameters.is_object() {
204            return Err(OpenApiError::Validation(format!(
205                "Tool '{tool_name}' has invalid parameters schema - must be an object"
206            )));
207        }
208
209        let schema_obj = tool_metadata.parameters.as_object().unwrap();
210
211        // Check for required properties field
212        if let Some(properties) = schema_obj.get("properties") {
213            if !properties.is_object() {
214                return Err(OpenApiError::Validation(format!(
215                    "Tool '{tool_name}' properties field must be an object"
216                )));
217            }
218        } else {
219            return Err(OpenApiError::Validation(format!(
220                "Tool '{tool_name}' is missing properties field in parameters schema"
221            )));
222        }
223
224        // Validate required field if present
225        if let Some(required) = schema_obj.get("required") {
226            if !required.is_array() {
227                return Err(OpenApiError::Validation(format!(
228                    "Tool '{tool_name}' required field must be an array"
229                )));
230            }
231        }
232
233        // Check HTTP method is valid
234        let valid_methods = ["GET", "POST", "PUT", "DELETE", "PATCH", "HEAD", "OPTIONS"];
235        if !valid_methods.contains(&tool_metadata.method.to_uppercase().as_str()) {
236            return Err(OpenApiError::Validation(format!(
237                "Tool '{}' has invalid HTTP method: {}",
238                tool_name, tool_metadata.method
239            )));
240        }
241
242        // Check path is not empty
243        if tool_metadata.path.is_empty() {
244            return Err(OpenApiError::Validation(format!(
245                "Tool '{tool_name}' has empty path"
246            )));
247        }
248
249        Ok(())
250    }
251}
252
253impl Default for ToolRegistry {
254    fn default() -> Self {
255        Self::new()
256    }
257}
258
259/// Statistics about the tool registry
260#[derive(Debug, Clone)]
261pub struct ToolRegistryStats {
262    pub total_tools: usize,
263    pub method_distribution: HashMap<String, usize>,
264    pub unique_paths: usize,
265    pub has_spec: bool,
266}
267
268impl ToolRegistryStats {
269    /// Get a summary string of the registry stats
270    pub fn summary(&self) -> String {
271        let methods: Vec<String> = self
272            .method_distribution
273            .iter()
274            .map(|(method, count)| format!("{}: {}", method.to_uppercase(), count))
275            .collect();
276
277        format!(
278            "Tools: {}, Methods: [{}], Paths: {}, Spec: {}",
279            self.total_tools,
280            methods.join(", "),
281            self.unique_paths,
282            if self.has_spec { "loaded" } else { "none" }
283        )
284    }
285}