Skip to main content

a3s_code_core/tools/
mod.rs

1//! Extensible Tool System
2//!
3//! Provides a trait-based abstraction for tools.
4//!
5//! ## Architecture
6//!
7//! ```text
8//! ToolRegistry
9//!   └── builtin tools (bash, read, write, edit, grep, glob, ls, patch, web_fetch, web_search)
10//! ```
11
12mod agent_dir_script_tool;
13mod artifacts;
14pub(crate) mod builtin;
15mod invocation;
16pub(crate) mod process;
17mod program_tool;
18mod registry;
19mod selector;
20pub mod skill;
21pub mod task;
22mod types;
23
24pub use crate::dynamic_workflow::register_dynamic_workflow;
25pub use agent_dir_script_tool::AgentDirScriptTool;
26pub use artifacts::{ArtifactStore, ArtifactStoreLimits, ToolArtifact};
27pub(crate) use builtin::register_skill;
28pub use builtin::{
29    register_generate_object, register_program, register_program_with_catalog, register_task,
30    register_task_with_mcp, register_task_with_mcp_managers,
31};
32pub(crate) use invocation::{
33    registry_tool_invoker, HostDirectPolicy, InvocationOrigin, ToolInvocation, ToolInvoker,
34};
35pub use program_tool::ProgramTool;
36pub use registry::ToolRegistry;
37pub use selector::{select_tools_for_messages, select_tools_for_prompt};
38pub use task::{
39    parallel_task_params_schema, task_params_schema, ParallelTaskParams, ParallelTaskTool,
40    TaskExecutor, TaskParams, TaskResult, TaskTool,
41};
42pub(crate) use types::{AgentEventBarrier, AgentEventBarrierReceiver};
43pub use types::{
44    InvocationRuntime, Tool, ToolContext, ToolErrorKind, ToolEventSender, ToolOutput,
45    ToolStreamEvent,
46};
47
48use crate::llm::ToolDefinition;
49use crate::text::truncate_utf8;
50use anyhow::Result;
51use serde::{Deserialize, Serialize};
52use std::collections::HashMap;
53use std::path::PathBuf;
54use std::sync::Arc;
55
56/// Maximum output size in bytes before truncation
57pub const MAX_OUTPUT_SIZE: usize = 100 * 1024; // 100KB
58
59/// Maximum lines to read from a file
60pub const MAX_READ_LINES: usize = 2000;
61
62/// Maximum line length before truncation
63pub const MAX_LINE_LENGTH: usize = 2000;
64
65#[derive(Debug, Clone, PartialEq, Eq)]
66pub(crate) struct ToolOutputArtifact {
67    pub artifact_id: String,
68    pub artifact_uri: String,
69    pub original_bytes: usize,
70    pub shown_bytes: usize,
71}
72
73#[derive(Debug, Clone)]
74pub(crate) struct TruncatedToolOutput {
75    pub content: String,
76    pub artifact: Option<ToolOutputArtifact>,
77}
78
79pub(crate) fn truncate_tool_output_with_artifact(
80    tool_name: &str,
81    output: &str,
82) -> TruncatedToolOutput {
83    if output.len() <= MAX_OUTPUT_SIZE {
84        return TruncatedToolOutput {
85            content: output.to_string(),
86            artifact: None,
87        };
88    }
89
90    let shown = truncate_utf8(output, MAX_OUTPUT_SIZE);
91    let artifact = tool_output_artifact(tool_name, output, shown.len());
92    let artifact_uri = artifact.artifact_uri.clone();
93    let content = format!(
94        "{}\n\n[tool output truncated: showing the first {} of {} bytes. Full output artifact: {}. Use narrower arguments such as offset/limit or filtering when possible.]",
95        shown,
96        shown.len(),
97        output.len(),
98        artifact_uri,
99    );
100
101    TruncatedToolOutput {
102        content,
103        artifact: Some(artifact),
104    }
105}
106
107pub(crate) fn tool_output_artifact(
108    tool_name: &str,
109    output: &str,
110    shown_bytes: usize,
111) -> ToolOutputArtifact {
112    use std::hash::{Hash, Hasher};
113
114    let mut hasher = std::collections::hash_map::DefaultHasher::new();
115    tool_name.hash(&mut hasher);
116    output.len().hash(&mut hasher);
117    output.hash(&mut hasher);
118    let digest = hasher.finish();
119    let sanitized_tool = tool_name
120        .chars()
121        .map(|ch| {
122            if ch.is_ascii_alphanumeric() || ch == '_' || ch == '-' {
123                ch
124            } else {
125                '_'
126            }
127        })
128        .collect::<String>();
129    let artifact_id = format!("tool-output:{sanitized_tool}:{digest:016x}");
130    let artifact_uri = format!("a3s://tool-output/{sanitized_tool}/{digest:016x}");
131
132    ToolOutputArtifact {
133        artifact_id,
134        artifact_uri,
135        original_bytes: output.len(),
136        shown_bytes,
137    }
138}
139
140pub(crate) fn merge_tool_output_artifact_metadata(
141    metadata: Option<serde_json::Value>,
142    artifact: &ToolOutputArtifact,
143) -> serde_json::Value {
144    let artifact_json = serde_json::json!({
145        "artifact_id": artifact.artifact_id,
146        "artifact_uri": artifact.artifact_uri,
147        "original_bytes": artifact.original_bytes,
148        "shown_bytes": artifact.shown_bytes,
149    });
150
151    match metadata {
152        Some(serde_json::Value::Object(mut object)) => {
153            object.insert("artifact".to_string(), artifact_json);
154            serde_json::Value::Object(object)
155        }
156        Some(value) => serde_json::json!({
157            "artifact": artifact_json,
158            "previous_metadata": value,
159        }),
160        None => serde_json::json!({
161            "artifact": artifact_json,
162        }),
163    }
164}
165
166/// Tool execution result returned by direct tool execution.
167#[derive(Debug, Clone, Serialize, Deserialize)]
168pub struct ToolResult {
169    pub name: String,
170    pub output: String,
171    pub exit_code: i32,
172    #[serde(skip_serializing_if = "Option::is_none")]
173    pub metadata: Option<serde_json::Value>,
174    /// Image attachments from tool execution (multi-modal output).
175    #[serde(skip)]
176    pub images: Vec<crate::llm::Attachment>,
177    /// Structured discriminant for tool failures. Populated by built-in
178    /// tools that can map their failure into a typed [`ToolErrorKind`]
179    /// (e.g. `edit`/`patch` setting `VersionConflict` on a CAS rejection
180    /// from `WorkspaceError`). Forwarded to the SDK so callers can react
181    /// programmatically without parsing `output`.
182    #[serde(skip_serializing_if = "Option::is_none")]
183    pub error_kind: Option<types::ToolErrorKind>,
184}
185
186impl ToolResult {
187    pub fn success(name: &str, output: String) -> Self {
188        Self {
189            name: name.to_string(),
190            output,
191            exit_code: 0,
192            metadata: None,
193            images: Vec::new(),
194            error_kind: None,
195        }
196    }
197
198    pub fn error(name: &str, message: String) -> Self {
199        Self {
200            name: name.to_string(),
201            output: message,
202            exit_code: 1,
203            metadata: None,
204            images: Vec::new(),
205            error_kind: None,
206        }
207    }
208}
209
210impl From<ToolOutput> for ToolResult {
211    fn from(output: ToolOutput) -> Self {
212        Self {
213            name: String::new(),
214            output: output.content,
215            exit_code: if output.success { 0 } else { 1 },
216            metadata: output.metadata,
217            images: output.images,
218            error_kind: output.error_kind,
219        }
220    }
221}
222
223/// Tool executor with workspace sandboxing
224///
225/// This is the main entry point for tool execution. It wraps the ToolRegistry.
226pub struct ToolExecutor {
227    workspace: PathBuf,
228    registry: Arc<ToolRegistry>,
229    command_env: Option<Arc<HashMap<String, String>>>,
230}
231
232/// Build a log line for a tool invocation that excludes argument *values*.
233///
234/// Argument values (full bash commands, file contents written by `write`/`edit`)
235/// can contain secrets, so the summary records only the tool name, the sorted
236/// argument field names, and the serialized payload size — never the values. This
237/// keeps the always-on `info!` tool trace (also exported to OTLP) compliant with
238/// the "never log secrets" boundary. Use `trace!` for full args when debugging.
239fn redacted_tool_log_summary(name: &str, args: &serde_json::Value) -> String {
240    let arg_keys: Vec<&str> = match args.as_object() {
241        Some(map) => {
242            let mut keys: Vec<&str> = map.keys().map(String::as_str).collect();
243            keys.sort_unstable();
244            keys
245        }
246        None => Vec::new(),
247    };
248    format!(
249        "Executing tool: {} (arg_keys={:?}, {} bytes)",
250        name,
251        arg_keys,
252        args.to_string().len()
253    )
254}
255
256/// Log a tool invocation without leaking argument values. See
257/// [`redacted_tool_log_summary`] for the redaction rationale.
258fn log_tool_invocation(name: &str, args: &serde_json::Value) {
259    tracing::info!("{}", redacted_tool_log_summary(name, args));
260    tracing::trace!("Tool {} full args: {}", name, args);
261}
262
263impl ToolExecutor {
264    pub fn new(workspace: String) -> Self {
265        let workspace_services =
266            crate::workspace::WorkspaceServices::local(PathBuf::from(&workspace));
267        Self::build(
268            workspace,
269            None,
270            ArtifactStoreLimits::default(),
271            workspace_services,
272        )
273    }
274
275    pub fn new_with_artifact_limits(
276        workspace: String,
277        artifact_limits: ArtifactStoreLimits,
278    ) -> Self {
279        let workspace_services =
280            crate::workspace::WorkspaceServices::local(PathBuf::from(&workspace));
281        Self::build(workspace, None, artifact_limits, workspace_services)
282    }
283
284    pub fn new_with_workspace_services(
285        workspace: String,
286        workspace_services: Arc<crate::workspace::WorkspaceServices>,
287    ) -> Self {
288        Self::build(
289            workspace,
290            None,
291            ArtifactStoreLimits::default(),
292            workspace_services,
293        )
294    }
295
296    pub fn new_with_workspace_services_and_artifact_limits(
297        workspace: String,
298        workspace_services: Arc<crate::workspace::WorkspaceServices>,
299        artifact_limits: ArtifactStoreLimits,
300    ) -> Self {
301        Self::build(workspace, None, artifact_limits, workspace_services)
302    }
303
304    fn build(
305        workspace: String,
306        command_env: Option<HashMap<String, String>>,
307        artifact_limits: ArtifactStoreLimits,
308        workspace_services: Arc<crate::workspace::WorkspaceServices>,
309    ) -> Self {
310        let workspace_path = PathBuf::from(&workspace);
311        let command_env = command_env.map(Arc::new);
312        let registry = Arc::new(ToolRegistry::with_artifact_limits_and_workspace_services(
313            workspace_path.clone(),
314            artifact_limits,
315            Arc::clone(&workspace_services),
316        ));
317        if let Some(env) = command_env.clone() {
318            registry.set_command_env(env);
319        }
320
321        // Register native Rust built-in tools — only those whose required
322        // workspace capability is available, so the model never sees a tool
323        // the backend cannot service.
324        builtin::register_builtins(&registry, &workspace_services.capabilities());
325        // Batch tool requires Arc<ToolRegistry>, registered separately
326        builtin::register_batch(&registry);
327        builtin::register_program(&registry);
328
329        Self {
330            workspace: workspace_path,
331            registry,
332            command_env,
333        }
334    }
335
336    fn check_workspace_boundary(
337        name: &str,
338        args: &serde_json::Value,
339        ctx: &ToolContext,
340    ) -> Result<()> {
341        let path_field = match name {
342            "read" | "write" | "edit" | "patch" => Some("file_path"),
343            "ls" | "grep" | "glob" => Some("path"),
344            _ => None,
345        };
346
347        if let Some(field) = path_field {
348            if let Some(path_str) = args.get(field).and_then(|v| v.as_str()) {
349                ctx.resolve_workspace_path(path_str).map_err(|e| {
350                    anyhow::anyhow!(
351                        "Workspace boundary check failed for tool '{}' path '{}': {}",
352                        name,
353                        path_str,
354                        e
355                    )
356                })?;
357            }
358        }
359
360        Ok(())
361    }
362
363    pub fn workspace(&self) -> &PathBuf {
364        &self.workspace
365    }
366
367    pub fn registry(&self) -> &Arc<ToolRegistry> {
368        &self.registry
369    }
370
371    /// Get a stored tool artifact by URI.
372    pub fn get_artifact(&self, artifact_uri: &str) -> Option<ToolArtifact> {
373        self.registry.get_artifact(artifact_uri)
374    }
375
376    /// Return a clone of the executor's artifact store handle.
377    pub fn artifact_store(&self) -> ArtifactStore {
378        self.registry.artifact_store()
379    }
380
381    /// Replace the sink used for compact execution trace events.
382    pub fn set_trace_sink(&self, sink: Arc<dyn crate::trace::TraceSink>) {
383        self.registry.set_trace_sink(sink);
384    }
385
386    /// Return the currently configured execution trace sink.
387    pub fn trace_sink(&self) -> Arc<dyn crate::trace::TraceSink> {
388        self.registry.trace_sink()
389    }
390
391    pub fn command_env(&self) -> Option<Arc<HashMap<String, String>>> {
392        self.command_env.clone()
393    }
394
395    pub fn register_dynamic_tool(&self, tool: Arc<dyn Tool>) {
396        self.registry.register(tool);
397    }
398
399    pub(crate) fn register_dynamic_tool_with_shadow(
400        &self,
401        tool: Arc<dyn Tool>,
402    ) -> (bool, Option<Arc<dyn Tool>>) {
403        self.registry.register_with_shadow(tool)
404    }
405
406    pub(crate) fn restore_dynamic_tool_if_same(
407        &self,
408        name: &str,
409        expected: &Arc<dyn Tool>,
410        replacement: Option<Arc<dyn Tool>>,
411    ) -> bool {
412        self.registry.restore_if_same(name, expected, replacement)
413    }
414
415    pub(crate) fn register_dynamic_tool_if_absent(&self, tool: Arc<dyn Tool>) -> bool {
416        self.registry.register_if_absent(tool)
417    }
418
419    pub fn unregister_dynamic_tool(&self, name: &str) {
420        self.registry.unregister(name);
421    }
422
423    /// Unregister all dynamic tools whose names start with the given prefix.
424    pub fn unregister_tools_by_prefix(&self, prefix: &str) {
425        self.registry.unregister_by_prefix(prefix);
426    }
427
428    /// Replace the model-visible `program` tool with a custom PTC catalog.
429    pub fn register_program_catalog(&self, catalog: crate::program::ProgramCatalog) {
430        builtin::register_program_with_catalog(&self.registry, catalog);
431    }
432
433    /// Execute directly against this low-level executor.
434    ///
435    /// This API intentionally does not install agent/session permission, HITL,
436    /// hook, budget, queue, timeout, cancellation, or sanitization policy.
437    /// Session hosts should use [`crate::AgentSession::tool`] (or its typed
438    /// helpers), and agent runtimes must dispatch through their scoped tool
439    /// invocation gateway.
440    pub async fn execute(&self, name: &str, args: &serde_json::Value) -> Result<ToolResult> {
441        let ctx = self.registry.context();
442        if let Err(e) = Self::check_workspace_boundary(name, args, &ctx) {
443            return Ok(ToolResult::error(name, e.to_string()));
444        }
445
446        log_tool_invocation(name, args);
447        let mut result = self.registry.execute_with_context(name, args, &ctx).await;
448        if let Ok(ref mut r) = result {
449            self.attach_diff_metadata(name, args, r);
450        }
451        match &result {
452            Ok(r) => tracing::info!("Tool {} completed with exit_code={}", name, r.exit_code),
453            Err(e) => tracing::error!("Tool {} failed: {}", name, e),
454        }
455        result
456    }
457
458    /// Execute directly with a caller-owned context.
459    ///
460    /// Like [`Self::execute`], this is an ungoverned standalone boundary. A
461    /// `ToolContext` supplies capabilities to the tool but is not itself a
462    /// substitute for the agent/session invocation gateway.
463    pub async fn execute_with_context(
464        &self,
465        name: &str,
466        args: &serde_json::Value,
467        ctx: &ToolContext,
468    ) -> Result<ToolResult> {
469        Self::check_workspace_boundary(name, args, ctx)?;
470        log_tool_invocation(name, args);
471        let mut result = self.registry.execute_with_context(name, args, ctx).await;
472        if let Ok(ref mut r) = result {
473            self.attach_diff_metadata(name, args, r);
474        }
475        match &result {
476            Ok(r) => tracing::info!("Tool {} completed with exit_code={}", name, r.exit_code),
477            Err(e) => tracing::error!("Tool {} failed: {}", name, e),
478        }
479        result
480    }
481
482    fn attach_diff_metadata(&self, name: &str, args: &serde_json::Value, result: &mut ToolResult) {
483        if !matches!(name, "write" | "edit" | "patch") {
484            return;
485        }
486        let Some(file_path) = args.get("file_path").and_then(serde_json::Value::as_str) else {
487            return;
488        };
489        // Only store file_path in metadata, let translate_event read the actual content
490        // using the session's correct workspace
491        let meta = result.metadata.get_or_insert_with(|| serde_json::json!({}));
492        meta["file_path"] = serde_json::Value::String(file_path.to_string());
493    }
494
495    pub fn definitions(&self) -> Vec<ToolDefinition> {
496        self.registry.definitions()
497    }
498}
499
500#[cfg(test)]
501#[path = "tests.rs"]
502mod tests;