ricecoder_mcp/
registry.rs

1//! Tool Registry for managing available tools
2
3use crate::error::Result;
4use crate::metadata::ToolMetadata;
5use std::collections::HashMap;
6
7/// Tool Registry for managing all available tools
8#[derive(Debug, Clone)]
9pub struct ToolRegistry {
10    tools: HashMap<String, ToolMetadata>,
11}
12
13impl ToolRegistry {
14    /// Creates a new tool registry
15    pub fn new() -> Self {
16        Self {
17            tools: HashMap::new(),
18        }
19    }
20
21    /// Registers a tool in the registry
22    pub fn register_tool(&mut self, tool: ToolMetadata) -> Result<()> {
23        // Check for naming conflicts
24        if self.tools.contains_key(&tool.id) {
25            return Err(crate::error::Error::NamingConflict(format!(
26                "Tool with ID '{}' already exists",
27                tool.id
28            )));
29        }
30
31        self.tools.insert(tool.id.clone(), tool);
32        Ok(())
33    }
34
35    /// Gets a tool by ID
36    pub fn get_tool(&self, id: &str) -> Option<&ToolMetadata> {
37        self.tools.get(id)
38    }
39
40    /// Lists all tools
41    pub fn list_tools(&self) -> Vec<&ToolMetadata> {
42        self.tools.values().collect()
43    }
44
45    /// Lists tools by category
46    pub fn list_tools_by_category(&self, category: &str) -> Vec<&ToolMetadata> {
47        self.tools
48            .values()
49            .filter(|t| t.category == category)
50            .collect()
51    }
52
53    /// Lists tools by server
54    pub fn list_tools_by_server(&self, server_id: &str) -> Vec<&ToolMetadata> {
55        self.tools
56            .values()
57            .filter(|t| t.server_id.as_deref() == Some(server_id))
58            .collect()
59    }
60
61    /// Gets the number of registered tools
62    pub fn tool_count(&self) -> usize {
63        self.tools.len()
64    }
65}
66
67impl Default for ToolRegistry {
68    fn default() -> Self {
69        Self::new()
70    }
71}
72
73#[cfg(test)]
74mod tests {
75    use super::*;
76    use crate::metadata::ToolSource;
77
78    #[test]
79    fn test_register_tool() {
80        let mut registry = ToolRegistry::new();
81        let tool = ToolMetadata {
82            id: "test-tool".to_string(),
83            name: "Test Tool".to_string(),
84            description: "A test tool".to_string(),
85            category: "test".to_string(),
86            parameters: vec![],
87            return_type: "string".to_string(),
88            source: ToolSource::Custom,
89            server_id: None,
90        };
91
92        let result = registry.register_tool(tool);
93        assert!(result.is_ok());
94        assert_eq!(registry.tool_count(), 1);
95    }
96
97    #[test]
98    fn test_register_duplicate_tool() {
99        let mut registry = ToolRegistry::new();
100        let tool = ToolMetadata {
101            id: "test-tool".to_string(),
102            name: "Test Tool".to_string(),
103            description: "A test tool".to_string(),
104            category: "test".to_string(),
105            parameters: vec![],
106            return_type: "string".to_string(),
107            source: ToolSource::Custom,
108            server_id: None,
109        };
110
111        registry.register_tool(tool.clone()).unwrap();
112        let result = registry.register_tool(tool);
113        assert!(result.is_err());
114    }
115
116    #[test]
117    fn test_get_tool() {
118        let mut registry = ToolRegistry::new();
119        let tool = ToolMetadata {
120            id: "test-tool".to_string(),
121            name: "Test Tool".to_string(),
122            description: "A test tool".to_string(),
123            category: "test".to_string(),
124            parameters: vec![],
125            return_type: "string".to_string(),
126            source: ToolSource::Custom,
127            server_id: None,
128        };
129
130        registry.register_tool(tool).unwrap();
131        let retrieved = registry.get_tool("test-tool");
132        assert!(retrieved.is_some());
133        assert_eq!(retrieved.unwrap().name, "Test Tool");
134    }
135
136    #[test]
137    fn test_list_tools_by_category() {
138        let mut registry = ToolRegistry::new();
139        let tool1 = ToolMetadata {
140            id: "tool1".to_string(),
141            name: "Tool 1".to_string(),
142            description: "Tool 1".to_string(),
143            category: "math".to_string(),
144            parameters: vec![],
145            return_type: "number".to_string(),
146            source: ToolSource::Custom,
147            server_id: None,
148        };
149
150        let tool2 = ToolMetadata {
151            id: "tool2".to_string(),
152            name: "Tool 2".to_string(),
153            description: "Tool 2".to_string(),
154            category: "string".to_string(),
155            parameters: vec![],
156            return_type: "string".to_string(),
157            source: ToolSource::Custom,
158            server_id: None,
159        };
160
161        registry.register_tool(tool1).unwrap();
162        registry.register_tool(tool2).unwrap();
163
164        let math_tools = registry.list_tools_by_category("math");
165        assert_eq!(math_tools.len(), 1);
166        assert_eq!(math_tools[0].id, "tool1");
167    }
168}
ricecoder_mcp/registry.rs

ricecoder_mcp/
registry.rs
