Skip to main content

bamboo_agent/agent/tools/
executor.rs

1use std::sync::Arc;
2
3use crate::agent::core::tools::{
4    normalize_tool_name, Tool, ToolCall, ToolError, ToolExecutionContext, ToolExecutor, ToolResult,
5    ToolSchema,
6};
7use async_trait::async_trait;
8use serde_json::json;
9
10use crate::agent::tools::guide::{context::GuideBuildContext, EnhancedPromptBuilder, ToolGuide};
11use crate::agent::tools::permission::{check_permissions, PermissionChecker, PermissionError};
12use crate::agent::tools::tools::{
13    ApplyPatchTool, AskUserTool, CreateTodoListTool, ExecuteCommandTool, FileExistsTool,
14    GetCurrentDirTool, GetFileInfoTool, GitDiffTool, GitStatusTool, GitWriteTool, GlobSearchTool,
15    HttpRequestTool, ListDirectoryTool, ReadFileRangeTool, ReadFileTool, SearchInFileTool,
16    SearchInProjectTool, SetWorkspaceTool, SleepTool, TerminalSessionTool, ToolRegistry,
17    UpdateTodoItemTool, WriteFileTool,
18};
19use crate::core::Config;
20use tokio::sync::RwLock;
21
22/// List of all built-in tool names.
23///
24/// This list intentionally includes only tools that are always registered by
25/// `BuiltinToolExecutor::new()`. Optional tools (for example integrations that
26/// depend on host binaries) should NOT be added here.
27pub const BUILTIN_TOOL_NAMES: [&str; 22] = [
28    "read_file",
29    "write_file",
30    "list_directory",
31    "file_exists",
32    "get_file_info",
33    "execute_command",
34    "ask_user",
35    "get_current_dir",
36    "set_workspace",
37    "read_file_range",
38    "search_in_file",
39    "apply_patch",
40    "search_in_project",
41    "git_status",
42    "git_diff",
43    "git_write",
44    "create_todo_list",
45    "update_todo_item",
46    "glob_search",
47    "http_request",
48    "sleep",
49    "terminal_session",
50];
51
52/// Normalizes a tool reference to a standard tool name
53///
54/// Handles legacy aliases like "run_command" -> "execute_command"
55/// Returns None if the tool name is not recognized
56pub fn normalize_tool_ref(value: &str) -> Option<String> {
57    let trimmed = value.trim();
58    if trimmed.is_empty() {
59        return None;
60    }
61    let raw_tool_name = trimmed.split("::").last().unwrap_or(trimmed);
62    let tool_name = match raw_tool_name {
63        "run_command" => "execute_command",
64        _ => raw_tool_name,
65    };
66    // `claude_code` is an optional integration tool that may be registered at runtime.
67    if BUILTIN_TOOL_NAMES.iter().any(|name| name == &tool_name) || tool_name == "claude_code" {
68        Some(tool_name.to_string())
69    } else {
70        None
71    }
72}
73
74/// Checks if a tool reference is a built-in tool
75pub fn is_builtin_tool(value: &str) -> bool {
76    normalize_tool_ref(value).is_some()
77}
78
79/// Built-in tool executor that uses ToolRegistry for dynamic dispatch
80pub struct BuiltinToolExecutor {
81    registry: ToolRegistry,
82    permission_checker: Option<Arc<dyn PermissionChecker>>,
83}
84
85impl BuiltinToolExecutor {
86    /// Creates a new executor with all built-in tools registered
87    pub fn new() -> Self {
88        let registry = ToolRegistry::new();
89        Self::register_builtin_tools(&registry, None);
90        Self {
91            registry,
92            permission_checker: None,
93        }
94    }
95
96    /// Creates a new executor with a permission checker
97    pub fn new_with_permissions(permission_checker: Arc<dyn PermissionChecker>) -> Self {
98        let registry = ToolRegistry::new();
99        Self::register_builtin_tools(&registry, None);
100        Self {
101            registry,
102            permission_checker: Some(permission_checker),
103        }
104    }
105
106    /// Creates a new executor that can read the shared, hot-reloadable config.
107    ///
108    /// Use this when running inside the Bamboo server so tools (notably
109    /// `http_request`) honor proxy settings from `config.json`.
110    pub fn new_with_config(config: Arc<RwLock<Config>>) -> Self {
111        let registry = ToolRegistry::new();
112        Self::register_builtin_tools(&registry, Some(config));
113        Self {
114            registry,
115            permission_checker: None,
116        }
117    }
118
119    /// Creates a new executor with both shared config and a permission checker.
120    pub fn new_with_config_and_permissions(
121        config: Arc<RwLock<Config>>,
122        permission_checker: Arc<dyn PermissionChecker>,
123    ) -> Self {
124        let registry = ToolRegistry::new();
125        Self::register_builtin_tools(&registry, Some(config));
126        Self {
127            registry,
128            permission_checker: Some(permission_checker),
129        }
130    }
131
132    /// Creates a new executor from an existing registry
133    pub fn with_registry(registry: ToolRegistry) -> Self {
134        Self {
135            registry,
136            permission_checker: None,
137        }
138    }
139
140    /// Returns a reference to the internal registry
141    pub fn registry(&self) -> &ToolRegistry {
142        &self.registry
143    }
144
145    /// Registers all built-in tools to the given registry
146    fn register_builtin_tools(registry: &ToolRegistry, config: Option<Arc<RwLock<Config>>>) {
147        // Register filesystem tools
148        let _ = registry.register(ReadFileTool::new());
149        let _ = registry.register(WriteFileTool::new());
150        let _ = registry.register(ListDirectoryTool::new());
151        let _ = registry.register(FileExistsTool::new());
152        let _ = registry.register(GetFileInfoTool::new());
153
154        // Register command tools
155        let _ = registry.register(ExecuteCommandTool::new());
156        let _ = registry.register(AskUserTool::new());
157        let _ = registry.register(GetCurrentDirTool::new());
158
159        // Register workspace tools
160        let _ = registry.register(SetWorkspaceTool::new());
161
162        // Register advanced file tools
163        let _ = registry.register(ReadFileRangeTool::new());
164        let _ = registry.register(SearchInFileTool::new());
165        let _ = registry.register(ApplyPatchTool::new());
166
167        // Register project-wide tools
168        let _ = registry.register(SearchInProjectTool::new());
169
170        // Register git tools
171        let _ = registry.register(GitStatusTool::new());
172        let _ = registry.register(GitDiffTool::new());
173        let _ = registry.register(GitWriteTool::new());
174
175        // Register todo list tools
176        let _ = registry.register(CreateTodoListTool::new());
177        let _ = registry.register(UpdateTodoItemTool::new());
178
179        // Register new utility tools
180        let _ = registry.register(GlobSearchTool::new());
181        let _ = match config {
182            Some(config) => registry.register(HttpRequestTool::new_with_config(config)),
183            None => registry.register(HttpRequestTool::new()),
184        };
185        let _ = registry.register(SleepTool::new());
186        let _ = registry.register(TerminalSessionTool::new());
187    }
188
189    /// Returns all built-in tool schemas
190    pub fn tool_schemas() -> Vec<ToolSchema> {
191        let registry = ToolRegistry::new();
192        Self::register_builtin_tools(&registry, None);
193        registry.list_tools()
194    }
195
196    /// Registers a custom tool to this executor
197    pub fn register_tool<T: Tool + 'static>(&self, tool: T) -> Result<(), ToolError> {
198        self.registry
199            .register(tool)
200            .map_err(|e| ToolError::Execution(e.to_string()))
201    }
202
203    /// Register a tool with its guide
204    pub fn register_tool_with_guide<T, G>(&self, tool: T, guide: G) -> Result<(), ToolError>
205    where
206        T: Tool + 'static,
207        G: ToolGuide + 'static,
208    {
209        self.registry
210            .register_with_guide(tool, guide)
211            .map_err(|e| ToolError::Execution(e.to_string()))
212    }
213
214    /// Get guide for a tool
215    pub fn get_guide(&self, tool_name: &str) -> Option<Arc<dyn ToolGuide>> {
216        self.registry.get_guide(tool_name)
217    }
218
219    /// Build enhanced prompt for all registered tools
220    pub fn build_enhanced_prompt(&self, context: GuideBuildContext) -> String {
221        EnhancedPromptBuilder::build(Some(&self.registry), &self.registry.list_tools(), &context)
222    }
223}
224
225fn permission_error_to_tool_error(error: PermissionError) -> ToolError {
226    match error {
227        PermissionError::CheckFailed(_) => ToolError::InvalidArguments(error.to_string()),
228        _ => ToolError::Execution(error.to_string()),
229    }
230}
231
232impl Default for BuiltinToolExecutor {
233    fn default() -> Self {
234        Self::new()
235    }
236}
237
238#[async_trait]
239impl ToolExecutor for BuiltinToolExecutor {
240    async fn execute(&self, call: &ToolCall) -> Result<ToolResult, ToolError> {
241        self.execute_with_context(call, ToolExecutionContext::none(&call.id))
242            .await
243    }
244
245    async fn execute_with_context(
246        &self,
247        call: &ToolCall,
248        ctx: ToolExecutionContext<'_>,
249    ) -> Result<ToolResult, ToolError> {
250        let args_raw = call.function.arguments.trim();
251        let args: serde_json::Value = if args_raw.is_empty() {
252            json!({})
253        } else {
254            serde_json::from_str(args_raw).map_err(|e| {
255                ToolError::InvalidArguments(format!("Invalid JSON arguments: {}", e))
256            })?
257        };
258
259        let tool_name = normalize_tool_name(&call.function.name);
260
261        // Look up the tool in the registry
262        let tool = self
263            .registry
264            .get(tool_name)
265            .ok_or_else(|| ToolError::NotFound(format!("Tool '{}' not found", tool_name)))?;
266
267        if let Some(permission_checker) = &self.permission_checker {
268            if let Some(contexts) =
269                check_permissions(tool_name, &args).map_err(permission_error_to_tool_error)?
270            {
271                for context in contexts {
272                    let resource = context.resource.clone();
273                    let allowed = permission_checker
274                        .check_or_request(context)
275                        .await
276                        .map_err(permission_error_to_tool_error)?;
277                    if !allowed {
278                        return Err(ToolError::Execution(format!(
279                            "Permission denied for: {}",
280                            resource
281                        )));
282                    }
283                }
284            }
285        }
286
287        tool.execute_with_context(args, ctx).await
288    }
289
290    fn list_tools(&self) -> Vec<ToolSchema> {
291        self.registry.list_tools()
292    }
293}
294
295/// Builder for constructing a BuiltinToolExecutor with custom tool configurations
296pub struct BuiltinToolExecutorBuilder {
297    registry: ToolRegistry,
298    permission_checker: Option<Arc<dyn PermissionChecker>>,
299}
300
301impl BuiltinToolExecutorBuilder {
302    /// Creates a new builder with no tools registered
303    pub fn new() -> Self {
304        Self {
305            registry: ToolRegistry::new(),
306            permission_checker: None,
307        }
308    }
309
310    /// Registers all default built-in tools
311    pub fn with_default_tools(self) -> Self {
312        BuiltinToolExecutor::register_builtin_tools(&self.registry, None);
313        self
314    }
315
316    /// Registers a specific filesystem tool by name
317    pub fn with_filesystem_tool(self, name: &str) -> Result<Self, ToolError> {
318        match name {
319            "read_file" => self.registry.register(ReadFileTool::new()),
320            "write_file" => self.registry.register(WriteFileTool::new()),
321            "list_directory" => self.registry.register(ListDirectoryTool::new()),
322            "file_exists" => self.registry.register(FileExistsTool::new()),
323            "get_file_info" => self.registry.register(GetFileInfoTool::new()),
324            _ => return Err(ToolError::NotFound(format!("Unknown tool: {}", name))),
325        }
326        .map_err(|e| ToolError::Execution(e.to_string()))?;
327        Ok(self)
328    }
329
330    /// Registers a specific command tool by name
331    pub fn with_command_tool(self, name: &str) -> Result<Self, ToolError> {
332        match name {
333            "execute_command" => self.registry.register(ExecuteCommandTool::new()),
334            "get_current_dir" => self.registry.register(GetCurrentDirTool::new()),
335            _ => return Err(ToolError::NotFound(format!("Unknown tool: {}", name))),
336        }
337        .map_err(|e| ToolError::Execution(e.to_string()))?;
338        Ok(self)
339    }
340
341    /// Registers a custom tool
342    pub fn with_tool<T: Tool + 'static>(self, tool: T) -> Result<Self, ToolError> {
343        self.registry
344            .register(tool)
345            .map_err(|e| ToolError::Execution(e.to_string()))?;
346        Ok(self)
347    }
348
349    /// Sets a permission checker for this executor
350    pub fn with_permission_checker(mut self, checker: Arc<dyn PermissionChecker>) -> Self {
351        self.permission_checker = Some(checker);
352        self
353    }
354
355    /// Builds the executor
356    pub fn build(self) -> BuiltinToolExecutor {
357        BuiltinToolExecutor {
358            registry: self.registry,
359            permission_checker: self.permission_checker,
360        }
361    }
362}
363
364impl Default for BuiltinToolExecutorBuilder {
365    fn default() -> Self {
366        Self::new()
367    }
368}
369
370#[cfg(test)]
371mod tests {
372    use super::*;
373    use crate::agent::core::tools::FunctionCall;
374    use crate::agent::core::tools::ToolExecutionContext;
375    use crate::agent::core::AgentEvent;
376    use serde_json::json;
377    use std::sync::Arc;
378    use tokio::fs;
379    use tokio::sync::mpsc;
380
381    use crate::agent::tools::tools::WriteFileTool;
382
383    fn make_tool_call(name: &str, args: serde_json::Value) -> ToolCall {
384        ToolCall {
385            id: "call_1".to_string(),
386            tool_type: "function".to_string(),
387            function: FunctionCall {
388                name: name.to_string(),
389                arguments: args.to_string(),
390            },
391        }
392    }
393
394    fn make_executor(
395        permission_checker: Option<Arc<dyn PermissionChecker>>,
396    ) -> BuiltinToolExecutor {
397        let builder = BuiltinToolExecutorBuilder::new()
398            .with_tool(WriteFileTool::new())
399            .expect("register write_file tool");
400
401        let builder = match permission_checker {
402            Some(checker) => builder.with_permission_checker(checker),
403            None => builder,
404        };
405
406        builder.build()
407    }
408
409    #[test]
410    fn test_normalize_tool_ref_supports_legacy_run_command_alias() {
411        assert_eq!(
412            normalize_tool_ref("default::run_command"),
413            Some("execute_command".to_string())
414        );
415    }
416
417    #[test]
418    fn test_normalize_tool_ref_rejects_unknown_tool() {
419        assert_eq!(normalize_tool_ref("default::search"), None);
420    }
421
422    #[test]
423    fn test_executor_has_all_builtin_tools() {
424        let executor = BuiltinToolExecutor::new();
425        let tools = executor.list_tools();
426
427        assert_eq!(tools.len(), BUILTIN_TOOL_NAMES.len());
428
429        let tool_names: Vec<String> = tools.iter().map(|t| t.function.name.clone()).collect();
430        for tool_name in BUILTIN_TOOL_NAMES {
431            assert!(tool_names.contains(&tool_name.to_string()));
432        }
433    }
434
435    #[test]
436    fn test_executor_builds_enhanced_prompt() {
437        let executor = BuiltinToolExecutor::new();
438        let prompt = executor.build_enhanced_prompt(GuideBuildContext::default());
439        assert!(prompt.contains("## Tool Usage Guidelines"));
440        assert!(prompt.contains("**read_file**"));
441    }
442
443    #[test]
444    fn test_executor_builder_empty() {
445        let executor = BuiltinToolExecutorBuilder::new().build();
446        assert!(executor.list_tools().is_empty());
447    }
448
449    #[test]
450    fn test_executor_builder_with_default_tools() {
451        let executor = BuiltinToolExecutorBuilder::new()
452            .with_default_tools()
453            .build();
454        assert_eq!(executor.list_tools().len(), BUILTIN_TOOL_NAMES.len());
455    }
456
457    #[test]
458    fn test_executor_builder_with_specific_tool() {
459        let executor = BuiltinToolExecutorBuilder::new()
460            .with_filesystem_tool("read_file")
461            .unwrap()
462            .build();
463
464        let tools = executor.list_tools();
465        assert_eq!(tools.len(), 1);
466        assert_eq!(tools[0].function.name, "read_file");
467    }
468
469    #[tokio::test]
470    async fn test_executor_skips_permission_checks_without_checker() {
471        let executor = make_executor(None);
472        let path = "/tmp/executor_permission_none.txt";
473        let _ = fs::remove_file(path).await;
474
475        let call = make_tool_call("write_file", json!({"path": path, "content": "ok"}));
476        let result = executor.execute(&call).await.expect("execute tool");
477
478        assert!(result.success);
479        let _ = fs::remove_file(path).await;
480    }
481
482    #[tokio::test]
483    async fn test_executor_with_permission_checker_enforces_checks() {
484        let checker = Arc::new(crate::agent::tools::permission::DenyDangerousPermissionChecker);
485        let executor = make_executor(Some(checker));
486        let path = "/tmp/executor_permission_denied.txt";
487        let _ = fs::remove_file(path).await;
488
489        let call = make_tool_call("write_file", json!({"path": path, "content": "nope"}));
490        let result = executor.execute(&call).await;
491
492        assert!(matches!(result, Err(ToolError::Execution(_))));
493        assert!(fs::metadata(path).await.is_err());
494    }
495
496    #[tokio::test]
497    async fn tool_can_stream_events_via_execute_with_context() {
498        struct StreamingTool;
499
500        #[async_trait]
501        impl Tool for StreamingTool {
502            fn name(&self) -> &str {
503                "streaming_tool"
504            }
505
506            fn description(&self) -> &str {
507                "streams one token"
508            }
509
510            fn parameters_schema(&self) -> serde_json::Value {
511                json!({"type":"object","properties":{}})
512            }
513
514            async fn execute(&self, _args: serde_json::Value) -> Result<ToolResult, ToolError> {
515                Ok(ToolResult {
516                    success: true,
517                    result: "ok".to_string(),
518                    display_preference: None,
519                })
520            }
521
522            async fn execute_with_context(
523                &self,
524                args: serde_json::Value,
525                ctx: ToolExecutionContext<'_>,
526            ) -> Result<ToolResult, ToolError> {
527                ctx.emit(AgentEvent::Token {
528                    content: "stream".to_string(),
529                })
530                .await;
531                self.execute(args).await
532            }
533        }
534
535        let executor = BuiltinToolExecutor::new();
536        executor
537            .register_tool(StreamingTool)
538            .expect("register streaming tool");
539
540        let (tx, mut rx) = mpsc::channel(8);
541        let call = make_tool_call("streaming_tool", json!({}));
542
543        let result = executor
544            .execute_with_context(
545                &call,
546                ToolExecutionContext {
547                    session_id: Some("s1"),
548                    tool_call_id: &call.id,
549                    event_tx: Some(&tx),
550                },
551            )
552            .await
553            .expect("execute tool");
554
555        assert!(result.success);
556        assert_eq!(result.result, "ok");
557
558        let ev = rx.recv().await.expect("expected streamed event");
559        assert!(
560            matches!(ev, AgentEvent::ToolToken { tool_call_id, content } if tool_call_id == "call_1" && content == "stream")
561        );
562    }
563}