1#![allow(clippy::items_after_test_module)]
2
3use std::collections::{BTreeMap, BTreeSet};
4use std::sync::{Arc, RwLock};
5
6use serde::{Deserialize, Serialize};
7
8use crate::{
9 PreparedToolCall, ProgressSender, ToolCall, ToolContext, ToolContract, ToolManifest,
10 ToolPrepareCall, ToolProvider, ToolResult,
11};
12
13const PLUGIN_SOURCE_ID: &str = "plugins";
14
15#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
16#[serde(transparent)]
17pub struct ToolSourceHandle {
18 id: String,
19}
20
21impl ToolSourceHandle {
22 pub(crate) fn new(id: impl Into<String>) -> Self {
23 Self { id: id.into() }
24 }
25
26 pub(crate) fn as_str(&self) -> &str {
27 &self.id
28 }
29}
30
31#[derive(Clone, Debug, Serialize, Deserialize)]
32pub struct ToolStateEntry {
33 manifest: ToolManifest,
34}
35
36impl ToolStateEntry {
37 pub(crate) fn new(manifest: ToolManifest) -> Self {
38 Self { manifest }
39 }
40
41 pub fn manifest(&self) -> &ToolManifest {
42 &self.manifest
43 }
44
45 pub fn manifest_mut(&mut self) -> &mut ToolManifest {
46 &mut self.manifest
47 }
48}
49
50#[derive(Clone, Debug, Default)]
51pub struct ToolState {
52 generation: u64,
53 tools: Arc<BTreeMap<String, ToolStateEntry>>,
54}
55
56impl ToolState {
57 pub(crate) fn new(generation: u64, tools: BTreeMap<String, ToolStateEntry>) -> Self {
58 Self {
59 generation,
60 tools: Arc::new(tools),
61 }
62 }
63
64 pub fn generation(&self) -> u64 {
65 self.generation
66 }
67
68 pub fn with_generation(mut self, generation: u64) -> Self {
69 self.generation = generation;
70 self
71 }
72
73 pub fn tool_manifests(&self) -> Vec<ToolManifest> {
74 self.tools
75 .values()
76 .map(|entry| entry.manifest.clone())
77 .collect()
78 }
79
80 pub fn get(&self, name: &str) -> Option<&ToolStateEntry> {
81 self.tools.get(name)
82 }
83
84 pub fn manifest_mut(&mut self, name: &str) -> Option<&mut ToolManifest> {
85 Arc::make_mut(&mut self.tools)
86 .get_mut(name)
87 .map(|entry| &mut entry.manifest)
88 }
89
90 pub fn contains(&self, name: &str) -> bool {
91 self.tools.contains_key(name)
92 }
93
94 pub fn is_empty(&self) -> bool {
95 self.tools.is_empty()
96 }
97
98 pub fn len(&self) -> usize {
99 self.tools.len()
100 }
101
102 pub fn iter(&self) -> impl Iterator<Item = (&str, &ToolStateEntry)> {
103 self.tools
104 .iter()
105 .map(|(name, entry)| (name.as_str(), entry))
106 }
107
108 pub fn set_availability(
109 &mut self,
110 name: &str,
111 availability: Option<crate::ToolAvailability>,
112 ) -> Result<(), ReconfigureError> {
113 let Some(entry) = Arc::make_mut(&mut self.tools).get_mut(name) else {
114 return Err(ReconfigureError::Validation(format!(
115 "unknown tool `{name}`"
116 )));
117 };
118 entry.manifest.availability_override = availability;
119 Ok(())
120 }
121
122 pub fn retain(&mut self, mut keep: impl FnMut(&str, &ToolStateEntry) -> bool) {
123 Arc::make_mut(&mut self.tools).retain(|name, entry| keep(name, entry));
124 }
125
126 pub fn remove(&mut self, name: &str) -> Option<ToolStateEntry> {
127 Arc::make_mut(&mut self.tools).remove(name)
128 }
129
130 pub(crate) fn entries(&self) -> &BTreeMap<String, ToolStateEntry> {
131 self.tools.as_ref()
132 }
133}
134
135impl Serialize for ToolState {
136 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
137 where
138 S: serde::Serializer,
139 {
140 #[derive(Serialize)]
141 struct ToolStateRef<'a> {
142 generation: u64,
143 tools: &'a BTreeMap<String, ToolStateEntry>,
144 }
145
146 ToolStateRef {
147 generation: self.generation,
148 tools: self.tools.as_ref(),
149 }
150 .serialize(serializer)
151 }
152}
153
154impl<'de> Deserialize<'de> for ToolState {
155 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
156 where
157 D: serde::Deserializer<'de>,
158 {
159 #[derive(Deserialize)]
160 struct ToolStateOwned {
161 generation: u64,
162 tools: BTreeMap<String, ToolStateEntry>,
163 }
164
165 let owned = ToolStateOwned::deserialize(deserializer)?;
166 Ok(Self {
167 generation: owned.generation,
168 tools: Arc::new(owned.tools),
169 })
170 }
171}
172
173#[async_trait::async_trait]
174pub(crate) trait ToolSourceExecutor: Send + Sync + 'static {
175 fn id(&self) -> &str;
176 fn advertised_tools(&self) -> Vec<ToolManifest>;
177 fn resolve_manifest(&self, name: &str) -> Option<ToolManifest> {
178 self.advertised_tools()
179 .into_iter()
180 .find(|manifest| manifest.name == name)
181 }
182 fn resolve_contract(&self, name: &str) -> Option<Arc<ToolContract>>;
183 async fn prepare_tool_call(
184 &self,
185 call: ToolPrepareCall<'_>,
186 ) -> Result<PreparedToolCall, ToolResult> {
187 Ok(PreparedToolCall::identity(call.pending))
188 }
189 async fn execute(
190 &self,
191 tool: &str,
192 args: &serde_json::Value,
193 context: &ToolContext<'_>,
194 progress: Option<&ProgressSender>,
195 ) -> ToolResult;
196}
197
198struct ToolProviderSource {
199 id: String,
200 provider: Arc<dyn ToolProvider>,
201}
202
203impl ToolProviderSource {
204 fn new(id: impl Into<String>, provider: Arc<dyn ToolProvider>) -> Self {
205 Self {
206 id: id.into(),
207 provider,
208 }
209 }
210}
211
212struct ToolProviderGroupSource {
213 id: String,
214 tools: RwLock<BTreeMap<String, (ToolManifest, usize)>>,
215 providers: Vec<Arc<dyn ToolProvider>>,
216}
217
218impl ToolProviderGroupSource {
219 fn new(id: impl Into<String>, providers: Vec<Arc<dyn ToolProvider>>) -> Self {
220 let mut tools = BTreeMap::new();
221 for (provider_idx, provider) in providers.iter().enumerate() {
222 for manifest in provider.tool_manifests() {
223 tools.insert(manifest.name.clone(), (manifest, provider_idx));
224 }
225 }
226 Self {
227 id: id.into(),
228 tools: RwLock::new(tools),
229 providers,
230 }
231 }
232
233 fn read_advertised_tools(&self) -> Vec<ToolManifest> {
234 let mut tools = BTreeMap::new();
235 for (provider_idx, provider) in self.providers.iter().enumerate() {
236 for manifest in provider.tool_manifests() {
237 tools.insert(manifest.name.clone(), (manifest, provider_idx));
238 }
239 }
240 let manifests = tools
241 .values()
242 .map(|(manifest, _)| manifest.clone())
243 .collect::<Vec<_>>();
244 *self
245 .tools
246 .write()
247 .expect("tool provider group lock poisoned") = tools;
248 manifests
249 }
250
251 fn provider_index_for(&self, name: &str) -> Option<usize> {
252 self.resolve_manifest(name).and_then(|_| {
253 self.tools
254 .read()
255 .expect("tool provider group lock poisoned")
256 .get(name)
257 .map(|(_, provider_idx)| *provider_idx)
258 })
259 }
260}
261
262#[async_trait::async_trait]
263impl ToolSourceExecutor for ToolProviderGroupSource {
264 fn id(&self) -> &str {
265 &self.id
266 }
267
268 fn advertised_tools(&self) -> Vec<ToolManifest> {
269 self.read_advertised_tools()
270 }
271
272 fn resolve_manifest(&self, name: &str) -> Option<ToolManifest> {
273 if let Some((manifest, _)) = self
274 .tools
275 .read()
276 .expect("tool provider group lock poisoned")
277 .get(name)
278 {
279 return Some(manifest.clone());
280 }
281 for (provider_idx, provider) in self.providers.iter().enumerate() {
282 if let Some(manifest) = provider.resolve_manifest(name) {
283 self.tools
284 .write()
285 .expect("tool provider group lock poisoned")
286 .insert(name.to_string(), (manifest.clone(), provider_idx));
287 return Some(manifest);
288 }
289 }
290 None
291 }
292
293 fn resolve_contract(&self, name: &str) -> Option<Arc<ToolContract>> {
294 let provider_idx = self.provider_index_for(name)?;
295 self.providers[provider_idx].resolve_contract(name)
296 }
297
298 async fn prepare_tool_call(
299 &self,
300 call: ToolPrepareCall<'_>,
301 ) -> Result<PreparedToolCall, ToolResult> {
302 let name = call.pending.tool_name.clone();
303 let Some(provider_idx) = self.provider_index_for(&name) else {
304 return Err(ToolResult::err_fmt(format_args!("Unknown tool: {name}")));
305 };
306 self.providers[provider_idx].prepare_tool_call(call).await
307 }
308
309 async fn execute(
310 &self,
311 tool: &str,
312 args: &serde_json::Value,
313 context: &ToolContext<'_>,
314 progress: Option<&ProgressSender>,
315 ) -> ToolResult {
316 let Some(provider_idx) = self.provider_index_for(tool) else {
317 return ToolResult::err_fmt(format_args!("Unknown tool: {tool}"));
318 };
319 self.providers[provider_idx]
320 .execute(ToolCall {
321 name: tool,
322 args,
323 context,
324 progress,
325 })
326 .await
327 }
328}
329
330#[async_trait::async_trait]
331impl ToolSourceExecutor for ToolProviderSource {
332 fn id(&self) -> &str {
333 &self.id
334 }
335
336 fn advertised_tools(&self) -> Vec<ToolManifest> {
337 self.provider.tool_manifests()
338 }
339
340 fn resolve_manifest(&self, name: &str) -> Option<ToolManifest> {
341 self.provider.resolve_manifest(name)
342 }
343
344 fn resolve_contract(&self, name: &str) -> Option<Arc<ToolContract>> {
345 self.provider.resolve_contract(name)
346 }
347
348 async fn prepare_tool_call(
349 &self,
350 call: ToolPrepareCall<'_>,
351 ) -> Result<PreparedToolCall, ToolResult> {
352 self.provider.prepare_tool_call(call).await
353 }
354
355 async fn execute(
356 &self,
357 tool: &str,
358 args: &serde_json::Value,
359 context: &ToolContext<'_>,
360 progress: Option<&ProgressSender>,
361 ) -> ToolResult {
362 self.provider
363 .execute(ToolCall {
364 name: tool,
365 args,
366 context,
367 progress,
368 })
369 .await
370 }
371}
372
373#[derive(Clone)]
374struct ToolRegistryEntry {
375 manifest: ToolManifest,
376 source_id: String,
377}
378
379impl ToolRegistryEntry {
380 fn new(manifest: ToolManifest, source_id: impl Into<String>) -> Self {
381 Self {
382 manifest,
383 source_id: source_id.into(),
384 }
385 }
386
387 fn export(&self) -> ToolStateEntry {
388 ToolStateEntry::new(self.manifest.clone())
389 }
390}
391
392#[derive(Clone)]
393struct ToolRegistryState {
394 generation: u64,
395 tools: BTreeMap<String, ToolRegistryEntry>,
396 next_live_source_id: u64,
397}
398
399#[derive(Debug, thiserror::Error)]
400pub enum ReconfigureError {
401 #[error("validation error: {0}")]
402 Validation(String),
403 #[error("unknown tool source: {0}")]
404 UnknownSource(String),
405 #[error("generation mismatch: expected {expected}, actual {actual}")]
406 GenerationMismatch { expected: u64, actual: u64 },
407}
408
409#[derive(Clone)]
410pub struct ToolRegistry {
411 sources: Arc<RwLock<BTreeMap<String, Arc<dyn ToolSourceExecutor>>>>,
412 state: Arc<RwLock<ToolRegistryState>>,
413}
414
415impl ToolRegistry {
416 pub fn from_tool_provider(provider: Arc<dyn ToolProvider>) -> Result<Self, ReconfigureError> {
417 let registry = Self::empty();
418 registry.upsert_source(Arc::new(ToolProviderSource::new(
419 PLUGIN_SOURCE_ID,
420 provider,
421 )))?;
422 Ok(registry)
423 }
424
425 pub(crate) fn from_tool_providers(
426 providers: Vec<Arc<dyn ToolProvider>>,
427 ) -> Result<Self, ReconfigureError> {
428 let registry = Self::empty();
429 registry.upsert_source(Arc::new(ToolProviderGroupSource::new(
430 PLUGIN_SOURCE_ID,
431 providers,
432 )))?;
433 Ok(registry)
434 }
435
436 pub(crate) fn empty() -> Self {
437 Self {
438 sources: Arc::new(RwLock::new(BTreeMap::new())),
439 state: Arc::new(RwLock::new(ToolRegistryState {
440 generation: 0,
441 tools: BTreeMap::new(),
442 next_live_source_id: 0,
443 })),
444 }
445 }
446
447 pub fn generation(&self) -> u64 {
448 self.state
449 .read()
450 .expect("tool registry state lock poisoned")
451 .generation
452 }
453
454 pub fn export_state(&self) -> ToolState {
455 let state = self
456 .state
457 .read()
458 .expect("tool registry state lock poisoned");
459 ToolState::new(state.generation, export_tool_state_entries(&state.tools))
460 }
461
462 pub fn apply_state(&self, next: ToolState) -> Result<u64, ReconfigureError> {
463 let current_generation = self.generation();
464 if next.generation != current_generation {
465 return Err(ReconfigureError::GenerationMismatch {
466 expected: next.generation,
467 actual: current_generation,
468 });
469 }
470
471 validate_unique_manifest_entries(next.entries().values())?;
472 let rebound_tools = {
473 let sources = self.sources.read().expect("tool source lock poisoned");
474 rebind_tool_state_entries(next.entries(), &sources)?
475 };
476
477 let mut state = self
478 .state
479 .write()
480 .expect("tool registry state lock poisoned");
481 if state.generation != next.generation {
482 return Err(ReconfigureError::GenerationMismatch {
483 expected: next.generation,
484 actual: state.generation,
485 });
486 }
487 state.tools = rebound_tools;
488 state.generation += 1;
489 Ok(state.generation)
490 }
491
492 pub fn restore_state(&self, snapshot: ToolState) -> Result<u64, ReconfigureError> {
506 validate_unique_manifest_entries(snapshot.entries().values())?;
507 let rebound_tools = {
508 let sources = self.sources.read().expect("tool source lock poisoned");
509 rebind_tool_state_entries(snapshot.entries(), &sources)?
510 };
511
512 let mut state = self
513 .state
514 .write()
515 .expect("tool registry state lock poisoned");
516 state.tools = rebound_tools;
517 state.generation = snapshot.generation();
518 Ok(state.generation)
519 }
520
521 pub fn add_tool_provider(
522 &self,
523 provider: Arc<dyn ToolProvider>,
524 ) -> Result<ToolSourceHandle, ReconfigureError> {
525 let source_id = {
526 let mut state = self
527 .state
528 .write()
529 .expect("tool registry state lock poisoned");
530 state.next_live_source_id += 1;
531 format!("live:{}", state.next_live_source_id)
532 };
533 self.upsert_source(Arc::new(ToolProviderSource::new(
534 source_id.clone(),
535 provider,
536 )))?;
537 Ok(ToolSourceHandle::new(source_id))
538 }
539
540 pub(crate) fn compose_session_surface(
541 &self,
542 include_base_tools: bool,
543 context_providers: Vec<Arc<dyn ToolProvider>>,
544 ) -> Result<Self, ReconfigureError> {
545 let registry = if include_base_tools {
546 self.fork_with_state(self.export_state())?
547 } else {
548 Self::empty()
549 };
550 registry.upsert_overlay_source(Arc::new(ToolProviderGroupSource::new(
551 "context",
552 context_providers,
553 )))?;
554 Ok(registry)
555 }
556
557 pub(crate) fn upsert_source(
558 &self,
559 source: Arc<dyn ToolSourceExecutor>,
560 ) -> Result<u64, ReconfigureError> {
561 let source_id = source.id().to_string();
562 let advertised_tools = source
563 .advertised_tools()
564 .into_iter()
565 .map(|manifest| manifest_with_compact_contract(source.as_ref(), manifest))
566 .collect::<Vec<_>>();
567 validate_unique_manifests(&advertised_tools)?;
568
569 let mut state = self
570 .state
571 .write()
572 .expect("tool registry state lock poisoned");
573 let previous_overrides = state
574 .tools
575 .iter()
576 .map(|(name, entry)| (name.clone(), entry.manifest.availability_override))
577 .collect::<BTreeMap<_, _>>();
578 let same_source_names = state
579 .tools
580 .iter()
581 .filter_map(|(name, entry)| (entry.source_id == source_id).then_some(name.clone()))
582 .collect::<BTreeSet<_>>();
583 for manifest in &advertised_tools {
584 if let Some(existing) = state.tools.get(&manifest.name)
585 && existing.source_id != source_id
586 {
587 return Err(ReconfigureError::Validation(format!(
588 "duplicate tool name `{}` from source `{}` conflicts with source `{}`",
589 manifest.name, source_id, existing.source_id
590 )));
591 }
592 if let Some((existing_name, existing)) = state
593 .tools
594 .iter()
595 .find(|(_, entry)| entry.source_id != source_id && entry.manifest.id == manifest.id)
596 {
597 return Err(ReconfigureError::Validation(format!(
598 "duplicate tool id `{}` from source `{}` conflicts with tool `{}` from source `{}`",
599 manifest.id, source_id, existing_name, existing.source_id
600 )));
601 }
602 }
603 state.tools.retain(|name, entry| {
604 entry.source_id != source_id || !same_source_names.contains(name)
605 });
606
607 for mut manifest in advertised_tools {
608 let name = manifest.name.clone();
609 manifest.availability_override = previous_overrides
610 .get(&name)
611 .copied()
612 .flatten()
613 .or(manifest.availability_override);
614 state
615 .tools
616 .insert(name, ToolRegistryEntry::new(manifest, source_id.clone()));
617 }
618
619 self.sources
620 .write()
621 .expect("tool source lock poisoned")
622 .insert(source_id, source);
623 state.generation += 1;
624 Ok(state.generation)
625 }
626
627 fn upsert_overlay_source(
628 &self,
629 source: Arc<dyn ToolSourceExecutor>,
630 ) -> Result<u64, ReconfigureError> {
631 let source_id = source.id().to_string();
632 let advertised_tools = source
633 .advertised_tools()
634 .into_iter()
635 .map(|manifest| manifest_with_compact_contract(source.as_ref(), manifest))
636 .collect::<Vec<_>>();
637 validate_unique_manifests(&advertised_tools)?;
638
639 let advertised_names = advertised_tools
640 .iter()
641 .map(|manifest| manifest.name.clone())
642 .collect::<BTreeSet<_>>();
643 let advertised_ids = advertised_tools
644 .iter()
645 .map(|manifest| manifest.id.clone())
646 .collect::<BTreeSet<_>>();
647 let mut state = self
648 .state
649 .write()
650 .expect("tool registry state lock poisoned");
651 let previous_overrides = state
652 .tools
653 .iter()
654 .map(|(name, entry)| (name.clone(), entry.manifest.availability_override))
655 .collect::<BTreeMap<_, _>>();
656 state.tools.retain(|name, entry| {
657 entry.source_id != source_id
658 && !advertised_names.contains(name)
659 && !advertised_ids.contains(&entry.manifest.id)
660 });
661 for mut manifest in advertised_tools {
662 let name = manifest.name.clone();
663 manifest.availability_override = previous_overrides
664 .get(&name)
665 .copied()
666 .flatten()
667 .or(manifest.availability_override);
668 state
669 .tools
670 .insert(name, ToolRegistryEntry::new(manifest, source_id.clone()));
671 }
672 self.sources
673 .write()
674 .expect("tool source lock poisoned")
675 .insert(source_id, source);
676 state.generation += 1;
677 Ok(state.generation)
678 }
679
680 pub fn remove_source(&self, handle: &ToolSourceHandle) -> Result<u64, ReconfigureError> {
681 self.remove_source_id(handle.as_str())
682 }
683
684 pub fn refresh_sources(&self) -> Result<u64, ReconfigureError> {
685 let sources = self
686 .sources
687 .read()
688 .expect("tool source lock poisoned")
689 .values()
690 .cloned()
691 .collect::<Vec<_>>();
692 let mut generation = self.generation();
693 for source in sources {
694 generation = self.upsert_source(source)?;
695 }
696 Ok(generation)
697 }
698
699 pub(crate) fn remove_source_id(&self, source_id: &str) -> Result<u64, ReconfigureError> {
700 {
701 let mut sources = self.sources.write().expect("tool source lock poisoned");
702 if sources.remove(source_id).is_none() {
703 return Err(ReconfigureError::UnknownSource(source_id.to_string()));
704 }
705 }
706 let mut state = self
707 .state
708 .write()
709 .expect("tool registry state lock poisoned");
710 state.tools.retain(|_, entry| entry.source_id != source_id);
711 state.generation += 1;
712 Ok(state.generation)
713 }
714
715 pub(crate) fn fork_with_state(&self, snapshot: ToolState) -> Result<Self, ReconfigureError> {
716 let sources = self
717 .sources
718 .read()
719 .expect("tool source lock poisoned")
720 .iter()
721 .map(|(k, v)| (k.clone(), Arc::clone(v)))
722 .collect::<BTreeMap<_, _>>();
723 validate_unique_manifest_entries(snapshot.entries().values())?;
724 let tools = rebind_tool_state_entries(snapshot.entries(), &sources)?;
725 let generation = snapshot.generation.max(1);
726 Ok(Self {
727 sources: Arc::new(RwLock::new(sources)),
728 state: Arc::new(RwLock::new(ToolRegistryState {
729 generation,
730 tools,
731 next_live_source_id: 0,
732 })),
733 })
734 }
735}
736
737#[async_trait::async_trait]
738impl ToolProvider for ToolRegistry {
739 fn tool_manifests(&self) -> Vec<ToolManifest> {
740 let state = self
741 .state
742 .read()
743 .expect("tool registry state lock poisoned");
744 state
745 .tools
746 .values()
747 .map(|entry| entry.manifest.clone())
748 .collect()
749 }
750
751 fn resolve_manifest(&self, name: &str) -> Option<ToolManifest> {
752 if let Some(manifest) = {
753 let state = self
754 .state
755 .read()
756 .expect("tool registry state lock poisoned");
757 state.tools.get(name).map(|entry| entry.manifest.clone())
758 } {
759 return Some(manifest);
760 }
761
762 let sources = self
763 .sources
764 .read()
765 .expect("tool source lock poisoned")
766 .iter()
767 .map(|(source_id, source)| (source_id.clone(), Arc::clone(source)))
768 .collect::<Vec<_>>();
769 for (source_id, source) in sources {
770 let Some(manifest) = source.resolve_manifest(name) else {
771 continue;
772 };
773 let mut manifest = manifest_with_compact_contract(source.as_ref(), manifest);
774 let previous_override = {
775 let state = self
776 .state
777 .read()
778 .expect("tool registry state lock poisoned");
779 state
780 .tools
781 .get(&manifest.name)
782 .and_then(|entry| entry.manifest.availability_override)
783 };
784 manifest.availability_override = previous_override.or(manifest.availability_override);
785 let mut state = self
786 .state
787 .write()
788 .expect("tool registry state lock poisoned");
789 if let Some(existing) = state.tools.get(&manifest.name) {
790 return (existing.source_id == source_id).then(|| existing.manifest.clone());
791 }
792 if let Some((_, existing)) = state
793 .tools
794 .iter()
795 .find(|(_, entry)| entry.manifest.id == manifest.id)
796 {
797 return (existing.source_id == source_id).then(|| existing.manifest.clone());
798 }
799 state.tools.insert(
800 manifest.name.clone(),
801 ToolRegistryEntry::new(manifest.clone(), source_id),
802 );
803 state.generation += 1;
804 return Some(manifest);
805 }
806 None
807 }
808
809 fn resolve_contract(&self, name: &str) -> Option<Arc<ToolContract>> {
810 let source_id = self.resolve_manifest(name).and_then(|_| {
811 let state = self
812 .state
813 .read()
814 .expect("tool registry state lock poisoned");
815 state.tools.get(name).map(|entry| entry.source_id.clone())
816 })?;
817 self.sources
818 .read()
819 .expect("tool source lock poisoned")
820 .get(&source_id)?
821 .resolve_contract(name)
822 }
823
824 async fn prepare_tool_call(
825 &self,
826 call: ToolPrepareCall<'_>,
827 ) -> Result<PreparedToolCall, ToolResult> {
828 let name = call.pending.tool_name.clone();
829 let source_id = self.resolve_manifest(&name).and_then(|_| {
830 let state = self
831 .state
832 .read()
833 .expect("tool registry state lock poisoned");
834 state.tools.get(&name).map(|entry| entry.source_id.clone())
835 });
836 let Some(source_id) = source_id else {
837 return Err(ToolResult::err_fmt(format_args!("Unknown tool: {name}")));
838 };
839 let source = {
840 self.sources
841 .read()
842 .expect("tool source lock poisoned")
843 .get(&source_id)
844 .cloned()
845 };
846 let Some(source) = source else {
847 return Err(ToolResult::err_fmt(format_args!(
848 "Tool source missing for tool `{name}`"
849 )));
850 };
851 source.prepare_tool_call(call).await
852 }
853
854 async fn execute(&self, call: ToolCall<'_>) -> ToolResult {
855 let name = call.name;
856 let source_id = self.resolve_manifest(name).and_then(|_| {
857 let state = self
858 .state
859 .read()
860 .expect("tool registry state lock poisoned");
861 state.tools.get(name).map(|entry| entry.source_id.clone())
862 });
863 let Some(source_id) = source_id else {
864 return ToolResult::err_fmt(format_args!("Unknown tool: {name}"));
865 };
866 let source = {
867 self.sources
868 .read()
869 .expect("tool source lock poisoned")
870 .get(&source_id)
871 .cloned()
872 };
873 let Some(source) = source else {
874 return ToolResult::err_fmt(format_args!("Tool source missing for tool `{name}`"));
875 };
876 source
877 .execute(name, call.args, call.context, call.progress)
878 .await
879 }
880}
881
882fn validate_unique_manifests(manifests: &[ToolManifest]) -> Result<(), ReconfigureError> {
883 let mut names = BTreeSet::new();
884 let mut ids = BTreeSet::new();
885 for manifest in manifests {
886 if manifest.id.as_str().trim().is_empty() {
887 return Err(ReconfigureError::Validation(
888 "tool id cannot be empty".to_string(),
889 ));
890 }
891 if !ids.insert(manifest.id.clone()) {
892 return Err(ReconfigureError::Validation(format!(
893 "duplicate tool id `{}` in source",
894 manifest.id
895 )));
896 }
897 if manifest.name.trim().is_empty() {
898 return Err(ReconfigureError::Validation(
899 "tool name cannot be empty".to_string(),
900 ));
901 }
902 if !names.insert(manifest.name.clone()) {
903 return Err(ReconfigureError::Validation(format!(
904 "duplicate tool name `{}` in source",
905 manifest.name
906 )));
907 }
908 }
909 Ok(())
910}
911
912fn manifest_with_compact_contract(
913 source: &dyn ToolSourceExecutor,
914 mut manifest: ToolManifest,
915) -> ToolManifest {
916 if manifest.compact_contract.is_none()
917 && let Some(contract) = source.resolve_contract(&manifest.name)
918 {
919 manifest.compact_contract = Some(contract.compact_contract(&manifest));
920 }
921 manifest
922}
923
924fn export_tool_state_entries(
925 entries: &BTreeMap<String, ToolRegistryEntry>,
926) -> BTreeMap<String, ToolStateEntry> {
927 entries
928 .iter()
929 .map(|(name, entry)| (name.clone(), entry.export()))
930 .collect()
931}
932
933fn rebind_tool_state_entries(
934 entries: &BTreeMap<String, ToolStateEntry>,
935 sources: &BTreeMap<String, Arc<dyn ToolSourceExecutor>>,
936) -> Result<BTreeMap<String, ToolRegistryEntry>, ReconfigureError> {
937 let mut rebound = BTreeMap::new();
938 for (name, entry) in entries {
939 if name != &entry.manifest.name {
940 return Err(ReconfigureError::Validation(format!(
941 "tool state key `{}` does not match manifest name `{}`",
942 name, entry.manifest.name
943 )));
944 }
945
946 let mut name_matches = Vec::new();
947 for (source_id, source) in sources {
948 let Some(manifest) = source.resolve_manifest(name) else {
949 continue;
950 };
951 name_matches.push((
952 source_id.clone(),
953 manifest_with_compact_contract(source.as_ref(), manifest),
954 ));
955 }
956
957 if name_matches.is_empty() {
958 return Err(ReconfigureError::Validation(format!(
959 "no registered tool source resolves tool `{name}`"
960 )));
961 }
962
963 let matching_id = name_matches
964 .iter()
965 .filter(|(_, manifest)| manifest.id == entry.manifest.id)
966 .collect::<Vec<_>>();
967
968 if matching_id.len() == 1 {
969 let source_id = matching_id[0].0.clone();
970 rebound.insert(
971 name.clone(),
972 ToolRegistryEntry::new(entry.manifest.clone(), source_id),
973 );
974 } else if matching_id.is_empty() {
975 let resolved_ids = name_matches
976 .iter()
977 .map(|(_, manifest)| manifest.id.as_str())
978 .collect::<Vec<_>>()
979 .join(", ");
980 return Err(ReconfigureError::Validation(format!(
981 "tool `{name}` resolved with id(s) `{resolved_ids}`, expected `{}`",
982 entry.manifest.id
983 )));
984 } else {
985 return Err(ReconfigureError::Validation(format!(
986 "tool `{name}` with id `{}` is resolved by multiple registered sources",
987 entry.manifest.id
988 )));
989 }
990 }
991 Ok(rebound)
992}
993
994fn validate_unique_manifest_entries<'a>(
995 entries: impl IntoIterator<Item = &'a ToolStateEntry>,
996) -> Result<(), ReconfigureError> {
997 let manifests = entries
998 .into_iter()
999 .map(|entry| entry.manifest.clone())
1000 .collect::<Vec<_>>();
1001 validate_unique_manifests(&manifests)
1002}
1003
1004#[cfg(test)]
1005mod tests {
1006 use super::*;
1007 use crate::ToolDefinition;
1008 use serde_json::json;
1009 use std::sync::atomic::{AtomicUsize, Ordering};
1010
1011 struct MockTool;
1012 struct MixedEnabledTool;
1013 struct ExternalMockSource;
1014 struct ExactResolvingSource {
1015 manifest_resolutions: Arc<AtomicUsize>,
1016 contract_resolutions: Arc<AtomicUsize>,
1017 executions: Arc<AtomicUsize>,
1018 }
1019 struct NamedExactSource {
1020 id: &'static str,
1021 }
1022 struct DynamicToolProvider {
1023 names: Arc<std::sync::Mutex<Vec<String>>>,
1024 }
1025
1026 fn test_tool(
1027 name: &str,
1028 description: &str,
1029 availability: crate::ToolAvailabilityConfig,
1030 ) -> ToolDefinition {
1031 ToolDefinition::raw_with_id(
1032 format!("tool:{name}"),
1033 name,
1034 description,
1035 ToolDefinition::default_input_schema(),
1036 json!({ "type": "string" }),
1037 )
1038 .with_availability(availability)
1039 }
1040
1041 fn manifests(definitions: Vec<ToolDefinition>) -> Vec<ToolManifest> {
1042 definitions
1043 .into_iter()
1044 .map(|tool| tool.manifest())
1045 .collect()
1046 }
1047
1048 fn contract_from(definitions: Vec<ToolDefinition>, name: &str) -> Option<Arc<ToolContract>> {
1049 definitions
1050 .into_iter()
1051 .find(|tool| tool.name() == name)
1052 .map(|tool| Arc::new(tool.contract()))
1053 }
1054
1055 fn dynamic_definition(name: &str) -> ToolDefinition {
1056 test_tool(name, "dynamic", crate::ToolAvailabilityConfig::callable())
1057 }
1058
1059 fn test_tool_context() -> crate::ToolContext<'static> {
1060 crate::ToolContext::builder(
1061 "registry-test".to_string(),
1062 Arc::new(crate::testing::MockSessionManager::default()),
1063 Arc::new(crate::testing::MockSessionManager::default()),
1064 Arc::new(crate::testing::MockSessionManager::default()),
1065 Arc::new(crate::UnavailableProcessService),
1066 Arc::new(crate::DefaultProcessCancelAbility),
1067 crate::runtime::RuntimeEffectControllerHandle::shared(Arc::new(
1068 crate::InlineRuntimeEffectController,
1069 )),
1070 Arc::new(crate::InMemoryAttachmentStore::new()),
1071 crate::DirectCompletionClient::unavailable(
1072 "direct completions are unavailable in this test context",
1073 ),
1074 )
1075 .build()
1076 }
1077
1078 #[async_trait::async_trait]
1079 impl ToolProvider for MockTool {
1080 fn tool_manifests(&self) -> Vec<ToolManifest> {
1081 manifests(vec![test_tool(
1082 "mock_tool",
1083 "mock",
1084 crate::ToolAvailabilityConfig::callable(),
1085 )])
1086 }
1087
1088 fn resolve_contract(&self, name: &str) -> Option<Arc<ToolContract>> {
1089 contract_from(
1090 vec![test_tool(
1091 "mock_tool",
1092 "mock",
1093 crate::ToolAvailabilityConfig::callable(),
1094 )],
1095 name,
1096 )
1097 }
1098
1099 async fn execute(&self, _call: ToolCall<'_>) -> ToolResult {
1100 ToolResult::ok(serde_json::json!("ok"))
1101 }
1102 }
1103
1104 #[async_trait::async_trait]
1105 impl ToolProvider for MixedEnabledTool {
1106 fn tool_manifests(&self) -> Vec<ToolManifest> {
1107 manifests(vec![
1108 test_tool(
1109 "enabled_tool",
1110 "enabled",
1111 crate::ToolAvailabilityConfig::callable(),
1112 ),
1113 test_tool(
1114 "disabled_tool",
1115 "disabled",
1116 crate::ToolAvailabilityConfig::off(),
1117 ),
1118 ])
1119 }
1120
1121 fn resolve_contract(&self, name: &str) -> Option<Arc<ToolContract>> {
1122 contract_from(
1123 vec![
1124 test_tool(
1125 "enabled_tool",
1126 "enabled",
1127 crate::ToolAvailabilityConfig::callable(),
1128 ),
1129 test_tool(
1130 "disabled_tool",
1131 "disabled",
1132 crate::ToolAvailabilityConfig::off(),
1133 ),
1134 ],
1135 name,
1136 )
1137 }
1138
1139 async fn execute(&self, _call: ToolCall<'_>) -> ToolResult {
1140 ToolResult::ok(serde_json::json!("ok"))
1141 }
1142 }
1143
1144 #[async_trait::async_trait]
1145 impl ToolSourceExecutor for ExternalMockSource {
1146 fn id(&self) -> &str {
1147 "external"
1148 }
1149
1150 fn advertised_tools(&self) -> Vec<ToolManifest> {
1151 manifests(vec![ToolDefinition::raw_with_id(
1152 "tool:mcp__demo__search",
1153 "mcp__demo__search",
1154 "search",
1155 json!({
1156 "type": "object",
1157 "properties": {
1158 "query": { "type": "string" }
1159 },
1160 "required": ["query"],
1161 "additionalProperties": false
1162 }),
1163 json!({ "type": "object", "additionalProperties": true }),
1164 )])
1165 }
1166
1167 fn resolve_contract(&self, name: &str) -> Option<Arc<ToolContract>> {
1168 contract_from(
1169 vec![ToolDefinition::raw_with_id(
1170 "tool:mcp__demo__search",
1171 "mcp__demo__search",
1172 "search",
1173 json!({
1174 "type": "object",
1175 "properties": {
1176 "query": { "type": "string" }
1177 },
1178 "required": ["query"],
1179 "additionalProperties": false
1180 }),
1181 json!({ "type": "object", "additionalProperties": true }),
1182 )],
1183 name,
1184 )
1185 }
1186
1187 async fn execute(
1188 &self,
1189 tool: &str,
1190 args: &serde_json::Value,
1191 _context: &ToolContext<'_>,
1192 _progress: Option<&ProgressSender>,
1193 ) -> ToolResult {
1194 ToolResult::ok(json!({
1195 "tool": tool,
1196 "args": args
1197 }))
1198 }
1199 }
1200
1201 #[async_trait::async_trait]
1202 impl ToolSourceExecutor for ExactResolvingSource {
1203 fn id(&self) -> &str {
1204 "exact"
1205 }
1206
1207 fn advertised_tools(&self) -> Vec<ToolManifest> {
1208 Vec::new()
1209 }
1210
1211 fn resolve_manifest(&self, name: &str) -> Option<ToolManifest> {
1212 self.manifest_resolutions.fetch_add(1, Ordering::SeqCst);
1213 (name == "host_only").then(|| {
1214 test_tool(
1215 "host_only",
1216 "host-only",
1217 crate::ToolAvailabilityConfig::callable(),
1218 )
1219 .manifest()
1220 })
1221 }
1222
1223 fn resolve_contract(&self, name: &str) -> Option<Arc<ToolContract>> {
1224 self.contract_resolutions.fetch_add(1, Ordering::SeqCst);
1225 contract_from(
1226 vec![test_tool(
1227 "host_only",
1228 "host-only",
1229 crate::ToolAvailabilityConfig::callable(),
1230 )],
1231 name,
1232 )
1233 }
1234
1235 async fn execute(
1236 &self,
1237 tool: &str,
1238 _args: &serde_json::Value,
1239 _context: &ToolContext<'_>,
1240 _progress: Option<&ProgressSender>,
1241 ) -> ToolResult {
1242 self.executions.fetch_add(1, Ordering::SeqCst);
1243 ToolResult::ok(json!(tool))
1244 }
1245 }
1246
1247 #[async_trait::async_trait]
1248 impl ToolSourceExecutor for NamedExactSource {
1249 fn id(&self) -> &str {
1250 self.id
1251 }
1252
1253 fn advertised_tools(&self) -> Vec<ToolManifest> {
1254 Vec::new()
1255 }
1256
1257 fn resolve_manifest(&self, name: &str) -> Option<ToolManifest> {
1258 (name == "host_only").then(|| {
1259 test_tool(
1260 "host_only",
1261 "host-only",
1262 crate::ToolAvailabilityConfig::callable(),
1263 )
1264 .manifest()
1265 })
1266 }
1267
1268 fn resolve_contract(&self, _name: &str) -> Option<Arc<ToolContract>> {
1269 None
1270 }
1271
1272 async fn execute(
1273 &self,
1274 tool: &str,
1275 _args: &serde_json::Value,
1276 _context: &ToolContext<'_>,
1277 _progress: Option<&ProgressSender>,
1278 ) -> ToolResult {
1279 ToolResult::ok(json!(tool))
1280 }
1281 }
1282
1283 #[async_trait::async_trait]
1284 impl ToolProvider for DynamicToolProvider {
1285 fn tool_manifests(&self) -> Vec<ToolManifest> {
1286 self.names
1287 .lock()
1288 .expect("dynamic tool names lock")
1289 .iter()
1290 .map(|name| dynamic_definition(name).manifest())
1291 .collect()
1292 }
1293
1294 fn resolve_contract(&self, name: &str) -> Option<Arc<ToolContract>> {
1295 self.names
1296 .lock()
1297 .expect("dynamic tool names lock")
1298 .iter()
1299 .any(|tool_name| tool_name == name)
1300 .then(|| Arc::new(dynamic_definition(name).contract()))
1301 }
1302
1303 async fn execute(&self, call: ToolCall<'_>) -> ToolResult {
1304 ToolResult::ok(json!(call.name))
1305 }
1306 }
1307
1308 #[test]
1309 fn registry_preserves_initial_availability_state() {
1310 let registry =
1311 ToolRegistry::from_tool_provider(Arc::new(MixedEnabledTool)).expect("registry");
1312 let snapshot = registry.export_state();
1313 assert_eq!(
1314 snapshot
1315 .get("enabled_tool")
1316 .unwrap()
1317 .manifest()
1318 .effective_availability(),
1319 crate::ToolAvailability::Callable
1320 );
1321 assert_eq!(
1322 snapshot
1323 .get("disabled_tool")
1324 .unwrap()
1325 .manifest()
1326 .effective_availability(),
1327 crate::ToolAvailability::Off
1328 );
1329 }
1330
1331 #[test]
1332 fn exported_tool_state_is_source_free() {
1333 let registry = ToolRegistry::from_tool_provider(Arc::new(MockTool)).expect("registry");
1334 registry
1335 .add_tool_provider(Arc::new(MixedEnabledTool))
1336 .expect("live provider registered");
1337
1338 let value = serde_json::to_value(registry.export_state()).expect("serialized tool state");
1339 let serialized = value.to_string();
1340
1341 assert!(!serialized.contains("source_id"));
1342 assert!(!serialized.contains(PLUGIN_SOURCE_ID));
1343 assert!(!serialized.contains("live:"));
1344 }
1345
1346 #[test]
1347 fn apply_state_rebinds_source_free_snapshot_to_current_sources() {
1348 let source_registry =
1349 ToolRegistry::from_tool_provider(Arc::new(MixedEnabledTool)).expect("source registry");
1350 let snapshot = source_registry.export_state();
1351
1352 let target_registry =
1353 ToolRegistry::from_tool_provider(Arc::new(MixedEnabledTool)).expect("target registry");
1354 let next_generation = target_registry
1355 .apply_state(snapshot.with_generation(target_registry.generation()))
1356 .expect("state rebound");
1357
1358 assert_eq!(next_generation, target_registry.generation());
1359 assert!(target_registry.resolve_contract("enabled_tool").is_some());
1360 }
1361
1362 #[test]
1363 fn apply_state_rejects_tools_not_advertised_by_source() {
1364 let registry = ToolRegistry::from_tool_provider(Arc::new(MockTool)).expect("registry");
1365 let snapshot = registry.export_state();
1366 let generation = snapshot.generation();
1367 let mut tools = snapshot.entries().clone();
1368 tools.insert(
1369 "missing".to_string(),
1370 ToolStateEntry::new(
1371 test_tool(
1372 "missing",
1373 "missing",
1374 crate::ToolAvailabilityConfig::callable(),
1375 )
1376 .manifest(),
1377 ),
1378 );
1379 let snapshot = ToolState::new(generation, tools);
1380 assert!(matches!(
1381 registry.apply_state(snapshot),
1382 Err(ReconfigureError::Validation(_))
1383 ));
1384 }
1385
1386 #[test]
1387 fn apply_state_rejects_snapshot_when_provider_is_absent() {
1388 let source_registry =
1389 ToolRegistry::from_tool_provider(Arc::new(MockTool)).expect("source registry");
1390 source_registry
1391 .upsert_source(Arc::new(ExternalMockSource))
1392 .expect("source registered");
1393 let snapshot = source_registry.export_state();
1394
1395 let target_registry =
1396 ToolRegistry::from_tool_provider(Arc::new(MockTool)).expect("target registry");
1397 let err = target_registry
1398 .apply_state(snapshot.with_generation(target_registry.generation()))
1399 .expect_err("missing provider should fail");
1400
1401 assert!(matches!(err, ReconfigureError::Validation(_)));
1402 }
1403
1404 #[test]
1405 fn apply_state_rejects_ambiguous_current_source_binding() {
1406 let registry = ToolRegistry::empty();
1407 registry
1408 .upsert_source(Arc::new(NamedExactSource { id: "exact-a" }))
1409 .expect("source a registered");
1410 registry
1411 .upsert_source(Arc::new(NamedExactSource { id: "exact-b" }))
1412 .expect("source b registered");
1413
1414 let mut tools = BTreeMap::new();
1415 tools.insert(
1416 "host_only".to_string(),
1417 ToolStateEntry::new(
1418 test_tool(
1419 "host_only",
1420 "host-only",
1421 crate::ToolAvailabilityConfig::callable(),
1422 )
1423 .manifest(),
1424 ),
1425 );
1426
1427 let err = registry
1428 .apply_state(ToolState::new(registry.generation(), tools))
1429 .expect_err("ambiguous source binding should fail");
1430
1431 assert!(matches!(err, ReconfigureError::Validation(_)));
1432 }
1433
1434 #[test]
1435 fn advertised_manifest_resolves_without_exact_host_lookup() {
1436 let manifest_resolutions = Arc::new(AtomicUsize::new(0));
1437 let registry = ToolRegistry::from_tool_provider(Arc::new(MockTool)).expect("registry");
1438 registry
1439 .upsert_source(Arc::new(ExactResolvingSource {
1440 manifest_resolutions: Arc::clone(&manifest_resolutions),
1441 contract_resolutions: Arc::new(AtomicUsize::new(0)),
1442 executions: Arc::new(AtomicUsize::new(0)),
1443 }))
1444 .expect("source registered");
1445
1446 assert_eq!(
1447 registry
1448 .resolve_manifest("mock_tool")
1449 .map(|manifest| manifest.name),
1450 Some("mock_tool".to_string())
1451 );
1452 assert_eq!(manifest_resolutions.load(Ordering::SeqCst), 0);
1453 }
1454
1455 #[test]
1456 fn refresh_sources_re_reads_group_provider_manifests() {
1457 let names = Arc::new(std::sync::Mutex::new(vec!["dynamic_one".to_string()]));
1458 let provider: Arc<dyn ToolProvider> = Arc::new(DynamicToolProvider {
1459 names: Arc::clone(&names),
1460 });
1461 let registry = ToolRegistry::from_tool_providers(vec![provider]).expect("registry");
1462
1463 let tool_names = || {
1464 registry
1465 .tool_manifests()
1466 .into_iter()
1467 .map(|manifest| manifest.name)
1468 .collect::<BTreeSet<_>>()
1469 };
1470
1471 assert!(tool_names().contains("dynamic_one"));
1472 assert!(!tool_names().contains("dynamic_two"));
1473
1474 names
1475 .lock()
1476 .expect("dynamic tool names lock")
1477 .push("dynamic_two".to_string());
1478 registry.refresh_sources().expect("refresh sources");
1479 let refreshed = tool_names();
1480 assert!(refreshed.contains("dynamic_one"));
1481 assert!(refreshed.contains("dynamic_two"));
1482
1483 names
1484 .lock()
1485 .expect("dynamic tool names lock")
1486 .retain(|name| name != "dynamic_one");
1487 registry.refresh_sources().expect("refresh sources");
1488 let refreshed = tool_names();
1489 assert!(!refreshed.contains("dynamic_one"));
1490 assert!(refreshed.contains("dynamic_two"));
1491 }
1492
1493 #[tokio::test]
1494 async fn unknown_manifest_exact_resolves_and_routes_to_owner() {
1495 let manifest_resolutions = Arc::new(AtomicUsize::new(0));
1496 let contract_resolutions = Arc::new(AtomicUsize::new(0));
1497 let executions = Arc::new(AtomicUsize::new(0));
1498 let registry = ToolRegistry::from_tool_provider(Arc::new(MockTool)).expect("registry");
1499 registry
1500 .upsert_source(Arc::new(ExactResolvingSource {
1501 manifest_resolutions: Arc::clone(&manifest_resolutions),
1502 contract_resolutions: Arc::clone(&contract_resolutions),
1503 executions: Arc::clone(&executions),
1504 }))
1505 .expect("source registered");
1506
1507 assert_eq!(
1508 registry
1509 .resolve_manifest("host_only")
1510 .map(|manifest| manifest.name),
1511 Some("host_only".to_string())
1512 );
1513 assert_eq!(manifest_resolutions.load(Ordering::SeqCst), 1);
1514
1515 let contract = registry.resolve_contract("host_only");
1516 assert!(contract.is_some());
1517 assert_eq!(manifest_resolutions.load(Ordering::SeqCst), 1);
1518 assert_eq!(contract_resolutions.load(Ordering::SeqCst), 1);
1519
1520 let context = test_tool_context();
1521 let args = json!({});
1522 let result = registry
1523 .execute(crate::ToolCall {
1524 name: "host_only",
1525 args: &args,
1526 context: &context,
1527 progress: None,
1528 })
1529 .await;
1530 assert!(result.is_success());
1531 assert_eq!(result.value_for_projection(), json!("host_only"));
1532 assert_eq!(executions.load(Ordering::SeqCst), 1);
1533 }
1534
1535 #[test]
1536 fn unknown_manifest_without_host_resolver_is_unavailable() {
1537 let registry = ToolRegistry::from_tool_provider(Arc::new(MockTool)).expect("registry");
1538
1539 assert!(registry.resolve_manifest("missing").is_none());
1540 assert!(registry.resolve_contract("missing").is_none());
1541 }
1542
1543 #[tokio::test]
1544 async fn upsert_source_registers_and_executes_external_tools() {
1545 let registry = ToolRegistry::from_tool_provider(Arc::new(MockTool)).expect("registry");
1546 registry
1547 .upsert_source(Arc::new(ExternalMockSource))
1548 .expect("source registered");
1549
1550 let defs = registry.tool_manifests();
1551 assert!(defs.iter().any(|def| def.name == "mcp__demo__search"));
1552
1553 let context = test_tool_context();
1554 let args = json!({ "query": "hello" });
1555 let result = registry
1556 .execute(crate::ToolCall {
1557 name: "mcp__demo__search",
1558 args: &args,
1559 context: &context,
1560 progress: None,
1561 })
1562 .await;
1563 assert!(result.is_success());
1564 assert_eq!(
1565 result.value_for_projection()["tool"],
1566 json!("mcp__demo__search")
1567 );
1568 assert_eq!(
1569 result.value_for_projection()["args"]["query"],
1570 json!("hello")
1571 );
1572 }
1573
1574 #[test]
1575 fn upsert_source_preserves_availability_override_on_refresh() {
1576 let registry = ToolRegistry::from_tool_provider(Arc::new(MockTool)).expect("registry");
1577 registry
1578 .upsert_source(Arc::new(ExternalMockSource))
1579 .expect("source registered");
1580 let mut snapshot = registry.export_state();
1581 snapshot
1582 .set_availability("mcp__demo__search", Some(crate::ToolAvailability::Off))
1583 .unwrap();
1584 registry.apply_state(snapshot).unwrap();
1585 registry
1586 .upsert_source(Arc::new(ExternalMockSource))
1587 .expect("source refreshed");
1588 let snapshot = registry.export_state();
1589 assert_eq!(
1590 snapshot
1591 .get("mcp__demo__search")
1592 .unwrap()
1593 .manifest()
1594 .effective_availability(),
1595 crate::ToolAvailability::Off
1596 );
1597 }
1598
1599 #[test]
1600 fn restore_state_adopts_generation_at_or_above_three() {
1601 let source = ToolRegistry::from_tool_provider(Arc::new(MockTool)).expect("source registry");
1607 let snapshot = source.export_state().with_generation(3);
1608
1609 let target = ToolRegistry::from_tool_provider(Arc::new(MockTool)).expect("target registry");
1610 assert_eq!(
1611 target.generation(),
1612 1,
1613 "a fresh registry starts at generation 1"
1614 );
1615 let restored = target
1616 .restore_state(snapshot.clone())
1617 .expect("restore adopts the snapshot generation");
1618 assert_eq!(restored, 3, "restore returns the adopted generation");
1619 assert_eq!(
1620 target.generation(),
1621 3,
1622 "restore adopts gen 3 onto a base-1 registry without bumping"
1623 );
1624 assert_eq!(target.export_state().generation(), 3);
1626
1627 let fresh = ToolRegistry::from_tool_provider(Arc::new(MockTool)).expect("fresh registry");
1630 assert!(
1631 matches!(
1632 fresh.apply_state(snapshot),
1633 Err(ReconfigureError::GenerationMismatch {
1634 expected: 3,
1635 actual: 1
1636 })
1637 ),
1638 "apply_state must reject a gen-3 snapshot on a base-1 registry"
1639 );
1640 }
1641
1642 #[test]
1643 fn remove_source_removes_all_source_tools() {
1644 let registry = ToolRegistry::from_tool_provider(Arc::new(MockTool)).expect("registry");
1645 registry
1646 .upsert_source(Arc::new(ExternalMockSource))
1647 .expect("source registered");
1648 registry
1649 .remove_source_id("external")
1650 .expect("source removed");
1651 let defs = registry.tool_manifests();
1652 assert!(!defs.iter().any(|def| def.name == "mcp__demo__search"));
1653 }
1654
1655 #[test]
1656 fn project_tool_catalog_keeps_searchable_tools_with_surface_metadata() {
1657 fn dummy_tool(name: &str) -> crate::ToolDefinition {
1658 let tool = crate::ToolDefinition::raw_with_id(
1659 format!("tool:{name}"),
1660 name,
1661 format!("desc for {name}"),
1662 crate::ToolDefinition::default_input_schema(),
1663 serde_json::json!({}),
1664 );
1665 match name {
1666 "read_file" => {
1667 tool.with_agent_surface(crate::ToolAgentSurface::new(["files"], "read"))
1668 }
1669 "search_tools" => {
1670 tool.with_agent_surface(crate::ToolAgentSurface::new(["tools"], "search"))
1671 }
1672 _ => tool,
1673 }
1674 }
1675 let catalog = project_tool_catalog([
1676 crate::ToolSurfaceEntry {
1677 manifest: dummy_tool("read_file").manifest(),
1678 availability: crate::ToolAvailability::Showcased,
1679 },
1680 crate::ToolSurfaceEntry {
1681 manifest: dummy_tool("search_tools").manifest(),
1682 availability: crate::ToolAvailability::Callable,
1683 },
1684 ]);
1685 assert_eq!(catalog.len(), 2);
1686 assert_eq!(catalog[0]["name"], serde_json::json!("read_file"));
1687 assert_eq!(
1688 catalog[0]["contract"]["signature"],
1689 serde_json::json!("await files.read({})?")
1690 );
1691 assert_eq!(catalog[0]["showcased"], serde_json::json!(true));
1692 assert_eq!(catalog[1]["callable"], serde_json::json!(true));
1693 }
1694
1695 #[test]
1696 fn project_tool_catalog_preserves_dynamic_output_contracts() {
1697 fn dummy_tool(name: &str) -> crate::ToolDefinition {
1698 crate::ToolDefinition::raw_with_id(
1699 format!("tool:{name}"),
1700 name,
1701 format!("desc for {name}"),
1702 crate::ToolDefinition::default_input_schema(),
1703 serde_json::json!({}),
1704 )
1705 .with_agent_surface(crate::ToolAgentSurface::new(["llm"], "query"))
1706 }
1707 let catalog = project_tool_catalog([crate::ToolSurfaceEntry {
1708 manifest: dummy_tool("llm_query")
1709 .with_output_from_input_schema(
1710 "output",
1711 Some(serde_json::json!({ "type": "string" })),
1712 )
1713 .manifest(),
1714 availability: crate::ToolAvailability::Searchable,
1715 }]);
1716
1717 assert_eq!(
1718 catalog[0]["contract"]["signature"],
1719 serde_json::json!("await llm.query<T = str>({})?")
1720 );
1721 assert_eq!(catalog[0]["contract"]["returns"], serde_json::json!("T"));
1722 }
1723}
1724
1725pub(crate) fn project_tool_catalog<I>(entries: I) -> Vec<serde_json::Value>
1726where
1727 I: IntoIterator<Item = crate::ToolSurfaceEntry>,
1728{
1729 entries
1730 .into_iter()
1731 .filter(|entry| entry.availability.is_searchable())
1732 .map(|entry| {
1733 let manifest = entry.manifest;
1734 let availability = entry.availability;
1735 let agent_surface = manifest.agent_surface.executable_for(&manifest.name);
1736 let call = agent_surface.call_path();
1737 let mut projected = serde_json::json!({
1738 "id": manifest.id,
1739 "name": manifest.name,
1740 "module_path": agent_surface.module_path,
1741 "operation": agent_surface.operation,
1742 "authority_type": agent_surface.authority_type,
1743 "call": call,
1744 "description": manifest.description,
1745 "aliases": agent_surface.aliases,
1746 "availability": availability,
1747 "callable": availability.is_callable(),
1748 "showcased": availability.is_showcased(),
1749 "searchable": availability.is_searchable(),
1750 "activation": manifest.activation,
1751 });
1752 if let Some(contract) = manifest.compact_contract {
1753 projected
1754 .as_object_mut()
1755 .expect("projected tool catalog entry is an object")
1756 .insert("contract".to_string(), serde_json::json!(contract));
1757 }
1758 projected
1759 })
1760 .collect()
1761}