Skip to main content

a3s_code_core/tools/
mod.rs

1//! Extensible Tool System
2//!
3//! Provides a trait-based abstraction for tools.
4//!
5//! ## Architecture
6//!
7//! ```text
8//! ToolRegistry
9//!   └── builtin tools (bash, read, write, edit, grep, glob, ls, patch, web_fetch, web_search)
10//! ```
11
12mod builtin;
13mod process;
14mod registry;
15pub mod task;
16mod types;
17
18pub use registry::ToolRegistry;
19pub use task::{
20    parallel_task_params_schema, task_params_schema, ParallelTaskParams, ParallelTaskTool,
21    TaskExecutor, TaskParams, TaskResult,
22};
23pub use types::{Tool, ToolContext, ToolEventSender, ToolOutput, ToolStreamEvent};
24
25use crate::file_history::{self, FileHistory};
26use crate::llm::ToolDefinition;
27use crate::permissions::{PermissionChecker, PermissionDecision};
28use anyhow::Result;
29use serde::{Deserialize, Serialize};
30use std::path::PathBuf;
31use std::sync::Arc;
32
33/// Maximum output size in bytes before truncation
34pub const MAX_OUTPUT_SIZE: usize = 100 * 1024; // 100KB
35
36/// Maximum lines to read from a file
37pub const MAX_READ_LINES: usize = 2000;
38
39/// Maximum line length before truncation
40pub const MAX_LINE_LENGTH: usize = 2000;
41
42/// Tool execution result (legacy format for backward compatibility)
43#[derive(Debug, Clone, Serialize, Deserialize)]
44pub struct ToolResult {
45    pub name: String,
46    pub output: String,
47    pub exit_code: i32,
48    #[serde(skip_serializing_if = "Option::is_none")]
49    pub metadata: Option<serde_json::Value>,
50    /// Image attachments from tool execution (multi-modal output).
51    #[serde(skip)]
52    pub images: Vec<crate::llm::Attachment>,
53}
54
55impl ToolResult {
56    pub fn success(name: &str, output: String) -> Self {
57        Self {
58            name: name.to_string(),
59            output,
60            exit_code: 0,
61            metadata: None,
62            images: Vec::new(),
63        }
64    }
65
66    pub fn error(name: &str, message: String) -> Self {
67        Self {
68            name: name.to_string(),
69            output: message,
70            exit_code: 1,
71            metadata: None,
72            images: Vec::new(),
73        }
74    }
75}
76
77impl From<ToolOutput> for ToolResult {
78    fn from(output: ToolOutput) -> Self {
79        Self {
80            name: String::new(),
81            output: output.content,
82            exit_code: if output.success { 0 } else { 1 },
83            metadata: output.metadata,
84            images: output.images,
85        }
86    }
87}
88
89/// Tool executor with workspace sandboxing
90///
91/// This is the main entry point for tool execution. It wraps the ToolRegistry
92/// and provides backward-compatible API. Includes file version history tracking
93/// for write/edit/patch operations.
94///
95/// Defense-in-depth: An optional permission policy can be set to block
96/// denied tools even if the caller bypasses the agent loop's authorization.
97pub struct ToolExecutor {
98    workspace: PathBuf,
99    registry: Arc<ToolRegistry>,
100    file_history: Arc<FileHistory>,
101    guard_policy: Option<Arc<dyn PermissionChecker>>,
102}
103
104impl ToolExecutor {
105    pub fn new(workspace: String) -> Self {
106        let workspace_path = PathBuf::from(&workspace);
107        let registry = Arc::new(ToolRegistry::new(workspace_path.clone()));
108
109        // Register native Rust built-in tools
110        builtin::register_builtins(&registry);
111        // Batch tool requires Arc<ToolRegistry>, registered separately
112        builtin::register_batch(&registry);
113
114        Self {
115            workspace: workspace_path,
116            registry,
117            file_history: Arc::new(FileHistory::new(500)),
118            guard_policy: None,
119        }
120    }
121
122    pub fn set_guard_policy(&mut self, policy: Arc<dyn PermissionChecker>) {
123        self.guard_policy = Some(policy);
124    }
125
126    fn check_guard(&self, name: &str, args: &serde_json::Value) -> Result<()> {
127        if let Some(checker) = &self.guard_policy {
128            if checker.check(name, args) == PermissionDecision::Deny {
129                anyhow::bail!(
130                    "Defense-in-depth: Tool '{}' is blocked by guard permission policy",
131                    name
132                );
133            }
134        }
135        Ok(())
136    }
137
138    fn check_workspace_boundary(
139        name: &str,
140        args: &serde_json::Value,
141        ctx: &ToolContext,
142    ) -> Result<()> {
143        let path_field = match name {
144            "read" | "write" | "edit" | "patch" => Some("file_path"),
145            "ls" | "grep" | "glob" => Some("path"),
146            _ => None,
147        };
148
149        if let Some(field) = path_field {
150            if let Some(path_str) = args.get(field).and_then(|v| v.as_str()) {
151                let target = if std::path::Path::new(path_str).is_absolute() {
152                    std::path::PathBuf::from(path_str)
153                } else {
154                    ctx.workspace.join(path_str)
155                };
156
157                if let (Ok(canonical_target), Ok(canonical_workspace)) = (
158                    target.canonicalize().or_else(|_| {
159                        target
160                            .parent()
161                            .and_then(|p| p.canonicalize().ok())
162                            .ok_or_else(|| {
163                                std::io::Error::new(
164                                    std::io::ErrorKind::NotFound,
165                                    "parent not found",
166                                )
167                            })
168                    }),
169                    ctx.workspace.canonicalize(),
170                ) {
171                    if !canonical_target.starts_with(&canonical_workspace) {
172                        anyhow::bail!(
173                            "Workspace boundary violation: tool '{}' path '{}' escapes workspace '{}'",
174                            name,
175                            path_str,
176                            ctx.workspace.display()
177                        );
178                    }
179                }
180            }
181        }
182
183        Ok(())
184    }
185
186    pub fn workspace(&self) -> &PathBuf {
187        &self.workspace
188    }
189
190    pub fn registry(&self) -> &Arc<ToolRegistry> {
191        &self.registry
192    }
193
194    pub fn register_dynamic_tool(&self, tool: Arc<dyn Tool>) {
195        self.registry.register(tool);
196    }
197
198    pub fn unregister_dynamic_tool(&self, name: &str) {
199        self.registry.unregister(name);
200    }
201
202    pub fn file_history(&self) -> &Arc<FileHistory> {
203        &self.file_history
204    }
205
206    fn capture_snapshot(&self, name: &str, args: &serde_json::Value) {
207        if let Some(file_path) = file_history::extract_file_path(name, args) {
208            let resolved = self.workspace.join(&file_path);
209            let path_to_read = if resolved.exists() {
210                resolved
211            } else if std::path::Path::new(&file_path).exists() {
212                std::path::PathBuf::from(&file_path)
213            } else {
214                self.file_history.save_snapshot(&file_path, "", name);
215                return;
216            };
217
218            match std::fs::read_to_string(&path_to_read) {
219                Ok(content) => {
220                    self.file_history.save_snapshot(&file_path, &content, name);
221                    tracing::debug!(
222                        "Captured file snapshot for {} before {} (version {})",
223                        file_path,
224                        name,
225                        self.file_history.list_versions(&file_path).len() - 1,
226                    );
227                }
228                Err(e) => {
229                    tracing::warn!("Failed to capture snapshot for {}: {}", file_path, e);
230                }
231            }
232        }
233    }
234
235    pub async fn execute(&self, name: &str, args: &serde_json::Value) -> Result<ToolResult> {
236        self.check_guard(name, args)?;
237        tracing::info!("Executing tool: {} with args: {}", name, args);
238        self.capture_snapshot(name, args);
239        let result = self.registry.execute(name, args).await;
240        match &result {
241            Ok(r) => tracing::info!("Tool {} completed with exit_code={}", name, r.exit_code),
242            Err(e) => tracing::error!("Tool {} failed: {}", name, e),
243        }
244        result
245    }
246
247    pub async fn execute_with_context(
248        &self,
249        name: &str,
250        args: &serde_json::Value,
251        ctx: &ToolContext,
252    ) -> Result<ToolResult> {
253        self.check_guard(name, args)?;
254        Self::check_workspace_boundary(name, args, ctx)?;
255        tracing::info!("Executing tool: {} with args: {}", name, args);
256        self.capture_snapshot(name, args);
257        let result = self.registry.execute_with_context(name, args, ctx).await;
258        match &result {
259            Ok(r) => tracing::info!("Tool {} completed with exit_code={}", name, r.exit_code),
260            Err(e) => tracing::error!("Tool {} failed: {}", name, e),
261        }
262        result
263    }
264
265    pub fn definitions(&self) -> Vec<ToolDefinition> {
266        self.registry.definitions()
267    }
268}
269
270#[cfg(test)]
271mod tests {
272    use super::*;
273
274    #[tokio::test]
275    async fn test_tool_executor_creation() {
276        let executor = ToolExecutor::new("/tmp".to_string());
277        assert_eq!(executor.registry.len(), 12);
278    }
279
280    #[tokio::test]
281    async fn test_unknown_tool() {
282        let executor = ToolExecutor::new("/tmp".to_string());
283        let result = executor
284            .execute("unknown", &serde_json::json!({}))
285            .await
286            .unwrap();
287        assert_eq!(result.exit_code, 1);
288        assert!(result.output.contains("Unknown tool"));
289    }
290
291    #[tokio::test]
292    async fn test_builtin_tools_registered() {
293        let executor = ToolExecutor::new("/tmp".to_string());
294        let definitions = executor.definitions();
295
296        assert!(definitions.iter().any(|t| t.name == "bash"));
297        assert!(definitions.iter().any(|t| t.name == "read"));
298        assert!(definitions.iter().any(|t| t.name == "write"));
299        assert!(definitions.iter().any(|t| t.name == "edit"));
300        assert!(definitions.iter().any(|t| t.name == "grep"));
301        assert!(definitions.iter().any(|t| t.name == "glob"));
302        assert!(definitions.iter().any(|t| t.name == "ls"));
303        assert!(definitions.iter().any(|t| t.name == "patch"));
304        assert!(definitions.iter().any(|t| t.name == "web_fetch"));
305        assert!(definitions.iter().any(|t| t.name == "web_search"));
306        assert!(definitions.iter().any(|t| t.name == "batch"));
307    }
308
309    #[test]
310    fn test_tool_result_success() {
311        let result = ToolResult::success("test_tool", "output text".to_string());
312        assert_eq!(result.name, "test_tool");
313        assert_eq!(result.output, "output text");
314        assert_eq!(result.exit_code, 0);
315        assert!(result.metadata.is_none());
316    }
317
318    #[test]
319    fn test_tool_result_error() {
320        let result = ToolResult::error("test_tool", "error message".to_string());
321        assert_eq!(result.name, "test_tool");
322        assert_eq!(result.output, "error message");
323        assert_eq!(result.exit_code, 1);
324        assert!(result.metadata.is_none());
325    }
326
327    #[test]
328    fn test_tool_result_from_tool_output_success() {
329        let output = ToolOutput {
330            content: "success content".to_string(),
331            success: true,
332            metadata: None,
333            images: Vec::new(),
334        };
335        let result: ToolResult = output.into();
336        assert_eq!(result.output, "success content");
337        assert_eq!(result.exit_code, 0);
338        assert!(result.metadata.is_none());
339    }
340
341    #[test]
342    fn test_tool_result_from_tool_output_failure() {
343        let output = ToolOutput {
344            content: "failure content".to_string(),
345            success: false,
346            metadata: Some(serde_json::json!({"error": "test"})),
347            images: Vec::new(),
348        };
349        let result: ToolResult = output.into();
350        assert_eq!(result.output, "failure content");
351        assert_eq!(result.exit_code, 1);
352        assert_eq!(result.metadata, Some(serde_json::json!({"error": "test"})));
353    }
354
355    #[test]
356    fn test_tool_result_metadata_propagation() {
357        let output = ToolOutput::success("content")
358            .with_metadata(serde_json::json!({"_load_skill": true, "skill_name": "test"}));
359        let result: ToolResult = output.into();
360        assert_eq!(result.exit_code, 0);
361        let meta = result.metadata.unwrap();
362        assert_eq!(meta["_load_skill"], true);
363        assert_eq!(meta["skill_name"], "test");
364    }
365
366    #[test]
367    fn test_tool_executor_workspace() {
368        let executor = ToolExecutor::new("/test/workspace".to_string());
369        assert_eq!(executor.workspace().to_str().unwrap(), "/test/workspace");
370    }
371
372    #[test]
373    fn test_tool_executor_registry() {
374        let executor = ToolExecutor::new("/tmp".to_string());
375        let registry = executor.registry();
376        assert_eq!(registry.len(), 12);
377    }
378
379    #[test]
380    fn test_tool_executor_file_history() {
381        let executor = ToolExecutor::new("/tmp".to_string());
382        let history = executor.file_history();
383        assert_eq!(history.list_versions("nonexistent.txt").len(), 0);
384    }
385
386    #[test]
387    fn test_max_output_size_constant() {
388        assert_eq!(MAX_OUTPUT_SIZE, 100 * 1024);
389    }
390
391    #[test]
392    fn test_max_read_lines_constant() {
393        assert_eq!(MAX_READ_LINES, 2000);
394    }
395
396    #[test]
397    fn test_max_line_length_constant() {
398        assert_eq!(MAX_LINE_LENGTH, 2000);
399    }
400
401    #[test]
402    fn test_tool_result_clone() {
403        let result = ToolResult::success("test", "output".to_string());
404        let cloned = result.clone();
405        assert_eq!(result.name, cloned.name);
406        assert_eq!(result.output, cloned.output);
407        assert_eq!(result.exit_code, cloned.exit_code);
408        assert_eq!(result.metadata, cloned.metadata);
409    }
410
411    #[test]
412    fn test_tool_result_debug() {
413        let result = ToolResult::success("test", "output".to_string());
414        let debug_str = format!("{:?}", result);
415        assert!(debug_str.contains("test"));
416        assert!(debug_str.contains("output"));
417    }
418}