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: 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 = ToolRegistry::new(workspace_path.clone());
108
109        // Register native Rust built-in tools
110        builtin::register_builtins(&registry);
111
112        Self {
113            workspace: workspace_path,
114            registry,
115            file_history: Arc::new(FileHistory::new(500)),
116            guard_policy: None,
117        }
118    }
119
120    pub fn set_guard_policy(&mut self, policy: Arc<dyn PermissionChecker>) {
121        self.guard_policy = Some(policy);
122    }
123
124    fn check_guard(&self, name: &str, args: &serde_json::Value) -> Result<()> {
125        if let Some(checker) = &self.guard_policy {
126            if checker.check(name, args) == PermissionDecision::Deny {
127                anyhow::bail!(
128                    "Defense-in-depth: Tool '{}' is blocked by guard permission policy",
129                    name
130                );
131            }
132        }
133        Ok(())
134    }
135
136    fn check_workspace_boundary(
137        name: &str,
138        args: &serde_json::Value,
139        ctx: &ToolContext,
140    ) -> Result<()> {
141        let path_field = match name {
142            "read" | "write" | "edit" | "patch" => Some("file_path"),
143            "ls" | "grep" | "glob" => Some("path"),
144            _ => None,
145        };
146
147        if let Some(field) = path_field {
148            if let Some(path_str) = args.get(field).and_then(|v| v.as_str()) {
149                let target = if std::path::Path::new(path_str).is_absolute() {
150                    std::path::PathBuf::from(path_str)
151                } else {
152                    ctx.workspace.join(path_str)
153                };
154
155                if let (Ok(canonical_target), Ok(canonical_workspace)) = (
156                    target.canonicalize().or_else(|_| {
157                        target
158                            .parent()
159                            .and_then(|p| p.canonicalize().ok())
160                            .ok_or_else(|| {
161                                std::io::Error::new(
162                                    std::io::ErrorKind::NotFound,
163                                    "parent not found",
164                                )
165                            })
166                    }),
167                    ctx.workspace.canonicalize(),
168                ) {
169                    if !canonical_target.starts_with(&canonical_workspace) {
170                        anyhow::bail!(
171                            "Workspace boundary violation: tool '{}' path '{}' escapes workspace '{}'",
172                            name,
173                            path_str,
174                            ctx.workspace.display()
175                        );
176                    }
177                }
178            }
179        }
180
181        Ok(())
182    }
183
184    pub fn workspace(&self) -> &PathBuf {
185        &self.workspace
186    }
187
188    pub fn registry(&self) -> &ToolRegistry {
189        &self.registry
190    }
191
192    pub fn register_dynamic_tool(&self, tool: Arc<dyn Tool>) {
193        self.registry.register(tool);
194    }
195
196    pub fn unregister_dynamic_tool(&self, name: &str) {
197        self.registry.unregister(name);
198    }
199
200    pub fn file_history(&self) -> &Arc<FileHistory> {
201        &self.file_history
202    }
203
204    fn capture_snapshot(&self, name: &str, args: &serde_json::Value) {
205        if let Some(file_path) = file_history::extract_file_path(name, args) {
206            let resolved = self.workspace.join(&file_path);
207            let path_to_read = if resolved.exists() {
208                resolved
209            } else if std::path::Path::new(&file_path).exists() {
210                std::path::PathBuf::from(&file_path)
211            } else {
212                self.file_history.save_snapshot(&file_path, "", name);
213                return;
214            };
215
216            match std::fs::read_to_string(&path_to_read) {
217                Ok(content) => {
218                    self.file_history.save_snapshot(&file_path, &content, name);
219                    tracing::debug!(
220                        "Captured file snapshot for {} before {} (version {})",
221                        file_path,
222                        name,
223                        self.file_history.list_versions(&file_path).len() - 1,
224                    );
225                }
226                Err(e) => {
227                    tracing::warn!("Failed to capture snapshot for {}: {}", file_path, e);
228                }
229            }
230        }
231    }
232
233    pub async fn execute(&self, name: &str, args: &serde_json::Value) -> Result<ToolResult> {
234        self.check_guard(name, args)?;
235        tracing::info!("Executing tool: {} with args: {}", name, args);
236        self.capture_snapshot(name, args);
237        let result = self.registry.execute(name, args).await;
238        match &result {
239            Ok(r) => tracing::info!("Tool {} completed with exit_code={}", name, r.exit_code),
240            Err(e) => tracing::error!("Tool {} failed: {}", name, e),
241        }
242        result
243    }
244
245    pub async fn execute_with_context(
246        &self,
247        name: &str,
248        args: &serde_json::Value,
249        ctx: &ToolContext,
250    ) -> Result<ToolResult> {
251        self.check_guard(name, args)?;
252        Self::check_workspace_boundary(name, args, ctx)?;
253        tracing::info!("Executing tool: {} with args: {}", name, args);
254        self.capture_snapshot(name, args);
255        let result = self.registry.execute_with_context(name, args, ctx).await;
256        match &result {
257            Ok(r) => tracing::info!("Tool {} completed with exit_code={}", name, r.exit_code),
258            Err(e) => tracing::error!("Tool {} failed: {}", name, e),
259        }
260        result
261    }
262
263    pub fn definitions(&self) -> Vec<ToolDefinition> {
264        self.registry.definitions()
265    }
266}
267
268#[cfg(test)]
269mod tests {
270    use super::*;
271
272    #[tokio::test]
273    async fn test_tool_executor_creation() {
274        let executor = ToolExecutor::new("/tmp".to_string());
275        assert_eq!(executor.registry.len(), 10);
276    }
277
278    #[tokio::test]
279    async fn test_unknown_tool() {
280        let executor = ToolExecutor::new("/tmp".to_string());
281        let result = executor
282            .execute("unknown", &serde_json::json!({}))
283            .await
284            .unwrap();
285        assert_eq!(result.exit_code, 1);
286        assert!(result.output.contains("Unknown tool"));
287    }
288
289    #[tokio::test]
290    async fn test_builtin_tools_registered() {
291        let executor = ToolExecutor::new("/tmp".to_string());
292        let definitions = executor.definitions();
293
294        assert!(definitions.iter().any(|t| t.name == "bash"));
295        assert!(definitions.iter().any(|t| t.name == "read"));
296        assert!(definitions.iter().any(|t| t.name == "write"));
297        assert!(definitions.iter().any(|t| t.name == "edit"));
298        assert!(definitions.iter().any(|t| t.name == "grep"));
299        assert!(definitions.iter().any(|t| t.name == "glob"));
300        assert!(definitions.iter().any(|t| t.name == "ls"));
301        assert!(definitions.iter().any(|t| t.name == "patch"));
302        assert!(definitions.iter().any(|t| t.name == "web_fetch"));
303        assert!(definitions.iter().any(|t| t.name == "web_search"));
304    }
305
306    #[test]
307    fn test_tool_result_success() {
308        let result = ToolResult::success("test_tool", "output text".to_string());
309        assert_eq!(result.name, "test_tool");
310        assert_eq!(result.output, "output text");
311        assert_eq!(result.exit_code, 0);
312        assert!(result.metadata.is_none());
313    }
314
315    #[test]
316    fn test_tool_result_error() {
317        let result = ToolResult::error("test_tool", "error message".to_string());
318        assert_eq!(result.name, "test_tool");
319        assert_eq!(result.output, "error message");
320        assert_eq!(result.exit_code, 1);
321        assert!(result.metadata.is_none());
322    }
323
324    #[test]
325    fn test_tool_result_from_tool_output_success() {
326        let output = ToolOutput {
327            content: "success content".to_string(),
328            success: true,
329            metadata: None,
330            images: Vec::new(),
331        };
332        let result: ToolResult = output.into();
333        assert_eq!(result.output, "success content");
334        assert_eq!(result.exit_code, 0);
335        assert!(result.metadata.is_none());
336    }
337
338    #[test]
339    fn test_tool_result_from_tool_output_failure() {
340        let output = ToolOutput {
341            content: "failure content".to_string(),
342            success: false,
343            metadata: Some(serde_json::json!({"error": "test"})),
344            images: Vec::new(),
345        };
346        let result: ToolResult = output.into();
347        assert_eq!(result.output, "failure content");
348        assert_eq!(result.exit_code, 1);
349        assert_eq!(result.metadata, Some(serde_json::json!({"error": "test"})));
350    }
351
352    #[test]
353    fn test_tool_result_metadata_propagation() {
354        let output = ToolOutput::success("content")
355            .with_metadata(serde_json::json!({"_load_skill": true, "skill_name": "test"}));
356        let result: ToolResult = output.into();
357        assert_eq!(result.exit_code, 0);
358        let meta = result.metadata.unwrap();
359        assert_eq!(meta["_load_skill"], true);
360        assert_eq!(meta["skill_name"], "test");
361    }
362
363    #[test]
364    fn test_tool_executor_workspace() {
365        let executor = ToolExecutor::new("/test/workspace".to_string());
366        assert_eq!(executor.workspace().to_str().unwrap(), "/test/workspace");
367    }
368
369    #[test]
370    fn test_tool_executor_registry() {
371        let executor = ToolExecutor::new("/tmp".to_string());
372        let registry = executor.registry();
373        assert_eq!(registry.len(), 10);
374    }
375
376    #[test]
377    fn test_tool_executor_file_history() {
378        let executor = ToolExecutor::new("/tmp".to_string());
379        let history = executor.file_history();
380        assert_eq!(history.list_versions("nonexistent.txt").len(), 0);
381    }
382
383    #[test]
384    fn test_max_output_size_constant() {
385        assert_eq!(MAX_OUTPUT_SIZE, 100 * 1024);
386    }
387
388    #[test]
389    fn test_max_read_lines_constant() {
390        assert_eq!(MAX_READ_LINES, 2000);
391    }
392
393    #[test]
394    fn test_max_line_length_constant() {
395        assert_eq!(MAX_LINE_LENGTH, 2000);
396    }
397
398    #[test]
399    fn test_tool_result_clone() {
400        let result = ToolResult::success("test", "output".to_string());
401        let cloned = result.clone();
402        assert_eq!(result.name, cloned.name);
403        assert_eq!(result.output, cloned.output);
404        assert_eq!(result.exit_code, cloned.exit_code);
405        assert_eq!(result.metadata, cloned.metadata);
406    }
407
408    #[test]
409    fn test_tool_result_debug() {
410        let result = ToolResult::success("test", "output".to_string());
411        let debug_str = format!("{:?}", result);
412        assert!(debug_str.contains("test"));
413        assert!(debug_str.contains("output"));
414    }
415}