a3s_code_core/tools/
registry.rs

1//! Tool Registry
2//!
3//! Central registry for all tools (built-in and dynamic).
4//! Provides thread-safe registration, lookup, and execution.
5
6use super::artifacts::{ArtifactStore, ArtifactStoreLimits, ToolArtifact};
7use super::types::{Tool, ToolContext, ToolOutput};
8use super::ToolResult;
9use super::{
10    merge_tool_output_artifact_metadata, truncate_tool_output_with_artifact, ToolOutputArtifact,
11};
12use crate::llm::ToolDefinition;
13use crate::trace::{InMemoryTraceSink, TraceEvent, TraceSink};
14use anyhow::Result;
15use std::collections::HashMap;
16use std::path::PathBuf;
17use std::sync::{Arc, RwLock};
18
19/// Tool registry for managing all available tools
20pub struct ToolRegistry {
21    tools: RwLock<HashMap<String, Arc<dyn Tool>>>,
22    /// Names of builtin tools that cannot be overridden
23    builtins: RwLock<std::collections::HashSet<String>>,
24    context: RwLock<ToolContext>,
25    artifact_store: ArtifactStore,
26    trace_sink: RwLock<Arc<dyn TraceSink>>,
27}
28
29impl ToolRegistry {
30    /// Create a new tool registry
31    pub fn new(workspace: PathBuf) -> Self {
32        Self::with_artifact_limits(workspace, ArtifactStoreLimits::default())
33    }
34
35    /// Create a new tool registry with custom artifact retention limits.
36    pub fn with_artifact_limits(workspace: PathBuf, artifact_limits: ArtifactStoreLimits) -> Self {
37        Self::with_artifact_limits_and_workspace_services(
38            workspace.clone(),
39            artifact_limits,
40            crate::workspace::WorkspaceServices::local(workspace),
41        )
42    }
43
44    /// Create a new tool registry with custom artifact limits and workspace backend.
45    pub fn with_artifact_limits_and_workspace_services(
46        workspace: PathBuf,
47        artifact_limits: ArtifactStoreLimits,
48        workspace_services: Arc<crate::workspace::WorkspaceServices>,
49    ) -> Self {
50        let context = ToolContext::new(workspace).with_workspace_services(workspace_services);
51        Self {
52            tools: RwLock::new(HashMap::new()),
53            builtins: RwLock::new(std::collections::HashSet::new()),
54            context: RwLock::new(context),
55            artifact_store: ArtifactStore::with_limits(artifact_limits),
56            trace_sink: RwLock::new(Arc::new(InMemoryTraceSink::default())),
57        }
58    }
59
60    /// Register a builtin tool (cannot be overridden by dynamic tools)
61    pub fn register_builtin(&self, tool: Arc<dyn Tool>) {
62        let name = tool.name().to_string();
63        let mut tools = self.tools.write().unwrap();
64        let mut builtins = self.builtins.write().unwrap();
65        tracing::debug!("Registering builtin tool: {}", name);
66        tools.insert(name.clone(), tool);
67        builtins.insert(name);
68    }
69
70    /// Register a tool
71    ///
72    /// If a tool with the same name already exists as a builtin, the registration
73    /// is rejected to prevent shadowing of core tools.
74    pub fn register(&self, tool: Arc<dyn Tool>) {
75        let name = tool.name().to_string();
76        let builtins = self.builtins.read().unwrap();
77        if builtins.contains(&name) {
78            tracing::warn!(
79                "Rejected registration of tool '{}': cannot shadow builtin",
80                name
81            );
82            return;
83        }
84        drop(builtins);
85        let mut tools = self.tools.write().unwrap();
86        tracing::debug!("Registering tool: {}", name);
87        tools.insert(name, tool);
88    }
89
90    /// Unregister a tool by name
91    ///
92    /// Returns true if the tool was found and removed.
93    pub fn unregister(&self, name: &str) -> bool {
94        if self.builtins.read().unwrap().contains(name) {
95            tracing::warn!(
96                "Rejected unregister of tool '{}': builtin tools cannot be removed through dynamic unregister",
97                name
98            );
99            return false;
100        }
101        let mut tools = self.tools.write().unwrap();
102        tracing::debug!("Unregistering tool: {}", name);
103        tools.remove(name).is_some()
104    }
105
106    /// Unregister all tools whose names start with the given prefix.
107    pub fn unregister_by_prefix(&self, prefix: &str) {
108        let builtins = self.builtins.read().unwrap().clone();
109        let mut tools = self.tools.write().unwrap();
110        tools.retain(|name, _| builtins.contains(name) || !name.starts_with(prefix));
111        tracing::debug!("Unregistered tools with prefix: {}", prefix);
112    }
113
114    /// Get a tool by name
115    pub fn get(&self, name: &str) -> Option<Arc<dyn Tool>> {
116        let tools = self.tools.read().unwrap();
117        tools.get(name).cloned()
118    }
119
120    /// Check if a tool exists
121    pub fn contains(&self, name: &str) -> bool {
122        let tools = self.tools.read().unwrap();
123        tools.contains_key(name)
124    }
125
126    /// Get all tool definitions for LLM
127    pub fn definitions(&self) -> Vec<ToolDefinition> {
128        let tools = self.tools.read().unwrap();
129        let mut definitions = tools
130            .values()
131            .map(|tool| ToolDefinition {
132                name: tool.name().to_string(),
133                description: tool.description().to_string(),
134                parameters: tool.parameters(),
135            })
136            .collect::<Vec<_>>();
137        definitions.sort_by(|a, b| a.name.cmp(&b.name));
138        definitions
139    }
140
141    /// List all registered tool names
142    pub fn list(&self) -> Vec<String> {
143        let tools = self.tools.read().unwrap();
144        let mut names = tools.keys().cloned().collect::<Vec<_>>();
145        names.sort();
146        names
147    }
148
149    /// Get the number of registered tools
150    pub fn len(&self) -> usize {
151        let tools = self.tools.read().unwrap();
152        tools.len()
153    }
154
155    /// Check if registry is empty
156    pub fn is_empty(&self) -> bool {
157        self.len() == 0
158    }
159
160    /// Get the tool context
161    pub fn context(&self) -> ToolContext {
162        self.context.read().unwrap().clone()
163    }
164
165    /// Return a clone of the registry's artifact store handle.
166    pub fn artifact_store(&self) -> ArtifactStore {
167        self.artifact_store.clone()
168    }
169
170    /// Get a stored tool artifact by URI.
171    pub fn get_artifact(&self, artifact_uri: &str) -> Option<ToolArtifact> {
172        self.artifact_store.get(artifact_uri)
173    }
174
175    /// Replace the trace sink used for compact tool/program execution events.
176    pub fn set_trace_sink(&self, sink: Arc<dyn TraceSink>) {
177        *self.trace_sink.write().unwrap() = sink;
178    }
179
180    /// Return the current trace sink.
181    pub fn trace_sink(&self) -> Arc<dyn TraceSink> {
182        Arc::clone(&self.trace_sink.read().unwrap())
183    }
184
185    /// Set the search configuration for the tool context
186    pub fn set_search_config(&self, config: crate::config::SearchConfig) {
187        let mut ctx = self.context.write().unwrap();
188        *ctx = ctx.clone().with_search_config(config);
189    }
190
191    /// Set a sandbox executor so that `bash` tool calls use the sandbox even
192    /// when executed without an explicit `ToolContext` (i.e., via `execute()`).
193    pub fn set_sandbox(&self, sandbox: std::sync::Arc<dyn crate::sandbox::BashSandbox>) {
194        let mut ctx = self.context.write().unwrap();
195        *ctx = ctx.clone().with_sandbox(sandbox);
196    }
197
198    /// Set environment overrides used by subprocess-backed tools when executed
199    /// without an explicit context.
200    pub fn set_command_env(&self, env: Arc<HashMap<String, String>>) {
201        let mut ctx = self.context.write().unwrap();
202        *ctx = ctx.clone().with_command_env(env);
203    }
204
205    /// Execute a tool by name using the registry's default context
206    pub async fn execute(&self, name: &str, args: &serde_json::Value) -> Result<ToolResult> {
207        let ctx = self.context();
208        self.execute_with_context(name, args, &ctx).await
209    }
210
211    /// Execute a tool by name with an external context
212    pub async fn execute_with_context(
213        &self,
214        name: &str,
215        args: &serde_json::Value,
216        ctx: &ToolContext,
217    ) -> Result<ToolResult> {
218        let start = std::time::Instant::now();
219
220        let tool = self.get(name);
221
222        let result = match tool {
223            Some(tool) => {
224                let mut output = tool.execute(args, ctx).await?;
225                let original_content = output.content.clone();
226                let truncated = truncate_tool_output_with_artifact(name, &output.content);
227                output.content = truncated.content;
228                if let Some(artifact) = truncated.artifact {
229                    self.store_tool_artifact(name, &original_content, &artifact);
230                    output.metadata = Some(merge_tool_output_artifact_metadata(
231                        output.metadata,
232                        &artifact,
233                    ));
234                }
235                Ok(ToolResult {
236                    name: name.to_string(),
237                    output: output.content,
238                    exit_code: if output.success { 0 } else { 1 },
239                    metadata: output.metadata,
240                    images: output.images,
241                    error_kind: output.error_kind,
242                })
243            }
244            None => Ok(ToolResult::error(name, format!("Unknown tool: {}", name))),
245        };
246
247        if let Ok(ref r) = result {
248            crate::telemetry::record_tool_result(r.exit_code, start.elapsed());
249            self.record_trace_event(name, r, start.elapsed());
250        }
251
252        result
253    }
254
255    /// Execute a tool and return raw output using the registry's default context
256    pub async fn execute_raw(
257        &self,
258        name: &str,
259        args: &serde_json::Value,
260    ) -> Result<Option<ToolOutput>> {
261        let ctx = self.context();
262        self.execute_raw_with_context(name, args, &ctx).await
263    }
264
265    /// Execute a tool and return raw output with an external context
266    pub async fn execute_raw_with_context(
267        &self,
268        name: &str,
269        args: &serde_json::Value,
270        ctx: &ToolContext,
271    ) -> Result<Option<ToolOutput>> {
272        let tool = self.get(name);
273
274        match tool {
275            Some(tool) => {
276                let mut output = tool.execute(args, ctx).await?;
277                let original_content = output.content.clone();
278                let truncated = truncate_tool_output_with_artifact(name, &output.content);
279                output.content = truncated.content;
280                if let Some(artifact) = truncated.artifact {
281                    self.store_tool_artifact(name, &original_content, &artifact);
282                    output.metadata = Some(merge_tool_output_artifact_metadata(
283                        output.metadata,
284                        &artifact,
285                    ));
286                }
287                Ok(Some(output))
288            }
289            None => Ok(None),
290        }
291    }
292
293    fn store_tool_artifact(&self, tool_name: &str, content: &str, artifact: &ToolOutputArtifact) {
294        self.artifact_store.put(ToolArtifact {
295            artifact_id: artifact.artifact_id.clone(),
296            artifact_uri: artifact.artifact_uri.clone(),
297            tool_name: tool_name.to_string(),
298            content: content.to_string(),
299            original_bytes: artifact.original_bytes,
300            shown_bytes: artifact.shown_bytes,
301        });
302    }
303
304    fn record_trace_event(&self, name: &str, result: &ToolResult, duration: std::time::Duration) {
305        let sink = self.trace_sink();
306        sink.record(TraceEvent::tool_execution(
307            name,
308            result.exit_code == 0,
309            result.exit_code,
310            duration,
311            result.output.len(),
312            result.metadata.as_ref(),
313        ));
314
315        if name == "program" {
316            sink.record(TraceEvent::program_execution(
317                name,
318                result.exit_code == 0,
319                result.exit_code,
320                duration,
321                result.output.len(),
322                result.metadata.as_ref(),
323            ));
324        }
325    }
326}
327
328#[cfg(test)]
329mod tests {
330    use super::*;
331    use crate::trace::{InMemoryTraceSink, TraceEventKind};
332    use async_trait::async_trait;
333
334    struct MockTool {
335        name: String,
336    }
337
338    #[async_trait]
339    impl Tool for MockTool {
340        fn name(&self) -> &str {
341            &self.name
342        }
343
344        fn description(&self) -> &str {
345            "A mock tool for testing"
346        }
347
348        fn parameters(&self) -> serde_json::Value {
349            serde_json::json!({
350                "type": "object",
351                "additionalProperties": false,
352                "properties": {},
353                "required": []
354            })
355        }
356
357        async fn execute(
358            &self,
359            _args: &serde_json::Value,
360            _ctx: &ToolContext,
361        ) -> Result<ToolOutput> {
362            Ok(ToolOutput::success("mock output"))
363        }
364    }
365
366    #[test]
367    fn test_registry_register_and_get() {
368        let registry = ToolRegistry::new(PathBuf::from("/tmp"));
369
370        let tool = Arc::new(MockTool {
371            name: "test".to_string(),
372        });
373        registry.register(tool);
374
375        assert!(registry.contains("test"));
376        assert!(!registry.contains("nonexistent"));
377
378        let retrieved = registry.get("test");
379        assert!(retrieved.is_some());
380        assert_eq!(retrieved.unwrap().name(), "test");
381    }
382
383    #[test]
384    fn test_registry_unregister() {
385        let registry = ToolRegistry::new(PathBuf::from("/tmp"));
386
387        let tool = Arc::new(MockTool {
388            name: "test".to_string(),
389        });
390        registry.register(tool);
391
392        assert!(registry.contains("test"));
393        assert!(registry.unregister("test"));
394        assert!(!registry.contains("test"));
395        assert!(!registry.unregister("test")); // Already removed
396    }
397
398    #[test]
399    fn test_registry_unregister_preserves_builtins() {
400        let registry = ToolRegistry::new(PathBuf::from("/tmp"));
401        registry.register_builtin(Arc::new(MockTool {
402            name: "read".to_string(),
403        }));
404
405        assert!(!registry.unregister("read"));
406        assert!(registry.contains("read"));
407    }
408
409    #[test]
410    fn test_registry_unregister_by_prefix_preserves_builtins() {
411        let registry = ToolRegistry::new(PathBuf::from("/tmp"));
412        registry.register_builtin(Arc::new(MockTool {
413            name: "mcp__builtin".to_string(),
414        }));
415        registry.register(Arc::new(MockTool {
416            name: "mcp__dynamic".to_string(),
417        }));
418
419        registry.unregister_by_prefix("mcp__");
420
421        assert!(registry.contains("mcp__builtin"));
422        assert!(!registry.contains("mcp__dynamic"));
423    }
424
425    #[test]
426    fn test_registry_definitions() {
427        let registry = ToolRegistry::new(PathBuf::from("/tmp"));
428
429        registry.register(Arc::new(MockTool {
430            name: "tool2".to_string(),
431        }));
432        registry.register(Arc::new(MockTool {
433            name: "tool1".to_string(),
434        }));
435
436        let definitions = registry.definitions();
437        assert_eq!(definitions.len(), 2);
438        let names: Vec<&str> = definitions
439            .iter()
440            .map(|definition| definition.name.as_str())
441            .collect();
442        assert_eq!(names, vec!["tool1", "tool2"]);
443    }
444
445    #[tokio::test]
446    async fn test_registry_execute() {
447        let registry = ToolRegistry::new(PathBuf::from("/tmp"));
448
449        registry.register(Arc::new(MockTool {
450            name: "test".to_string(),
451        }));
452
453        let result = registry
454            .execute("test", &serde_json::json!({}))
455            .await
456            .unwrap();
457        assert_eq!(result.exit_code, 0);
458        assert_eq!(result.output, "mock output");
459    }
460
461    #[tokio::test]
462    async fn test_registry_execute_unknown() {
463        let registry = ToolRegistry::new(PathBuf::from("/tmp"));
464
465        let result = registry
466            .execute("unknown", &serde_json::json!({}))
467            .await
468            .unwrap();
469        assert_eq!(result.exit_code, 1);
470        assert!(result.output.contains("Unknown tool"));
471    }
472
473    #[tokio::test]
474    async fn test_registry_execute_with_context_success() {
475        let registry = ToolRegistry::new(PathBuf::from("/tmp"));
476        let ctx = ToolContext::new(PathBuf::from("/tmp"));
477        let trace_sink = InMemoryTraceSink::default();
478        registry.set_trace_sink(Arc::new(trace_sink.clone()));
479
480        registry.register(Arc::new(MockTool {
481            name: "my_tool".to_string(),
482        }));
483
484        let result = registry
485            .execute_with_context("my_tool", &serde_json::json!({}), &ctx)
486            .await
487            .unwrap();
488        assert_eq!(result.name, "my_tool");
489        assert_eq!(result.exit_code, 0);
490        assert_eq!(result.output, "mock output");
491
492        let events = trace_sink.events();
493        assert_eq!(events.len(), 1);
494        assert_eq!(events[0].kind, TraceEventKind::ToolExecution);
495        assert_eq!(events[0].name, "my_tool");
496        assert!(events[0].success);
497        assert_eq!(events[0].output_bytes, "mock output".len());
498    }
499
500    #[tokio::test]
501    async fn test_registry_execute_with_context_unknown_tool() {
502        let registry = ToolRegistry::new(PathBuf::from("/tmp"));
503        let ctx = ToolContext::new(PathBuf::from("/tmp"));
504
505        let result = registry
506            .execute_with_context("nonexistent", &serde_json::json!({}), &ctx)
507            .await
508            .unwrap();
509        assert_eq!(result.exit_code, 1);
510        assert!(result.output.contains("Unknown tool: nonexistent"));
511    }
512
513    struct FailingTool;
514
515    #[async_trait]
516    impl Tool for FailingTool {
517        fn name(&self) -> &str {
518            "failing"
519        }
520
521        fn description(&self) -> &str {
522            "A tool that returns failure"
523        }
524
525        fn parameters(&self) -> serde_json::Value {
526            serde_json::json!({
527                "type": "object",
528                "additionalProperties": false,
529                "properties": {},
530                "required": []
531            })
532        }
533
534        async fn execute(
535            &self,
536            _args: &serde_json::Value,
537            _ctx: &ToolContext,
538        ) -> Result<ToolOutput> {
539            Ok(ToolOutput::error("something went wrong"))
540        }
541    }
542
543    #[tokio::test]
544    async fn test_registry_execute_failing_tool() {
545        let registry = ToolRegistry::new(PathBuf::from("/tmp"));
546        registry.register(Arc::new(FailingTool));
547
548        let result = registry
549            .execute("failing", &serde_json::json!({}))
550            .await
551            .unwrap();
552        assert_eq!(result.exit_code, 1);
553        assert_eq!(result.output, "something went wrong");
554    }
555
556    struct LargeOutputTool;
557
558    #[async_trait]
559    impl Tool for LargeOutputTool {
560        fn name(&self) -> &str {
561            "large_output"
562        }
563
564        fn description(&self) -> &str {
565            "A tool that returns more than the maximum output size"
566        }
567
568        fn parameters(&self) -> serde_json::Value {
569            serde_json::json!({
570                "type": "object",
571                "additionalProperties": false,
572                "properties": {},
573                "required": []
574            })
575        }
576
577        async fn execute(
578            &self,
579            _args: &serde_json::Value,
580            _ctx: &ToolContext,
581        ) -> Result<ToolOutput> {
582            Ok(ToolOutput::success(
583                "x".repeat(super::super::MAX_OUTPUT_SIZE + 1),
584            ))
585        }
586    }
587
588    #[tokio::test]
589    async fn test_registry_truncates_large_tool_output() {
590        let registry = ToolRegistry::new(PathBuf::from("/tmp"));
591        let trace_sink = InMemoryTraceSink::default();
592        registry.set_trace_sink(Arc::new(trace_sink.clone()));
593        registry.register(Arc::new(LargeOutputTool));
594
595        let result = registry
596            .execute("large_output", &serde_json::json!({}))
597            .await
598            .unwrap();
599
600        assert_eq!(result.exit_code, 0);
601        assert!(result.output.contains("[tool output truncated:"));
602        assert!(result
603            .output
604            .contains("Full output artifact: a3s://tool-output/large_output/"));
605        assert!(result.output.len() < super::super::MAX_OUTPUT_SIZE + 512);
606        let metadata = result.metadata.expect("artifact metadata");
607        assert_eq!(
608            metadata["artifact"]["original_bytes"],
609            serde_json::json!(super::super::MAX_OUTPUT_SIZE + 1)
610        );
611        assert_eq!(
612            metadata["artifact"]["shown_bytes"],
613            serde_json::json!(super::super::MAX_OUTPUT_SIZE)
614        );
615        assert!(metadata["artifact"]["artifact_id"]
616            .as_str()
617            .unwrap()
618            .starts_with("tool-output:large_output:"));
619        assert!(metadata["artifact"]["artifact_uri"]
620            .as_str()
621            .unwrap()
622            .starts_with("a3s://tool-output/large_output/"));
623
624        let artifact_uri = metadata["artifact"]["artifact_uri"].as_str().unwrap();
625        let artifact = registry
626            .get_artifact(artifact_uri)
627            .expect("full output artifact");
628        assert_eq!(artifact.tool_name, "large_output");
629        assert_eq!(artifact.original_bytes, super::super::MAX_OUTPUT_SIZE + 1);
630        assert_eq!(artifact.shown_bytes, super::super::MAX_OUTPUT_SIZE);
631        assert_eq!(
632            artifact.content,
633            "x".repeat(super::super::MAX_OUTPUT_SIZE + 1)
634        );
635
636        let events = trace_sink.events();
637        assert_eq!(events.len(), 1);
638        assert_eq!(events[0].artifact_uris, vec![artifact_uri]);
639    }
640
641    #[tokio::test]
642    async fn test_registry_execute_raw_success() {
643        let registry = ToolRegistry::new(PathBuf::from("/tmp"));
644        registry.register(Arc::new(MockTool {
645            name: "raw_test".to_string(),
646        }));
647
648        let output = registry
649            .execute_raw("raw_test", &serde_json::json!({}))
650            .await
651            .unwrap();
652        assert!(output.is_some());
653        let output = output.unwrap();
654        assert!(output.success);
655        assert_eq!(output.content, "mock output");
656    }
657
658    #[tokio::test]
659    async fn test_registry_execute_raw_stores_truncated_artifact() {
660        let registry = ToolRegistry::new(PathBuf::from("/tmp"));
661        registry.register(Arc::new(LargeOutputTool));
662
663        let output = registry
664            .execute_raw("large_output", &serde_json::json!({}))
665            .await
666            .unwrap()
667            .expect("raw output");
668
669        assert!(output.content.contains("[tool output truncated:"));
670        let metadata = output.metadata.expect("artifact metadata");
671        let artifact_uri = metadata["artifact"]["artifact_uri"].as_str().unwrap();
672        let artifact = registry
673            .get_artifact(artifact_uri)
674            .expect("full output artifact");
675        assert_eq!(artifact.tool_name, "large_output");
676        assert_eq!(artifact.content.len(), super::super::MAX_OUTPUT_SIZE + 1);
677    }
678
679    #[tokio::test]
680    async fn test_registry_execute_raw_unknown() {
681        let registry = ToolRegistry::new(PathBuf::from("/tmp"));
682
683        let output = registry
684            .execute_raw("missing", &serde_json::json!({}))
685            .await
686            .unwrap();
687        assert!(output.is_none());
688    }
689
690    #[test]
691    fn test_registry_list() {
692        let registry = ToolRegistry::new(PathBuf::from("/tmp"));
693        registry.register(Arc::new(MockTool {
694            name: "beta".to_string(),
695        }));
696        registry.register(Arc::new(MockTool {
697            name: "alpha".to_string(),
698        }));
699
700        let names = registry.list();
701        assert_eq!(names, vec!["alpha".to_string(), "beta".to_string()]);
702    }
703
704    #[test]
705    fn test_registry_len_and_is_empty() {
706        let registry = ToolRegistry::new(PathBuf::from("/tmp"));
707        assert!(registry.is_empty());
708        assert_eq!(registry.len(), 0);
709
710        registry.register(Arc::new(MockTool {
711            name: "t".to_string(),
712        }));
713        assert!(!registry.is_empty());
714        assert_eq!(registry.len(), 1);
715    }
716
717    #[test]
718    fn test_registry_replace_tool() {
719        let registry = ToolRegistry::new(PathBuf::from("/tmp"));
720        registry.register(Arc::new(MockTool {
721            name: "dup".to_string(),
722        }));
723        registry.register(Arc::new(MockTool {
724            name: "dup".to_string(),
725        }));
726        // Should still have only 1 tool (replaced)
727        assert_eq!(registry.len(), 1);
728    }
729}
a3s_code_core/tools/registry.rs

a3s_code_core/tools/
registry.rs
