Skip to main content

hh_cli/tool/
registry.rs

1use crate::config::Settings;
2use crate::core::ToolExecutor;
3use crate::tool::bash::BashTool;
4use crate::tool::edit::EditTool;
5use crate::tool::fs::{FsGlob, FsGrep, FsList, FsRead, FsWrite};
6use crate::tool::question::QuestionTool;
7use crate::tool::skill::SkillTool;
8use crate::tool::task::{TaskTool, TaskToolRuntimeContext};
9use crate::tool::todo::{TodoReadTool, TodoWriteTool};
10use crate::tool::web::{WebFetchTool, WebSearchTool};
11use crate::tool::{Tool, ToolSchema};
12use async_trait::async_trait;
13use std::collections::HashMap;
14use std::path::Path;
15use std::sync::Arc;
16
17pub struct ToolRegistry {
18    tools: HashMap<String, Arc<dyn Tool>>,
19    non_blocking_tools: std::collections::HashSet<String>,
20}
21
22#[derive(Clone, Default)]
23pub struct ToolRegistryContext {
24    pub task: Option<TaskToolRuntimeContext>,
25}
26
27impl ToolRegistry {
28    pub fn new(settings: &Settings, workspace_root: &Path) -> Self {
29        Self::new_with_context(settings, workspace_root, ToolRegistryContext::default())
30    }
31
32    pub fn new_with_context(
33        settings: &Settings,
34        workspace_root: &Path,
35        context: ToolRegistryContext,
36    ) -> Self {
37        let mut tools: HashMap<String, Arc<dyn Tool>> = HashMap::new();
38        let mut non_blocking_tools = std::collections::HashSet::new();
39
40        if settings.tools.fs {
41            register(&mut tools, "read", FsRead);
42            register(
43                &mut tools,
44                "write",
45                FsWrite::new(workspace_root.to_path_buf()),
46            );
47            register(&mut tools, "list", FsList);
48            register(&mut tools, "glob", FsGlob);
49            register(&mut tools, "grep", FsGrep);
50            register(&mut tools, "todo_read", TodoReadTool);
51            register(&mut tools, "todo_write", TodoWriteTool);
52            register(&mut tools, "question", QuestionTool);
53            register(
54                &mut tools,
55                "edit",
56                EditTool::new(workspace_root.to_path_buf()),
57            );
58            register(
59                &mut tools,
60                "skill",
61                SkillTool::new(workspace_root.to_path_buf()),
62            );
63
64            if let Some(task_context) = context.task {
65                register(&mut tools, "task", TaskTool::new(task_context));
66                non_blocking_tools.insert("task".to_string());
67            }
68        }
69
70        if settings.tools.bash {
71            register(&mut tools, "bash", BashTool::new());
72        }
73
74        if settings.tools.web {
75            register(&mut tools, "web_fetch", WebFetchTool::new());
76            register(&mut tools, "web_search", WebSearchTool::new());
77        }
78
79        Self {
80            tools,
81            non_blocking_tools,
82        }
83    }
84
85    pub fn schemas(&self) -> Vec<ToolSchema> {
86        self.tools.values().map(|t| t.schema()).collect()
87    }
88
89    pub async fn execute(&self, name: &str, args: serde_json::Value) -> crate::tool::ToolResult {
90        match self.tools.get(name) {
91            Some(tool) => tool.execute(args).await,
92            None => {
93                crate::tool::ToolResult::err_text("unknown_tool", format!("unknown tool: {}", name))
94            }
95        }
96    }
97
98    pub fn names(&self) -> Vec<String> {
99        let mut names = self.tools.keys().cloned().collect::<Vec<_>>();
100        names.sort();
101        names
102    }
103}
104
105#[async_trait]
106impl ToolExecutor for ToolRegistry {
107    fn schemas(&self) -> Vec<ToolSchema> {
108        self.schemas()
109    }
110
111    async fn execute(&self, name: &str, args: serde_json::Value) -> crate::tool::ToolResult {
112        self.execute(name, args).await
113    }
114
115    fn is_non_blocking(&self, name: &str) -> bool {
116        self.non_blocking_tools.contains(name)
117    }
118}
119
120fn register<T: Tool + 'static>(tools: &mut HashMap<String, Arc<dyn Tool>>, name: &str, tool: T) {
121    tools.insert(name.to_string(), Arc::new(tool));
122}