Skip to main content

salvor_tools/
registry.rs

1//! [`ToolSet`]: the name-keyed registry of tools the runtime dispatches
2//! through, plus [`RegistryError`] for a duplicate registration.
3
4use std::collections::BTreeMap;
5
6use thiserror::Error;
7
8use crate::erased::{DynTool, ToolDescriptor, TypedTool};
9use crate::handler::ToolHandler;
10
11/// A collection of tools keyed by name.
12///
13/// This is what an agent is configured with and what the runtime loop
14/// dispatches through: the model names a tool, the loop looks it up here and
15/// calls it.
16/// Tools are stored type-erased as `Box<dyn DynTool>`, so native
17/// [`ToolHandler`]s and (later) MCP-backed tools live side by side.
18///
19/// Names are unique. Registering a second tool under a name already present is
20/// a [`RegistryError`], not a silent overwrite: a shadowed tool is a
21/// configuration bug the caller should hear about, since the model would be
22/// told about one tool and dispatched to another.
23///
24/// The backing map is a `BTreeMap`, so [`descriptors`](Self::descriptors) and
25/// [`tools`](Self::tools) enumerate in a stable, name-sorted order. A
26/// deterministic tool list is worth having for a runtime built on
27/// deterministic replay: the order the model is shown its tools does not wobble
28/// between runs.
29#[derive(Default)]
30pub struct ToolSet {
31    tools: BTreeMap<String, Box<dyn DynTool>>,
32}
33
34impl ToolSet {
35    /// An empty registry.
36    pub fn new() -> Self {
37        Self::default()
38    }
39
40    /// Registers a typed [`ToolHandler`], wrapping it as a
41    /// [`DynTool`](crate::DynTool) automatically.
42    ///
43    /// Returns [`RegistryError::DuplicateName`] if a tool is already registered
44    /// under this handler's [`NAME`](crate::ToolMeta::NAME); the existing tool
45    /// is left in place.
46    pub fn register<H: ToolHandler + 'static>(&mut self, handler: H) -> Result<(), RegistryError> {
47        self.register_dyn(Box::new(TypedTool::new(handler)))
48    }
49
50    /// Registers an already type-erased tool.
51    ///
52    /// This is the path a runtime-defined tool (such as an MCP-backed one)
53    /// takes, since it implements [`DynTool`](crate::DynTool) directly rather
54    /// than through [`TypedTool`](crate::TypedTool). Same duplicate-name rule
55    /// as [`register`](Self::register).
56    pub fn register_dyn(&mut self, tool: Box<dyn DynTool>) -> Result<(), RegistryError> {
57        let name = tool.name().to_owned();
58        if self.tools.contains_key(&name) {
59            return Err(RegistryError::DuplicateName { name });
60        }
61        self.tools.insert(name, tool);
62        Ok(())
63    }
64
65    /// Looks up a tool by the name the model called it by.
66    pub fn get(&self, name: &str) -> Option<&dyn DynTool> {
67        self.tools.get(name).map(Box::as_ref)
68    }
69
70    /// How many tools are registered.
71    pub fn len(&self) -> usize {
72        self.tools.len()
73    }
74
75    /// Whether the registry holds no tools.
76    pub fn is_empty(&self) -> bool {
77        self.tools.is_empty()
78    }
79
80    /// Every registered tool, name-sorted, as `dyn DynTool`.
81    pub fn tools(&self) -> impl Iterator<Item = &dyn DynTool> {
82        self.tools.values().map(Box::as_ref)
83    }
84
85    /// A [`ToolDescriptor`] for every registered tool, name-sorted.
86    ///
87    /// This is the list handed to a model: each tool's name, description,
88    /// effect, and input schema.
89    pub fn descriptors(&self) -> Vec<ToolDescriptor> {
90        self.tools().map(DynTool::descriptor).collect()
91    }
92}
93
94/// The error [`ToolSet`] registration returns.
95#[derive(Debug, Error, PartialEq, Eq)]
96pub enum RegistryError {
97    /// A tool is already registered under this name.
98    #[error("a tool named `{name}` is already registered")]
99    DuplicateName {
100        /// The name that collided.
101        name: String,
102    },
103}