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 artifacts;
13mod builtin;
14mod process;
15mod program_tool;
16mod registry;
17mod selector;
18pub mod skill;
19pub mod task;
20mod types;
21
22pub use artifacts::{ArtifactStore, ArtifactStoreLimits, ToolArtifact};
23pub(crate) use builtin::register_skill;
24pub use builtin::{
25    register_program, register_program_with_catalog, register_task, register_task_with_mcp,
26};
27pub use program_tool::ProgramTool;
28pub use registry::ToolRegistry;
29pub use selector::{select_tools_for_messages, select_tools_for_prompt};
30pub use task::{
31    parallel_task_params_schema, task_params_schema, ParallelTaskParams, ParallelTaskTool,
32    TaskExecutor, TaskParams, TaskResult, TaskTool,
33};
34pub use types::{Tool, ToolContext, ToolEventSender, ToolOutput, ToolStreamEvent};
35
36use crate::file_history::{self, FileHistory};
37use crate::llm::ToolDefinition;
38use crate::permissions::{PermissionChecker, PermissionDecision};
39use crate::text::truncate_utf8;
40use anyhow::Result;
41use serde::{Deserialize, Serialize};
42use std::collections::HashMap;
43use std::path::PathBuf;
44use std::sync::Arc;
45
46/// Maximum output size in bytes before truncation
47pub const MAX_OUTPUT_SIZE: usize = 100 * 1024; // 100KB
48
49/// Maximum lines to read from a file
50pub const MAX_READ_LINES: usize = 2000;
51
52/// Maximum line length before truncation
53pub const MAX_LINE_LENGTH: usize = 2000;
54
55#[derive(Debug, Clone, PartialEq, Eq)]
56pub(crate) struct ToolOutputArtifact {
57    pub artifact_id: String,
58    pub artifact_uri: String,
59    pub original_bytes: usize,
60    pub shown_bytes: usize,
61}
62
63#[derive(Debug, Clone)]
64pub(crate) struct TruncatedToolOutput {
65    pub content: String,
66    pub artifact: Option<ToolOutputArtifact>,
67}
68
69pub(crate) fn truncate_tool_output_with_artifact(
70    tool_name: &str,
71    output: &str,
72) -> TruncatedToolOutput {
73    if output.len() <= MAX_OUTPUT_SIZE {
74        return TruncatedToolOutput {
75            content: output.to_string(),
76            artifact: None,
77        };
78    }
79
80    let shown = truncate_utf8(output, MAX_OUTPUT_SIZE);
81    let artifact = tool_output_artifact(tool_name, output, shown.len());
82    let artifact_uri = artifact.artifact_uri.clone();
83    let content = format!(
84        "{}\n\n[tool output truncated: showing the first {} of {} bytes. Full output artifact: {}. Use narrower arguments such as offset/limit or filtering when possible.]",
85        shown,
86        shown.len(),
87        output.len(),
88        artifact_uri,
89    );
90
91    TruncatedToolOutput {
92        content,
93        artifact: Some(artifact),
94    }
95}
96
97pub(crate) fn tool_output_artifact(
98    tool_name: &str,
99    output: &str,
100    shown_bytes: usize,
101) -> ToolOutputArtifact {
102    use std::hash::{Hash, Hasher};
103
104    let mut hasher = std::collections::hash_map::DefaultHasher::new();
105    tool_name.hash(&mut hasher);
106    output.len().hash(&mut hasher);
107    output.hash(&mut hasher);
108    let digest = hasher.finish();
109    let sanitized_tool = tool_name
110        .chars()
111        .map(|ch| {
112            if ch.is_ascii_alphanumeric() || ch == '_' || ch == '-' {
113                ch
114            } else {
115                '_'
116            }
117        })
118        .collect::<String>();
119    let artifact_id = format!("tool-output:{sanitized_tool}:{digest:016x}");
120    let artifact_uri = format!("a3s://tool-output/{sanitized_tool}/{digest:016x}");
121
122    ToolOutputArtifact {
123        artifact_id,
124        artifact_uri,
125        original_bytes: output.len(),
126        shown_bytes,
127    }
128}
129
130pub(crate) fn merge_tool_output_artifact_metadata(
131    metadata: Option<serde_json::Value>,
132    artifact: &ToolOutputArtifact,
133) -> serde_json::Value {
134    let artifact_json = serde_json::json!({
135        "artifact_id": artifact.artifact_id,
136        "artifact_uri": artifact.artifact_uri,
137        "original_bytes": artifact.original_bytes,
138        "shown_bytes": artifact.shown_bytes,
139    });
140
141    match metadata {
142        Some(serde_json::Value::Object(mut object)) => {
143            object.insert("artifact".to_string(), artifact_json);
144            serde_json::Value::Object(object)
145        }
146        Some(value) => serde_json::json!({
147            "artifact": artifact_json,
148            "previous_metadata": value,
149        }),
150        None => serde_json::json!({
151            "artifact": artifact_json,
152        }),
153    }
154}
155
156/// Tool execution result returned by direct tool execution.
157#[derive(Debug, Clone, Serialize, Deserialize)]
158pub struct ToolResult {
159    pub name: String,
160    pub output: String,
161    pub exit_code: i32,
162    #[serde(skip_serializing_if = "Option::is_none")]
163    pub metadata: Option<serde_json::Value>,
164    /// Image attachments from tool execution (multi-modal output).
165    #[serde(skip)]
166    pub images: Vec<crate::llm::Attachment>,
167}
168
169impl ToolResult {
170    pub fn success(name: &str, output: String) -> Self {
171        Self {
172            name: name.to_string(),
173            output,
174            exit_code: 0,
175            metadata: None,
176            images: Vec::new(),
177        }
178    }
179
180    pub fn error(name: &str, message: String) -> Self {
181        Self {
182            name: name.to_string(),
183            output: message,
184            exit_code: 1,
185            metadata: None,
186            images: Vec::new(),
187        }
188    }
189}
190
191impl From<ToolOutput> for ToolResult {
192    fn from(output: ToolOutput) -> Self {
193        Self {
194            name: String::new(),
195            output: output.content,
196            exit_code: if output.success { 0 } else { 1 },
197            metadata: output.metadata,
198            images: output.images,
199        }
200    }
201}
202
203/// Tool executor with workspace sandboxing
204///
205/// This is the main entry point for tool execution. It wraps the ToolRegistry
206/// and captures file snapshots before write/edit/patch operations.
207///
208/// Defense-in-depth: An optional permission policy can be set to block
209/// denied tools even if the caller bypasses the agent loop's authorization.
210pub struct ToolExecutor {
211    workspace: PathBuf,
212    registry: Arc<ToolRegistry>,
213    file_history: Arc<FileHistory>,
214    guard_policy: Option<Arc<dyn PermissionChecker>>,
215    command_env: Option<Arc<HashMap<String, String>>>,
216}
217
218impl ToolExecutor {
219    pub fn new(workspace: String) -> Self {
220        Self::new_with_options(workspace, None, ArtifactStoreLimits::default())
221    }
222
223    pub fn new_with_command_env(workspace: String, command_env: HashMap<String, String>) -> Self {
224        Self::new_with_options(workspace, Some(command_env), ArtifactStoreLimits::default())
225    }
226
227    pub fn new_with_artifact_limits(
228        workspace: String,
229        artifact_limits: ArtifactStoreLimits,
230    ) -> Self {
231        Self::new_with_options(workspace, None, artifact_limits)
232    }
233
234    pub fn new_with_command_env_and_artifact_limits(
235        workspace: String,
236        command_env: HashMap<String, String>,
237        artifact_limits: ArtifactStoreLimits,
238    ) -> Self {
239        Self::new_with_options(workspace, Some(command_env), artifact_limits)
240    }
241
242    fn new_with_options(
243        workspace: String,
244        command_env: Option<HashMap<String, String>>,
245        artifact_limits: ArtifactStoreLimits,
246    ) -> Self {
247        let workspace_path = PathBuf::from(&workspace);
248        let registry = Arc::new(ToolRegistry::with_artifact_limits(
249            workspace_path.clone(),
250            artifact_limits,
251        ));
252
253        // Register native Rust built-in tools
254        builtin::register_builtins(&registry);
255        // Batch tool requires Arc<ToolRegistry>, registered separately
256        builtin::register_batch(&registry);
257        builtin::register_program(&registry);
258
259        Self {
260            workspace: workspace_path,
261            registry,
262            file_history: Arc::new(FileHistory::new(500)),
263            guard_policy: None,
264            command_env: command_env.map(Arc::new),
265        }
266    }
267
268    pub fn set_guard_policy(&mut self, policy: Arc<dyn PermissionChecker>) {
269        self.guard_policy = Some(policy);
270    }
271
272    fn check_guard(&self, name: &str, args: &serde_json::Value) -> Result<()> {
273        if let Some(checker) = &self.guard_policy {
274            if checker.check(name, args) == PermissionDecision::Deny {
275                anyhow::bail!(
276                    "Defense-in-depth: Tool '{}' is blocked by guard permission policy",
277                    name
278                );
279            }
280        }
281        Ok(())
282    }
283
284    fn check_workspace_boundary(
285        name: &str,
286        args: &serde_json::Value,
287        ctx: &ToolContext,
288    ) -> Result<()> {
289        let path_field = match name {
290            "read" | "write" | "edit" | "patch" => Some("file_path"),
291            "ls" | "grep" | "glob" => Some("path"),
292            _ => None,
293        };
294
295        if let Some(field) = path_field {
296            if let Some(path_str) = args.get(field).and_then(|v| v.as_str()) {
297                let target = if std::path::Path::new(path_str).is_absolute() {
298                    std::path::PathBuf::from(path_str)
299                } else {
300                    ctx.workspace.join(path_str)
301                };
302
303                // Canonicalize workspace first — fail closed if it can't be resolved
304                let canonical_workspace = ctx.workspace.canonicalize().map_err(|e| {
305                    anyhow::anyhow!(
306                        "Workspace boundary check failed: cannot canonicalize workspace '{}': {}",
307                        ctx.workspace.display(),
308                        e
309                    )
310                })?;
311
312                // Try to canonicalize target; fall back to parent directory for new files
313                let canonical_target = target.canonicalize().or_else(|_| {
314                    target
315                        .parent()
316                        .and_then(|p| p.canonicalize().ok())
317                        .ok_or_else(|| {
318                            std::io::Error::new(std::io::ErrorKind::NotFound, "parent not found")
319                        })
320                });
321
322                match canonical_target {
323                    Ok(canonical) => {
324                        if !canonical.starts_with(&canonical_workspace) {
325                            anyhow::bail!(
326                                "Workspace boundary violation: tool '{}' path '{}' escapes workspace '{}'",
327                                name,
328                                path_str,
329                                ctx.workspace.display()
330                            );
331                        }
332                    }
333                    Err(_) => {
334                        // Fail closed: if we can't resolve the target path, deny the operation
335                        anyhow::bail!(
336                            "Workspace boundary check failed: cannot resolve path '{}' for tool '{}'",
337                            path_str,
338                            name
339                        );
340                    }
341                }
342            }
343        }
344
345        Ok(())
346    }
347
348    pub fn workspace(&self) -> &PathBuf {
349        &self.workspace
350    }
351
352    pub fn registry(&self) -> &Arc<ToolRegistry> {
353        &self.registry
354    }
355
356    /// Get a stored tool artifact by URI.
357    pub fn get_artifact(&self, artifact_uri: &str) -> Option<ToolArtifact> {
358        self.registry.get_artifact(artifact_uri)
359    }
360
361    /// Return a clone of the executor's artifact store handle.
362    pub fn artifact_store(&self) -> ArtifactStore {
363        self.registry.artifact_store()
364    }
365
366    /// Replace the sink used for compact execution trace events.
367    pub fn set_trace_sink(&self, sink: Arc<dyn crate::trace::TraceSink>) {
368        self.registry.set_trace_sink(sink);
369    }
370
371    /// Return the currently configured execution trace sink.
372    pub fn trace_sink(&self) -> Arc<dyn crate::trace::TraceSink> {
373        self.registry.trace_sink()
374    }
375
376    pub fn command_env(&self) -> Option<Arc<HashMap<String, String>>> {
377        self.command_env.clone()
378    }
379
380    pub fn register_dynamic_tool(&self, tool: Arc<dyn Tool>) {
381        self.registry.register(tool);
382    }
383
384    pub fn unregister_dynamic_tool(&self, name: &str) {
385        self.registry.unregister(name);
386    }
387
388    /// Unregister all dynamic tools whose names start with the given prefix.
389    pub fn unregister_tools_by_prefix(&self, prefix: &str) {
390        self.registry.unregister_by_prefix(prefix);
391    }
392
393    /// Replace the model-visible `program` tool with a custom PTC catalog.
394    pub fn register_program_catalog(&self, catalog: crate::program::ProgramCatalog) {
395        builtin::register_program_with_catalog(&self.registry, catalog);
396    }
397
398    fn capture_snapshot(&self, name: &str, args: &serde_json::Value) {
399        if let Some(file_path) = file_history::extract_file_path(name, args) {
400            let resolved = self.workspace.join(&file_path);
401            let path_to_read = if resolved.exists() {
402                resolved
403            } else if std::path::Path::new(&file_path).exists() {
404                std::path::PathBuf::from(&file_path)
405            } else {
406                self.file_history.save_snapshot(&file_path, "", name);
407                return;
408            };
409
410            match std::fs::read_to_string(&path_to_read) {
411                Ok(content) => {
412                    self.file_history.save_snapshot(&file_path, &content, name);
413                    tracing::debug!(
414                        "Captured file snapshot for {} before {} (version {})",
415                        file_path,
416                        name,
417                        self.file_history.list_versions(&file_path).len() - 1,
418                    );
419                }
420                Err(e) => {
421                    tracing::warn!("Failed to capture snapshot for {}: {}", file_path, e);
422                }
423            }
424        }
425    }
426
427    pub async fn execute(&self, name: &str, args: &serde_json::Value) -> Result<ToolResult> {
428        self.check_guard(name, args)?;
429        tracing::info!("Executing tool: {} with args: {}", name, args);
430        self.capture_snapshot(name, args);
431        let mut result = self.registry.execute(name, args).await;
432        if let Ok(ref mut r) = result {
433            self.attach_diff_metadata(name, args, r);
434        }
435        match &result {
436            Ok(r) => tracing::info!("Tool {} completed with exit_code={}", name, r.exit_code),
437            Err(e) => tracing::error!("Tool {} failed: {}", name, e),
438        }
439        result
440    }
441
442    pub async fn execute_with_context(
443        &self,
444        name: &str,
445        args: &serde_json::Value,
446        ctx: &ToolContext,
447    ) -> Result<ToolResult> {
448        self.check_guard(name, args)?;
449        Self::check_workspace_boundary(name, args, ctx)?;
450        tracing::info!("Executing tool: {} with args: {}", name, args);
451        self.capture_snapshot(name, args);
452        let mut result = self.registry.execute_with_context(name, args, ctx).await;
453        if let Ok(ref mut r) = result {
454            self.attach_diff_metadata(name, args, r);
455        }
456        match &result {
457            Ok(r) => tracing::info!("Tool {} completed with exit_code={}", name, r.exit_code),
458            Err(e) => tracing::error!("Tool {} failed: {}", name, e),
459        }
460        result
461    }
462
463    fn attach_diff_metadata(&self, name: &str, args: &serde_json::Value, result: &mut ToolResult) {
464        if !file_history::is_file_modifying_tool(name) {
465            return;
466        }
467        let Some(file_path) = file_history::extract_file_path(name, args) else {
468            return;
469        };
470        // Only store file_path in metadata, let translate_event read the actual content
471        // using the session's correct workspace
472        let meta = result.metadata.get_or_insert_with(|| serde_json::json!({}));
473        meta["file_path"] = serde_json::Value::String(file_path);
474    }
475
476    pub fn definitions(&self) -> Vec<ToolDefinition> {
477        self.registry.definitions()
478    }
479}
480
481#[cfg(test)]
482mod tests {
483    use super::*;
484    use async_trait::async_trait;
485
486    struct LargeArtifactTool;
487
488    #[async_trait]
489    impl Tool for LargeArtifactTool {
490        fn name(&self) -> &str {
491            "large_artifact"
492        }
493
494        fn description(&self) -> &str {
495            "Produces large output for artifact API tests"
496        }
497
498        fn parameters(&self) -> serde_json::Value {
499            serde_json::json!({
500                "type": "object",
501                "additionalProperties": false,
502                "properties": {},
503                "required": []
504            })
505        }
506
507        async fn execute(
508            &self,
509            args: &serde_json::Value,
510            _ctx: &ToolContext,
511        ) -> Result<ToolOutput> {
512            let suffix = args
513                .get("suffix")
514                .and_then(|value| value.as_str())
515                .unwrap_or_default();
516            Ok(ToolOutput::success(format!(
517                "{}{}",
518                "z".repeat(MAX_OUTPUT_SIZE + 1),
519                suffix
520            )))
521        }
522    }
523
524    struct EchoTool;
525
526    #[async_trait]
527    impl Tool for EchoTool {
528        fn name(&self) -> &str {
529            "echo"
530        }
531
532        fn description(&self) -> &str {
533            "Echoes the message argument"
534        }
535
536        fn parameters(&self) -> serde_json::Value {
537            serde_json::json!({
538                "type": "object",
539                "additionalProperties": false,
540                "properties": {
541                    "message": { "type": "string" }
542                },
543                "required": ["message"]
544            })
545        }
546
547        async fn execute(
548            &self,
549            args: &serde_json::Value,
550            _ctx: &ToolContext,
551        ) -> Result<ToolOutput> {
552            Ok(ToolOutput::success(
553                args["message"].as_str().unwrap_or_default(),
554            ))
555        }
556    }
557
558    #[tokio::test]
559    async fn test_tool_executor_creation() {
560        let executor = ToolExecutor::new("/tmp".to_string());
561        // Baseline tools on a raw ToolExecutor: 13
562        assert_eq!(executor.registry.len(), 13);
563    }
564
565    #[tokio::test]
566    async fn test_unknown_tool() {
567        let executor = ToolExecutor::new("/tmp".to_string());
568        let result = executor
569            .execute("unknown", &serde_json::json!({}))
570            .await
571            .unwrap();
572        assert_eq!(result.exit_code, 1);
573        assert!(result.output.contains("Unknown tool"));
574    }
575
576    #[tokio::test]
577    async fn test_builtin_tools_registered() {
578        let executor = ToolExecutor::new("/tmp".to_string());
579        let definitions = executor.definitions();
580
581        assert!(definitions.iter().any(|t| t.name == "bash"));
582        assert!(definitions.iter().any(|t| t.name == "read"));
583        assert!(definitions.iter().any(|t| t.name == "write"));
584        assert!(definitions.iter().any(|t| t.name == "edit"));
585        assert!(definitions.iter().any(|t| t.name == "grep"));
586        assert!(definitions.iter().any(|t| t.name == "glob"));
587        assert!(definitions.iter().any(|t| t.name == "ls"));
588        assert!(definitions.iter().any(|t| t.name == "patch"));
589        assert!(definitions.iter().any(|t| t.name == "web_fetch"));
590        assert!(definitions.iter().any(|t| t.name == "web_search"));
591        assert!(definitions.iter().any(|t| t.name == "batch"));
592    }
593
594    #[test]
595    fn test_tool_result_success() {
596        let result = ToolResult::success("test_tool", "output text".to_string());
597        assert_eq!(result.name, "test_tool");
598        assert_eq!(result.output, "output text");
599        assert_eq!(result.exit_code, 0);
600        assert!(result.metadata.is_none());
601    }
602
603    #[test]
604    fn test_tool_result_error() {
605        let result = ToolResult::error("test_tool", "error message".to_string());
606        assert_eq!(result.name, "test_tool");
607        assert_eq!(result.output, "error message");
608        assert_eq!(result.exit_code, 1);
609        assert!(result.metadata.is_none());
610    }
611
612    #[test]
613    fn test_tool_result_from_tool_output_success() {
614        let output = ToolOutput {
615            content: "success content".to_string(),
616            success: true,
617            metadata: None,
618            images: Vec::new(),
619        };
620        let result: ToolResult = output.into();
621        assert_eq!(result.output, "success content");
622        assert_eq!(result.exit_code, 0);
623        assert!(result.metadata.is_none());
624    }
625
626    #[test]
627    fn test_tool_result_from_tool_output_failure() {
628        let output = ToolOutput {
629            content: "failure content".to_string(),
630            success: false,
631            metadata: Some(serde_json::json!({"error": "test"})),
632            images: Vec::new(),
633        };
634        let result: ToolResult = output.into();
635        assert_eq!(result.output, "failure content");
636        assert_eq!(result.exit_code, 1);
637        assert_eq!(result.metadata, Some(serde_json::json!({"error": "test"})));
638    }
639
640    #[test]
641    fn test_tool_result_metadata_propagation() {
642        let output = ToolOutput::success("content")
643            .with_metadata(serde_json::json!({"_load_skill": true, "skill_name": "test"}));
644        let result: ToolResult = output.into();
645        assert_eq!(result.exit_code, 0);
646        let meta = result.metadata.unwrap();
647        assert_eq!(meta["_load_skill"], true);
648        assert_eq!(meta["skill_name"], "test");
649    }
650
651    #[test]
652    fn test_tool_executor_workspace() {
653        let executor = ToolExecutor::new("/test/workspace".to_string());
654        assert_eq!(executor.workspace().to_str().unwrap(), "/test/workspace");
655    }
656
657    #[test]
658    fn test_tool_executor_registry() {
659        let executor = ToolExecutor::new("/tmp".to_string());
660        let registry = executor.registry();
661        // Baseline tools on a raw ToolExecutor: 13
662        assert_eq!(registry.len(), 13);
663    }
664
665    #[tokio::test]
666    async fn test_tool_executor_get_artifact() {
667        let executor = ToolExecutor::new("/tmp".to_string());
668        executor.register_dynamic_tool(Arc::new(LargeArtifactTool));
669
670        let result = executor
671            .execute("large_artifact", &serde_json::json!({}))
672            .await
673            .unwrap();
674
675        let artifact_uri = result.metadata.as_ref().unwrap()["artifact"]["artifact_uri"]
676            .as_str()
677            .unwrap();
678        let artifact = executor.get_artifact(artifact_uri).expect("artifact");
679        assert_eq!(artifact.tool_name, "large_artifact");
680        assert_eq!(artifact.content.len(), MAX_OUTPUT_SIZE + 1);
681        assert!(executor.artifact_store().get(artifact_uri).is_some());
682    }
683
684    #[tokio::test]
685    async fn test_tool_executor_respects_artifact_limits() {
686        let executor = ToolExecutor::new_with_artifact_limits(
687            "/tmp".to_string(),
688            ArtifactStoreLimits {
689                max_artifacts: 1,
690                max_bytes: usize::MAX,
691            },
692        );
693        executor.register_dynamic_tool(Arc::new(LargeArtifactTool));
694
695        let first = executor
696            .execute("large_artifact", &serde_json::json!({}))
697            .await
698            .unwrap();
699        let first_uri = first.metadata.as_ref().unwrap()["artifact"]["artifact_uri"]
700            .as_str()
701            .unwrap()
702            .to_string();
703
704        executor
705            .execute("large_artifact", &serde_json::json!({ "suffix": "again" }))
706            .await
707            .unwrap();
708
709        assert_eq!(executor.artifact_store().limits().max_artifacts, 1);
710        assert_eq!(executor.artifact_store().len(), 1);
711        assert!(executor.get_artifact(&first_uri).is_none());
712    }
713
714    #[tokio::test]
715    async fn test_tool_executor_register_program_catalog() {
716        let executor = ToolExecutor::new("/tmp".to_string());
717        let trace_sink = crate::trace::InMemoryTraceSink::default();
718        executor.set_trace_sink(Arc::new(trace_sink.clone()));
719        executor.register_dynamic_tool(Arc::new(EchoTool));
720        let mut catalog = crate::program::ProgramCatalog::new();
721        catalog.register(
722            crate::program::ProgramTemplate::new("custom_echo", "Run a custom echo program")
723                .with_parameter(crate::program::ProgramParameter::required(
724                    "message",
725                    "Message to echo",
726                ))
727                .with_step(
728                    crate::program::ProgramStepTemplate::new(
729                        "echo",
730                        serde_json::json!({ "message": "{{message}}" }),
731                    )
732                    .with_label("echo_message"),
733                ),
734        );
735        executor.register_program_catalog(catalog);
736
737        let result = executor
738            .execute(
739                "program",
740                &serde_json::json!({
741                    "name": "custom_echo",
742                    "inputs": {
743                        "message": "hello from catalog"
744                    }
745                }),
746            )
747            .await
748            .unwrap();
749
750        assert_eq!(result.exit_code, 0);
751        assert!(result.output.contains("hello from catalog"));
752        let metadata = result.metadata.expect("metadata");
753        assert_eq!(metadata["trace"]["program_name"], "custom_echo");
754        assert_eq!(metadata["trace"]["steps"][0]["label"], "echo_message");
755
756        let events = trace_sink.events();
757        assert!(events.iter().any(|event| {
758            event.kind == crate::trace::TraceEventKind::ToolExecution && event.name == "echo"
759        }));
760        assert!(events.iter().any(|event| {
761            event.kind == crate::trace::TraceEventKind::ToolExecution && event.name == "program"
762        }));
763        let program_event = events
764            .iter()
765            .find(|event| event.kind == crate::trace::TraceEventKind::ProgramExecution)
766            .expect("program trace event");
767        assert_eq!(
768            program_event.details.as_ref().unwrap()["program_name"],
769            "custom_echo"
770        );
771    }
772
773    #[test]
774    fn test_max_output_size_constant() {
775        assert_eq!(MAX_OUTPUT_SIZE, 100 * 1024);
776    }
777
778    #[test]
779    fn test_max_read_lines_constant() {
780        assert_eq!(MAX_READ_LINES, 2000);
781    }
782
783    #[test]
784    fn test_max_line_length_constant() {
785        assert_eq!(MAX_LINE_LENGTH, 2000);
786    }
787
788    #[test]
789    fn test_truncate_tool_output_with_artifact_reference() {
790        let output = "x".repeat(MAX_OUTPUT_SIZE + 1);
791        let truncated = truncate_tool_output_with_artifact("test/tool", &output);
792
793        let artifact = truncated.artifact.expect("artifact");
794        assert!(truncated.content.contains("Full output artifact:"));
795        assert_eq!(artifact.original_bytes, MAX_OUTPUT_SIZE + 1);
796        assert_eq!(artifact.shown_bytes, MAX_OUTPUT_SIZE);
797        assert!(artifact.artifact_id.starts_with("tool-output:test_tool:"));
798        assert!(artifact
799            .artifact_uri
800            .starts_with("a3s://tool-output/test_tool/"));
801    }
802
803    #[test]
804    fn test_tool_result_clone() {
805        let result = ToolResult::success("test", "output".to_string());
806        let cloned = result.clone();
807        assert_eq!(result.name, cloned.name);
808        assert_eq!(result.output, cloned.output);
809        assert_eq!(result.exit_code, cloned.exit_code);
810        assert_eq!(result.metadata, cloned.metadata);
811    }
812
813    #[test]
814    fn test_tool_result_debug() {
815        let result = ToolResult::success("test", "output".to_string());
816        let debug_str = format!("{:?}", result);
817        assert!(debug_str.contains("test"));
818        assert!(debug_str.contains("output"));
819    }
820
821    #[tokio::test]
822    async fn test_execute_attaches_diff_metadata() {
823        use tempfile::TempDir;
824        let dir = TempDir::new().unwrap();
825        let file = dir.path().join("hello.txt");
826        std::fs::write(&file, "before content\n").unwrap();
827
828        let executor = ToolExecutor::new(dir.path().to_str().unwrap().to_string());
829        let args = serde_json::json!({
830            "file_path": "hello.txt",
831            "content": "after content\n"
832        });
833        let result = executor.execute("write", &args).await.unwrap();
834
835        let meta = result.metadata.expect("metadata should be present");
836        assert_eq!(meta["before"], "before content\n");
837        assert_eq!(meta["after"], "after content\n");
838        assert_eq!(meta["file_path"], "hello.txt");
839    }
840
841    #[tokio::test]
842    async fn test_execute_with_context_attaches_diff_metadata() {
843        use tempfile::TempDir;
844        let dir = TempDir::new().unwrap();
845        let canonical_dir = dir.path().canonicalize().unwrap();
846        let file = canonical_dir.join("ctx.txt");
847        std::fs::write(&file, "original\n").unwrap();
848
849        let executor = ToolExecutor::new(canonical_dir.to_str().unwrap().to_string());
850        let ctx = ToolContext {
851            workspace: canonical_dir.clone(),
852            session_id: None,
853            event_tx: None,
854            agent_event_tx: None,
855            search_config: None,
856            sandbox: None,
857            command_env: None,
858        };
859        let args = serde_json::json!({
860            "file_path": "ctx.txt",
861            "content": "updated\n"
862        });
863        let result = executor
864            .execute_with_context("write", &args, &ctx)
865            .await
866            .unwrap();
867        assert_eq!(result.exit_code, 0, "write tool failed: {}", result.output);
868
869        let meta = result.metadata.expect("metadata should be present");
870        assert_eq!(meta["before"], "original\n");
871        assert_eq!(meta["after"], "updated\n");
872        assert_eq!(meta["file_path"], "ctx.txt");
873    }
874}