Skip to main content

lash_core/
tool_registry.rs

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    /// Restore a persisted [`ToolState`] snapshot onto a freshly-built registry,
493    /// adopting the snapshot's generation verbatim.
494    ///
495    /// Unlike [`apply_state`](Self::apply_state) — which applies an incremental
496    /// *delta* expected at the current generation and bumps it by one — a
497    /// restore reconstructs the exact persisted state regardless of the fresh
498    /// registry's base generation, and does **not** bump. This is idempotent: a
499    /// snapshot exported at generation `G` restores to generation `G`, so a
500    /// re-export round-trips. Cold rebuilds (the durable process worker, session
501    /// resume) restore a session whose tool surface reached generation `G ≥ 2`
502    /// onto a base registry at generation 1 — `apply_state` would reject that
503    /// (`expected G, actual 1`); `restore_state` adopts `G`. Entries are still
504    /// rebound to the live sources, so source identity is reconnected.
505    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        // Cold rebuild ratchet: a session whose tool surface advanced to
1602        // generation >= 3 restores onto a fresh base-1 registry. `restore_state`
1603        // adopts the snapshot's generation verbatim; `apply_state` (a gen-matched
1604        // delta) rejects it. This is the exact divergence the durable worker /
1605        // session resume rebuild relies on `restore_state` to absorb.
1606        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        // A re-export round-trips at the same generation (idempotent).
1625        assert_eq!(target.export_state().generation(), 3);
1626
1627        // apply_state on the same high-generation snapshot is rejected — proving
1628        // the rebuild would have failed without restore_state.
1629        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}