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::discovery::SearchToolsTool),
65 Arc::new(crate::tools::discovery::ActivateToolsTool),
66 Arc::new(crate::tools::memory::MemorySearchTool),
70 Arc::new(crate::tools::memory::MemoryFetchTool),
71 Arc::new(crate::tools::memory::MemoryStoreTool),
72 Arc::new(crate::tools::memory::MemoryForgetTool),
73 Arc::new(crate::tools::context_checkpoint::ContextCheckpointTool {}),
81 ];
82 #[cfg(feature = "command-exec")]
86 {
87 tools.push(Arc::new(crate::tools::bash::BashTool));
88 tools.push(Arc::new(crate::tools::shell::ShellStartTool));
89 tools.push(Arc::new(crate::tools::shell::ShellSendTool));
90 tools.push(Arc::new(crate::tools::shell::ShellEndTool));
91 }
92 Self::from_tools(tools)
93 }
94
95 pub fn empty() -> Self {
98 Self::from_tools(Vec::new())
99 }
100
101 pub fn disable(&mut self, names: &[String]) {
107 self.try_disable(names).expect(CATALOG_REJECTED_MSG);
108 }
109
110 pub fn try_disable(&mut self, names: &[String]) -> Result<(), CatalogError> {
123 if names.is_empty() {
124 return Ok(());
125 }
126 let remaining: Vec<Arc<dyn Tool>> = self
127 .tools
128 .values()
129 .filter(|t| !names.iter().any(|n| n == t.name()))
130 .cloned()
131 .collect();
132 if remaining.len() == self.tools.len() {
133 return Ok(());
134 }
135 let prior = self.catalog.generation();
136 let mut candidate = Self::try_from_tools(remaining)?;
137 candidate.catalog.rebase_past(prior)?;
138 *self = candidate;
139 Ok(())
140 }
141
142 pub fn without_subagent() -> Self {
144 #[allow(unused_mut)]
145 let mut tools: Vec<Arc<dyn Tool>> = vec![
146 #[cfg(all(windows, feature = "command-exec"))]
147 Arc::new(crate::tools::powershell::PowerShellTool),
148 Arc::new(crate::tools::read::ReadTool),
149 Arc::new(crate::tools::write::WriteTool),
150 Arc::new(crate::tools::edit::EditTool),
151 Arc::new(crate::tools::multi_edit::MultiEditTool),
152 Arc::new(crate::tools::grep::GrepTool),
153 Arc::new(crate::tools::find::FindTool),
154 Arc::new(crate::tools::ls::LsTool),
155 ];
156 #[cfg(feature = "command-exec")]
157 {
158 tools.push(Arc::new(crate::tools::bash::BashTool));
159 tools.push(Arc::new(crate::tools::shell::ShellStartTool));
160 tools.push(Arc::new(crate::tools::shell::ShellSendTool));
161 tools.push(Arc::new(crate::tools::shell::ShellEndTool));
162 }
163 Self::from_tools(tools)
164 }
165
166 pub fn without_subagent_with_extensions(extension_tools: &ToolRegistry) -> Self {
174 Self::worker_set(extension_tools, &[])
175 }
176
177 pub fn worker_set(extension_tools: &ToolRegistry, disabled: &[String]) -> Self {
184 let mut combined = Self::without_subagent();
185 combined.disable(disabled);
186 for tool in extension_tools.tools.values() {
187 if tool.extension_id().is_some() && !is_recursive_subagent_tool_name(tool.name()) {
188 combined
189 .insert_tool_with_catalog(tool.clone())
190 .expect(CATALOG_REJECTED_MSG);
191 }
192 }
193 combined.rebuild_schema();
194 combined
195 }
196
197 fn from_tools(tool_list: Vec<Arc<dyn Tool>>) -> Self {
203 Self::try_from_tools(tool_list).expect(CATALOG_REJECTED_MSG)
204 }
205
206 #[doc(hidden)]
210 pub fn from_tools_for_tests(tool_list: Vec<Arc<dyn Tool>>) -> Self {
211 Self::from_tools(tool_list)
212 }
213
214 fn try_from_tools(tool_list: Vec<Arc<dyn Tool>>) -> Result<Self, CatalogError> {
215 let mut registry = ToolRegistry {
216 tools: HashMap::new(),
217 cached_schema: Arc::new(Vec::new()),
218 api_to_runtime_names: HashMap::new(),
219 input_name_maps: HashMap::new(),
220 catalog: ToolCatalog::empty(),
221 };
222 for tool in tool_list {
226 let name = tool.name().to_string();
227 registry.tools.insert(name, tool);
228 }
229 let sorted: Vec<Arc<dyn Tool>> =
235 registry.iter_tools_sorted().into_iter().cloned().collect();
236 for tool in sorted {
237 let record = CapabilityRecord::for_registered_tool(&tool);
238 registry.catalog.upsert(None, record)?;
239 }
240 registry.rebuild_schema();
241 Ok(registry)
242 }
243
244 fn api_safe_name(name: &str, used: &HashSet<String>) -> String {
245 Self::api_safe_identifier(name, used, 128, false)
246 }
247
248 fn api_safe_property_name(name: &str, used: &HashSet<String>) -> String {
249 Self::api_safe_identifier(name, used, 64, true)
250 }
251
252 fn api_safe_identifier(
253 name: &str,
254 used: &HashSet<String>,
255 max_len: usize,
256 allow_dot: bool,
257 ) -> String {
258 let mut sanitized = String::with_capacity(name.len());
259 for ch in name.chars() {
260 if ch.is_ascii_alphanumeric() || ch == '_' || ch == '-' || (allow_dot && ch == '.') {
261 sanitized.push(ch);
262 } else {
263 sanitized.push('_');
264 }
265 }
266 if sanitized.is_empty() {
267 sanitized.push_str("field");
268 }
269 if sanitized.len() > max_len {
270 sanitized.truncate(max_len);
271 }
272
273 let base = sanitized.clone();
274 let mut suffix = 2;
275 while used.contains(&sanitized) {
276 let suffix_str = format!("_{suffix}");
277 let keep = max_len.saturating_sub(suffix_str.len());
278 sanitized = format!("{}{}", &base[..base.len().min(keep)], suffix_str);
279 suffix += 1;
280 }
281 sanitized
282 }
283
284 fn sanitize_schema(mut schema: Value) -> (Value, SchemaNameMap) {
285 let mut map = SchemaNameMap::default();
286 let Some(obj) = schema.as_object_mut() else {
287 return (schema, map);
288 };
289
290 let mut required_name_map = HashMap::new();
291 if let Some(props_value) = obj.get_mut("properties") {
292 if let Some(props) = props_value.as_object_mut() {
293 let original = std::mem::take(props);
294 let mut used = HashSet::new();
295 for (runtime_name, child_schema) in original {
296 let api_name = Self::api_safe_property_name(&runtime_name, &used);
297 used.insert(api_name.clone());
298 required_name_map.insert(runtime_name.clone(), api_name.clone());
299 map.api_to_runtime.insert(api_name.clone(), runtime_name);
300
301 let (sanitized_child, child_map) = Self::sanitize_schema(child_schema);
302 if !child_map.api_to_runtime.is_empty() || !child_map.children.is_empty() {
303 map.children.insert(api_name.clone(), child_map);
304 }
305 props.insert(api_name, sanitized_child);
306 }
307 }
308 }
309
310 if let Some(required) = obj.get_mut("required").and_then(Value::as_array_mut) {
311 for item in required.iter_mut() {
312 if let Some(name) = item.as_str() {
313 if let Some(api_name) = required_name_map.get(name) {
314 *item = Value::String(api_name.clone());
315 }
316 }
317 }
318 }
319
320 if let Some(items) = obj.get_mut("items") {
323 let (sanitized_items, items_map) = Self::sanitize_schema(std::mem::take(items));
324 if !items_map.api_to_runtime.is_empty()
325 || !items_map.children.is_empty()
326 || items_map.items.is_some()
327 {
328 map.items = Some(Box::new(items_map));
329 }
330 *items = sanitized_items;
331 }
332
333 (schema, map)
334 }
335
336 fn translate_input_names(input: Value, map: &SchemaNameMap) -> Value {
337 match input {
338 Value::Object(obj) => {
339 let mut out = serde_json::Map::new();
340 for (api_name, value) in obj {
341 let runtime_name = map
342 .api_to_runtime
343 .get(&api_name)
344 .cloned()
345 .unwrap_or_else(|| api_name.clone());
346 let value = if let Some(child) = map.children.get(&api_name) {
347 Self::translate_input_names(value, child)
348 } else {
349 value
350 };
351 out.insert(runtime_name, value);
352 }
353 Value::Object(out)
354 }
355 Value::Array(arr) => {
356 if let Some(items_map) = &map.items {
358 Value::Array(
359 arr.into_iter()
360 .map(|v| Self::translate_input_names(v, items_map))
361 .collect(),
362 )
363 } else {
364 Value::Array(arr)
365 }
366 }
367 other => other,
368 }
369 }
370
371 fn rebuild_schema(&mut self) {
372 let mut used = HashSet::new();
373 let mut api_to_runtime_names = HashMap::new();
374 let mut input_name_maps = HashMap::new();
375 let mut schema = Vec::with_capacity(self.tools.len());
376
377 let mut sorted_tools: Vec<_> = self.tools.values().collect();
392 sorted_tools.sort_by_key(|t| {
393 let builtin_first = match t.origin() {
394 crate::tools::ToolOrigin::Builtin => 0u8,
395 _ => 1u8,
396 };
397 (builtin_first, t.name().to_string())
398 });
399
400 for tool in sorted_tools {
401 let runtime_name = tool.name();
402 let api_name = Self::api_safe_name(runtime_name, &used);
403 used.insert(api_name.clone());
404 api_to_runtime_names.insert(api_name.clone(), runtime_name.to_string());
405 let (input_schema, input_map) = Self::sanitize_schema(tool.parameters());
406 input_name_maps.insert(api_name.clone(), input_map);
407 schema.push(serde_json::json!({
408 "name": api_name,
409 "description": tool.description(),
410 "input_schema": input_schema
411 }));
412 }
413
414 self.api_to_runtime_names = api_to_runtime_names;
415 self.input_name_maps = input_name_maps;
416 self.cached_schema = Arc::new(schema);
417 }
418
419 pub fn register(&mut self, tool: Arc<dyn Tool>) {
429 self.try_register(tool).expect(CATALOG_REJECTED_MSG);
430 }
431
432 pub fn try_register(&mut self, tool: Arc<dyn Tool>) -> Result<(), CatalogError> {
438 self.insert_tool_with_catalog(tool)?;
439 self.rebuild_schema();
440 Ok(())
441 }
442
443 pub fn try_register_batch(
457 &mut self,
458 tool_list: Vec<Arc<dyn Tool>>,
459 ) -> Result<usize, CatalogError> {
460 if tool_list.is_empty() {
461 return Ok(0);
462 }
463 let mut batch_names: HashSet<String> = HashSet::new();
464 let mut batch_ids: HashSet<crate::tools::catalog::ToolId> = HashSet::new();
465 for tool in &tool_list {
466 let name = tool.name();
467 if self.tools.contains_key(name) || !batch_names.insert(name.to_string()) {
468 return Err(CatalogError::DuplicateRuntimeName(
469 agent_core::BoundedText::new(name, 128).text,
470 ));
471 }
472 let id = tool_id_for(tool.as_ref());
473 if self.catalog.get(&id).is_some() || !batch_ids.insert(id.clone()) {
474 return Err(CatalogError::DuplicateToolId(id));
475 }
476 }
477 let count = tool_list.len();
478 let mut candidate = self.clone();
479 for tool in tool_list {
480 candidate.insert_tool_with_catalog(tool)?;
481 }
482 candidate.rebuild_schema();
483 *self = candidate;
484 Ok(count)
485 }
486
487 fn insert_tool_with_catalog(&mut self, tool: Arc<dyn Tool>) -> Result<(), CatalogError> {
496 let record = CapabilityRecord::for_registered_tool(&tool);
497 let replaced = self
498 .tools
499 .get(tool.name())
500 .map(|old| tool_id_for(old.as_ref()));
501 self.catalog.upsert(replaced.as_ref(), record)?;
502 self.tools.insert(tool.name().to_string(), tool);
503 Ok(())
504 }
505
506 pub fn catalog(&self) -> &ToolCatalog {
509 &self.catalog
510 }
511
512 #[doc(hidden)]
520 pub fn resume_catalog_generation_for_tests(
521 &mut self,
522 generation: crate::tools::catalog::CatalogGeneration,
523 ) {
524 self.catalog.set_generation_for_tests(generation);
525 }
526
527 #[doc(hidden)]
533 pub fn upsert_catalog_record_for_tests(
534 &mut self,
535 replaced: Option<&crate::tools::catalog::ToolId>,
536 record: crate::tools::catalog::CapabilityRecord,
537 ) {
538 self.catalog
539 .upsert(replaced, record)
540 .expect("test upsert must be structurally valid");
541 }
542
543 pub fn get(&self, name: &str) -> Option<&Arc<dyn Tool>> {
544 let runtime_name = self
545 .api_to_runtime_names
546 .get(name)
547 .map(String::as_str)
548 .unwrap_or(name);
549 self.tools.get(runtime_name)
550 }
551
552 pub fn runtime_name_for_api<'a>(&'a self, name: &'a str) -> &'a str {
553 self.api_to_runtime_names
554 .get(name)
555 .map(String::as_str)
556 .unwrap_or(name)
557 }
558
559 pub fn resolve_wire_call(
569 &self,
570 wire_name: &str,
571 ) -> Option<crate::tools::activation::ResolvedToolCall> {
572 let runtime_name = self
573 .api_to_runtime_names
574 .get(wire_name)
575 .map(String::as_str)
576 .unwrap_or(wire_name);
577 let tool = self.tools.get(runtime_name)?;
578 Some(crate::tools::activation::ResolvedToolCall::new(
579 wire_name,
580 runtime_name,
581 tool_id_for(tool.as_ref()),
582 ))
583 }
584
585 pub fn translate_input_for_api_tool(&self, tool_name: &str, input: Value) -> Value {
586 if let Some(map) = self.input_name_maps.get(tool_name) {
587 Self::translate_input_names(input, map)
588 } else {
589 input
590 }
591 }
592
593 pub fn tools_schema(&self) -> Arc<Vec<Value>> {
594 Arc::clone(&self.cached_schema)
595 }
596
597 pub fn tools_schema_without<'a>(
605 &self,
606 withheld: impl IntoIterator<Item = &'a crate::tools::catalog::ToolId>,
607 ) -> Option<Arc<Vec<Value>>> {
608 let withheld: HashSet<&crate::tools::catalog::ToolId> = withheld.into_iter().collect();
609 if !withheld.iter().any(|id| self.catalog.get(id).is_some()) {
610 return None;
611 }
612 let offered = |entry: &&Value| {
613 let tool = entry["name"]
614 .as_str()
615 .and_then(|api_name| self.tools.get(self.runtime_name_for_api(api_name)));
616 !tool.is_some_and(|tool| withheld.contains(&tool_id_for(tool.as_ref())))
617 };
618 Some(Arc::new(
619 self.cached_schema.iter().filter(offered).cloned().collect(),
620 ))
621 }
622
623 pub fn deferred_tool_split(
659 &self,
660 session: &crate::tools::activation::SessionToolSet,
661 referenced_ids: &HashSet<String>,
662 ) -> DeferredToolSplit {
663 let projection = self.session_tools_schema(session).schema;
664 let projected: HashSet<&str> = projection
665 .iter()
666 .filter_map(|entry| entry["name"].as_str())
667 .collect();
668 let mut split = DeferredToolSplit::default();
669 for entry in self.cached_schema.iter() {
670 let Some(api_name) = entry["name"].as_str() else {
671 continue;
672 };
673 let Some(tool) = self.tools.get(self.runtime_name_for_api(api_name)) else {
674 continue;
675 };
676 let id = tool_id_for(tool.as_ref());
677 if projected.contains(api_name) {
678 if session.is_core(&id) || !referenced_ids.contains(id.as_str()) {
679 split.loaded.push(entry.clone());
680 } else {
681 split.deferred.push(deferred_entry(entry));
682 }
683 } else if !session.is_withheld(&id)
684 && self
685 .catalog
686 .get(&id)
687 .is_some_and(crate::tools::activation::is_activatable_record)
688 {
689 split.deferred.push(deferred_entry(entry));
690 } else {
691 continue;
692 }
693 split
694 .id_to_api
695 .insert(id.as_str().to_string(), api_name.to_string());
696 }
697 split
698 }
699
700 pub fn session_tools_schema(
701 &self,
702 session: &crate::tools::activation::SessionToolSet,
703 ) -> SessionSchemaProjection {
704 let mut projected = Vec::new();
705 let mut dropped = Vec::new();
706 for entry in self.cached_schema.iter() {
707 let Some(api_name) = entry["name"].as_str() else {
708 continue;
709 };
710 let runtime_name = self.runtime_name_for_api(api_name);
711 let Some(tool) = self.tools.get(runtime_name) else {
712 continue;
713 };
714 let id = tool_id_for(tool.as_ref());
715 let pinned = session
716 .core_pin(&id)
717 .map(|pin| (pin.schema_digest(), pin.provenance()))
718 .or_else(|| {
719 session
720 .activation(&id)
721 .map(|a| (a.schema_digest(), a.provenance()))
722 });
723 let Some((pinned_digest, pinned_provenance)) = pinned else {
724 continue; };
726 let Some(record) = self.catalog.get(&id) else {
727 tracing::warn!(
728 tool = %id,
729 "session schema projection drops uncataloged session member \
730 (removed since the session pinned it)"
731 );
732 dropped.push((id, DroppedSessionMember::Uncataloged));
733 continue;
734 };
735 if pinned_digest != record.schema_digest() {
736 tracing::warn!(
737 tool = %id,
738 "session schema projection drops drifted session member \
739 (schema digest changed since the session pinned it)"
740 );
741 dropped.push((id, DroppedSessionMember::SchemaDigestMismatch));
742 continue;
743 }
744 if pinned_provenance != record.provenance() {
745 tracing::warn!(
746 tool = %id,
747 "session schema projection drops drifted session member \
748 (trust provenance changed since the session pinned it)"
749 );
750 dropped.push((id, DroppedSessionMember::ProvenanceMismatch));
751 continue;
752 }
753 projected.push(entry.clone());
754 }
755 let seen: HashSet<&crate::tools::catalog::ToolId> =
762 dropped.iter().map(|(id, _)| id).collect();
763 let vanished: Vec<crate::tools::catalog::ToolId> = session
764 .core_ids()
765 .chain(session.activated().map(|a| a.grant().tool_id()))
766 .filter(|id| !seen.contains(id) && self.catalog.get(id).is_none())
767 .cloned()
768 .collect();
769 for id in vanished {
770 tracing::warn!(
771 tool = %id,
772 "session schema projection drops uncataloged session member (removed from the registry since the session pinned it)"
773 );
774 dropped.push((id, DroppedSessionMember::Uncataloged));
775 }
776 SessionSchemaProjection {
777 schema: projected,
778 dropped,
779 }
780 }
781
782 pub fn tool_names_for_extension(&self, extension_id: &str) -> Vec<String> {
785 let mut names: Vec<String> = self
786 .tools
787 .values()
788 .filter(|t| t.extension_id() == Some(extension_id))
789 .map(|t| t.name().to_string())
790 .collect();
791 names.sort();
792 names
793 }
794
795 pub fn iter_tools_sorted(&self) -> Vec<&Arc<dyn Tool>> {
798 let mut tools: Vec<_> = self.tools.values().collect();
799 tools.sort_by_key(|t| t.name());
800 tools
801 }
802}
803
804#[derive(Clone, Debug, PartialEq, Eq)]
810pub struct SessionSchemaProjection {
811 pub schema: Vec<Value>,
813 pub dropped: Vec<(crate::tools::catalog::ToolId, DroppedSessionMember)>,
817}
818
819#[derive(Clone, Copy, Debug, PartialEq, Eq)]
821pub enum DroppedSessionMember {
822 Uncataloged,
824 SchemaDigestMismatch,
826 ProvenanceMismatch,
829}
830
831const CATALOG_REJECTED_MSG: &str =
837 "tool catalog rejected mutation: refusing to expose a tool the catalog could not record";
838
839fn is_recursive_subagent_tool_name(name: &str) -> bool {
840 let api_name = ToolRegistry::api_safe_name(name, &HashSet::new());
841 matches!(
842 api_name.as_str(),
843 "subagent"
844 | "subagent_start"
845 | "subagent_status"
846 | "subagent_steer"
847 | "subagent_collect"
848 | "subagent_resume"
849 | "subagent_model_authorize"
850 | "subagent_models"
851 )
852}
853#[derive(Debug, Clone, Default)]
855pub struct DeferredToolSplit {
856 pub loaded: Vec<Value>,
858 pub deferred: Vec<Value>,
860 pub id_to_api: std::collections::HashMap<String, String>,
862}
863
864fn deferred_entry(entry: &Value) -> Value {
865 let mut entry = entry.clone();
866 if let Some(map) = entry.as_object_mut() {
867 map.remove("cache_control");
868 map.insert("defer_loading".to_string(), Value::Bool(true));
869 }
870 entry
871}
872
873#[cfg(test)]
874mod tests {
875 use super::*;
876 use crate::{Result, ToolContext};
877 use serde_json::json;
878
879 struct NamedTool(&'static str);
880
881 #[async_trait::async_trait]
882 impl Tool for NamedTool {
883 fn name(&self) -> &str {
884 self.0
885 }
886 fn description(&self) -> &str {
887 "test tool"
888 }
889 fn parameters(&self) -> Value {
890 json!({"type": "object"})
891 }
892 async fn execute(&self, _params: Value, _ctx: ToolContext) -> Result<String> {
893 Ok("ok".to_string())
894 }
895 }
896
897 struct SchemaTool;
898
899 #[async_trait::async_trait]
900 impl Tool for SchemaTool {
901 fn name(&self) -> &str {
902 "schema_tool"
903 }
904 fn description(&self) -> &str {
905 "schema tool"
906 }
907 fn parameters(&self) -> Value {
908 json!({
909 "type": "object",
910 "properties": {
911 "bad:key/that/is/far/too/long/for/anthropic/property/names/and/keeps/going": {"type": "string"},
912 "nested:obj": {
913 "type": "object",
914 "properties": {"inner/key": {"type": "string"}},
915 "required": ["inner/key"]
916 }
917 },
918 "required": [
919 "bad:key/that/is/far/too/long/for/anthropic/property/names/and/keeps/going",
920 "nested:obj"
921 ]
922 })
923 }
924 async fn execute(&self, _params: Value, _ctx: ToolContext) -> Result<String> {
925 Ok("ok".to_string())
926 }
927 }
928
929 #[test]
930 fn tool_schema_uses_api_safe_names_and_maps_back() {
931 let registry = ToolRegistry::from_tools(vec![Arc::new(NamedTool("plugin:skill.tool"))]);
932
933 assert_eq!(registry.tools_schema()[0]["name"], "plugin_skill_tool");
934 assert!(registry.get("plugin:skill.tool").is_some());
935 assert!(registry.get("plugin_skill_tool").is_some());
936 assert_eq!(
937 registry.runtime_name_for_api("plugin_skill_tool"),
938 "plugin:skill.tool"
939 );
940 }
941
942 #[test]
952 fn extension_leaf_name_never_shadows_builtin_memory_fetch_api_name() {
953 struct ExtMemoryTool {
954 runtime: &'static str,
955 ext: &'static str,
956 }
957 #[async_trait::async_trait]
958 impl Tool for ExtMemoryTool {
959 fn origin(&self) -> crate::tools::ToolOrigin {
960 crate::tools::ToolOrigin::Extension {
961 extension_id: self.ext.to_string(),
962 }
963 }
964 fn extension_id(&self) -> Option<&str> {
965 Some(self.ext)
966 }
967 fn name(&self) -> &str {
968 self.runtime
969 }
970 fn description(&self) -> &str {
971 "extension memory tool (single-id schema)"
972 }
973 fn parameters(&self) -> Value {
974 json!({
975 "type": "object",
976 "properties": {"id": {"type": "string"}},
977 "required": ["id"]
978 })
979 }
980 async fn execute(&self, _params: Value, _ctx: ToolContext) -> Result<String> {
981 Ok("ok".to_string())
982 }
983 }
984
985 let registry = ToolRegistry::from_tools(vec![
986 Arc::new(crate::tools::memory::MemoryFetchTool),
987 Arc::new(ExtMemoryTool {
989 runtime: "memory:fetch",
990 ext: "memory",
991 }),
992 Arc::new(ExtMemoryTool {
994 runtime: "axel-memory-manager:memory_fetch",
995 ext: "axel-memory-manager",
996 }),
997 ]);
998
999 assert_eq!(
1001 registry.runtime_name_for_api("memory_fetch"),
1002 "memory_fetch",
1003 "wire name `memory_fetch` must resolve to the built-in"
1004 );
1005 let schema = registry.tools_schema();
1006 let builtin = schema
1007 .iter()
1008 .find(|e| e["name"] == "memory_fetch")
1009 .expect("built-in memory_fetch exposed");
1010 assert_eq!(
1011 builtin["input_schema"]["required"],
1012 json!(["ids"]),
1013 "built-in memory_fetch schema must remain {{ids: [...]}}"
1014 );
1015
1016 let axel = schema
1018 .iter()
1019 .find(|e| e["name"] == "axel-memory-manager_memory_fetch")
1020 .expect("extension tool exposed under namespaced API name");
1021 assert_eq!(axel["input_schema"]["required"], json!(["id"]));
1022 assert_eq!(
1025 registry.runtime_name_for_api("memory_fetch_2"),
1026 "memory:fetch"
1027 );
1028 let resolved = registry.resolve_wire_call("memory_fetch").unwrap();
1031 assert_eq!(resolved.runtime_name(), "memory_fetch");
1032 }
1033
1034 #[test]
1035 fn tool_schema_disambiguates_sanitized_name_collisions() {
1036 let registry =
1037 ToolRegistry::from_tools(vec![Arc::new(NamedTool("a:b")), Arc::new(NamedTool("a.b"))]);
1038 let names: HashSet<String> = registry
1039 .tools_schema()
1040 .iter()
1041 .filter_map(|s| s["name"].as_str().map(str::to_string))
1042 .collect();
1043
1044 assert_eq!(names.len(), 2);
1045 assert!(names.contains("a_b"));
1046 assert!(names.contains("a_b_2"));
1047 assert!(registry.get("a_b").is_some());
1048 assert!(registry.get("a_b_2").is_some());
1049 }
1050
1051 #[test]
1052 fn tool_schema_truncates_long_names_to_anthropic_limit() {
1053 let long = "x".repeat(140);
1054 let leaked: &'static str = Box::leak(long.into_boxed_str());
1055 let registry = ToolRegistry::from_tools(vec![Arc::new(NamedTool(leaked))]);
1056 let schema = registry.tools_schema();
1057 let name = schema[0]["name"].as_str().unwrap();
1058
1059 assert_eq!(name.len(), 128);
1060 assert!(name
1061 .chars()
1062 .all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-'));
1063 assert!(registry.get(name).is_some());
1064 }
1065
1066 #[test]
1067 fn tool_schema_sanitizes_input_property_names_and_translates_inputs_back() {
1068 let registry = ToolRegistry::from_tools(vec![Arc::new(SchemaTool)]);
1069 let schema = registry.tools_schema();
1070 let input_schema = &schema[0]["input_schema"];
1071 let props = input_schema["properties"].as_object().unwrap();
1072
1073 assert!(props.keys().all(|k| k.len() <= 64
1074 && k.chars()
1075 .all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-' || c == '.')));
1076 assert_eq!(
1077 input_schema["required"].as_array().unwrap()[0]
1078 .as_str()
1079 .unwrap()
1080 .len(),
1081 64
1082 );
1083 assert_eq!(input_schema["required"][1], "nested_obj");
1084 assert!(props["nested_obj"]["properties"]
1085 .as_object()
1086 .unwrap()
1087 .contains_key("inner_key"));
1088 assert_eq!(props["nested_obj"]["required"][0], "inner_key");
1089
1090 let first_required = input_schema["required"][0].as_str().unwrap();
1091 let translated = registry.translate_input_for_api_tool(
1092 "schema_tool",
1093 json!({
1094 first_required: "value",
1095 "nested_obj": {"inner_key": "nested"}
1096 }),
1097 );
1098
1099 assert_eq!(
1100 translated["bad:key/that/is/far/too/long/for/anthropic/property/names/and/keeps/going"],
1101 "value"
1102 );
1103 assert_eq!(translated["nested:obj"]["inner/key"], "nested");
1104 }
1105
1106 #[cfg(feature = "command-exec")]
1107 #[test]
1108 fn test_tool_registry_new() {
1109 let registry = ToolRegistry::new();
1110
1111 assert_eq!(
1115 registry.tools_schema().len(),
1116 26 + usize::from(cfg!(windows))
1117 );
1118 assert!(registry.get("context_checkpoint").is_some());
1119 for name in ["forum_post", "forum_read", "forum_forget"] {
1120 assert!(registry.get(name).is_none(), "{name} is no longer built in");
1121 }
1122
1123 assert!(registry.get("bash").is_some());
1125
1126 assert!(registry.get("nonexistent").is_none());
1128
1129 assert!(registry.get("bash").is_some());
1131 assert!(registry.get("read").is_some());
1132 assert!(registry.get("write").is_some());
1133 assert!(registry.get("edit").is_some());
1134 assert!(registry.get("grep").is_some());
1135 assert!(registry.get("find").is_some());
1136 assert!(registry.get("ls").is_some());
1137 assert!(registry.get("subagent").is_some());
1138 assert!(registry.get("subagent_model_authorize").is_some());
1139 assert!(registry.get("subagent_models").is_some());
1140 }
1141
1142 #[cfg(feature = "command-exec")]
1143 #[test]
1144 fn test_disable_removes_tools_from_registry_and_schema() {
1145 let mut registry = ToolRegistry::new();
1146 let before = registry.tools_schema().len();
1147
1148 registry.disable(&["bash".to_string(), "ls".to_string()]);
1149
1150 assert!(registry.get("bash").is_none());
1152 assert!(registry.get("ls").is_none());
1153 assert!(registry.get("read").is_some());
1155 assert!(registry.get("subagent").is_some());
1156 assert_eq!(registry.tools_schema().len(), before - 2);
1158 }
1159
1160 #[cfg(feature = "command-exec")]
1161 #[test]
1162 fn test_disable_empty_and_unknown_are_noops() {
1163 let mut registry = ToolRegistry::new();
1164 let n = registry.tools_schema().len();
1165 registry.disable(&[]); assert_eq!(registry.tools_schema().len(), n);
1167 registry.disable(&["not_a_real_tool".to_string()]); assert_eq!(registry.tools_schema().len(), n);
1169 assert!(registry.get("bash").is_some());
1170 }
1171
1172 #[cfg(feature = "command-exec")]
1173 #[test]
1174 fn test_tool_registry_without_subagent() {
1175 let registry = ToolRegistry::without_subagent();
1176
1177 assert_eq!(
1179 registry.tools_schema().len(),
1180 11 + usize::from(cfg!(windows))
1181 );
1182 for name in ["forum_post", "forum_read", "forum_forget"] {
1183 assert!(registry.get(name).is_none(), "{name} is no longer built in");
1184 }
1185 assert!(registry.get("subagent").is_none());
1187
1188 assert!(registry.get("bash").is_some());
1190
1191 assert!(registry.get("bash").is_some());
1193 assert!(registry.get("read").is_some());
1194 assert!(registry.get("write").is_some());
1195 assert!(registry.get("edit").is_some());
1196 assert!(registry.get("grep").is_some());
1197 assert!(registry.get("find").is_some());
1198 assert!(registry.get("ls").is_some());
1199 }
1200
1201 #[cfg(not(feature = "command-exec"))]
1206 #[test]
1207 fn hardened_build_excludes_command_exec_tools() {
1208 let registry = ToolRegistry::new();
1209 assert!(registry.get("bash").is_none(), "bash must be absent in a hardened build");
1210 assert!(registry.get("shell_start").is_none());
1211 assert!(registry.get("shell_send").is_none());
1212 assert!(registry.get("shell_end").is_none());
1213 assert!(registry.get("read").is_some());
1216 assert!(registry.get("ls").is_some());
1217 assert!(registry.get("grep").is_some());
1218 assert!(registry.get("find").is_some());
1219 assert!(registry.get("subagent").is_some());
1220
1221 let without_sub = ToolRegistry::without_subagent();
1222 assert!(without_sub.get("bash").is_none());
1223 assert!(without_sub.get("shell_start").is_none());
1224 assert!(without_sub.get("read").is_some());
1225 }
1226
1227 #[test]
1228 fn test_tool_registry_register() {
1229 let mut registry = ToolRegistry::without_subagent();
1230 let initial_count = registry.tools_schema().len();
1231
1232 struct TestTool;
1234 #[async_trait::async_trait]
1235 impl Tool for TestTool {
1236 fn name(&self) -> &str {
1237 "test_tool"
1238 }
1239 fn description(&self) -> &str {
1240 "A test tool"
1241 }
1242 fn parameters(&self) -> Value {
1243 json!({"type": "object"})
1244 }
1245 async fn execute(&self, _params: Value, _ctx: ToolContext) -> Result<String> {
1246 Ok("test result".to_string())
1247 }
1248 }
1249
1250 registry.register(Arc::new(TestTool));
1251
1252 assert_eq!(registry.tools_schema().len(), initial_count + 1);
1254
1255 assert!(registry.get("test_tool").is_some());
1257 }
1258
1259 #[test]
1260 fn tool_names_for_extension_filters_by_owner_and_sorts() {
1261 struct OwnedTool(&'static str, Option<&'static str>);
1262 #[async_trait::async_trait]
1263 impl Tool for OwnedTool {
1264 fn name(&self) -> &str {
1265 self.0
1266 }
1267 fn description(&self) -> &str {
1268 "owned"
1269 }
1270 fn parameters(&self) -> Value {
1271 json!({"type": "object"})
1272 }
1273 async fn execute(&self, _params: Value, _ctx: ToolContext) -> Result<String> {
1274 Ok("ok".to_string())
1275 }
1276 fn extension_id(&self) -> Option<&str> {
1277 self.1
1278 }
1279 }
1280
1281 let mut registry = ToolRegistry::without_subagent();
1282 registry.register(Arc::new(OwnedTool("alpha:zed", Some("alpha"))));
1283 registry.register(Arc::new(OwnedTool("alpha:bar", Some("alpha"))));
1284 registry.register(Arc::new(OwnedTool("beta:thing", Some("beta"))));
1285
1286 assert_eq!(
1287 registry.tool_names_for_extension("alpha"),
1288 vec!["alpha:bar".to_string(), "alpha:zed".to_string()]
1289 );
1290 assert_eq!(
1291 registry.tool_names_for_extension("beta"),
1292 vec!["beta:thing".to_string()]
1293 );
1294 assert!(registry.tool_names_for_extension("ghost").is_empty());
1295 assert!(registry.tool_names_for_extension("bash").is_empty());
1297 }
1298
1299 struct OwnedTool(&'static str, Option<&'static str>);
1300 #[async_trait::async_trait]
1301 impl Tool for OwnedTool {
1302 fn name(&self) -> &str {
1303 self.0
1304 }
1305 fn description(&self) -> &str {
1306 "owned"
1307 }
1308 fn parameters(&self) -> Value {
1309 json!({"type": "object"})
1310 }
1311 async fn execute(&self, _params: Value, _ctx: ToolContext) -> Result<String> {
1312 Ok("ok".to_string())
1313 }
1314 fn extension_id(&self) -> Option<&str> {
1315 self.1
1316 }
1317 }
1318
1319 #[test]
1320 fn without_subagent_excludes_subagent_tools() {
1321 let registry = ToolRegistry::without_subagent();
1322 assert!(registry.get("subagent").is_none());
1323 assert!(registry.get("subagent_start").is_none());
1324 assert!(registry.get("subagent_status").is_none());
1325 assert!(registry.get("subagent_steer").is_none());
1326 assert!(registry.get("subagent_collect").is_none());
1327 assert!(registry.get("subagent_resume").is_none());
1328 assert!(registry.get("ls").is_some());
1330 assert!(registry.get("read").is_some());
1331 }
1332
1333 #[test]
1334 fn without_subagent_with_extensions_includes_extension_tools() {
1335 let mut other = ToolRegistry::empty();
1336 other.register(Arc::new(OwnedTool("alpha:do_thing", Some("alpha"))));
1337
1338 let merged = ToolRegistry::without_subagent_with_extensions(&other);
1339
1340 assert!(merged.get("alpha:do_thing").is_some());
1342 assert!(merged.get("read").is_some());
1344 assert!(merged.get("subagent_start").is_none());
1346 }
1347
1348 #[test]
1353 fn worker_set_disables_the_listed_built_ins_before_the_extension_merge() {
1354 let mut shared = ToolRegistry::empty();
1355 shared.register(Arc::new(OwnedTool("alpha:do_thing", Some("alpha"))));
1356 let disabled: Vec<String> = ["bash", "shell_start", "ls", "alpha:do_thing"]
1357 .map(String::from)
1358 .to_vec();
1359
1360 let all = ToolRegistry::worker_set(&shared, &[]);
1361 assert!(all.get("ls").is_some());
1362 let worker = ToolRegistry::worker_set(&shared, &disabled);
1363 let names: Vec<String> = worker
1364 .tools_schema()
1365 .iter()
1366 .filter_map(|t| t["name"].as_str().map(String::from))
1367 .collect();
1368 for name in ["bash", "shell_start", "ls"] {
1369 assert!(worker.get(name).is_none(), "{name} kept");
1370 assert!(!names.iter().any(|n| n == name), "{name} advertised: {names:?}");
1371 }
1372 assert!(worker.get("alpha:do_thing").is_some(), "extension tool kept");
1373 assert!(worker.get("read").is_some() && worker.get("subagent_start").is_none());
1374 assert_eq!(names.len(), worker.tools.len());
1375
1376 let mut foreground = ToolRegistry::new();
1378 foreground.disable(&disabled);
1379 for tool in shared.tools.values() {
1380 foreground.insert_tool_with_catalog(tool.clone()).unwrap();
1381 }
1382 for name in worker.tools.keys() {
1383 assert!(foreground.get(name).is_some(), "worker has {name}, the foreground does not");
1384 }
1385 }
1386
1387 #[tokio::test]
1391 async fn disabled_multi_edit_is_gone_from_workers_too() {
1392 let disabled = vec!["multi_edit".to_string()];
1393 let worker = ToolRegistry::worker_set(&ToolRegistry::empty(), &disabled);
1394 let legacy = crate::tools::subagent::subagent_tools(&disabled).await;
1395 for registry in [worker, legacy] {
1396 assert!(registry.get("multi_edit").is_none());
1397 assert!(!registry.tools_schema().iter().any(|t| t["name"] == "multi_edit"));
1398 assert!(registry.get("edit").is_some() && registry.get("write").is_some());
1399 }
1400 assert!(ToolRegistry::worker_set(&ToolRegistry::empty(), &[]).get("multi_edit").is_some());
1401 }
1402
1403 #[cfg(feature = "command-exec")]
1404 #[test]
1405 fn worker_registries_keep_their_bounds_with_extensions_merged() {
1406 let mut extensions = ToolRegistry::empty();
1407 extensions.register(Arc::new(OwnedTool("peer:probe", Some("peer"))));
1408 let merged = ToolRegistry::without_subagent_with_extensions(&extensions);
1409 assert_eq!(merged.tools_schema().len(), 12 + usize::from(cfg!(windows)));
1410 assert!(merged.get("peer:probe").is_some());
1411 for registry in [ToolRegistry::without_subagent(), merged] {
1412 for name in [
1413 "subagent",
1414 "subagent_start",
1415 "subagent_status",
1416 "subagent_steer",
1417 "subagent_collect",
1418 "subagent_resume",
1419 "subagent_model_authorize",
1420 "subagent_models",
1421 "search_tools",
1422 "activate_tools",
1423 "memory_store",
1424 "memory_search",
1425 "memory_fetch",
1426 "memory_forget",
1427 ] {
1428 assert!(registry.get(name).is_none(), "worker must not gain {name}");
1429 }
1430 for name in ["bash", "write", "edit", "multi_edit"] {
1431 assert!(registry.get(name).is_some(), "worker must retain {name}");
1432 }
1433 }
1434 }
1435
1436 #[test]
1437 fn without_subagent_with_extensions_excludes_built_ins_from_other_registry() {
1438 let other = ToolRegistry::new();
1443
1444 let merged = ToolRegistry::without_subagent_with_extensions(&other);
1445
1446 let read = merged.get("read").expect("read present");
1448 assert!(read.extension_id().is_none());
1449 assert!(merged.get("subagent_start").is_none());
1451 assert!(merged.get("subagent").is_none());
1452 }
1453
1454 #[test]
1455 fn without_subagent_with_extensions_rejects_recursive_tool_name_collisions() {
1456 let mut other = ToolRegistry::empty();
1457 for name in [
1458 "subagent",
1459 "subagent_start",
1460 "subagent_status",
1461 "subagent_steer",
1462 "subagent_collect",
1463 "subagent_resume",
1464 "subagent_model_authorize",
1465 "subagent_models",
1466 ] {
1467 other.register(Arc::new(OwnedTool(name, Some("malicious-extension"))));
1468 }
1469
1470 let merged = ToolRegistry::without_subagent_with_extensions(&other);
1471 for name in [
1472 "subagent",
1473 "subagent_start",
1474 "subagent_status",
1475 "subagent_steer",
1476 "subagent_collect",
1477 "subagent_resume",
1478 "subagent_model_authorize",
1479 "subagent_models",
1480 ] {
1481 assert!(merged.get(name).is_none(), "{name} must stay unavailable");
1482 }
1483 }
1484
1485 #[test]
1486 fn without_subagent_with_extensions_rejects_api_sanitized_recursive_collisions() {
1487 let mut other = ToolRegistry::empty();
1488 for runtime_name in [
1489 "subagent:start",
1490 "subagent:status",
1491 "subagent:steer",
1492 "subagent:collect",
1493 "subagent:resume",
1494 "subagent:model:authorize",
1495 "subagent:models",
1496 ] {
1497 other.register(Arc::new(OwnedTool(runtime_name, Some("subagent"))));
1498 }
1499
1500 let merged = ToolRegistry::without_subagent_with_extensions(&other);
1501 for api_name in [
1502 "subagent_start",
1503 "subagent_status",
1504 "subagent_steer",
1505 "subagent_collect",
1506 "subagent_resume",
1507 "subagent_model_authorize",
1508 "subagent_models",
1509 ] {
1510 assert!(
1511 merged.get(api_name).is_none(),
1512 "API-safe recursive name {api_name} must stay unavailable"
1513 );
1514 }
1515 }
1516
1517 #[test]
1518 fn without_subagent_with_extensions_does_not_overwrite_existing_builtin() {
1519 let mut other = ToolRegistry::empty();
1525 other.register(Arc::new(OwnedTool("ext:custom", Some("ext"))));
1526
1527 let merged = ToolRegistry::without_subagent_with_extensions(&other);
1528 assert!(merged.get("ext:custom").is_some());
1529 assert!(merged.get("read").is_some());
1530 assert!(merged.get("read").unwrap().extension_id().is_none());
1531 }
1532
1533 struct McpTwin(&'static str);
1535
1536 #[async_trait::async_trait]
1537 impl Tool for McpTwin {
1538 fn name(&self) -> &str {
1539 self.0
1540 }
1541 fn description(&self) -> &str {
1542 "mcp twin"
1543 }
1544 fn parameters(&self) -> Value {
1545 json!({"type": "object"})
1546 }
1547 async fn execute(&self, _params: Value, _ctx: ToolContext) -> Result<String> {
1548 Ok("ok".to_string())
1549 }
1550 fn origin(&self) -> crate::tools::ToolOrigin {
1551 crate::tools::ToolOrigin::Mcp {
1552 server_id: "srv".to_string(),
1553 server_tool_name: "shared".to_string(),
1554 }
1555 }
1556 }
1557
1558 #[test]
1559 #[should_panic(expected = "tool catalog rejected mutation")]
1560 fn from_tools_with_duplicate_capability_identity_fails_closed() {
1561 let _ = ToolRegistry::from_tools(vec![
1565 Arc::new(McpTwin("twin_one")),
1566 Arc::new(McpTwin("twin_two")),
1567 ]);
1568 }
1569}