1use std::collections::{BTreeSet, HashMap};
2use std::path::{Path, PathBuf};
3use std::sync::{Arc, LazyLock, Mutex, Weak};
4
5use crate::auth_store::{ProviderKind, StoredProvider};
6use crate::config_hub::{
7 AuthModelCacheCommit, AuthModelCacheUpdate, AuthProviderInsertCommit,
8 AuthProviderRuntimeCommit, AuthProviderRuntimeState, ConfigError, ConfigHub,
9 ProviderConfigUpdate, ProviderConfigWriteMode,
10};
11use crate::model_registry::{
12 CatalogDelta, CatalogError, PreparedProviderCatalog, ProviderDescriptor,
13 commit_prepared_provider_catalog, prepare_provider_catalog, provider_catalog_namespace,
14 remove_provider_catalog, shortest_unique_provider_id,
15};
16use crate::provider::{
17 ModelDiscoveryError, Provider, ProviderRegistry, ReasoningWireProfile, WeakProviderRegistry,
18};
19
20#[derive(Debug, thiserror::Error)]
21#[non_exhaustive]
22pub enum ProviderLifecycleError {
23 #[error(transparent)]
24 Config(#[from] ConfigError),
25 #[error(transparent)]
26 Catalog(#[from] CatalogError),
27 #[error(transparent)]
28 Discovery(#[from] ModelDiscoveryError),
29 #[error("auth provider `{id}` does not exist")]
30 ProviderNotFound { id: String },
31 #[error("auth provider `{id}` is disabled")]
32 ProviderDisabled { id: String },
33 #[error("auth provider `{id}` has no live provider")]
34 LiveProviderMissing { id: String },
35 #[error("live provider name `{actual}` does not match auth provider id `{expected}`")]
36 LiveProviderNameMismatch { expected: String, actual: String },
37 #[error("auth provider `{id}` changed while the operation was running")]
38 Stale { id: String },
39}
40
41#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
42pub struct ProviderStateChange {
43 pub auth_changed: bool,
44 pub live_changed: bool,
45 pub catalog_changed: bool,
46}
47
48#[derive(Debug, Clone, Default, PartialEq, Eq)]
49pub struct ProviderLifecycleOutcome {
50 pub state: ProviderStateChange,
51 pub catalog: Option<CatalogDelta>,
52}
53
54#[derive(Debug, Clone, Default, PartialEq, Eq)]
55pub struct ProviderReconcileOutcome {
56 pub providers: Vec<(String, ProviderStateChange)>,
57}
58
59#[derive(Debug, Clone, Copy, PartialEq, Eq)]
61#[non_exhaustive]
62pub enum ProviderCatalogRefreshOutcome {
63 NotNeeded,
65 AlreadyInFlight,
67 CatalogUpdated(CatalogDelta),
69}
70
71#[derive(Clone)]
72pub struct ProviderLifecycle {
73 hub: ConfigHub,
74 providers: ProviderRegistry,
75 state: Arc<Mutex<LifecycleState>>,
76 config_state: Arc<Mutex<ConfigProviderLifecycleState>>,
77}
78
79#[derive(Clone)]
80pub(crate) struct ProviderLifecycleOwner {
81 hub: ConfigHub,
82 state: Arc<Mutex<LifecycleState>>,
83 config_state: Arc<Mutex<ConfigProviderLifecycleState>>,
84}
85
86#[derive(Default)]
87struct LifecycleState {
88 generations: HashMap<String, u64>,
89 operation_locks: HashMap<String, Arc<tokio::sync::Mutex<()>>>,
90 refresh_locks: HashMap<String, Arc<tokio::sync::Mutex<()>>>,
91 live_providers: HashMap<String, LiveProviderEntry>,
92 registries: Vec<WeakProviderRegistry>,
93}
94
95#[derive(Default)]
96struct ConfigProviderLifecycleState {
97 config_providers: HashMap<String, crate::model_registry::ProviderEntry>,
98 registries: Vec<WeakProviderRegistry>,
99}
100
101#[derive(Clone)]
102struct LiveProviderEntry {
103 kind: ProviderKind,
104 provider: Arc<dyn Provider>,
105}
106
107struct LiveProviderRegistration {
108 changed: bool,
109 replaced: Vec<Arc<dyn Provider>>,
110}
111
112struct LiveProviderRemoval {
113 changed: bool,
114 removed: Vec<Arc<dyn Provider>>,
115}
116
117struct ProviderOperationFence {
118 operation: Option<tokio::sync::OwnedMutexGuard<()>>,
119 providers: Vec<Arc<dyn Provider>>,
120}
121
122struct ProviderRefreshAttempt {
123 generation: u64,
124 runtime: AuthProviderRuntimeState,
125 live_provider: Arc<dyn Provider>,
126}
127
128enum ProviderRefreshPreparation {
129 Hydrated { delta: CatalogDelta, changed: bool },
130 Discover(Box<ProviderRefreshAttempt>),
131}
132
133type ProviderRefreshPreparationResult = (
134 Result<ProviderRefreshPreparation, ProviderLifecycleError>,
135 Vec<Arc<dyn Provider>>,
136);
137
138type ProviderRefreshAttemptResult = (
139 Result<ProviderRefreshAttempt, ProviderLifecycleError>,
140 Vec<Arc<dyn Provider>>,
141);
142
143type ProviderCatalogHydrationResult = (
144 Result<(CatalogDelta, bool), ProviderLifecycleError>,
145 Vec<Arc<dyn Provider>>,
146);
147
148impl ProviderOperationFence {
149 fn new(operation: tokio::sync::OwnedMutexGuard<()>) -> Self {
150 Self {
151 operation: Some(operation),
152 providers: Vec::new(),
153 }
154 }
155
156 fn protect(&mut self, provider: Arc<dyn Provider>) {
157 self.providers.push(provider);
158 }
159
160 fn extend(&mut self, providers: Vec<Arc<dyn Provider>>) {
161 self.providers.extend(providers);
162 }
163}
164
165impl Drop for ProviderOperationFence {
166 fn drop(&mut self) {
167 drop(self.operation.take());
168 self.providers.clear();
169 }
170}
171
172static LIFECYCLE_COORDINATORS: LazyLock<Mutex<HashMap<PathBuf, Weak<Mutex<LifecycleState>>>>> =
173 LazyLock::new(|| Mutex::new(HashMap::new()));
174static CONFIG_PROVIDER_COORDINATORS: LazyLock<
175 Mutex<HashMap<PathBuf, Weak<Mutex<ConfigProviderLifecycleState>>>>,
176> = LazyLock::new(|| Mutex::new(HashMap::new()));
177
178impl ProviderLifecycle {
179 pub fn new(hub: ConfigHub, providers: ProviderRegistry) -> Self {
180 let (lifecycle, replaced) = Self::new_deferred(hub, providers);
181 drop(replaced);
182 lifecycle
183 }
184
185 pub(crate) fn new_deferred(
186 hub: ConfigHub,
187 providers: ProviderRegistry,
188 ) -> (Self, Vec<Arc<dyn Provider>>) {
189 let state = shared_lifecycle_state(&hub);
190 let config_state = shared_config_provider_state(&hub);
191 let mut replaced = Vec::new();
192 {
193 let mut shared = state
194 .lock()
195 .unwrap_or_else(std::sync::PoisonError::into_inner);
196 let mut registered = false;
197 shared.registries.retain(|registry| {
198 let Some(registry) = registry.upgrade() else {
199 return false;
200 };
201 registered |= registry.shares_storage_with(&providers);
202 true
203 });
204 if !registered {
205 for (provider_id, live) in &shared.live_providers {
206 replaced.extend(
207 providers.register_named(provider_id.clone(), live.provider.clone()),
208 );
209 }
210 shared.registries.push(providers.downgrade());
211 }
212 }
213 {
214 let mut shared = config_state
215 .lock()
216 .unwrap_or_else(std::sync::PoisonError::into_inner);
217 let mut registered = false;
218 shared.registries.retain(|registry| {
219 let Some(registry) = registry.upgrade() else {
220 return false;
221 };
222 registered |= registry.shares_storage_with(&providers);
223 true
224 });
225 if !registered {
226 let mut config_providers = shared
227 .config_providers
228 .iter()
229 .map(|(name, entry)| (name.clone(), entry.clone()))
230 .collect::<Vec<_>>();
231 config_providers.sort_by(|left, right| left.0.cmp(&right.0));
232 for (name, entry) in config_providers {
233 let (_, previous) = crate::config_provider::reconcile_config_provider_deferred(
234 &providers, &name, &entry,
235 );
236 replaced.extend(previous);
237 }
238 shared.registries.push(providers.downgrade());
239 }
240 }
241 (
242 Self {
243 hub,
244 providers,
245 state,
246 config_state,
247 },
248 replaced,
249 )
250 }
251
252 pub fn config_hub(&self) -> &ConfigHub {
253 &self.hub
254 }
255
256 pub fn provider_registry(&self) -> &ProviderRegistry {
257 &self.providers
258 }
259
260 pub(crate) fn owner(&self) -> ProviderLifecycleOwner {
261 ProviderLifecycleOwner {
262 hub: self.hub.clone(),
263 state: self.state.clone(),
264 config_state: self.config_state.clone(),
265 }
266 }
267
268 pub(crate) fn from_owner(owner: ProviderLifecycleOwner, providers: ProviderRegistry) -> Self {
269 Self {
270 hub: owner.hub,
271 providers,
272 state: owner.state,
273 config_state: owner.config_state,
274 }
275 }
276
277 pub fn reload_config_providers(&self) -> Result<(), ProviderLifecycleError> {
279 reload_config_providers_for_hub(&self.hub)?;
280 Ok(())
281 }
282
283 pub fn create_config_provider(
285 &self,
286 update: ProviderConfigUpdate<'_>,
287 ) -> Result<(), ProviderLifecycleError> {
288 self.commit_config_provider(update, ProviderConfigWriteMode::Create)
289 }
290
291 pub fn update_config_provider(
293 &self,
294 update: ProviderConfigUpdate<'_>,
295 ) -> Result<(), ProviderLifecycleError> {
296 self.commit_config_provider(update, ProviderConfigWriteMode::Update)
297 }
298
299 pub fn remove_config_provider(&self, name: &str) -> Result<(), ProviderLifecycleError> {
300 remove_config_provider_for_hub(&self.hub, name)?;
301 Ok(())
302 }
303
304 fn commit_config_provider(
305 &self,
306 update: ProviderConfigUpdate<'_>,
307 mode: ProviderConfigWriteMode,
308 ) -> Result<(), ProviderLifecycleError> {
309 mutate_config_provider_for_hub(&self.hub, update, mode)?;
310 Ok(())
311 }
312
313 pub async fn install_provider(
314 &self,
315 provider_record: StoredProvider,
316 live_provider: Arc<dyn Provider>,
317 ) -> Result<CatalogDelta, ProviderLifecycleError> {
318 self.validate_new_provider(&provider_record, &live_provider)?;
319 let provider_id = provider_record.id.clone();
320 let operation_lock = self.operation_lock(&provider_id);
321 let mut fence = ProviderOperationFence::new(operation_lock.lock_owned().await);
322 fence.protect(live_provider.clone());
323 let generation = self.generation(&provider_id);
324 let result = match live_provider.try_discover_models().await {
325 Ok(models) => {
326 self.commit_new_provider(provider_record, live_provider, &models, generation)
327 }
328 Err(error) => Err(error.into()),
329 };
330 match result {
331 Ok((delta, replaced)) => {
332 fence.extend(replaced);
333 drop(fence);
334 Ok(delta)
335 }
336 Err(error) => {
337 drop(fence);
338 Err(error)
339 }
340 }
341 }
342
343 pub async fn install_pre_discovered_provider(
344 &self,
345 provider_record: StoredProvider,
346 live_provider: Arc<dyn Provider>,
347 models: Vec<crate::provider::DiscoveredModelDetails>,
348 ) -> Result<CatalogDelta, ProviderLifecycleError> {
349 self.validate_new_provider(&provider_record, &live_provider)?;
350 let provider_id = provider_record.id.clone();
351 let operation_lock = self.operation_lock(&provider_id);
352 let mut fence = ProviderOperationFence::new(operation_lock.lock_owned().await);
353 fence.protect(live_provider.clone());
354 let generation = self.generation(&provider_id);
355 let result = self.commit_new_provider(provider_record, live_provider, &models, generation);
356 match result {
357 Ok((delta, replaced)) => {
358 fence.extend(replaced);
359 drop(fence);
360 Ok(delta)
361 }
362 Err(error) => {
363 drop(fence);
364 Err(error)
365 }
366 }
367 }
368
369 fn commit_new_provider(
370 &self,
371 provider_record: StoredProvider,
372 live_provider: Arc<dyn Provider>,
373 models: &[crate::provider::DiscoveredModelDetails],
374 generation: u64,
375 ) -> Result<(CatalogDelta, Vec<Arc<dyn Provider>>), ProviderLifecycleError> {
376 let provider_id = provider_record.id.clone();
377 let mut state = self.lock_state();
378 self.ensure_generation(&state, &provider_id, generation)?;
379 let mut store = self.hub.load_auth()?;
380 if store
381 .providers
382 .iter()
383 .any(|provider| provider.id == provider_id)
384 {
385 return Err(ConfigError::Invalid(format!(
386 "auth provider id {provider_id:?} already exists"
387 ))
388 .into());
389 }
390 let mut expected_provider_ids = store
391 .providers
392 .iter()
393 .map(|provider| provider.id.clone())
394 .collect::<Vec<_>>();
395 expected_provider_ids.sort();
396 store.providers.push(provider_record.clone());
397 let descriptor = self.provider_descriptor(&provider_record, &store.providers)?;
398 let prepared = prepare_provider_catalog(descriptor, models)?;
399 let namespace = prepared.namespace().to_string();
400 let provider_kind = provider_record.kind.clone();
401 let (commit, applied) = self
402 .hub
403 .add_auth_provider_with_model_cache_details_if_provider_ids_and_then(
404 provider_record,
405 Some(&expected_provider_ids),
406 &namespace,
407 chrono::Utc::now().timestamp(),
408 models,
409 || {
410 let registration = Self::register_live_provider(
411 &mut state,
412 &provider_id,
413 provider_kind,
414 live_provider,
415 );
416 (
417 commit_prepared_provider_catalog(prepared),
418 registration.replaced,
419 )
420 },
421 )?;
422 match (commit, applied) {
423 (AuthProviderInsertCommit::Inserted, Some((delta, replaced_providers))) => {
424 Self::bump_generation(&mut state, &provider_id);
425 drop(state);
426 Ok((delta, replaced_providers))
427 }
428 (AuthProviderInsertCommit::Changed, _) => {
429 drop(state);
430 Err(ProviderLifecycleError::Stale { id: provider_id })
431 }
432 (AuthProviderInsertCommit::Inserted, None) => {
433 unreachable!("provider insertion callback was not run")
434 }
435 }
436 }
437
438 pub fn catalog_refresh_plan(&self) -> Vec<String> {
440 let state = self.lock_state();
441 let mut plan = state.live_providers.keys().cloned().collect::<Vec<_>>();
442 plan.sort();
443 plan
444 }
445
446 pub async fn refresh_models(
447 &self,
448 provider_id: &str,
449 ) -> Result<CatalogDelta, ProviderLifecycleError> {
450 let refresh_lock = self.refresh_lock(provider_id);
451 let _refresh = refresh_lock.lock_owned().await;
452 match self.run_catalog_refresh(provider_id, false).await? {
453 ProviderCatalogRefreshOutcome::CatalogUpdated(delta) => Ok(delta),
454 ProviderCatalogRefreshOutcome::NotNeeded
455 | ProviderCatalogRefreshOutcome::AlreadyInFlight => {
456 unreachable!("forced catalog refresh did not run")
457 }
458 }
459 }
460
461 pub async fn refresh_models_if_stale(
463 &self,
464 provider_id: &str,
465 ) -> Result<ProviderCatalogRefreshOutcome, ProviderLifecycleError> {
466 let refresh_lock = self.refresh_lock(provider_id);
467 let Ok(_refresh) = refresh_lock.try_lock_owned() else {
468 return Ok(ProviderCatalogRefreshOutcome::AlreadyInFlight);
469 };
470 self.run_catalog_refresh(provider_id, true).await
471 }
472
473 async fn run_catalog_refresh(
474 &self,
475 provider_id: &str,
476 only_if_stale: bool,
477 ) -> Result<ProviderCatalogRefreshOutcome, ProviderLifecycleError> {
478 let operation_lock = self.operation_lock(provider_id);
479 let mut fence = ProviderOperationFence::new(operation_lock.lock_owned().await);
480 let (preparation, deferred) = self.prepare_catalog_refresh(
481 provider_id,
482 chrono::Utc::now().timestamp(),
483 only_if_stale,
484 None,
485 None,
486 );
487 fence.extend(deferred);
488 let attempt = match preparation? {
489 ProviderRefreshPreparation::Hydrated { delta, changed } => {
490 drop(fence);
491 return Ok(if changed {
492 ProviderCatalogRefreshOutcome::CatalogUpdated(delta)
493 } else {
494 ProviderCatalogRefreshOutcome::NotNeeded
495 });
496 }
497 ProviderRefreshPreparation::Discover(attempt) => attempt,
498 };
499 drop(fence);
500
501 let discovery = attempt.live_provider.try_discover_models().await;
502 let operation_lock = self.operation_lock(provider_id);
503 let mut fence = ProviderOperationFence::new(operation_lock.lock_owned().await);
504 fence.protect(attempt.live_provider.clone());
505 let (result, deferred) = match discovery {
506 Ok(models) => self.finish_catalog_refresh(provider_id, &attempt, &models),
507 Err(error) => {
508 self.reconcile_refresh_failure(provider_id, &attempt, error.into(), Vec::new())
509 }
510 };
511 fence.extend(deferred);
512 drop(fence);
513 result.map(ProviderCatalogRefreshOutcome::CatalogUpdated)
514 }
515
516 fn prepare_catalog_refresh(
517 &self,
518 provider_id: &str,
519 now: i64,
520 hydrate_if_fresh: bool,
521 expected_live_provider: Option<&Arc<dyn Provider>>,
522 expected_catalog_snapshot: Option<&crate::auth_store::AuthProviderCatalogSnapshot>,
523 ) -> ProviderRefreshPreparationResult {
524 for _ in 0..3 {
525 let (attempt, mut deferred) = self.load_catalog_refresh_attempt(provider_id, now);
526 let attempt = match attempt {
527 Ok(attempt) => attempt,
528 Err(error) => return (Err(error), deferred),
529 };
530 if expected_live_provider
531 .is_some_and(|expected| !Arc::ptr_eq(expected, &attempt.live_provider))
532 {
533 return (
534 Err(ProviderLifecycleError::Stale {
535 id: provider_id.to_string(),
536 }),
537 deferred,
538 );
539 }
540 if !hydrate_if_fresh
541 || attempt.runtime.model_cache_freshness
542 != crate::auth_store::ModelCacheFreshness::Fresh
543 || expected_catalog_snapshot
544 .is_some_and(|expected| expected == &attempt.runtime.catalog_snapshot)
545 {
546 return (
547 Ok(ProviderRefreshPreparation::Discover(Box::new(attempt))),
548 deferred,
549 );
550 }
551 let (hydrated, hydrate_deferred) =
552 self.hydrate_cached_catalog_if_current(provider_id, &attempt);
553 deferred.extend(hydrate_deferred);
554 match hydrated {
555 Ok((delta, changed)) => {
556 return (
557 Ok(ProviderRefreshPreparation::Hydrated { delta, changed }),
558 deferred,
559 );
560 }
561 Err(ProviderLifecycleError::Stale { .. }) => continue,
562 Err(error) => return (Err(error), deferred),
563 }
564 }
565 let (authoritative, deferred) = self.load_catalog_refresh_attempt(provider_id, now);
566 match authoritative {
567 Ok(attempt)
568 if expected_live_provider
569 .is_none_or(|expected| Arc::ptr_eq(expected, &attempt.live_provider)) =>
570 {
571 (
572 Err(ProviderLifecycleError::Stale {
573 id: provider_id.to_string(),
574 }),
575 deferred,
576 )
577 }
578 Ok(_) => (
579 Err(ProviderLifecycleError::Stale {
580 id: provider_id.to_string(),
581 }),
582 deferred,
583 ),
584 Err(error) => (Err(error), deferred),
585 }
586 }
587
588 fn load_catalog_refresh_attempt(
589 &self,
590 provider_id: &str,
591 now: i64,
592 ) -> ProviderRefreshAttemptResult {
593 let mut state = self.lock_state();
594 let runtime = self
595 .hub
596 .load_or_create_auth_provider_runtime_state_at(provider_id, now);
597 let runtime = match runtime {
598 Ok(Some(runtime)) => runtime,
599 Ok(None) => {
600 Self::bump_generation(&mut state, provider_id);
601 let (_, removed) = Self::remove_runtime_provider(&mut state, provider_id, false);
602 drop(state);
603 return (
604 Err(ProviderLifecycleError::ProviderNotFound {
605 id: provider_id.to_string(),
606 }),
607 removed,
608 );
609 }
610 Err(error) => {
611 Self::bump_generation(&mut state, provider_id);
612 let (_, removed) = Self::remove_runtime_provider(&mut state, provider_id, false);
613 drop(state);
614 return (Err(error.into()), removed);
615 }
616 };
617 if !runtime.provider.enabled {
618 Self::bump_generation(&mut state, provider_id);
619 let (_, removed) = Self::remove_runtime_provider(&mut state, provider_id, false);
620 drop(state);
621 return (
622 Err(ProviderLifecycleError::ProviderDisabled {
623 id: provider_id.to_string(),
624 }),
625 removed,
626 );
627 }
628 let Some(live) = state.live_providers.get(provider_id) else {
629 Self::bump_generation(&mut state, provider_id);
630 let (_, removed) = Self::remove_runtime_provider(&mut state, provider_id, false);
631 drop(state);
632 return (
633 Err(ProviderLifecycleError::LiveProviderMissing {
634 id: provider_id.to_string(),
635 }),
636 removed,
637 );
638 };
639 if live.kind != runtime.provider.kind {
640 Self::bump_generation(&mut state, provider_id);
641 let (_, removed) = Self::remove_runtime_provider(&mut state, provider_id, false);
642 drop(state);
643 return (
644 Err(ProviderLifecycleError::Stale {
645 id: provider_id.to_string(),
646 }),
647 removed,
648 );
649 }
650 let attempt = ProviderRefreshAttempt {
651 generation: Self::generation_in(&state, provider_id),
652 runtime,
653 live_provider: live.provider.clone(),
654 };
655 (Ok(attempt), Vec::new())
656 }
657
658 fn hydrate_cached_catalog_if_current(
659 &self,
660 provider_id: &str,
661 attempt: &ProviderRefreshAttempt,
662 ) -> ProviderCatalogHydrationResult {
663 let mut state = self.lock_state();
664 if let Err(error) = self.ensure_generation(&state, provider_id, attempt.generation) {
665 return (Err(error), Vec::new());
666 }
667 let Some(current_live) = state.live_providers.get(provider_id) else {
668 return (
669 Err(ProviderLifecycleError::Stale {
670 id: provider_id.to_string(),
671 }),
672 Vec::new(),
673 );
674 };
675 if current_live.kind != attempt.runtime.provider.kind
676 || !Arc::ptr_eq(&attempt.live_provider, ¤t_live.provider)
677 {
678 return (
679 Err(ProviderLifecycleError::Stale {
680 id: provider_id.to_string(),
681 }),
682 Vec::new(),
683 );
684 }
685 let Some(models) = attempt.runtime.model_cache.as_deref() else {
686 return (
687 Err(ProviderLifecycleError::Stale {
688 id: provider_id.to_string(),
689 }),
690 Vec::new(),
691 );
692 };
693 let descriptor = match self.provider_descriptor_from_state(
694 &attempt.runtime.provider,
695 &attempt.runtime.provider_ids,
696 attempt.runtime.model_namespace.as_deref(),
697 ) {
698 Ok(descriptor) => descriptor,
699 Err(error) => return (Err(error), Vec::new()),
700 };
701 let prepared = match prepare_provider_catalog(descriptor, models) {
702 Ok(prepared) => prepared,
703 Err(error) => return (Err(error.into()), Vec::new()),
704 };
705 let namespace = prepared.namespace().to_string();
706 let revision = crate::model_registry::model_catalog_revision();
707 let expected_provider_ids = attempt
708 .runtime
709 .model_namespace
710 .is_none()
711 .then_some(attempt.runtime.provider_ids.as_slice());
712 let commit = self.hub.commit_auth_provider_runtime_if_current_and_then(
713 &attempt.runtime.provider,
714 &attempt.runtime.catalog_snapshot,
715 false,
716 Some(&namespace),
717 expected_provider_ids,
718 || {
719 let delta = commit_prepared_provider_catalog(prepared);
720 let changed = crate::model_registry::model_catalog_revision() != revision;
721 (delta, changed)
722 },
723 );
724 let (commit, hydrated) = match commit {
725 Ok(result) => result,
726 Err(error) => return (Err(error.into()), Vec::new()),
727 };
728 match commit {
729 AuthProviderRuntimeCommit::Applied { .. } => {
730 drop(state);
731 (
732 Ok(hydrated.expect("catalog hydration callback was not run")),
733 Vec::new(),
734 )
735 }
736 AuthProviderRuntimeCommit::Missing => {
737 Self::bump_generation(&mut state, provider_id);
738 let (_, removed) = Self::remove_runtime_provider(&mut state, provider_id, false);
739 drop(state);
740 (
741 Err(ProviderLifecycleError::ProviderNotFound {
742 id: provider_id.to_string(),
743 }),
744 removed,
745 )
746 }
747 AuthProviderRuntimeCommit::Disabled => {
748 Self::bump_generation(&mut state, provider_id);
749 let (_, removed) = Self::remove_runtime_provider(&mut state, provider_id, false);
750 drop(state);
751 (
752 Err(ProviderLifecycleError::ProviderDisabled {
753 id: provider_id.to_string(),
754 }),
755 removed,
756 )
757 }
758 AuthProviderRuntimeCommit::Changed => {
759 drop(state);
760 (
761 Err(ProviderLifecycleError::Stale {
762 id: provider_id.to_string(),
763 }),
764 Vec::new(),
765 )
766 }
767 }
768 }
769
770 fn finish_catalog_refresh(
771 &self,
772 provider_id: &str,
773 attempt: &ProviderRefreshAttempt,
774 models: &[crate::provider::DiscoveredModelDetails],
775 ) -> (
776 Result<CatalogDelta, ProviderLifecycleError>,
777 Vec<Arc<dyn Provider>>,
778 ) {
779 let ProviderRefreshAttempt {
780 generation,
781 runtime,
782 live_provider,
783 } = attempt;
784 let deferred = Vec::new();
785 let state = self.lock_state();
786 if let Err(error) = self.ensure_generation(&state, provider_id, *generation) {
787 drop(state);
788 return self.finish_refresh_failure(
789 provider_id,
790 &runtime.provider.kind,
791 *generation,
792 live_provider,
793 error,
794 deferred,
795 );
796 }
797 let Some(current_live) = state.live_providers.get(provider_id) else {
798 drop(state);
799 return self.finish_refresh_failure(
800 provider_id,
801 &runtime.provider.kind,
802 *generation,
803 live_provider,
804 ProviderLifecycleError::Stale {
805 id: provider_id.to_string(),
806 },
807 deferred,
808 );
809 };
810 if current_live.kind != runtime.provider.kind
811 || !Arc::ptr_eq(live_provider, ¤t_live.provider)
812 {
813 drop(state);
814 return self.finish_refresh_failure(
815 provider_id,
816 &runtime.provider.kind,
817 *generation,
818 live_provider,
819 ProviderLifecycleError::Stale {
820 id: provider_id.to_string(),
821 },
822 deferred,
823 );
824 }
825
826 let descriptor = match self.provider_descriptor_from_state(
827 &runtime.provider,
828 &runtime.provider_ids,
829 runtime.model_namespace.as_deref(),
830 ) {
831 Ok(descriptor) => descriptor,
832 Err(error) => {
833 drop(state);
834 return self.finish_refresh_failure(
835 provider_id,
836 &runtime.provider.kind,
837 *generation,
838 live_provider,
839 error,
840 deferred,
841 );
842 }
843 };
844 let prepared = match prepare_provider_catalog(descriptor, models) {
845 Ok(prepared) => prepared,
846 Err(error) => {
847 drop(state);
848 return self.finish_refresh_failure(
849 provider_id,
850 &runtime.provider.kind,
851 *generation,
852 live_provider,
853 error.into(),
854 deferred,
855 );
856 }
857 };
858 let expected_provider_ids = runtime
859 .model_namespace
860 .is_none()
861 .then_some(runtime.provider_ids.as_slice());
862 let result = self.commit_cache_and_catalog_if_current(
863 &runtime.provider,
864 &runtime.catalog_snapshot,
865 expected_provider_ids,
866 chrono::Utc::now().timestamp(),
867 models,
868 prepared,
869 );
870 drop(state);
871 match result {
872 Ok(delta) => (Ok(delta), deferred),
873 Err(error @ ProviderLifecycleError::Stale { .. }) => {
874 self.reconcile_refresh_failure(provider_id, attempt, error, deferred)
875 }
876 Err(error) => self.finish_refresh_failure(
877 provider_id,
878 &runtime.provider.kind,
879 *generation,
880 live_provider,
881 error,
882 deferred,
883 ),
884 }
885 }
886
887 fn reconcile_refresh_failure(
888 &self,
889 provider_id: &str,
890 attempt: &ProviderRefreshAttempt,
891 original: ProviderLifecycleError,
892 mut deferred: Vec<Arc<dyn Provider>>,
893 ) -> (
894 Result<CatalogDelta, ProviderLifecycleError>,
895 Vec<Arc<dyn Provider>>,
896 ) {
897 let (preparation, reconciled_deferred) = self.prepare_catalog_refresh(
898 provider_id,
899 chrono::Utc::now().timestamp(),
900 true,
901 Some(&attempt.live_provider),
902 Some(&attempt.runtime.catalog_snapshot),
903 );
904 deferred.extend(reconciled_deferred);
905 match preparation {
906 Ok(ProviderRefreshPreparation::Hydrated { delta, .. }) => (Ok(delta), deferred),
907 Ok(ProviderRefreshPreparation::Discover(_)) => self.finish_refresh_failure(
908 provider_id,
909 &attempt.runtime.provider.kind,
910 attempt.generation,
911 &attempt.live_provider,
912 original,
913 deferred,
914 ),
915 Err(error) => self.finish_refresh_failure(
916 provider_id,
917 &attempt.runtime.provider.kind,
918 attempt.generation,
919 &attempt.live_provider,
920 error,
921 deferred,
922 ),
923 }
924 }
925
926 fn finish_refresh_failure(
927 &self,
928 provider_id: &str,
929 expected_kind: &ProviderKind,
930 expected_generation: u64,
931 expected_provider: &Arc<dyn Provider>,
932 original: ProviderLifecycleError,
933 mut deferred: Vec<Arc<dyn Provider>>,
934 ) -> (
935 Result<CatalogDelta, ProviderLifecycleError>,
936 Vec<Arc<dyn Provider>>,
937 ) {
938 let catalog_identity_invalid = matches!(
939 &original,
940 ProviderLifecycleError::Catalog(
941 CatalogError::NamespaceChanged { .. }
942 | CatalogError::NamespaceInUse { .. }
943 | CatalogError::NamespaceStore { .. }
944 )
945 );
946 let mut state = self.lock_state();
947 let authoritative = self
948 .hub
949 .load_or_create_auth_provider_runtime_state(provider_id);
950 let error = match authoritative {
951 Err(error) => error.into(),
952 Ok(None) => ProviderLifecycleError::Stale {
953 id: provider_id.to_string(),
954 },
955 Ok(Some(runtime))
956 if !runtime.provider.enabled || runtime.provider.kind != *expected_kind =>
957 {
958 ProviderLifecycleError::Stale {
959 id: provider_id.to_string(),
960 }
961 }
962 Ok(Some(_)) => {
963 let same_generation =
964 Self::generation_in(&state, provider_id) == expected_generation;
965 let current = state.live_providers.get(provider_id);
966 let same_provider = current.is_some_and(|current| {
967 current.kind == *expected_kind
968 && Arc::ptr_eq(¤t.provider, expected_provider)
969 });
970 if same_generation && same_provider && !catalog_identity_invalid {
971 drop(state);
972 return (Err(original), deferred);
973 }
974 if catalog_identity_invalid {
975 original
976 } else if current.is_some_and(|current| current.kind == *expected_kind) {
977 drop(state);
978 return (
979 Err(ProviderLifecycleError::Stale {
980 id: provider_id.to_string(),
981 }),
982 deferred,
983 );
984 } else {
985 ProviderLifecycleError::Stale {
986 id: provider_id.to_string(),
987 }
988 }
989 }
990 };
991 Self::bump_generation(&mut state, provider_id);
992 let (_, removed) = Self::remove_runtime_provider(&mut state, provider_id, false);
993 deferred.extend(removed);
994 drop(state);
995 (Err(error), deferred)
996 }
997
998 pub async fn restore_provider(
999 &self,
1000 provider_id: &str,
1001 expected_kind: ProviderKind,
1002 live_provider: Arc<dyn Provider>,
1003 ) -> Result<ProviderLifecycleOutcome, ProviderLifecycleError> {
1004 self.validate_live_name(provider_id, &live_provider)?;
1005 let operation_lock = self.operation_lock(provider_id);
1006 let mut fence = ProviderOperationFence::new(operation_lock.lock_owned().await);
1007 fence.protect(live_provider.clone());
1008 let (result, deferred) =
1009 self.commit_existing_provider_runtime(provider_id, expected_kind, live_provider, false);
1010 fence.extend(deferred);
1011 drop(fence);
1012 result
1013 }
1014
1015 pub async fn enable_provider(
1016 &self,
1017 provider_id: &str,
1018 expected_kind: ProviderKind,
1019 live_provider: Arc<dyn Provider>,
1020 ) -> Result<ProviderLifecycleOutcome, ProviderLifecycleError> {
1021 self.validate_live_name(provider_id, &live_provider)?;
1022 let operation_lock = self.operation_lock(provider_id);
1023 let mut fence = ProviderOperationFence::new(operation_lock.lock_owned().await);
1024 fence.protect(live_provider.clone());
1025 let (result, deferred) =
1026 self.commit_existing_provider_runtime(provider_id, expected_kind, live_provider, true);
1027 fence.extend(deferred);
1028 drop(fence);
1029 result
1030 }
1031
1032 pub fn reconcile_inactive_providers(
1033 &self,
1034 ) -> Result<ProviderReconcileOutcome, ProviderLifecycleError> {
1035 let mut state = self.lock_state();
1036 let removals = self
1037 .hub
1038 .with_auth_provider_runtime_descriptors(|providers| {
1039 let active = providers
1040 .iter()
1041 .filter(|provider| provider.enabled)
1042 .map(|provider| (provider.id.as_str(), &provider.kind))
1043 .collect::<HashMap<_, _>>();
1044 let mut inactive = state
1045 .live_providers
1046 .iter()
1047 .filter(|(provider_id, live)| {
1048 active
1049 .get(provider_id.as_str())
1050 .is_none_or(|kind| **kind != live.kind)
1051 })
1052 .map(|(provider_id, _)| provider_id.clone())
1053 .collect::<BTreeSet<_>>();
1054 inactive.extend(
1055 providers
1056 .iter()
1057 .filter(|provider| !provider.enabled)
1058 .map(|provider| provider.id.clone()),
1059 );
1060
1061 inactive
1062 .into_iter()
1063 .map(|provider_id| {
1064 Self::bump_generation(&mut state, &provider_id);
1065 let removal = Self::remove_live_provider(&mut state, &provider_id);
1066 (provider_id, removal)
1067 })
1068 .collect::<Vec<_>>()
1069 });
1070 let removals = match removals {
1071 Ok(removals) => removals,
1072 Err(error) => {
1073 let provider_ids = state
1074 .live_providers
1075 .keys()
1076 .cloned()
1077 .collect::<BTreeSet<_>>();
1078 let mut removed_providers = Vec::new();
1079 for provider_id in provider_ids {
1080 Self::bump_generation(&mut state, &provider_id);
1081 let removal = Self::remove_live_provider(&mut state, &provider_id);
1082 removed_providers.extend(removal.removed);
1083 remove_provider_catalog(&provider_id);
1084 }
1085 drop(state);
1086 drop(removed_providers);
1087 return Err(error.into());
1088 }
1089 };
1090 let mut removed_providers = Vec::new();
1091 let providers = removals
1092 .into_iter()
1093 .map(|(provider_id, removal)| {
1094 removed_providers.extend(removal.removed);
1095 let change = ProviderStateChange {
1096 auth_changed: false,
1097 live_changed: removal.changed,
1098 catalog_changed: remove_provider_catalog(&provider_id),
1099 };
1100 (provider_id, change)
1101 })
1102 .collect();
1103 drop(state);
1104 drop(removed_providers);
1105 Ok(ProviderReconcileOutcome { providers })
1106 }
1107
1108 pub fn disable_provider(
1109 &self,
1110 provider_id: &str,
1111 ) -> Result<ProviderStateChange, ProviderLifecycleError> {
1112 let mut state = self.lock_state();
1113 let auth_changed = match self
1114 .hub
1115 .set_auth_provider_enabled_with_change(provider_id, false)
1116 {
1117 Ok(Some(auth_changed)) => auth_changed,
1118 Ok(None) => {
1119 Self::bump_generation(&mut state, provider_id);
1120 let (_, removed) = Self::remove_runtime_provider(&mut state, provider_id, false);
1121 drop(state);
1122 drop(removed);
1123 return Err(ProviderLifecycleError::ProviderNotFound {
1124 id: provider_id.to_string(),
1125 });
1126 }
1127 Err(error) => {
1128 Self::bump_generation(&mut state, provider_id);
1129 let (_, removed) = Self::remove_runtime_provider(&mut state, provider_id, false);
1130 drop(state);
1131 drop(removed);
1132 return Err(error.into());
1133 }
1134 };
1135 Self::bump_generation(&mut state, provider_id);
1136 let (change, removed) =
1137 Self::remove_runtime_provider(&mut state, provider_id, auth_changed);
1138 drop(state);
1139 drop(removed);
1140 Ok(change)
1141 }
1142
1143 pub fn remove_provider(
1144 &self,
1145 provider_id: &str,
1146 ) -> Result<ProviderStateChange, ProviderLifecycleError> {
1147 let mut state = self.lock_state();
1148 let auth_changed = match self.hub.remove_auth_provider(provider_id) {
1149 Ok(auth_changed) => auth_changed,
1150 Err(error) => {
1151 Self::bump_generation(&mut state, provider_id);
1152 let (_, removed) = Self::remove_runtime_provider(&mut state, provider_id, false);
1153 drop(state);
1154 drop(removed);
1155 return Err(error.into());
1156 }
1157 };
1158 Self::bump_generation(&mut state, provider_id);
1159 let (change, removed) =
1160 Self::remove_runtime_provider(&mut state, provider_id, auth_changed);
1161 drop(state);
1162 drop(removed);
1163 Ok(change)
1164 }
1165
1166 fn commit_existing_provider_runtime(
1167 &self,
1168 provider_id: &str,
1169 expected_kind: ProviderKind,
1170 live_provider: Arc<dyn Provider>,
1171 enable_if_disabled: bool,
1172 ) -> (
1173 Result<ProviderLifecycleOutcome, ProviderLifecycleError>,
1174 Vec<Arc<dyn Provider>>,
1175 ) {
1176 let mut state = self.lock_state();
1177 let runtime = self
1178 .hub
1179 .load_or_create_auth_provider_runtime_state(provider_id);
1180 let runtime = match runtime {
1181 Ok(Some(runtime)) => runtime,
1182 Ok(None) => {
1183 Self::bump_generation(&mut state, provider_id);
1184 let (_, removed) = Self::remove_runtime_provider(&mut state, provider_id, false);
1185 drop(state);
1186 return (
1187 Err(ProviderLifecycleError::ProviderNotFound {
1188 id: provider_id.to_string(),
1189 }),
1190 removed,
1191 );
1192 }
1193 Err(error) => {
1194 Self::bump_generation(&mut state, provider_id);
1195 let (_, removed) = Self::remove_runtime_provider(&mut state, provider_id, false);
1196 drop(state);
1197 return (Err(error.into()), removed);
1198 }
1199 };
1200 let provider = runtime.provider;
1201 if provider.kind != expected_kind {
1202 Self::bump_generation(&mut state, provider_id);
1203 let (_, removed) = Self::remove_runtime_provider(&mut state, provider_id, false);
1204 drop(state);
1205 return (
1206 Err(ProviderLifecycleError::Stale {
1207 id: provider_id.to_string(),
1208 }),
1209 removed,
1210 );
1211 }
1212 if !provider.enabled && !enable_if_disabled {
1213 let validation = self.hub.commit_auth_provider_runtime_if_current_and_then(
1214 &provider,
1215 &runtime.catalog_snapshot,
1216 false,
1217 None,
1218 None,
1219 || (),
1220 );
1221 let error = match validation {
1222 Ok((AuthProviderRuntimeCommit::Missing, _)) => {
1223 ProviderLifecycleError::ProviderNotFound {
1224 id: provider_id.to_string(),
1225 }
1226 }
1227 Ok((AuthProviderRuntimeCommit::Disabled, _)) => {
1228 ProviderLifecycleError::ProviderDisabled {
1229 id: provider_id.to_string(),
1230 }
1231 }
1232 Ok((
1233 AuthProviderRuntimeCommit::Changed | AuthProviderRuntimeCommit::Applied { .. },
1234 _,
1235 )) => ProviderLifecycleError::Stale {
1236 id: provider_id.to_string(),
1237 },
1238 Err(error) => error.into(),
1239 };
1240 Self::bump_generation(&mut state, provider_id);
1241 let (_, removed) = Self::remove_runtime_provider(&mut state, provider_id, false);
1242 drop(state);
1243 return (Err(error), removed);
1244 }
1245 let prepared = match runtime.model_cache.as_deref() {
1246 Some(models) => self
1247 .provider_descriptor_from_state(
1248 &provider,
1249 &runtime.provider_ids,
1250 runtime.model_namespace.as_deref(),
1251 )
1252 .and_then(|descriptor| {
1253 prepare_provider_catalog(descriptor, models).map_err(Into::into)
1254 })
1255 .map(Some),
1256 None => Ok(None),
1257 };
1258 let prepared = match prepared {
1259 Ok(prepared) => prepared,
1260 Err(error) => {
1261 Self::bump_generation(&mut state, provider_id);
1262 let (_, removed) = Self::remove_runtime_provider(&mut state, provider_id, false);
1263 drop(state);
1264 return (Err(error), removed);
1265 }
1266 };
1267 let model_namespace = prepared
1268 .as_ref()
1269 .map(|prepared| prepared.namespace().to_string());
1270 let selected_provider = state
1271 .live_providers
1272 .get(provider_id)
1273 .filter(|current| current.kind == provider.kind)
1274 .map(|current| current.provider.clone())
1275 .unwrap_or_else(|| live_provider.clone());
1276 let catalog_revision = prepared
1277 .is_some()
1278 .then(crate::model_registry::model_catalog_revision);
1279 let expected_provider_ids = runtime
1280 .model_namespace
1281 .is_none()
1282 .then_some(runtime.provider_ids.as_slice());
1283 let commit = self.hub.commit_auth_provider_runtime_if_current_and_then(
1284 &provider,
1285 &runtime.catalog_snapshot,
1286 enable_if_disabled,
1287 model_namespace.as_deref(),
1288 expected_provider_ids,
1289 || {
1290 let registration = Self::register_live_provider(
1291 &mut state,
1292 provider_id,
1293 provider.kind.clone(),
1294 selected_provider,
1295 );
1296 let (catalog, catalog_changed) = match prepared {
1297 Some(prepared) => {
1298 let catalog = commit_prepared_provider_catalog(prepared);
1299 let changed = catalog_revision.is_some_and(|revision| {
1300 crate::model_registry::model_catalog_revision() != revision
1301 });
1302 (Some(catalog), changed)
1303 }
1304 None => (None, remove_provider_catalog(provider_id)),
1305 };
1306 (catalog, catalog_changed, registration)
1307 },
1308 );
1309 let (commit, applied) = match commit {
1310 Ok(commit) => commit,
1311 Err(error) => {
1312 Self::bump_generation(&mut state, provider_id);
1313 let (_, removed) = Self::remove_runtime_provider(&mut state, provider_id, false);
1314 drop(state);
1315 return (Err(error.into()), removed);
1316 }
1317 };
1318
1319 match commit {
1320 AuthProviderRuntimeCommit::Applied { auth_changed } => {
1321 let (catalog, catalog_changed, registration) =
1322 applied.expect("applied provider runtime callback was not run");
1323 if auth_changed || catalog_changed || registration.changed {
1324 Self::bump_generation(&mut state, provider_id);
1325 }
1326 drop(state);
1327 (
1328 Ok(ProviderLifecycleOutcome {
1329 state: ProviderStateChange {
1330 auth_changed,
1331 live_changed: registration.changed,
1332 catalog_changed,
1333 },
1334 catalog,
1335 }),
1336 registration.replaced,
1337 )
1338 }
1339 AuthProviderRuntimeCommit::Missing => {
1340 Self::bump_generation(&mut state, provider_id);
1341 let (_, removed) = Self::remove_runtime_provider(&mut state, provider_id, false);
1342 drop(state);
1343 (
1344 Err(ProviderLifecycleError::ProviderNotFound {
1345 id: provider_id.to_string(),
1346 }),
1347 removed,
1348 )
1349 }
1350 AuthProviderRuntimeCommit::Disabled => {
1351 Self::bump_generation(&mut state, provider_id);
1352 let (_, removed) = Self::remove_runtime_provider(&mut state, provider_id, false);
1353 drop(state);
1354 (
1355 Err(ProviderLifecycleError::ProviderDisabled {
1356 id: provider_id.to_string(),
1357 }),
1358 removed,
1359 )
1360 }
1361 AuthProviderRuntimeCommit::Changed => {
1362 Self::bump_generation(&mut state, provider_id);
1363 let (_, removed) = Self::remove_runtime_provider(&mut state, provider_id, false);
1364 drop(state);
1365 (
1366 Err(ProviderLifecycleError::Stale {
1367 id: provider_id.to_string(),
1368 }),
1369 removed,
1370 )
1371 }
1372 }
1373 }
1374
1375 fn provider_descriptor(
1376 &self,
1377 provider: &StoredProvider,
1378 providers: &[StoredProvider],
1379 ) -> Result<ProviderDescriptor, ProviderLifecycleError> {
1380 let provider_ids = providers
1381 .iter()
1382 .map(|provider| provider.id.clone())
1383 .collect::<Vec<_>>();
1384 let persisted_namespace = self.hub.load_auth_model_namespace(&provider.id)?;
1385 self.provider_descriptor_from_state(provider, &provider_ids, persisted_namespace.as_deref())
1386 }
1387
1388 fn provider_descriptor_from_state(
1389 &self,
1390 provider: &StoredProvider,
1391 provider_ids: &[String],
1392 persisted_namespace: Option<&str>,
1393 ) -> Result<ProviderDescriptor, ProviderLifecycleError> {
1394 let in_memory_namespace = provider_catalog_namespace(&provider.id);
1395 if let (Some(current), Some(persisted)) = (&in_memory_namespace, persisted_namespace)
1396 && current != persisted
1397 {
1398 return Err(CatalogError::NamespaceChanged {
1399 provider_key: provider.id.clone(),
1400 current: current.clone(),
1401 requested: persisted.to_string(),
1402 }
1403 .into());
1404 }
1405 let namespace = in_memory_namespace
1406 .or_else(|| persisted_namespace.map(str::to_string))
1407 .unwrap_or_else(|| {
1408 let short_id = shortest_unique_provider_id(&provider.id, provider_ids);
1409 format!("{short_id}@{}", provider.name)
1410 });
1411 Ok(ProviderDescriptor {
1412 provider_key: provider.id.clone(),
1413 provider_name: provider.name.clone(),
1414 namespace,
1415 wire_profile: wire_profile_for_kind(&provider.kind),
1416 })
1417 }
1418
1419 fn commit_cache_and_catalog_if_current(
1420 &self,
1421 provider: &StoredProvider,
1422 expected_catalog: &crate::auth_store::AuthProviderCatalogSnapshot,
1423 expected_provider_ids: Option<&[String]>,
1424 fetched_at: i64,
1425 models: &[crate::provider::DiscoveredModelDetails],
1426 prepared: PreparedProviderCatalog,
1427 ) -> Result<CatalogDelta, ProviderLifecycleError> {
1428 let namespace = prepared.namespace().to_string();
1429 let (commit, delta) = self
1430 .hub
1431 .update_auth_model_cache_details_if_enabled_and_then(
1432 AuthModelCacheUpdate {
1433 expected: provider,
1434 expected_catalog,
1435 expected_provider_ids,
1436 model_namespace: &namespace,
1437 fetched_at,
1438 models,
1439 },
1440 || commit_prepared_provider_catalog(prepared),
1441 )?;
1442 match (commit, delta) {
1443 (AuthModelCacheCommit::Updated, Some(delta)) => Ok(delta),
1444 (
1445 AuthModelCacheCommit::Missing
1446 | AuthModelCacheCommit::Disabled
1447 | AuthModelCacheCommit::Changed,
1448 _,
1449 ) => Err(ProviderLifecycleError::Stale {
1450 id: provider.id.clone(),
1451 }),
1452 (AuthModelCacheCommit::Updated, None) => unreachable!("updated callback was not run"),
1453 }
1454 }
1455
1456 fn validate_live_name(
1457 &self,
1458 provider_id: &str,
1459 live_provider: &Arc<dyn Provider>,
1460 ) -> Result<(), ProviderLifecycleError> {
1461 if live_provider.name() == provider_id {
1462 return Ok(());
1463 }
1464 Err(ProviderLifecycleError::LiveProviderNameMismatch {
1465 expected: provider_id.to_string(),
1466 actual: live_provider.name().to_string(),
1467 })
1468 }
1469
1470 fn validate_new_provider(
1471 &self,
1472 provider: &StoredProvider,
1473 live_provider: &Arc<dyn Provider>,
1474 ) -> Result<(), ProviderLifecycleError> {
1475 self.validate_live_name(&provider.id, live_provider)?;
1476 if provider.enabled {
1477 return Ok(());
1478 }
1479 Err(ProviderLifecycleError::ProviderDisabled {
1480 id: provider.id.clone(),
1481 })
1482 }
1483
1484 fn operation_lock(&self, provider_id: &str) -> Arc<tokio::sync::Mutex<()>> {
1485 self.lock_state()
1486 .operation_locks
1487 .entry(provider_id.to_string())
1488 .or_default()
1489 .clone()
1490 }
1491
1492 fn refresh_lock(&self, provider_id: &str) -> Arc<tokio::sync::Mutex<()>> {
1493 self.lock_state()
1494 .refresh_locks
1495 .entry(provider_id.to_string())
1496 .or_default()
1497 .clone()
1498 }
1499
1500 fn generation(&self, provider_id: &str) -> u64 {
1501 Self::generation_in(&self.lock_state(), provider_id)
1502 }
1503
1504 fn generation_in(state: &LifecycleState, provider_id: &str) -> u64 {
1505 state.generations.get(provider_id).copied().unwrap_or(0)
1506 }
1507
1508 fn ensure_generation(
1509 &self,
1510 state: &LifecycleState,
1511 provider_id: &str,
1512 expected: u64,
1513 ) -> Result<(), ProviderLifecycleError> {
1514 if Self::generation_in(state, provider_id) == expected {
1515 return Ok(());
1516 }
1517 Err(ProviderLifecycleError::Stale {
1518 id: provider_id.to_string(),
1519 })
1520 }
1521
1522 fn bump_generation(state: &mut LifecycleState, provider_id: &str) {
1523 let generation = state
1524 .generations
1525 .entry(provider_id.to_string())
1526 .or_default();
1527 *generation = generation.wrapping_add(1);
1528 }
1529
1530 fn register_live_provider(
1531 state: &mut LifecycleState,
1532 provider_id: &str,
1533 kind: ProviderKind,
1534 provider: Arc<dyn Provider>,
1535 ) -> LiveProviderRegistration {
1536 let mut changed = state.live_providers.get(provider_id).is_none_or(|current| {
1537 current.kind != kind || !Arc::ptr_eq(¤t.provider, &provider)
1538 });
1539 let mut replaced = state
1540 .live_providers
1541 .insert(
1542 provider_id.to_string(),
1543 LiveProviderEntry {
1544 kind,
1545 provider: provider.clone(),
1546 },
1547 )
1548 .into_iter()
1549 .map(|live| live.provider)
1550 .collect::<Vec<_>>();
1551 state.registries.retain(|registry| {
1552 let Some(registry) = registry.upgrade() else {
1553 return false;
1554 };
1555 let previous = registry.register_named(provider_id.to_string(), provider.clone());
1556 changed |= previous
1557 .as_ref()
1558 .is_none_or(|previous| !Arc::ptr_eq(previous, &provider));
1559 replaced.extend(previous);
1560 true
1561 });
1562 LiveProviderRegistration { changed, replaced }
1563 }
1564
1565 fn remove_live_provider(state: &mut LifecycleState, provider_id: &str) -> LiveProviderRemoval {
1566 let mut removed = state
1567 .live_providers
1568 .remove(provider_id)
1569 .map(|live| live.provider)
1570 .into_iter()
1571 .collect::<Vec<_>>();
1572 state.registries.retain(|registry| {
1573 let Some(registry) = registry.upgrade() else {
1574 return false;
1575 };
1576 removed.extend(registry.take_named(provider_id));
1577 true
1578 });
1579 LiveProviderRemoval {
1580 changed: !removed.is_empty(),
1581 removed,
1582 }
1583 }
1584
1585 fn remove_runtime_provider(
1586 state: &mut LifecycleState,
1587 provider_id: &str,
1588 auth_changed: bool,
1589 ) -> (ProviderStateChange, Vec<Arc<dyn Provider>>) {
1590 let removal = Self::remove_live_provider(state, provider_id);
1591 (
1592 ProviderStateChange {
1593 auth_changed,
1594 live_changed: removal.changed,
1595 catalog_changed: remove_provider_catalog(provider_id),
1596 },
1597 removal.removed,
1598 )
1599 }
1600
1601 fn lock_state(&self) -> std::sync::MutexGuard<'_, LifecycleState> {
1602 self.state
1603 .lock()
1604 .unwrap_or_else(std::sync::PoisonError::into_inner)
1605 }
1606}
1607
1608pub(crate) fn reload_config_providers_for_hub(hub: &ConfigHub) -> Result<(), ConfigError> {
1609 let state = shared_config_provider_state(hub);
1610 let mut state = state
1611 .lock()
1612 .unwrap_or_else(std::sync::PoisonError::into_inner);
1613 let mut replaced = Vec::new();
1614 hub.reload_and_then(|snapshot| {
1615 let mut providers = snapshot.providers.into_iter().collect::<Vec<_>>();
1616 providers.sort_by(|left, right| left.0.cmp(&right.0));
1617 let current_names = providers
1618 .iter()
1619 .map(|(name, _)| name.clone())
1620 .collect::<BTreeSet<_>>();
1621 let stale_names = state
1622 .config_providers
1623 .keys()
1624 .filter(|name| !current_names.contains(*name))
1625 .cloned()
1626 .collect::<Vec<_>>();
1627 state.registries.retain(|registry| {
1628 let Some(registry) = registry.upgrade() else {
1629 return false;
1630 };
1631 for name in &stale_names {
1632 replaced.extend(registry.take_named(&format!("config:{name}")));
1633 }
1634 for (name, entry) in &providers {
1635 let (_, previous) = crate::config_provider::reconcile_config_provider_deferred(
1636 ®istry, name, entry,
1637 );
1638 replaced.extend(previous);
1639 }
1640 true
1641 });
1642 state.config_providers = providers.into_iter().collect();
1643 })?;
1644 drop(state);
1645 drop(replaced);
1646 Ok(())
1647}
1648
1649pub(crate) fn upsert_config_provider_for_hub(
1650 hub: &ConfigHub,
1651 update: ProviderConfigUpdate<'_>,
1652) -> Result<(), ConfigError> {
1653 mutate_config_provider_for_hub(hub, update, ProviderConfigWriteMode::Upsert)
1654}
1655
1656pub(crate) fn remove_config_provider_for_hub(
1657 hub: &ConfigHub,
1658 name: &str,
1659) -> Result<(), ConfigError> {
1660 let state = shared_config_provider_state(hub);
1661 let mut state = state
1662 .lock()
1663 .unwrap_or_else(std::sync::PoisonError::into_inner);
1664 let mut replaced = Vec::new();
1665 hub.remove_provider_config_and_then(name, || {
1666 state.registries.retain(|registry| {
1667 let Some(registry) = registry.upgrade() else {
1668 return false;
1669 };
1670 replaced.extend(registry.take_named(&format!("config:{name}")));
1671 true
1672 });
1673 state.config_providers.remove(name);
1674 })?;
1675 drop(state);
1676 drop(replaced);
1677 Ok(())
1678}
1679
1680fn mutate_config_provider_for_hub(
1681 hub: &ConfigHub,
1682 update: ProviderConfigUpdate<'_>,
1683 mode: ProviderConfigWriteMode,
1684) -> Result<(), ConfigError> {
1685 let state = shared_config_provider_state(hub);
1686 let mut state = state
1687 .lock()
1688 .unwrap_or_else(std::sync::PoisonError::into_inner);
1689 let mut replaced = Vec::new();
1690 hub.write_provider_config_and_then(update, mode, |entry| {
1691 state.registries.retain(|registry| {
1692 let Some(registry) = registry.upgrade() else {
1693 return false;
1694 };
1695 let (_, previous) = crate::config_provider::reconcile_config_provider_deferred(
1696 ®istry,
1697 &entry.name,
1698 entry,
1699 );
1700 replaced.extend(previous);
1701 true
1702 });
1703 state
1704 .config_providers
1705 .insert(entry.name.clone(), entry.clone());
1706 })?;
1707 drop(state);
1708 drop(replaced);
1709 Ok(())
1710}
1711
1712fn wire_profile_for_kind(kind: &ProviderKind) -> ReasoningWireProfile {
1713 match kind {
1714 ProviderKind::Codex => ReasoningWireProfile::CodexResponses,
1715 ProviderKind::AnthropicOauth => ReasoningWireProfile::AnthropicMessages,
1716 ProviderKind::GitHubCopilot | ProviderKind::Custom => ReasoningWireProfile::Unknown,
1717 }
1718}
1719
1720fn shared_lifecycle_state(hub: &ConfigHub) -> Arc<Mutex<LifecycleState>> {
1721 let key = coordinator_key(hub.auth_path());
1722 let mut coordinators = LIFECYCLE_COORDINATORS
1723 .lock()
1724 .unwrap_or_else(std::sync::PoisonError::into_inner);
1725 coordinators.retain(|_, coordinator| coordinator.strong_count() > 0);
1726 if let Some(state) = coordinators.get(&key).and_then(Weak::upgrade) {
1727 return state;
1728 }
1729 let state = Arc::new(Mutex::new(LifecycleState::default()));
1730 coordinators.insert(key, Arc::downgrade(&state));
1731 state
1732}
1733
1734fn shared_config_provider_state(hub: &ConfigHub) -> Arc<Mutex<ConfigProviderLifecycleState>> {
1735 let key = coordinator_key(&hub.config_toml_path());
1736 let mut coordinators = CONFIG_PROVIDER_COORDINATORS
1737 .lock()
1738 .unwrap_or_else(std::sync::PoisonError::into_inner);
1739 coordinators.retain(|_, coordinator| coordinator.strong_count() > 0);
1740 if let Some(state) = coordinators.get(&key).and_then(Weak::upgrade) {
1741 return state;
1742 }
1743 let state = Arc::new(Mutex::new(ConfigProviderLifecycleState::default()));
1744 coordinators.insert(key, Arc::downgrade(&state));
1745 state
1746}
1747
1748fn coordinator_key(path: &Path) -> PathBuf {
1749 if let Some(parent) = path.parent()
1750 && let Ok(parent) = std::fs::canonicalize(parent)
1751 && let Some(file_name) = path.file_name()
1752 {
1753 return parent.join(file_name);
1754 }
1755 if path.is_absolute() {
1756 return path.to_path_buf();
1757 }
1758 std::env::current_dir()
1759 .map(|current| current.join(path))
1760 .unwrap_or_else(|_| path.to_path_buf())
1761}
1762
1763#[cfg(test)]
1764mod tests {
1765 use std::sync::Mutex as StdMutex;
1766 use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
1767
1768 use tokio::sync::oneshot;
1769
1770 use super::*;
1771 use crate::error::RuntimeError;
1772 use crate::event::Observable;
1773 use crate::provider::{
1774 AssistantMessage, CapabilityKnowledge, DiscoveredModelDetails, LlmRequest,
1775 ModelCapabilities, ReasoningEffort,
1776 };
1777 use crate::tool::BoxFut;
1778
1779 struct TestProvider {
1780 name: String,
1781 models: Arc<StdMutex<Vec<DiscoveredModelDetails>>>,
1782 block: Arc<StdMutex<Option<DiscoveryBlock>>>,
1783 fail: Arc<StdMutex<Option<ModelDiscoveryError>>>,
1784 before_discovery: Arc<StdMutex<Option<DiscoveryHook>>>,
1785 discovery_calls: Arc<AtomicUsize>,
1786 on_drop: StdMutex<Option<DiscoveryHook>>,
1787 }
1788
1789 type DiscoveryHook = Box<dyn FnOnce() + Send>;
1790
1791 struct DiscoveryBlock {
1792 started: oneshot::Sender<()>,
1793 proceed: oneshot::Receiver<()>,
1794 }
1795
1796 impl TestProvider {
1797 fn new(name: &str, models: Vec<DiscoveredModelDetails>) -> Self {
1798 Self {
1799 name: name.into(),
1800 models: Arc::new(StdMutex::new(models)),
1801 block: Arc::new(StdMutex::new(None)),
1802 fail: Arc::new(StdMutex::new(None)),
1803 before_discovery: Arc::new(StdMutex::new(None)),
1804 discovery_calls: Arc::new(AtomicUsize::new(0)),
1805 on_drop: StdMutex::new(None),
1806 }
1807 }
1808
1809 fn set_models(&self, models: Vec<DiscoveredModelDetails>) {
1810 *self.models.lock().unwrap() = models;
1811 }
1812
1813 fn fail_once(&self, error: ModelDiscoveryError) {
1814 *self.fail.lock().unwrap() = Some(error);
1815 }
1816
1817 fn block_once(&self) -> (oneshot::Receiver<()>, oneshot::Sender<()>) {
1818 let (started_tx, started_rx) = oneshot::channel();
1819 let (proceed_tx, proceed_rx) = oneshot::channel();
1820 *self.block.lock().unwrap() = Some(DiscoveryBlock {
1821 started: started_tx,
1822 proceed: proceed_rx,
1823 });
1824 (started_rx, proceed_tx)
1825 }
1826
1827 fn before_discovery_once(&self, hook: impl FnOnce() + Send + 'static) {
1828 *self.before_discovery.lock().unwrap() = Some(Box::new(hook));
1829 }
1830
1831 fn discovery_calls(&self) -> usize {
1832 self.discovery_calls.load(Ordering::SeqCst)
1833 }
1834
1835 fn on_drop(&self, hook: impl FnOnce() + Send + 'static) {
1836 *self.on_drop.lock().unwrap() = Some(Box::new(hook));
1837 }
1838 }
1839
1840 impl Drop for TestProvider {
1841 fn drop(&mut self) {
1842 if let Some(on_drop) = self.on_drop.lock().unwrap().take() {
1843 on_drop();
1844 }
1845 }
1846 }
1847
1848 impl Provider for TestProvider {
1849 fn name(&self) -> &str {
1850 &self.name
1851 }
1852
1853 fn call<'a>(
1854 &'a self,
1855 _req: LlmRequest,
1856 ) -> BoxFut<'a, Result<AssistantMessage, RuntimeError>> {
1857 Box::pin(async { unreachable!("not used by lifecycle tests") })
1858 }
1859
1860 fn call_streaming(&self, _req: LlmRequest) -> Observable<AssistantMessage> {
1861 unreachable!("not used by lifecycle tests")
1862 }
1863
1864 fn try_discover_models(
1865 &self,
1866 ) -> BoxFut<'static, Result<Vec<DiscoveredModelDetails>, ModelDiscoveryError>> {
1867 self.discovery_calls.fetch_add(1, Ordering::SeqCst);
1868 let models = self.models.lock().unwrap().clone();
1869 let block = self.block.lock().unwrap().take();
1870 let error = self.fail.lock().unwrap().take();
1871 let before_discovery = self.before_discovery.lock().unwrap().take();
1872 Box::pin(async move {
1873 if let Some(before_discovery) = before_discovery {
1874 before_discovery();
1875 }
1876 if let Some(block) = block {
1877 let _ = block.started.send(());
1878 let _ = block.proceed.await;
1879 }
1880 if let Some(error) = error {
1881 return Err(error);
1882 }
1883 Ok(models)
1884 })
1885 }
1886 }
1887
1888 fn provider_record(id: &str) -> StoredProvider {
1889 StoredProvider {
1890 id: id.into(),
1891 name: "Account".into(),
1892 kind: ProviderKind::Codex,
1893 access_token: "access".into(),
1894 refresh_token: Some("refresh".into()),
1895 expires_at: i64::MAX,
1896 account: Some("account@example.com".into()),
1897 enabled: true,
1898 model_cache: None,
1899 }
1900 }
1901
1902 fn model(slug: &str) -> DiscoveredModelDetails {
1903 DiscoveredModelDetails {
1904 slug: slug.into(),
1905 context_budget: Some(128_000),
1906 capability_knowledge: CapabilityKnowledge::Advertised(ModelCapabilities::default()),
1907 }
1908 }
1909
1910 fn model_with_effort(slug: &str, effort: ReasoningEffort) -> DiscoveredModelDetails {
1911 DiscoveredModelDetails {
1912 slug: slug.into(),
1913 context_budget: Some(128_000),
1914 capability_knowledge: CapabilityKnowledge::Advertised(ModelCapabilities {
1915 reasoning_efforts: vec![effort],
1916 ..Default::default()
1917 }),
1918 }
1919 }
1920
1921 fn seed_cached_provider(
1922 lifecycle: &ProviderLifecycle,
1923 id: &str,
1924 enabled: bool,
1925 models: &[DiscoveredModelDetails],
1926 ) -> String {
1927 let mut record = provider_record(id);
1928 record.enabled = enabled;
1929 let namespace = format!("{id}@account");
1930 lifecycle
1931 .hub
1932 .add_auth_provider_with_model_cache_details(record, &namespace, 1, models)
1933 .unwrap();
1934 namespace
1935 }
1936
1937 fn seed_legacy_cached_provider(lifecycle: &ProviderLifecycle, id: &str, enabled: bool) {
1938 let mut record = provider_record(id);
1939 record.enabled = enabled;
1940 record.model_cache = Some(crate::auth_store::ModelCache {
1941 fetched_at: 1,
1942 models: vec![crate::auth_store::CachedModel {
1943 slug: "cached".into(),
1944 context_budget: Some(128_000),
1945 thinking: true,
1946 }],
1947 });
1948 lifecycle.hub.add_auth_provider(record).unwrap();
1949 }
1950
1951 fn fixture(id: &str) -> (tempfile::TempDir, ProviderLifecycle, Arc<TestProvider>) {
1952 let dir = tempfile::tempdir().unwrap();
1953 let lifecycle = ProviderLifecycle::new(
1954 ConfigHub::from_config_dir(dir.path()),
1955 ProviderRegistry::new(),
1956 );
1957 let provider = Arc::new(TestProvider::new(id, vec![model("initial")]));
1958 (dir, lifecycle, provider)
1959 }
1960
1961 fn test_runtime() -> tokio::runtime::Runtime {
1962 tokio::runtime::Builder::new_current_thread()
1963 .enable_all()
1964 .build()
1965 .unwrap()
1966 }
1967
1968 #[test]
1969 fn concurrent_config_provider_updates_leave_config_and_live_registries_in_sync() {
1970 struct ConfigReset;
1971
1972 impl Drop for ConfigReset {
1973 fn drop(&mut self) {
1974 crate::model_registry::set_provider_config(Default::default());
1975 }
1976 }
1977
1978 let _registry_lock = crate::model_registry::MODEL_CONFIG_LOCK
1979 .lock()
1980 .unwrap_or_else(std::sync::PoisonError::into_inner);
1981 let _reset = ConfigReset;
1982 let dir = tempfile::tempdir().unwrap();
1983 let hub = ConfigHub::from_config_dir(dir.path());
1984 let registry = ProviderRegistry::new();
1985 let lifecycle = ProviderLifecycle::new(hub.clone(), registry.clone());
1986 lifecycle
1987 .create_config_provider(ProviderConfigUpdate {
1988 name: "gateway",
1989 kind: "openai-compat",
1990 api_key: Some("initial-key"),
1991 api_key_env: None,
1992 base_url: Some("https://gateway.example/v1"),
1993 max_tokens: Some(8_192),
1994 reasoning_format: None,
1995 prompt_cache_key: None,
1996 enabled: true,
1997 })
1998 .unwrap();
1999
2000 let peer_registry = ProviderRegistry::new();
2001 let _peer = ProviderLifecycle::new(hub.clone(), peer_registry.clone());
2002 let writers = 16;
2003 let barrier = Arc::new(std::sync::Barrier::new(writers));
2004 let handles = (0..writers)
2005 .map(|index| {
2006 let lifecycle = lifecycle.clone();
2007 let barrier = barrier.clone();
2008 std::thread::spawn(move || {
2009 let api_key = format!("key-{index}");
2010 barrier.wait();
2011 lifecycle
2012 .update_config_provider(ProviderConfigUpdate {
2013 name: "gateway",
2014 kind: "openai-compat",
2015 api_key: Some(&api_key),
2016 api_key_env: None,
2017 base_url: Some("https://gateway.example/v1"),
2018 max_tokens: Some(8_192 + index as u32),
2019 reasoning_format: None,
2020 prompt_cache_key: None,
2021 enabled: index % 2 == 0,
2022 })
2023 .unwrap()
2024 })
2025 })
2026 .collect::<Vec<_>>();
2027 for handle in handles {
2028 handle.join().unwrap();
2029 }
2030 let committed = hub.model_config().unwrap().unwrap().providers["gateway"].clone();
2031 let global = crate::model_registry::all_provider_entries()
2032 .into_iter()
2033 .find(|(name, _)| name == "gateway")
2034 .map(|(_, entry)| entry)
2035 .unwrap();
2036 assert_eq!(global.api_key, committed.api_key);
2037 assert_eq!(global.max_tokens, committed.max_tokens);
2038 assert_eq!(global.enabled, committed.enabled);
2039 let available = crate::config_provider::config_provider_availability(&committed)
2040 == crate::config_provider::ConfigProviderAvailability::Available;
2041 assert_eq!(registry.contains("config:gateway"), available);
2042 assert_eq!(peer_registry.contains("config:gateway"), available);
2043 }
2044
2045 #[test]
2046 fn public_config_mutations_sync_existing_and_new_registries() {
2047 struct ConfigReset;
2048
2049 impl Drop for ConfigReset {
2050 fn drop(&mut self) {
2051 crate::model_registry::set_provider_config(Default::default());
2052 }
2053 }
2054
2055 let _registry_lock = crate::model_registry::MODEL_CONFIG_LOCK
2056 .lock()
2057 .unwrap_or_else(std::sync::PoisonError::into_inner);
2058 let _reset = ConfigReset;
2059 let dir = tempfile::tempdir().unwrap();
2060 let hub = ConfigHub::from_config_dir(dir.path());
2061 let active = ProviderConfigUpdate {
2062 name: "gateway",
2063 kind: "openai-compat",
2064 api_key: Some("test-key"),
2065 api_key_env: None,
2066 base_url: Some("https://gateway.example/v1"),
2067 max_tokens: None,
2068 reasoning_format: None,
2069 prompt_cache_key: None,
2070 enabled: true,
2071 };
2072 hub.upsert_provider(active).unwrap();
2073 let registry = ProviderRegistry::new();
2074 let lifecycle = ProviderLifecycle::new(hub.clone(), registry.clone());
2075 lifecycle.reload_config_providers().unwrap();
2076 assert!(registry.contains("config:gateway"));
2077
2078 hub.upsert_provider(ProviderConfigUpdate {
2079 enabled: false,
2080 ..active
2081 })
2082 .unwrap();
2083 assert!(!registry.contains("config:gateway"));
2084 let before_attach_registry = ProviderRegistry::new();
2085 let _before_attach = ProviderLifecycle::new(hub.clone(), before_attach_registry.clone());
2086 assert!(!before_attach_registry.contains("config:gateway"));
2087
2088 hub.upsert_provider(active).unwrap();
2089 assert!(registry.contains("config:gateway"));
2090 assert!(before_attach_registry.contains("config:gateway"));
2091 let after_attach_registry = ProviderRegistry::new();
2092 let _after_attach = ProviderLifecycle::new(hub.clone(), after_attach_registry.clone());
2093 assert!(after_attach_registry.contains("config:gateway"));
2094
2095 std::fs::write(
2096 hub.config_toml_path(),
2097 "[providers.gateway]\nkind = \"openai-compat\"\napi_key = \"reloaded-key\"\nenabled = false\n",
2098 )
2099 .unwrap();
2100 hub.reload().unwrap();
2101 assert!(!registry.contains("config:gateway"));
2102 assert!(!before_attach_registry.contains("config:gateway"));
2103 assert!(!after_attach_registry.contains("config:gateway"));
2104 let reloaded_registry = ProviderRegistry::new();
2105 let _reloaded = ProviderLifecycle::new(hub.clone(), reloaded_registry.clone());
2106 assert!(!reloaded_registry.contains("config:gateway"));
2107 }
2108
2109 #[test]
2110 fn public_config_mutations_fan_out_across_distinct_auth_stores() {
2111 struct ConfigReset;
2112
2113 impl Drop for ConfigReset {
2114 fn drop(&mut self) {
2115 crate::model_registry::set_provider_config(Default::default());
2116 }
2117 }
2118
2119 let _registry_lock = crate::model_registry::MODEL_CONFIG_LOCK
2120 .lock()
2121 .unwrap_or_else(std::sync::PoisonError::into_inner);
2122 let _reset = ConfigReset;
2123 let dir = tempfile::tempdir().unwrap();
2124 let hub_a = ConfigHub::from_auth_path(dir.path().join("auth-a.json"));
2125 let hub_b = ConfigHub::from_auth_path(dir.path().join("auth-b.json"));
2126 let registry_a = ProviderRegistry::new();
2127 let registry_b = ProviderRegistry::new();
2128 let lifecycle_a = ProviderLifecycle::new(hub_a.clone(), registry_a.clone());
2129 let lifecycle_b = ProviderLifecycle::new(hub_b.clone(), registry_b.clone());
2130
2131 assert!(!Arc::ptr_eq(&lifecycle_a.state, &lifecycle_b.state));
2132 assert!(Arc::ptr_eq(
2133 &lifecycle_a.config_state,
2134 &lifecycle_b.config_state
2135 ));
2136
2137 let active = ProviderConfigUpdate {
2138 name: "gateway",
2139 kind: "openai-compat",
2140 api_key: Some("test-key"),
2141 api_key_env: None,
2142 base_url: Some("https://gateway.example/v1"),
2143 max_tokens: None,
2144 reasoning_format: None,
2145 prompt_cache_key: None,
2146 enabled: true,
2147 };
2148 hub_a.upsert_provider(active).unwrap();
2149 assert!(registry_a.contains("config:gateway"));
2150 assert!(registry_b.contains("config:gateway"));
2151
2152 hub_b
2153 .upsert_provider(ProviderConfigUpdate {
2154 enabled: false,
2155 ..active
2156 })
2157 .unwrap();
2158 assert!(!registry_a.contains("config:gateway"));
2159 assert!(!registry_b.contains("config:gateway"));
2160
2161 let registry_c = ProviderRegistry::new();
2162 let lifecycle_c = ProviderLifecycle::new(hub_a.clone(), registry_c.clone());
2163 assert!(Arc::ptr_eq(
2164 &lifecycle_a.config_state,
2165 &lifecycle_c.config_state
2166 ));
2167 assert!(!registry_c.contains("config:gateway"));
2168 }
2169
2170 #[test]
2171 fn public_config_reload_fans_out_across_distinct_auth_stores() {
2172 struct ConfigReset;
2173
2174 impl Drop for ConfigReset {
2175 fn drop(&mut self) {
2176 crate::model_registry::set_provider_config(Default::default());
2177 }
2178 }
2179
2180 let _registry_lock = crate::model_registry::MODEL_CONFIG_LOCK
2181 .lock()
2182 .unwrap_or_else(std::sync::PoisonError::into_inner);
2183 let _reset = ConfigReset;
2184 let dir = tempfile::tempdir().unwrap();
2185 let hub_a = ConfigHub::from_auth_path(dir.path().join("auth-a.json"));
2186 let hub_b = ConfigHub::from_auth_path(dir.path().join("auth-b.json"));
2187 let registry_a = ProviderRegistry::new();
2188 let registry_b = ProviderRegistry::new();
2189 let _lifecycle_a = ProviderLifecycle::new(hub_a.clone(), registry_a.clone());
2190 let _lifecycle_b = ProviderLifecycle::new(hub_b.clone(), registry_b.clone());
2191
2192 std::fs::write(
2193 hub_a.config_toml_path(),
2194 "[providers.gateway]\nkind = \"openai-compat\"\napi_key = \"first-key\"\nenabled = true\n",
2195 )
2196 .unwrap();
2197 hub_a.reload().unwrap();
2198 assert!(registry_a.contains("config:gateway"));
2199 assert!(registry_b.contains("config:gateway"));
2200
2201 std::fs::write(
2202 hub_b.config_toml_path(),
2203 "[providers.gateway]\nkind = \"openai-compat\"\napi_key = \"second-key\"\nenabled = false\n",
2204 )
2205 .unwrap();
2206 hub_b.reload().unwrap();
2207 assert!(!registry_a.contains("config:gateway"));
2208 assert!(!registry_b.contains("config:gateway"));
2209
2210 let projected = crate::model_registry::all_provider_entries()
2211 .into_iter()
2212 .find(|(name, _)| name == "gateway")
2213 .map(|(_, entry)| entry)
2214 .unwrap();
2215 assert_eq!(projected.api_key.as_deref(), Some("second-key"));
2216 assert_eq!(projected.enabled, Some(false));
2217 }
2218
2219 #[test]
2220 fn install_commits_auth_live_provider_and_catalog_together() {
2221 let _registry_lock = crate::model_registry::MODEL_CONFIG_LOCK
2222 .lock()
2223 .unwrap_or_else(std::sync::PoisonError::into_inner);
2224 test_runtime().block_on(async {
2225 let (_dir, lifecycle, provider) = fixture("lifecycle-install");
2226 remove_provider_catalog("lifecycle-install");
2227
2228 let delta = lifecycle
2229 .install_provider(provider_record("lifecycle-install"), provider)
2230 .await
2231 .unwrap();
2232
2233 assert_eq!(delta.total, 1);
2234 assert!(lifecycle.providers.contains("lifecycle-install"));
2235 assert_eq!(
2236 lifecycle.hub.load_auth().unwrap().providers[0].id,
2237 "lifecycle-install"
2238 );
2239 assert!(provider_catalog_namespace("lifecycle-install").is_some());
2240 lifecycle.remove_provider("lifecycle-install").unwrap();
2241 });
2242 }
2243
2244 #[test]
2245 fn pre_discovered_install_does_not_run_discovery() {
2246 let _registry_lock = crate::model_registry::MODEL_CONFIG_LOCK
2247 .lock()
2248 .unwrap_or_else(std::sync::PoisonError::into_inner);
2249 test_runtime().block_on(async {
2250 let (_dir, lifecycle, provider) = fixture("lifecycle-pre-discovered");
2251 remove_provider_catalog("lifecycle-pre-discovered");
2252 provider.fail_once(ModelDiscoveryError::Transport("unused".into()));
2253
2254 lifecycle
2255 .install_pre_discovered_provider(
2256 provider_record("lifecycle-pre-discovered"),
2257 provider.clone(),
2258 vec![model("provided")],
2259 )
2260 .await
2261 .unwrap();
2262
2263 assert_eq!(provider.discovery_calls(), 0);
2264 let namespace = lifecycle
2265 .hub
2266 .load_auth_model_namespace("lifecycle-pre-discovered")
2267 .unwrap()
2268 .unwrap();
2269 assert!(crate::model_registry::model_entry(&format!("{namespace}:provided")).is_some());
2270 assert!(matches!(
2271 lifecycle.refresh_models("lifecycle-pre-discovered").await,
2272 Err(ProviderLifecycleError::Discovery(
2273 ModelDiscoveryError::Transport(message)
2274 )) if message == "unused"
2275 ));
2276 lifecycle
2277 .remove_provider("lifecycle-pre-discovered")
2278 .unwrap();
2279 });
2280 }
2281
2282 #[test]
2283 fn cached_restore_is_offline_and_reuses_the_shared_provider() {
2284 let _registry_lock = crate::model_registry::MODEL_CONFIG_LOCK
2285 .lock()
2286 .unwrap_or_else(std::sync::PoisonError::into_inner);
2287 test_runtime().block_on(async {
2288 let (_dir, lifecycle, provider) = fixture("lifecycle-restore");
2289 remove_provider_catalog("lifecycle-restore");
2290 seed_legacy_cached_provider(&lifecycle, "lifecycle-restore", true);
2291 provider.fail_once(ModelDiscoveryError::Transport("offline".into()));
2292
2293 let first = lifecycle
2294 .restore_provider("lifecycle-restore", ProviderKind::Codex, provider.clone())
2295 .await
2296 .unwrap();
2297
2298 assert_eq!(provider.discovery_calls(), 0);
2299 assert!(first.state.auth_changed);
2300 assert!(first.state.live_changed);
2301 assert!(first.state.catalog_changed);
2302 let namespace = lifecycle
2303 .hub
2304 .load_auth_model_namespace("lifecycle-restore")
2305 .unwrap()
2306 .unwrap();
2307 assert!(crate::model_registry::model_entry(&format!("{namespace}:cached")).is_some());
2308 let expected: Arc<dyn Provider> = provider.clone();
2309 assert!(Arc::ptr_eq(
2310 &lifecycle.providers.get("lifecycle-restore").unwrap(),
2311 &expected
2312 ));
2313
2314 let peer_registry = ProviderRegistry::new();
2315 let peer = ProviderLifecycle::new(lifecycle.hub.clone(), peer_registry.clone());
2316 let replacement = Arc::new(TestProvider::new(
2317 "lifecycle-restore",
2318 vec![model("replacement")],
2319 ));
2320 let before_revision = crate::model_registry::model_catalog_revision();
2321 let second = peer
2322 .restore_provider(
2323 "lifecycle-restore",
2324 ProviderKind::Codex,
2325 replacement.clone(),
2326 )
2327 .await
2328 .unwrap();
2329
2330 assert_eq!(replacement.discovery_calls(), 0);
2331 assert_eq!(second.state, ProviderStateChange::default());
2332 assert_eq!(
2333 crate::model_registry::model_catalog_revision(),
2334 before_revision
2335 );
2336 assert!(Arc::ptr_eq(
2337 &peer_registry.get("lifecycle-restore").unwrap(),
2338 &expected
2339 ));
2340 assert!(matches!(
2341 lifecycle.refresh_models("lifecycle-restore").await,
2342 Err(ProviderLifecycleError::Discovery(
2343 ModelDiscoveryError::Transport(message)
2344 )) if message == "offline"
2345 ));
2346 assert_eq!(provider.discovery_calls(), 1);
2347 lifecycle.remove_provider("lifecycle-restore").unwrap();
2348 });
2349 }
2350
2351 #[test]
2352 fn catalog_refresh_plan_contains_restored_live_providers() {
2353 let _registry_lock = crate::model_registry::MODEL_CONFIG_LOCK
2354 .lock()
2355 .unwrap_or_else(std::sync::PoisonError::into_inner);
2356 test_runtime().block_on(async {
2357 let (_dir, lifecycle, provider) = fixture("catalog-refresh-plan");
2358 remove_provider_catalog("catalog-refresh-plan");
2359 seed_legacy_cached_provider(&lifecycle, "catalog-refresh-plan", true);
2360 lifecycle
2361 .hub
2362 .add_auth_provider(provider_record("catalog-refresh-no-live"))
2363 .unwrap();
2364 lifecycle
2365 .restore_provider(
2366 "catalog-refresh-plan",
2367 ProviderKind::Codex,
2368 provider.clone(),
2369 )
2370 .await
2371 .unwrap();
2372 assert_eq!(
2373 lifecycle.catalog_refresh_plan(),
2374 vec!["catalog-refresh-plan"]
2375 );
2376 provider.set_models(vec![model("fresh")]);
2377 lifecycle
2378 .refresh_models("catalog-refresh-plan")
2379 .await
2380 .unwrap();
2381 assert_eq!(
2382 lifecycle.catalog_refresh_plan(),
2383 vec!["catalog-refresh-plan"]
2384 );
2385 assert_eq!(
2386 lifecycle
2387 .refresh_models_if_stale("catalog-refresh-plan")
2388 .await
2389 .unwrap(),
2390 ProviderCatalogRefreshOutcome::NotNeeded
2391 );
2392 assert_eq!(provider.discovery_calls(), 1);
2393
2394 lifecycle.remove_provider("catalog-refresh-plan").unwrap();
2395 lifecycle
2396 .remove_provider("catalog-refresh-no-live")
2397 .unwrap();
2398 });
2399 }
2400
2401 #[test]
2402 fn catalog_refresh_plan_revalidates_external_provider_mutations() {
2403 let _registry_lock = crate::model_registry::MODEL_CONFIG_LOCK
2404 .lock()
2405 .unwrap_or_else(std::sync::PoisonError::into_inner);
2406 test_runtime().block_on(async {
2407 for mutation in ["remove", "disable", "kind"] {
2408 let provider_id = format!("catalog-plan-{mutation}");
2409 let (_dir, lifecycle, provider) = fixture(&provider_id);
2410 remove_provider_catalog(&provider_id);
2411 seed_legacy_cached_provider(&lifecycle, &provider_id, true);
2412 lifecycle
2413 .restore_provider(&provider_id, ProviderKind::Codex, provider)
2414 .await
2415 .unwrap();
2416
2417 let external_hub = ConfigHub::from_config_dir(lifecycle.hub.config_dir());
2418 match mutation {
2419 "remove" => {
2420 assert!(external_hub.remove_auth_provider(&provider_id).unwrap());
2421 }
2422 "disable" => {
2423 assert_eq!(
2424 external_hub
2425 .set_auth_provider_enabled_with_change(&provider_id, false)
2426 .unwrap(),
2427 Some(true)
2428 );
2429 }
2430 "kind" => external_hub
2431 .update_auth(|store| {
2432 store.providers[0].kind = ProviderKind::AnthropicOauth;
2433 Ok(())
2434 })
2435 .unwrap(),
2436 _ => unreachable!(),
2437 }
2438
2439 assert_eq!(lifecycle.catalog_refresh_plan(), vec![provider_id.clone()]);
2440 let result = lifecycle.refresh_models_if_stale(&provider_id).await;
2441 match mutation {
2442 "remove" => assert!(matches!(
2443 result,
2444 Err(ProviderLifecycleError::ProviderNotFound { .. })
2445 )),
2446 "disable" => assert!(matches!(
2447 result,
2448 Err(ProviderLifecycleError::ProviderDisabled { .. })
2449 )),
2450 "kind" => assert!(matches!(result, Err(ProviderLifecycleError::Stale { .. }))),
2451 _ => unreachable!(),
2452 }
2453 assert!(!lifecycle.providers.contains(&provider_id));
2454 assert!(provider_catalog_namespace(&provider_id).is_none());
2455 if mutation != "remove" {
2456 lifecycle.remove_provider(&provider_id).unwrap();
2457 }
2458 }
2459 });
2460 }
2461
2462 #[test]
2463 fn fresh_external_cache_is_hydrated_when_plan_observes_fresh_storage() {
2464 let _registry_lock = crate::model_registry::MODEL_CONFIG_LOCK
2465 .lock()
2466 .unwrap_or_else(std::sync::PoisonError::into_inner);
2467 test_runtime().block_on(async {
2468 let (_dir, lifecycle, provider) = fixture("catalog-refresh-external-cache");
2469 remove_provider_catalog("catalog-refresh-external-cache");
2470 seed_legacy_cached_provider(&lifecycle, "catalog-refresh-external-cache", true);
2471 lifecycle
2472 .restore_provider(
2473 "catalog-refresh-external-cache",
2474 ProviderKind::Codex,
2475 provider.clone(),
2476 )
2477 .await
2478 .unwrap();
2479 let namespace = lifecycle
2480 .hub
2481 .load_auth_model_namespace("catalog-refresh-external-cache")
2482 .unwrap()
2483 .unwrap();
2484 let external_hub = ConfigHub::from_config_dir(lifecycle.hub.config_dir());
2485 external_hub
2486 .update_auth_model_cache_details(
2487 "catalog-refresh-external-cache",
2488 &namespace,
2489 chrono::Utc::now().timestamp(),
2490 &[model_with_effort("cached", ReasoningEffort::High)],
2491 )
2492 .unwrap();
2493 assert_eq!(
2494 lifecycle.catalog_refresh_plan(),
2495 vec!["catalog-refresh-external-cache"]
2496 );
2497 let registry_key = format!("{namespace}:cached");
2498 assert!(
2499 crate::model_registry::model_entry(®istry_key)
2500 .unwrap()
2501 .reasoning_efforts
2502 .is_empty()
2503 );
2504
2505 assert!(matches!(
2506 lifecycle
2507 .refresh_models_if_stale("catalog-refresh-external-cache")
2508 .await
2509 .unwrap(),
2510 ProviderCatalogRefreshOutcome::CatalogUpdated(delta)
2511 if delta.updated == 1 && delta.total == 1
2512 ));
2513 assert_eq!(provider.discovery_calls(), 0);
2514 assert_eq!(
2515 crate::model_registry::model_entry(®istry_key)
2516 .unwrap()
2517 .reasoning_efforts,
2518 vec![ReasoningEffort::High]
2519 );
2520
2521 lifecycle
2522 .remove_provider("catalog-refresh-external-cache")
2523 .unwrap();
2524 });
2525 }
2526
2527 #[test]
2528 fn stale_network_commit_hydrates_the_external_cache_winner() {
2529 let _registry_lock = crate::model_registry::MODEL_CONFIG_LOCK
2530 .lock()
2531 .unwrap_or_else(std::sync::PoisonError::into_inner);
2532 test_runtime().block_on(async {
2533 let (_dir, lifecycle, provider) = fixture("catalog-refresh-external-winner");
2534 remove_provider_catalog("catalog-refresh-external-winner");
2535 seed_legacy_cached_provider(&lifecycle, "catalog-refresh-external-winner", true);
2536 lifecycle
2537 .restore_provider(
2538 "catalog-refresh-external-winner",
2539 ProviderKind::Codex,
2540 provider.clone(),
2541 )
2542 .await
2543 .unwrap();
2544 let namespace = lifecycle
2545 .hub
2546 .load_auth_model_namespace("catalog-refresh-external-winner")
2547 .unwrap()
2548 .unwrap();
2549 provider.set_models(vec![model_with_effort("loser", ReasoningEffort::Low)]);
2550 let (started, proceed) = provider.block_once();
2551 let background_lifecycle = lifecycle.clone();
2552 let background = tokio::spawn(async move {
2553 background_lifecycle
2554 .refresh_models_if_stale("catalog-refresh-external-winner")
2555 .await
2556 });
2557 started.await.unwrap();
2558
2559 let external_hub = ConfigHub::from_config_dir(lifecycle.hub.config_dir());
2560 external_hub
2561 .update_auth_model_cache_details(
2562 "catalog-refresh-external-winner",
2563 &namespace,
2564 chrono::Utc::now().timestamp(),
2565 &[model_with_effort("winner", ReasoningEffort::High)],
2566 )
2567 .unwrap();
2568 proceed.send(()).unwrap();
2569
2570 assert!(matches!(
2571 background.await.unwrap().unwrap(),
2572 ProviderCatalogRefreshOutcome::CatalogUpdated(delta)
2573 if delta.added == 1 && delta.removed == 1 && delta.total == 1
2574 ));
2575 assert_eq!(provider.discovery_calls(), 1);
2576 assert!(crate::model_registry::model_entry(&format!("{namespace}:loser")).is_none());
2577 assert_eq!(
2578 crate::model_registry::model_entry(&format!("{namespace}:winner"))
2579 .unwrap()
2580 .reasoning_efforts,
2581 vec![ReasoningEffort::High]
2582 );
2583 assert_eq!(
2584 lifecycle.hub.load_auth().unwrap().providers[0]
2585 .model_cache
2586 .as_ref()
2587 .unwrap()
2588 .models[0]
2589 .slug,
2590 "winner"
2591 );
2592
2593 lifecycle
2594 .remove_provider("catalog-refresh-external-winner")
2595 .unwrap();
2596 });
2597 }
2598
2599 #[test]
2600 fn failed_network_discovery_hydrates_the_external_cache_winner() {
2601 let _registry_lock = crate::model_registry::MODEL_CONFIG_LOCK
2602 .lock()
2603 .unwrap_or_else(std::sync::PoisonError::into_inner);
2604 test_runtime().block_on(async {
2605 let (_dir, lifecycle, provider) = fixture("catalog-refresh-error-winner");
2606 remove_provider_catalog("catalog-refresh-error-winner");
2607 seed_legacy_cached_provider(&lifecycle, "catalog-refresh-error-winner", true);
2608 lifecycle
2609 .restore_provider(
2610 "catalog-refresh-error-winner",
2611 ProviderKind::Codex,
2612 provider.clone(),
2613 )
2614 .await
2615 .unwrap();
2616 let namespace = lifecycle
2617 .hub
2618 .load_auth_model_namespace("catalog-refresh-error-winner")
2619 .unwrap()
2620 .unwrap();
2621 provider.fail_once(ModelDiscoveryError::Transport("offline".into()));
2622 let (started, proceed) = provider.block_once();
2623 let background_lifecycle = lifecycle.clone();
2624 let background = tokio::spawn(async move {
2625 background_lifecycle
2626 .refresh_models_if_stale("catalog-refresh-error-winner")
2627 .await
2628 });
2629 started.await.unwrap();
2630
2631 ConfigHub::from_config_dir(lifecycle.hub.config_dir())
2632 .update_auth_model_cache_details(
2633 "catalog-refresh-error-winner",
2634 &namespace,
2635 chrono::Utc::now().timestamp(),
2636 &[model_with_effort("winner", ReasoningEffort::High)],
2637 )
2638 .unwrap();
2639 proceed.send(()).unwrap();
2640
2641 assert!(matches!(
2642 background.await.unwrap().unwrap(),
2643 ProviderCatalogRefreshOutcome::CatalogUpdated(delta)
2644 if delta.added == 1 && delta.removed == 1 && delta.total == 1
2645 ));
2646 assert_eq!(provider.discovery_calls(), 1);
2647 assert_eq!(
2648 crate::model_registry::model_entry(&format!("{namespace}:winner"))
2649 .unwrap()
2650 .reasoning_efforts,
2651 vec![ReasoningEffort::High]
2652 );
2653
2654 lifecycle
2655 .remove_provider("catalog-refresh-error-winner")
2656 .unwrap();
2657 });
2658 }
2659
2660 #[test]
2661 fn cache_hydration_cas_remove_and_disable_prune_local_runtime() {
2662 let _registry_lock = crate::model_registry::MODEL_CONFIG_LOCK
2663 .lock()
2664 .unwrap_or_else(std::sync::PoisonError::into_inner);
2665 test_runtime().block_on(async {
2666 for (provider_id, remove) in [
2667 ("catalog-hydrate-removed", true),
2668 ("catalog-hydrate-disabled", false),
2669 ] {
2670 let (_dir, lifecycle, provider) = fixture(provider_id);
2671 remove_provider_catalog(provider_id);
2672 seed_cached_provider(&lifecycle, provider_id, true, &[model("cached")]);
2673 let expected_live: Arc<dyn Provider> = provider.clone();
2674 lifecycle
2675 .restore_provider(provider_id, ProviderKind::Codex, provider)
2676 .await
2677 .unwrap();
2678 let (attempt, deferred) = lifecycle.load_catalog_refresh_attempt(provider_id, 1);
2679 assert!(deferred.is_empty());
2680 let attempt = attempt.unwrap();
2681 let external_hub = ConfigHub::from_config_dir(lifecycle.hub.config_dir());
2682 if remove {
2683 external_hub.remove_auth_provider(provider_id).unwrap();
2684 } else {
2685 external_hub
2686 .set_auth_provider_enabled(provider_id, false)
2687 .unwrap();
2688 }
2689
2690 let (result, deferred) =
2691 lifecycle.hydrate_cached_catalog_if_current(provider_id, &attempt);
2692 drop(deferred);
2693 if remove {
2694 assert!(matches!(
2695 result,
2696 Err(ProviderLifecycleError::ProviderNotFound { .. })
2697 ));
2698 } else {
2699 assert!(matches!(result, Err(ProviderLifecycleError::Stale { .. })));
2700 let (result, deferred) = lifecycle.prepare_catalog_refresh(
2701 provider_id,
2702 1,
2703 true,
2704 Some(&expected_live),
2705 Some(&attempt.runtime.catalog_snapshot),
2706 );
2707 drop(deferred);
2708 assert!(matches!(
2709 result,
2710 Err(ProviderLifecycleError::ProviderDisabled { .. })
2711 ));
2712 }
2713 assert!(!lifecycle.providers.contains(provider_id));
2714 assert!(provider_catalog_namespace(provider_id).is_none());
2715 }
2716 });
2717 }
2718
2719 #[test]
2720 fn cache_hydration_cas_kind_change_is_reconciled_fail_closed() {
2721 let _registry_lock = crate::model_registry::MODEL_CONFIG_LOCK
2722 .lock()
2723 .unwrap_or_else(std::sync::PoisonError::into_inner);
2724 test_runtime().block_on(async {
2725 let provider_id = "catalog-hydrate-kind-change";
2726 let (_dir, lifecycle, provider) = fixture(provider_id);
2727 remove_provider_catalog(provider_id);
2728 seed_cached_provider(&lifecycle, provider_id, true, &[model("cached")]);
2729 lifecycle
2730 .restore_provider(provider_id, ProviderKind::Codex, provider.clone())
2731 .await
2732 .unwrap();
2733 let expected_live: Arc<dyn Provider> = provider;
2734 let (attempt, deferred) = lifecycle.load_catalog_refresh_attempt(provider_id, 1);
2735 assert!(deferred.is_empty());
2736 let attempt = attempt.unwrap();
2737 lifecycle
2738 .hub
2739 .update_auth(|store| {
2740 store.providers[0].kind = ProviderKind::AnthropicOauth;
2741 Ok(())
2742 })
2743 .unwrap();
2744
2745 let (result, deferred) =
2746 lifecycle.hydrate_cached_catalog_if_current(provider_id, &attempt);
2747 assert!(deferred.is_empty());
2748 assert!(matches!(result, Err(ProviderLifecycleError::Stale { .. })));
2749 let (result, deferred) = lifecycle.prepare_catalog_refresh(
2750 provider_id,
2751 1,
2752 true,
2753 Some(&expected_live),
2754 Some(&attempt.runtime.catalog_snapshot),
2755 );
2756 drop(deferred);
2757 assert!(matches!(result, Err(ProviderLifecycleError::Stale { .. })));
2758 assert!(!lifecycle.providers.contains(provider_id));
2759 assert!(provider_catalog_namespace(provider_id).is_none());
2760 lifecycle.remove_provider(provider_id).unwrap();
2761 });
2762 }
2763
2764 #[test]
2765 fn background_refresh_is_deduplicated_without_blocking_runtime_restore() {
2766 let _registry_lock = crate::model_registry::MODEL_CONFIG_LOCK
2767 .lock()
2768 .unwrap_or_else(std::sync::PoisonError::into_inner);
2769 test_runtime().block_on(async {
2770 let (_dir, lifecycle, provider) = fixture("catalog-refresh-background");
2771 remove_provider_catalog("catalog-refresh-background");
2772 seed_legacy_cached_provider(&lifecycle, "catalog-refresh-background", true);
2773 lifecycle
2774 .restore_provider(
2775 "catalog-refresh-background",
2776 ProviderKind::Codex,
2777 provider.clone(),
2778 )
2779 .await
2780 .unwrap();
2781 provider.set_models(vec![model("fresh")]);
2782 let (started, proceed) = provider.block_once();
2783 let background_lifecycle = lifecycle.clone();
2784 let background = tokio::spawn(async move {
2785 background_lifecycle
2786 .refresh_models_if_stale("catalog-refresh-background")
2787 .await
2788 });
2789 started.await.unwrap();
2790
2791 assert!(
2792 lifecycle
2793 .operation_lock("catalog-refresh-background")
2794 .try_lock()
2795 .is_ok()
2796 );
2797 assert_eq!(
2798 lifecycle
2799 .refresh_models_if_stale("catalog-refresh-background")
2800 .await
2801 .unwrap(),
2802 ProviderCatalogRefreshOutcome::AlreadyInFlight
2803 );
2804 let restored = lifecycle
2805 .restore_provider(
2806 "catalog-refresh-background",
2807 ProviderKind::Codex,
2808 provider.clone(),
2809 )
2810 .await
2811 .unwrap();
2812 assert_eq!(restored.state, ProviderStateChange::default());
2813
2814 proceed.send(()).unwrap();
2815 assert!(matches!(
2816 background.await.unwrap().unwrap(),
2817 ProviderCatalogRefreshOutcome::CatalogUpdated(delta) if delta.total == 1
2818 ));
2819 assert_eq!(provider.discovery_calls(), 1);
2820 assert_eq!(
2821 lifecycle
2822 .refresh_models_if_stale("catalog-refresh-background")
2823 .await
2824 .unwrap(),
2825 ProviderCatalogRefreshOutcome::NotNeeded
2826 );
2827 assert_eq!(provider.discovery_calls(), 1);
2828
2829 lifecycle
2830 .remove_provider("catalog-refresh-background")
2831 .unwrap();
2832 });
2833 }
2834
2835 #[test]
2836 fn concurrent_background_refresh_returns_already_in_flight() {
2837 let _registry_lock = crate::model_registry::MODEL_CONFIG_LOCK
2838 .lock()
2839 .unwrap_or_else(std::sync::PoisonError::into_inner);
2840 test_runtime().block_on(async {
2841 let (_dir, lifecycle, provider) = fixture("catalog-refresh-failure");
2842 remove_provider_catalog("catalog-refresh-failure");
2843 seed_legacy_cached_provider(&lifecycle, "catalog-refresh-failure", true);
2844 lifecycle
2845 .restore_provider(
2846 "catalog-refresh-failure",
2847 ProviderKind::Codex,
2848 provider.clone(),
2849 )
2850 .await
2851 .unwrap();
2852 provider.fail_once(ModelDiscoveryError::Transport("offline".into()));
2853 let (started, proceed) = provider.block_once();
2854 let background_lifecycle = lifecycle.clone();
2855 let background = tokio::spawn(async move {
2856 background_lifecycle
2857 .refresh_models_if_stale("catalog-refresh-failure")
2858 .await
2859 });
2860 started.await.unwrap();
2861
2862 assert_eq!(
2863 lifecycle
2864 .refresh_models_if_stale("catalog-refresh-failure")
2865 .await
2866 .unwrap(),
2867 ProviderCatalogRefreshOutcome::AlreadyInFlight
2868 );
2869 assert_eq!(provider.discovery_calls(), 1);
2870 proceed.send(()).unwrap();
2871 assert!(matches!(
2872 background.await.unwrap(),
2873 Err(ProviderLifecycleError::Discovery(
2874 ModelDiscoveryError::Transport(message)
2875 )) if message == "offline"
2876 ));
2877 assert_eq!(provider.discovery_calls(), 1);
2878
2879 lifecycle
2880 .remove_provider("catalog-refresh-failure")
2881 .unwrap();
2882 });
2883 }
2884
2885 #[test]
2886 fn enabling_a_cached_provider_is_atomic_and_offline() {
2887 let _registry_lock = crate::model_registry::MODEL_CONFIG_LOCK
2888 .lock()
2889 .unwrap_or_else(std::sync::PoisonError::into_inner);
2890 test_runtime().block_on(async {
2891 let (_dir, lifecycle, provider) = fixture("lifecycle-enable");
2892 remove_provider_catalog("lifecycle-enable");
2893 let namespace =
2894 seed_cached_provider(&lifecycle, "lifecycle-enable", false, &[model("cached")]);
2895 provider.fail_once(ModelDiscoveryError::Transport("unused".into()));
2896
2897 let outcome = lifecycle
2898 .enable_provider("lifecycle-enable", ProviderKind::Codex, provider.clone())
2899 .await
2900 .unwrap();
2901
2902 assert_eq!(provider.discovery_calls(), 0);
2903 assert_eq!(
2904 outcome.state,
2905 ProviderStateChange {
2906 auth_changed: true,
2907 live_changed: true,
2908 catalog_changed: true,
2909 }
2910 );
2911 assert!(lifecycle.hub.load_auth().unwrap().providers[0].enabled);
2912 assert!(lifecycle.providers.contains("lifecycle-enable"));
2913 assert!(crate::model_registry::model_entry(&format!("{namespace}:cached")).is_some());
2914 lifecycle.remove_provider("lifecycle-enable").unwrap();
2915 });
2916 }
2917
2918 #[test]
2919 fn disabled_restore_ignores_an_invalid_cached_catalog() {
2920 let _registry_lock = crate::model_registry::MODEL_CONFIG_LOCK
2921 .lock()
2922 .unwrap_or_else(std::sync::PoisonError::into_inner);
2923 test_runtime().block_on(async {
2924 let (_dir, lifecycle, provider) = fixture("lifecycle-disabled-cache");
2925 remove_provider_catalog("lifecycle-disabled-cache");
2926 let mut record = provider_record("lifecycle-disabled-cache");
2927 record.enabled = false;
2928 record.model_cache = Some(crate::auth_store::ModelCache {
2929 fetched_at: 1,
2930 models: vec![
2931 crate::auth_store::CachedModel {
2932 slug: "duplicate".into(),
2933 context_budget: None,
2934 thinking: false,
2935 },
2936 crate::auth_store::CachedModel {
2937 slug: "duplicate".into(),
2938 context_budget: None,
2939 thinking: false,
2940 },
2941 ],
2942 });
2943 lifecycle.hub.add_auth_provider(record).unwrap();
2944
2945 assert!(matches!(
2946 lifecycle
2947 .restore_provider(
2948 "lifecycle-disabled-cache",
2949 ProviderKind::Codex,
2950 provider.clone(),
2951 )
2952 .await,
2953 Err(ProviderLifecycleError::ProviderDisabled { id })
2954 if id == "lifecycle-disabled-cache"
2955 ));
2956 assert_eq!(provider.discovery_calls(), 0);
2957 assert!(!lifecycle.providers.contains("lifecycle-disabled-cache"));
2958 assert!(provider_catalog_namespace("lifecycle-disabled-cache").is_none());
2959 lifecycle
2960 .remove_provider("lifecycle-disabled-cache")
2961 .unwrap();
2962 });
2963 }
2964
2965 #[test]
2966 fn restoring_without_cache_removes_the_stale_catalog() {
2967 let _registry_lock = crate::model_registry::MODEL_CONFIG_LOCK
2968 .lock()
2969 .unwrap_or_else(std::sync::PoisonError::into_inner);
2970 test_runtime().block_on(async {
2971 let (_dir, lifecycle, provider) = fixture("lifecycle-empty-cache");
2972 remove_provider_catalog("lifecycle-empty-cache");
2973 lifecycle
2974 .install_provider(provider_record("lifecycle-empty-cache"), provider.clone())
2975 .await
2976 .unwrap();
2977 lifecycle
2978 .hub
2979 .update_auth(|store| {
2980 store.providers[0].model_cache = None;
2981 Ok(())
2982 })
2983 .unwrap();
2984 let replacement = Arc::new(TestProvider::new(
2985 "lifecycle-empty-cache",
2986 vec![model("unused")],
2987 ));
2988
2989 let outcome = lifecycle
2990 .restore_provider(
2991 "lifecycle-empty-cache",
2992 ProviderKind::Codex,
2993 replacement.clone(),
2994 )
2995 .await
2996 .unwrap();
2997
2998 assert_eq!(provider.discovery_calls(), 1);
2999 assert_eq!(replacement.discovery_calls(), 0);
3000 assert!(!outcome.state.live_changed);
3001 assert!(outcome.state.catalog_changed);
3002 assert!(outcome.catalog.is_none());
3003 assert!(lifecycle.providers.contains("lifecycle-empty-cache"));
3004 assert!(provider_catalog_namespace("lifecycle-empty-cache").is_none());
3005 lifecycle.remove_provider("lifecycle-empty-cache").unwrap();
3006 });
3007 }
3008
3009 #[test]
3010 fn reconcile_prunes_external_disable_remove_and_kind_change() {
3011 let _registry_lock = crate::model_registry::MODEL_CONFIG_LOCK
3012 .lock()
3013 .unwrap_or_else(std::sync::PoisonError::into_inner);
3014 test_runtime().block_on(async {
3015 let dir = tempfile::tempdir().unwrap();
3016 let lifecycle = ProviderLifecycle::new(
3017 ConfigHub::from_config_dir(dir.path()),
3018 ProviderRegistry::new(),
3019 );
3020 let ids = ["lifecycle-disabled", "lifecycle-removed", "lifecycle-kind"];
3021 for id in ids {
3022 remove_provider_catalog(id);
3023 lifecycle
3024 .install_provider(
3025 provider_record(id),
3026 Arc::new(TestProvider::new(id, vec![model("cached")])),
3027 )
3028 .await
3029 .unwrap();
3030 }
3031 let peer_registry = ProviderRegistry::new();
3032 let _peer = ProviderLifecycle::new(lifecycle.hub.clone(), peer_registry.clone());
3033
3034 lifecycle
3035 .hub
3036 .set_auth_provider_enabled("lifecycle-disabled", false)
3037 .unwrap();
3038 lifecycle
3039 .hub
3040 .remove_auth_provider("lifecycle-removed")
3041 .unwrap();
3042 lifecycle
3043 .hub
3044 .update_auth(|store| {
3045 store
3046 .providers
3047 .iter_mut()
3048 .find(|provider| provider.id == "lifecycle-kind")
3049 .unwrap()
3050 .kind = ProviderKind::AnthropicOauth;
3051 Ok(())
3052 })
3053 .unwrap();
3054
3055 let outcome = lifecycle.reconcile_inactive_providers().unwrap();
3056
3057 assert_eq!(outcome.providers.len(), 3);
3058 for id in ids {
3059 assert!(!lifecycle.providers.contains(id));
3060 assert!(!peer_registry.contains(id));
3061 assert!(provider_catalog_namespace(id).is_none());
3062 }
3063 let store = lifecycle.hub.load_auth().unwrap();
3064 assert_eq!(store.providers.len(), 2);
3065 assert!(
3066 !store
3067 .providers
3068 .iter()
3069 .find(|provider| provider.id == "lifecycle-disabled")
3070 .unwrap()
3071 .enabled
3072 );
3073 assert_eq!(
3074 store
3075 .providers
3076 .iter()
3077 .find(|provider| provider.id == "lifecycle-kind")
3078 .unwrap()
3079 .kind,
3080 ProviderKind::AnthropicOauth
3081 );
3082 lifecycle.remove_provider("lifecycle-disabled").unwrap();
3083 lifecycle.remove_provider("lifecycle-kind").unwrap();
3084 });
3085 }
3086
3087 #[test]
3088 fn reconcile_auth_parse_error_prunes_live_runtime() {
3089 let _registry_lock = crate::model_registry::MODEL_CONFIG_LOCK
3090 .lock()
3091 .unwrap_or_else(std::sync::PoisonError::into_inner);
3092 test_runtime().block_on(async {
3093 let (_dir, lifecycle, provider) = fixture("lifecycle-reconcile-invalid");
3094 remove_provider_catalog("lifecycle-reconcile-invalid");
3095 lifecycle
3096 .install_provider(provider_record("lifecycle-reconcile-invalid"), provider)
3097 .await
3098 .unwrap();
3099 std::fs::write(lifecycle.hub.auth_path(), b"{").unwrap();
3100
3101 assert!(matches!(
3102 lifecycle.reconcile_inactive_providers(),
3103 Err(ProviderLifecycleError::Config(_))
3104 ));
3105 assert!(!lifecycle.providers.contains("lifecycle-reconcile-invalid"));
3106 assert!(provider_catalog_namespace("lifecycle-reconcile-invalid").is_none());
3107 });
3108 }
3109
3110 #[test]
3111 fn provider_mutation_auth_errors_fail_closed() {
3112 let _registry_lock = crate::model_registry::MODEL_CONFIG_LOCK
3113 .lock()
3114 .unwrap_or_else(std::sync::PoisonError::into_inner);
3115 test_runtime().block_on(async {
3116 for (id, remove) in [
3117 ("lifecycle-disable-invalid", false),
3118 ("lifecycle-remove-invalid", true),
3119 ] {
3120 let (_dir, lifecycle, provider) = fixture(id);
3121 remove_provider_catalog(id);
3122 lifecycle
3123 .install_provider(provider_record(id), provider)
3124 .await
3125 .unwrap();
3126 std::fs::write(lifecycle.hub.auth_path(), b"{").unwrap();
3127
3128 let result = if remove {
3129 lifecycle.remove_provider(id)
3130 } else {
3131 lifecycle.disable_provider(id)
3132 };
3133 assert!(matches!(result, Err(ProviderLifecycleError::Config(_))));
3134 assert!(!lifecycle.providers.contains(id));
3135 assert!(provider_catalog_namespace(id).is_none());
3136 }
3137 });
3138 }
3139
3140 #[test]
3141 fn kind_change_after_construction_fails_closed() {
3142 let _registry_lock = crate::model_registry::MODEL_CONFIG_LOCK
3143 .lock()
3144 .unwrap_or_else(std::sync::PoisonError::into_inner);
3145 test_runtime().block_on(async {
3146 let (_dir, lifecycle, provider) = fixture("lifecycle-kind-race");
3147 remove_provider_catalog("lifecycle-kind-race");
3148 seed_cached_provider(&lifecycle, "lifecycle-kind-race", true, &[model("cached")]);
3149 lifecycle
3150 .restore_provider("lifecycle-kind-race", ProviderKind::Codex, provider)
3151 .await
3152 .unwrap();
3153 lifecycle
3154 .hub
3155 .update_auth(|store| {
3156 store.providers[0].kind = ProviderKind::AnthropicOauth;
3157 Ok(())
3158 })
3159 .unwrap();
3160 let stale_candidate = Arc::new(TestProvider::new(
3161 "lifecycle-kind-race",
3162 vec![model("unused")],
3163 ));
3164
3165 assert!(matches!(
3166 lifecycle
3167 .restore_provider(
3168 "lifecycle-kind-race",
3169 ProviderKind::Codex,
3170 stale_candidate,
3171 )
3172 .await,
3173 Err(ProviderLifecycleError::Stale { id }) if id == "lifecycle-kind-race"
3174 ));
3175 assert!(!lifecycle.providers.contains("lifecycle-kind-race"));
3176 assert!(provider_catalog_namespace("lifecycle-kind-race").is_none());
3177 lifecycle.remove_provider("lifecycle-kind-race").unwrap();
3178 });
3179 }
3180
3181 #[test]
3182 fn invalid_cached_catalog_fails_closed() {
3183 let _registry_lock = crate::model_registry::MODEL_CONFIG_LOCK
3184 .lock()
3185 .unwrap_or_else(std::sync::PoisonError::into_inner);
3186 test_runtime().block_on(async {
3187 let (_dir, lifecycle, provider) = fixture("lifecycle-invalid-cache");
3188 remove_provider_catalog("lifecycle-invalid-cache");
3189 seed_cached_provider(
3190 &lifecycle,
3191 "lifecycle-invalid-cache",
3192 true,
3193 &[model("cached")],
3194 );
3195 let operation_lock = lifecycle.operation_lock("lifecycle-invalid-cache");
3196 let dropped = Arc::new(AtomicBool::new(false));
3197 let dropped_from_callback = dropped.clone();
3198 provider.on_drop(move || {
3199 assert!(operation_lock.try_lock().is_ok());
3200 dropped_from_callback.store(true, Ordering::SeqCst);
3201 });
3202 lifecycle
3203 .restore_provider("lifecycle-invalid-cache", ProviderKind::Codex, provider)
3204 .await
3205 .unwrap();
3206 lifecycle
3207 .hub
3208 .update_auth(|store| {
3209 store.providers[0].model_cache = Some(crate::auth_store::ModelCache {
3210 fetched_at: 2,
3211 models: vec![
3212 crate::auth_store::CachedModel {
3213 slug: "duplicate".into(),
3214 context_budget: None,
3215 thinking: false,
3216 },
3217 crate::auth_store::CachedModel {
3218 slug: "duplicate".into(),
3219 context_budget: None,
3220 thinking: false,
3221 },
3222 ],
3223 });
3224 Ok(())
3225 })
3226 .unwrap();
3227
3228 assert!(matches!(
3229 lifecycle
3230 .restore_provider(
3231 "lifecycle-invalid-cache",
3232 ProviderKind::Codex,
3233 Arc::new(TestProvider::new(
3234 "lifecycle-invalid-cache",
3235 vec![model("unused")],
3236 )),
3237 )
3238 .await,
3239 Err(ProviderLifecycleError::Catalog(CatalogError::DuplicateModel { model }))
3240 if model == "duplicate"
3241 ));
3242 assert!(!lifecycle.providers.contains("lifecycle-invalid-cache"));
3243 assert!(provider_catalog_namespace("lifecycle-invalid-cache").is_none());
3244 assert!(dropped.load(Ordering::SeqCst));
3245 lifecycle
3246 .remove_provider("lifecycle-invalid-cache")
3247 .unwrap();
3248 });
3249 }
3250
3251 #[test]
3252 fn refresh_cancellation_releases_the_operation_before_provider_drop() {
3253 let _registry_lock = crate::model_registry::MODEL_CONFIG_LOCK
3254 .lock()
3255 .unwrap_or_else(std::sync::PoisonError::into_inner);
3256 test_runtime().block_on(async {
3257 let (_dir, lifecycle, provider) = fixture("lifecycle-refresh-cancel");
3258 remove_provider_catalog("lifecycle-refresh-cancel");
3259 lifecycle
3260 .install_provider(
3261 provider_record("lifecycle-refresh-cancel"),
3262 provider.clone(),
3263 )
3264 .await
3265 .unwrap();
3266 let operation_lock = lifecycle.operation_lock("lifecycle-refresh-cancel");
3267 let dropped = Arc::new(AtomicBool::new(false));
3268 let dropped_from_callback = dropped.clone();
3269 provider.on_drop(move || {
3270 assert!(operation_lock.try_lock().is_ok());
3271 dropped_from_callback.store(true, Ordering::SeqCst);
3272 });
3273 let (started, _proceed) = provider.block_once();
3274 drop(provider);
3275 let refresh_lifecycle = lifecycle.clone();
3276 let refresh = tokio::spawn(async move {
3277 refresh_lifecycle
3278 .refresh_models("lifecycle-refresh-cancel")
3279 .await
3280 });
3281 started.await.unwrap();
3282 lifecycle
3283 .disable_provider("lifecycle-refresh-cancel")
3284 .unwrap();
3285
3286 refresh.abort();
3287 assert!(refresh.await.unwrap_err().is_cancelled());
3288 assert!(dropped.load(Ordering::SeqCst));
3289 lifecycle
3290 .remove_provider("lifecycle-refresh-cancel")
3291 .unwrap();
3292 });
3293 }
3294
3295 #[test]
3296 fn provider_id_compare_and_swap_rejects_stale_namespaces() {
3297 let dir = tempfile::tempdir().unwrap();
3298 let hub = ConfigHub::from_config_dir(dir.path());
3299 let mut target = provider_record("namespace-target");
3300 target.enabled = false;
3301 hub.add_auth_provider(target).unwrap();
3302 let runtime = hub
3303 .load_or_create_auth_provider_runtime_state("namespace-target")
3304 .unwrap()
3305 .unwrap();
3306 hub.add_auth_provider(provider_record("namespace-peer"))
3307 .unwrap();
3308 let callback_ran = AtomicBool::new(false);
3309
3310 let (commit, callback) = hub
3311 .commit_auth_provider_runtime_if_current_and_then(
3312 &runtime.provider,
3313 &runtime.catalog_snapshot,
3314 true,
3315 Some("namespace-target@account"),
3316 Some(&runtime.provider_ids),
3317 || callback_ran.store(true, Ordering::SeqCst),
3318 )
3319 .unwrap();
3320
3321 assert_eq!(commit, AuthProviderRuntimeCommit::Changed);
3322 assert!(callback.is_none());
3323 assert!(!callback_ran.load(Ordering::SeqCst));
3324 let stored = hub
3325 .load_auth()
3326 .unwrap()
3327 .providers
3328 .into_iter()
3329 .find(|provider| provider.id == "namespace-target")
3330 .unwrap();
3331 assert!(!stored.enabled);
3332 assert_eq!(
3333 hub.load_auth_model_namespace("namespace-target").unwrap(),
3334 None
3335 );
3336 }
3337
3338 #[test]
3339 fn provider_insert_rejects_a_changed_provider_id_set() {
3340 let dir = tempfile::tempdir().unwrap();
3341 let hub = ConfigHub::from_config_dir(dir.path());
3342 let expected_provider_ids = Vec::new();
3343 hub.add_auth_provider(provider_record("insert-peer"))
3344 .unwrap();
3345 let callback_ran = AtomicBool::new(false);
3346
3347 let (commit, callback) = hub
3348 .add_auth_provider_with_model_cache_details_if_provider_ids_and_then(
3349 provider_record("insert-target"),
3350 Some(&expected_provider_ids),
3351 "insert-target@account",
3352 1,
3353 &[model("cached")],
3354 || callback_ran.store(true, Ordering::SeqCst),
3355 )
3356 .unwrap();
3357
3358 assert_eq!(commit, AuthProviderInsertCommit::Changed);
3359 assert!(callback.is_none());
3360 assert!(!callback_ran.load(Ordering::SeqCst));
3361 assert!(
3362 hub.load_auth()
3363 .unwrap()
3364 .providers
3365 .iter()
3366 .all(|provider| provider.id != "insert-target")
3367 );
3368 }
3369
3370 #[test]
3371 fn provider_drop_runs_after_lifecycle_locks_are_released() {
3372 let _registry_lock = crate::model_registry::MODEL_CONFIG_LOCK
3373 .lock()
3374 .unwrap_or_else(std::sync::PoisonError::into_inner);
3375 test_runtime().block_on(async {
3376 let (_dir, lifecycle, provider) = fixture("lifecycle-drop");
3377 remove_provider_catalog("lifecycle-drop");
3378 lifecycle
3379 .install_provider(provider_record("lifecycle-drop"), provider.clone())
3380 .await
3381 .unwrap();
3382 let state = Arc::downgrade(&lifecycle.state);
3383 let registry = lifecycle.providers.clone();
3384 let hub = lifecycle.hub.clone();
3385 let dropped = Arc::new(AtomicBool::new(false));
3386 let dropped_from_callback = dropped.clone();
3387 provider.on_drop(move || {
3388 assert!(state.upgrade().unwrap().try_lock().is_ok());
3389 assert!(registry.get("lifecycle-drop").is_none());
3390 assert!(
3391 !hub.set_auth_provider_enabled("lifecycle-drop", false)
3392 .unwrap()
3393 );
3394 dropped_from_callback.store(true, Ordering::SeqCst);
3395 });
3396 drop(provider);
3397
3398 lifecycle.remove_provider("lifecycle-drop").unwrap();
3399
3400 assert!(dropped.load(Ordering::SeqCst));
3401 });
3402 }
3403
3404 #[test]
3405 fn failed_refresh_preserves_cache_and_catalog_revision() {
3406 let _registry_lock = crate::model_registry::MODEL_CONFIG_LOCK
3407 .lock()
3408 .unwrap_or_else(std::sync::PoisonError::into_inner);
3409 test_runtime().block_on(async {
3410 let (_dir, lifecycle, provider) = fixture("lifecycle-failure");
3411 remove_provider_catalog("lifecycle-failure");
3412 lifecycle
3413 .install_provider(provider_record("lifecycle-failure"), provider.clone())
3414 .await
3415 .unwrap();
3416 let auth_path = lifecycle.hub.config_dir().join("auth.json");
3417 let before_auth = std::fs::read(&auth_path).unwrap();
3418 let before_revision = crate::model_registry::model_catalog_revision();
3419 provider.fail_once(ModelDiscoveryError::Transport("offline".into()));
3420
3421 assert!(matches!(
3422 lifecycle.refresh_models("lifecycle-failure").await,
3423 Err(ProviderLifecycleError::Discovery(
3424 ModelDiscoveryError::Transport(message)
3425 )) if message == "offline"
3426 ));
3427 assert_eq!(std::fs::read(auth_path).unwrap(), before_auth);
3428 assert_eq!(
3429 crate::model_registry::model_catalog_revision(),
3430 before_revision
3431 );
3432 lifecycle.remove_provider("lifecycle-failure").unwrap();
3433 });
3434 }
3435
3436 #[test]
3437 fn refresh_auth_parse_error_prunes_live_runtime() {
3438 let _registry_lock = crate::model_registry::MODEL_CONFIG_LOCK
3439 .lock()
3440 .unwrap_or_else(std::sync::PoisonError::into_inner);
3441 test_runtime().block_on(async {
3442 let (_dir, lifecycle, provider) = fixture("lifecycle-refresh-invalid-auth");
3443 remove_provider_catalog("lifecycle-refresh-invalid-auth");
3444 lifecycle
3445 .install_provider(
3446 provider_record("lifecycle-refresh-invalid-auth"),
3447 provider.clone(),
3448 )
3449 .await
3450 .unwrap();
3451 let auth_path = lifecycle.hub.auth_path().to_path_buf();
3452 provider.before_discovery_once(move || {
3453 std::fs::write(auth_path, b"{").unwrap();
3454 });
3455
3456 assert!(matches!(
3457 lifecycle
3458 .refresh_models("lifecycle-refresh-invalid-auth")
3459 .await,
3460 Err(ProviderLifecycleError::Config(_))
3461 ));
3462 assert!(
3463 !lifecycle
3464 .providers
3465 .contains("lifecycle-refresh-invalid-auth")
3466 );
3467 assert!(provider_catalog_namespace("lifecycle-refresh-invalid-auth").is_none());
3468 });
3469 }
3470
3471 #[test]
3472 fn disable_invalidates_in_flight_refresh_without_resurrection() {
3473 let _registry_lock = crate::model_registry::MODEL_CONFIG_LOCK
3474 .lock()
3475 .unwrap_or_else(std::sync::PoisonError::into_inner);
3476 test_runtime().block_on(async {
3477 let (_dir, lifecycle, provider) = fixture("lifecycle-race");
3478 remove_provider_catalog("lifecycle-race");
3479 lifecycle
3480 .install_provider(provider_record("lifecycle-race"), provider.clone())
3481 .await
3482 .unwrap();
3483 let peer_registry = ProviderRegistry::new();
3484 let peer = ProviderLifecycle::new(lifecycle.hub.clone(), peer_registry.clone());
3485 assert!(Arc::ptr_eq(&lifecycle.state, &peer.state));
3486 assert!(peer_registry.contains("lifecycle-race"));
3487 provider.set_models(vec![model("replacement")]);
3488 let (started, proceed) = provider.block_once();
3489 let refresh_lifecycle = lifecycle.clone();
3490 let refresh =
3491 tokio::spawn(
3492 async move { refresh_lifecycle.refresh_models("lifecycle-race").await },
3493 );
3494 started.await.unwrap();
3495
3496 peer.disable_provider("lifecycle-race").unwrap();
3497 proceed.send(()).unwrap();
3498 assert!(matches!(
3499 refresh.await.unwrap(),
3500 Err(ProviderLifecycleError::Stale { id }) if id == "lifecycle-race"
3501 ));
3502 assert!(!lifecycle.providers.contains("lifecycle-race"));
3503 assert!(!peer_registry.contains("lifecycle-race"));
3504 assert!(provider_catalog_namespace("lifecycle-race").is_none());
3505 let stored = lifecycle.hub.load_auth().unwrap().providers.remove(0);
3506 assert!(!stored.enabled);
3507 assert_eq!(stored.model_cache.unwrap().models[0].slug, "initial");
3508 lifecycle.remove_provider("lifecycle-race").unwrap();
3509 });
3510 }
3511
3512 #[test]
3513 fn independent_lifecycles_serialize_refreshes_for_the_same_auth_store() {
3514 let _registry_lock = crate::model_registry::MODEL_CONFIG_LOCK
3515 .lock()
3516 .unwrap_or_else(std::sync::PoisonError::into_inner);
3517 test_runtime().block_on(async {
3518 let (_dir, lifecycle, provider) = fixture("lifecycle-shared-gate");
3519 remove_provider_catalog("lifecycle-shared-gate");
3520 lifecycle
3521 .install_provider(provider_record("lifecycle-shared-gate"), provider.clone())
3522 .await
3523 .unwrap();
3524 let peer_registry = ProviderRegistry::new();
3525 let peer = ProviderLifecycle::new(lifecycle.hub.clone(), peer_registry.clone());
3526 assert!(Arc::ptr_eq(&lifecycle.state, &peer.state));
3527 assert!(peer_registry.contains("lifecycle-shared-gate"));
3528
3529 provider.set_models(vec![model("slow")]);
3530 let (started, proceed) = provider.block_once();
3531 let slow_lifecycle = lifecycle.clone();
3532 let slow = tokio::spawn(async move {
3533 slow_lifecycle.refresh_models("lifecycle-shared-gate").await
3534 });
3535 started.await.unwrap();
3536 provider.set_models(vec![model("fresh")]);
3537 let fast =
3538 tokio::spawn(async move { peer.refresh_models("lifecycle-shared-gate").await });
3539 tokio::task::yield_now().await;
3540 assert!(!fast.is_finished());
3541
3542 proceed.send(()).unwrap();
3543 slow.await.unwrap().unwrap();
3544 fast.await.unwrap().unwrap();
3545
3546 let stored = lifecycle.hub.load_auth().unwrap().providers.remove(0);
3547 assert_eq!(stored.model_cache.unwrap().models[0].slug, "fresh");
3548 let namespace = lifecycle
3549 .hub
3550 .load_auth_model_namespace("lifecycle-shared-gate")
3551 .unwrap()
3552 .unwrap();
3553 assert!(crate::model_registry::model_entry(&format!("{namespace}:fresh")).is_some());
3554 assert!(crate::model_registry::model_entry(&format!("{namespace}:slow")).is_none());
3555 lifecycle.remove_provider("lifecycle-shared-gate").unwrap();
3556 });
3557 }
3558
3559 #[test]
3560 fn external_cache_aba_hydrates_the_authoritative_catalog() {
3561 let _registry_lock = crate::model_registry::MODEL_CONFIG_LOCK
3562 .lock()
3563 .unwrap_or_else(std::sync::PoisonError::into_inner);
3564 test_runtime().block_on(async {
3565 let (_dir, lifecycle, provider) = fixture("lifecycle-cache-cas");
3566 remove_provider_catalog("lifecycle-cache-cas");
3567 lifecycle
3568 .install_provider(provider_record("lifecycle-cache-cas"), provider.clone())
3569 .await
3570 .unwrap();
3571 let namespace = lifecycle
3572 .hub
3573 .load_auth_model_namespace("lifecycle-cache-cas")
3574 .unwrap()
3575 .unwrap();
3576 let initial_fetched_at = lifecycle.hub.load_auth().unwrap().providers[0]
3577 .model_cache
3578 .as_ref()
3579 .unwrap()
3580 .fetched_at;
3581 let initial = lifecycle
3582 .hub
3583 .load_auth_model_cache_details("lifecycle-cache-cas")
3584 .unwrap()
3585 .unwrap();
3586 provider.set_models(vec![model("slow")]);
3587 let (started, proceed) = provider.block_once();
3588 let refresh_lifecycle = lifecycle.clone();
3589 let refresh = tokio::spawn(async move {
3590 refresh_lifecycle
3591 .refresh_models("lifecycle-cache-cas")
3592 .await
3593 });
3594 started.await.unwrap();
3595
3596 assert!(
3597 lifecycle
3598 .hub
3599 .update_auth_model_cache_details(
3600 "lifecycle-cache-cas",
3601 &namespace,
3602 99,
3603 &[model("external")],
3604 )
3605 .unwrap()
3606 );
3607 assert!(
3608 lifecycle
3609 .hub
3610 .update_auth_model_cache_details(
3611 "lifecycle-cache-cas",
3612 &namespace,
3613 initial_fetched_at,
3614 &initial,
3615 )
3616 .unwrap()
3617 );
3618 proceed.send(()).unwrap();
3619 assert_eq!(
3620 refresh.await.unwrap().unwrap(),
3621 CatalogDelta {
3622 total: 1,
3623 ..Default::default()
3624 }
3625 );
3626
3627 let stored = lifecycle.hub.load_auth().unwrap().providers.remove(0);
3628 assert_eq!(stored.model_cache.unwrap().models[0].slug, "initial");
3629 assert!(crate::model_registry::model_entry(&format!("{namespace}:initial")).is_some());
3630 assert!(crate::model_registry::model_entry(&format!("{namespace}:slow")).is_none());
3631 lifecycle.remove_provider("lifecycle-cache-cas").unwrap();
3632 });
3633 }
3634
3635 #[test]
3636 fn external_enable_aba_preserves_the_authoritative_cache() {
3637 let _registry_lock = crate::model_registry::MODEL_CONFIG_LOCK
3638 .lock()
3639 .unwrap_or_else(std::sync::PoisonError::into_inner);
3640 test_runtime().block_on(async {
3641 let (_dir, lifecycle, provider) = fixture("lifecycle-enable-aba");
3642 remove_provider_catalog("lifecycle-enable-aba");
3643 lifecycle
3644 .install_provider(provider_record("lifecycle-enable-aba"), provider.clone())
3645 .await
3646 .unwrap();
3647 provider.set_models(vec![model("slow")]);
3648 let (started, proceed) = provider.block_once();
3649 let refresh_lifecycle = lifecycle.clone();
3650 let refresh = tokio::spawn(async move {
3651 refresh_lifecycle
3652 .refresh_models("lifecycle-enable-aba")
3653 .await
3654 });
3655 started.await.unwrap();
3656
3657 lifecycle
3658 .hub
3659 .set_auth_provider_enabled("lifecycle-enable-aba", false)
3660 .unwrap();
3661 lifecycle
3662 .hub
3663 .set_auth_provider_enabled("lifecycle-enable-aba", true)
3664 .unwrap();
3665 proceed.send(()).unwrap();
3666
3667 assert_eq!(
3668 refresh.await.unwrap().unwrap(),
3669 CatalogDelta {
3670 total: 1,
3671 ..Default::default()
3672 }
3673 );
3674 let stored = lifecycle.hub.load_auth().unwrap().providers.remove(0);
3675 assert!(stored.enabled);
3676 assert_eq!(stored.model_cache.unwrap().models[0].slug, "initial");
3677 lifecycle.remove_provider("lifecycle-enable-aba").unwrap();
3678 });
3679 }
3680
3681 #[test]
3682 fn external_kind_change_during_refresh_prunes_the_old_runtime() {
3683 let _registry_lock = crate::model_registry::MODEL_CONFIG_LOCK
3684 .lock()
3685 .unwrap_or_else(std::sync::PoisonError::into_inner);
3686 test_runtime().block_on(async {
3687 let (_dir, lifecycle, provider) = fixture("lifecycle-refresh-kind");
3688 remove_provider_catalog("lifecycle-refresh-kind");
3689 lifecycle
3690 .install_provider(provider_record("lifecycle-refresh-kind"), provider.clone())
3691 .await
3692 .unwrap();
3693 provider.set_models(vec![model("replacement")]);
3694 let (started, proceed) = provider.block_once();
3695 let refresh_lifecycle = lifecycle.clone();
3696 let refresh = tokio::spawn(async move {
3697 refresh_lifecycle
3698 .refresh_models("lifecycle-refresh-kind")
3699 .await
3700 });
3701 started.await.unwrap();
3702 lifecycle
3703 .hub
3704 .update_auth(|store| {
3705 store.providers[0].kind = ProviderKind::AnthropicOauth;
3706 Ok(())
3707 })
3708 .unwrap();
3709 proceed.send(()).unwrap();
3710
3711 assert!(matches!(
3712 refresh.await.unwrap(),
3713 Err(ProviderLifecycleError::Stale { id }) if id == "lifecycle-refresh-kind"
3714 ));
3715 assert!(!lifecycle.providers.contains("lifecycle-refresh-kind"));
3716 assert!(provider_catalog_namespace("lifecycle-refresh-kind").is_none());
3717 lifecycle.remove_provider("lifecycle-refresh-kind").unwrap();
3718 });
3719 }
3720
3721 #[test]
3722 fn discovery_failure_observes_an_external_disable() {
3723 let _registry_lock = crate::model_registry::MODEL_CONFIG_LOCK
3724 .lock()
3725 .unwrap_or_else(std::sync::PoisonError::into_inner);
3726 test_runtime().block_on(async {
3727 let (_dir, lifecycle, provider) = fixture("lifecycle-refresh-disabled");
3728 remove_provider_catalog("lifecycle-refresh-disabled");
3729 lifecycle
3730 .install_provider(
3731 provider_record("lifecycle-refresh-disabled"),
3732 provider.clone(),
3733 )
3734 .await
3735 .unwrap();
3736 let hub = lifecycle.hub.clone();
3737 provider.before_discovery_once(move || {
3738 assert!(
3739 hub.set_auth_provider_enabled("lifecycle-refresh-disabled", false)
3740 .unwrap()
3741 );
3742 });
3743 provider.fail_once(ModelDiscoveryError::Transport("offline".into()));
3744
3745 assert!(matches!(
3746 lifecycle.refresh_models("lifecycle-refresh-disabled").await,
3747 Err(ProviderLifecycleError::Stale { id }) if id == "lifecycle-refresh-disabled"
3748 ));
3749 assert!(!lifecycle.providers.contains("lifecycle-refresh-disabled"));
3750 assert!(provider_catalog_namespace("lifecycle-refresh-disabled").is_none());
3751 lifecycle
3752 .remove_provider("lifecycle-refresh-disabled")
3753 .unwrap();
3754 });
3755 }
3756
3757 #[test]
3758 fn namespace_mismatch_during_refresh_fails_closed() {
3759 let _registry_lock = crate::model_registry::MODEL_CONFIG_LOCK
3760 .lock()
3761 .unwrap_or_else(std::sync::PoisonError::into_inner);
3762 test_runtime().block_on(async {
3763 let (_dir, lifecycle, provider) = fixture("lifecycle-refresh-namespace");
3764 remove_provider_catalog("lifecycle-refresh-namespace");
3765 lifecycle
3766 .install_provider(
3767 provider_record("lifecycle-refresh-namespace"),
3768 provider.clone(),
3769 )
3770 .await
3771 .unwrap();
3772 remove_provider_catalog("lifecycle-refresh-namespace");
3773 crate::model_registry::replace_provider_catalog(
3774 ProviderDescriptor {
3775 provider_key: "lifecycle-refresh-namespace".into(),
3776 provider_name: "Account".into(),
3777 namespace: "unexpected@account".into(),
3778 wire_profile: ReasoningWireProfile::CodexResponses,
3779 },
3780 &[model("replacement")],
3781 )
3782 .unwrap();
3783
3784 assert!(matches!(
3785 lifecycle
3786 .refresh_models("lifecycle-refresh-namespace")
3787 .await,
3788 Err(ProviderLifecycleError::Catalog(
3789 CatalogError::NamespaceChanged { .. }
3790 ))
3791 ));
3792 assert!(!lifecycle.providers.contains("lifecycle-refresh-namespace"));
3793 assert!(provider_catalog_namespace("lifecycle-refresh-namespace").is_none());
3794 lifecycle
3795 .remove_provider("lifecycle-refresh-namespace")
3796 .unwrap();
3797 });
3798 }
3799
3800 #[test]
3801 fn registry_key_conflict_during_refresh_preserves_last_good_catalog() {
3802 struct ConfigReset;
3803
3804 impl Drop for ConfigReset {
3805 fn drop(&mut self) {
3806 crate::model_registry::set_provider_config(
3807 crate::model_registry::ProviderConfig::default(),
3808 );
3809 }
3810 }
3811
3812 let _registry_lock = crate::model_registry::MODEL_CONFIG_LOCK
3813 .lock()
3814 .unwrap_or_else(std::sync::PoisonError::into_inner);
3815 let _config_reset = ConfigReset;
3816 test_runtime().block_on(async {
3817 let (_dir, lifecycle, provider) = fixture("lifecycle-refresh-registry-key");
3818 remove_provider_catalog("lifecycle-refresh-registry-key");
3819 lifecycle
3820 .hub
3821 .add_auth_provider_with_model_cache_details(
3822 provider_record("lifecycle-refresh-registry-key"),
3823 "stable",
3824 1,
3825 &[model("initial")],
3826 )
3827 .unwrap();
3828 lifecycle
3829 .restore_provider(
3830 "lifecycle-refresh-registry-key",
3831 ProviderKind::Codex,
3832 provider.clone(),
3833 )
3834 .await
3835 .unwrap();
3836 let mut config = crate::model_registry::ProviderConfig::default();
3837 config.providers.insert(
3838 "stable".into(),
3839 crate::model_registry::ProviderEntry {
3840 kind: "openai".into(),
3841 base_url: Some("https://api.openai.com/v1".into()),
3842 ..Default::default()
3843 },
3844 );
3845 crate::model_registry::set_provider_config(config);
3846 let before_revision = crate::model_registry::model_catalog_revision();
3847 provider.set_models(vec![model("gpt-4o")]);
3848
3849 assert!(matches!(
3850 lifecycle
3851 .refresh_models("lifecycle-refresh-registry-key")
3852 .await,
3853 Err(ProviderLifecycleError::Catalog(
3854 CatalogError::RegistryKeyInUse { .. }
3855 ))
3856 ));
3857 assert!(
3858 lifecycle
3859 .providers
3860 .contains("lifecycle-refresh-registry-key")
3861 );
3862 assert!(crate::model_registry::model_entry("stable:initial").is_some());
3863 assert_eq!(
3864 crate::model_registry::model_catalog_revision(),
3865 before_revision
3866 );
3867 lifecycle
3868 .remove_provider("lifecycle-refresh-registry-key")
3869 .unwrap();
3870 });
3871 }
3872
3873 #[test]
3874 fn credential_rotation_during_discovery_does_not_invalidate_catalog_refresh() {
3875 let _registry_lock = crate::model_registry::MODEL_CONFIG_LOCK
3876 .lock()
3877 .unwrap_or_else(std::sync::PoisonError::into_inner);
3878 test_runtime().block_on(async {
3879 let (_dir, lifecycle, provider) = fixture("lifecycle-token-refresh");
3880 remove_provider_catalog("lifecycle-token-refresh");
3881 lifecycle
3882 .install_provider(provider_record("lifecycle-token-refresh"), provider.clone())
3883 .await
3884 .unwrap();
3885 provider.set_models(vec![model("replacement")]);
3886 let hub = lifecycle.hub.clone();
3887 provider.before_discovery_once(move || {
3888 assert!(
3889 hub.update_auth_tokens(
3890 "lifecycle-token-refresh",
3891 crate::config_hub::AuthTokenUpdate {
3892 access_token: "rotated-access".into(),
3893 refresh_token: Some("rotated-refresh".into()),
3894 expires_at: 123,
3895 account: Some("rotated@example.com".into()),
3896 },
3897 )
3898 .unwrap()
3899 );
3900 });
3901
3902 let delta = lifecycle
3903 .refresh_models("lifecycle-token-refresh")
3904 .await
3905 .unwrap();
3906
3907 assert_eq!(delta.total, 1);
3908 let stored = lifecycle.hub.load_auth().unwrap().providers.remove(0);
3909 assert_eq!(stored.access_token, "rotated-access");
3910 assert_eq!(stored.refresh_token.as_deref(), Some("rotated-refresh"));
3911 assert_eq!(stored.model_cache.unwrap().models[0].slug, "replacement");
3912 lifecycle
3913 .remove_provider("lifecycle-token-refresh")
3914 .unwrap();
3915 });
3916 }
3917}