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