atman_runtime/
auth_store.rs

1use std::collections::{HashMap, VecDeque};
2
3use anyhow::Result;
4use serde::{Deserialize, Serialize};
5
6use crate::storage::config_dir;
7
8const AUTH_FILENAME: &str = "auth.json";
9const MODEL_CACHE_SCHEMA_VERSION: u32 = 1;
10pub(crate) const MODEL_CACHE_FRESHNESS_WINDOW_SECONDS: i64 = 15 * 60;
11
12#[derive(Debug, Clone, Copy, PartialEq, Eq)]
13pub(crate) enum ModelCacheFreshness {
14    Fresh,
15    Missing,
16    LegacySchema,
17    MissingCapabilities,
18    Expired,
19}
20
21#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
22#[serde(rename_all = "kebab-case")]
23pub enum ProviderKind {
24    Codex,
25    AnthropicOauth,
26    GitHubCopilot,
27    Custom,
28}
29
30#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
31pub struct ModelCache {
32    pub fetched_at: i64,
33    pub models: Vec<CachedModel>,
34}
35
36#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
37pub struct CachedModel {
38    pub slug: String,
39    #[serde(skip_serializing_if = "Option::is_none")]
40    pub context_budget: Option<u64>,
41    pub thinking: bool,
42}
43
44#[derive(Clone, Serialize, Deserialize)]
45pub struct StoredProvider {
46    pub id: String,
47    pub name: String,
48    pub kind: ProviderKind,
49    pub access_token: String,
50    #[serde(skip_serializing_if = "Option::is_none")]
51    pub refresh_token: Option<String>,
52    pub expires_at: i64,
53    #[serde(skip_serializing_if = "Option::is_none")]
54    pub account: Option<String>,
55    pub enabled: bool,
56    #[serde(skip_serializing_if = "Option::is_none")]
57    pub model_cache: Option<ModelCache>,
58}
59
60impl std::fmt::Debug for StoredProvider {
61    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
62        formatter
63            .debug_struct("StoredProvider")
64            .field("id", &self.id)
65            .field("name", &self.name)
66            .field("kind", &self.kind)
67            .field("access_token", &"[redacted]")
68            .field(
69                "refresh_token",
70                &self.refresh_token.as_ref().map(|_| "[redacted]"),
71            )
72            .field("expires_at", &self.expires_at)
73            .field("account", &self.account.as_ref().map(|_| "[redacted]"))
74            .field("enabled", &self.enabled)
75            .field("model_cache", &self.model_cache)
76            .finish()
77    }
78}
79
80#[derive(Debug, Clone, Default, Serialize, Deserialize)]
81pub struct AuthStore {
82    #[serde(default, skip_serializing_if = "Vec::is_empty")]
83    pub providers: Vec<StoredProvider>,
84}
85
86#[derive(Debug, Clone, Default, Serialize, Deserialize)]
87pub(crate) struct AuthStoreDocument {
88    #[serde(default, skip_serializing_if = "Vec::is_empty")]
89    providers: Vec<StoredProviderDocument>,
90}
91
92#[derive(Clone, Serialize, Deserialize)]
93struct StoredProviderDocument {
94    id: String,
95    name: String,
96    kind: ProviderKind,
97    access_token: String,
98    #[serde(skip_serializing_if = "Option::is_none")]
99    refresh_token: Option<String>,
100    expires_at: i64,
101    #[serde(skip_serializing_if = "Option::is_none")]
102    account: Option<String>,
103    enabled: bool,
104    #[serde(default, skip_serializing_if = "Option::is_none")]
105    credential_revision: Option<String>,
106    #[serde(default, skip_serializing_if = "Option::is_none")]
107    catalog_revision: Option<String>,
108    #[serde(default, skip_serializing_if = "Option::is_none")]
109    model_namespace: Option<String>,
110    #[serde(skip_serializing_if = "Option::is_none")]
111    model_cache: Option<ModelCacheDocument>,
112}
113
114impl std::fmt::Debug for StoredProviderDocument {
115    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
116        formatter
117            .debug_struct("StoredProviderDocument")
118            .field("id", &self.id)
119            .field("name", &self.name)
120            .field("kind", &self.kind)
121            .field("access_token", &"[redacted]")
122            .field(
123                "refresh_token",
124                &self.refresh_token.as_ref().map(|_| "[redacted]"),
125            )
126            .field("expires_at", &self.expires_at)
127            .field("account", &self.account.as_ref().map(|_| "[redacted]"))
128            .field("enabled", &self.enabled)
129            .field("credential_revision", &self.credential_revision)
130            .field("catalog_revision", &self.catalog_revision)
131            .field("model_namespace", &self.model_namespace)
132            .field("model_cache", &self.model_cache)
133            .finish()
134    }
135}
136
137#[derive(Debug, Clone, Serialize, Deserialize)]
138struct ModelCacheDocument {
139    #[serde(default, skip_serializing_if = "is_zero")]
140    schema_version: u32,
141    fetched_at: i64,
142    models: Vec<CachedModelDocument>,
143}
144
145#[derive(Debug, Clone, Serialize, Deserialize)]
146struct CachedModelDocument {
147    slug: String,
148    #[serde(skip_serializing_if = "Option::is_none")]
149    context_budget: Option<u64>,
150    thinking: bool,
151    #[serde(default, skip_serializing_if = "Option::is_none")]
152    capabilities: Option<crate::provider::ModelCapabilities>,
153}
154
155#[derive(Debug, Clone, PartialEq, Eq)]
156pub(crate) struct AuthProviderCatalogSnapshot {
157    revision: String,
158}
159
160#[derive(Debug, Clone, PartialEq, Eq)]
161pub(crate) struct AuthProviderCredentialSnapshot {
162    revision: String,
163}
164
165#[derive(Debug, Clone)]
166pub(crate) enum AuthCredentialCommit {
167    Updated { provider: StoredProvider },
168    Missing,
169    Changed,
170}
171
172fn is_zero(value: &u32) -> bool {
173    *value == 0
174}
175
176impl AuthStoreDocument {
177    pub(crate) fn legacy_view(&self) -> AuthStore {
178        AuthStore {
179            providers: self
180                .providers
181                .iter()
182                .map(StoredProviderDocument::legacy_view)
183                .collect(),
184        }
185    }
186
187    pub(crate) fn merge_legacy_view(&mut self, store: AuthStore) {
188        let mut previous: HashMap<String, VecDeque<StoredProviderDocument>> = HashMap::new();
189        for provider in self.providers.drain(..) {
190            previous
191                .entry(provider.id.clone())
192                .or_default()
193                .push_back(provider);
194        }
195        self.providers = store
196            .providers
197            .into_iter()
198            .map(
199                |provider| match previous.get_mut(&provider.id).and_then(VecDeque::pop_front) {
200                    Some(document) => document.merge(provider),
201                    None => provider.into(),
202                },
203            )
204            .collect();
205    }
206
207    pub(crate) fn update_model_cache_details(
208        &mut self,
209        provider_id: &str,
210        model_namespace: &str,
211        fetched_at: i64,
212        models: &[crate::provider::DiscoveredModelDetails],
213    ) -> Result<bool, String> {
214        let Some(provider) = self
215            .providers
216            .iter_mut()
217            .find(|provider| provider.id == provider_id)
218        else {
219            return Ok(false);
220        };
221        assign_model_namespace(provider, provider_id, model_namespace)?;
222        provider.model_cache = Some(ModelCacheDocument::from_details(fetched_at, models));
223        provider.bump_catalog_revision();
224        Ok(true)
225    }
226
227    pub(crate) fn ensure_model_namespace(
228        &mut self,
229        provider_id: &str,
230        model_namespace: &str,
231    ) -> Result<bool, String> {
232        let Some(provider) = self
233            .providers
234            .iter_mut()
235            .find(|provider| provider.id == provider_id)
236        else {
237            return Err(format!("auth provider `{provider_id}` does not exist"));
238        };
239        let changed = assign_model_namespace(provider, provider_id, model_namespace)?;
240        if changed {
241            provider.bump_catalog_revision();
242        }
243        Ok(changed)
244    }
245
246    pub(crate) fn set_provider_enabled(
247        &mut self,
248        provider_id: &str,
249        enabled: bool,
250    ) -> Option<bool> {
251        let provider = self
252            .providers
253            .iter_mut()
254            .find(|provider| provider.id == provider_id)?;
255        if provider.enabled == enabled {
256            return Some(false);
257        }
258        provider.enabled = enabled;
259        provider.bump_catalog_revision();
260        Some(true)
261    }
262
263    pub(crate) fn update_model_cache(&mut self, provider_id: &str, cache: ModelCache) -> bool {
264        let Some(provider) = self
265            .providers
266            .iter_mut()
267            .find(|provider| provider.id == provider_id)
268        else {
269            return false;
270        };
271        provider.model_cache = Some(cache.into());
272        provider.bump_catalog_revision();
273        true
274    }
275
276    pub(crate) fn model_namespace(&self, provider_id: &str) -> Option<String> {
277        self.providers
278            .iter()
279            .find(|provider| provider.id == provider_id)?
280            .model_namespace
281            .clone()
282    }
283
284    pub(crate) fn model_cache_details(
285        &self,
286        provider_id: &str,
287    ) -> Option<Vec<crate::provider::DiscoveredModelDetails>> {
288        self.providers
289            .iter()
290            .find(|provider| provider.id == provider_id)?
291            .model_cache
292            .as_ref()
293            .map(ModelCacheDocument::details)
294    }
295
296    pub(crate) fn model_cache_freshness(
297        &self,
298        provider_id: &str,
299        now: i64,
300        max_age_seconds: i64,
301    ) -> Option<ModelCacheFreshness> {
302        let provider = self
303            .providers
304            .iter()
305            .find(|provider| provider.id == provider_id)?;
306        Some(match provider.model_cache.as_ref() {
307            Some(cache) => cache.freshness(now, max_age_seconds),
308            None => ModelCacheFreshness::Missing,
309        })
310    }
311
312    pub(crate) fn provider_catalog_snapshot(
313        &self,
314        provider_id: &str,
315    ) -> Option<AuthProviderCatalogSnapshot> {
316        let provider = self
317            .providers
318            .iter()
319            .find(|provider| provider.id == provider_id)?;
320        Some(AuthProviderCatalogSnapshot {
321            revision: provider.catalog_revision.clone()?,
322        })
323    }
324
325    pub(crate) fn provider_credential_state(
326        &self,
327        provider_id: &str,
328    ) -> Option<(StoredProvider, Option<AuthProviderCredentialSnapshot>)> {
329        let provider = self
330            .providers
331            .iter()
332            .find(|provider| provider.id == provider_id)?;
333        Some((
334            provider.legacy_view(),
335            provider
336                .credential_revision
337                .clone()
338                .map(|revision| AuthProviderCredentialSnapshot { revision }),
339        ))
340    }
341
342    pub(crate) fn ensure_provider_catalog_state(
343        &mut self,
344        provider_id: &str,
345    ) -> Option<((StoredProvider, AuthProviderCatalogSnapshot), bool)> {
346        let provider = self
347            .providers
348            .iter_mut()
349            .find(|provider| provider.id == provider_id)?;
350        let changed = provider.catalog_revision.is_none();
351        let revision = provider
352            .catalog_revision
353            .get_or_insert_with(new_catalog_revision)
354            .clone();
355        Some((
356            (
357                provider.legacy_view(),
358                AuthProviderCatalogSnapshot { revision },
359            ),
360            changed,
361        ))
362    }
363
364    pub(crate) fn ensure_provider_credential_state(
365        &mut self,
366        provider_id: &str,
367    ) -> Option<((StoredProvider, AuthProviderCredentialSnapshot), bool)> {
368        let provider = self
369            .providers
370            .iter_mut()
371            .find(|provider| provider.id == provider_id)?;
372        let changed = provider.credential_revision.is_none();
373        let revision = provider
374            .credential_revision
375            .get_or_insert_with(new_credential_revision)
376            .clone();
377        Some((
378            (
379                provider.legacy_view(),
380                AuthProviderCredentialSnapshot { revision },
381            ),
382            changed,
383        ))
384    }
385
386    pub(crate) fn update_provider_credentials(
387        &mut self,
388        provider_id: &str,
389        expected: &AuthProviderCredentialSnapshot,
390        access_token: String,
391        refresh_token: Option<String>,
392        expires_at: i64,
393        account: Option<String>,
394    ) -> AuthCredentialCommit {
395        let Some(provider) = self
396            .providers
397            .iter_mut()
398            .find(|provider| provider.id == provider_id)
399        else {
400            return AuthCredentialCommit::Missing;
401        };
402        if provider.credential_revision.as_deref() != Some(expected.revision.as_str()) {
403            return AuthCredentialCommit::Changed;
404        }
405
406        provider.access_token = access_token;
407        provider.expires_at = expires_at;
408        if refresh_token.is_some() {
409            provider.refresh_token = refresh_token;
410        }
411        if account.is_some() {
412            provider.account = account;
413        }
414        provider.bump_credential_revision();
415        AuthCredentialCommit::Updated {
416            provider: provider.legacy_view(),
417        }
418    }
419}
420
421fn assign_model_namespace(
422    provider: &mut StoredProviderDocument,
423    provider_id: &str,
424    model_namespace: &str,
425) -> Result<bool, String> {
426    if let Some(existing) = provider.model_namespace.as_deref() {
427        if existing == model_namespace {
428            return Ok(false);
429        }
430        return Err(format!(
431            "provider `{provider_id}` model namespace is already `{existing}`"
432        ));
433    }
434    provider.model_namespace = Some(model_namespace.to_string());
435    Ok(true)
436}
437
438impl StoredProviderDocument {
439    fn legacy_view(&self) -> StoredProvider {
440        StoredProvider {
441            id: self.id.clone(),
442            name: self.name.clone(),
443            kind: self.kind.clone(),
444            access_token: self.access_token.clone(),
445            refresh_token: self.refresh_token.clone(),
446            expires_at: self.expires_at,
447            account: self.account.clone(),
448            enabled: self.enabled,
449            model_cache: self
450                .model_cache
451                .as_ref()
452                .map(ModelCacheDocument::legacy_view),
453        }
454    }
455
456    fn merge(mut self, provider: StoredProvider) -> Self {
457        let credential_changed = self.kind != provider.kind
458            || self.access_token != provider.access_token
459            || self.refresh_token != provider.refresh_token
460            || self.expires_at != provider.expires_at
461            || self.account != provider.account;
462        let catalog_changed = self.name != provider.name
463            || self.kind != provider.kind
464            || self.enabled != provider.enabled
465            || self
466                .model_cache
467                .as_ref()
468                .map(ModelCacheDocument::legacy_view)
469                != provider.model_cache;
470        let model_cache = match (provider.model_cache, self.model_cache.take()) {
471            (Some(cache), Some(document)) if document.legacy_view() == cache => Some(document),
472            (Some(cache), _) => Some(cache.into()),
473            (None, _) => None,
474        };
475        let catalog_revision = if catalog_changed {
476            Some(new_catalog_revision())
477        } else {
478            self.catalog_revision
479        };
480        let credential_revision = if credential_changed {
481            Some(new_credential_revision())
482        } else {
483            self.credential_revision
484        };
485        Self {
486            id: provider.id,
487            name: provider.name,
488            kind: provider.kind,
489            access_token: provider.access_token,
490            refresh_token: provider.refresh_token,
491            expires_at: provider.expires_at,
492            account: provider.account,
493            enabled: provider.enabled,
494            credential_revision,
495            catalog_revision,
496            model_namespace: self.model_namespace,
497            model_cache,
498        }
499    }
500
501    fn bump_catalog_revision(&mut self) {
502        self.catalog_revision = Some(new_catalog_revision());
503    }
504
505    fn bump_credential_revision(&mut self) {
506        let revision = new_credential_revision();
507        self.credential_revision = Some(revision);
508    }
509}
510
511impl From<StoredProvider> for StoredProviderDocument {
512    fn from(provider: StoredProvider) -> Self {
513        Self {
514            id: provider.id,
515            name: provider.name,
516            kind: provider.kind,
517            access_token: provider.access_token,
518            refresh_token: provider.refresh_token,
519            expires_at: provider.expires_at,
520            account: provider.account,
521            enabled: provider.enabled,
522            credential_revision: Some(new_credential_revision()),
523            catalog_revision: Some(new_catalog_revision()),
524            model_namespace: None,
525            model_cache: provider.model_cache.map(ModelCacheDocument::from),
526        }
527    }
528}
529
530fn new_catalog_revision() -> String {
531    uuid::Uuid::new_v4().simple().to_string()
532}
533
534fn new_credential_revision() -> String {
535    uuid::Uuid::new_v4().simple().to_string()
536}
537
538impl ModelCacheDocument {
539    fn legacy_view(&self) -> ModelCache {
540        ModelCache {
541            fetched_at: self.fetched_at,
542            models: self
543                .models
544                .iter()
545                .map(CachedModelDocument::legacy_view)
546                .collect(),
547        }
548    }
549
550    fn from_details(fetched_at: i64, models: &[crate::provider::DiscoveredModelDetails]) -> Self {
551        Self {
552            schema_version: MODEL_CACHE_SCHEMA_VERSION,
553            fetched_at,
554            models: models
555                .iter()
556                .map(|model| CachedModelDocument {
557                    slug: model.slug.clone(),
558                    context_budget: model.context_budget,
559                    thinking: model.capability_knowledge.thinking(),
560                    capabilities: model.capability_knowledge.advertised().cloned(),
561                })
562                .collect(),
563        }
564    }
565
566    fn details(&self) -> Vec<crate::provider::DiscoveredModelDetails> {
567        let capabilities_are_current = self.schema_version == MODEL_CACHE_SCHEMA_VERSION;
568        self.models
569            .iter()
570            .map(|model| crate::provider::DiscoveredModelDetails {
571                slug: model.slug.clone(),
572                context_budget: model.context_budget,
573                capability_knowledge: if capabilities_are_current {
574                    model
575                        .capabilities
576                        .clone()
577                        .map(crate::provider::CapabilityKnowledge::Advertised)
578                        .unwrap_or(crate::provider::CapabilityKnowledge::Legacy {
579                            thinking: model.thinking,
580                        })
581                } else {
582                    crate::provider::CapabilityKnowledge::Legacy {
583                        thinking: model.thinking,
584                    }
585                },
586            })
587            .collect()
588    }
589
590    fn freshness(&self, now: i64, max_age_seconds: i64) -> ModelCacheFreshness {
591        if self.schema_version != MODEL_CACHE_SCHEMA_VERSION {
592            return ModelCacheFreshness::LegacySchema;
593        }
594        if self.models.iter().any(|model| model.capabilities.is_none()) {
595            return ModelCacheFreshness::MissingCapabilities;
596        }
597        if now < self.fetched_at || now.saturating_sub(self.fetched_at) >= max_age_seconds {
598            return ModelCacheFreshness::Expired;
599        }
600        ModelCacheFreshness::Fresh
601    }
602}
603
604impl From<ModelCache> for ModelCacheDocument {
605    fn from(cache: ModelCache) -> Self {
606        Self {
607            schema_version: 0,
608            fetched_at: cache.fetched_at,
609            models: cache
610                .models
611                .into_iter()
612                .map(CachedModelDocument::from)
613                .collect(),
614        }
615    }
616}
617
618impl CachedModelDocument {
619    fn legacy_view(&self) -> CachedModel {
620        CachedModel {
621            slug: self.slug.clone(),
622            context_budget: self.context_budget,
623            thinking: self.thinking,
624        }
625    }
626}
627
628impl From<CachedModel> for CachedModelDocument {
629    fn from(model: CachedModel) -> Self {
630        Self {
631            slug: model.slug,
632            context_budget: model.context_budget,
633            thinking: model.thinking,
634            capabilities: None,
635        }
636    }
637}
638
639impl AuthStore {
640    pub fn load() -> Result<Self> {
641        Ok(crate::config_hub::ConfigHub::global()?.load_auth()?)
642    }
643
644    pub fn save_to(&self, path: &std::path::Path) -> Result<()> {
645        let hub = crate::config_hub::ConfigHub::from_auth_path(path);
646        hub.update_auth(|store| {
647            *store = self.clone();
648            Ok(())
649        })?;
650        Ok(())
651    }
652
653    pub fn save(&self) -> Result<()> {
654        let dir = config_dir()?;
655        self.save_to(&dir.join(AUTH_FILENAME))
656    }
657
658    pub fn add(&mut self, p: StoredProvider) {
659        self.providers.push(p);
660    }
661
662    pub fn remove(&mut self, id: &str) -> bool {
663        let len_before = self.providers.len();
664        self.providers.retain(|p| p.id != id);
665        self.providers.len() < len_before
666    }
667
668    /// Update the model cache for a provider by ID. Returns false if provider not found.
669    pub fn update_model_cache(&mut self, provider_id: &str, cache: ModelCache) -> bool {
670        if let Some(p) = self.providers.iter_mut().find(|p| p.id == provider_id) {
671            p.model_cache = Some(cache);
672            true
673        } else {
674            false
675        }
676    }
677}
678
679/// Save discovered models as cache for a provider. Reads auth.json, updates, writes back.
680pub fn save_provider_model_cache(
681    provider_id: &str,
682    models: &[crate::provider::DiscoveredModel],
683) -> Result<()> {
684    let cache = ModelCache {
685        fetched_at: chrono::Utc::now().timestamp(),
686        models: models
687            .iter()
688            .map(|model| CachedModel {
689                slug: model.slug.clone(),
690                context_budget: model.context_budget,
691                thinking: model.thinking,
692            })
693            .collect(),
694    };
695    let _ = crate::config_hub::ConfigHub::global()?.update_auth_model_cache(provider_id, cache)?;
696    Ok(())
697}
698
699/// Save discovered capability metadata in the versioned auth cache.
700pub fn save_provider_model_cache_details(
701    provider_id: &str,
702    model_namespace: &str,
703    models: &[crate::provider::DiscoveredModelDetails],
704) -> Result<()> {
705    let updated = crate::config_hub::ConfigHub::global()?.update_auth_model_cache_details(
706        provider_id,
707        model_namespace,
708        chrono::Utc::now().timestamp(),
709        models,
710    )?;
711    if !updated {
712        anyhow::bail!("auth provider `{provider_id}` does not exist");
713    }
714    Ok(())
715}
716
717/// Load the stable model namespace assigned to one provider.
718pub fn load_provider_model_namespace(provider_id: &str) -> Result<Option<String>> {
719    Ok(crate::config_hub::ConfigHub::global()?.load_auth_model_namespace(provider_id)?)
720}
721
722/// Persist a provider namespace without changing its model cache freshness.
723pub fn ensure_provider_model_namespace(provider_id: &str, model_namespace: &str) -> Result<()> {
724    crate::config_hub::ConfigHub::global()?
725        .ensure_auth_model_namespace(provider_id, model_namespace)?;
726    Ok(())
727}
728
729/// Load cached models with capability provenance.
730pub fn load_provider_model_cache_details(
731    provider_id: &str,
732) -> Result<Option<Vec<crate::provider::DiscoveredModelDetails>>> {
733    Ok(crate::config_hub::ConfigHub::global()?.load_auth_model_cache_details(provider_id)?)
734}
735
736/// Convert cached models to discovered models for registry hydration.
737pub fn cached_to_discovered(cache: &ModelCache) -> Vec<crate::provider::DiscoveredModel> {
738    cache
739        .models
740        .iter()
741        .map(|m| crate::provider::DiscoveredModel {
742            slug: m.slug.clone(),
743            context_budget: m.context_budget,
744            thinking: m.thinking,
745        })
746        .collect()
747}
748
749/// Adapt the public cache format with legacy capability provenance.
750pub fn cached_to_discovered_details(
751    cache: &ModelCache,
752) -> Vec<crate::provider::DiscoveredModelDetails> {
753    cached_to_discovered(cache)
754        .into_iter()
755        .map(crate::provider::DiscoveredModelDetails::from)
756        .collect()
757}
758
759#[cfg(test)]
760mod tests {
761    use super::*;
762    use tempfile::TempDir;
763
764    fn provider(id: &str) -> StoredProvider {
765        StoredProvider {
766            id: id.into(),
767            name: "OAuth account".into(),
768            kind: ProviderKind::Codex,
769            access_token: "old-access".into(),
770            refresh_token: Some("old-refresh".into()),
771            expires_at: 1,
772            account: Some("account@example.test".into()),
773            enabled: true,
774            model_cache: None,
775        }
776    }
777
778    #[test]
779    fn load_returns_empty_when_file_missing() {
780        let tmp = TempDir::new().unwrap();
781        let path = tmp.path().join("auth.json");
782        let store: AuthStore = std::fs::read(&path)
783            .ok()
784            .and_then(|b| serde_json::from_slice(&b).ok())
785            .unwrap_or_default();
786        assert!(store.providers.is_empty());
787    }
788
789    #[test]
790    fn save_then_load_round_trips() {
791        let tmp = TempDir::new().unwrap();
792        let path = tmp.path().join("auth.json");
793        let mut store = AuthStore::default();
794        store.add(StoredProvider {
795            id: "test-1".into(),
796            name: "Personal Codex".into(),
797            kind: ProviderKind::Codex,
798            access_token: "tok1".into(),
799            refresh_token: Some("rt1".into()),
800            expires_at: 1761735358,
801            account: Some("x@example.com".into()),
802            enabled: true,
803            model_cache: None,
804        });
805        store.save_to(&path).unwrap();
806
807        let bytes = std::fs::read(&path).unwrap();
808        let loaded: AuthStore = serde_json::from_slice(&bytes).unwrap();
809        assert_eq!(loaded.providers.len(), 1);
810        assert_eq!(loaded.providers[0].name, "Personal Codex");
811    }
812
813    #[test]
814    fn stored_provider_debug_redacts_credentials() {
815        let provider = provider("debug");
816        let debug = format!("{provider:?}");
817
818        assert!(debug.contains("[redacted]"));
819        assert!(!debug.contains("old-access"));
820        assert!(!debug.contains("old-refresh"));
821        assert!(!debug.contains("account@example.test"));
822    }
823
824    #[test]
825    fn auth_store_document_debug_redacts_credentials() {
826        let document = AuthStoreDocument {
827            providers: vec![provider("debug-document").into()],
828        };
829        let debug = format!("{document:?}");
830
831        assert_eq!(debug.matches("[redacted]").count(), 3);
832        assert!(!debug.contains("old-access"));
833        assert!(!debug.contains("old-refresh"));
834        assert!(!debug.contains("account@example.test"));
835    }
836
837    #[test]
838    fn remove_existing_id_returns_true() {
839        let mut store = AuthStore::default();
840        store.add(StoredProvider {
841            id: "keep".into(),
842            name: "A".into(),
843            kind: ProviderKind::Custom,
844            access_token: "t".into(),
845            refresh_token: None,
846            expires_at: 0,
847            account: None,
848            enabled: true,
849            model_cache: None,
850        });
851        store.add(StoredProvider {
852            id: "del".into(),
853            name: "B".into(),
854            kind: ProviderKind::Custom,
855            access_token: "t".into(),
856            refresh_token: None,
857            expires_at: 0,
858            account: None,
859            enabled: true,
860            model_cache: None,
861        });
862        assert!(store.remove("del"));
863        assert_eq!(store.providers.len(), 1);
864        assert_eq!(store.providers[0].id, "keep");
865    }
866
867    #[test]
868    fn remove_missing_id_returns_false() {
869        let mut store = AuthStore::default();
870        assert!(!store.remove("nope"));
871    }
872
873    #[test]
874    fn provider_kind_serde_round_trip() {
875        let kinds = vec![
876            ProviderKind::Codex,
877            ProviderKind::AnthropicOauth,
878            ProviderKind::GitHubCopilot,
879            ProviderKind::Custom,
880        ];
881        for k in kinds {
882            let json = serde_json::to_string(&k).unwrap();
883            let back: ProviderKind = serde_json::from_str(&json).unwrap();
884            assert_eq!(back, k);
885        }
886    }
887
888    #[test]
889    fn auth_store_serde_backward_compat_empty_json() {
890        let store: AuthStore = serde_json::from_str("{}").unwrap();
891        assert!(store.providers.is_empty());
892    }
893
894    #[test]
895    fn auth_store_serde_backward_compat_null_providers() {
896        let json = r#"{"providers": []}"#;
897        let store: AuthStore = serde_json::from_str(json).unwrap();
898        assert!(store.providers.is_empty());
899    }
900
901    #[test]
902    fn public_model_cache_retains_the_v1_9_1_shape() {
903        let cache = ModelCache {
904            fetched_at: 10,
905            models: vec![CachedModel {
906                slug: "codex/test".into(),
907                context_budget: Some(8_192),
908                thinking: true,
909            }],
910        };
911
912        let value = serde_json::to_value(cache).unwrap();
913        assert!(value.get("schema_version").is_none());
914        assert!(value["models"][0].get("capabilities").is_none());
915    }
916
917    #[test]
918    fn legacy_model_cache_document_preserves_unknown_capabilities() {
919        let cache: ModelCacheDocument = serde_json::from_str(
920            r#"{"fetched_at":10,"models":[{"slug":"codex/test","thinking":true}]}"#,
921        )
922        .unwrap();
923
924        assert_eq!(cache.schema_version, 0);
925        assert_eq!(cache.models[0].capabilities, None);
926        assert_eq!(
927            cache.details()[0].capability_knowledge,
928            crate::provider::CapabilityKnowledge::Legacy { thinking: true }
929        );
930    }
931
932    #[test]
933    fn unversioned_cache_does_not_trust_a_serialized_empty_capability_object() {
934        let cache: ModelCacheDocument = serde_json::from_str(
935            r#"{
936                "fetched_at": 10,
937                "models": [{
938                    "slug": "codex/test",
939                    "thinking": true,
940                    "capabilities": {
941                        "reasoning_efforts": [],
942                        "reasoning_modes": [],
943                        "input_modalities": []
944                    }
945                }]
946            }"#,
947        )
948        .unwrap();
949
950        assert_eq!(cache.schema_version, 0);
951        assert!(cache.models[0].capabilities.is_some());
952        assert_eq!(
953            cache.details()[0].capability_knowledge,
954            crate::provider::CapabilityKnowledge::Legacy { thinking: true }
955        );
956    }
957
958    #[test]
959    fn versioned_model_cache_distinguishes_explicit_empty_capabilities() {
960        let cache = ModelCacheDocument::from_details(
961            10,
962            &[crate::provider::DiscoveredModelDetails {
963                slug: "codex/test".into(),
964                context_budget: Some(8_192),
965                capability_knowledge: crate::provider::CapabilityKnowledge::Advertised(
966                    crate::provider::ModelCapabilities::default(),
967                ),
968            }],
969        );
970        let json = serde_json::to_string(&cache).unwrap();
971        let decoded: ModelCacheDocument = serde_json::from_str(&json).unwrap();
972
973        assert_eq!(decoded.schema_version, MODEL_CACHE_SCHEMA_VERSION);
974        assert_eq!(
975            decoded.models[0].capabilities,
976            Some(crate::provider::ModelCapabilities::default())
977        );
978        assert_eq!(
979            decoded.details()[0].capability_knowledge,
980            crate::provider::CapabilityKnowledge::Advertised(
981                crate::provider::ModelCapabilities::default()
982            )
983        );
984    }
985
986    #[test]
987    fn model_cache_freshness_requires_current_complete_recent_metadata() {
988        let explicit_empty = crate::provider::DiscoveredModelDetails {
989            slug: "codex/test".into(),
990            context_budget: Some(8_192),
991            capability_knowledge: crate::provider::CapabilityKnowledge::Advertised(
992                crate::provider::ModelCapabilities::default(),
993            ),
994        };
995        let fresh = ModelCacheDocument::from_details(100, &[explicit_empty]);
996        let empty = ModelCacheDocument::from_details(100, &[]);
997        let legacy: ModelCacheDocument = serde_json::from_str(
998            r#"{"fetched_at":100,"models":[{"slug":"codex/test","thinking":true}]}"#,
999        )
1000        .unwrap();
1001        let incomplete: ModelCacheDocument = serde_json::from_str(
1002            r#"{
1003                "schema_version": 1,
1004                "fetched_at": 100,
1005                "models": [{"slug":"codex/test","thinking":true}]
1006            }"#,
1007        )
1008        .unwrap();
1009        let missing = AuthStoreDocument {
1010            providers: vec![provider("missing").into()],
1011        };
1012
1013        assert_eq!(
1014            missing.model_cache_freshness("missing", 100, 900),
1015            Some(ModelCacheFreshness::Missing)
1016        );
1017        assert_eq!(fresh.freshness(999, 900), ModelCacheFreshness::Fresh);
1018        assert_eq!(empty.freshness(999, 900), ModelCacheFreshness::Fresh);
1019        assert_eq!(fresh.freshness(1_000, 900), ModelCacheFreshness::Expired);
1020        assert_eq!(fresh.freshness(99, 900), ModelCacheFreshness::Expired);
1021        assert_eq!(
1022            legacy.freshness(1_000, 900),
1023            ModelCacheFreshness::LegacySchema
1024        );
1025        assert_eq!(
1026            incomplete.freshness(1_000, 900),
1027            ModelCacheFreshness::MissingCapabilities
1028        );
1029    }
1030
1031    #[test]
1032    fn versioned_wire_cache_is_readable_as_the_public_dto() {
1033        let wire = ModelCacheDocument::from_details(
1034            10,
1035            &[crate::provider::DiscoveredModelDetails {
1036                slug: "codex/test".into(),
1037                context_budget: Some(8_192),
1038                capability_knowledge: crate::provider::CapabilityKnowledge::Advertised(
1039                    crate::provider::ModelCapabilities::default(),
1040                ),
1041            }],
1042        );
1043
1044        let json = serde_json::to_string(&wire).unwrap();
1045        let legacy: ModelCache = serde_json::from_str(&json).unwrap();
1046
1047        assert_eq!(legacy.fetched_at, 10);
1048        assert_eq!(legacy.models[0].slug, "codex/test");
1049        assert!(!legacy.models[0].thinking);
1050    }
1051
1052    #[test]
1053    fn legacy_view_updates_preserve_namespace_and_capability_metadata() {
1054        let mut document = AuthStoreDocument::default();
1055        document.merge_legacy_view(AuthStore {
1056            providers: vec![StoredProvider {
1057                id: "provider".into(),
1058                name: "Provider".into(),
1059                kind: ProviderKind::Codex,
1060                access_token: "old".into(),
1061                refresh_token: None,
1062                expires_at: 1,
1063                account: None,
1064                enabled: true,
1065                model_cache: None,
1066            }],
1067        });
1068        assert!(
1069            document
1070                .update_model_cache_details(
1071                    "provider",
1072                    "stable-provider",
1073                    10,
1074                    &[crate::provider::DiscoveredModelDetails {
1075                        slug: "codex/test".into(),
1076                        context_budget: Some(8_192),
1077                        capability_knowledge: crate::provider::CapabilityKnowledge::Advertised(
1078                            crate::provider::ModelCapabilities::default(),
1079                        ),
1080                    }],
1081                )
1082                .unwrap()
1083        );
1084
1085        let mut legacy = document.legacy_view();
1086        legacy.providers[0].access_token = "new".into();
1087        document.merge_legacy_view(legacy);
1088
1089        assert_eq!(document.legacy_view().providers[0].access_token, "new");
1090        assert_eq!(
1091            document.model_namespace("provider").as_deref(),
1092            Some("stable-provider")
1093        );
1094        assert!(matches!(
1095            document.model_cache_details("provider").unwrap()[0].capability_knowledge,
1096            crate::provider::CapabilityKnowledge::Advertised(_)
1097        ));
1098        assert!(
1099            document
1100                .update_model_cache_details("provider", "changed", 11, &[])
1101                .unwrap_err()
1102                .contains("already")
1103        );
1104        assert_eq!(
1105            document.model_namespace("provider").as_deref(),
1106            Some("stable-provider")
1107        );
1108    }
1109
1110    #[test]
1111    fn provider_rename_preserves_credential_cas_and_invalidates_catalog() {
1112        let mut document = AuthStoreDocument::default();
1113        document.merge_legacy_view(AuthStore {
1114            providers: vec![provider("provider")],
1115        });
1116        let credential = document
1117            .provider_credential_state("provider")
1118            .unwrap()
1119            .1
1120            .unwrap();
1121        let catalog = document.provider_catalog_snapshot("provider").unwrap();
1122
1123        let mut legacy = document.legacy_view();
1124        legacy.providers[0].name = "Renamed OAuth account".into();
1125        document.merge_legacy_view(legacy);
1126
1127        assert_eq!(
1128            document
1129                .provider_credential_state("provider")
1130                .unwrap()
1131                .1
1132                .unwrap(),
1133            credential
1134        );
1135        assert_ne!(
1136            document.provider_catalog_snapshot("provider").unwrap(),
1137            catalog
1138        );
1139        assert!(matches!(
1140            document.update_provider_credentials(
1141                "provider",
1142                &credential,
1143                "fresh-access".into(),
1144                Some("fresh-refresh".into()),
1145                2,
1146                Some("fresh@example.test".into()),
1147            ),
1148            AuthCredentialCommit::Updated { .. }
1149        ));
1150    }
1151
1152    #[test]
1153    fn legacy_duplicate_ids_do_not_cross_wire_capability_metadata() {
1154        let provider = |name: &str| StoredProvider {
1155            id: "duplicate".into(),
1156            name: name.into(),
1157            kind: ProviderKind::Codex,
1158            access_token: "old".into(),
1159            refresh_token: None,
1160            expires_at: 1,
1161            account: None,
1162            enabled: true,
1163            model_cache: None,
1164        };
1165        let details = |capabilities: crate::provider::ModelCapabilities| {
1166            ModelCacheDocument::from_details(
1167                10,
1168                &[crate::provider::DiscoveredModelDetails {
1169                    slug: "api/same".into(),
1170                    context_budget: Some(8_192),
1171                    capability_knowledge: crate::provider::CapabilityKnowledge::Advertised(
1172                        capabilities,
1173                    ),
1174                }],
1175            )
1176        };
1177        let mut first: StoredProviderDocument = provider("first").into();
1178        first.model_cache = Some(details(crate::provider::ModelCapabilities::default()));
1179        let mut second: StoredProviderDocument = provider("second").into();
1180        second.model_cache = Some(details(crate::provider::ModelCapabilities {
1181            input_modalities: vec![crate::provider::InputModality::Image],
1182            ..Default::default()
1183        }));
1184        let mut document = AuthStoreDocument {
1185            providers: vec![first, second],
1186        };
1187
1188        let mut legacy = document.legacy_view();
1189        legacy.providers[0].access_token = "new-first".into();
1190        legacy.providers[1].access_token = "new-second".into();
1191        document.merge_legacy_view(legacy);
1192
1193        assert_eq!(document.providers[0].access_token, "new-first");
1194        assert_eq!(document.providers[1].access_token, "new-second");
1195        assert_eq!(
1196            document.providers[0].model_cache.as_ref().unwrap().models[0].capabilities,
1197            Some(crate::provider::ModelCapabilities::default())
1198        );
1199        assert_eq!(
1200            document.providers[1].model_cache.as_ref().unwrap().models[0]
1201                .capabilities
1202                .as_ref()
1203                .unwrap()
1204                .input_modalities,
1205            [crate::provider::InputModality::Image]
1206        );
1207    }
1208}
atman_runtime/auth_store.rs

atman_runtime/
auth_store.rs
