kaish_kernel/tools/
registry.rs

1//! Tool registry for looking up and managing tools.
2
3use std::collections::HashMap;
4use std::sync::Arc;
5
6use super::traits::{Tool, ToolSchema};
7
8/// Registry of available tools.
9#[derive(Default)]
10pub struct ToolRegistry {
11    tools: HashMap<String, Arc<dyn Tool>>,
12}
13
14impl ToolRegistry {
15    /// Create a new empty registry.
16    pub fn new() -> Self {
17        Self::default()
18    }
19
20    /// Register a tool.
21    pub fn register(&mut self, tool: impl Tool + 'static) {
22        let name = tool.name().to_string();
23        self.tools.insert(name, Arc::new(tool));
24    }
25
26    /// Register a tool that's already in an Arc.
27    pub fn register_arc(&mut self, tool: Arc<dyn Tool>) {
28        let name = tool.name().to_string();
29        self.tools.insert(name, tool);
30    }
31
32    /// Look up a tool by name.
33    pub fn get(&self, name: &str) -> Option<Arc<dyn Tool>> {
34        self.tools.get(name).cloned()
35    }
36
37    /// Check if a tool exists.
38    pub fn contains(&self, name: &str) -> bool {
39        self.tools.contains_key(name)
40    }
41
42    /// List all tool names.
43    pub fn names(&self) -> Vec<&str> {
44        let mut names: Vec<_> = self.tools.keys().map(|s| s.as_str()).collect();
45        names.sort();
46        names
47    }
48
49    /// List all tool schemas.
50    pub fn schemas(&self) -> Vec<ToolSchema> {
51        let mut schemas: Vec<_> = self.tools.values().map(|t| t.schema()).collect();
52        schemas.sort_by(|a, b| a.name.cmp(&b.name));
53        schemas
54    }
55
56    /// Number of registered tools.
57    pub fn len(&self) -> usize {
58        self.tools.len()
59    }
60
61    /// Check if empty.
62    pub fn is_empty(&self) -> bool {
63        self.tools.is_empty()
64    }
65}
66
67impl std::fmt::Debug for ToolRegistry {
68    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
69        f.debug_struct("ToolRegistry")
70            .field("tools", &self.names())
71            .finish()
72    }
73}
74
75#[cfg(test)]
76mod tests {
77    use super::*;
78    use crate::interpreter::ExecResult;
79    use crate::tools::{ExecContext, ToolArgs};
80    use async_trait::async_trait;
81
82    struct DummyTool;
83
84    #[async_trait]
85    impl Tool for DummyTool {
86        fn name(&self) -> &str {
87            "dummy"
88        }
89
90        fn schema(&self) -> ToolSchema {
91            ToolSchema::new("dummy", "A test tool")
92        }
93
94        async fn execute(&self, _args: ToolArgs, _ctx: &mut ExecContext) -> ExecResult {
95            ExecResult::success("dummy output")
96        }
97    }
98
99    #[test]
100    fn test_register_and_get() {
101        let mut registry = ToolRegistry::new();
102        registry.register(DummyTool);
103
104        assert!(registry.contains("dummy"));
105        assert!(registry.get("dummy").is_some());
106        assert!(!registry.contains("nonexistent"));
107    }
108
109    #[test]
110    fn test_names_sorted() {
111        let mut registry = ToolRegistry::new();
112
113        struct ToolA;
114        struct ToolZ;
115
116        #[async_trait]
117        impl Tool for ToolA {
118            fn name(&self) -> &str { "aaa" }
119            fn schema(&self) -> ToolSchema { ToolSchema::new("aaa", "") }
120            async fn execute(&self, _: ToolArgs, _: &mut ExecContext) -> ExecResult {
121                ExecResult::success("")
122            }
123        }
124
125        #[async_trait]
126        impl Tool for ToolZ {
127            fn name(&self) -> &str { "zzz" }
128            fn schema(&self) -> ToolSchema { ToolSchema::new("zzz", "") }
129            async fn execute(&self, _: ToolArgs, _: &mut ExecContext) -> ExecResult {
130                ExecResult::success("")
131            }
132        }
133
134        registry.register(ToolZ);
135        registry.register(ToolA);
136
137        let names = registry.names();
138        assert_eq!(names, vec!["aaa", "zzz"]);
139    }
140}
kaish_kernel/tools/registry.rs

kaish_kernel/tools/
registry.rs
