aacode_rs/tools/
registry.rs

1// Copyright (c) 2026 xiefujin <490021684@qq.com>
2// Licensed under GPL-3.0, see LICENSE file for full license terms.
3
4//! Tool registry. Async Tool trait.
5
6use super::schema::ToolSchema;
7use crate::error::Result;
8use serde_json::Value;
9use std::collections::HashMap;
10use std::sync::atomic::AtomicBool;
11
12#[async_trait::async_trait]
13pub trait Tool: Send + Sync {
14    fn schema(&self) -> ToolSchema;
15    async fn call(&self, args: &Value, cancel: &AtomicBool) -> Result<String>;
16    fn name(&self) -> &'static str {
17        ""
18    }
19    /// Non-reentrant resource lane (e.g. `"shell"`, `"browser"`, `"python"`).
20    ///
21    /// Tools sharing a lane must not run concurrently. A batch that contains
22    /// any lane tool is executed sequentially (safe and simple); batches with
23    /// no lane still run concurrently.
24    fn lane(&self) -> Option<&'static str> {
25        None
26    }
27}
28
29/// Conservative operation classification for the read-only gate.
30///
31/// **Fail-closed**: anything not listed here is treated as mutating/executing.
32/// New read-only tools must be added explicitly; unknown tools are blocked in
33/// `read_only` mode rather than silently allowed.
34pub fn is_read_only(name: &str) -> bool {
35    const EXACT: &[&str] = &[
36        "snapshot",
37        "screenshot",
38        "fetch_url",
39        "fetch_rendered",
40        "browser_tools",
41        "mcp_status",
42        "run_skills",
43        "find_elements",
44        "navigate",
45        "open",
46        "back",
47        "forward",
48        "reload",
49        "stop",
50        "scroll",
51        "session_export",
52        "job_status",
53        "job_collect",
54        "job_list",
55        "ui_backends",
56        "ui_list_windows",
57        "ui_snapshot",
58        "ui_find",
59        "ui_wait",
60    ];
61    const PREFIXES: &[&str] = &[
62        "get_",
63        "list_",
64        "search_",
65        "read_",
66        "extract_",
67        "wait_",
68        "assert_",
69        "understand_",
70        "analyze_",
71    ];
72    EXACT.contains(&name) || PREFIXES.iter().any(|p| name.starts_with(p))
73}
74
75#[derive(Default)]
76pub struct ToolRegistry {
77    tools: HashMap<String, Box<dyn Tool>>,
78    schemas: HashMap<String, ToolSchema>,
79    /// When set, every tool observation is sanitized before it reaches the
80    /// model/session (see `crate::observation`). Optional so unit tests that
81    /// only exercise a single tool keep working with raw output.
82    artifact_store: Option<crate::artifacts::ArtifactStore>,
83    preview_budget: Option<std::sync::Arc<crate::observation::PreviewBudget>>,
84    sanitize_limits: crate::observation::SanitizeLimits,
85}
86
87impl ToolRegistry {
88    pub fn new() -> Self {
89        ToolRegistry::default()
90    }
91
92    /// Enable the observation choke point: all tool output is passed through
93    /// `observation::sanitize` (binary → artifact ref + bounded preview).
94    pub fn enable_sanitizer(
95        &mut self,
96        project_path: &std::path::Path,
97        limits: crate::observation::SanitizeLimits,
98    ) {
99        self.preview_budget = Some(std::sync::Arc::new(
100            crate::observation::PreviewBudget::default_for(&limits),
101        ));
102        self.artifact_store = Some(crate::artifacts::ArtifactStore::new(project_path));
103        self.sanitize_limits = limits;
104    }
105
106    pub fn artifact_store(&self) -> Option<&crate::artifacts::ArtifactStore> {
107        self.artifact_store.as_ref()
108    }
109
110    fn sanitize_observation(&self, obs: String) -> String {
111        match (self.artifact_store.as_ref(), self.preview_budget.as_ref()) {
112            (Some(store), Some(budget)) => {
113                crate::observation::sanitize(&obs, store, budget, &self.sanitize_limits)
114            }
115            _ => obs,
116        }
117    }
118    pub fn register(&mut self, tool: Box<dyn Tool>) {
119        let schema = tool.schema();
120        let name = schema.name.to_string();
121        self.schemas.insert(name.clone(), schema);
122        self.tools.insert(name, tool);
123    }
124    pub fn contains(&self, name: &str) -> bool {
125        self.tools.contains_key(name)
126    }
127
128    /// The non-reentrant lane for a tool, if any.
129    pub fn lane_of(&self, name: &str) -> Option<&'static str> {
130        self.tools.get(name).and_then(|t| t.lane())
131    }
132    pub fn names(&self) -> Vec<String> {
133        let mut v: Vec<String> = self.tools.keys().cloned().collect();
134        v.sort();
135        v
136    }
137    pub fn schema(&self, name: &str) -> Option<&ToolSchema> {
138        self.schemas.get(name)
139    }
140
141    /// Async execute for the loop.
142    pub async fn execute(&self, name: &str, raw_args: Value, cancel: &AtomicBool) -> String {
143        let tool = match self.tools.get(name) {
144            Some(t) => t,
145            None => return self.not_found_message(name),
146        };
147        let schema = self.schemas.get(name).unwrap();
148        let normalized = schema.normalize_params(&raw_args);
149        if let Err(msg) = schema.validate(&normalized) {
150            return format!(
151                "Parameter validation failed\n\n{msg}\n\nTool docs:\n{}",
152                schema.documentation()
153            );
154        }
155        match tool.call(&normalized, cancel).await {
156            Ok(obs) => self.sanitize_observation(obs),
157            Err(crate::error::AacodeError::Cancelled) => "Execution cancelled".to_string(),
158            Err(e) => format!("Execution error: {e}"),
159        }
160    }
161
162    pub fn openai_tools(&self) -> Vec<Value> {
163        let mut names: Vec<&String> = self.schemas.keys().collect();
164        names.sort();
165        names
166            .iter()
167            .filter_map(|n| self.schemas.get(*n))
168            .map(|s| s.to_openai_tool())
169            .collect()
170    }
171    pub fn anthropic_tools(&self) -> Vec<Value> {
172        let mut names: Vec<&String> = self.schemas.keys().collect();
173        names.sort();
174        names
175            .iter()
176            .filter_map(|n| self.schemas.get(*n))
177            .map(|s| s.to_anthropic_tool())
178            .collect()
179    }
180
181    fn not_found_message(&self, name: &str) -> String {
182        let mut msg = format!("Error: Unknown tool '{name}'\n\n");
183        let suggestions = self.suggest_similar(name, 3);
184        if !suggestions.is_empty() {
185            msg.push_str("Did you mean:\n");
186            for s in &suggestions {
187                msg.push_str(&format!("  - {s}\n"));
188            }
189            msg.push('\n');
190        }
191        msg.push_str("Available tools:\n");
192        for n in self.names() {
193            if let Some(s) = self.schemas.get(&n) {
194                let desc: String = s.description.chars().take(60).collect();
195                msg.push_str(&format!("  - {n}: {desc}\n"));
196            }
197        }
198        msg
199    }
200
201    fn suggest_similar(&self, name: &str, max: usize) -> Vec<String> {
202        let mut scored: Vec<(usize, &str)> = self
203            .tools
204            .keys()
205            .map(|k| (levenshtein(name, k), k.as_str()))
206            .filter(|(d, _)| *d <= 4)
207            .collect();
208        scored.sort_by_key(|(d, _)| *d);
209        scored
210            .iter()
211            .take(max)
212            .map(|(_, k)| k.to_string())
213            .collect()
214    }
215}
216
217fn levenshtein(a: &str, b: &str) -> usize {
218    let (a, b) = if a.len() > b.len() { (b, a) } else { (a, b) };
219    let n = a.len();
220    let m = b.len();
221    let mut prev: Vec<usize> = (0..=n).collect();
222    let mut curr = vec![0usize; n + 1];
223    for j in 1..=m {
224        curr[0] = j;
225        for i in 1..=n {
226            curr[i] = if a.as_bytes()[i - 1] == b.as_bytes()[j - 1] {
227                prev[i - 1]
228            } else {
229                1 + prev[i - 1].min(prev[i]).min(curr[i - 1])
230            };
231        }
232        std::mem::swap(&mut prev, &mut curr);
233    }
234    prev[n]
235}
236
237#[cfg(test)]
238mod tests {
239    use super::super::schema::{ParamType, ToolParameter};
240    use super::*;
241    use serde_json::json;
242
243    #[test]
244    fn read_only_classification_is_fail_closed() {
245        assert!(is_read_only("get_page_text"));
246        assert!(is_read_only("list_tabs"));
247        assert!(is_read_only("wait_for_element"));
248        assert!(is_read_only("search_history"));
249        assert!(is_read_only("navigate"));
250        assert!(is_read_only("snapshot"));
251        assert!(is_read_only("job_status"));
252        assert!(is_read_only("job_collect"));
253        assert!(!is_read_only("job_start"));
254        assert!(!is_read_only("job_cancel"));
255        assert!(!is_read_only("run_shell"));
256        assert!(!is_read_only("execute_python"));
257        assert!(!is_read_only("browser_call"));
258        assert!(!is_read_only("click"));
259        assert!(!is_read_only("delegate_task"));
260        // Unknown tools are blocked in read-only mode (fail-closed).
261        assert!(!is_read_only("some_future_mutating_tool"));
262    }
263
264    struct BlobTool;
265    #[async_trait::async_trait]
266    impl Tool for BlobTool {
267        fn schema(&self) -> ToolSchema {
268            ToolSchema::new("blob", "returns a big base64 blob", vec![])
269        }
270        async fn call(&self, _a: &Value, _c: &AtomicBool) -> Result<String> {
271            use base64::Engine;
272            let raw: Vec<u8> = (0..200_000u32).map(|i| (i % 256) as u8).collect();
273            let b64 = base64::engine::general_purpose::STANDARD.encode(&raw);
274            Ok(json!({"result": {"base64": b64, "width": 400, "height": 800}}).to_string())
275        }
276    }
277
278    /// The choke point must bound *any* tool output by default: a large base64
279    /// payload comes back as an artifact ref + tiny preview, never the bulk.
280    #[tokio::test]
281    async fn execute_sanitizes_bulk_binary_by_default() {
282        let dir = std::env::temp_dir().join(format!("aacode_reg_san_{}", uuid::Uuid::new_v4()));
283        std::fs::create_dir_all(&dir).unwrap();
284        let mut reg = ToolRegistry::new();
285        reg.enable_sanitizer(&dir, crate::observation::SanitizeLimits::default());
286        reg.register(Box::new(BlobTool));
287        let out = reg
288            .execute("blob", json!({}), &AtomicBool::new(false))
289            .await;
290        assert!(out.len() < 2_000, "output not bounded: {} chars", out.len());
291        assert!(out.contains("artifact"), "no artifact ref: {out}");
292        assert!(
293            out.contains(".aacode/artifacts/"),
294            "no artifact path: {out}"
295        );
296        assert!(!out.contains(&"A".repeat(200)), "bulk base64 leaked");
297    }
298
299    /// Without the sanitizer (single-tool unit tests) raw output is preserved.
300    #[tokio::test]
301    async fn execute_without_sanitizer_is_raw() {
302        let mut reg = ToolRegistry::new();
303        reg.register(Box::new(BlobTool));
304        let out = reg
305            .execute("blob", json!({}), &AtomicBool::new(false))
306            .await;
307        assert!(out.len() > 2_000, "raw output expected without sanitizer");
308    }
309
310    struct Echo;
311    #[async_trait::async_trait]
312    impl Tool for Echo {
313        fn schema(&self) -> ToolSchema {
314            ToolSchema::new(
315                "echo",
316                "echo text",
317                vec![ToolParameter::new(
318                    "text",
319                    ParamType::String,
320                    true,
321                    "text",
322                    &["msg"],
323                )],
324            )
325        }
326        async fn call(&self, args: &Value, _c: &AtomicBool) -> Result<String> {
327            Ok(args["text"].as_str().unwrap_or("").to_string())
328        }
329    }
330
331    fn reg() -> ToolRegistry {
332        let mut r = ToolRegistry::new();
333        r.register(Box::new(Echo));
334        r
335    }
336
337    #[tokio::test]
338    async fn executes_with_alias() {
339        let r = reg();
340        let cancel = AtomicBool::new(false);
341        let out = r.execute("echo", json!({"msg": "hi"}), &cancel).await;
342        assert_eq!(out, "hi");
343    }
344
345    #[tokio::test]
346    async fn validation_error_for_missing() {
347        let r = reg();
348        let cancel = AtomicBool::new(false);
349        let out = r.execute("echo", json!({}), &cancel).await;
350        assert!(out.contains("validation failed"));
351        assert!(out.contains("text"));
352    }
353
354    #[tokio::test]
355    async fn unknown_tool_suggests() {
356        let r = reg();
357        let cancel = AtomicBool::new(false);
358        let out = r.execute("ecko", json!({}), &cancel).await;
359        assert!(out.contains("Unknown tool 'ecko'"));
360        assert!(out.contains("echo"));
361    }
362
363    #[test]
364    fn levenshtein_basic() {
365        assert_eq!(levenshtein("cat", "bat"), 1);
366        assert_eq!(levenshtein("kitten", "sitting"), 3);
367        assert_eq!(levenshtein("echo", "echo"), 0);
368    }
369
370    #[test]
371    fn exports_native_tools() {
372        assert!(!reg().openai_tools().is_empty());
373    }
374}
aacode_rs/tools/registry.rs

aacode_rs/tools/
registry.rs
