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;
13pub(crate) mod builtin;
14pub(crate) mod 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_generate_object, register_program, register_program_with_catalog, register_task,
26    register_task_with_mcp,
27};
28pub use program_tool::ProgramTool;
29pub use registry::ToolRegistry;
30pub use selector::{select_tools_for_messages, select_tools_for_prompt};
31pub use task::{
32    parallel_task_params_schema, task_params_schema, ParallelTaskParams, ParallelTaskTool,
33    TaskExecutor, TaskParams, TaskResult, TaskTool,
34};
35pub use types::{Tool, ToolContext, ToolErrorKind, ToolEventSender, ToolOutput, ToolStreamEvent};
36
37use crate::file_history::{self, FileHistory};
38use crate::llm::ToolDefinition;
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    /// Structured discriminant for tool failures. Populated by built-in
168    /// tools that can map their failure into a typed [`ToolErrorKind`]
169    /// (e.g. `edit`/`patch` setting `VersionConflict` on a CAS rejection
170    /// from `WorkspaceError`). Forwarded to the SDK so callers can react
171    /// programmatically without parsing `output`.
172    #[serde(skip_serializing_if = "Option::is_none")]
173    pub error_kind: Option<types::ToolErrorKind>,
174}
175
176impl ToolResult {
177    pub fn success(name: &str, output: String) -> Self {
178        Self {
179            name: name.to_string(),
180            output,
181            exit_code: 0,
182            metadata: None,
183            images: Vec::new(),
184            error_kind: None,
185        }
186    }
187
188    pub fn error(name: &str, message: String) -> Self {
189        Self {
190            name: name.to_string(),
191            output: message,
192            exit_code: 1,
193            metadata: None,
194            images: Vec::new(),
195            error_kind: None,
196        }
197    }
198}
199
200impl From<ToolOutput> for ToolResult {
201    fn from(output: ToolOutput) -> Self {
202        Self {
203            name: String::new(),
204            output: output.content,
205            exit_code: if output.success { 0 } else { 1 },
206            metadata: output.metadata,
207            images: output.images,
208            error_kind: output.error_kind,
209        }
210    }
211}
212
213/// Tool executor with workspace sandboxing
214///
215/// This is the main entry point for tool execution. It wraps the ToolRegistry
216/// and captures file snapshots before write/edit/patch operations.
217pub struct ToolExecutor {
218    workspace: PathBuf,
219    registry: Arc<ToolRegistry>,
220    file_history: Arc<FileHistory>,
221    command_env: Option<Arc<HashMap<String, String>>>,
222    workspace_services: Arc<crate::workspace::WorkspaceServices>,
223}
224
225/// Build a log line for a tool invocation that excludes argument *values*.
226///
227/// Argument values (full bash commands, file contents written by `write`/`edit`)
228/// can contain secrets, so the summary records only the tool name, the sorted
229/// argument field names, and the serialized payload size — never the values. This
230/// keeps the always-on `info!` tool trace (also exported to OTLP) compliant with
231/// the "never log secrets" boundary. Use `trace!` for full args when debugging.
232fn redacted_tool_log_summary(name: &str, args: &serde_json::Value) -> String {
233    let arg_keys: Vec<&str> = match args.as_object() {
234        Some(map) => {
235            let mut keys: Vec<&str> = map.keys().map(String::as_str).collect();
236            keys.sort_unstable();
237            keys
238        }
239        None => Vec::new(),
240    };
241    format!(
242        "Executing tool: {} (arg_keys={:?}, {} bytes)",
243        name,
244        arg_keys,
245        args.to_string().len()
246    )
247}
248
249/// Log a tool invocation without leaking argument values. See
250/// [`redacted_tool_log_summary`] for the redaction rationale.
251fn log_tool_invocation(name: &str, args: &serde_json::Value) {
252    tracing::info!("{}", redacted_tool_log_summary(name, args));
253    tracing::trace!("Tool {} full args: {}", name, args);
254}
255
256impl ToolExecutor {
257    pub fn new(workspace: String) -> Self {
258        let workspace_services =
259            crate::workspace::WorkspaceServices::local(PathBuf::from(&workspace));
260        Self::build(
261            workspace,
262            None,
263            ArtifactStoreLimits::default(),
264            workspace_services,
265        )
266    }
267
268    pub fn new_with_artifact_limits(
269        workspace: String,
270        artifact_limits: ArtifactStoreLimits,
271    ) -> Self {
272        let workspace_services =
273            crate::workspace::WorkspaceServices::local(PathBuf::from(&workspace));
274        Self::build(workspace, None, artifact_limits, workspace_services)
275    }
276
277    pub fn new_with_workspace_services(
278        workspace: String,
279        workspace_services: Arc<crate::workspace::WorkspaceServices>,
280    ) -> Self {
281        Self::build(
282            workspace,
283            None,
284            ArtifactStoreLimits::default(),
285            workspace_services,
286        )
287    }
288
289    pub fn new_with_workspace_services_and_artifact_limits(
290        workspace: String,
291        workspace_services: Arc<crate::workspace::WorkspaceServices>,
292        artifact_limits: ArtifactStoreLimits,
293    ) -> Self {
294        Self::build(workspace, None, artifact_limits, workspace_services)
295    }
296
297    fn build(
298        workspace: String,
299        command_env: Option<HashMap<String, String>>,
300        artifact_limits: ArtifactStoreLimits,
301        workspace_services: Arc<crate::workspace::WorkspaceServices>,
302    ) -> Self {
303        let workspace_path = PathBuf::from(&workspace);
304        let command_env = command_env.map(Arc::new);
305        let registry = Arc::new(ToolRegistry::with_artifact_limits_and_workspace_services(
306            workspace_path.clone(),
307            artifact_limits,
308            Arc::clone(&workspace_services),
309        ));
310        if let Some(env) = command_env.clone() {
311            registry.set_command_env(env);
312        }
313
314        // Register native Rust built-in tools — only those whose required
315        // workspace capability is available, so the model never sees a tool
316        // the backend cannot service.
317        builtin::register_builtins(&registry, &workspace_services.capabilities());
318        // Batch tool requires Arc<ToolRegistry>, registered separately
319        builtin::register_batch(&registry);
320        builtin::register_program(&registry);
321
322        Self {
323            workspace: workspace_path,
324            registry,
325            file_history: Arc::new(FileHistory::new(500)),
326            command_env,
327            workspace_services,
328        }
329    }
330
331    fn check_workspace_boundary(
332        name: &str,
333        args: &serde_json::Value,
334        ctx: &ToolContext,
335    ) -> Result<()> {
336        let path_field = match name {
337            "read" | "write" | "edit" | "patch" => Some("file_path"),
338            "ls" | "grep" | "glob" => Some("path"),
339            _ => None,
340        };
341
342        if let Some(field) = path_field {
343            if let Some(path_str) = args.get(field).and_then(|v| v.as_str()) {
344                ctx.resolve_workspace_path(path_str).map_err(|e| {
345                    anyhow::anyhow!(
346                        "Workspace boundary check failed for tool '{}' path '{}': {}",
347                        name,
348                        path_str,
349                        e
350                    )
351                })?;
352            }
353        }
354
355        Ok(())
356    }
357
358    pub fn workspace(&self) -> &PathBuf {
359        &self.workspace
360    }
361
362    pub fn registry(&self) -> &Arc<ToolRegistry> {
363        &self.registry
364    }
365
366    /// Get a stored tool artifact by URI.
367    pub fn get_artifact(&self, artifact_uri: &str) -> Option<ToolArtifact> {
368        self.registry.get_artifact(artifact_uri)
369    }
370
371    /// Return a clone of the executor's artifact store handle.
372    pub fn artifact_store(&self) -> ArtifactStore {
373        self.registry.artifact_store()
374    }
375
376    /// Replace the sink used for compact execution trace events.
377    pub fn set_trace_sink(&self, sink: Arc<dyn crate::trace::TraceSink>) {
378        self.registry.set_trace_sink(sink);
379    }
380
381    /// Return the currently configured execution trace sink.
382    pub fn trace_sink(&self) -> Arc<dyn crate::trace::TraceSink> {
383        self.registry.trace_sink()
384    }
385
386    pub fn command_env(&self) -> Option<Arc<HashMap<String, String>>> {
387        self.command_env.clone()
388    }
389
390    pub fn register_dynamic_tool(&self, tool: Arc<dyn Tool>) {
391        self.registry.register(tool);
392    }
393
394    pub fn unregister_dynamic_tool(&self, name: &str) {
395        self.registry.unregister(name);
396    }
397
398    /// Unregister all dynamic tools whose names start with the given prefix.
399    pub fn unregister_tools_by_prefix(&self, prefix: &str) {
400        self.registry.unregister_by_prefix(prefix);
401    }
402
403    /// Replace the model-visible `program` tool with a custom PTC catalog.
404    pub fn register_program_catalog(&self, catalog: crate::program::ProgramCatalog) {
405        builtin::register_program_with_catalog(&self.registry, catalog);
406    }
407
408    fn capture_snapshot(&self, name: &str, args: &serde_json::Value) {
409        let Some(local_root) = self.workspace_services.local_root() else {
410            return;
411        };
412
413        if let Some(file_path) = file_history::extract_file_path(name, args) {
414            let workspace_path = match self.workspace_services.normalize_path(&file_path) {
415                Ok(path) => path,
416                Err(e) => {
417                    tracing::warn!(
418                        "Skipping file snapshot for invalid path {}: {}",
419                        file_path,
420                        e
421                    );
422                    return;
423                }
424            };
425            let path_to_read = if workspace_path.is_root() {
426                local_root.to_path_buf()
427            } else {
428                local_root.join(workspace_path.as_str())
429            };
430
431            if !path_to_read.exists() {
432                self.file_history.save_snapshot(&file_path, "", name);
433                return;
434            }
435
436            match std::fs::read_to_string(&path_to_read) {
437                Ok(content) => {
438                    self.file_history.save_snapshot(&file_path, &content, name);
439                    tracing::debug!(
440                        "Captured file snapshot for {} before {} (version {})",
441                        file_path,
442                        name,
443                        self.file_history.list_versions(&file_path).len() - 1,
444                    );
445                }
446                Err(e) => {
447                    tracing::warn!("Failed to capture snapshot for {}: {}", file_path, e);
448                }
449            }
450        }
451    }
452
453    pub async fn execute(&self, name: &str, args: &serde_json::Value) -> Result<ToolResult> {
454        let ctx = self.registry.context();
455        if let Err(e) = Self::check_workspace_boundary(name, args, &ctx) {
456            return Ok(ToolResult::error(name, e.to_string()));
457        }
458
459        log_tool_invocation(name, args);
460        self.capture_snapshot(name, args);
461        let mut result = self.registry.execute_with_context(name, args, &ctx).await;
462        if let Ok(ref mut r) = result {
463            self.attach_diff_metadata(name, args, r);
464        }
465        match &result {
466            Ok(r) => tracing::info!("Tool {} completed with exit_code={}", name, r.exit_code),
467            Err(e) => tracing::error!("Tool {} failed: {}", name, e),
468        }
469        result
470    }
471
472    pub async fn execute_with_context(
473        &self,
474        name: &str,
475        args: &serde_json::Value,
476        ctx: &ToolContext,
477    ) -> Result<ToolResult> {
478        Self::check_workspace_boundary(name, args, ctx)?;
479        log_tool_invocation(name, args);
480        self.capture_snapshot(name, args);
481        let mut result = self.registry.execute_with_context(name, args, ctx).await;
482        if let Ok(ref mut r) = result {
483            self.attach_diff_metadata(name, args, r);
484        }
485        match &result {
486            Ok(r) => tracing::info!("Tool {} completed with exit_code={}", name, r.exit_code),
487            Err(e) => tracing::error!("Tool {} failed: {}", name, e),
488        }
489        result
490    }
491
492    fn attach_diff_metadata(&self, name: &str, args: &serde_json::Value, result: &mut ToolResult) {
493        if !file_history::is_file_modifying_tool(name) {
494            return;
495        }
496        let Some(file_path) = file_history::extract_file_path(name, args) else {
497            return;
498        };
499        // Only store file_path in metadata, let translate_event read the actual content
500        // using the session's correct workspace
501        let meta = result.metadata.get_or_insert_with(|| serde_json::json!({}));
502        meta["file_path"] = serde_json::Value::String(file_path);
503    }
504
505    pub fn definitions(&self) -> Vec<ToolDefinition> {
506        self.registry.definitions()
507    }
508}
509
510#[cfg(test)]
511mod tests {
512    use super::*;
513    use crate::workspace::{
514        CommandOutput, CommandRequest, WorkspaceCommandRunner, WorkspaceDirEntry, WorkspaceError,
515        WorkspaceFileSystem, WorkspaceFileType, WorkspacePath, WorkspaceRef, WorkspaceResult,
516        WorkspaceServices, WorkspaceWriteOutcome,
517    };
518    use async_trait::async_trait;
519    use std::sync::RwLock;
520
521    #[test]
522    fn test_redacted_tool_log_summary_omits_values() {
523        let args = serde_json::json!({
524            "command": "export AWS_SECRET_ACCESS_KEY=AKIAIOSFODNN7EXAMPLE && deploy",
525            "timeout": 30
526        });
527        let summary = redacted_tool_log_summary("bash", &args);
528        // Field names and size are logged...
529        assert!(summary.contains("bash"));
530        assert!(summary.contains("command"));
531        assert!(summary.contains("timeout"));
532        assert!(summary.contains("bytes"));
533        // ...but never the values (the secret must not appear).
534        assert!(!summary.contains("AKIAIOSFODNN7EXAMPLE"));
535        assert!(!summary.contains("deploy"));
536    }
537
538    #[test]
539    fn test_redacted_tool_log_summary_handles_non_object_args() {
540        let summary = redacted_tool_log_summary("noop", &serde_json::json!("raw string"));
541        assert!(summary.contains("noop"));
542        assert!(summary.contains("arg_keys=[]"));
543        assert!(!summary.contains("raw string"));
544    }
545
546    struct LargeArtifactTool;
547
548    #[async_trait]
549    impl Tool for LargeArtifactTool {
550        fn name(&self) -> &str {
551            "large_artifact"
552        }
553
554        fn description(&self) -> &str {
555            "Produces large output for artifact API tests"
556        }
557
558        fn parameters(&self) -> serde_json::Value {
559            serde_json::json!({
560                "type": "object",
561                "additionalProperties": false,
562                "properties": {},
563                "required": []
564            })
565        }
566
567        async fn execute(
568            &self,
569            args: &serde_json::Value,
570            _ctx: &ToolContext,
571        ) -> Result<ToolOutput> {
572            let suffix = args
573                .get("suffix")
574                .and_then(|value| value.as_str())
575                .unwrap_or_default();
576            Ok(ToolOutput::success(format!(
577                "{}{}",
578                "z".repeat(MAX_OUTPUT_SIZE + 1),
579                suffix
580            )))
581        }
582    }
583
584    struct EchoTool;
585
586    #[async_trait]
587    impl Tool for EchoTool {
588        fn name(&self) -> &str {
589            "echo"
590        }
591
592        fn description(&self) -> &str {
593            "Echoes the message argument"
594        }
595
596        fn parameters(&self) -> serde_json::Value {
597            serde_json::json!({
598                "type": "object",
599                "additionalProperties": false,
600                "properties": {
601                    "message": { "type": "string" }
602                },
603                "required": ["message"]
604            })
605        }
606
607        async fn execute(
608            &self,
609            args: &serde_json::Value,
610            _ctx: &ToolContext,
611        ) -> Result<ToolOutput> {
612            Ok(ToolOutput::success(
613                args["message"].as_str().unwrap_or_default(),
614            ))
615        }
616    }
617
618    #[derive(Default)]
619    struct MemoryWorkspaceFs {
620        files: RwLock<HashMap<String, String>>,
621    }
622
623    impl MemoryWorkspaceFs {
624        fn insert(&self, path: &str, content: &str) {
625            self.files
626                .write()
627                .unwrap()
628                .insert(path.to_string(), content.to_string());
629        }
630
631        fn get(&self, path: &str) -> Option<String> {
632            self.files.read().unwrap().get(path).cloned()
633        }
634    }
635
636    #[async_trait]
637    impl WorkspaceFileSystem for MemoryWorkspaceFs {
638        async fn read_text(&self, path: &WorkspacePath) -> WorkspaceResult<String> {
639            self.files
640                .read()
641                .unwrap()
642                .get(path.as_str())
643                .cloned()
644                .ok_or_else(|| WorkspaceError::NotFound {
645                    path: path.as_str().to_string(),
646                })
647        }
648
649        async fn write_text(
650            &self,
651            path: &WorkspacePath,
652            content: &str,
653        ) -> WorkspaceResult<WorkspaceWriteOutcome> {
654            self.insert(path.as_str(), content);
655            Ok(WorkspaceWriteOutcome {
656                bytes: content.len(),
657                lines: content.lines().count(),
658            })
659        }
660
661        async fn list_dir(&self, path: &WorkspacePath) -> WorkspaceResult<Vec<WorkspaceDirEntry>> {
662            let prefix = if path.is_root() {
663                String::new()
664            } else {
665                format!("{}/", path.as_str())
666            };
667            let files = self.files.read().unwrap();
668            let mut entries = Vec::new();
669            for name in files.keys() {
670                if !name.starts_with(&prefix) {
671                    continue;
672                }
673                let remaining = &name[prefix.len()..];
674                if remaining.is_empty() || remaining.contains('/') {
675                    continue;
676                }
677                entries.push(WorkspaceDirEntry {
678                    name: remaining.to_string(),
679                    kind: WorkspaceFileType::File,
680                    size: files
681                        .get(name)
682                        .map(|content| content.len() as u64)
683                        .unwrap_or(0),
684                });
685            }
686            Ok(entries)
687        }
688    }
689
690    struct MockCommandRunner;
691
692    #[async_trait]
693    impl WorkspaceCommandRunner for MockCommandRunner {
694        async fn exec(&self, request: CommandRequest) -> Result<CommandOutput> {
695            Ok(CommandOutput {
696                output: format!("remote: {}\n", request.command),
697                exit_code: 0,
698                timed_out: false,
699            })
700        }
701    }
702
703    #[tokio::test]
704    async fn test_tool_executor_creation() {
705        let executor = ToolExecutor::new("/tmp".to_string());
706        // Baseline tools on a raw ToolExecutor: 13
707        assert_eq!(executor.registry.len(), 13);
708    }
709
710    #[tokio::test]
711    async fn test_unknown_tool() {
712        let executor = ToolExecutor::new("/tmp".to_string());
713        let result = executor
714            .execute("unknown", &serde_json::json!({}))
715            .await
716            .unwrap();
717        assert_eq!(result.exit_code, 1);
718        assert!(result.output.contains("Unknown tool"));
719    }
720
721    #[tokio::test]
722    async fn test_builtin_tools_registered() {
723        let executor = ToolExecutor::new("/tmp".to_string());
724        let definitions = executor.definitions();
725
726        assert!(definitions.iter().any(|t| t.name == "bash"));
727        assert!(definitions.iter().any(|t| t.name == "read"));
728        assert!(definitions.iter().any(|t| t.name == "write"));
729        assert!(definitions.iter().any(|t| t.name == "edit"));
730        assert!(definitions.iter().any(|t| t.name == "grep"));
731        assert!(definitions.iter().any(|t| t.name == "glob"));
732        assert!(definitions.iter().any(|t| t.name == "ls"));
733        assert!(definitions.iter().any(|t| t.name == "patch"));
734        assert!(definitions.iter().any(|t| t.name == "web_fetch"));
735        assert!(definitions.iter().any(|t| t.name == "web_search"));
736        assert!(definitions.iter().any(|t| t.name == "batch"));
737    }
738
739    #[tokio::test]
740    async fn test_builtin_file_tools_use_workspace_services() {
741        let fs = Arc::new(MemoryWorkspaceFs::default());
742        fs.insert("remote.txt", "first\nsecond\n");
743        let services = WorkspaceServices::builder(
744            WorkspaceRef::new("browser-workspace", "browser://workspace"),
745            fs.clone(),
746        )
747        .build();
748        let executor = ToolExecutor::new_with_workspace_services_and_artifact_limits(
749            "/server/local-placeholder".to_string(),
750            services,
751            ArtifactStoreLimits::default(),
752        );
753        let definitions = executor.definitions();
754        assert!(definitions.iter().any(|tool| tool.name == "read"));
755        assert!(definitions.iter().any(|tool| tool.name == "write"));
756        assert!(definitions.iter().any(|tool| tool.name == "ls"));
757        assert!(!definitions.iter().any(|tool| tool.name == "bash"));
758        assert!(!definitions.iter().any(|tool| tool.name == "grep"));
759        assert!(definitions.iter().any(|tool| tool.name == "edit"));
760        assert!(definitions.iter().any(|tool| tool.name == "patch"));
761
762        let read = executor
763            .execute("read", &serde_json::json!({"file_path": "remote.txt"}))
764            .await
765            .unwrap();
766        assert_eq!(read.exit_code, 0);
767        assert!(read.output.contains("first"));
768
769        let write = executor
770            .execute(
771                "write",
772                &serde_json::json!({"file_path": "created.txt", "content": "remote write\n"}),
773            )
774            .await
775            .unwrap();
776        assert_eq!(write.exit_code, 0);
777        assert_eq!(fs.get("created.txt").unwrap(), "remote write\n");
778
779        let ls = executor
780            .execute("ls", &serde_json::json!({}))
781            .await
782            .unwrap();
783        assert_eq!(ls.exit_code, 0);
784        assert!(ls.output.contains("created.txt"));
785        assert!(ls.output.contains("remote.txt"));
786    }
787
788    #[tokio::test]
789    async fn test_bash_uses_workspace_command_runner() {
790        let fs = Arc::new(MemoryWorkspaceFs::default());
791        let fs_backend: Arc<dyn WorkspaceFileSystem> = fs;
792        let services = WorkspaceServices::builder(
793            WorkspaceRef::new("remote-workspace", "remote://workspace"),
794            fs_backend,
795        )
796        .command_runner(Arc::new(MockCommandRunner))
797        .build();
798        let executor = ToolExecutor::new_with_workspace_services_and_artifact_limits(
799            "/server/local-placeholder".to_string(),
800            services,
801            ArtifactStoreLimits::default(),
802        );
803        assert!(executor
804            .definitions()
805            .iter()
806            .any(|tool| tool.name == "bash"));
807
808        let result = executor
809            .execute("bash", &serde_json::json!({"command": "pwd"}))
810            .await
811            .unwrap();
812
813        assert_eq!(result.exit_code, 0);
814        assert_eq!(result.output, "remote: pwd\n");
815    }
816
817    #[tokio::test]
818    async fn test_command_env_is_available_on_default_context() {
819        let temp = tempfile::tempdir().unwrap();
820        let mut env = HashMap::new();
821        env.insert(
822            "A3S_COMMAND_ENV_TEST".to_string(),
823            "registry-env".to_string(),
824        );
825
826        let executor = ToolExecutor::new(temp.path().to_string_lossy().to_string());
827        executor.registry().set_command_env(Arc::new(env));
828        let context = executor.registry().context();
829        assert_eq!(
830            context
831                .command_env
832                .as_ref()
833                .and_then(|env| env.get("A3S_COMMAND_ENV_TEST"))
834                .map(String::as_str),
835            Some("registry-env")
836        );
837
838        #[cfg(windows)]
839        let command = "Write-Output $env:A3S_COMMAND_ENV_TEST";
840        #[cfg(not(windows))]
841        let command = "printf '%s' \"$A3S_COMMAND_ENV_TEST\"";
842
843        let result = executor
844            .execute("bash", &serde_json::json!({ "command": command }))
845            .await
846            .unwrap();
847
848        assert_eq!(result.exit_code, 0, "{}", result.output);
849        assert!(result.output.contains("registry-env"));
850    }
851
852    #[tokio::test]
853    async fn test_execute_applies_workspace_boundary_for_default_context() {
854        let workspace = tempfile::tempdir().unwrap();
855        let outside = tempfile::tempdir().unwrap();
856        std::fs::write(outside.path().join("secret.txt"), "secret").unwrap();
857
858        let executor = ToolExecutor::new(workspace.path().to_string_lossy().to_string());
859        let result = executor
860            .execute(
861                "grep",
862                &serde_json::json!({
863                    "pattern": "secret",
864                    "path": outside.path().to_string_lossy()
865                }),
866            )
867            .await
868            .unwrap();
869
870        assert_eq!(result.exit_code, 1);
871        assert!(result.output.contains("Workspace boundary"));
872        assert!(result.output.contains("escapes workspace"));
873    }
874
875    #[test]
876    fn test_tool_result_success() {
877        let result = ToolResult::success("test_tool", "output text".to_string());
878        assert_eq!(result.name, "test_tool");
879        assert_eq!(result.output, "output text");
880        assert_eq!(result.exit_code, 0);
881        assert!(result.metadata.is_none());
882    }
883
884    #[test]
885    fn test_tool_result_error() {
886        let result = ToolResult::error("test_tool", "error message".to_string());
887        assert_eq!(result.name, "test_tool");
888        assert_eq!(result.output, "error message");
889        assert_eq!(result.exit_code, 1);
890        assert!(result.metadata.is_none());
891    }
892
893    #[test]
894    fn test_tool_result_from_tool_output_success() {
895        let output = ToolOutput {
896            content: "success content".to_string(),
897            success: true,
898            metadata: None,
899            images: Vec::new(),
900            error_kind: None,
901        };
902        let result: ToolResult = output.into();
903        assert_eq!(result.output, "success content");
904        assert_eq!(result.exit_code, 0);
905        assert!(result.metadata.is_none());
906    }
907
908    #[test]
909    fn test_tool_result_from_tool_output_failure() {
910        let output = ToolOutput {
911            content: "failure content".to_string(),
912            success: false,
913            metadata: Some(serde_json::json!({"error": "test"})),
914            images: Vec::new(),
915            error_kind: None,
916        };
917        let result: ToolResult = output.into();
918        assert_eq!(result.output, "failure content");
919        assert_eq!(result.exit_code, 1);
920        assert_eq!(result.metadata, Some(serde_json::json!({"error": "test"})));
921    }
922
923    #[test]
924    fn test_tool_result_metadata_propagation() {
925        let output = ToolOutput::success("content")
926            .with_metadata(serde_json::json!({"_load_skill": true, "skill_name": "test"}));
927        let result: ToolResult = output.into();
928        assert_eq!(result.exit_code, 0);
929        let meta = result.metadata.unwrap();
930        assert_eq!(meta["_load_skill"], true);
931        assert_eq!(meta["skill_name"], "test");
932    }
933
934    #[test]
935    fn test_tool_executor_workspace() {
936        let executor = ToolExecutor::new("/test/workspace".to_string());
937        assert_eq!(executor.workspace().to_str().unwrap(), "/test/workspace");
938    }
939
940    #[test]
941    fn test_tool_executor_registry() {
942        let executor = ToolExecutor::new("/tmp".to_string());
943        let registry = executor.registry();
944        // Baseline tools on a raw ToolExecutor: 13
945        assert_eq!(registry.len(), 13);
946    }
947
948    #[tokio::test]
949    async fn test_tool_executor_get_artifact() {
950        let executor = ToolExecutor::new("/tmp".to_string());
951        executor.register_dynamic_tool(Arc::new(LargeArtifactTool));
952
953        let result = executor
954            .execute("large_artifact", &serde_json::json!({}))
955            .await
956            .unwrap();
957
958        let artifact_uri = result.metadata.as_ref().unwrap()["artifact"]["artifact_uri"]
959            .as_str()
960            .unwrap();
961        let artifact = executor.get_artifact(artifact_uri).expect("artifact");
962        assert_eq!(artifact.tool_name, "large_artifact");
963        assert_eq!(artifact.content.len(), MAX_OUTPUT_SIZE + 1);
964        assert!(executor.artifact_store().get(artifact_uri).is_some());
965    }
966
967    #[tokio::test]
968    async fn test_tool_executor_respects_artifact_limits() {
969        let executor = ToolExecutor::new_with_artifact_limits(
970            "/tmp".to_string(),
971            ArtifactStoreLimits {
972                max_artifacts: 1,
973                max_bytes: usize::MAX,
974            },
975        );
976        executor.register_dynamic_tool(Arc::new(LargeArtifactTool));
977
978        let first = executor
979            .execute("large_artifact", &serde_json::json!({}))
980            .await
981            .unwrap();
982        let first_uri = first.metadata.as_ref().unwrap()["artifact"]["artifact_uri"]
983            .as_str()
984            .unwrap()
985            .to_string();
986
987        executor
988            .execute("large_artifact", &serde_json::json!({ "suffix": "again" }))
989            .await
990            .unwrap();
991
992        assert_eq!(executor.artifact_store().limits().max_artifacts, 1);
993        assert_eq!(executor.artifact_store().len(), 1);
994        assert!(executor.get_artifact(&first_uri).is_none());
995    }
996
997    #[tokio::test]
998    async fn test_tool_executor_register_program_catalog_keeps_script_only_program_tool() {
999        let executor = ToolExecutor::new("/tmp".to_string());
1000        let trace_sink = crate::trace::InMemoryTraceSink::default();
1001        executor.set_trace_sink(Arc::new(trace_sink.clone()));
1002        executor.register_dynamic_tool(Arc::new(EchoTool));
1003        let mut catalog = crate::program::ProgramCatalog::new();
1004        catalog.register(
1005            crate::program::ProgramTemplate::new("custom_echo", "Run a custom echo program")
1006                .with_parameter(crate::program::ProgramParameter::required(
1007                    "message",
1008                    "Message to echo",
1009                ))
1010                .with_step(
1011                    crate::program::ProgramStepTemplate::new(
1012                        "echo",
1013                        serde_json::json!({ "message": "{{message}}" }),
1014                    )
1015                    .with_label("echo_message"),
1016                ),
1017        );
1018        executor.register_program_catalog(catalog);
1019
1020        let result = executor
1021            .execute(
1022                "program",
1023                &serde_json::json!({
1024                    "name": "custom_echo",
1025                    "inputs": {
1026                        "message": "hello from catalog"
1027                    }
1028                }),
1029            )
1030            .await
1031            .unwrap();
1032
1033        assert_eq!(result.exit_code, 1);
1034        assert!(result.output.contains("type parameter is required"));
1035
1036        let events = trace_sink.events();
1037        assert!(events.iter().any(|event| {
1038            event.kind == crate::trace::TraceEventKind::ToolExecution && event.name == "program"
1039        }));
1040        assert!(!events.iter().any(|event| {
1041            event.kind == crate::trace::TraceEventKind::ToolExecution && event.name == "echo"
1042        }));
1043    }
1044
1045    #[test]
1046    fn test_max_output_size_constant() {
1047        assert_eq!(MAX_OUTPUT_SIZE, 100 * 1024);
1048    }
1049
1050    #[test]
1051    fn test_max_read_lines_constant() {
1052        assert_eq!(MAX_READ_LINES, 2000);
1053    }
1054
1055    #[test]
1056    fn test_max_line_length_constant() {
1057        assert_eq!(MAX_LINE_LENGTH, 2000);
1058    }
1059
1060    #[test]
1061    fn test_truncate_tool_output_with_artifact_reference() {
1062        let output = "x".repeat(MAX_OUTPUT_SIZE + 1);
1063        let truncated = truncate_tool_output_with_artifact("test/tool", &output);
1064
1065        let artifact = truncated.artifact.expect("artifact");
1066        assert!(truncated.content.contains("Full output artifact:"));
1067        assert_eq!(artifact.original_bytes, MAX_OUTPUT_SIZE + 1);
1068        assert_eq!(artifact.shown_bytes, MAX_OUTPUT_SIZE);
1069        assert!(artifact.artifact_id.starts_with("tool-output:test_tool:"));
1070        assert!(artifact
1071            .artifact_uri
1072            .starts_with("a3s://tool-output/test_tool/"));
1073    }
1074
1075    #[test]
1076    fn test_tool_result_clone() {
1077        let result = ToolResult::success("test", "output".to_string());
1078        let cloned = result.clone();
1079        assert_eq!(result.name, cloned.name);
1080        assert_eq!(result.output, cloned.output);
1081        assert_eq!(result.exit_code, cloned.exit_code);
1082        assert_eq!(result.metadata, cloned.metadata);
1083    }
1084
1085    #[test]
1086    fn test_tool_result_debug() {
1087        let result = ToolResult::success("test", "output".to_string());
1088        let debug_str = format!("{:?}", result);
1089        assert!(debug_str.contains("test"));
1090        assert!(debug_str.contains("output"));
1091    }
1092
1093    #[tokio::test]
1094    async fn test_execute_attaches_diff_metadata() {
1095        use tempfile::TempDir;
1096        let dir = TempDir::new().unwrap();
1097        let file = dir.path().join("hello.txt");
1098        std::fs::write(&file, "before content\n").unwrap();
1099
1100        let executor = ToolExecutor::new(dir.path().to_str().unwrap().to_string());
1101        let args = serde_json::json!({
1102            "file_path": "hello.txt",
1103            "content": "after content\n"
1104        });
1105        let result = executor.execute("write", &args).await.unwrap();
1106
1107        let meta = result.metadata.expect("metadata should be present");
1108        assert_eq!(meta["before"], "before content\n");
1109        assert_eq!(meta["after"], "after content\n");
1110        assert_eq!(meta["file_path"], "hello.txt");
1111    }
1112
1113    #[tokio::test]
1114    async fn test_execute_with_context_attaches_diff_metadata() {
1115        use tempfile::TempDir;
1116        let dir = TempDir::new().unwrap();
1117        let canonical_dir = dir.path().canonicalize().unwrap();
1118        let file = canonical_dir.join("ctx.txt");
1119        std::fs::write(&file, "original\n").unwrap();
1120
1121        let executor = ToolExecutor::new(canonical_dir.to_str().unwrap().to_string());
1122        let ctx = ToolContext::new(canonical_dir.clone());
1123        let args = serde_json::json!({
1124            "file_path": "ctx.txt",
1125            "content": "updated\n"
1126        });
1127        let result = executor
1128            .execute_with_context("write", &args, &ctx)
1129            .await
1130            .unwrap();
1131        assert_eq!(result.exit_code, 0, "write tool failed: {}", result.output);
1132
1133        let meta = result.metadata.expect("metadata should be present");
1134        assert_eq!(meta["before"], "original\n");
1135        assert_eq!(meta["after"], "updated\n");
1136        assert_eq!(meta["file_path"], "ctx.txt");
1137    }
1138}