entelix_core/tools/
toolset.rs1use std::collections::{BTreeMap, HashSet};
16use std::sync::Arc;
17
18use crate::error::{Error, Result};
19use crate::identity::validate_config_identifier;
20use crate::ir::ToolSpec;
21use crate::tools::metadata::ToolMetadata;
22use crate::tools::registry::ToolRegistry;
23use crate::tools::tool::Tool;
24
25pub struct Toolset<D = ()>
32where
33 D: Send + Sync + 'static,
34{
35 id: String,
36 by_name: BTreeMap<String, Arc<dyn Tool<D>>>,
37}
38
39impl<D> Clone for Toolset<D>
40where
41 D: Send + Sync + 'static,
42{
43 fn clone(&self) -> Self {
44 Self {
45 id: self.id.clone(),
46 by_name: self.by_name.clone(),
47 }
48 }
49}
50
51impl<D> Toolset<D>
52where
53 D: Send + Sync + 'static,
54{
55 pub fn new(id: impl Into<String>) -> Result<Self> {
61 let id = id.into();
62 validate_identifier("Toolset::new", "id", &id)?;
63 Ok(Self {
64 id,
65 by_name: BTreeMap::new(),
66 })
67 }
68
69 #[must_use]
71 pub fn id(&self) -> &str {
72 &self.id
73 }
74
75 #[must_use]
77 pub fn len(&self) -> usize {
78 self.by_name.len()
79 }
80
81 #[must_use]
83 pub fn is_empty(&self) -> bool {
84 self.by_name.is_empty()
85 }
86
87 pub fn names(&self) -> impl Iterator<Item = &str> {
89 self.by_name.keys().map(String::as_str)
90 }
91
92 #[must_use]
94 pub fn get(&self, name: &str) -> Option<&Arc<dyn Tool<D>>> {
95 self.by_name.get(name)
96 }
97
98 #[must_use]
104 pub fn tool_specs(&self) -> Arc<[ToolSpec]> {
105 self.by_name
106 .values()
107 .map(|tool| tool.metadata().to_tool_spec())
108 .collect()
109 }
110
111 pub fn register(mut self, tool: Arc<dyn Tool<D>>) -> Result<Self> {
121 validate_metadata("Toolset::register", tool.metadata())?;
122 let name = tool.metadata().name.clone();
123 if self.by_name.contains_key(&name) {
124 return Err(Error::config(format!(
125 "Toolset::register: tool '{name}' is already registered in toolset '{}'",
126 self.id
127 )));
128 }
129 self.by_name.insert(name, tool);
130 Ok(self)
131 }
132
133 pub fn restricted_to(&self, allowed: &[&str]) -> Result<Self> {
139 validate_allowed_names("Toolset::restricted_to", allowed)?;
140 let missing: Vec<&str> = allowed
141 .iter()
142 .copied()
143 .filter(|name| !self.by_name.contains_key(*name))
144 .collect();
145 if !missing.is_empty() {
146 return Err(Error::config(format!(
147 "Toolset::restricted_to: tool name(s) not in toolset '{}': {}",
148 self.id,
149 missing.join(", ")
150 )));
151 }
152
153 let allowed: HashSet<&str> = allowed.iter().copied().collect();
154 let by_name = self
155 .by_name
156 .iter()
157 .filter(|(name, _)| allowed.contains(name.as_str()))
158 .map(|(name, tool)| (name.clone(), Arc::clone(tool)))
159 .collect();
160 Ok(Self {
161 id: self.id.clone(),
162 by_name,
163 })
164 }
165}
166
167impl<D> Toolset<D>
168where
169 D: Clone + Send + Sync + 'static,
170{
171 pub fn install_into(&self, mut registry: ToolRegistry<D>) -> Result<ToolRegistry<D>> {
178 for tool in self.by_name.values() {
179 registry = registry.register(Arc::clone(tool))?;
180 }
181 Ok(registry)
182 }
183}
184
185impl<D> std::fmt::Debug for Toolset<D>
186where
187 D: Send + Sync + 'static,
188{
189 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
190 f.debug_struct("Toolset")
191 .field("id", &self.id)
192 .field("tools", &self.by_name.keys().collect::<Vec<_>>())
193 .finish()
194 }
195}
196
197fn validate_identifier(surface: &str, field: &str, value: &str) -> Result<()> {
198 validate_config_identifier(surface, field, value)
199}
200
201fn validate_metadata(surface: &str, metadata: &ToolMetadata) -> Result<()> {
202 validate_identifier(surface, "tool name", &metadata.name)?;
203 if metadata.description.trim().is_empty() {
204 return Err(Error::config(format!(
205 "{surface}: tool '{}' description must not be empty",
206 metadata.name
207 )));
208 }
209 jsonschema::options()
210 .build(&metadata.input_schema)
211 .map_err(|err| {
212 Error::config(format!(
213 "{surface}: tool '{}' input schema is invalid: {err}",
214 metadata.name
215 ))
216 })?;
217 if let Some(output_schema) = &metadata.output_schema {
218 jsonschema::options().build(output_schema).map_err(|err| {
219 Error::config(format!(
220 "{surface}: tool '{}' output schema is invalid: {err}",
221 metadata.name
222 ))
223 })?;
224 }
225 Ok(())
226}
227
228fn validate_allowed_names(surface: &str, allowed: &[&str]) -> Result<()> {
229 for name in allowed {
230 validate_identifier(surface, "requested tool name", name)?;
231 }
232 let mut seen = HashSet::with_capacity(allowed.len());
233 let duplicates: Vec<&str> = allowed
234 .iter()
235 .copied()
236 .filter(|name| !seen.insert(*name))
237 .collect();
238 if !duplicates.is_empty() {
239 return Err(Error::config(format!(
240 "{surface}: duplicate tool name(s): {}",
241 duplicates.join(", ")
242 )));
243 }
244 Ok(())
245}
246
247#[cfg(test)]
248#[allow(clippy::unwrap_used)]
249mod tests {
250 use async_trait::async_trait;
251 use serde_json::{Value, json};
252
253 use super::*;
254 use crate::agent_context::AgentContext;
255
256 struct EchoTool {
257 metadata: ToolMetadata,
258 }
259
260 impl EchoTool {
261 fn new(name: &str) -> Self {
262 Self {
263 metadata: ToolMetadata::function(
264 name,
265 format!("Echo tool {name}"),
266 json!({"type": "object", "properties": {}}),
267 ),
268 }
269 }
270 }
271
272 #[async_trait]
273 impl Tool for EchoTool {
274 fn metadata(&self) -> &ToolMetadata {
275 &self.metadata
276 }
277
278 async fn execute(&self, input: Value, _ctx: &AgentContext<()>) -> Result<Value> {
279 Ok(input)
280 }
281 }
282
283 #[test]
284 fn toolset_rejects_empty_id() {
285 let err = Toolset::<()>::new(" ").unwrap_err();
286 assert!(format!("{err}").contains("id must not be empty"));
287 }
288
289 #[test]
290 fn toolset_rejects_ambiguous_ids_and_tool_names() {
291 let err = Toolset::<()>::new("core ").unwrap_err();
292 assert!(format!("{err}").contains("leading or trailing whitespace"));
293
294 let err = Toolset::<()>::new("core\nnext").unwrap_err();
295 assert!(format!("{err}").contains("control characters"));
296
297 let err = Toolset::new("core")
298 .unwrap()
299 .register(Arc::new(EchoTool::new("echo ")))
300 .unwrap_err();
301 assert!(format!("{err}").contains("leading or trailing whitespace"));
302
303 let err = Toolset::new("core")
304 .unwrap()
305 .register(Arc::new(EchoTool::new("echo\nnext")))
306 .unwrap_err();
307 assert!(format!("{err}").contains("control characters"));
308 }
309
310 #[test]
311 fn toolset_accepts_free_form_tool_descriptions() {
312 let tool = EchoTool {
313 metadata: ToolMetadata::function(
314 "summarize",
315 "Summarize the supplied content in two concise sentences.",
316 json!({"type": "object", "properties": {}}),
317 ),
318 };
319
320 let set = Toolset::new("core")
321 .unwrap()
322 .register(Arc::new(tool))
323 .unwrap();
324 assert_eq!(set.names().collect::<Vec<_>>(), vec!["summarize"]);
325 }
326
327 #[test]
328 fn toolset_rejects_empty_tool_descriptions() {
329 let tool = EchoTool {
330 metadata: ToolMetadata::function(
331 "summarize",
332 " ",
333 json!({"type": "object", "properties": {}}),
334 ),
335 };
336
337 let err = Toolset::new("core")
338 .unwrap()
339 .register(Arc::new(tool))
340 .unwrap_err();
341 assert!(format!("{err}").contains("description must not be empty"));
342 }
343
344 #[test]
345 fn toolset_rejects_duplicate_tool_names() {
346 let err = Toolset::new("core")
347 .unwrap()
348 .register(Arc::new(EchoTool::new("echo")))
349 .unwrap()
350 .register(Arc::new(EchoTool::new("echo")))
351 .unwrap_err();
352 assert!(format!("{err}").contains("already registered"));
353 }
354
355 #[test]
356 fn restricted_to_is_strict_and_stable() {
357 let set = Toolset::new("core")
358 .unwrap()
359 .register(Arc::new(EchoTool::new("beta")))
360 .unwrap()
361 .register(Arc::new(EchoTool::new("alpha")))
362 .unwrap();
363
364 let narrowed = set.restricted_to(&["alpha"]).unwrap();
365 assert_eq!(narrowed.names().collect::<Vec<_>>(), vec!["alpha"]);
366
367 let err = set.restricted_to(&["alpha", "ghost"]).unwrap_err();
368 assert!(format!("{err}").contains("ghost"));
369
370 let err = set.restricted_to(&["alpha "]).unwrap_err();
371 assert!(format!("{err}").contains("leading or trailing whitespace"));
372
373 let err = set.restricted_to(&["alpha\nnext"]).unwrap_err();
374 assert!(format!("{err}").contains("control characters"));
375
376 let err = set.restricted_to(&["alpha", "alpha"]).unwrap_err();
377 assert!(format!("{err}").contains("duplicate tool name"));
378 }
379
380 #[tokio::test]
381 async fn install_into_preserves_registry_dispatch() {
382 let set = Toolset::new("core")
383 .unwrap()
384 .register(Arc::new(EchoTool::new("echo")))
385 .unwrap();
386 let registry = set.install_into(ToolRegistry::new()).unwrap();
387 let output = registry
388 .dispatch(
389 "tool_use_1",
390 "echo",
391 json!({"value": 1}),
392 &crate::ExecutionContext::new(),
393 )
394 .await
395 .unwrap();
396 assert_eq!(output, json!({"value": 1}));
397 }
398}