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