llm_tool/registry/
mod.rs

1//! Tool registry: registry and concurrent dispatch of named tools.
2
3mod batch;
4
5use alloc::{boxed::Box, sync::Arc, vec::Vec};
6
7pub use batch::{
8    AllowDestructiveBatch, BatchExecutionMode, BatchToolCall, BatchToolCallOutcome,
9    BatchToolCallStr,
10};
11
12use super::{
13    rust_tool::{ErasedTool, RustTool, definition_of},
14    types::{OutputSanitization, RegistryItem, ToolContext, ToolDefinition, ToolError, ToolOutput},
15};
16use crate::{
17    compat::{HashMap, RwLock, read_lock_recover, write_lock_recover},
18    rust_tool::BoxToolFuture,
19    types::ToolEffect,
20};
21
22/// Entry holding a cached [`ToolDefinition`] alongside the type-erased tool.
23///
24/// The definition is computed once at registration time so that
25/// [`ToolRegistry::definitions`] and [`ToolRegistry::iter`] never
26/// regenerate JSON schemas.
27struct RegisteredTool {
28    definition: ToolDefinition,
29    erased: Box<dyn ErasedTool>,
30}
31
32/// A registry of named tools available for dynamic dispatch.
33///
34/// Holds type-erased tool implementations and cached [`ToolDefinition`](super::types::ToolDefinition)
35/// schemas for fast lookup and execution.
36///
37/// Every dispatched output is sanitized with at least the registry's
38/// [`output_floor`](Self::output_floor) (default [`OutputSanitization::SanitizeTokens`]),
39/// even if an individual tool opted out.
40pub struct ToolRegistry {
41    tools: HashMap<&'static str, RegisteredTool>,
42    cached_definitions: RwLock<Option<Arc<[ToolDefinition]>>>,
43    output_floor: OutputSanitization,
44}
45
46impl core::fmt::Debug for ToolRegistry {
47    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
48        let names: Vec<&str> = self
49            .tools
50            .values()
51            .map(|r| r.definition.name.as_str())
52            .collect();
53        f.debug_struct("ToolRegistry")
54            .field("tool_count", &self.tools.len())
55            .field("tool_names", &names)
56            .field("output_floor", &self.output_floor)
57            .finish_non_exhaustive()
58    }
59}
60
61impl Default for ToolRegistry {
62    fn default() -> Self {
63        Self::new()
64    }
65}
66
67impl ToolRegistry {
68    /// Create an empty registry.
69    #[must_use]
70    pub fn new() -> Self {
71        Self {
72            tools: HashMap::new(),
73            cached_definitions: RwLock::new(None),
74            output_floor: OutputSanitization::default(),
75        }
76    }
77
78    /// Set the minimum [`OutputSanitization`] enforced on every dispatched tool output.
79    ///
80    /// The effective policy per call is the tool's own policy raised to at least this
81    /// floor ([`OutputSanitization::at_least`]). Lowering it to
82    /// [`OutputSanitization::None`] restores raw passthrough for tools that opt out.
83    #[must_use]
84    pub const fn with_output_floor(mut self, floor: OutputSanitization) -> Self {
85        self.output_floor = floor;
86        self
87    }
88
89    /// In-place variant of [`with_output_floor`](Self::with_output_floor).
90    pub const fn set_output_floor(&mut self, floor: OutputSanitization) -> &mut Self {
91        self.output_floor = floor;
92        self
93    }
94
95    /// The minimum [`OutputSanitization`] enforced on every dispatched tool output.
96    #[must_use]
97    pub const fn output_floor(&self) -> OutputSanitization {
98        self.output_floor
99    }
100
101    /// The policy actually applied to `name`'s output: its own policy raised to the floor.
102    ///
103    /// Returns `None` if no tool named `name` is registered.
104    #[must_use]
105    pub fn effective_output_sanitization(&self, name: &str) -> Option<OutputSanitization> {
106        self.definition(name)
107            .map(|def| def.output_sanitization.at_least(self.output_floor))
108    }
109
110    /// Raise `future`'s output to the registry floor when the tool's own policy is weaker.
111    fn enforce_output_floor<'a>(
112        &self,
113        tool_policy: OutputSanitization,
114        future: BoxToolFuture<'a>,
115    ) -> BoxToolFuture<'a> {
116        let effective = tool_policy.at_least(self.output_floor);
117        if effective == tool_policy {
118            return future;
119        }
120        Box::pin(async move { future.await.map(|output| effective.apply(output)) })
121    }
122
123    fn invalidate_cache(&self) {
124        let mut guard = write_lock_recover(&self.cached_definitions);
125        *guard = None;
126    }
127
128    /// Register a [`RustTool`]. Returns `&mut Self` for chaining.
129    ///
130    /// The tool's [`ToolDefinition`] (including JSON schema) is computed once
131    /// here and cached for the lifetime of the registration.
132    ///
133    /// If a tool with the same name was already registered, it is replaced.
134    ///
135    /// # Panics
136    ///
137    /// Panics if the tool's JSON schema cannot be serialized. This indicates a
138    /// bug in the tool's `Params` type (e.g. a broken `JsonSchema` impl). Use
139    /// [`try_register`](Self::try_register) for the non-panicking variant.
140    pub fn register<T: RustTool + 'static>(&mut self, tool: T) -> &mut Self {
141        if let Err(e) = self.try_register(tool) {
142            panic!("Failed to build definition for tool '{}': {e}", T::NAME);
143        }
144        self
145    }
146
147    /// Register a [`RustTool`], returning an error instead of panicking if the
148    /// tool's JSON schema cannot be built.
149    ///
150    /// This is the fallible counterpart to [`register`](Self::register); prefer
151    /// it when tool types are supplied dynamically and a broken `JsonSchema`
152    /// impl should not abort the process. If a tool with the same name was
153    /// already registered, it is replaced.
154    ///
155    /// # Errors
156    ///
157    /// Returns [`ToolError`] if the tool's `Params` type fails to produce a
158    /// JSON schema.
159    pub fn try_register<T: RustTool + 'static>(&mut self, tool: T) -> Result<&mut Self, ToolError> {
160        let definition = definition_of(&tool)?;
161        self.tools.insert(
162            T::NAME,
163            RegisteredTool {
164                definition,
165                erased: Box::new(tool),
166            },
167        );
168        self.invalidate_cache();
169        Ok(self)
170    }
171
172    /// Register a [`RustTool`], consuming and returning `Self` for owned chaining.
173    #[must_use]
174    pub fn with_tool<T: RustTool + 'static>(mut self, tool: T) -> Self {
175        self.register(tool);
176        self
177    }
178
179    /// Return a shared, deterministically sorted slice of all [`ToolDefinition`]s,
180    /// computed once and cached until the registry is mutated.
181    #[must_use]
182    pub fn definitions_slice(&self) -> Arc<[ToolDefinition]> {
183        {
184            let guard = read_lock_recover(&self.cached_definitions);
185            if let Some(ref slice) = *guard {
186                return Arc::clone(slice);
187            }
188        }
189        let mut defs: Vec<ToolDefinition> = self
190            .tools
191            .values()
192            .map(|entry| entry.definition.clone())
193            .collect();
194        defs.sort_by(|a, b| a.name.cmp(&b.name));
195        let arc_slice: Arc<[ToolDefinition]> = Arc::from(defs);
196        {
197            let mut guard = write_lock_recover(&self.cached_definitions);
198            *guard = Some(Arc::clone(&arc_slice));
199        }
200        arc_slice
201    }
202
203    /// Collect [`ToolDefinition`]s for all registered tools.
204    ///
205    /// Returns clones of the cached definitions computed at registration time.
206    #[must_use]
207    pub fn definitions(&self) -> Vec<ToolDefinition> {
208        self.definitions_slice().to_vec()
209    }
210
211    /// Dispatch a tool call by name with raw JSON arguments and a context,
212    /// returning a boxed future.
213    ///
214    /// # Errors
215    ///
216    /// Returns [`ToolError::not_found`] if no tool named `name` is registered.
217    pub fn dispatch_boxed<'a>(
218        &'a self,
219        name: &str,
220        args: serde_json::Value,
221        ctx: &'a ToolContext,
222    ) -> Result<BoxToolFuture<'a>, ToolError> {
223        let Some(entry) = self.tools.get(name) else {
224            return Err(ToolError::not_found(RegistryItem::Tool, name));
225        };
226        let future = entry.erased.call_erased(args, ctx);
227        Ok(self.enforce_output_floor(entry.definition.output_sanitization, future))
228    }
229
230    /// Dispatch a tool call by name with a raw JSON arguments string,
231    /// returning a boxed future without intermediate DOM allocation.
232    ///
233    /// # Errors
234    ///
235    /// Returns [`ToolError::not_found`] if no tool named `name` is registered.
236    pub fn dispatch_boxed_str<'a>(
237        &'a self,
238        name: &str,
239        args_json: &'a str,
240        ctx: &'a ToolContext,
241    ) -> Result<BoxToolFuture<'a>, ToolError> {
242        let Some(entry) = self.tools.get(name) else {
243            return Err(ToolError::not_found(RegistryItem::Tool, name));
244        };
245        let future = entry.erased.call_erased_str(args_json, ctx);
246        Ok(self.enforce_output_floor(entry.definition.output_sanitization, future))
247    }
248
249    /// Dispatch a tool call by name with raw JSON arguments and a context.
250    ///
251    /// # Errors
252    ///
253    /// Returns [`ToolError::not_found`] if no tool named `name` is registered
254    /// (carrying `error_kind = "not_registered"` metadata), or the tool's own
255    /// error if argument deserialization fails or the handler returns an error.
256    pub async fn dispatch(
257        &self,
258        name: &str,
259        args: serde_json::Value,
260        ctx: &ToolContext,
261    ) -> Result<ToolOutput, ToolError> {
262        self.dispatch_boxed(name, args, ctx)?.await
263    }
264
265    /// Dispatch a tool call by name with a raw JSON string argument.
266    ///
267    /// # Errors
268    ///
269    /// Returns [`ToolError::not_found`] if no tool named `name` is registered,
270    /// or the tool's own error if JSON parsing or the handler fails.
271    pub async fn dispatch_str(
272        &self,
273        name: &str,
274        args_json: &str,
275        ctx: &ToolContext,
276    ) -> Result<ToolOutput, ToolError> {
277        self.dispatch_boxed_str(name, args_json, ctx)?.await
278    }
279
280    /// Return `true` if the registered tool has [`ToolEffect::Destructive`].
281    #[must_use]
282    pub fn is_destructive(&self, name: &str) -> bool {
283        self.definition(name)
284            .is_some_and(|def| matches!(def.effect, ToolEffect::Destructive))
285    }
286
287    /// Remove a tool by name, returning `true` if it was present.
288    pub fn remove(&mut self, name: &str) -> bool {
289        let removed = self.tools.remove(name).is_some();
290        if removed {
291            self.invalidate_cache();
292        }
293        removed
294    }
295
296    /// Clear all registered tools.
297    pub fn clear(&mut self) {
298        self.tools.clear();
299        self.invalidate_cache();
300    }
301
302    /// Number of registered tools.
303    #[must_use]
304    pub fn len(&self) -> usize {
305        self.tools.len()
306    }
307
308    /// Whether the registry has no registered tools.
309    #[must_use]
310    pub fn is_empty(&self) -> bool {
311        self.tools.is_empty()
312    }
313
314    /// Whether a tool with the given name is registered.
315    #[must_use]
316    pub fn contains(&self, name: &str) -> bool {
317        self.tools.contains_key(name)
318    }
319
320    /// Borrow the cached [`ToolDefinition`] for a registered tool by name.
321    ///
322    /// Returns `None` if no tool named `name` is registered. Unlike
323    /// [`definitions`](Self::definitions), this clones nothing.
324    #[must_use]
325    pub fn definition(&self, name: &str) -> Option<&ToolDefinition> {
326        self.tools.get(name).map(|entry| &entry.definition)
327    }
328
329    /// Iterate over `(name, definition)` pairs for every registered tool.
330    ///
331    /// Yields clones of the cached definitions computed at registration time.
332    #[must_use]
333    pub fn iter(&self) -> ToolDefinitions<'_> {
334        ToolDefinitions {
335            inner: self.tools.iter(),
336        }
337    }
338}
339
340/// Borrowing iterator over `(name, definition)` pairs, yielded by
341/// [`ToolRegistry::iter`] and by `&ToolRegistry`'s [`IntoIterator`] impl.
342///
343/// Unlike a boxed trait object, this named iterator allocates nothing to
344/// construct and forwards `size_hint`/`len` from the underlying map iterator.
345/// Each cached [`ToolDefinition`] is cloned lazily as it is yielded.
346pub struct ToolDefinitions<'a> {
347    inner: crate::compat::HashMapIter<'a, &'static str, RegisteredTool>,
348}
349
350impl Iterator for ToolDefinitions<'_> {
351    type Item = (&'static str, ToolDefinition);
352
353    fn next(&mut self) -> Option<Self::Item> {
354        self.inner
355            .next()
356            .map(|(name, entry)| (*name, entry.definition.clone()))
357    }
358
359    fn size_hint(&self) -> (usize, Option<usize>) {
360        self.inner.size_hint()
361    }
362}
363
364impl ExactSizeIterator for ToolDefinitions<'_> {
365    fn len(&self) -> usize {
366        self.inner.len()
367    }
368}
369
370/// Iterate over `(name, definition)` pairs for every registered tool.
371///
372/// Yields `(&'static str, ToolDefinition)` for each tool in the registry.
373impl<'a> IntoIterator for &'a ToolRegistry {
374    type Item = (&'static str, ToolDefinition);
375    type IntoIter = ToolDefinitions<'a>;
376
377    fn into_iter(self) -> Self::IntoIter {
378        self.iter()
379    }
380}
381
382#[cfg(all(test, feature = "std"))]
383mod tests;
llm_tool/registry/mod.rs

llm_tool/registry/
mod.rs
