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 Arc::new(crate::tools::forum::ForumPostTool),
83 Arc::new(crate::tools::forum::ForumReadTool),
84 Arc::new(crate::tools::forum::ForumForgetTool),
85 ];
86 Self::from_tools(tools)
87 }
88
89 pub fn empty() -> Self {
92 Self::from_tools(Vec::new())
93 }
94
95 pub fn disable(&mut self, names: &[String]) {
101 self.try_disable(names).expect(CATALOG_REJECTED_MSG);
102 }
103
104 pub fn try_disable(&mut self, names: &[String]) -> Result<(), CatalogError> {
117 if names.is_empty() {
118 return Ok(());
119 }
120 let remaining: Vec<Arc<dyn Tool>> = self
121 .tools
122 .values()
123 .filter(|t| !names.iter().any(|n| n == t.name()))
124 .cloned()
125 .collect();
126 if remaining.len() == self.tools.len() {
127 return Ok(());
128 }
129 let prior = self.catalog.generation();
130 let mut candidate = Self::try_from_tools(remaining)?;
131 candidate.catalog.rebase_past(prior)?;
132 *self = candidate;
133 Ok(())
134 }
135
136 pub fn without_subagent() -> Self {
138 let tools: Vec<Arc<dyn Tool>> = vec![
139 Arc::new(crate::tools::bash::BashTool),
140 #[cfg(windows)]
141 Arc::new(crate::tools::powershell::PowerShellTool),
142 Arc::new(crate::tools::read::ReadTool),
143 Arc::new(crate::tools::write::WriteTool),
144 Arc::new(crate::tools::edit::EditTool),
145 Arc::new(crate::tools::grep::GrepTool),
146 Arc::new(crate::tools::find::FindTool),
147 Arc::new(crate::tools::ls::LsTool),
148 Arc::new(crate::tools::shell::ShellStartTool),
149 Arc::new(crate::tools::shell::ShellSendTool),
150 Arc::new(crate::tools::shell::ShellEndTool),
151 Arc::new(crate::tools::forum::ForumPostTool),
152 Arc::new(crate::tools::forum::ForumReadTool),
153 Arc::new(crate::tools::forum::ForumForgetTool),
154 ];
155 Self::from_tools(tools)
156 }
157
158 pub fn without_subagent_with_extensions(extension_tools: &ToolRegistry) -> Self {
166 let mut combined = Self::without_subagent();
167 for tool in extension_tools.tools.values() {
168 if tool.extension_id().is_some() && !is_recursive_subagent_tool_name(tool.name()) {
169 combined
170 .insert_tool_with_catalog(tool.clone())
171 .expect(CATALOG_REJECTED_MSG);
172 }
173 }
174 combined.rebuild_schema();
175 combined
176 }
177
178 fn from_tools(tool_list: Vec<Arc<dyn Tool>>) -> Self {
184 Self::try_from_tools(tool_list).expect(CATALOG_REJECTED_MSG)
185 }
186
187 #[doc(hidden)]
191 pub fn from_tools_for_tests(tool_list: Vec<Arc<dyn Tool>>) -> Self {
192 Self::from_tools(tool_list)
193 }
194
195 fn try_from_tools(tool_list: Vec<Arc<dyn Tool>>) -> Result<Self, CatalogError> {
196 let mut registry = ToolRegistry {
197 tools: HashMap::new(),
198 cached_schema: Arc::new(Vec::new()),
199 api_to_runtime_names: HashMap::new(),
200 input_name_maps: HashMap::new(),
201 catalog: ToolCatalog::empty(),
202 };
203 for tool in tool_list {
207 let name = tool.name().to_string();
208 registry.tools.insert(name, tool);
209 }
210 let sorted: Vec<Arc<dyn Tool>> =
216 registry.iter_tools_sorted().into_iter().cloned().collect();
217 for tool in sorted {
218 let record = CapabilityRecord::for_registered_tool(&tool);
219 registry.catalog.upsert(None, record)?;
220 }
221 registry.rebuild_schema();
222 Ok(registry)
223 }
224
225 fn api_safe_name(name: &str, used: &HashSet<String>) -> String {
226 Self::api_safe_identifier(name, used, 128, false)
227 }
228
229 fn api_safe_property_name(name: &str, used: &HashSet<String>) -> String {
230 Self::api_safe_identifier(name, used, 64, true)
231 }
232
233 fn api_safe_identifier(
234 name: &str,
235 used: &HashSet<String>,
236 max_len: usize,
237 allow_dot: bool,
238 ) -> String {
239 let mut sanitized = String::with_capacity(name.len());
240 for ch in name.chars() {
241 if ch.is_ascii_alphanumeric() || ch == '_' || ch == '-' || (allow_dot && ch == '.') {
242 sanitized.push(ch);
243 } else {
244 sanitized.push('_');
245 }
246 }
247 if sanitized.is_empty() {
248 sanitized.push_str("field");
249 }
250 if sanitized.len() > max_len {
251 sanitized.truncate(max_len);
252 }
253
254 let base = sanitized.clone();
255 let mut suffix = 2;
256 while used.contains(&sanitized) {
257 let suffix_str = format!("_{suffix}");
258 let keep = max_len.saturating_sub(suffix_str.len());
259 sanitized = format!("{}{}", &base[..base.len().min(keep)], suffix_str);
260 suffix += 1;
261 }
262 sanitized
263 }
264
265 fn sanitize_schema(mut schema: Value) -> (Value, SchemaNameMap) {
266 let mut map = SchemaNameMap::default();
267 let Some(obj) = schema.as_object_mut() else {
268 return (schema, map);
269 };
270
271 let mut required_name_map = HashMap::new();
272 if let Some(props_value) = obj.get_mut("properties") {
273 if let Some(props) = props_value.as_object_mut() {
274 let original = std::mem::take(props);
275 let mut used = HashSet::new();
276 for (runtime_name, child_schema) in original {
277 let api_name = Self::api_safe_property_name(&runtime_name, &used);
278 used.insert(api_name.clone());
279 required_name_map.insert(runtime_name.clone(), api_name.clone());
280 map.api_to_runtime.insert(api_name.clone(), runtime_name);
281
282 let (sanitized_child, child_map) = Self::sanitize_schema(child_schema);
283 if !child_map.api_to_runtime.is_empty() || !child_map.children.is_empty() {
284 map.children.insert(api_name.clone(), child_map);
285 }
286 props.insert(api_name, sanitized_child);
287 }
288 }
289 }
290
291 if let Some(required) = obj.get_mut("required").and_then(Value::as_array_mut) {
292 for item in required.iter_mut() {
293 if let Some(name) = item.as_str() {
294 if let Some(api_name) = required_name_map.get(name) {
295 *item = Value::String(api_name.clone());
296 }
297 }
298 }
299 }
300
301 if let Some(items) = obj.get_mut("items") {
304 let (sanitized_items, items_map) = Self::sanitize_schema(std::mem::take(items));
305 if !items_map.api_to_runtime.is_empty()
306 || !items_map.children.is_empty()
307 || items_map.items.is_some()
308 {
309 map.items = Some(Box::new(items_map));
310 }
311 *items = sanitized_items;
312 }
313
314 (schema, map)
315 }
316
317 fn translate_input_names(input: Value, map: &SchemaNameMap) -> Value {
318 match input {
319 Value::Object(obj) => {
320 let mut out = serde_json::Map::new();
321 for (api_name, value) in obj {
322 let runtime_name = map
323 .api_to_runtime
324 .get(&api_name)
325 .cloned()
326 .unwrap_or_else(|| api_name.clone());
327 let value = if let Some(child) = map.children.get(&api_name) {
328 Self::translate_input_names(value, child)
329 } else {
330 value
331 };
332 out.insert(runtime_name, value);
333 }
334 Value::Object(out)
335 }
336 Value::Array(arr) => {
337 if let Some(items_map) = &map.items {
339 Value::Array(
340 arr.into_iter()
341 .map(|v| Self::translate_input_names(v, items_map))
342 .collect(),
343 )
344 } else {
345 Value::Array(arr)
346 }
347 }
348 other => other,
349 }
350 }
351
352 fn rebuild_schema(&mut self) {
353 let mut used = HashSet::new();
354 let mut api_to_runtime_names = HashMap::new();
355 let mut input_name_maps = HashMap::new();
356 let mut schema = Vec::with_capacity(self.tools.len());
357
358 let mut sorted_tools: Vec<_> = self.tools.values().collect();
373 sorted_tools.sort_by_key(|t| {
374 let builtin_first = match t.origin() {
375 crate::tools::ToolOrigin::Builtin => 0u8,
376 _ => 1u8,
377 };
378 (builtin_first, t.name().to_string())
379 });
380
381 for tool in sorted_tools {
382 let runtime_name = tool.name();
383 let api_name = Self::api_safe_name(runtime_name, &used);
384 used.insert(api_name.clone());
385 api_to_runtime_names.insert(api_name.clone(), runtime_name.to_string());
386 let (input_schema, input_map) = Self::sanitize_schema(tool.parameters());
387 input_name_maps.insert(api_name.clone(), input_map);
388 schema.push(serde_json::json!({
389 "name": api_name,
390 "description": tool.description(),
391 "input_schema": input_schema
392 }));
393 }
394
395 self.api_to_runtime_names = api_to_runtime_names;
396 self.input_name_maps = input_name_maps;
397 self.cached_schema = Arc::new(schema);
398 }
399
400 pub fn register(&mut self, tool: Arc<dyn Tool>) {
410 self.try_register(tool).expect(CATALOG_REJECTED_MSG);
411 }
412
413 pub fn try_register(&mut self, tool: Arc<dyn Tool>) -> Result<(), CatalogError> {
419 self.insert_tool_with_catalog(tool)?;
420 self.rebuild_schema();
421 Ok(())
422 }
423
424 pub fn try_register_batch(
438 &mut self,
439 tool_list: Vec<Arc<dyn Tool>>,
440 ) -> Result<usize, CatalogError> {
441 if tool_list.is_empty() {
442 return Ok(0);
443 }
444 let mut batch_names: HashSet<String> = HashSet::new();
445 let mut batch_ids: HashSet<crate::tools::catalog::ToolId> = HashSet::new();
446 for tool in &tool_list {
447 let name = tool.name();
448 if self.tools.contains_key(name) || !batch_names.insert(name.to_string()) {
449 return Err(CatalogError::DuplicateRuntimeName(
450 agent_core::BoundedText::new(name, 128).text,
451 ));
452 }
453 let id = tool_id_for(tool.as_ref());
454 if self.catalog.get(&id).is_some() || !batch_ids.insert(id.clone()) {
455 return Err(CatalogError::DuplicateToolId(id));
456 }
457 }
458 let count = tool_list.len();
459 let mut candidate = self.clone();
460 for tool in tool_list {
461 candidate.insert_tool_with_catalog(tool)?;
462 }
463 candidate.rebuild_schema();
464 *self = candidate;
465 Ok(count)
466 }
467
468 fn insert_tool_with_catalog(&mut self, tool: Arc<dyn Tool>) -> Result<(), CatalogError> {
477 let record = CapabilityRecord::for_registered_tool(&tool);
478 let replaced = self
479 .tools
480 .get(tool.name())
481 .map(|old| tool_id_for(old.as_ref()));
482 self.catalog.upsert(replaced.as_ref(), record)?;
483 self.tools.insert(tool.name().to_string(), tool);
484 Ok(())
485 }
486
487 pub fn catalog(&self) -> &ToolCatalog {
490 &self.catalog
491 }
492
493 #[doc(hidden)]
501 pub fn resume_catalog_generation_for_tests(
502 &mut self,
503 generation: crate::tools::catalog::CatalogGeneration,
504 ) {
505 self.catalog.set_generation_for_tests(generation);
506 }
507
508 #[doc(hidden)]
514 pub fn upsert_catalog_record_for_tests(
515 &mut self,
516 replaced: Option<&crate::tools::catalog::ToolId>,
517 record: crate::tools::catalog::CapabilityRecord,
518 ) {
519 self.catalog
520 .upsert(replaced, record)
521 .expect("test upsert must be structurally valid");
522 }
523
524 pub fn get(&self, name: &str) -> Option<&Arc<dyn Tool>> {
525 let runtime_name = self
526 .api_to_runtime_names
527 .get(name)
528 .map(String::as_str)
529 .unwrap_or(name);
530 self.tools.get(runtime_name)
531 }
532
533 pub fn runtime_name_for_api<'a>(&'a self, name: &'a str) -> &'a str {
534 self.api_to_runtime_names
535 .get(name)
536 .map(String::as_str)
537 .unwrap_or(name)
538 }
539
540 pub fn resolve_wire_call(
550 &self,
551 wire_name: &str,
552 ) -> Option<crate::tools::activation::ResolvedToolCall> {
553 let runtime_name = self
554 .api_to_runtime_names
555 .get(wire_name)
556 .map(String::as_str)
557 .unwrap_or(wire_name);
558 let tool = self.tools.get(runtime_name)?;
559 Some(crate::tools::activation::ResolvedToolCall::new(
560 wire_name,
561 runtime_name,
562 tool_id_for(tool.as_ref()),
563 ))
564 }
565
566 pub fn translate_input_for_api_tool(&self, tool_name: &str, input: Value) -> Value {
567 if let Some(map) = self.input_name_maps.get(tool_name) {
568 Self::translate_input_names(input, map)
569 } else {
570 input
571 }
572 }
573
574 pub fn tools_schema(&self) -> Arc<Vec<Value>> {
575 Arc::clone(&self.cached_schema)
576 }
577
578 pub fn deferred_tool_split(
614 &self,
615 session: &crate::tools::activation::SessionToolSet,
616 referenced_ids: &HashSet<String>,
617 ) -> DeferredToolSplit {
618 let projection = self.session_tools_schema(session).schema;
619 let projected: HashSet<&str> = projection
620 .iter()
621 .filter_map(|entry| entry["name"].as_str())
622 .collect();
623 let mut split = DeferredToolSplit::default();
624 for entry in self.cached_schema.iter() {
625 let Some(api_name) = entry["name"].as_str() else {
626 continue;
627 };
628 let Some(tool) = self.tools.get(self.runtime_name_for_api(api_name)) else {
629 continue;
630 };
631 let id = tool_id_for(tool.as_ref());
632 if projected.contains(api_name) {
633 if session.is_core(&id) || !referenced_ids.contains(id.as_str()) {
634 split.loaded.push(entry.clone());
635 } else {
636 split.deferred.push(deferred_entry(entry));
637 }
638 } else if self
639 .catalog
640 .get(&id)
641 .is_some_and(crate::tools::activation::is_activatable_record)
642 {
643 split.deferred.push(deferred_entry(entry));
644 } else {
645 continue;
646 }
647 split
648 .id_to_api
649 .insert(id.as_str().to_string(), api_name.to_string());
650 }
651 split
652 }
653
654 pub fn session_tools_schema(
655 &self,
656 session: &crate::tools::activation::SessionToolSet,
657 ) -> SessionSchemaProjection {
658 let mut projected = Vec::new();
659 let mut dropped = Vec::new();
660 for entry in self.cached_schema.iter() {
661 let Some(api_name) = entry["name"].as_str() else {
662 continue;
663 };
664 let runtime_name = self.runtime_name_for_api(api_name);
665 let Some(tool) = self.tools.get(runtime_name) else {
666 continue;
667 };
668 let id = tool_id_for(tool.as_ref());
669 let pinned = session
670 .core_pin(&id)
671 .map(|pin| (pin.schema_digest(), pin.provenance()))
672 .or_else(|| {
673 session
674 .activation(&id)
675 .map(|a| (a.schema_digest(), a.provenance()))
676 });
677 let Some((pinned_digest, pinned_provenance)) = pinned else {
678 continue; };
680 let Some(record) = self.catalog.get(&id) else {
681 tracing::warn!(
682 tool = %id,
683 "session schema projection drops uncataloged session member \
684 (removed since the session pinned it)"
685 );
686 dropped.push((id, DroppedSessionMember::Uncataloged));
687 continue;
688 };
689 if pinned_digest != record.schema_digest() {
690 tracing::warn!(
691 tool = %id,
692 "session schema projection drops drifted session member \
693 (schema digest changed since the session pinned it)"
694 );
695 dropped.push((id, DroppedSessionMember::SchemaDigestMismatch));
696 continue;
697 }
698 if pinned_provenance != record.provenance() {
699 tracing::warn!(
700 tool = %id,
701 "session schema projection drops drifted session member \
702 (trust provenance changed since the session pinned it)"
703 );
704 dropped.push((id, DroppedSessionMember::ProvenanceMismatch));
705 continue;
706 }
707 projected.push(entry.clone());
708 }
709 let seen: HashSet<&crate::tools::catalog::ToolId> =
716 dropped.iter().map(|(id, _)| id).collect();
717 let vanished: Vec<crate::tools::catalog::ToolId> = session
718 .core_ids()
719 .chain(session.activated().map(|a| a.grant().tool_id()))
720 .filter(|id| !seen.contains(id) && self.catalog.get(id).is_none())
721 .cloned()
722 .collect();
723 for id in vanished {
724 tracing::warn!(
725 tool = %id,
726 "session schema projection drops uncataloged session member (removed from the registry since the session pinned it)"
727 );
728 dropped.push((id, DroppedSessionMember::Uncataloged));
729 }
730 SessionSchemaProjection {
731 schema: projected,
732 dropped,
733 }
734 }
735
736 pub fn tool_names_for_extension(&self, extension_id: &str) -> Vec<String> {
739 let mut names: Vec<String> = self
740 .tools
741 .values()
742 .filter(|t| t.extension_id() == Some(extension_id))
743 .map(|t| t.name().to_string())
744 .collect();
745 names.sort();
746 names
747 }
748
749 pub fn iter_tools_sorted(&self) -> Vec<&Arc<dyn Tool>> {
752 let mut tools: Vec<_> = self.tools.values().collect();
753 tools.sort_by_key(|t| t.name());
754 tools
755 }
756}
757
758#[derive(Clone, Debug, PartialEq, Eq)]
764pub struct SessionSchemaProjection {
765 pub schema: Vec<Value>,
767 pub dropped: Vec<(crate::tools::catalog::ToolId, DroppedSessionMember)>,
771}
772
773#[derive(Clone, Copy, Debug, PartialEq, Eq)]
775pub enum DroppedSessionMember {
776 Uncataloged,
778 SchemaDigestMismatch,
780 ProvenanceMismatch,
783}
784
785const CATALOG_REJECTED_MSG: &str =
791 "tool catalog rejected mutation: refusing to expose a tool the catalog could not record";
792
793fn is_recursive_subagent_tool_name(name: &str) -> bool {
794 let api_name = ToolRegistry::api_safe_name(name, &HashSet::new());
795 matches!(
796 api_name.as_str(),
797 "subagent"
798 | "subagent_start"
799 | "subagent_status"
800 | "subagent_steer"
801 | "subagent_collect"
802 | "subagent_resume"
803 | "subagent_model_authorize"
804 | "subagent_models"
805 | "memory_context"
809 )
810}
811#[derive(Debug, Clone, Default)]
813pub struct DeferredToolSplit {
814 pub loaded: Vec<Value>,
816 pub deferred: Vec<Value>,
818 pub id_to_api: std::collections::HashMap<String, String>,
820}
821
822fn deferred_entry(entry: &Value) -> Value {
823 let mut entry = entry.clone();
824 if let Some(map) = entry.as_object_mut() {
825 map.remove("cache_control");
826 map.insert("defer_loading".to_string(), Value::Bool(true));
827 }
828 entry
829}
830
831#[cfg(test)]
832mod tests {
833 use super::*;
834 use crate::{Result, ToolContext};
835 use serde_json::json;
836
837 struct NamedTool(&'static str);
838
839 #[async_trait::async_trait]
840 impl Tool for NamedTool {
841 fn name(&self) -> &str {
842 self.0
843 }
844 fn description(&self) -> &str {
845 "test tool"
846 }
847 fn parameters(&self) -> Value {
848 json!({"type": "object"})
849 }
850 async fn execute(&self, _params: Value, _ctx: ToolContext) -> Result<String> {
851 Ok("ok".to_string())
852 }
853 }
854
855 struct SchemaTool;
856
857 #[async_trait::async_trait]
858 impl Tool for SchemaTool {
859 fn name(&self) -> &str {
860 "schema_tool"
861 }
862 fn description(&self) -> &str {
863 "schema tool"
864 }
865 fn parameters(&self) -> Value {
866 json!({
867 "type": "object",
868 "properties": {
869 "bad:key/that/is/far/too/long/for/anthropic/property/names/and/keeps/going": {"type": "string"},
870 "nested:obj": {
871 "type": "object",
872 "properties": {"inner/key": {"type": "string"}},
873 "required": ["inner/key"]
874 }
875 },
876 "required": [
877 "bad:key/that/is/far/too/long/for/anthropic/property/names/and/keeps/going",
878 "nested:obj"
879 ]
880 })
881 }
882 async fn execute(&self, _params: Value, _ctx: ToolContext) -> Result<String> {
883 Ok("ok".to_string())
884 }
885 }
886
887 #[test]
888 fn tool_schema_uses_api_safe_names_and_maps_back() {
889 let registry = ToolRegistry::from_tools(vec![Arc::new(NamedTool("plugin:skill.tool"))]);
890
891 assert_eq!(registry.tools_schema()[0]["name"], "plugin_skill_tool");
892 assert!(registry.get("plugin:skill.tool").is_some());
893 assert!(registry.get("plugin_skill_tool").is_some());
894 assert_eq!(
895 registry.runtime_name_for_api("plugin_skill_tool"),
896 "plugin:skill.tool"
897 );
898 }
899
900 #[test]
910 fn extension_leaf_name_never_shadows_builtin_memory_fetch_api_name() {
911 struct ExtMemoryTool {
912 runtime: &'static str,
913 ext: &'static str,
914 }
915 #[async_trait::async_trait]
916 impl Tool for ExtMemoryTool {
917 fn origin(&self) -> crate::tools::ToolOrigin {
918 crate::tools::ToolOrigin::Extension {
919 extension_id: self.ext.to_string(),
920 }
921 }
922 fn extension_id(&self) -> Option<&str> {
923 Some(self.ext)
924 }
925 fn name(&self) -> &str {
926 self.runtime
927 }
928 fn description(&self) -> &str {
929 "extension memory tool (single-id schema)"
930 }
931 fn parameters(&self) -> Value {
932 json!({
933 "type": "object",
934 "properties": {"id": {"type": "string"}},
935 "required": ["id"]
936 })
937 }
938 async fn execute(&self, _params: Value, _ctx: ToolContext) -> Result<String> {
939 Ok("ok".to_string())
940 }
941 }
942
943 let registry = ToolRegistry::from_tools(vec![
944 Arc::new(crate::tools::memory::MemoryFetchTool),
945 Arc::new(ExtMemoryTool {
947 runtime: "memory:fetch",
948 ext: "memory",
949 }),
950 Arc::new(ExtMemoryTool {
952 runtime: "axel-memory-manager:memory_fetch",
953 ext: "axel-memory-manager",
954 }),
955 ]);
956
957 assert_eq!(
959 registry.runtime_name_for_api("memory_fetch"),
960 "memory_fetch",
961 "wire name `memory_fetch` must resolve to the built-in"
962 );
963 let schema = registry.tools_schema();
964 let builtin = schema
965 .iter()
966 .find(|e| e["name"] == "memory_fetch")
967 .expect("built-in memory_fetch exposed");
968 assert_eq!(
969 builtin["input_schema"]["required"],
970 json!(["ids"]),
971 "built-in memory_fetch schema must remain {{ids: [...]}}"
972 );
973
974 let axel = schema
976 .iter()
977 .find(|e| e["name"] == "axel-memory-manager_memory_fetch")
978 .expect("extension tool exposed under namespaced API name");
979 assert_eq!(axel["input_schema"]["required"], json!(["id"]));
980 assert_eq!(
983 registry.runtime_name_for_api("memory_fetch_2"),
984 "memory:fetch"
985 );
986 let resolved = registry.resolve_wire_call("memory_fetch").unwrap();
989 assert_eq!(resolved.runtime_name(), "memory_fetch");
990 }
991
992 #[test]
993 fn tool_schema_disambiguates_sanitized_name_collisions() {
994 let registry =
995 ToolRegistry::from_tools(vec![Arc::new(NamedTool("a:b")), Arc::new(NamedTool("a.b"))]);
996 let names: HashSet<String> = registry
997 .tools_schema()
998 .iter()
999 .filter_map(|s| s["name"].as_str().map(str::to_string))
1000 .collect();
1001
1002 assert_eq!(names.len(), 2);
1003 assert!(names.contains("a_b"));
1004 assert!(names.contains("a_b_2"));
1005 assert!(registry.get("a_b").is_some());
1006 assert!(registry.get("a_b_2").is_some());
1007 }
1008
1009 #[test]
1010 fn tool_schema_truncates_long_names_to_anthropic_limit() {
1011 let long = "x".repeat(140);
1012 let leaked: &'static str = Box::leak(long.into_boxed_str());
1013 let registry = ToolRegistry::from_tools(vec![Arc::new(NamedTool(leaked))]);
1014 let schema = registry.tools_schema();
1015 let name = schema[0]["name"].as_str().unwrap();
1016
1017 assert_eq!(name.len(), 128);
1018 assert!(name
1019 .chars()
1020 .all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-'));
1021 assert!(registry.get(name).is_some());
1022 }
1023
1024 #[test]
1025 fn tool_schema_sanitizes_input_property_names_and_translates_inputs_back() {
1026 let registry = ToolRegistry::from_tools(vec![Arc::new(SchemaTool)]);
1027 let schema = registry.tools_schema();
1028 let input_schema = &schema[0]["input_schema"];
1029 let props = input_schema["properties"].as_object().unwrap();
1030
1031 assert!(props.keys().all(|k| k.len() <= 64
1032 && k.chars()
1033 .all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-' || c == '.')));
1034 assert_eq!(
1035 input_schema["required"].as_array().unwrap()[0]
1036 .as_str()
1037 .unwrap()
1038 .len(),
1039 64
1040 );
1041 assert_eq!(input_schema["required"][1], "nested_obj");
1042 assert!(props["nested_obj"]["properties"]
1043 .as_object()
1044 .unwrap()
1045 .contains_key("inner_key"));
1046 assert_eq!(props["nested_obj"]["required"][0], "inner_key");
1047
1048 let first_required = input_schema["required"][0].as_str().unwrap();
1049 let translated = registry.translate_input_for_api_tool(
1050 "schema_tool",
1051 json!({
1052 first_required: "value",
1053 "nested_obj": {"inner_key": "nested"}
1054 }),
1055 );
1056
1057 assert_eq!(
1058 translated["bad:key/that/is/far/too/long/for/anthropic/property/names/and/keeps/going"],
1059 "value"
1060 );
1061 assert_eq!(translated["nested:obj"]["inner/key"], "nested");
1062 }
1063
1064 #[test]
1065 fn test_tool_registry_new() {
1066 let registry = ToolRegistry::new();
1067
1068 assert_eq!(
1072 registry.tools_schema().len(),
1073 28 + usize::from(cfg!(windows))
1074 );
1075 for name in ["forum_post", "forum_read", "forum_forget"] {
1076 assert!(registry.get(name).is_some(), "missing {name}");
1077 }
1078
1079 assert!(registry.get("bash").is_some());
1081
1082 assert!(registry.get("nonexistent").is_none());
1084
1085 assert!(registry.get("bash").is_some());
1087 assert!(registry.get("read").is_some());
1088 assert!(registry.get("write").is_some());
1089 assert!(registry.get("edit").is_some());
1090 assert!(registry.get("grep").is_some());
1091 assert!(registry.get("find").is_some());
1092 assert!(registry.get("ls").is_some());
1093 assert!(registry.get("subagent").is_some());
1094 assert!(registry.get("subagent_model_authorize").is_some());
1095 assert!(registry.get("subagent_models").is_some());
1096 }
1097
1098 #[test]
1099 fn test_disable_removes_tools_from_registry_and_schema() {
1100 let mut registry = ToolRegistry::new();
1101 let before = registry.tools_schema().len();
1102
1103 registry.disable(&["bash".to_string(), "ls".to_string()]);
1104
1105 assert!(registry.get("bash").is_none());
1107 assert!(registry.get("ls").is_none());
1108 assert!(registry.get("read").is_some());
1110 assert!(registry.get("subagent").is_some());
1111 assert_eq!(registry.tools_schema().len(), before - 2);
1113 }
1114
1115 #[test]
1116 fn test_disable_empty_and_unknown_are_noops() {
1117 let mut registry = ToolRegistry::new();
1118 let n = registry.tools_schema().len();
1119 registry.disable(&[]); assert_eq!(registry.tools_schema().len(), n);
1121 registry.disable(&["not_a_real_tool".to_string()]); assert_eq!(registry.tools_schema().len(), n);
1123 assert!(registry.get("bash").is_some());
1124 }
1125
1126 #[test]
1127 fn test_tool_registry_without_subagent() {
1128 let registry = ToolRegistry::without_subagent();
1129
1130 assert_eq!(
1132 registry.tools_schema().len(),
1133 13 + usize::from(cfg!(windows))
1134 );
1135 for name in ["forum_post", "forum_read", "forum_forget"] {
1136 assert!(registry.get(name).is_some(), "missing {name}");
1137 }
1138 assert!(registry.get("subagent").is_none());
1140
1141 assert!(registry.get("bash").is_some());
1143
1144 assert!(registry.get("bash").is_some());
1146 assert!(registry.get("read").is_some());
1147 assert!(registry.get("write").is_some());
1148 assert!(registry.get("edit").is_some());
1149 assert!(registry.get("grep").is_some());
1150 assert!(registry.get("find").is_some());
1151 assert!(registry.get("ls").is_some());
1152 }
1153
1154 #[test]
1155 fn test_tool_registry_register() {
1156 let mut registry = ToolRegistry::without_subagent();
1157 let initial_count = registry.tools_schema().len();
1158
1159 struct TestTool;
1161 #[async_trait::async_trait]
1162 impl Tool for TestTool {
1163 fn name(&self) -> &str {
1164 "test_tool"
1165 }
1166 fn description(&self) -> &str {
1167 "A test tool"
1168 }
1169 fn parameters(&self) -> Value {
1170 json!({"type": "object"})
1171 }
1172 async fn execute(&self, _params: Value, _ctx: ToolContext) -> Result<String> {
1173 Ok("test result".to_string())
1174 }
1175 }
1176
1177 registry.register(Arc::new(TestTool));
1178
1179 assert_eq!(registry.tools_schema().len(), initial_count + 1);
1181
1182 assert!(registry.get("test_tool").is_some());
1184 }
1185
1186 #[test]
1187 fn tool_names_for_extension_filters_by_owner_and_sorts() {
1188 struct OwnedTool(&'static str, Option<&'static str>);
1189 #[async_trait::async_trait]
1190 impl Tool for OwnedTool {
1191 fn name(&self) -> &str {
1192 self.0
1193 }
1194 fn description(&self) -> &str {
1195 "owned"
1196 }
1197 fn parameters(&self) -> Value {
1198 json!({"type": "object"})
1199 }
1200 async fn execute(&self, _params: Value, _ctx: ToolContext) -> Result<String> {
1201 Ok("ok".to_string())
1202 }
1203 fn extension_id(&self) -> Option<&str> {
1204 self.1
1205 }
1206 }
1207
1208 let mut registry = ToolRegistry::without_subagent();
1209 registry.register(Arc::new(OwnedTool("alpha:zed", Some("alpha"))));
1210 registry.register(Arc::new(OwnedTool("alpha:bar", Some("alpha"))));
1211 registry.register(Arc::new(OwnedTool("beta:thing", Some("beta"))));
1212
1213 assert_eq!(
1214 registry.tool_names_for_extension("alpha"),
1215 vec!["alpha:bar".to_string(), "alpha:zed".to_string()]
1216 );
1217 assert_eq!(
1218 registry.tool_names_for_extension("beta"),
1219 vec!["beta:thing".to_string()]
1220 );
1221 assert!(registry.tool_names_for_extension("ghost").is_empty());
1222 assert!(registry.tool_names_for_extension("bash").is_empty());
1224 }
1225
1226 struct OwnedTool(&'static str, Option<&'static str>);
1227 #[async_trait::async_trait]
1228 impl Tool for OwnedTool {
1229 fn name(&self) -> &str {
1230 self.0
1231 }
1232 fn description(&self) -> &str {
1233 "owned"
1234 }
1235 fn parameters(&self) -> Value {
1236 json!({"type": "object"})
1237 }
1238 async fn execute(&self, _params: Value, _ctx: ToolContext) -> Result<String> {
1239 Ok("ok".to_string())
1240 }
1241 fn extension_id(&self) -> Option<&str> {
1242 self.1
1243 }
1244 }
1245
1246 #[test]
1247 fn without_subagent_excludes_subagent_tools() {
1248 let registry = ToolRegistry::without_subagent();
1249 assert!(registry.get("subagent").is_none());
1250 assert!(registry.get("subagent_start").is_none());
1251 assert!(registry.get("subagent_status").is_none());
1252 assert!(registry.get("subagent_steer").is_none());
1253 assert!(registry.get("subagent_collect").is_none());
1254 assert!(registry.get("subagent_resume").is_none());
1255 assert!(registry.get("bash").is_some());
1257 assert!(registry.get("read").is_some());
1258 }
1259
1260 #[test]
1261 fn without_subagent_with_extensions_includes_extension_tools() {
1262 let mut other = ToolRegistry::empty();
1263 other.register(Arc::new(OwnedTool("alpha:do_thing", Some("alpha"))));
1264
1265 let merged = ToolRegistry::without_subagent_with_extensions(&other);
1266
1267 assert!(merged.get("alpha:do_thing").is_some());
1269 assert!(merged.get("bash").is_some());
1271 assert!(merged.get("read").is_some());
1272 assert!(merged.get("subagent_start").is_none());
1274 }
1275
1276 #[test]
1277 fn forum_membership_preserves_worker_bounds_and_foreground_disclosure() {
1278 use crate::tools::activation::{SessionId, SessionToolSet};
1279 use crate::tools::catalog::{ToolEffect, ToolId};
1280 let main = ToolRegistry::new();
1281 let mut extensions = ToolRegistry::empty();
1282 extensions.register(Arc::new(OwnedTool("peer:probe", Some("peer"))));
1283 let merged = ToolRegistry::without_subagent_with_extensions(&extensions);
1284 assert_eq!(merged.tools_schema().len(), 14 + usize::from(cfg!(windows)));
1285 assert!(merged.get("peer:probe").is_some());
1286 for registry in [&main, &ToolRegistry::without_subagent(), &merged] {
1287 for (name, effect) in [
1288 ("forum_post", ToolEffect::NonIdempotent),
1289 ("forum_read", ToolEffect::ReadOnly),
1290 ("forum_forget", ToolEffect::NonIdempotent),
1291 ] {
1292 let tool = registry.get(name).unwrap();
1293 assert_eq!(tool.origin(), crate::tools::ToolOrigin::Builtin);
1294 assert_eq!(tool.effect(), effect);
1295 assert!(registry.catalog().get(&ToolId::builtin(name)).is_some());
1296 assert_eq!(
1297 registry
1298 .tools_schema()
1299 .iter()
1300 .filter(|schema| schema["name"] == name)
1301 .count(),
1302 1
1303 );
1304 }
1305 }
1306 for registry in [ToolRegistry::without_subagent(), merged] {
1307 for name in [
1308 "subagent",
1309 "subagent_start",
1310 "subagent_status",
1311 "subagent_steer",
1312 "subagent_collect",
1313 "subagent_resume",
1314 "subagent_model_authorize",
1315 "subagent_models",
1316 "search_tools",
1317 "activate_tools",
1318 "memory_context",
1319 "memory_store",
1320 "memory_search",
1321 "memory_fetch",
1322 "memory_forget",
1323 ] {
1324 assert!(registry.get(name).is_none(), "worker must not gain {name}");
1325 }
1326 for name in ["bash", "write", "edit"] {
1327 assert!(registry.get(name).is_some(), "worker must retain {name}");
1328 }
1329 }
1330 let session = SessionToolSet::progressive_core_for_catalog(
1331 SessionId::parse("forum-foreground-test").unwrap(),
1332 main.catalog(),
1333 );
1334 let projection = main.session_tools_schema(&session);
1335 for name in ["forum_post", "forum_read", "forum_forget"] {
1336 assert!(!session.is_core(&ToolId::builtin(name)));
1337 assert!(!projection
1338 .schema
1339 .iter()
1340 .any(|schema| schema["name"] == name));
1341 }
1342 }
1343
1344 #[test]
1345 fn without_subagent_with_extensions_excludes_built_ins_from_other_registry() {
1346 let other = ToolRegistry::new();
1351
1352 let merged = ToolRegistry::without_subagent_with_extensions(&other);
1353
1354 let bash = merged.get("bash").expect("bash present");
1356 assert!(bash.extension_id().is_none());
1357 assert!(merged.get("subagent_start").is_none());
1359 assert!(merged.get("subagent").is_none());
1360 }
1361
1362 #[test]
1363 fn without_subagent_with_extensions_rejects_recursive_tool_name_collisions() {
1364 let mut other = ToolRegistry::empty();
1365 for name in [
1366 "subagent",
1367 "subagent_start",
1368 "subagent_status",
1369 "subagent_steer",
1370 "subagent_collect",
1371 "subagent_resume",
1372 "subagent_model_authorize",
1373 "subagent_models",
1374 "memory_context",
1375 ] {
1376 other.register(Arc::new(OwnedTool(name, Some("malicious-extension"))));
1377 }
1378
1379 let merged = ToolRegistry::without_subagent_with_extensions(&other);
1380 for name in [
1381 "subagent",
1382 "subagent_start",
1383 "subagent_status",
1384 "subagent_steer",
1385 "subagent_collect",
1386 "subagent_resume",
1387 "subagent_model_authorize",
1388 "subagent_models",
1389 "memory_context",
1390 ] {
1391 assert!(merged.get(name).is_none(), "{name} must stay unavailable");
1392 }
1393 }
1394
1395 #[test]
1396 fn without_subagent_with_extensions_rejects_api_sanitized_recursive_collisions() {
1397 let mut other = ToolRegistry::empty();
1398 for runtime_name in [
1399 "subagent:start",
1400 "subagent:status",
1401 "subagent:steer",
1402 "subagent:collect",
1403 "subagent:resume",
1404 "subagent:model:authorize",
1405 "subagent:models",
1406 "memory:context",
1407 ] {
1408 other.register(Arc::new(OwnedTool(runtime_name, Some("subagent"))));
1409 }
1410
1411 let merged = ToolRegistry::without_subagent_with_extensions(&other);
1412 for api_name in [
1413 "subagent_start",
1414 "subagent_status",
1415 "subagent_steer",
1416 "subagent_collect",
1417 "subagent_resume",
1418 "subagent_model_authorize",
1419 "subagent_models",
1420 "memory_context",
1421 ] {
1422 assert!(
1423 merged.get(api_name).is_none(),
1424 "API-safe recursive name {api_name} must stay unavailable"
1425 );
1426 }
1427 }
1428
1429 #[test]
1433 fn memory_context_present_in_primary_registry_absent_in_recursive_ones() {
1434 assert!(ToolRegistry::new().get("memory_context").is_some());
1435 assert!(ToolRegistry::without_subagent()
1436 .get("memory_context")
1437 .is_none());
1438
1439 let merged = ToolRegistry::without_subagent_with_extensions(&ToolRegistry::new());
1440 assert!(
1441 merged.get("memory_context").is_none(),
1442 "recursive registries must never expose memory_context"
1443 );
1444 }
1445
1446 #[test]
1447 fn without_subagent_with_extensions_does_not_overwrite_existing_builtin() {
1448 let mut other = ToolRegistry::empty();
1454 other.register(Arc::new(OwnedTool("ext:custom", Some("ext"))));
1455
1456 let merged = ToolRegistry::without_subagent_with_extensions(&other);
1457 assert!(merged.get("ext:custom").is_some());
1458 assert!(merged.get("bash").is_some());
1459 assert!(merged.get("bash").unwrap().extension_id().is_none());
1460 }
1461
1462 struct McpTwin(&'static str);
1464
1465 #[async_trait::async_trait]
1466 impl Tool for McpTwin {
1467 fn name(&self) -> &str {
1468 self.0
1469 }
1470 fn description(&self) -> &str {
1471 "mcp twin"
1472 }
1473 fn parameters(&self) -> Value {
1474 json!({"type": "object"})
1475 }
1476 async fn execute(&self, _params: Value, _ctx: ToolContext) -> Result<String> {
1477 Ok("ok".to_string())
1478 }
1479 fn origin(&self) -> crate::tools::ToolOrigin {
1480 crate::tools::ToolOrigin::Mcp {
1481 server_id: "srv".to_string(),
1482 server_tool_name: "shared".to_string(),
1483 }
1484 }
1485 }
1486
1487 #[test]
1488 #[should_panic(expected = "tool catalog rejected mutation")]
1489 fn from_tools_with_duplicate_capability_identity_fails_closed() {
1490 let _ = ToolRegistry::from_tools(vec![
1494 Arc::new(McpTwin("twin_one")),
1495 Arc::new(McpTwin("twin_two")),
1496 ]);
1497 }
1498}