rmcp_openapi/
tool_registry.rs1use crate::error::OpenApiError;
2use crate::openapi_spec::{OpenApiOperation, OpenApiSpec};
3use crate::server::ToolMetadata;
4use std::collections::HashMap;
5
6#[derive(Debug, Clone)]
8pub struct ToolRegistry {
9 tools: HashMap<String, ToolMetadata>,
11 operations: HashMap<String, OpenApiOperation>,
13 spec: Option<OpenApiSpec>,
15}
16
17impl ToolRegistry {
18 pub fn new() -> Self {
20 Self {
21 tools: HashMap::new(),
22 operations: HashMap::new(),
23 spec: None,
24 }
25 }
26
27 pub fn register_from_spec(&mut self, spec: OpenApiSpec) -> Result<usize, OpenApiError> {
29 self.clear();
31
32 let tools_metadata = spec.to_tool_metadata()?;
34 let mut registered_count = 0;
35
36 for tool in tools_metadata {
38 if let Some(operation) = spec.get_operation(&tool.name) {
40 self.register_tool(tool, operation.clone())?;
41 registered_count += 1;
42 }
43 }
44
45 self.spec = Some(spec);
47
48 Ok(registered_count)
49 }
50
51 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 self.validate_tool(&tool)?;
61
62 self.tools.insert(tool_name.clone(), tool);
64 self.operations.insert(tool_name, operation);
65
66 Ok(())
67 }
68
69 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 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 pub fn get_tool(&self, name: &str) -> Option<&ToolMetadata> {
102 self.tools.get(name)
103 }
104
105 pub fn get_operation(&self, tool_name: &str) -> Option<&OpenApiOperation> {
107 self.operations.get(tool_name)
108 }
109
110 pub fn get_tool_names(&self) -> Vec<String> {
112 self.tools.keys().cloned().collect()
113 }
114
115 pub fn get_all_tools(&self) -> Vec<&ToolMetadata> {
117 self.tools.values().collect()
118 }
119
120 pub fn tool_count(&self) -> usize {
122 self.tools.len()
123 }
124
125 pub fn has_tool(&self, name: &str) -> bool {
127 self.tools.contains_key(name)
128 }
129
130 pub fn remove_tool(&mut self, name: &str) -> Option<ToolMetadata> {
132 self.operations.remove(name);
133 self.tools.remove(name)
134 }
135
136 pub fn clear(&mut self) {
138 self.tools.clear();
139 self.operations.clear();
140 self.spec = None;
141 }
142
143 pub fn get_spec(&self) -> Option<&OpenApiSpec> {
145 self.spec.as_ref()
146 }
147
148 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 pub fn validate_registry(&self) -> Result<(), OpenApiError> {
168 for tool in self.tools.values() {
169 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 self.validate_tool_metadata(&tool.name, tool)?;
179 }
180
181 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 fn validate_tool_metadata(
195 &self,
196 tool_name: &str,
197 tool_metadata: &ToolMetadata,
198 ) -> Result<(), OpenApiError> {
199 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 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 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 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 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#[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 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}