talos_core/tool/
registry.rs1use std::collections::HashMap;
2use std::sync::Arc;
3
4use serde_json::Value;
5use thiserror::Error;
6
7use super::AgentTool;
8
9#[derive(Debug, Error)]
11pub enum ToolError {
12 #[error("tool not found: {0}")]
14 ToolNotFound(String),
15
16 #[error("invalid input for tool: {0}")]
18 InvalidInput(String),
19
20 #[error("tool execution error: {0}")]
22 ExecutionError(String),
23}
24
25#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
31pub struct ToolContributionSource(String);
32
33impl ToolContributionSource {
34 pub fn new(source: impl Into<String>) -> Self {
36 Self(source.into())
37 }
38
39 #[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#[derive(Clone)]
54pub struct ToolContribution {
55 source: ToolContributionSource,
56 tool: Arc<dyn AgentTool>,
57}
58
59impl ToolContribution {
60 pub fn new(source: ToolContributionSource, tool: Arc<dyn AgentTool>) -> Self {
62 Self { source, tool }
63 }
64
65 #[must_use]
67 pub fn source(&self) -> &ToolContributionSource {
68 &self.source
69 }
70
71 #[must_use]
73 pub fn name(&self) -> &str {
74 self.tool.name()
75 }
76
77 #[must_use]
79 pub fn tool(&self) -> &Arc<dyn AgentTool> {
80 &self.tool
81 }
82
83 #[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#[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 pub tool_name: String,
99 pub existing_source: ToolContributionSource,
101 pub incoming_source: ToolContributionSource,
103}
104
105const LEGACY_TOOL_REGISTRATION_SOURCE: &str = "legacy:unchecked";
106
107#[derive(Default)]
112pub struct ToolRegistry {
113 tools: HashMap<String, Arc<dyn AgentTool>>,
114 contribution_sources: HashMap<String, ToolContributionSource>,
115}
116
117impl ToolRegistry {
118 pub fn new() -> Self {
120 Self::default()
121 }
122
123 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 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 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 pub fn get(&self, name: &str) -> Option<&dyn AgentTool> {
210 self.tools.get(name).map(|t| t.as_ref())
211 }
212
213 pub fn list(&self) -> Vec<&dyn AgentTool> {
215 self.tools.values().map(|t| t.as_ref()).collect()
216 }
217
218 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 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 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
260fn 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}