agent_code_lib/tools/
registry.rs

1//! Tool registry: collects tools and dispatches by name.
2
3use std::sync::Arc;
4
5use super::{Tool, ToolSchema};
6
7/// Registry of available tools, keyed by name.
8pub struct ToolRegistry {
9    tools: Vec<Arc<dyn Tool>>,
10}
11
12impl ToolRegistry {
13    /// Create an empty registry.
14    pub fn new() -> Self {
15        Self { tools: Vec::new() }
16    }
17
18    /// Create a registry with all default built-in tools.
19    pub fn default_tools() -> Self {
20        let mut registry = Self::new();
21        registry.register(Arc::new(super::agent::AgentTool));
22        registry.register(Arc::new(super::bash::BashTool));
23        registry.register(Arc::new(super::file_read::FileReadTool));
24        registry.register(Arc::new(super::file_write::FileWriteTool));
25        registry.register(Arc::new(super::file_edit::FileEditTool));
26        registry.register(Arc::new(super::multi_edit::MultiEditTool));
27        registry.register(Arc::new(super::grep::GrepTool));
28        registry.register(Arc::new(super::glob::GlobTool));
29        registry.register(Arc::new(super::notebook_edit::NotebookEditTool));
30        registry.register(Arc::new(super::lsp_tool::LspTool));
31        registry.register(Arc::new(super::mcp_resources::ListMcpResourcesTool));
32        registry.register(Arc::new(super::mcp_resources::ReadMcpResourceTool));
33        registry.register(Arc::new(super::plan_mode::EnterPlanModeTool));
34        registry.register(Arc::new(super::plan_mode::ExitPlanModeTool));
35        registry.register(Arc::new(super::repl_tool::ReplTool));
36        registry.register(Arc::new(super::send_message::SendMessageTool));
37        registry.register(Arc::new(super::skill_tool::SkillTool));
38        registry.register(Arc::new(super::sleep_tool::SleepTool));
39        registry.register(Arc::new(super::tasks::TaskCreateTool));
40        registry.register(Arc::new(super::tasks::TaskUpdateTool));
41        registry.register(Arc::new(super::tasks::TaskGetTool));
42        registry.register(Arc::new(super::tasks::TaskListTool));
43        registry.register(Arc::new(super::tasks::TaskStopTool));
44        registry.register(Arc::new(super::tasks::TaskOutputTool));
45        registry.register(Arc::new(super::todo_write::TodoWriteTool));
46        registry.register(Arc::new(super::tool_search::ToolSearchTool));
47        registry.register(Arc::new(super::worktree::EnterWorktreeTool));
48        registry.register(Arc::new(super::worktree::ExitWorktreeTool));
49        registry.register(Arc::new(super::web_fetch::WebFetchTool));
50        registry.register(Arc::new(super::web_search::WebSearchTool));
51        registry.register(Arc::new(super::ask_user::AskUserQuestionTool));
52        registry.register(Arc::new(super::powershell::PowerShellTool));
53        registry
54    }
55
56    /// Register a new tool.
57    pub fn register(&mut self, tool: Arc<dyn Tool>) {
58        self.tools.push(tool);
59    }
60
61    /// Look up a tool by name.
62    pub fn get(&self, name: &str) -> Option<Arc<dyn Tool>> {
63        self.tools.iter().find(|t| t.name() == name).cloned()
64    }
65
66    /// Get all registered tools.
67    pub fn all(&self) -> &[Arc<dyn Tool>] {
68        &self.tools
69    }
70
71    /// Get tool schemas for the API request.
72    pub fn schemas(&self) -> Vec<ToolSchema> {
73        self.tools
74            .iter()
75            .filter(|t| t.is_enabled())
76            .map(|t| ToolSchema::from(t.as_ref()))
77            .collect()
78    }
79
80    /// Get only always-loaded (core) tool schemas.
81    /// Deferred tools are discoverable via ToolSearch but not sent on every request.
82    pub fn core_schemas(&self) -> Vec<ToolSchema> {
83        self.tools
84            .iter()
85            .filter(|t| t.is_enabled() && !is_deferred(t.name()))
86            .map(|t| ToolSchema::from(t.as_ref()))
87            .collect()
88    }
89
90    /// Get deferred tool names (for the ToolSearch system prompt).
91    pub fn deferred_names(&self) -> Vec<&str> {
92        self.tools
93            .iter()
94            .filter(|t| t.is_enabled() && is_deferred(t.name()))
95            .map(|t| t.name())
96            .collect()
97    }
98}
99
100/// Tools that are deferred — not sent on every request to save prompt tokens.
101/// These are discoverable via ToolSearch and loaded on demand.
102const DEFERRED_TOOLS: &[&str] = &[
103    "NotebookEdit",
104    "LSP",
105    "ListMcpResources",
106    "ReadMcpResource",
107    "EnterWorktree",
108    "ExitWorktree",
109    "Sleep",
110    "TodoWrite",
111    "REPL",
112    "SendMessage",
113    "TaskStop",
114    "TaskOutput",
115    "TaskGet",
116];
117
118fn is_deferred(name: &str) -> bool {
119    DEFERRED_TOOLS.contains(&name)
120}
121
122#[cfg(test)]
123mod tests {
124    use super::*;
125
126    #[test]
127    fn test_default_tools_count() {
128        let reg = ToolRegistry::default_tools();
129        assert!(reg.all().len() >= 30);
130    }
131
132    #[test]
133    fn test_get_by_name() {
134        let reg = ToolRegistry::default_tools();
135        assert!(reg.get("Bash").is_some());
136        assert!(reg.get("FileRead").is_some());
137        assert!(reg.get("NonExistent").is_none());
138    }
139
140    #[test]
141    fn test_schemas_returns_enabled() {
142        let reg = ToolRegistry::default_tools();
143        let schemas = reg.schemas();
144        assert!(!schemas.is_empty());
145    }
146
147    #[test]
148    fn test_core_schemas_excludes_deferred() {
149        let reg = ToolRegistry::default_tools();
150        let core = reg.core_schemas();
151        let all = reg.schemas();
152        assert!(core.len() < all.len());
153    }
154
155    #[test]
156    fn test_deferred_names() {
157        let reg = ToolRegistry::default_tools();
158        let deferred = reg.deferred_names();
159        assert!(deferred.contains(&"NotebookEdit"));
160        assert!(deferred.contains(&"Sleep"));
161        assert!(!deferred.contains(&"Bash"));
162        assert!(!deferred.contains(&"FileRead"));
163    }
164
165    #[test]
166    fn test_register_custom_tool() {
167        let mut reg = ToolRegistry::new();
168        assert_eq!(reg.all().len(), 0);
169        reg.register(Arc::new(super::super::bash::BashTool));
170        assert_eq!(reg.all().len(), 1);
171        assert!(reg.get("Bash").is_some());
172    }
173}
agent_code_lib/tools/registry.rs

agent_code_lib/tools/
registry.rs
