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