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 #[doc(hidden)]
506 pub fn upsert_catalog_record_for_tests(
507 &mut self,
508 replaced: Option<&crate::tools::catalog::ToolId>,
509 record: crate::tools::catalog::CapabilityRecord,
510 ) {
511 self.catalog
512 .upsert(replaced, record)
513 .expect("test upsert must be structurally valid");
514 }
515
516 pub fn get(&self, name: &str) -> Option<&Arc<dyn Tool>> {
517 let runtime_name = self
518 .api_to_runtime_names
519 .get(name)
520 .map(String::as_str)
521 .unwrap_or(name);
522 self.tools.get(runtime_name)
523 }
524
525 pub fn runtime_name_for_api<'a>(&'a self, name: &'a str) -> &'a str {
526 self.api_to_runtime_names
527 .get(name)
528 .map(String::as_str)
529 .unwrap_or(name)
530 }
531
532 pub fn resolve_wire_call(
542 &self,
543 wire_name: &str,
544 ) -> Option<crate::tools::activation::ResolvedToolCall> {
545 let runtime_name = self
546 .api_to_runtime_names
547 .get(wire_name)
548 .map(String::as_str)
549 .unwrap_or(wire_name);
550 let tool = self.tools.get(runtime_name)?;
551 Some(crate::tools::activation::ResolvedToolCall::new(
552 wire_name,
553 runtime_name,
554 tool_id_for(tool.as_ref()),
555 ))
556 }
557
558 pub fn translate_input_for_api_tool(&self, tool_name: &str, input: Value) -> Value {
559 if let Some(map) = self.input_name_maps.get(tool_name) {
560 Self::translate_input_names(input, map)
561 } else {
562 input
563 }
564 }
565
566 pub fn tools_schema(&self) -> Arc<Vec<Value>> {
567 Arc::clone(&self.cached_schema)
568 }
569
570 pub fn session_tools_schema(
589 &self,
590 session: &crate::tools::activation::SessionToolSet,
591 ) -> SessionSchemaProjection {
592 let mut projected = Vec::new();
593 let mut dropped = Vec::new();
594 for entry in self.cached_schema.iter() {
595 let Some(api_name) = entry["name"].as_str() else {
596 continue;
597 };
598 let runtime_name = self.runtime_name_for_api(api_name);
599 let Some(tool) = self.tools.get(runtime_name) else {
600 continue;
601 };
602 let id = tool_id_for(tool.as_ref());
603 let pinned = session
604 .core_pin(&id)
605 .map(|pin| (pin.schema_digest(), pin.provenance()))
606 .or_else(|| {
607 session
608 .activation(&id)
609 .map(|a| (a.schema_digest(), a.provenance()))
610 });
611 let Some((pinned_digest, pinned_provenance)) = pinned else {
612 continue; };
614 let Some(record) = self.catalog.get(&id) else {
615 tracing::warn!(
616 tool = %id,
617 "session schema projection drops uncataloged session member \
618 (removed since the session pinned it)"
619 );
620 dropped.push((id, DroppedSessionMember::Uncataloged));
621 continue;
622 };
623 if pinned_digest != record.schema_digest() {
624 tracing::warn!(
625 tool = %id,
626 "session schema projection drops drifted session member \
627 (schema digest changed since the session pinned it)"
628 );
629 dropped.push((id, DroppedSessionMember::SchemaDigestMismatch));
630 continue;
631 }
632 if pinned_provenance != record.provenance() {
633 tracing::warn!(
634 tool = %id,
635 "session schema projection drops drifted session member \
636 (trust provenance changed since the session pinned it)"
637 );
638 dropped.push((id, DroppedSessionMember::ProvenanceMismatch));
639 continue;
640 }
641 projected.push(entry.clone());
642 }
643 let seen: HashSet<&crate::tools::catalog::ToolId> =
650 dropped.iter().map(|(id, _)| id).collect();
651 let vanished: Vec<crate::tools::catalog::ToolId> = session
652 .core_ids()
653 .chain(session.activated().map(|a| a.grant().tool_id()))
654 .filter(|id| !seen.contains(id) && self.catalog.get(id).is_none())
655 .cloned()
656 .collect();
657 for id in vanished {
658 tracing::warn!(
659 tool = %id,
660 "session schema projection drops uncataloged session member (removed from the registry since the session pinned it)"
661 );
662 dropped.push((id, DroppedSessionMember::Uncataloged));
663 }
664 SessionSchemaProjection {
665 schema: projected,
666 dropped,
667 }
668 }
669
670 pub fn tool_names_for_extension(&self, extension_id: &str) -> Vec<String> {
673 let mut names: Vec<String> = self
674 .tools
675 .values()
676 .filter(|t| t.extension_id() == Some(extension_id))
677 .map(|t| t.name().to_string())
678 .collect();
679 names.sort();
680 names
681 }
682
683 pub fn iter_tools_sorted(&self) -> Vec<&Arc<dyn Tool>> {
686 let mut tools: Vec<_> = self.tools.values().collect();
687 tools.sort_by_key(|t| t.name());
688 tools
689 }
690}
691
692#[derive(Clone, Debug, PartialEq, Eq)]
698pub struct SessionSchemaProjection {
699 pub schema: Vec<Value>,
701 pub dropped: Vec<(crate::tools::catalog::ToolId, DroppedSessionMember)>,
705}
706
707#[derive(Clone, Copy, Debug, PartialEq, Eq)]
709pub enum DroppedSessionMember {
710 Uncataloged,
712 SchemaDigestMismatch,
714 ProvenanceMismatch,
717}
718
719const CATALOG_REJECTED_MSG: &str =
725 "tool catalog rejected mutation: refusing to expose a tool the catalog could not record";
726
727fn is_recursive_subagent_tool_name(name: &str) -> bool {
728 let api_name = ToolRegistry::api_safe_name(name, &HashSet::new());
729 matches!(
730 api_name.as_str(),
731 "subagent"
732 | "subagent_start"
733 | "subagent_status"
734 | "subagent_steer"
735 | "subagent_collect"
736 | "subagent_resume"
737 | "subagent_model_authorize"
738 | "subagent_models"
739 | "memory_context"
743 )
744}
745#[cfg(test)]
746mod tests {
747 use super::*;
748 use crate::{Result, ToolContext};
749 use serde_json::json;
750
751 struct NamedTool(&'static str);
752
753 #[async_trait::async_trait]
754 impl Tool for NamedTool {
755 fn name(&self) -> &str {
756 self.0
757 }
758 fn description(&self) -> &str {
759 "test tool"
760 }
761 fn parameters(&self) -> Value {
762 json!({"type": "object"})
763 }
764 async fn execute(&self, _params: Value, _ctx: ToolContext) -> Result<String> {
765 Ok("ok".to_string())
766 }
767 }
768
769 struct SchemaTool;
770
771 #[async_trait::async_trait]
772 impl Tool for SchemaTool {
773 fn name(&self) -> &str {
774 "schema_tool"
775 }
776 fn description(&self) -> &str {
777 "schema tool"
778 }
779 fn parameters(&self) -> Value {
780 json!({
781 "type": "object",
782 "properties": {
783 "bad:key/that/is/far/too/long/for/anthropic/property/names/and/keeps/going": {"type": "string"},
784 "nested:obj": {
785 "type": "object",
786 "properties": {"inner/key": {"type": "string"}},
787 "required": ["inner/key"]
788 }
789 },
790 "required": [
791 "bad:key/that/is/far/too/long/for/anthropic/property/names/and/keeps/going",
792 "nested:obj"
793 ]
794 })
795 }
796 async fn execute(&self, _params: Value, _ctx: ToolContext) -> Result<String> {
797 Ok("ok".to_string())
798 }
799 }
800
801 #[test]
802 fn tool_schema_uses_api_safe_names_and_maps_back() {
803 let registry = ToolRegistry::from_tools(vec![Arc::new(NamedTool("plugin:skill.tool"))]);
804
805 assert_eq!(registry.tools_schema()[0]["name"], "plugin_skill_tool");
806 assert!(registry.get("plugin:skill.tool").is_some());
807 assert!(registry.get("plugin_skill_tool").is_some());
808 assert_eq!(
809 registry.runtime_name_for_api("plugin_skill_tool"),
810 "plugin:skill.tool"
811 );
812 }
813
814 #[test]
824 fn extension_leaf_name_never_shadows_builtin_memory_fetch_api_name() {
825 struct ExtMemoryTool {
826 runtime: &'static str,
827 ext: &'static str,
828 }
829 #[async_trait::async_trait]
830 impl Tool for ExtMemoryTool {
831 fn origin(&self) -> crate::tools::ToolOrigin {
832 crate::tools::ToolOrigin::Extension {
833 extension_id: self.ext.to_string(),
834 }
835 }
836 fn extension_id(&self) -> Option<&str> {
837 Some(self.ext)
838 }
839 fn name(&self) -> &str {
840 self.runtime
841 }
842 fn description(&self) -> &str {
843 "extension memory tool (single-id schema)"
844 }
845 fn parameters(&self) -> Value {
846 json!({
847 "type": "object",
848 "properties": {"id": {"type": "string"}},
849 "required": ["id"]
850 })
851 }
852 async fn execute(&self, _params: Value, _ctx: ToolContext) -> Result<String> {
853 Ok("ok".to_string())
854 }
855 }
856
857 let registry = ToolRegistry::from_tools(vec![
858 Arc::new(crate::tools::memory::MemoryFetchTool),
859 Arc::new(ExtMemoryTool {
861 runtime: "memory:fetch",
862 ext: "memory",
863 }),
864 Arc::new(ExtMemoryTool {
866 runtime: "axel-memory-manager:memory_fetch",
867 ext: "axel-memory-manager",
868 }),
869 ]);
870
871 assert_eq!(
873 registry.runtime_name_for_api("memory_fetch"),
874 "memory_fetch",
875 "wire name `memory_fetch` must resolve to the built-in"
876 );
877 let schema = registry.tools_schema();
878 let builtin = schema
879 .iter()
880 .find(|e| e["name"] == "memory_fetch")
881 .expect("built-in memory_fetch exposed");
882 assert_eq!(
883 builtin["input_schema"]["required"],
884 json!(["ids"]),
885 "built-in memory_fetch schema must remain {{ids: [...]}}"
886 );
887
888 let axel = schema
890 .iter()
891 .find(|e| e["name"] == "axel-memory-manager_memory_fetch")
892 .expect("extension tool exposed under namespaced API name");
893 assert_eq!(axel["input_schema"]["required"], json!(["id"]));
894 assert_eq!(
897 registry.runtime_name_for_api("memory_fetch_2"),
898 "memory:fetch"
899 );
900 let resolved = registry.resolve_wire_call("memory_fetch").unwrap();
903 assert_eq!(resolved.runtime_name(), "memory_fetch");
904 }
905
906 #[test]
907 fn tool_schema_disambiguates_sanitized_name_collisions() {
908 let registry =
909 ToolRegistry::from_tools(vec![Arc::new(NamedTool("a:b")), Arc::new(NamedTool("a.b"))]);
910 let names: HashSet<String> = registry
911 .tools_schema()
912 .iter()
913 .filter_map(|s| s["name"].as_str().map(str::to_string))
914 .collect();
915
916 assert_eq!(names.len(), 2);
917 assert!(names.contains("a_b"));
918 assert!(names.contains("a_b_2"));
919 assert!(registry.get("a_b").is_some());
920 assert!(registry.get("a_b_2").is_some());
921 }
922
923 #[test]
924 fn tool_schema_truncates_long_names_to_anthropic_limit() {
925 let long = "x".repeat(140);
926 let leaked: &'static str = Box::leak(long.into_boxed_str());
927 let registry = ToolRegistry::from_tools(vec![Arc::new(NamedTool(leaked))]);
928 let schema = registry.tools_schema();
929 let name = schema[0]["name"].as_str().unwrap();
930
931 assert_eq!(name.len(), 128);
932 assert!(name
933 .chars()
934 .all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-'));
935 assert!(registry.get(name).is_some());
936 }
937
938 #[test]
939 fn tool_schema_sanitizes_input_property_names_and_translates_inputs_back() {
940 let registry = ToolRegistry::from_tools(vec![Arc::new(SchemaTool)]);
941 let schema = registry.tools_schema();
942 let input_schema = &schema[0]["input_schema"];
943 let props = input_schema["properties"].as_object().unwrap();
944
945 assert!(props.keys().all(|k| k.len() <= 64
946 && k.chars()
947 .all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-' || c == '.')));
948 assert_eq!(
949 input_schema["required"].as_array().unwrap()[0]
950 .as_str()
951 .unwrap()
952 .len(),
953 64
954 );
955 assert_eq!(input_schema["required"][1], "nested_obj");
956 assert!(props["nested_obj"]["properties"]
957 .as_object()
958 .unwrap()
959 .contains_key("inner_key"));
960 assert_eq!(props["nested_obj"]["required"][0], "inner_key");
961
962 let first_required = input_schema["required"][0].as_str().unwrap();
963 let translated = registry.translate_input_for_api_tool(
964 "schema_tool",
965 json!({
966 first_required: "value",
967 "nested_obj": {"inner_key": "nested"}
968 }),
969 );
970
971 assert_eq!(
972 translated["bad:key/that/is/far/too/long/for/anthropic/property/names/and/keeps/going"],
973 "value"
974 );
975 assert_eq!(translated["nested:obj"]["inner/key"], "nested");
976 }
977
978 #[test]
979 fn test_tool_registry_new() {
980 let registry = ToolRegistry::new();
981
982 assert_eq!(registry.tools_schema().len(), 25);
988
989 assert!(registry.get("bash").is_some());
991
992 assert!(registry.get("nonexistent").is_none());
994
995 assert!(registry.get("bash").is_some());
997 assert!(registry.get("read").is_some());
998 assert!(registry.get("write").is_some());
999 assert!(registry.get("edit").is_some());
1000 assert!(registry.get("grep").is_some());
1001 assert!(registry.get("find").is_some());
1002 assert!(registry.get("ls").is_some());
1003 assert!(registry.get("subagent").is_some());
1004 assert!(registry.get("subagent_model_authorize").is_some());
1005 assert!(registry.get("subagent_models").is_some());
1006 }
1007
1008 #[test]
1009 fn test_disable_removes_tools_from_registry_and_schema() {
1010 let mut registry = ToolRegistry::new();
1011 let before = registry.tools_schema().len();
1012
1013 registry.disable(&["bash".to_string(), "ls".to_string()]);
1014
1015 assert!(registry.get("bash").is_none());
1017 assert!(registry.get("ls").is_none());
1018 assert!(registry.get("read").is_some());
1020 assert!(registry.get("subagent").is_some());
1021 assert_eq!(registry.tools_schema().len(), before - 2);
1023 }
1024
1025 #[test]
1026 fn test_disable_empty_and_unknown_are_noops() {
1027 let mut registry = ToolRegistry::new();
1028 let n = registry.tools_schema().len();
1029 registry.disable(&[]); assert_eq!(registry.tools_schema().len(), n);
1031 registry.disable(&["not_a_real_tool".to_string()]); assert_eq!(registry.tools_schema().len(), n);
1033 assert!(registry.get("bash").is_some());
1034 }
1035
1036 #[test]
1037 fn test_tool_registry_without_subagent() {
1038 let registry = ToolRegistry::without_subagent();
1039
1040 assert_eq!(registry.tools_schema().len(), 10);
1042 assert!(registry.get("subagent").is_none());
1044
1045 assert!(registry.get("bash").is_some());
1047
1048 assert!(registry.get("bash").is_some());
1050 assert!(registry.get("read").is_some());
1051 assert!(registry.get("write").is_some());
1052 assert!(registry.get("edit").is_some());
1053 assert!(registry.get("grep").is_some());
1054 assert!(registry.get("find").is_some());
1055 assert!(registry.get("ls").is_some());
1056 }
1057
1058 #[test]
1059 fn test_tool_registry_register() {
1060 let mut registry = ToolRegistry::without_subagent();
1061 let initial_count = registry.tools_schema().len();
1062
1063 struct TestTool;
1065 #[async_trait::async_trait]
1066 impl Tool for TestTool {
1067 fn name(&self) -> &str {
1068 "test_tool"
1069 }
1070 fn description(&self) -> &str {
1071 "A test tool"
1072 }
1073 fn parameters(&self) -> Value {
1074 json!({"type": "object"})
1075 }
1076 async fn execute(&self, _params: Value, _ctx: ToolContext) -> Result<String> {
1077 Ok("test result".to_string())
1078 }
1079 }
1080
1081 registry.register(Arc::new(TestTool));
1082
1083 assert_eq!(registry.tools_schema().len(), initial_count + 1);
1085
1086 assert!(registry.get("test_tool").is_some());
1088 }
1089
1090 #[test]
1091 fn tool_names_for_extension_filters_by_owner_and_sorts() {
1092 struct OwnedTool(&'static str, Option<&'static str>);
1093 #[async_trait::async_trait]
1094 impl Tool for OwnedTool {
1095 fn name(&self) -> &str {
1096 self.0
1097 }
1098 fn description(&self) -> &str {
1099 "owned"
1100 }
1101 fn parameters(&self) -> Value {
1102 json!({"type": "object"})
1103 }
1104 async fn execute(&self, _params: Value, _ctx: ToolContext) -> Result<String> {
1105 Ok("ok".to_string())
1106 }
1107 fn extension_id(&self) -> Option<&str> {
1108 self.1
1109 }
1110 }
1111
1112 let mut registry = ToolRegistry::without_subagent();
1113 registry.register(Arc::new(OwnedTool("alpha:zed", Some("alpha"))));
1114 registry.register(Arc::new(OwnedTool("alpha:bar", Some("alpha"))));
1115 registry.register(Arc::new(OwnedTool("beta:thing", Some("beta"))));
1116
1117 assert_eq!(
1118 registry.tool_names_for_extension("alpha"),
1119 vec!["alpha:bar".to_string(), "alpha:zed".to_string()]
1120 );
1121 assert_eq!(
1122 registry.tool_names_for_extension("beta"),
1123 vec!["beta:thing".to_string()]
1124 );
1125 assert!(registry.tool_names_for_extension("ghost").is_empty());
1126 assert!(registry.tool_names_for_extension("bash").is_empty());
1128 }
1129
1130 struct OwnedTool(&'static str, Option<&'static str>);
1131 #[async_trait::async_trait]
1132 impl Tool for OwnedTool {
1133 fn name(&self) -> &str {
1134 self.0
1135 }
1136 fn description(&self) -> &str {
1137 "owned"
1138 }
1139 fn parameters(&self) -> Value {
1140 json!({"type": "object"})
1141 }
1142 async fn execute(&self, _params: Value, _ctx: ToolContext) -> Result<String> {
1143 Ok("ok".to_string())
1144 }
1145 fn extension_id(&self) -> Option<&str> {
1146 self.1
1147 }
1148 }
1149
1150 #[test]
1151 fn without_subagent_excludes_subagent_tools() {
1152 let registry = ToolRegistry::without_subagent();
1153 assert!(registry.get("subagent").is_none());
1154 assert!(registry.get("subagent_start").is_none());
1155 assert!(registry.get("subagent_status").is_none());
1156 assert!(registry.get("subagent_steer").is_none());
1157 assert!(registry.get("subagent_collect").is_none());
1158 assert!(registry.get("subagent_resume").is_none());
1159 assert!(registry.get("bash").is_some());
1161 assert!(registry.get("read").is_some());
1162 }
1163
1164 #[test]
1165 fn without_subagent_with_extensions_includes_extension_tools() {
1166 let mut other = ToolRegistry::empty();
1167 other.register(Arc::new(OwnedTool("alpha:do_thing", Some("alpha"))));
1168
1169 let merged = ToolRegistry::without_subagent_with_extensions(&other);
1170
1171 assert!(merged.get("alpha:do_thing").is_some());
1173 assert!(merged.get("bash").is_some());
1175 assert!(merged.get("read").is_some());
1176 assert!(merged.get("subagent_start").is_none());
1178 }
1179
1180 #[test]
1181 fn without_subagent_with_extensions_excludes_built_ins_from_other_registry() {
1182 let other = ToolRegistry::new();
1187
1188 let merged = ToolRegistry::without_subagent_with_extensions(&other);
1189
1190 let bash = merged.get("bash").expect("bash present");
1192 assert!(bash.extension_id().is_none());
1193 assert!(merged.get("subagent_start").is_none());
1195 assert!(merged.get("subagent").is_none());
1196 }
1197
1198 #[test]
1199 fn without_subagent_with_extensions_rejects_recursive_tool_name_collisions() {
1200 let mut other = ToolRegistry::empty();
1201 for name in [
1202 "subagent",
1203 "subagent_start",
1204 "subagent_status",
1205 "subagent_steer",
1206 "subagent_collect",
1207 "subagent_resume",
1208 "subagent_model_authorize",
1209 "subagent_models",
1210 "memory_context",
1211 ] {
1212 other.register(Arc::new(OwnedTool(name, Some("malicious-extension"))));
1213 }
1214
1215 let merged = ToolRegistry::without_subagent_with_extensions(&other);
1216 for name in [
1217 "subagent",
1218 "subagent_start",
1219 "subagent_status",
1220 "subagent_steer",
1221 "subagent_collect",
1222 "subagent_resume",
1223 "subagent_model_authorize",
1224 "subagent_models",
1225 "memory_context",
1226 ] {
1227 assert!(merged.get(name).is_none(), "{name} must stay unavailable");
1228 }
1229 }
1230
1231 #[test]
1232 fn without_subagent_with_extensions_rejects_api_sanitized_recursive_collisions() {
1233 let mut other = ToolRegistry::empty();
1234 for runtime_name in [
1235 "subagent:start",
1236 "subagent:status",
1237 "subagent:steer",
1238 "subagent:collect",
1239 "subagent:resume",
1240 "subagent:model:authorize",
1241 "subagent:models",
1242 "memory:context",
1243 ] {
1244 other.register(Arc::new(OwnedTool(runtime_name, Some("subagent"))));
1245 }
1246
1247 let merged = ToolRegistry::without_subagent_with_extensions(&other);
1248 for api_name in [
1249 "subagent_start",
1250 "subagent_status",
1251 "subagent_steer",
1252 "subagent_collect",
1253 "subagent_resume",
1254 "subagent_model_authorize",
1255 "subagent_models",
1256 "memory_context",
1257 ] {
1258 assert!(
1259 merged.get(api_name).is_none(),
1260 "API-safe recursive name {api_name} must stay unavailable"
1261 );
1262 }
1263 }
1264
1265 #[test]
1269 fn memory_context_present_in_primary_registry_absent_in_recursive_ones() {
1270 assert!(ToolRegistry::new().get("memory_context").is_some());
1271 assert!(ToolRegistry::without_subagent()
1272 .get("memory_context")
1273 .is_none());
1274
1275 let merged = ToolRegistry::without_subagent_with_extensions(&ToolRegistry::new());
1276 assert!(
1277 merged.get("memory_context").is_none(),
1278 "recursive registries must never expose memory_context"
1279 );
1280 }
1281
1282 #[test]
1283 fn without_subagent_with_extensions_does_not_overwrite_existing_builtin() {
1284 let mut other = ToolRegistry::empty();
1290 other.register(Arc::new(OwnedTool("ext:custom", Some("ext"))));
1291
1292 let merged = ToolRegistry::without_subagent_with_extensions(&other);
1293 assert!(merged.get("ext:custom").is_some());
1294 assert!(merged.get("bash").is_some());
1295 assert!(merged.get("bash").unwrap().extension_id().is_none());
1296 }
1297
1298 struct McpTwin(&'static str);
1300
1301 #[async_trait::async_trait]
1302 impl Tool for McpTwin {
1303 fn name(&self) -> &str {
1304 self.0
1305 }
1306 fn description(&self) -> &str {
1307 "mcp twin"
1308 }
1309 fn parameters(&self) -> Value {
1310 json!({"type": "object"})
1311 }
1312 async fn execute(&self, _params: Value, _ctx: ToolContext) -> Result<String> {
1313 Ok("ok".to_string())
1314 }
1315 fn origin(&self) -> crate::tools::ToolOrigin {
1316 crate::tools::ToolOrigin::Mcp {
1317 server_id: "srv".to_string(),
1318 server_tool_name: "shared".to_string(),
1319 }
1320 }
1321 }
1322
1323 #[test]
1324 #[should_panic(expected = "tool catalog rejected mutation")]
1325 fn from_tools_with_duplicate_capability_identity_fails_closed() {
1326 let _ = ToolRegistry::from_tools(vec![
1330 Arc::new(McpTwin("twin_one")),
1331 Arc::new(McpTwin("twin_two")),
1332 ]);
1333 }
1334}