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

llm_tool/registry/
mod.rs
