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}
156
157impl Default for Toolbox {
158    fn default() -> Self {
159        Self::new()
160    }
161}
162
163impl Toolbox {
164    pub fn new() -> Self {
165        Self {
166            tools: BTreeMap::new(),
167            allow: HashMap::new(),
168            admission: None,
169            journals: vec![],
170            counters: Counters::default(),
171        }
172    }
173
174    /// Register a tool. Names that the hard exposure rules forbid are refused
175    /// here, so they cannot even be listed.
176    pub fn register(&mut self, tool: Arc<dyn Tool>) -> Result<(), ToolError> {
177        let name = tool.spec().name.clone();
178        if let Some(why) = exposure::forbidden(&name) {
179            return Err(ToolError::new(ERR_NOT_EXPOSED, why));
180        }
181        self.tools.insert(name, tool);
182        Ok(())
183    }
184
185    pub fn register_all(&mut self, tools: Vec<Arc<dyn Tool>>) -> Result<(), ToolError> {
186        tools.into_iter().try_for_each(|t| self.register(t))
187    }
188
189    /// Let `caller` use tools matching `patterns` (`*`, `prefix*`, exact; see [`pattern_matches`]).
190    pub fn allow<I: IntoIterator<Item = S>, S: Into<String>>(
191        &mut self,
192        caller: &str,
193        patterns: I,
194    ) -> &mut Self {
195        self.allow
196            .entry(caller.to_string())
197            .or_default()
198            .extend(patterns.into_iter().map(Into::into));
199        self
200    }
201
202    pub fn set_admission(&mut self, hook: Arc<dyn ToolAdmission>) -> &mut Self {
203        self.admission = Some(hook);
204        self
205    }
206
207    pub fn add_journal(&mut self, j: Arc<dyn Journal>) -> &mut Self {
208        self.journals.push(j);
209        self
210    }
211
212    pub fn metrics(&self) -> Metrics {
213        Metrics {
214            calls: self.counters.calls.load(Ordering::SeqCst),
215            ok: self.counters.ok.load(Ordering::SeqCst),
216            failed: self.counters.failed.load(Ordering::SeqCst),
217            denied: self.counters.denied.load(Ordering::SeqCst),
218        }
219    }
220
221    fn permitted(&self, caller: &str, tool: &str, subject: Option<&str>) -> bool {
222        let Some(patterns) = self.allow.get(caller) else {
223            return false;
224        };
225        match subject {
226            Some(s) => {
227                let full = format!("{tool}:{s}");
228                patterns.iter().any(|p| pattern_matches(p, &full))
229            }
230            None => patterns.iter().any(|p| pattern_matches(p, tool)),
231        }
232    }
233
234    /// Specs visible to `caller`: only tools it may call at all (subject-carrying
235    /// tools are listed when any pattern could match them).
236    pub fn list_for(&self, caller: &str) -> Vec<ToolSpec> {
237        let Some(patterns) = self.allow.get(caller) else {
238            return vec![];
239        };
240        self.tools
241            .values()
242            .filter(|t| {
243                let n = &t.spec().name;
244                patterns.iter().any(|p| {
245                    pattern_matches(p, n)
246                        || p.starts_with(&format!("{n}:"))
247                        || (p.ends_with('*') && n.starts_with(p.trim_end_matches('*')))
248                })
249            })
250            .map(|t| t.spec().clone())
251            .collect()
252    }
253
254    pub fn list_all(&self) -> Vec<ToolSpec> {
255        self.tools.values().map(|t| t.spec().clone()).collect()
256    }
257
258    fn emit(&self, ev: ToolEvent) {
259        for j in &self.journals {
260            j.record(&ev);
261        }
262    }
263
264    fn deny(
265        &self,
266        caller: &str,
267        tool: &str,
268        by: DeniedBy,
269        code: &str,
270        reason: String,
271    ) -> ToolError {
272        self.counters.denied.fetch_add(1, Ordering::SeqCst);
273        self.emit(ToolEvent::Denied {
274            caller: caller.into(),
275            tool: tool.into(),
276            by,
277            reason: reason.clone(),
278        });
279        ToolError::new(code, reason)
280    }
281
282    pub fn call(&self, caller: &Caller, name: &str, args: Value) -> Result<Value, ToolError> {
283        self.call_cancellable(caller, name, args, &CancelToken::new())
284    }
285
286    pub fn call_cancellable(
287        &self,
288        caller: &Caller,
289        name: &str,
290        args: Value,
291        cancel: &CancelToken,
292    ) -> Result<Value, ToolError> {
293        self.counters.calls.fetch_add(1, Ordering::SeqCst);
294        let tool = self
295            .tools
296            .get(name)
297            .cloned()
298            .ok_or_else(|| ToolError::new(ERR_UNKNOWN_TOOL, format!("no tool `{name}`")))?;
299        let subject = tool.subject(&args);
300
301        if let Some(why) =
302            exposure::forbidden(name).or_else(|| subject.as_deref().and_then(exposure::forbidden))
303        {
304            return Err(self.deny(&caller.app, name, DeniedBy::Exposure, ERR_NOT_EXPOSED, why));
305        }
306        if !self.permitted(&caller.app, name, subject.as_deref()) {
307            let what = subject
308                .as_deref()
309                .map(|s| format!("{name}:{s}"))
310                .unwrap_or_else(|| name.to_string());
311            return Err(self.deny(
312                &caller.app,
313                name,
314                DeniedBy::Allowlist,
315                ERR_NOT_ALLOWED,
316                format!("`{}` may not call `{what}`", caller.app),
317            ));
318        }
319        if let Err(e) = validate_args(&tool.spec().input_schema, &args) {
320            return Err(self.deny(&caller.app, name, DeniedBy::Arguments, ERR_INVALID_ARGS, e));
321        }
322
323        let req = ToolRequest {
324            caller: caller.app.clone(),
325            tool: name.to_string(),
326            subject,
327            effect: tool.spec().effect,
328            arg_keys: args
329                .as_object()
330                .map(|o| o.keys().cloned().collect())
331                .unwrap_or_default(),
332            arg_bytes: args.to_string().len(),
333        };
334        if let Some(h) = &self.admission {
335            if let Verdict::Deny(reason) = h.approve(&req) {
336                return Err(self.deny(&caller.app, name, DeniedBy::Admission, ERR_DENIED, reason));
337            }
338        }
339
340        self.emit(ToolEvent::Started {
341            caller: caller.app.clone(),
342            tool: name.into(),
343        });
344        let started = Instant::now();
345        let ctx = CallCtx { caller, cancel };
346        let res = if cancel.is_cancelled() {
347            Err(ToolError::new(ERR_CANCELLED, "cancelled before start"))
348        } else {
349            tool.call(&ctx, args)
350        };
351        let (ok, code) = match &res {
352            Ok(_) => (true, None),
353            Err(e) => (false, Some(e.code.clone())),
354        };
355        if ok {
356            self.counters.ok.fetch_add(1, Ordering::SeqCst);
357        } else {
358            self.counters.failed.fetch_add(1, Ordering::SeqCst);
359        }
360        if let Some(h) = &self.admission {
361            h.settle(&req, ok, code.as_deref());
362        }
363        self.emit(ToolEvent::Settled {
364            caller: caller.app.clone(),
365            tool: name.into(),
366            ok,
367            error_code: code,
368            millis: started.elapsed().as_millis() as u64,
369        });
370        res
371    }
372}
rightkit_agent_tools/toolbox.rs

rightkit_agent_tools/
toolbox.rs
