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

agent_code_lib/tools/
registry.rs
