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