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