rightkit_agent_tools/
toolbox.rs

1//! [`Toolbox`]: the registry and the one gate every call passes.
2//!
3//! Order: lookup, hard exposure rules, per-caller allowlist, argument schema,
4//! the caller's admission hook (typed, content-free, may deny before any side
5//! effect), execute, settle, journal.
6
7use crate::exposure;
8use crate::spec::*;
9use serde_json::{json, Value};
10use std::collections::{BTreeMap, HashMap};
11use std::sync::atomic::{AtomicU64, Ordering};
12use std::sync::{Arc, Mutex};
13use std::time::Instant;
14
15/// Content-free description of one intended call (no argument values).
16#[derive(Debug, Clone)]
17pub struct ToolRequest {
18    pub caller: String,
19    pub tool: String,
20    pub subject: Option<String>,
21    pub effect: Effect,
22    /// Top-level argument names, never values.
23    pub arg_keys: Vec<String>,
24    pub arg_bytes: usize,
25}
26
27#[derive(Debug, Clone, PartialEq, Eq)]
28pub enum Verdict {
29    Allow,
30    Deny(String),
31}
32
33/// The caller-supplied pre-effect hook (CodeRight EFF-001: validate, approve,
34/// execute, settle). `approve` runs after the toolbox's own validation and
35/// before the tool does anything; a denial executes nothing.
36pub trait ToolAdmission: Send + Sync {
37    fn approve(&self, req: &ToolRequest) -> Verdict;
38    fn settle(&self, _req: &ToolRequest, _ok: bool, _error_code: Option<&str>) {}
39}
40
41/// Lifecycle events for the caller's journal. Content-free.
42#[derive(Debug, Clone, PartialEq)]
43pub enum ToolEvent {
44    Started {
45        caller: String,
46        tool: String,
47    },
48    Denied {
49        caller: String,
50        tool: String,
51        by: DeniedBy,
52        reason: String,
53    },
54    Settled {
55        caller: String,
56        tool: String,
57        ok: bool,
58        error_code: Option<String>,
59        millis: u64,
60    },
61}
62
63#[derive(Debug, Clone, Copy, PartialEq, Eq)]
64pub enum DeniedBy {
65    Exposure,
66    Allowlist,
67    Arguments,
68    Admission,
69}
70
71impl ToolEvent {
72    pub fn to_value(&self) -> Value {
73        match self {
74            ToolEvent::Started { caller, tool } => {
75                json!({"event": "started", "caller": caller, "tool": tool})
76            }
77            ToolEvent::Denied {
78                caller,
79                tool,
80                by,
81                reason,
82            } => {
83                json!({"event": "denied", "caller": caller, "tool": tool, "by": format!("{by:?}").to_lowercase(), "reason": reason})
84            }
85            ToolEvent::Settled {
86                caller,
87                tool,
88                ok,
89                error_code,
90                millis,
91            } => {
92                json!({"event": "settled", "caller": caller, "tool": tool, "ok": ok, "error": error_code, "ms": millis})
93            }
94        }
95    }
96}
97
98pub trait Journal: Send + Sync {
99    fn record(&self, event: &ToolEvent);
100}
101
102#[derive(Default)]
103pub struct MemoryJournal(Mutex<Vec<ToolEvent>>);
104impl MemoryJournal {
105    pub fn new() -> Arc<Self> {
106        Arc::new(Self::default())
107    }
108    pub fn events(&self) -> Vec<ToolEvent> {
109        self.0.lock().unwrap_or_else(|e| e.into_inner()).clone()
110    }
111}
112impl Journal for MemoryJournal {
113    fn record(&self, event: &ToolEvent) {
114        self.0
115            .lock()
116            .unwrap_or_else(|e| e.into_inner())
117            .push(event.clone());
118    }
119}
120
121/// Content-free counters.
122#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
123pub struct Metrics {
124    pub calls: u64,
125    pub ok: u64,
126    pub failed: u64,
127    pub denied: u64,
128}
129
130#[derive(Default)]
131struct Counters {
132    calls: AtomicU64,
133    ok: AtomicU64,
134    failed: AtomicU64,
135    denied: AtomicU64,
136}
137
138/// `*` matches everything, `prefix*` a prefix, anything else exactly. Patterns
139/// are matched against `<tool>` and, for tools that carry a subject, against
140/// `<tool>:<subject>`; a tool with a subject needs a pattern that matches the
141/// full `<tool>:<subject>` form.
142fn pattern_matches(pattern: &str, candidate: &str) -> bool {
143    match pattern.strip_suffix('*') {
144        Some(prefix) => candidate.starts_with(prefix),
145        None => pattern == candidate,
146    }
147}
148
149pub struct Toolbox {
150    tools: BTreeMap<String, Arc<dyn Tool>>,
151    allow: HashMap<String, Vec<String>>,
152    admission: Option<Arc<dyn ToolAdmission>>,
153    journals: Vec<Arc<dyn Journal>>,
154    counters: Counters,
155    /// Apps whose `<app>.*` tools also get the dictation exposure layer.
156    dictation_apps: std::collections::HashSet<String>,
157}
158
159impl Default for Toolbox {
160    fn default() -> Self {
161        Self::new()
162    }
163}
164
165impl Toolbox {
166    pub fn new() -> Self {
167        Self {
168            tools: BTreeMap::new(),
169            allow: HashMap::new(),
170            admission: None,
171            journals: vec![],
172            counters: Counters::default(),
173            dictation_apps: Default::default(),
174        }
175    }
176
177    /// Apply the dictation exposure layer ([`exposure::forbidden_with`]) to every `<app>.*`
178    /// tool and its subject. Call before registering that app's tools.
179    pub fn protect_dictation(&mut self, app: impl Into<String>) {
180        self.dictation_apps.insert(app.into());
181    }
182
183    /// The hard exposure verdict for `name` (a tool name, or that tool's subject).
184    fn exposure_forbidden(&self, tool: &str, name: &str) -> Option<String> {
185        let app = tool.split('.').next().unwrap_or_default();
186        exposure::forbidden_with(name, self.dictation_apps.contains(app))
187    }
188
189    /// Register a tool. Names that the hard exposure rules forbid are refused
190    /// here, so they cannot even be listed.
191    pub fn register(&mut self, tool: Arc<dyn Tool>) -> Result<(), ToolError> {
192        let name = tool.spec().name.clone();
193        if let Some(why) = self.exposure_forbidden(&name, &name) {
194            return Err(ToolError::new(ERR_NOT_EXPOSED, why));
195        }
196        self.tools.insert(name, tool);
197        Ok(())
198    }
199
200    pub fn register_all(&mut self, tools: Vec<Arc<dyn Tool>>) -> Result<(), ToolError> {
201        tools.into_iter().try_for_each(|t| self.register(t))
202    }
203
204    /// Let `caller` use tools matching `patterns` (`*`, `prefix*`, exact; see [`pattern_matches`]).
205    pub fn allow<I: IntoIterator<Item = S>, S: Into<String>>(
206        &mut self,
207        caller: &str,
208        patterns: I,
209    ) -> &mut Self {
210        self.allow
211            .entry(caller.to_string())
212            .or_default()
213            .extend(patterns.into_iter().map(Into::into));
214        self
215    }
216
217    pub fn set_admission(&mut self, hook: Arc<dyn ToolAdmission>) -> &mut Self {
218        self.admission = Some(hook);
219        self
220    }
221
222    pub fn add_journal(&mut self, j: Arc<dyn Journal>) -> &mut Self {
223        self.journals.push(j);
224        self
225    }
226
227    pub fn metrics(&self) -> Metrics {
228        Metrics {
229            calls: self.counters.calls.load(Ordering::SeqCst),
230            ok: self.counters.ok.load(Ordering::SeqCst),
231            failed: self.counters.failed.load(Ordering::SeqCst),
232            denied: self.counters.denied.load(Ordering::SeqCst),
233        }
234    }
235
236    fn permitted(&self, caller: &str, tool: &str, subject: Option<&str>) -> bool {
237        let Some(patterns) = self.allow.get(caller) else {
238            return false;
239        };
240        match subject {
241            Some(s) => {
242                let full = format!("{tool}:{s}");
243                patterns.iter().any(|p| pattern_matches(p, &full))
244            }
245            None => patterns.iter().any(|p| pattern_matches(p, tool)),
246        }
247    }
248
249    /// Specs visible to `caller`: only tools it may call at all (subject-carrying
250    /// tools are listed when any pattern could match them).
251    pub fn list_for(&self, caller: &str) -> Vec<ToolSpec> {
252        let Some(patterns) = self.allow.get(caller) else {
253            return vec![];
254        };
255        self.tools
256            .values()
257            .filter(|t| {
258                let n = &t.spec().name;
259                patterns.iter().any(|p| {
260                    pattern_matches(p, n)
261                        || p.starts_with(&format!("{n}:"))
262                        || (p.ends_with('*') && n.starts_with(p.trim_end_matches('*')))
263                })
264            })
265            .map(|t| t.spec().clone())
266            .collect()
267    }
268
269    pub fn list_all(&self) -> Vec<ToolSpec> {
270        self.tools.values().map(|t| t.spec().clone()).collect()
271    }
272
273    fn emit(&self, ev: ToolEvent) {
274        for j in &self.journals {
275            j.record(&ev);
276        }
277    }
278
279    fn deny(
280        &self,
281        caller: &str,
282        tool: &str,
283        by: DeniedBy,
284        code: &str,
285        reason: String,
286    ) -> ToolError {
287        self.counters.denied.fetch_add(1, Ordering::SeqCst);
288        self.emit(ToolEvent::Denied {
289            caller: caller.into(),
290            tool: tool.into(),
291            by,
292            reason: reason.clone(),
293        });
294        ToolError::new(code, reason)
295    }
296
297    pub fn call(&self, caller: &Caller, name: &str, args: Value) -> Result<Value, ToolError> {
298        self.call_cancellable(caller, name, args, &CancelToken::new())
299    }
300
301    pub fn call_cancellable(
302        &self,
303        caller: &Caller,
304        name: &str,
305        args: Value,
306        cancel: &CancelToken,
307    ) -> Result<Value, ToolError> {
308        self.counters.calls.fetch_add(1, Ordering::SeqCst);
309        let tool = self
310            .tools
311            .get(name)
312            .cloned()
313            .ok_or_else(|| ToolError::new(ERR_UNKNOWN_TOOL, format!("no tool `{name}`")))?;
314        let subject = tool.subject(&args);
315
316        if let Some(why) = self.exposure_forbidden(name, name).or_else(|| {
317            subject
318                .as_deref()
319                .and_then(|s| self.exposure_forbidden(name, s))
320        }) {
321            return Err(self.deny(&caller.app, name, DeniedBy::Exposure, ERR_NOT_EXPOSED, why));
322        }
323        if !self.permitted(&caller.app, name, subject.as_deref()) {
324            let what = subject
325                .as_deref()
326                .map(|s| format!("{name}:{s}"))
327                .unwrap_or_else(|| name.to_string());
328            return Err(self.deny(
329                &caller.app,
330                name,
331                DeniedBy::Allowlist,
332                ERR_NOT_ALLOWED,
333                format!("`{}` may not call `{what}`", caller.app),
334            ));
335        }
336        if let Err(e) = validate_args(&tool.spec().input_schema, &args) {
337            return Err(self.deny(&caller.app, name, DeniedBy::Arguments, ERR_INVALID_ARGS, e));
338        }
339
340        let req = ToolRequest {
341            caller: caller.app.clone(),
342            tool: name.to_string(),
343            subject,
344            effect: tool.spec().effect,
345            arg_keys: args
346                .as_object()
347                .map(|o| o.keys().cloned().collect())
348                .unwrap_or_default(),
349            arg_bytes: args.to_string().len(),
350        };
351        if let Some(h) = &self.admission {
352            if let Verdict::Deny(reason) = h.approve(&req) {
353                return Err(self.deny(&caller.app, name, DeniedBy::Admission, ERR_DENIED, reason));
354            }
355        }
356
357        self.emit(ToolEvent::Started {
358            caller: caller.app.clone(),
359            tool: name.into(),
360        });
361        let started = Instant::now();
362        let ctx = CallCtx { caller, cancel };
363        let res = if cancel.is_cancelled() {
364            Err(ToolError::new(ERR_CANCELLED, "cancelled before start"))
365        } else {
366            tool.call(&ctx, args)
367        };
368        let (ok, code) = match &res {
369            Ok(_) => (true, None),
370            Err(e) => (false, Some(e.code.clone())),
371        };
372        if ok {
373            self.counters.ok.fetch_add(1, Ordering::SeqCst);
374        } else {
375            self.counters.failed.fetch_add(1, Ordering::SeqCst);
376        }
377        if let Some(h) = &self.admission {
378            h.settle(&req, ok, code.as_deref());
379        }
380        self.emit(ToolEvent::Settled {
381            caller: caller.app.clone(),
382            tool: name.into(),
383            ok,
384            error_code: code,
385            millis: started.elapsed().as_millis() as u64,
386        });
387        res
388    }
389}
rightkit_agent_tools/toolbox.rs

rightkit_agent_tools/
toolbox.rs
