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