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 #[allow(unused_mut)]
42 let mut tools: Vec<Arc<dyn Tool>> = vec![
43 #[cfg(all(windows, feature = "command-exec"))]
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::multi_edit::MultiEditTool),
49 Arc::new(crate::tools::grep::GrepTool),
50 Arc::new(crate::tools::find::FindTool),
51 Arc::new(crate::tools::ls::LsTool),
52 Arc::new(crate::tools::subagent::SubagentTool),
53 Arc::new(crate::tools::subagent::authorize_model::SubagentModelAuthorizeTool),
54 Arc::new(crate::tools::subagent::models::SubagentModelsTool),
55 Arc::new(crate::tools::subagent::start::SubagentStartTool),
56 Arc::new(crate::tools::subagent::status::SubagentStatusTool),
57 Arc::new(crate::tools::subagent::steer::SubagentSteerTool),
58 Arc::new(crate::tools::subagent::collect::SubagentCollectTool),
59 Arc::new(crate::tools::subagent::resume::SubagentResumeTool),
60 Arc::new(crate::tools::memory_context::MemoryContextTool),
65 Arc::new(crate::tools::discovery::SearchToolsTool),
70 Arc::new(crate::tools::discovery::ActivateToolsTool),
71 Arc::new(crate::tools::memory::MemorySearchTool),
75 Arc::new(crate::tools::memory::MemoryFetchTool),
76 Arc::new(crate::tools::memory::MemoryStoreTool),
77 Arc::new(crate::tools::memory::MemoryForgetTool),
78 Arc::new(crate::tools::forum::ForumPostTool),
81 Arc::new(crate::tools::forum::ForumReadTool),
82 Arc::new(crate::tools::forum::ForumForgetTool),
83 ];
84 #[cfg(feature = "command-exec")]
88 {
89 tools.push(Arc::new(crate::tools::bash::BashTool));
90 tools.push(Arc::new(crate::tools::shell::ShellStartTool));
91 tools.push(Arc::new(crate::tools::shell::ShellSendTool));
92 tools.push(Arc::new(crate::tools::shell::ShellEndTool));
93 }
94 Self::from_tools(tools)
95 }
96
97 pub fn empty() -> Self {
100 Self::from_tools(Vec::new())
101 }
102
103 pub fn disable(&mut self, names: &[String]) {
109 self.try_disable(names).expect(CATALOG_REJECTED_MSG);
110 }
111
112 pub fn try_disable(&mut self, names: &[String]) -> Result<(), CatalogError> {
125 if names.is_empty() {
126 return Ok(());
127 }
128 let remaining: Vec<Arc<dyn Tool>> = self
129 .tools
130 .values()
131 .filter(|t| !names.iter().any(|n| n == t.name()))
132 .cloned()
133 .collect();
134 if remaining.len() == self.tools.len() {
135 return Ok(());
136 }
137 let prior = self.catalog.generation();
138 let mut candidate = Self::try_from_tools(remaining)?;
139 candidate.catalog.rebase_past(prior)?;
140 *self = candidate;
141 Ok(())
142 }
143
144 pub fn without_subagent() -> Self {
146 #[allow(unused_mut)]
147 let mut tools: Vec<Arc<dyn Tool>> = vec![
148 #[cfg(all(windows, feature = "command-exec"))]
149 Arc::new(crate::tools::powershell::PowerShellTool),
150 Arc::new(crate::tools::read::ReadTool),
151 Arc::new(crate::tools::write::WriteTool),
152 Arc::new(crate::tools::edit::EditTool),
153 Arc::new(crate::tools::multi_edit::MultiEditTool),
154 Arc::new(crate::tools::grep::GrepTool),
155 Arc::new(crate::tools::find::FindTool),
156 Arc::new(crate::tools::ls::LsTool),
157 Arc::new(crate::tools::forum::ForumPostTool),
158 Arc::new(crate::tools::forum::ForumReadTool),
159 Arc::new(crate::tools::forum::ForumForgetTool),
160 ];
161 #[cfg(feature = "command-exec")]
162 {
163 tools.push(Arc::new(crate::tools::bash::BashTool));
164 tools.push(Arc::new(crate::tools::shell::ShellStartTool));
165 tools.push(Arc::new(crate::tools::shell::ShellSendTool));
166 tools.push(Arc::new(crate::tools::shell::ShellEndTool));
167 }
168 Self::from_tools(tools)
169 }
170
171 pub fn without_subagent_with_extensions(extension_tools: &ToolRegistry) -> Self {
179 Self::worker_set(extension_tools, &[])
180 }
181
182 pub fn worker_set(extension_tools: &ToolRegistry, disabled: &[String]) -> Self {
189 let mut combined = Self::without_subagent();
190 combined.disable(disabled);
191 for tool in extension_tools.tools.values() {
192 if tool.extension_id().is_some() && !is_recursive_subagent_tool_name(tool.name()) {
193 combined
194 .insert_tool_with_catalog(tool.clone())
195 .expect(CATALOG_REJECTED_MSG);
196 }
197 }
198 combined.rebuild_schema();
199 combined
200 }
201
202 fn from_tools(tool_list: Vec<Arc<dyn Tool>>) -> Self {
208 Self::try_from_tools(tool_list).expect(CATALOG_REJECTED_MSG)
209 }
210
211 #[doc(hidden)]
215 pub fn from_tools_for_tests(tool_list: Vec<Arc<dyn Tool>>) -> Self {
216 Self::from_tools(tool_list)
217 }
218
219 fn try_from_tools(tool_list: Vec<Arc<dyn Tool>>) -> Result<Self, CatalogError> {
220 let mut registry = ToolRegistry {
221 tools: HashMap::new(),
222 cached_schema: Arc::new(Vec::new()),
223 api_to_runtime_names: HashMap::new(),
224 input_name_maps: HashMap::new(),
225 catalog: ToolCatalog::empty(),
226 };
227 for tool in tool_list {
231 let name = tool.name().to_string();
232 registry.tools.insert(name, tool);
233 }
234 let sorted: Vec<Arc<dyn Tool>> =
240 registry.iter_tools_sorted().into_iter().cloned().collect();
241 for tool in sorted {
242 let record = CapabilityRecord::for_registered_tool(&tool);
243 registry.catalog.upsert(None, record)?;
244 }
245 registry.rebuild_schema();
246 Ok(registry)
247 }
248
249 fn api_safe_name(name: &str, used: &HashSet<String>) -> String {
250 Self::api_safe_identifier(name, used, 128, false)
251 }
252
253 pub(crate) fn unclaimed_wire_name(runtime_name: &str) -> String {
256 Self::api_safe_name(runtime_name, &HashSet::new())
257 }
258
259 fn api_safe_property_name(name: &str, used: &HashSet<String>) -> String {
260 Self::api_safe_identifier(name, used, 64, true)
261 }
262
263 fn api_safe_identifier(
264 name: &str,
265 used: &HashSet<String>,
266 max_len: usize,
267 allow_dot: bool,
268 ) -> String {
269 let mut sanitized = String::with_capacity(name.len());
270 for ch in name.chars() {
271 if ch.is_ascii_alphanumeric() || ch == '_' || ch == '-' || (allow_dot && ch == '.') {
272 sanitized.push(ch);
273 } else {
274 sanitized.push('_');
275 }
276 }
277 if sanitized.is_empty() {
278 sanitized.push_str("field");
279 }
280 if sanitized.len() > max_len {
281 sanitized.truncate(max_len);
282 }
283
284 let base = sanitized.clone();
285 let mut suffix = 2;
286 while used.contains(&sanitized) {
287 let suffix_str = format!("_{suffix}");
288 let keep = max_len.saturating_sub(suffix_str.len());
289 sanitized = format!("{}{}", &base[..base.len().min(keep)], suffix_str);
290 suffix += 1;
291 }
292 sanitized
293 }
294
295 fn sanitize_schema(mut schema: Value) -> (Value, SchemaNameMap) {
296 let mut map = SchemaNameMap::default();
297 let Some(obj) = schema.as_object_mut() else {
298 return (schema, map);
299 };
300
301 let mut required_name_map = HashMap::new();
302 if let Some(props_value) = obj.get_mut("properties") {
303 if let Some(props) = props_value.as_object_mut() {
304 let original = std::mem::take(props);
305 let mut used = HashSet::new();
306 for (runtime_name, child_schema) in original {
307 let api_name = Self::api_safe_property_name(&runtime_name, &used);
308 used.insert(api_name.clone());
309 required_name_map.insert(runtime_name.clone(), api_name.clone());
310 map.api_to_runtime.insert(api_name.clone(), runtime_name);
311
312 let (sanitized_child, child_map) = Self::sanitize_schema(child_schema);
313 if !child_map.api_to_runtime.is_empty() || !child_map.children.is_empty() {
314 map.children.insert(api_name.clone(), child_map);
315 }
316 props.insert(api_name, sanitized_child);
317 }
318 }
319 }
320
321 if let Some(required) = obj.get_mut("required").and_then(Value::as_array_mut) {
322 for item in required.iter_mut() {
323 if let Some(name) = item.as_str() {
324 if let Some(api_name) = required_name_map.get(name) {
325 *item = Value::String(api_name.clone());
326 }
327 }
328 }
329 }
330
331 if let Some(items) = obj.get_mut("items") {
334 let (sanitized_items, items_map) = Self::sanitize_schema(std::mem::take(items));
335 if !items_map.api_to_runtime.is_empty()
336 || !items_map.children.is_empty()
337 || items_map.items.is_some()
338 {
339 map.items = Some(Box::new(items_map));
340 }
341 *items = sanitized_items;
342 }
343
344 (schema, map)
345 }
346
347 fn translate_input_names(input: Value, map: &SchemaNameMap) -> Value {
348 match input {
349 Value::Object(obj) => {
350 let mut out = serde_json::Map::new();
351 for (api_name, value) in obj {
352 let runtime_name = map
353 .api_to_runtime
354 .get(&api_name)
355 .cloned()
356 .unwrap_or_else(|| api_name.clone());
357 let value = if let Some(child) = map.children.get(&api_name) {
358 Self::translate_input_names(value, child)
359 } else {
360 value
361 };
362 out.insert(runtime_name, value);
363 }
364 Value::Object(out)
365 }
366 Value::Array(arr) => {
367 if let Some(items_map) = &map.items {
369 Value::Array(
370 arr.into_iter()
371 .map(|v| Self::translate_input_names(v, items_map))
372 .collect(),
373 )
374 } else {
375 Value::Array(arr)
376 }
377 }
378 other => other,
379 }
380 }
381
382 fn rebuild_schema(&mut self) {
383 let mut used = HashSet::new();
384 let mut api_to_runtime_names = HashMap::new();
385 let mut input_name_maps = HashMap::new();
386 let mut schema = Vec::with_capacity(self.tools.len());
387
388 let mut sorted_tools: Vec<_> = self.tools.values().collect();
403 sorted_tools.sort_by_key(|t| {
404 let builtin_first = match t.origin() {
405 crate::tools::ToolOrigin::Builtin => 0u8,
406 _ => 1u8,
407 };
408 (builtin_first, t.name().to_string())
409 });
410
411 for tool in sorted_tools {
412 let runtime_name = tool.name();
413 let api_name = Self::api_safe_name(runtime_name, &used);
414 used.insert(api_name.clone());
415 api_to_runtime_names.insert(api_name.clone(), runtime_name.to_string());
416 let (input_schema, input_map) = Self::sanitize_schema(tool.parameters());
417 input_name_maps.insert(api_name.clone(), input_map);
418 schema.push(serde_json::json!({
419 "name": api_name,
420 "description": tool.description(),
421 "input_schema": input_schema
422 }));
423 }
424
425 self.api_to_runtime_names = api_to_runtime_names;
426 self.input_name_maps = input_name_maps;
427 self.cached_schema = Arc::new(schema);
428 }
429
430 pub fn register(&mut self, tool: Arc<dyn Tool>) {
440 self.try_register(tool).expect(CATALOG_REJECTED_MSG);
441 }
442
443 pub fn try_register(&mut self, tool: Arc<dyn Tool>) -> Result<(), CatalogError> {
449 self.insert_tool_with_catalog(tool)?;
450 self.rebuild_schema();
451 Ok(())
452 }
453
454 pub fn try_register_batch(
468 &mut self,
469 tool_list: Vec<Arc<dyn Tool>>,
470 ) -> Result<usize, CatalogError> {
471 if tool_list.is_empty() {
472 return Ok(0);
473 }
474 let mut batch_names: HashSet<String> = HashSet::new();
475 let mut batch_ids: HashSet<crate::tools::catalog::ToolId> = HashSet::new();
476 for tool in &tool_list {
477 let name = tool.name();
478 if self.tools.contains_key(name) || !batch_names.insert(name.to_string()) {
479 return Err(CatalogError::DuplicateRuntimeName(
480 agent_core::BoundedText::new(name, 128).text,
481 ));
482 }
483 let id = tool_id_for(tool.as_ref());
484 if self.catalog.get(&id).is_some() || !batch_ids.insert(id.clone()) {
485 return Err(CatalogError::DuplicateToolId(id));
486 }
487 }
488 let count = tool_list.len();
489 let mut candidate = self.clone();
490 for tool in tool_list {
491 candidate.insert_tool_with_catalog(tool)?;
492 }
493 candidate.rebuild_schema();
494 *self = candidate;
495 Ok(count)
496 }
497
498 fn insert_tool_with_catalog(&mut self, tool: Arc<dyn Tool>) -> Result<(), CatalogError> {
507 let record = CapabilityRecord::for_registered_tool(&tool);
508 let replaced = self
509 .tools
510 .get(tool.name())
511 .map(|old| tool_id_for(old.as_ref()));
512 self.catalog.upsert(replaced.as_ref(), record)?;
513 self.tools.insert(tool.name().to_string(), tool);
514 Ok(())
515 }
516
517 pub fn catalog(&self) -> &ToolCatalog {
520 &self.catalog
521 }
522
523 #[doc(hidden)]
531 pub fn resume_catalog_generation_for_tests(
532 &mut self,
533 generation: crate::tools::catalog::CatalogGeneration,
534 ) {
535 self.catalog.set_generation_for_tests(generation);
536 }
537
538 #[doc(hidden)]
544 pub fn upsert_catalog_record_for_tests(
545 &mut self,
546 replaced: Option<&crate::tools::catalog::ToolId>,
547 record: crate::tools::catalog::CapabilityRecord,
548 ) {
549 self.catalog
550 .upsert(replaced, record)
551 .expect("test upsert must be structurally valid");
552 }
553
554 pub fn get(&self, name: &str) -> Option<&Arc<dyn Tool>> {
555 let runtime_name = self
556 .api_to_runtime_names
557 .get(name)
558 .map(String::as_str)
559 .unwrap_or(name);
560 self.tools.get(runtime_name)
561 }
562
563 pub fn runtime_name_for_api<'a>(&'a self, name: &'a str) -> &'a str {
564 self.api_to_runtime_names
565 .get(name)
566 .map(String::as_str)
567 .unwrap_or(name)
568 }
569
570 pub fn resolve_wire_call(
580 &self,
581 wire_name: &str,
582 ) -> Option<crate::tools::activation::ResolvedToolCall> {
583 let runtime_name = self
584 .api_to_runtime_names
585 .get(wire_name)
586 .map(String::as_str)
587 .unwrap_or(wire_name);
588 let tool = self.tools.get(runtime_name)?;
589 Some(crate::tools::activation::ResolvedToolCall::new(
590 wire_name,
591 runtime_name,
592 tool_id_for(tool.as_ref()),
593 ))
594 }
595
596 pub fn translate_input_for_api_tool(&self, tool_name: &str, input: Value) -> Value {
597 if let Some(map) = self.input_name_maps.get(tool_name) {
598 Self::translate_input_names(input, map)
599 } else {
600 input
601 }
602 }
603
604 pub fn tools_schema(&self) -> Arc<Vec<Value>> {
605 Arc::clone(&self.cached_schema)
606 }
607
608 pub fn deferred_tool_split(
644 &self,
645 session: &crate::tools::activation::SessionToolSet,
646 referenced_ids: &HashSet<String>,
647 ) -> DeferredToolSplit {
648 let projection = self.session_tools_schema(session).schema;
649 let projected: HashSet<&str> = projection
650 .iter()
651 .filter_map(|entry| entry["name"].as_str())
652 .collect();
653 let mut split = DeferredToolSplit::default();
654 for entry in self.cached_schema.iter() {
655 let Some(api_name) = entry["name"].as_str() else {
656 continue;
657 };
658 let Some(tool) = self.tools.get(self.runtime_name_for_api(api_name)) else {
659 continue;
660 };
661 let id = tool_id_for(tool.as_ref());
662 if projected.contains(api_name) {
663 if session.is_core(&id) || !referenced_ids.contains(id.as_str()) {
664 split.loaded.push(entry.clone());
665 } else {
666 split.deferred.push(deferred_entry(entry));
667 }
668 } else if self
669 .catalog
670 .get(&id)
671 .is_some_and(crate::tools::activation::is_activatable_record)
672 {
673 split.deferred.push(deferred_entry(entry));
674 } else {
675 continue;
676 }
677 split
678 .id_to_api
679 .insert(id.as_str().to_string(), api_name.to_string());
680 }
681 split
682 }
683
684 pub fn session_tools_schema(
685 &self,
686 session: &crate::tools::activation::SessionToolSet,
687 ) -> SessionSchemaProjection {
688 let mut projected = Vec::new();
689 let mut dropped = Vec::new();
690 for entry in self.cached_schema.iter() {
691 let Some(api_name) = entry["name"].as_str() else {
692 continue;
693 };
694 let runtime_name = self.runtime_name_for_api(api_name);
695 let Some(tool) = self.tools.get(runtime_name) else {
696 continue;
697 };
698 let id = tool_id_for(tool.as_ref());
699 let pinned = session
700 .core_pin(&id)
701 .map(|pin| (pin.schema_digest(), pin.provenance()))
702 .or_else(|| {
703 session
704 .activation(&id)
705 .map(|a| (a.schema_digest(), a.provenance()))
706 });
707 let Some((pinned_digest, pinned_provenance)) = pinned else {
708 continue; };
710 let Some(record) = self.catalog.get(&id) else {
711 tracing::warn!(
712 tool = %id,
713 "session schema projection drops uncataloged session member \
714 (removed since the session pinned it)"
715 );
716 dropped.push((id, DroppedSessionMember::Uncataloged));
717 continue;
718 };
719 if pinned_digest != record.schema_digest() {
720 tracing::warn!(
721 tool = %id,
722 "session schema projection drops drifted session member \
723 (schema digest changed since the session pinned it)"
724 );
725 dropped.push((id, DroppedSessionMember::SchemaDigestMismatch));
726 continue;
727 }
728 if pinned_provenance != record.provenance() {
729 tracing::warn!(
730 tool = %id,
731 "session schema projection drops drifted session member \
732 (trust provenance changed since the session pinned it)"
733 );
734 dropped.push((id, DroppedSessionMember::ProvenanceMismatch));
735 continue;
736 }
737 projected.push(entry.clone());
738 }
739 let seen: HashSet<&crate::tools::catalog::ToolId> =
746 dropped.iter().map(|(id, _)| id).collect();
747 let vanished: Vec<crate::tools::catalog::ToolId> = session
748 .core_ids()
749 .chain(session.activated().map(|a| a.grant().tool_id()))
750 .filter(|id| !seen.contains(id) && self.catalog.get(id).is_none())
751 .cloned()
752 .collect();
753 for id in vanished {
754 tracing::warn!(
755 tool = %id,
756 "session schema projection drops uncataloged session member (removed from the registry since the session pinned it)"
757 );
758 dropped.push((id, DroppedSessionMember::Uncataloged));
759 }
760 SessionSchemaProjection {
761 schema: projected,
762 dropped,
763 }
764 }
765
766 pub fn tool_names_for_extension(&self, extension_id: &str) -> Vec<String> {
769 let mut names: Vec<String> = self
770 .tools
771 .values()
772 .filter(|t| t.extension_id() == Some(extension_id))
773 .map(|t| t.name().to_string())
774 .collect();
775 names.sort();
776 names
777 }
778
779 pub fn iter_tools_sorted(&self) -> Vec<&Arc<dyn Tool>> {
782 let mut tools: Vec<_> = self.tools.values().collect();
783 tools.sort_by_key(|t| t.name());
784 tools
785 }
786}
787
788#[derive(Clone, Debug, PartialEq, Eq)]
794pub struct SessionSchemaProjection {
795 pub schema: Vec<Value>,
797 pub dropped: Vec<(crate::tools::catalog::ToolId, DroppedSessionMember)>,
801}
802
803#[derive(Clone, Copy, Debug, PartialEq, Eq)]
805pub enum DroppedSessionMember {
806 Uncataloged,
808 SchemaDigestMismatch,
810 ProvenanceMismatch,
813}
814
815const CATALOG_REJECTED_MSG: &str =
821 "tool catalog rejected mutation: refusing to expose a tool the catalog could not record";
822
823fn is_recursive_subagent_tool_name(name: &str) -> bool {
824 let api_name = ToolRegistry::api_safe_name(name, &HashSet::new());
825 matches!(
826 api_name.as_str(),
827 "subagent"
828 | "subagent_start"
829 | "subagent_status"
830 | "subagent_steer"
831 | "subagent_collect"
832 | "subagent_resume"
833 | "subagent_model_authorize"
834 | "subagent_models"
835 | "memory_context"
839 )
840}
841#[derive(Debug, Clone, Default)]
843pub struct DeferredToolSplit {
844 pub loaded: Vec<Value>,
846 pub deferred: Vec<Value>,
848 pub id_to_api: std::collections::HashMap<String, String>,
850}
851
852fn deferred_entry(entry: &Value) -> Value {
853 let mut entry = entry.clone();
854 if let Some(map) = entry.as_object_mut() {
855 map.remove("cache_control");
856 map.insert("defer_loading".to_string(), Value::Bool(true));
857 }
858 entry
859}
860
861#[cfg(test)]
862mod tests {
863 use super::*;
864 use crate::{Result, ToolContext};
865 use serde_json::json;
866
867 struct NamedTool(&'static str);
868
869 #[async_trait::async_trait]
870 impl Tool for NamedTool {
871 fn name(&self) -> &str {
872 self.0
873 }
874 fn description(&self) -> &str {
875 "test tool"
876 }
877 fn parameters(&self) -> Value {
878 json!({"type": "object"})
879 }
880 async fn execute(&self, _params: Value, _ctx: ToolContext) -> Result<String> {
881 Ok("ok".to_string())
882 }
883 }
884
885 struct SchemaTool;
886
887 #[async_trait::async_trait]
888 impl Tool for SchemaTool {
889 fn name(&self) -> &str {
890 "schema_tool"
891 }
892 fn description(&self) -> &str {
893 "schema tool"
894 }
895 fn parameters(&self) -> Value {
896 json!({
897 "type": "object",
898 "properties": {
899 "bad:key/that/is/far/too/long/for/anthropic/property/names/and/keeps/going": {"type": "string"},
900 "nested:obj": {
901 "type": "object",
902 "properties": {"inner/key": {"type": "string"}},
903 "required": ["inner/key"]
904 }
905 },
906 "required": [
907 "bad:key/that/is/far/too/long/for/anthropic/property/names/and/keeps/going",
908 "nested:obj"
909 ]
910 })
911 }
912 async fn execute(&self, _params: Value, _ctx: ToolContext) -> Result<String> {
913 Ok("ok".to_string())
914 }
915 }
916
917 #[test]
918 fn tool_schema_uses_api_safe_names_and_maps_back() {
919 let registry = ToolRegistry::from_tools(vec![Arc::new(NamedTool("plugin:skill.tool"))]);
920
921 assert_eq!(registry.tools_schema()[0]["name"], "plugin_skill_tool");
922 assert!(registry.get("plugin:skill.tool").is_some());
923 assert!(registry.get("plugin_skill_tool").is_some());
924 assert_eq!(
925 registry.runtime_name_for_api("plugin_skill_tool"),
926 "plugin:skill.tool"
927 );
928 }
929
930 #[test]
940 fn extension_leaf_name_never_shadows_builtin_memory_fetch_api_name() {
941 struct ExtMemoryTool {
942 runtime: &'static str,
943 ext: &'static str,
944 }
945 #[async_trait::async_trait]
946 impl Tool for ExtMemoryTool {
947 fn origin(&self) -> crate::tools::ToolOrigin {
948 crate::tools::ToolOrigin::Extension {
949 extension_id: self.ext.to_string(),
950 }
951 }
952 fn extension_id(&self) -> Option<&str> {
953 Some(self.ext)
954 }
955 fn name(&self) -> &str {
956 self.runtime
957 }
958 fn description(&self) -> &str {
959 "extension memory tool (single-id schema)"
960 }
961 fn parameters(&self) -> Value {
962 json!({
963 "type": "object",
964 "properties": {"id": {"type": "string"}},
965 "required": ["id"]
966 })
967 }
968 async fn execute(&self, _params: Value, _ctx: ToolContext) -> Result<String> {
969 Ok("ok".to_string())
970 }
971 }
972
973 let registry = ToolRegistry::from_tools(vec![
974 Arc::new(crate::tools::memory::MemoryFetchTool),
975 Arc::new(ExtMemoryTool {
977 runtime: "memory:fetch",
978 ext: "memory",
979 }),
980 Arc::new(ExtMemoryTool {
982 runtime: "axel-memory-manager:memory_fetch",
983 ext: "axel-memory-manager",
984 }),
985 ]);
986
987 assert_eq!(
989 registry.runtime_name_for_api("memory_fetch"),
990 "memory_fetch",
991 "wire name `memory_fetch` must resolve to the built-in"
992 );
993 let schema = registry.tools_schema();
994 let builtin = schema
995 .iter()
996 .find(|e| e["name"] == "memory_fetch")
997 .expect("built-in memory_fetch exposed");
998 assert_eq!(
999 builtin["input_schema"]["required"],
1000 json!(["ids"]),
1001 "built-in memory_fetch schema must remain {{ids: [...]}}"
1002 );
1003
1004 let axel = schema
1006 .iter()
1007 .find(|e| e["name"] == "axel-memory-manager_memory_fetch")
1008 .expect("extension tool exposed under namespaced API name");
1009 assert_eq!(axel["input_schema"]["required"], json!(["id"]));
1010 assert_eq!(
1013 registry.runtime_name_for_api("memory_fetch_2"),
1014 "memory:fetch"
1015 );
1016 let resolved = registry.resolve_wire_call("memory_fetch").unwrap();
1019 assert_eq!(resolved.runtime_name(), "memory_fetch");
1020 }
1021
1022 #[test]
1023 fn tool_schema_disambiguates_sanitized_name_collisions() {
1024 let registry =
1025 ToolRegistry::from_tools(vec![Arc::new(NamedTool("a:b")), Arc::new(NamedTool("a.b"))]);
1026 let names: HashSet<String> = registry
1027 .tools_schema()
1028 .iter()
1029 .filter_map(|s| s["name"].as_str().map(str::to_string))
1030 .collect();
1031
1032 assert_eq!(names.len(), 2);
1033 assert!(names.contains("a_b"));
1034 assert!(names.contains("a_b_2"));
1035 assert!(registry.get("a_b").is_some());
1036 assert!(registry.get("a_b_2").is_some());
1037 }
1038
1039 #[test]
1040 fn tool_schema_truncates_long_names_to_anthropic_limit() {
1041 let long = "x".repeat(140);
1042 let leaked: &'static str = Box::leak(long.into_boxed_str());
1043 let registry = ToolRegistry::from_tools(vec![Arc::new(NamedTool(leaked))]);
1044 let schema = registry.tools_schema();
1045 let name = schema[0]["name"].as_str().unwrap();
1046
1047 assert_eq!(name.len(), 128);
1048 assert!(name
1049 .chars()
1050 .all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-'));
1051 assert!(registry.get(name).is_some());
1052 }
1053
1054 #[test]
1055 fn tool_schema_sanitizes_input_property_names_and_translates_inputs_back() {
1056 let registry = ToolRegistry::from_tools(vec![Arc::new(SchemaTool)]);
1057 let schema = registry.tools_schema();
1058 let input_schema = &schema[0]["input_schema"];
1059 let props = input_schema["properties"].as_object().unwrap();
1060
1061 assert!(props.keys().all(|k| k.len() <= 64
1062 && k.chars()
1063 .all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-' || c == '.')));
1064 assert_eq!(
1065 input_schema["required"].as_array().unwrap()[0]
1066 .as_str()
1067 .unwrap()
1068 .len(),
1069 64
1070 );
1071 assert_eq!(input_schema["required"][1], "nested_obj");
1072 assert!(props["nested_obj"]["properties"]
1073 .as_object()
1074 .unwrap()
1075 .contains_key("inner_key"));
1076 assert_eq!(props["nested_obj"]["required"][0], "inner_key");
1077
1078 let first_required = input_schema["required"][0].as_str().unwrap();
1079 let translated = registry.translate_input_for_api_tool(
1080 "schema_tool",
1081 json!({
1082 first_required: "value",
1083 "nested_obj": {"inner_key": "nested"}
1084 }),
1085 );
1086
1087 assert_eq!(
1088 translated["bad:key/that/is/far/too/long/for/anthropic/property/names/and/keeps/going"],
1089 "value"
1090 );
1091 assert_eq!(translated["nested:obj"]["inner/key"], "nested");
1092 }
1093
1094 #[cfg(feature = "command-exec")]
1095 #[test]
1096 fn test_tool_registry_new() {
1097 let registry = ToolRegistry::new();
1098
1099 assert_eq!(
1103 registry.tools_schema().len(),
1104 29 + usize::from(cfg!(windows))
1105 );
1106 for name in ["forum_post", "forum_read", "forum_forget"] {
1107 assert!(registry.get(name).is_some(), "missing {name}");
1108 }
1109
1110 assert!(registry.get("bash").is_some());
1112
1113 assert!(registry.get("nonexistent").is_none());
1115
1116 assert!(registry.get("bash").is_some());
1118 assert!(registry.get("read").is_some());
1119 assert!(registry.get("write").is_some());
1120 assert!(registry.get("edit").is_some());
1121 assert!(registry.get("grep").is_some());
1122 assert!(registry.get("find").is_some());
1123 assert!(registry.get("ls").is_some());
1124 assert!(registry.get("subagent").is_some());
1125 assert!(registry.get("subagent_model_authorize").is_some());
1126 assert!(registry.get("subagent_models").is_some());
1127 }
1128
1129 #[cfg(feature = "command-exec")]
1130 #[test]
1131 fn test_disable_removes_tools_from_registry_and_schema() {
1132 let mut registry = ToolRegistry::new();
1133 let before = registry.tools_schema().len();
1134
1135 registry.disable(&["bash".to_string(), "ls".to_string()]);
1136
1137 assert!(registry.get("bash").is_none());
1139 assert!(registry.get("ls").is_none());
1140 assert!(registry.get("read").is_some());
1142 assert!(registry.get("subagent").is_some());
1143 assert_eq!(registry.tools_schema().len(), before - 2);
1145 }
1146
1147 #[cfg(feature = "command-exec")]
1148 #[test]
1149 fn test_disable_empty_and_unknown_are_noops() {
1150 let mut registry = ToolRegistry::new();
1151 let n = registry.tools_schema().len();
1152 registry.disable(&[]); assert_eq!(registry.tools_schema().len(), n);
1154 registry.disable(&["not_a_real_tool".to_string()]); assert_eq!(registry.tools_schema().len(), n);
1156 assert!(registry.get("bash").is_some());
1157 }
1158
1159 #[cfg(feature = "command-exec")]
1160 #[test]
1161 fn test_tool_registry_without_subagent() {
1162 let registry = ToolRegistry::without_subagent();
1163
1164 assert_eq!(
1166 registry.tools_schema().len(),
1167 14 + usize::from(cfg!(windows))
1168 );
1169 for name in ["forum_post", "forum_read", "forum_forget"] {
1170 assert!(registry.get(name).is_some(), "missing {name}");
1171 }
1172 assert!(registry.get("subagent").is_none());
1174
1175 assert!(registry.get("bash").is_some());
1177
1178 assert!(registry.get("bash").is_some());
1180 assert!(registry.get("read").is_some());
1181 assert!(registry.get("write").is_some());
1182 assert!(registry.get("edit").is_some());
1183 assert!(registry.get("grep").is_some());
1184 assert!(registry.get("find").is_some());
1185 assert!(registry.get("ls").is_some());
1186 }
1187
1188 #[cfg(not(feature = "command-exec"))]
1193 #[test]
1194 fn hardened_build_excludes_command_exec_tools() {
1195 let registry = ToolRegistry::new();
1196 assert!(registry.get("bash").is_none(), "bash must be absent in a hardened build");
1197 assert!(registry.get("shell_start").is_none());
1198 assert!(registry.get("shell_send").is_none());
1199 assert!(registry.get("shell_end").is_none());
1200 assert!(registry.get("read").is_some());
1202 assert!(registry.get("ls").is_some());
1203 assert!(registry.get("subagent").is_some());
1204
1205 let without_sub = ToolRegistry::without_subagent();
1206 assert!(without_sub.get("bash").is_none());
1207 assert!(without_sub.get("shell_start").is_none());
1208 assert!(without_sub.get("read").is_some());
1209 }
1210
1211 #[test]
1212 fn test_tool_registry_register() {
1213 let mut registry = ToolRegistry::without_subagent();
1214 let initial_count = registry.tools_schema().len();
1215
1216 struct TestTool;
1218 #[async_trait::async_trait]
1219 impl Tool for TestTool {
1220 fn name(&self) -> &str {
1221 "test_tool"
1222 }
1223 fn description(&self) -> &str {
1224 "A test tool"
1225 }
1226 fn parameters(&self) -> Value {
1227 json!({"type": "object"})
1228 }
1229 async fn execute(&self, _params: Value, _ctx: ToolContext) -> Result<String> {
1230 Ok("test result".to_string())
1231 }
1232 }
1233
1234 registry.register(Arc::new(TestTool));
1235
1236 assert_eq!(registry.tools_schema().len(), initial_count + 1);
1238
1239 assert!(registry.get("test_tool").is_some());
1241 }
1242
1243 #[test]
1244 fn tool_names_for_extension_filters_by_owner_and_sorts() {
1245 struct OwnedTool(&'static str, Option<&'static str>);
1246 #[async_trait::async_trait]
1247 impl Tool for OwnedTool {
1248 fn name(&self) -> &str {
1249 self.0
1250 }
1251 fn description(&self) -> &str {
1252 "owned"
1253 }
1254 fn parameters(&self) -> Value {
1255 json!({"type": "object"})
1256 }
1257 async fn execute(&self, _params: Value, _ctx: ToolContext) -> Result<String> {
1258 Ok("ok".to_string())
1259 }
1260 fn extension_id(&self) -> Option<&str> {
1261 self.1
1262 }
1263 }
1264
1265 let mut registry = ToolRegistry::without_subagent();
1266 registry.register(Arc::new(OwnedTool("alpha:zed", Some("alpha"))));
1267 registry.register(Arc::new(OwnedTool("alpha:bar", Some("alpha"))));
1268 registry.register(Arc::new(OwnedTool("beta:thing", Some("beta"))));
1269
1270 assert_eq!(
1271 registry.tool_names_for_extension("alpha"),
1272 vec!["alpha:bar".to_string(), "alpha:zed".to_string()]
1273 );
1274 assert_eq!(
1275 registry.tool_names_for_extension("beta"),
1276 vec!["beta:thing".to_string()]
1277 );
1278 assert!(registry.tool_names_for_extension("ghost").is_empty());
1279 assert!(registry.tool_names_for_extension("bash").is_empty());
1281 }
1282
1283 struct OwnedTool(&'static str, Option<&'static str>);
1284 #[async_trait::async_trait]
1285 impl Tool for OwnedTool {
1286 fn name(&self) -> &str {
1287 self.0
1288 }
1289 fn description(&self) -> &str {
1290 "owned"
1291 }
1292 fn parameters(&self) -> Value {
1293 json!({"type": "object"})
1294 }
1295 async fn execute(&self, _params: Value, _ctx: ToolContext) -> Result<String> {
1296 Ok("ok".to_string())
1297 }
1298 fn extension_id(&self) -> Option<&str> {
1299 self.1
1300 }
1301 }
1302
1303 #[test]
1304 fn without_subagent_excludes_subagent_tools() {
1305 let registry = ToolRegistry::without_subagent();
1306 assert!(registry.get("subagent").is_none());
1307 assert!(registry.get("subagent_start").is_none());
1308 assert!(registry.get("subagent_status").is_none());
1309 assert!(registry.get("subagent_steer").is_none());
1310 assert!(registry.get("subagent_collect").is_none());
1311 assert!(registry.get("subagent_resume").is_none());
1312 assert!(registry.get("ls").is_some());
1314 assert!(registry.get("read").is_some());
1315 }
1316
1317 #[test]
1318 fn without_subagent_with_extensions_includes_extension_tools() {
1319 let mut other = ToolRegistry::empty();
1320 other.register(Arc::new(OwnedTool("alpha:do_thing", Some("alpha"))));
1321
1322 let merged = ToolRegistry::without_subagent_with_extensions(&other);
1323
1324 assert!(merged.get("alpha:do_thing").is_some());
1326 assert!(merged.get("read").is_some());
1328 assert!(merged.get("subagent_start").is_none());
1330 }
1331
1332 #[test]
1337 fn worker_set_disables_the_listed_built_ins_before_the_extension_merge() {
1338 let mut shared = ToolRegistry::empty();
1339 shared.register(Arc::new(OwnedTool("alpha:do_thing", Some("alpha"))));
1340 let disabled: Vec<String> = ["bash", "shell_start", "ls", "forum_read", "alpha:do_thing"]
1341 .map(String::from)
1342 .to_vec();
1343
1344 let all = ToolRegistry::worker_set(&shared, &[]);
1345 assert!(all.get("ls").is_some() && all.get("forum_read").is_some());
1346 let worker = ToolRegistry::worker_set(&shared, &disabled);
1347 let names: Vec<String> = worker
1348 .tools_schema()
1349 .iter()
1350 .filter_map(|t| t["name"].as_str().map(String::from))
1351 .collect();
1352 for name in ["bash", "shell_start", "ls", "forum_read"] {
1353 assert!(worker.get(name).is_none(), "{name} kept");
1354 assert!(!names.iter().any(|n| n == name), "{name} advertised: {names:?}");
1355 }
1356 assert!(worker.get("alpha:do_thing").is_some(), "extension tool kept");
1357 assert!(worker.get("read").is_some() && worker.get("subagent_start").is_none());
1358 assert_eq!(names.len(), worker.tools.len());
1359
1360 let mut foreground = ToolRegistry::new();
1362 foreground.disable(&disabled);
1363 for tool in shared.tools.values() {
1364 foreground.insert_tool_with_catalog(tool.clone()).unwrap();
1365 }
1366 for name in worker.tools.keys() {
1367 assert!(foreground.get(name).is_some(), "worker has {name}, the foreground does not");
1368 }
1369 }
1370
1371 #[tokio::test]
1375 async fn disabled_multi_edit_is_gone_from_workers_too() {
1376 let disabled = vec!["multi_edit".to_string()];
1377 let worker = ToolRegistry::worker_set(&ToolRegistry::empty(), &disabled);
1378 let legacy = crate::tools::subagent::subagent_tools(&disabled).await;
1379 for registry in [worker, legacy] {
1380 assert!(registry.get("multi_edit").is_none());
1381 assert!(!registry.tools_schema().iter().any(|t| t["name"] == "multi_edit"));
1382 assert!(registry.get("edit").is_some() && registry.get("write").is_some());
1383 }
1384 assert!(ToolRegistry::worker_set(&ToolRegistry::empty(), &[]).get("multi_edit").is_some());
1385 }
1386
1387 #[cfg(feature = "command-exec")]
1388 #[test]
1389 fn forum_membership_preserves_worker_bounds_and_foreground_disclosure() {
1390 use crate::tools::activation::{SessionId, SessionToolSet};
1391 use crate::tools::catalog::{ToolEffect, ToolId};
1392 let main = ToolRegistry::new();
1393 let mut extensions = ToolRegistry::empty();
1394 extensions.register(Arc::new(OwnedTool("peer:probe", Some("peer"))));
1395 let merged = ToolRegistry::without_subagent_with_extensions(&extensions);
1396 assert_eq!(merged.tools_schema().len(), 15 + usize::from(cfg!(windows)));
1397 assert!(merged.get("peer:probe").is_some());
1398 for registry in [&main, &ToolRegistry::without_subagent(), &merged] {
1399 for (name, effect) in [
1400 ("forum_post", ToolEffect::NonIdempotent),
1401 ("forum_read", ToolEffect::ReadOnly),
1402 ("forum_forget", ToolEffect::NonIdempotent),
1403 ] {
1404 let tool = registry.get(name).unwrap();
1405 assert_eq!(tool.origin(), crate::tools::ToolOrigin::Builtin);
1406 assert_eq!(tool.effect(), effect);
1407 assert!(registry.catalog().get(&ToolId::builtin(name)).is_some());
1408 assert_eq!(
1409 registry
1410 .tools_schema()
1411 .iter()
1412 .filter(|schema| schema["name"] == name)
1413 .count(),
1414 1
1415 );
1416 }
1417 }
1418 for registry in [ToolRegistry::without_subagent(), merged] {
1419 for name in [
1420 "subagent",
1421 "subagent_start",
1422 "subagent_status",
1423 "subagent_steer",
1424 "subagent_collect",
1425 "subagent_resume",
1426 "subagent_model_authorize",
1427 "subagent_models",
1428 "search_tools",
1429 "activate_tools",
1430 "memory_context",
1431 "memory_store",
1432 "memory_search",
1433 "memory_fetch",
1434 "memory_forget",
1435 ] {
1436 assert!(registry.get(name).is_none(), "worker must not gain {name}");
1437 }
1438 for name in ["bash", "write", "edit", "multi_edit"] {
1439 assert!(registry.get(name).is_some(), "worker must retain {name}");
1440 }
1441 }
1442 let session = SessionToolSet::progressive_core_for_catalog(
1443 SessionId::parse("forum-foreground-test").unwrap(),
1444 main.catalog(),
1445 );
1446 let projection = main.session_tools_schema(&session);
1447 for name in ["forum_post", "forum_read", "forum_forget"] {
1448 assert!(!session.is_core(&ToolId::builtin(name)));
1449 assert!(!projection
1450 .schema
1451 .iter()
1452 .any(|schema| schema["name"] == name));
1453 }
1454 }
1455
1456 #[test]
1457 fn without_subagent_with_extensions_excludes_built_ins_from_other_registry() {
1458 let other = ToolRegistry::new();
1463
1464 let merged = ToolRegistry::without_subagent_with_extensions(&other);
1465
1466 let read = merged.get("read").expect("read present");
1468 assert!(read.extension_id().is_none());
1469 assert!(merged.get("subagent_start").is_none());
1471 assert!(merged.get("subagent").is_none());
1472 }
1473
1474 #[test]
1475 fn without_subagent_with_extensions_rejects_recursive_tool_name_collisions() {
1476 let mut other = ToolRegistry::empty();
1477 for name in [
1478 "subagent",
1479 "subagent_start",
1480 "subagent_status",
1481 "subagent_steer",
1482 "subagent_collect",
1483 "subagent_resume",
1484 "subagent_model_authorize",
1485 "subagent_models",
1486 "memory_context",
1487 ] {
1488 other.register(Arc::new(OwnedTool(name, Some("malicious-extension"))));
1489 }
1490
1491 let merged = ToolRegistry::without_subagent_with_extensions(&other);
1492 for name in [
1493 "subagent",
1494 "subagent_start",
1495 "subagent_status",
1496 "subagent_steer",
1497 "subagent_collect",
1498 "subagent_resume",
1499 "subagent_model_authorize",
1500 "subagent_models",
1501 "memory_context",
1502 ] {
1503 assert!(merged.get(name).is_none(), "{name} must stay unavailable");
1504 }
1505 }
1506
1507 #[test]
1508 fn without_subagent_with_extensions_rejects_api_sanitized_recursive_collisions() {
1509 let mut other = ToolRegistry::empty();
1510 for runtime_name in [
1511 "subagent:start",
1512 "subagent:status",
1513 "subagent:steer",
1514 "subagent:collect",
1515 "subagent:resume",
1516 "subagent:model:authorize",
1517 "subagent:models",
1518 "memory:context",
1519 ] {
1520 other.register(Arc::new(OwnedTool(runtime_name, Some("subagent"))));
1521 }
1522
1523 let merged = ToolRegistry::without_subagent_with_extensions(&other);
1524 for api_name in [
1525 "subagent_start",
1526 "subagent_status",
1527 "subagent_steer",
1528 "subagent_collect",
1529 "subagent_resume",
1530 "subagent_model_authorize",
1531 "subagent_models",
1532 "memory_context",
1533 ] {
1534 assert!(
1535 merged.get(api_name).is_none(),
1536 "API-safe recursive name {api_name} must stay unavailable"
1537 );
1538 }
1539 }
1540
1541 #[test]
1545 fn memory_context_present_in_primary_registry_absent_in_recursive_ones() {
1546 assert!(ToolRegistry::new().get("memory_context").is_some());
1547 assert!(ToolRegistry::without_subagent()
1548 .get("memory_context")
1549 .is_none());
1550
1551 let merged = ToolRegistry::without_subagent_with_extensions(&ToolRegistry::new());
1552 assert!(
1553 merged.get("memory_context").is_none(),
1554 "recursive registries must never expose memory_context"
1555 );
1556 }
1557
1558 #[test]
1559 fn without_subagent_with_extensions_does_not_overwrite_existing_builtin() {
1560 let mut other = ToolRegistry::empty();
1566 other.register(Arc::new(OwnedTool("ext:custom", Some("ext"))));
1567
1568 let merged = ToolRegistry::without_subagent_with_extensions(&other);
1569 assert!(merged.get("ext:custom").is_some());
1570 assert!(merged.get("read").is_some());
1571 assert!(merged.get("read").unwrap().extension_id().is_none());
1572 }
1573
1574 struct McpTwin(&'static str);
1576
1577 #[async_trait::async_trait]
1578 impl Tool for McpTwin {
1579 fn name(&self) -> &str {
1580 self.0
1581 }
1582 fn description(&self) -> &str {
1583 "mcp twin"
1584 }
1585 fn parameters(&self) -> Value {
1586 json!({"type": "object"})
1587 }
1588 async fn execute(&self, _params: Value, _ctx: ToolContext) -> Result<String> {
1589 Ok("ok".to_string())
1590 }
1591 fn origin(&self) -> crate::tools::ToolOrigin {
1592 crate::tools::ToolOrigin::Mcp {
1593 server_id: "srv".to_string(),
1594 server_tool_name: "shared".to_string(),
1595 }
1596 }
1597 }
1598
1599 #[test]
1600 #[should_panic(expected = "tool catalog rejected mutation")]
1601 fn from_tools_with_duplicate_capability_identity_fails_closed() {
1602 let _ = ToolRegistry::from_tools(vec![
1606 Arc::new(McpTwin("twin_one")),
1607 Arc::new(McpTwin("twin_two")),
1608 ]);
1609 }
1610}