1use crate::tools::catalog::{tool_id_for, CapabilityRecord, CatalogError, ToolCatalog};
3use crate::tools::Tool;
4use serde_json::Value;
5use std::collections::{HashMap, HashSet};
6use std::sync::Arc;
7
8#[derive(Clone)]
11pub struct ToolRegistry {
12 tools: HashMap<String, Arc<dyn Tool>>,
13 cached_schema: Arc<Vec<Value>>,
15 api_to_runtime_names: HashMap<String, String>,
17 input_name_maps: HashMap<String, SchemaNameMap>,
19 catalog: ToolCatalog,
23}
24
25#[derive(Clone, Debug, Default)]
26struct SchemaNameMap {
27 api_to_runtime: HashMap<String, String>,
28 children: HashMap<String, SchemaNameMap>,
29 items: Option<Box<SchemaNameMap>>,
31}
32
33impl Default for ToolRegistry {
34 fn default() -> Self {
35 Self::new()
36 }
37}
38
39impl ToolRegistry {
40 pub fn new() -> Self {
41 let tools: Vec<Arc<dyn Tool>> = vec![
42 Arc::new(crate::tools::bash::BashTool),
43 #[cfg(windows)]
44 Arc::new(crate::tools::powershell::PowerShellTool),
45 Arc::new(crate::tools::read::ReadTool),
46 Arc::new(crate::tools::write::WriteTool),
47 Arc::new(crate::tools::edit::EditTool),
48 Arc::new(crate::tools::grep::GrepTool),
49 Arc::new(crate::tools::find::FindTool),
50 Arc::new(crate::tools::ls::LsTool),
51 Arc::new(crate::tools::subagent::SubagentTool),
52 Arc::new(crate::tools::subagent::authorize_model::SubagentModelAuthorizeTool),
53 Arc::new(crate::tools::subagent::models::SubagentModelsTool),
54 Arc::new(crate::tools::subagent::start::SubagentStartTool),
55 Arc::new(crate::tools::subagent::status::SubagentStatusTool),
56 Arc::new(crate::tools::subagent::steer::SubagentSteerTool),
57 Arc::new(crate::tools::subagent::collect::SubagentCollectTool),
58 Arc::new(crate::tools::subagent::resume::SubagentResumeTool),
59 Arc::new(crate::tools::memory_context::MemoryContextTool),
64 Arc::new(crate::tools::shell::ShellStartTool),
65 Arc::new(crate::tools::shell::ShellSendTool),
66 Arc::new(crate::tools::shell::ShellEndTool),
67 Arc::new(crate::tools::discovery::SearchToolsTool),
72 Arc::new(crate::tools::discovery::ActivateToolsTool),
73 Arc::new(crate::tools::memory::MemorySearchTool),
77 Arc::new(crate::tools::memory::MemoryFetchTool),
78 Arc::new(crate::tools::memory::MemoryStoreTool),
79 Arc::new(crate::tools::memory::MemoryForgetTool),
80 ];
81 Self::from_tools(tools)
82 }
83
84 pub fn empty() -> Self {
87 Self::from_tools(Vec::new())
88 }
89
90 pub fn disable(&mut self, names: &[String]) {
96 self.try_disable(names).expect(CATALOG_REJECTED_MSG);
97 }
98
99 pub fn try_disable(&mut self, names: &[String]) -> Result<(), CatalogError> {
112 if names.is_empty() {
113 return Ok(());
114 }
115 let remaining: Vec<Arc<dyn Tool>> = self
116 .tools
117 .values()
118 .filter(|t| !names.iter().any(|n| n == t.name()))
119 .cloned()
120 .collect();
121 if remaining.len() == self.tools.len() {
122 return Ok(());
123 }
124 let prior = self.catalog.generation();
125 let mut candidate = Self::try_from_tools(remaining)?;
126 candidate.catalog.rebase_past(prior)?;
127 *self = candidate;
128 Ok(())
129 }
130
131 pub fn without_subagent() -> Self {
133 let tools: Vec<Arc<dyn Tool>> = vec![
134 Arc::new(crate::tools::bash::BashTool),
135 #[cfg(windows)]
136 Arc::new(crate::tools::powershell::PowerShellTool),
137 Arc::new(crate::tools::read::ReadTool),
138 Arc::new(crate::tools::write::WriteTool),
139 Arc::new(crate::tools::edit::EditTool),
140 Arc::new(crate::tools::grep::GrepTool),
141 Arc::new(crate::tools::find::FindTool),
142 Arc::new(crate::tools::ls::LsTool),
143 Arc::new(crate::tools::shell::ShellStartTool),
144 Arc::new(crate::tools::shell::ShellSendTool),
145 Arc::new(crate::tools::shell::ShellEndTool),
146 ];
147 Self::from_tools(tools)
148 }
149
150 pub fn without_subagent_with_extensions(extension_tools: &ToolRegistry) -> Self {
158 let mut combined = Self::without_subagent();
159 for tool in extension_tools.tools.values() {
160 if tool.extension_id().is_some() && !is_recursive_subagent_tool_name(tool.name()) {
161 combined
162 .insert_tool_with_catalog(tool.clone())
163 .expect(CATALOG_REJECTED_MSG);
164 }
165 }
166 combined.rebuild_schema();
167 combined
168 }
169
170 fn from_tools(tool_list: Vec<Arc<dyn Tool>>) -> Self {
176 Self::try_from_tools(tool_list).expect(CATALOG_REJECTED_MSG)
177 }
178
179 #[doc(hidden)]
183 pub fn from_tools_for_tests(tool_list: Vec<Arc<dyn Tool>>) -> Self {
184 Self::from_tools(tool_list)
185 }
186
187 fn try_from_tools(tool_list: Vec<Arc<dyn Tool>>) -> Result<Self, CatalogError> {
188 let mut registry = ToolRegistry {
189 tools: HashMap::new(),
190 cached_schema: Arc::new(Vec::new()),
191 api_to_runtime_names: HashMap::new(),
192 input_name_maps: HashMap::new(),
193 catalog: ToolCatalog::empty(),
194 };
195 for tool in tool_list {
199 let name = tool.name().to_string();
200 registry.tools.insert(name, tool);
201 }
202 let sorted: Vec<Arc<dyn Tool>> =
208 registry.iter_tools_sorted().into_iter().cloned().collect();
209 for tool in sorted {
210 let record = CapabilityRecord::for_registered_tool(&tool);
211 registry.catalog.upsert(None, record)?;
212 }
213 registry.rebuild_schema();
214 Ok(registry)
215 }
216
217 fn api_safe_name(name: &str, used: &HashSet<String>) -> String {
218 Self::api_safe_identifier(name, used, 128, false)
219 }
220
221 fn api_safe_property_name(name: &str, used: &HashSet<String>) -> String {
222 Self::api_safe_identifier(name, used, 64, true)
223 }
224
225 fn api_safe_identifier(
226 name: &str,
227 used: &HashSet<String>,
228 max_len: usize,
229 allow_dot: bool,
230 ) -> String {
231 let mut sanitized = String::with_capacity(name.len());
232 for ch in name.chars() {
233 if ch.is_ascii_alphanumeric() || ch == '_' || ch == '-' || (allow_dot && ch == '.') {
234 sanitized.push(ch);
235 } else {
236 sanitized.push('_');
237 }
238 }
239 if sanitized.is_empty() {
240 sanitized.push_str("field");
241 }
242 if sanitized.len() > max_len {
243 sanitized.truncate(max_len);
244 }
245
246 let base = sanitized.clone();
247 let mut suffix = 2;
248 while used.contains(&sanitized) {
249 let suffix_str = format!("_{suffix}");
250 let keep = max_len.saturating_sub(suffix_str.len());
251 sanitized = format!("{}{}", &base[..base.len().min(keep)], suffix_str);
252 suffix += 1;
253 }
254 sanitized
255 }
256
257 fn sanitize_schema(mut schema: Value) -> (Value, SchemaNameMap) {
258 let mut map = SchemaNameMap::default();
259 let Some(obj) = schema.as_object_mut() else {
260 return (schema, map);
261 };
262
263 let mut required_name_map = HashMap::new();
264 if let Some(props_value) = obj.get_mut("properties") {
265 if let Some(props) = props_value.as_object_mut() {
266 let original = std::mem::take(props);
267 let mut used = HashSet::new();
268 for (runtime_name, child_schema) in original {
269 let api_name = Self::api_safe_property_name(&runtime_name, &used);
270 used.insert(api_name.clone());
271 required_name_map.insert(runtime_name.clone(), api_name.clone());
272 map.api_to_runtime.insert(api_name.clone(), runtime_name);
273
274 let (sanitized_child, child_map) = Self::sanitize_schema(child_schema);
275 if !child_map.api_to_runtime.is_empty() || !child_map.children.is_empty() {
276 map.children.insert(api_name.clone(), child_map);
277 }
278 props.insert(api_name, sanitized_child);
279 }
280 }
281 }
282
283 if let Some(required) = obj.get_mut("required").and_then(Value::as_array_mut) {
284 for item in required.iter_mut() {
285 if let Some(name) = item.as_str() {
286 if let Some(api_name) = required_name_map.get(name) {
287 *item = Value::String(api_name.clone());
288 }
289 }
290 }
291 }
292
293 if let Some(items) = obj.get_mut("items") {
296 let (sanitized_items, items_map) = Self::sanitize_schema(std::mem::take(items));
297 if !items_map.api_to_runtime.is_empty()
298 || !items_map.children.is_empty()
299 || items_map.items.is_some()
300 {
301 map.items = Some(Box::new(items_map));
302 }
303 *items = sanitized_items;
304 }
305
306 (schema, map)
307 }
308
309 fn translate_input_names(input: Value, map: &SchemaNameMap) -> Value {
310 match input {
311 Value::Object(obj) => {
312 let mut out = serde_json::Map::new();
313 for (api_name, value) in obj {
314 let runtime_name = map
315 .api_to_runtime
316 .get(&api_name)
317 .cloned()
318 .unwrap_or_else(|| api_name.clone());
319 let value = if let Some(child) = map.children.get(&api_name) {
320 Self::translate_input_names(value, child)
321 } else {
322 value
323 };
324 out.insert(runtime_name, value);
325 }
326 Value::Object(out)
327 }
328 Value::Array(arr) => {
329 if let Some(items_map) = &map.items {
331 Value::Array(
332 arr.into_iter()
333 .map(|v| Self::translate_input_names(v, items_map))
334 .collect(),
335 )
336 } else {
337 Value::Array(arr)
338 }
339 }
340 other => other,
341 }
342 }
343
344 fn rebuild_schema(&mut self) {
345 let mut used = HashSet::new();
346 let mut api_to_runtime_names = HashMap::new();
347 let mut input_name_maps = HashMap::new();
348 let mut schema = Vec::with_capacity(self.tools.len());
349
350 let mut sorted_tools: Vec<_> = self.tools.values().collect();
365 sorted_tools.sort_by_key(|t| {
366 let builtin_first = match t.origin() {
367 crate::tools::ToolOrigin::Builtin => 0u8,
368 _ => 1u8,
369 };
370 (builtin_first, t.name().to_string())
371 });
372
373 for tool in sorted_tools {
374 let runtime_name = tool.name();
375 let api_name = Self::api_safe_name(runtime_name, &used);
376 used.insert(api_name.clone());
377 api_to_runtime_names.insert(api_name.clone(), runtime_name.to_string());
378 let (input_schema, input_map) = Self::sanitize_schema(tool.parameters());
379 input_name_maps.insert(api_name.clone(), input_map);
380 schema.push(serde_json::json!({
381 "name": api_name,
382 "description": tool.description(),
383 "input_schema": input_schema
384 }));
385 }
386
387 self.api_to_runtime_names = api_to_runtime_names;
388 self.input_name_maps = input_name_maps;
389 self.cached_schema = Arc::new(schema);
390 }
391
392 pub fn register(&mut self, tool: Arc<dyn Tool>) {
402 self.try_register(tool).expect(CATALOG_REJECTED_MSG);
403 }
404
405 pub fn try_register(&mut self, tool: Arc<dyn Tool>) -> Result<(), CatalogError> {
411 self.insert_tool_with_catalog(tool)?;
412 self.rebuild_schema();
413 Ok(())
414 }
415
416 pub fn try_register_batch(
430 &mut self,
431 tool_list: Vec<Arc<dyn Tool>>,
432 ) -> Result<usize, CatalogError> {
433 if tool_list.is_empty() {
434 return Ok(0);
435 }
436 let mut batch_names: HashSet<String> = HashSet::new();
437 let mut batch_ids: HashSet<crate::tools::catalog::ToolId> = HashSet::new();
438 for tool in &tool_list {
439 let name = tool.name();
440 if self.tools.contains_key(name) || !batch_names.insert(name.to_string()) {
441 return Err(CatalogError::DuplicateRuntimeName(
442 agent_core::BoundedText::new(name, 128).text,
443 ));
444 }
445 let id = tool_id_for(tool.as_ref());
446 if self.catalog.get(&id).is_some() || !batch_ids.insert(id.clone()) {
447 return Err(CatalogError::DuplicateToolId(id));
448 }
449 }
450 let count = tool_list.len();
451 let mut candidate = self.clone();
452 for tool in tool_list {
453 candidate.insert_tool_with_catalog(tool)?;
454 }
455 candidate.rebuild_schema();
456 *self = candidate;
457 Ok(count)
458 }
459
460 fn insert_tool_with_catalog(&mut self, tool: Arc<dyn Tool>) -> Result<(), CatalogError> {
469 let record = CapabilityRecord::for_registered_tool(&tool);
470 let replaced = self
471 .tools
472 .get(tool.name())
473 .map(|old| tool_id_for(old.as_ref()));
474 self.catalog.upsert(replaced.as_ref(), record)?;
475 self.tools.insert(tool.name().to_string(), tool);
476 Ok(())
477 }
478
479 pub fn catalog(&self) -> &ToolCatalog {
482 &self.catalog
483 }
484
485 #[doc(hidden)]
493 pub fn resume_catalog_generation_for_tests(
494 &mut self,
495 generation: crate::tools::catalog::CatalogGeneration,
496 ) {
497 self.catalog.set_generation_for_tests(generation);
498 }
499
500 pub fn get(&self, name: &str) -> Option<&Arc<dyn Tool>> {
501 let runtime_name = self
502 .api_to_runtime_names
503 .get(name)
504 .map(String::as_str)
505 .unwrap_or(name);
506 self.tools.get(runtime_name)
507 }
508
509 pub fn runtime_name_for_api<'a>(&'a self, name: &'a str) -> &'a str {
510 self.api_to_runtime_names
511 .get(name)
512 .map(String::as_str)
513 .unwrap_or(name)
514 }
515
516 pub fn resolve_wire_call(
526 &self,
527 wire_name: &str,
528 ) -> Option<crate::tools::activation::ResolvedToolCall> {
529 let runtime_name = self
530 .api_to_runtime_names
531 .get(wire_name)
532 .map(String::as_str)
533 .unwrap_or(wire_name);
534 let tool = self.tools.get(runtime_name)?;
535 Some(crate::tools::activation::ResolvedToolCall::new(
536 wire_name,
537 runtime_name,
538 tool_id_for(tool.as_ref()),
539 ))
540 }
541
542 pub fn translate_input_for_api_tool(&self, tool_name: &str, input: Value) -> Value {
543 if let Some(map) = self.input_name_maps.get(tool_name) {
544 Self::translate_input_names(input, map)
545 } else {
546 input
547 }
548 }
549
550 pub fn tools_schema(&self) -> Arc<Vec<Value>> {
551 Arc::clone(&self.cached_schema)
552 }
553
554 pub fn session_tools_schema(
564 &self,
565 session: &crate::tools::activation::SessionToolSet,
566 ) -> Result<Vec<Value>, SessionSchemaError> {
567 if session.is_stale(&self.catalog) {
568 return Err(SessionSchemaError::StaleSessionSet {
569 set: session.catalog_generation(),
570 catalog: self.catalog.generation(),
571 });
572 }
573 let mut projected = Vec::new();
574 for entry in self.cached_schema.iter() {
575 let Some(api_name) = entry["name"].as_str() else {
576 continue;
577 };
578 let runtime_name = self.runtime_name_for_api(api_name);
579 let Some(tool) = self.tools.get(runtime_name) else {
580 continue;
581 };
582 let id = tool_id_for(tool.as_ref());
583 let pinned = session
584 .core_schema_digest(&id)
585 .or_else(|| session.activation(&id).map(|a| a.schema_digest()));
586 let Some(pinned) = pinned else {
587 continue; };
589 let record = self
590 .catalog
591 .get(&id)
592 .ok_or_else(|| SessionSchemaError::NotCataloged(id.clone()))?;
593 if pinned != record.schema_digest() {
594 return Err(SessionSchemaError::SchemaDigestMismatch(id));
595 }
596 projected.push(entry.clone());
597 }
598 Ok(projected)
599 }
600
601 pub fn tool_names_for_extension(&self, extension_id: &str) -> Vec<String> {
604 let mut names: Vec<String> = self
605 .tools
606 .values()
607 .filter(|t| t.extension_id() == Some(extension_id))
608 .map(|t| t.name().to_string())
609 .collect();
610 names.sort();
611 names
612 }
613
614 pub fn iter_tools_sorted(&self) -> Vec<&Arc<dyn Tool>> {
617 let mut tools: Vec<_> = self.tools.values().collect();
618 tools.sort_by_key(|t| t.name());
619 tools
620 }
621}
622
623#[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)]
626pub enum SessionSchemaError {
627 #[error(
628 "session tool set generation {} is stale against catalog generation {}",
629 set.value(),
630 catalog.value()
631 )]
632 StaleSessionSet {
633 set: crate::tools::catalog::CatalogGeneration,
634 catalog: crate::tools::catalog::CatalogGeneration,
635 },
636 #[error("session member is not cataloged: {0}")]
637 NotCataloged(crate::tools::catalog::ToolId),
638 #[error("schema digest changed since the session pinned it: {0}")]
639 SchemaDigestMismatch(crate::tools::catalog::ToolId),
640}
641
642const CATALOG_REJECTED_MSG: &str =
648 "tool catalog rejected mutation: refusing to expose a tool the catalog could not record";
649
650fn is_recursive_subagent_tool_name(name: &str) -> bool {
651 let api_name = ToolRegistry::api_safe_name(name, &HashSet::new());
652 matches!(
653 api_name.as_str(),
654 "subagent"
655 | "subagent_start"
656 | "subagent_status"
657 | "subagent_steer"
658 | "subagent_collect"
659 | "subagent_resume"
660 | "subagent_model_authorize"
661 | "subagent_models"
662 | "memory_context"
666 )
667}
668#[cfg(test)]
669mod tests {
670 use super::*;
671 use crate::{Result, ToolContext};
672 use serde_json::json;
673
674 struct NamedTool(&'static str);
675
676 #[async_trait::async_trait]
677 impl Tool for NamedTool {
678 fn name(&self) -> &str {
679 self.0
680 }
681 fn description(&self) -> &str {
682 "test tool"
683 }
684 fn parameters(&self) -> Value {
685 json!({"type": "object"})
686 }
687 async fn execute(&self, _params: Value, _ctx: ToolContext) -> Result<String> {
688 Ok("ok".to_string())
689 }
690 }
691
692 struct SchemaTool;
693
694 #[async_trait::async_trait]
695 impl Tool for SchemaTool {
696 fn name(&self) -> &str {
697 "schema_tool"
698 }
699 fn description(&self) -> &str {
700 "schema tool"
701 }
702 fn parameters(&self) -> Value {
703 json!({
704 "type": "object",
705 "properties": {
706 "bad:key/that/is/far/too/long/for/anthropic/property/names/and/keeps/going": {"type": "string"},
707 "nested:obj": {
708 "type": "object",
709 "properties": {"inner/key": {"type": "string"}},
710 "required": ["inner/key"]
711 }
712 },
713 "required": [
714 "bad:key/that/is/far/too/long/for/anthropic/property/names/and/keeps/going",
715 "nested:obj"
716 ]
717 })
718 }
719 async fn execute(&self, _params: Value, _ctx: ToolContext) -> Result<String> {
720 Ok("ok".to_string())
721 }
722 }
723
724 #[test]
725 fn tool_schema_uses_api_safe_names_and_maps_back() {
726 let registry = ToolRegistry::from_tools(vec![Arc::new(NamedTool("plugin:skill.tool"))]);
727
728 assert_eq!(registry.tools_schema()[0]["name"], "plugin_skill_tool");
729 assert!(registry.get("plugin:skill.tool").is_some());
730 assert!(registry.get("plugin_skill_tool").is_some());
731 assert_eq!(
732 registry.runtime_name_for_api("plugin_skill_tool"),
733 "plugin:skill.tool"
734 );
735 }
736
737 #[test]
747 fn extension_leaf_name_never_shadows_builtin_memory_fetch_api_name() {
748 struct ExtMemoryTool {
749 runtime: &'static str,
750 ext: &'static str,
751 }
752 #[async_trait::async_trait]
753 impl Tool for ExtMemoryTool {
754 fn origin(&self) -> crate::tools::ToolOrigin {
755 crate::tools::ToolOrigin::Extension {
756 extension_id: self.ext.to_string(),
757 }
758 }
759 fn extension_id(&self) -> Option<&str> {
760 Some(self.ext)
761 }
762 fn name(&self) -> &str {
763 self.runtime
764 }
765 fn description(&self) -> &str {
766 "extension memory tool (single-id schema)"
767 }
768 fn parameters(&self) -> Value {
769 json!({
770 "type": "object",
771 "properties": {"id": {"type": "string"}},
772 "required": ["id"]
773 })
774 }
775 async fn execute(&self, _params: Value, _ctx: ToolContext) -> Result<String> {
776 Ok("ok".to_string())
777 }
778 }
779
780 let registry = ToolRegistry::from_tools(vec![
781 Arc::new(crate::tools::memory::MemoryFetchTool),
782 Arc::new(ExtMemoryTool {
784 runtime: "memory:fetch",
785 ext: "memory",
786 }),
787 Arc::new(ExtMemoryTool {
789 runtime: "axel-memory-manager:memory_fetch",
790 ext: "axel-memory-manager",
791 }),
792 ]);
793
794 assert_eq!(
796 registry.runtime_name_for_api("memory_fetch"),
797 "memory_fetch",
798 "wire name `memory_fetch` must resolve to the built-in"
799 );
800 let schema = registry.tools_schema();
801 let builtin = schema
802 .iter()
803 .find(|e| e["name"] == "memory_fetch")
804 .expect("built-in memory_fetch exposed");
805 assert_eq!(
806 builtin["input_schema"]["required"],
807 json!(["ids"]),
808 "built-in memory_fetch schema must remain {{ids: [...]}}"
809 );
810
811 let axel = schema
813 .iter()
814 .find(|e| e["name"] == "axel-memory-manager_memory_fetch")
815 .expect("extension tool exposed under namespaced API name");
816 assert_eq!(axel["input_schema"]["required"], json!(["id"]));
817 assert_eq!(
820 registry.runtime_name_for_api("memory_fetch_2"),
821 "memory:fetch"
822 );
823 let resolved = registry.resolve_wire_call("memory_fetch").unwrap();
826 assert_eq!(resolved.runtime_name(), "memory_fetch");
827 }
828
829 #[test]
830 fn tool_schema_disambiguates_sanitized_name_collisions() {
831 let registry =
832 ToolRegistry::from_tools(vec![Arc::new(NamedTool("a:b")), Arc::new(NamedTool("a.b"))]);
833 let names: HashSet<String> = registry
834 .tools_schema()
835 .iter()
836 .filter_map(|s| s["name"].as_str().map(str::to_string))
837 .collect();
838
839 assert_eq!(names.len(), 2);
840 assert!(names.contains("a_b"));
841 assert!(names.contains("a_b_2"));
842 assert!(registry.get("a_b").is_some());
843 assert!(registry.get("a_b_2").is_some());
844 }
845
846 #[test]
847 fn tool_schema_truncates_long_names_to_anthropic_limit() {
848 let long = "x".repeat(140);
849 let leaked: &'static str = Box::leak(long.into_boxed_str());
850 let registry = ToolRegistry::from_tools(vec![Arc::new(NamedTool(leaked))]);
851 let schema = registry.tools_schema();
852 let name = schema[0]["name"].as_str().unwrap();
853
854 assert_eq!(name.len(), 128);
855 assert!(name
856 .chars()
857 .all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-'));
858 assert!(registry.get(name).is_some());
859 }
860
861 #[test]
862 fn tool_schema_sanitizes_input_property_names_and_translates_inputs_back() {
863 let registry = ToolRegistry::from_tools(vec![Arc::new(SchemaTool)]);
864 let schema = registry.tools_schema();
865 let input_schema = &schema[0]["input_schema"];
866 let props = input_schema["properties"].as_object().unwrap();
867
868 assert!(props.keys().all(|k| k.len() <= 64
869 && k.chars()
870 .all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-' || c == '.')));
871 assert_eq!(
872 input_schema["required"].as_array().unwrap()[0]
873 .as_str()
874 .unwrap()
875 .len(),
876 64
877 );
878 assert_eq!(input_schema["required"][1], "nested_obj");
879 assert!(props["nested_obj"]["properties"]
880 .as_object()
881 .unwrap()
882 .contains_key("inner_key"));
883 assert_eq!(props["nested_obj"]["required"][0], "inner_key");
884
885 let first_required = input_schema["required"][0].as_str().unwrap();
886 let translated = registry.translate_input_for_api_tool(
887 "schema_tool",
888 json!({
889 first_required: "value",
890 "nested_obj": {"inner_key": "nested"}
891 }),
892 );
893
894 assert_eq!(
895 translated["bad:key/that/is/far/too/long/for/anthropic/property/names/and/keeps/going"],
896 "value"
897 );
898 assert_eq!(translated["nested:obj"]["inner/key"], "nested");
899 }
900
901 #[test]
902 fn test_tool_registry_new() {
903 let registry = ToolRegistry::new();
904
905 assert_eq!(registry.tools_schema().len(), 25);
911
912 assert!(registry.get("bash").is_some());
914
915 assert!(registry.get("nonexistent").is_none());
917
918 assert!(registry.get("bash").is_some());
920 assert!(registry.get("read").is_some());
921 assert!(registry.get("write").is_some());
922 assert!(registry.get("edit").is_some());
923 assert!(registry.get("grep").is_some());
924 assert!(registry.get("find").is_some());
925 assert!(registry.get("ls").is_some());
926 assert!(registry.get("subagent").is_some());
927 assert!(registry.get("subagent_model_authorize").is_some());
928 assert!(registry.get("subagent_models").is_some());
929 }
930
931 #[test]
932 fn test_disable_removes_tools_from_registry_and_schema() {
933 let mut registry = ToolRegistry::new();
934 let before = registry.tools_schema().len();
935
936 registry.disable(&["bash".to_string(), "ls".to_string()]);
937
938 assert!(registry.get("bash").is_none());
940 assert!(registry.get("ls").is_none());
941 assert!(registry.get("read").is_some());
943 assert!(registry.get("subagent").is_some());
944 assert_eq!(registry.tools_schema().len(), before - 2);
946 }
947
948 #[test]
949 fn test_disable_empty_and_unknown_are_noops() {
950 let mut registry = ToolRegistry::new();
951 let n = registry.tools_schema().len();
952 registry.disable(&[]); assert_eq!(registry.tools_schema().len(), n);
954 registry.disable(&["not_a_real_tool".to_string()]); assert_eq!(registry.tools_schema().len(), n);
956 assert!(registry.get("bash").is_some());
957 }
958
959 #[test]
960 fn test_tool_registry_without_subagent() {
961 let registry = ToolRegistry::without_subagent();
962
963 assert_eq!(registry.tools_schema().len(), 10);
965 assert!(registry.get("subagent").is_none());
967
968 assert!(registry.get("bash").is_some());
970
971 assert!(registry.get("bash").is_some());
973 assert!(registry.get("read").is_some());
974 assert!(registry.get("write").is_some());
975 assert!(registry.get("edit").is_some());
976 assert!(registry.get("grep").is_some());
977 assert!(registry.get("find").is_some());
978 assert!(registry.get("ls").is_some());
979 }
980
981 #[test]
982 fn test_tool_registry_register() {
983 let mut registry = ToolRegistry::without_subagent();
984 let initial_count = registry.tools_schema().len();
985
986 struct TestTool;
988 #[async_trait::async_trait]
989 impl Tool for TestTool {
990 fn name(&self) -> &str {
991 "test_tool"
992 }
993 fn description(&self) -> &str {
994 "A test tool"
995 }
996 fn parameters(&self) -> Value {
997 json!({"type": "object"})
998 }
999 async fn execute(&self, _params: Value, _ctx: ToolContext) -> Result<String> {
1000 Ok("test result".to_string())
1001 }
1002 }
1003
1004 registry.register(Arc::new(TestTool));
1005
1006 assert_eq!(registry.tools_schema().len(), initial_count + 1);
1008
1009 assert!(registry.get("test_tool").is_some());
1011 }
1012
1013 #[test]
1014 fn tool_names_for_extension_filters_by_owner_and_sorts() {
1015 struct OwnedTool(&'static str, Option<&'static str>);
1016 #[async_trait::async_trait]
1017 impl Tool for OwnedTool {
1018 fn name(&self) -> &str {
1019 self.0
1020 }
1021 fn description(&self) -> &str {
1022 "owned"
1023 }
1024 fn parameters(&self) -> Value {
1025 json!({"type": "object"})
1026 }
1027 async fn execute(&self, _params: Value, _ctx: ToolContext) -> Result<String> {
1028 Ok("ok".to_string())
1029 }
1030 fn extension_id(&self) -> Option<&str> {
1031 self.1
1032 }
1033 }
1034
1035 let mut registry = ToolRegistry::without_subagent();
1036 registry.register(Arc::new(OwnedTool("alpha:zed", Some("alpha"))));
1037 registry.register(Arc::new(OwnedTool("alpha:bar", Some("alpha"))));
1038 registry.register(Arc::new(OwnedTool("beta:thing", Some("beta"))));
1039
1040 assert_eq!(
1041 registry.tool_names_for_extension("alpha"),
1042 vec!["alpha:bar".to_string(), "alpha:zed".to_string()]
1043 );
1044 assert_eq!(
1045 registry.tool_names_for_extension("beta"),
1046 vec!["beta:thing".to_string()]
1047 );
1048 assert!(registry.tool_names_for_extension("ghost").is_empty());
1049 assert!(registry.tool_names_for_extension("bash").is_empty());
1051 }
1052
1053 struct OwnedTool(&'static str, Option<&'static str>);
1054 #[async_trait::async_trait]
1055 impl Tool for OwnedTool {
1056 fn name(&self) -> &str {
1057 self.0
1058 }
1059 fn description(&self) -> &str {
1060 "owned"
1061 }
1062 fn parameters(&self) -> Value {
1063 json!({"type": "object"})
1064 }
1065 async fn execute(&self, _params: Value, _ctx: ToolContext) -> Result<String> {
1066 Ok("ok".to_string())
1067 }
1068 fn extension_id(&self) -> Option<&str> {
1069 self.1
1070 }
1071 }
1072
1073 #[test]
1074 fn without_subagent_excludes_subagent_tools() {
1075 let registry = ToolRegistry::without_subagent();
1076 assert!(registry.get("subagent").is_none());
1077 assert!(registry.get("subagent_start").is_none());
1078 assert!(registry.get("subagent_status").is_none());
1079 assert!(registry.get("subagent_steer").is_none());
1080 assert!(registry.get("subagent_collect").is_none());
1081 assert!(registry.get("subagent_resume").is_none());
1082 assert!(registry.get("bash").is_some());
1084 assert!(registry.get("read").is_some());
1085 }
1086
1087 #[test]
1088 fn without_subagent_with_extensions_includes_extension_tools() {
1089 let mut other = ToolRegistry::empty();
1090 other.register(Arc::new(OwnedTool("alpha:do_thing", Some("alpha"))));
1091
1092 let merged = ToolRegistry::without_subagent_with_extensions(&other);
1093
1094 assert!(merged.get("alpha:do_thing").is_some());
1096 assert!(merged.get("bash").is_some());
1098 assert!(merged.get("read").is_some());
1099 assert!(merged.get("subagent_start").is_none());
1101 }
1102
1103 #[test]
1104 fn without_subagent_with_extensions_excludes_built_ins_from_other_registry() {
1105 let other = ToolRegistry::new();
1110
1111 let merged = ToolRegistry::without_subagent_with_extensions(&other);
1112
1113 let bash = merged.get("bash").expect("bash present");
1115 assert!(bash.extension_id().is_none());
1116 assert!(merged.get("subagent_start").is_none());
1118 assert!(merged.get("subagent").is_none());
1119 }
1120
1121 #[test]
1122 fn without_subagent_with_extensions_rejects_recursive_tool_name_collisions() {
1123 let mut other = ToolRegistry::empty();
1124 for name in [
1125 "subagent",
1126 "subagent_start",
1127 "subagent_status",
1128 "subagent_steer",
1129 "subagent_collect",
1130 "subagent_resume",
1131 "subagent_model_authorize",
1132 "subagent_models",
1133 "memory_context",
1134 ] {
1135 other.register(Arc::new(OwnedTool(name, Some("malicious-extension"))));
1136 }
1137
1138 let merged = ToolRegistry::without_subagent_with_extensions(&other);
1139 for name in [
1140 "subagent",
1141 "subagent_start",
1142 "subagent_status",
1143 "subagent_steer",
1144 "subagent_collect",
1145 "subagent_resume",
1146 "subagent_model_authorize",
1147 "subagent_models",
1148 "memory_context",
1149 ] {
1150 assert!(merged.get(name).is_none(), "{name} must stay unavailable");
1151 }
1152 }
1153
1154 #[test]
1155 fn without_subagent_with_extensions_rejects_api_sanitized_recursive_collisions() {
1156 let mut other = ToolRegistry::empty();
1157 for runtime_name in [
1158 "subagent:start",
1159 "subagent:status",
1160 "subagent:steer",
1161 "subagent:collect",
1162 "subagent:resume",
1163 "subagent:model:authorize",
1164 "subagent:models",
1165 "memory:context",
1166 ] {
1167 other.register(Arc::new(OwnedTool(runtime_name, Some("subagent"))));
1168 }
1169
1170 let merged = ToolRegistry::without_subagent_with_extensions(&other);
1171 for api_name in [
1172 "subagent_start",
1173 "subagent_status",
1174 "subagent_steer",
1175 "subagent_collect",
1176 "subagent_resume",
1177 "subagent_model_authorize",
1178 "subagent_models",
1179 "memory_context",
1180 ] {
1181 assert!(
1182 merged.get(api_name).is_none(),
1183 "API-safe recursive name {api_name} must stay unavailable"
1184 );
1185 }
1186 }
1187
1188 #[test]
1192 fn memory_context_present_in_primary_registry_absent_in_recursive_ones() {
1193 assert!(ToolRegistry::new().get("memory_context").is_some());
1194 assert!(ToolRegistry::without_subagent()
1195 .get("memory_context")
1196 .is_none());
1197
1198 let merged = ToolRegistry::without_subagent_with_extensions(&ToolRegistry::new());
1199 assert!(
1200 merged.get("memory_context").is_none(),
1201 "recursive registries must never expose memory_context"
1202 );
1203 }
1204
1205 #[test]
1206 fn without_subagent_with_extensions_does_not_overwrite_existing_builtin() {
1207 let mut other = ToolRegistry::empty();
1213 other.register(Arc::new(OwnedTool("ext:custom", Some("ext"))));
1214
1215 let merged = ToolRegistry::without_subagent_with_extensions(&other);
1216 assert!(merged.get("ext:custom").is_some());
1217 assert!(merged.get("bash").is_some());
1218 assert!(merged.get("bash").unwrap().extension_id().is_none());
1219 }
1220
1221 struct McpTwin(&'static str);
1223
1224 #[async_trait::async_trait]
1225 impl Tool for McpTwin {
1226 fn name(&self) -> &str {
1227 self.0
1228 }
1229 fn description(&self) -> &str {
1230 "mcp twin"
1231 }
1232 fn parameters(&self) -> Value {
1233 json!({"type": "object"})
1234 }
1235 async fn execute(&self, _params: Value, _ctx: ToolContext) -> Result<String> {
1236 Ok("ok".to_string())
1237 }
1238 fn origin(&self) -> crate::tools::ToolOrigin {
1239 crate::tools::ToolOrigin::Mcp {
1240 server_id: "srv".to_string(),
1241 server_tool_name: "shared".to_string(),
1242 }
1243 }
1244 }
1245
1246 #[test]
1247 #[should_panic(expected = "tool catalog rejected mutation")]
1248 fn from_tools_with_duplicate_capability_identity_fails_closed() {
1249 let _ = ToolRegistry::from_tools(vec![
1253 Arc::new(McpTwin("twin_one")),
1254 Arc::new(McpTwin("twin_two")),
1255 ]);
1256 }
1257}