Skip to main content

talos_core/tool/
registry.rs

1use std::collections::HashMap;
2use std::sync::Arc;
3
4use serde_json::Value;
5use thiserror::Error;
6
7use super::AgentTool;
8
9/// Errors that can occur during tool registration, lookup, or execution.
10#[derive(Debug, Error)]
11pub enum ToolError {
12    /// The requested tool is not registered in the registry.
13    #[error("tool not found: {0}")]
14    ToolNotFound(String),
15
16    /// The input provided to a tool does not match its expected parameters.
17    #[error("invalid input for tool: {0}")]
18    InvalidInput(String),
19
20    /// An error occurred during tool execution.
21    #[error("tool execution error: {0}")]
22    ExecutionError(String),
23}
24
25/// Stable, display-safe identity for the crate, product profile, or plugin
26/// that contributed a tool.
27///
28/// Sources are diagnostics only: they do not grant permissions and must not
29/// contain credentials, workspace paths, tool arguments, or projected input.
30#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
31pub struct ToolContributionSource(String);
32
33impl ToolContributionSource {
34    /// Creates a stable contribution source identity.
35    pub fn new(source: impl Into<String>) -> Self {
36        Self(source.into())
37    }
38
39    /// Returns the source identity as a string slice.
40    #[must_use]
41    pub fn as_str(&self) -> &str {
42        &self.0
43    }
44}
45
46impl std::fmt::Display for ToolContributionSource {
47    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
48        formatter.write_str(&self.0)
49    }
50}
51
52/// One explicitly sourced tool instance ready for product composition.
53#[derive(Clone)]
54pub struct ToolContribution {
55    source: ToolContributionSource,
56    tool: Arc<dyn AgentTool>,
57}
58
59impl ToolContribution {
60    /// Creates a sourced tool contribution.
61    pub fn new(source: ToolContributionSource, tool: Arc<dyn AgentTool>) -> Self {
62        Self { source, tool }
63    }
64
65    /// Returns the stable source identity.
66    #[must_use]
67    pub fn source(&self) -> &ToolContributionSource {
68        &self.source
69    }
70
71    /// Returns the contributed tool name.
72    #[must_use]
73    pub fn name(&self) -> &str {
74        self.tool.name()
75    }
76
77    /// Returns the contributed tool instance.
78    #[must_use]
79    pub fn tool(&self) -> &Arc<dyn AgentTool> {
80        &self.tool
81    }
82
83    /// Applies an outer composition wrapper while preserving source identity.
84    #[must_use]
85    pub fn map_tool(mut self, wrap: impl FnOnce(Arc<dyn AgentTool>) -> Arc<dyn AgentTool>) -> Self {
86        self.tool = wrap(self.tool);
87        self
88    }
89}
90
91/// Deterministic duplicate-name diagnostic for checked tool composition.
92#[derive(Debug, Error, Clone, PartialEq, Eq)]
93#[error(
94    "duplicate tool registration '{tool_name}': existing source '{existing_source}', incoming source '{incoming_source}'"
95)]
96pub struct ToolRegistrationError {
97    /// Duplicate tool name.
98    pub tool_name: String,
99    /// Source that already owns the registered name.
100    pub existing_source: ToolContributionSource,
101    /// Source that attempted the duplicate registration.
102    pub incoming_source: ToolContributionSource,
103}
104
105const LEGACY_TOOL_REGISTRATION_SOURCE: &str = "legacy:unchecked";
106
107/// A registry for dynamically managing agent tools.
108///
109/// Tools are registered under their [`AgentTool::name`] and can be retrieved,
110/// listed, or have their inputs validated against their parameter schemas.
111#[derive(Default)]
112pub struct ToolRegistry {
113    tools: HashMap<String, Arc<dyn AgentTool>>,
114    contribution_sources: HashMap<String, ToolContributionSource>,
115}
116
117impl ToolRegistry {
118    /// Creates a new empty tool registry.
119    pub fn new() -> Self {
120        Self::default()
121    }
122
123    /// Registers a tool in the registry, replacing any existing tool with the
124    /// same name.
125    ///
126    /// This historical unchecked API remains temporarily source-compatible
127    /// during I158 migration. New product composition should use
128    /// [`register_contribution`](Self::register_contribution).
129    pub fn register(&mut self, tool: Arc<dyn AgentTool>) {
130        let name = tool.name().to_owned();
131        self.contribution_sources.remove(&name);
132        self.tools.insert(name, tool);
133    }
134
135    /// Registers one explicitly sourced contribution without replacing an
136    /// existing tool.
137    ///
138    /// A duplicate returns both source identities and leaves the current
139    /// registry entry unchanged.
140    pub fn register_contribution(
141        &mut self,
142        contribution: ToolContribution,
143    ) -> Result<(), ToolRegistrationError> {
144        let ToolContribution { source, tool } = contribution;
145        let tool_name = tool.name().to_owned();
146
147        if let Some(existing_source) = self.registered_source(&tool_name) {
148            return Err(ToolRegistrationError {
149                tool_name,
150                existing_source,
151                incoming_source: source,
152            });
153        }
154
155        self.contribution_sources.insert(tool_name.clone(), source);
156        self.tools.insert(tool_name, tool);
157        Ok(())
158    }
159
160    /// Registers a contribution batch transactionally.
161    ///
162    /// The complete batch is checked against the current registry and against
163    /// earlier entries in the same iteration order before any tool is inserted.
164    /// On the first duplicate, the registry remains unchanged.
165    pub fn register_contributions(
166        &mut self,
167        contributions: impl IntoIterator<Item = ToolContribution>,
168    ) -> Result<(), ToolRegistrationError> {
169        let contributions = contributions.into_iter().collect::<Vec<_>>();
170        let mut pending_sources = HashMap::<String, ToolContributionSource>::new();
171
172        for contribution in &contributions {
173            let tool_name = contribution.name().to_owned();
174            if let Some(existing_source) = self.registered_source(&tool_name) {
175                return Err(ToolRegistrationError {
176                    tool_name,
177                    existing_source,
178                    incoming_source: contribution.source().clone(),
179                });
180            }
181            if let Some(existing_source) = pending_sources.get(&tool_name) {
182                return Err(ToolRegistrationError {
183                    tool_name,
184                    existing_source: existing_source.clone(),
185                    incoming_source: contribution.source().clone(),
186                });
187            }
188            pending_sources.insert(tool_name, contribution.source().clone());
189        }
190
191        for ToolContribution { source, tool } in contributions {
192            let tool_name = tool.name().to_owned();
193            self.contribution_sources.insert(tool_name.clone(), source);
194            self.tools.insert(tool_name, tool);
195        }
196        Ok(())
197    }
198
199    fn registered_source(&self, tool_name: &str) -> Option<ToolContributionSource> {
200        self.tools.contains_key(tool_name).then(|| {
201            self.contribution_sources
202                .get(tool_name)
203                .cloned()
204                .unwrap_or_else(|| ToolContributionSource::new(LEGACY_TOOL_REGISTRATION_SOURCE))
205        })
206    }
207
208    /// Retrieves a tool by name, or `None` if not registered.
209    pub fn get(&self, name: &str) -> Option<&dyn AgentTool> {
210        self.tools.get(name).map(|t| t.as_ref())
211    }
212
213    /// Returns a list of all registered tools.
214    pub fn list(&self) -> Vec<&dyn AgentTool> {
215        self.tools.values().map(|t| t.as_ref()).collect()
216    }
217
218    /// Validates that the given input conforms to the tool's parameter schema.
219    ///
220    /// Returns `Ok(())` if the tool exists and the input is an object, or
221    /// `Err(ToolError)` if the tool is not found or the input is invalid.
222    ///
223    /// This performs a basic structural check (input must be a JSON object).
224    /// Full JSON Schema validation can be added later via the `jsonschema` crate.
225    pub fn validate_input(&self, name: &str, input: &Value) -> Result<(), ToolError> {
226        let tool = self
227            .get(name)
228            .ok_or_else(|| ToolError::ToolNotFound(name.to_owned()))?;
229
230        let params = tool.parameters();
231
232        // Basic validation: input must be an object
233        if !input.is_object() {
234            return Err(ToolError::InvalidInput(format!(
235                "expected object for tool '{name}', got {}",
236                input_type_name(input)
237            )));
238        }
239
240        // Check required fields if the schema specifies them
241        if let Some(schema_obj) = params.as_object()
242            && let Some(Value::Array(required)) = schema_obj.get("required")
243            && let Some(input_obj) = input.as_object()
244        {
245            for req in required {
246                if let Some(req_key) = req.as_str()
247                    && !input_obj.contains_key(req_key)
248                {
249                    return Err(ToolError::InvalidInput(format!(
250                        "missing required field '{req_key}' for tool '{name}'"
251                    )));
252                }
253            }
254        }
255
256        Ok(())
257    }
258}
259
260/// Returns a human-readable type name for a JSON value.
261fn input_type_name(value: &Value) -> &'static str {
262    match value {
263        Value::Null => "null",
264        Value::Bool(_) => "boolean",
265        Value::Number(_) => "number",
266        Value::String(_) => "string",
267        Value::Array(_) => "array",
268        Value::Object(_) => "object",
269    }
270}