imp_llm/
auth.rs

1use std::collections::HashMap;
2use std::path::PathBuf;
3use std::sync::Arc;
4
5use crate::truncate_chars_with_suffix;
6use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine};
7use serde::{Deserialize, Serialize};
8use serde_json::Value;
9
10use crate::error::Result;
11
12pub type ApiKey = String;
13const KEYRING_SERVICE: &str = "imp";
14const LEGACY_KEYRING_SERVICES: &[&str] = &["imp-cli", "impeccable", "mana"];
15
16fn provider_lookup_candidates(provider: &str) -> Vec<String> {
17    let mut candidates = vec![provider.to_string()];
18    let lower = provider.to_lowercase();
19    if lower != provider {
20        candidates.push(lower);
21    }
22    if provider == "render" {
23        candidates.push("Render".to_string());
24    }
25    dedupe_strings(candidates)
26}
27
28fn field_lookup_candidates(field: &str) -> Vec<String> {
29    let mut candidates = vec![field.to_string()];
30    if field == "secrets_key" {
31        candidates.push("secret_key".to_string());
32    }
33    if field == "secret_key" {
34        candidates.push("secrets_key".to_string());
35    }
36    dedupe_strings(candidates)
37}
38
39fn dedupe_strings(values: Vec<String>) -> Vec<String> {
40    let mut deduped = Vec::new();
41    for value in values {
42        if !deduped.contains(&value) {
43            deduped.push(value);
44        }
45    }
46    deduped
47}
48
49trait SecretBackend: Send + Sync {
50    fn get(&self, provider: &str, field: &str) -> Result<Option<String>>;
51    fn set(&self, provider: &str, field: &str, value: &str) -> Result<()>;
52    fn delete(&self, provider: &str, field: &str) -> Result<()>;
53}
54
55struct KeyringBackend;
56
57impl KeyringBackend {
58    fn entry(service: &str, provider: &str, field: &str) -> Result<keyring::Entry> {
59        keyring::Entry::new(service, &format!("{provider}:{field}"))
60            .map_err(|e| crate::error::Error::Auth(format!("Secure storage init failed: {e}")))
61    }
62
63    fn read_entry(service: &str, provider: &str, field: &str) -> Result<Option<String>> {
64        let entry = Self::entry(service, provider, field)?;
65        match entry.get_password() {
66            Ok(value) => Ok(Some(value)),
67            Err(keyring::Error::NoEntry) => Ok(None),
68            Err(error) => Err(Self::map_error("read", provider, field, error)),
69        }
70    }
71
72    fn lookup_secret(provider: &str, field: &str) -> Result<Option<String>> {
73        let providers = provider_lookup_candidates(provider);
74        let fields = field_lookup_candidates(field);
75        for candidate_provider in &providers {
76            for candidate_field in &fields {
77                if let Some(value) =
78                    Self::read_entry(KEYRING_SERVICE, candidate_provider, candidate_field)?
79                {
80                    return Ok(Some(value));
81                }
82            }
83        }
84        for service in LEGACY_KEYRING_SERVICES {
85            for candidate_provider in &providers {
86                for candidate_field in &fields {
87                    if let Some(value) =
88                        Self::read_entry(service, candidate_provider, candidate_field)?
89                    {
90                        return Ok(Some(value));
91                    }
92                }
93            }
94        }
95        Ok(None)
96    }
97
98    fn map_error(
99        action: &str,
100        provider: &str,
101        field: &str,
102        error: keyring::Error,
103    ) -> crate::error::Error {
104        crate::error::Error::Auth(format!(
105            "Secure storage {action} failed for {provider}.{field}: {error}"
106        ))
107    }
108}
109
110impl SecretBackend for KeyringBackend {
111    fn get(&self, provider: &str, field: &str) -> Result<Option<String>> {
112        Self::lookup_secret(provider, field)
113    }
114
115    fn set(&self, provider: &str, field: &str, value: &str) -> Result<()> {
116        let entry = Self::entry(KEYRING_SERVICE, provider, field)?;
117        entry
118            .set_password(value)
119            .map_err(|error| Self::map_error("write", provider, field, error))?;
120
121        match Self::read_entry(KEYRING_SERVICE, provider, field)? {
122            Some(stored) if stored == value => Ok(()),
123            Some(_) => Err(crate::error::Error::Auth(format!(
124                "Secure storage write verification failed for {provider}.{field}: readback did not match"
125            ))),
126            None => Err(crate::error::Error::Auth(format!(
127                "Secure storage write verification failed for {provider}.{field}: value was not readable after write"
128            ))),
129        }
130    }
131
132    fn delete(&self, provider: &str, field: &str) -> Result<()> {
133        let providers = provider_lookup_candidates(provider);
134        let fields = field_lookup_candidates(field);
135        let mut first_error = None;
136        for service in
137            std::iter::once(KEYRING_SERVICE).chain(LEGACY_KEYRING_SERVICES.iter().copied())
138        {
139            for candidate_provider in &providers {
140                for candidate_field in &fields {
141                    let entry = Self::entry(service, candidate_provider, candidate_field)?;
142                    match entry.delete_credential() {
143                        Ok(()) | Err(keyring::Error::NoEntry) => {}
144                        Err(error) if first_error.is_none() => {
145                            first_error = Some(Self::map_error(
146                                "delete",
147                                candidate_provider,
148                                candidate_field,
149                                error,
150                            ));
151                        }
152                        Err(_) => {}
153                    }
154                }
155            }
156        }
157        match first_error {
158            Some(error) => Err(error),
159            None => Ok(()),
160        }
161    }
162}
163
164#[derive(Debug, Clone, Serialize, Deserialize)]
165pub struct OAuthCredential {
166    pub access_token: String,
167    pub refresh_token: String,
168    pub expires_at: u64,
169}
170
171impl OAuthCredential {
172    /// Check whether this token has expired (or will within the next minute).
173    pub fn is_expired(&self) -> bool {
174        crate::now() >= self.expires_at
175    }
176}
177
178#[derive(Debug, Clone, Serialize, Deserialize)]
179#[serde(tag = "type")]
180pub enum StoredCredential {
181    ApiKey { key: String },
182    OAuth(OAuthCredential),
183    SecretFields { fields: Vec<String> },
184}
185
186#[derive(Debug, Clone, PartialEq, Eq)]
187pub enum SecretFieldStatus {
188    Present,
189    Missing,
190    Error(String),
191}
192
193impl SecretFieldStatus {
194    #[must_use]
195    pub fn is_present(&self) -> bool {
196        matches!(self, Self::Present)
197    }
198}
199
200#[derive(Debug, Clone, PartialEq, Eq)]
201pub struct SecretStatus {
202    pub provider: String,
203    pub fields: Vec<(String, SecretFieldStatus)>,
204}
205
206impl SecretStatus {
207    #[must_use]
208    pub fn is_usable(&self) -> bool {
209        self.fields.iter().all(|(_, status)| status.is_present())
210    }
211}
212
213#[derive(Debug, Clone, PartialEq, Eq)]
214pub struct OAuthDisplayInfo {
215    pub account_id: Option<String>,
216    pub plan: Option<String>,
217    pub using_subscription: bool,
218}
219
220impl OAuthDisplayInfo {
221    pub fn login_message(&self, provider: &str) -> String {
222        match provider {
223            "openai" | "openai-codex" => {
224                let mut message = String::from("Logged in to OpenAI / ChatGPT");
225                if let Some(account_id) = &self.account_id {
226                    message.push_str(&format!(" as account {account_id}"));
227                }
228                if let Some(plan) = &self.plan {
229                    message.push_str(&format!(", plan: {plan}"));
230                }
231                message.push('.');
232                message
233            }
234            "anthropic" => {
235                if let Some(plan) = &self.plan {
236                    format!("Logged in to Anthropic with {plan} subscription credentials.")
237                } else {
238                    "Logged in to Anthropic with OAuth subscription credentials.".into()
239                }
240            }
241            _ => format!("Logged in to {provider} with OAuth credentials."),
242        }
243    }
244
245    pub fn status_summary(&self) -> String {
246        match (&self.plan, self.short_account_id()) {
247            (Some(plan), Some(account_id)) => format!("{plan} · {account_id}"),
248            (Some(plan), None) => plan.clone(),
249            (None, Some(account_id)) => account_id,
250            (None, None) if self.using_subscription => "subscription".into(),
251            (None, None) => "oauth".into(),
252        }
253    }
254
255    pub fn short_account_id(&self) -> Option<String> {
256        self.account_id
257            .as_ref()
258            .map(|account_id| truncate_chars_with_suffix(account_id, 8, "…"))
259    }
260}
261
262/// Manages API keys and OAuth credentials.
263pub struct AuthStore {
264    runtime_keys: HashMap<String, String>,
265    pub stored: HashMap<String, StoredCredential>,
266    path: PathBuf,
267    backend: Arc<dyn SecretBackend>,
268}
269
270impl AuthStore {
271    pub fn new(path: PathBuf) -> Self {
272        Self::new_with_backend(path, Arc::new(KeyringBackend))
273    }
274
275    fn new_with_backend(path: PathBuf, backend: Arc<dyn SecretBackend>) -> Self {
276        Self {
277            runtime_keys: HashMap::new(),
278            stored: HashMap::new(),
279            path,
280            backend,
281        }
282    }
283
284    /// Load stored credentials from disk.
285    pub fn load(path: &std::path::Path) -> Result<Self> {
286        Self::load_with_backend(path, Arc::new(KeyringBackend))
287    }
288
289    fn load_with_backend(path: &std::path::Path, backend: Arc<dyn SecretBackend>) -> Result<Self> {
290        let stored = if path.exists() {
291            let data = std::fs::read_to_string(path)?;
292            serde_json::from_str(&data).map_err(|error| {
293                crate::error::Error::Auth(format!(
294                    "Failed to parse auth metadata at {}: {error}",
295                    path.display()
296                ))
297            })?
298        } else {
299            HashMap::new()
300        };
301        Ok(Self {
302            runtime_keys: HashMap::new(),
303            stored,
304            path: path.to_path_buf(),
305            backend,
306        })
307    }
308
309    /// Set a runtime override (not persisted).
310    /// Empty or whitespace-only values are treated as absent.
311    pub fn set_runtime_key(&mut self, provider: &str, key: String) {
312        let trimmed = key.trim();
313        if trimmed.is_empty() {
314            self.runtime_keys.remove(provider);
315            return;
316        }
317        self.runtime_keys
318            .insert(provider.to_string(), trimmed.to_string());
319    }
320
321    /// Check whether credentials are usable for a provider without producing an error.
322    /// Returns true if a runtime key, readable stored credential, or env var is available.
323    pub fn has_credentials(&self, provider: &str) -> bool {
324        self.resolve(provider).is_ok()
325    }
326
327    /// Resolution order: runtime override -> stored -> env var -> error.
328    pub fn resolve(&self, provider: &str) -> Result<ApiKey> {
329        if let Some(key) = self.runtime_keys.get(provider) {
330            return Ok(key.clone());
331        }
332
333        if let Some(StoredCredential::OAuth(oauth)) = self.stored.get(provider) {
334            return Ok(oauth.access_token.clone());
335        }
336
337        self.resolve_secret_field(provider, "api_key")
338    }
339
340    /// Resolve an API key without falling back to stored OAuth credentials.
341    pub fn resolve_api_key_only(&self, provider: &str) -> Result<ApiKey> {
342        self.resolve_secret_field(provider, "api_key")
343    }
344
345    /// Resolve a named secret field for any stored provider/service.
346    pub fn resolve_secret_field(&self, provider: &str, field: &str) -> Result<String> {
347        if field == "api_key" {
348            if let Some(key) = self.runtime_keys.get(provider) {
349                return Ok(key.clone());
350            }
351        }
352
353        if let Some((stored_provider, credential)) = self.stored_credential(provider) {
354            match credential {
355                StoredCredential::ApiKey { key } if field == "api_key" => return Ok(key.clone()),
356                StoredCredential::SecretFields { fields } => {
357                    if fields.iter().any(|name| name == field) {
358                        return self
359                            .backend
360                            .get(stored_provider, field)?
361                            .ok_or_else(|| missing_secret_error(stored_provider, field));
362                    }
363                }
364                StoredCredential::OAuth(_) => {}
365                StoredCredential::ApiKey { .. } => {}
366            }
367        }
368
369        if let Some(value) = resolve_env_secret(provider, field) {
370            return Ok(value);
371        }
372
373        Err(missing_secret_error(provider, field))
374    }
375
376    /// Store multiple named secret fields securely in the OS keychain and persist only metadata.
377    pub fn store_secret_fields(
378        &mut self,
379        provider: &str,
380        fields: HashMap<String, String>,
381    ) -> Result<()> {
382        if fields.is_empty() {
383            return Err(crate::error::Error::Auth(format!(
384                "No secret fields provided for {provider}."
385            )));
386        }
387
388        let mut field_names = Vec::with_capacity(fields.len());
389        for (field, value) in &fields {
390            let field = field.trim();
391            if field.is_empty() {
392                return Err(crate::error::Error::Auth(format!(
393                    "Secret field names for {provider} cannot be empty."
394                )));
395            }
396            if value.trim().is_empty() {
397                return Err(crate::error::Error::Auth(format!(
398                    "Secret value for {provider}.{field} cannot be empty."
399                )));
400            }
401            self.backend.set(provider, field, value)?;
402            field_names.push(field.to_string());
403        }
404
405        field_names.sort();
406        field_names.dedup();
407        self.stored.insert(
408            provider.to_string(),
409            StoredCredential::SecretFields {
410                fields: field_names,
411            },
412        );
413        self.save()
414    }
415
416    fn stored_credential(&self, provider: &str) -> Option<(&str, &StoredCredential)> {
417        provider_lookup_candidates(provider)
418            .into_iter()
419            .find_map(|candidate| {
420                self.stored
421                    .get_key_value(&candidate)
422                    .map(|(stored_provider, credential)| (stored_provider.as_str(), credential))
423            })
424    }
425
426    /// Check whether stored secret metadata points at readable secure-storage values.
427    pub fn secret_status(&self, provider: &str) -> Option<SecretStatus> {
428        let (stored_provider, credential) = self.stored_credential(provider)?;
429        let fields = match credential {
430            StoredCredential::SecretFields { fields } => fields
431                .iter()
432                .map(|field| {
433                    let status = match self.backend.get(stored_provider, field) {
434                        Ok(Some(value)) if !value.trim().is_empty() => SecretFieldStatus::Present,
435                        Ok(_) => SecretFieldStatus::Missing,
436                        Err(error) => SecretFieldStatus::Error(error.to_string()),
437                    };
438                    (field.clone(), status)
439                })
440                .collect(),
441            StoredCredential::ApiKey { key } => vec![(
442                "api_key".to_string(),
443                if key.trim().is_empty() {
444                    SecretFieldStatus::Missing
445                } else {
446                    SecretFieldStatus::Present
447                },
448            )],
449            StoredCredential::OAuth(oauth) => vec![(
450                "access_token".to_string(),
451                if oauth.access_token.trim().is_empty() {
452                    SecretFieldStatus::Missing
453                } else {
454                    SecretFieldStatus::Present
455                },
456            )],
457        };
458
459        Some(SecretStatus {
460            provider: stored_provider.to_string(),
461            fields,
462        })
463    }
464
465    /// Resolve all stored secret fields for a provider into a map.
466    pub fn resolve_secret_fields(&self, provider: &str) -> Result<HashMap<String, String>> {
467        match self.stored_credential(provider) {
468            Some((stored_provider, StoredCredential::SecretFields { fields })) => fields
469                .iter()
470                .map(|field| {
471                    self.resolve_secret_field(stored_provider, field)
472                        .map(|value| (field.clone(), value))
473                })
474                .collect(),
475            Some((_stored_provider, StoredCredential::ApiKey { key })) => {
476                Ok(HashMap::from([("api_key".to_string(), key.clone())]))
477            }
478            Some((_stored_provider, StoredCredential::OAuth(oauth))) => Ok(HashMap::from([(
479                "access_token".to_string(),
480                oauth.access_token.clone(),
481            )])),
482            None => {
483                if let Some(api_key) = resolve_env_secret(provider, "api_key") {
484                    Ok(HashMap::from([("api_key".to_string(), api_key)]))
485                } else {
486                    Err(missing_secret_error(provider, "api_key"))
487                }
488            }
489        }
490    }
491
492    /// Resolve a ChatGPT/OpenAI OAuth token, preferring `openai-codex` when present.
493    pub async fn resolve_chatgpt_oauth(&mut self) -> Result<ApiKey> {
494        for provider in ["openai-codex", "openai"] {
495            if self.get_oauth(provider).is_none() {
496                continue;
497            }
498
499            return self
500                .resolve_or_refresh(provider, |refresh_token| {
501                    let refresh_token = refresh_token.to_string();
502                    async move {
503                        crate::oauth::chatgpt::ChatGptOAuth::new()
504                            .refresh_token(&refresh_token)
505                            .await
506                    }
507                })
508                .await;
509        }
510
511        Err(crate::error::Error::Auth(
512            "No ChatGPT OAuth credential found. Run `imp login openai` or configure an OpenAI API key."
513                .into(),
514        ))
515    }
516
517    pub fn oauth_display_info(&self, provider: &str) -> Option<OAuthDisplayInfo> {
518        self.get_oauth(provider)
519            .and_then(|credential| oauth_display_info_for_credential(provider, credential))
520    }
521
522    /// Store a credential and persist to disk.
523    pub fn store(&mut self, provider: &str, credential: StoredCredential) -> Result<()> {
524        self.stored.insert(provider.to_string(), credential);
525        self.save()
526    }
527
528    /// Resolve API key, auto-refreshing expired OAuth tokens.
529    /// Persists the refreshed credential to disk on success.
530    pub async fn resolve_with_refresh(&mut self, provider: &str) -> Result<ApiKey> {
531        if let Some(StoredCredential::OAuth(oauth)) = self.stored.get(provider) {
532            if oauth.is_expired() {
533                let refresh_token = oauth.refresh_token.clone();
534                let result = match provider {
535                    "anthropic" => {
536                        crate::oauth::anthropic::AnthropicOAuth::new()
537                            .refresh_token(&refresh_token)
538                            .await
539                    }
540                    "kimi-code" => {
541                        crate::oauth::kimi_code::KimiCodeOAuth::new()
542                            .refresh_token(&refresh_token)
543                            .await
544                    }
545                    _ => {
546                        return Err(crate::error::Error::Auth(format!(
547                            "OAuth refresh not implemented for provider: {provider}"
548                        )));
549                    }
550                };
551                match result {
552                    Ok(new_cred) => {
553                        self.store(provider, StoredCredential::OAuth(new_cred))?;
554                    }
555                    Err(e) => {
556                        return Err(crate::error::Error::Auth(format!(
557                            "Token refresh failed: {e}. Run `imp login` to re-authenticate."
558                        )));
559                    }
560                }
561            }
562        }
563        self.resolve(provider)
564    }
565
566    /// Check if the stored OAuth credential for a provider is expired.
567    pub fn is_oauth_expired(&self, provider: &str) -> bool {
568        matches!(
569            self.stored.get(provider),
570            Some(StoredCredential::OAuth(oauth)) if oauth.is_expired()
571        )
572    }
573
574    /// Get the stored OAuth credential for a provider (if any).
575    pub fn get_oauth(&self, provider: &str) -> Option<&OAuthCredential> {
576        match self.stored.get(provider) {
577            Some(StoredCredential::OAuth(oauth)) => Some(oauth),
578            _ => None,
579        }
580    }
581
582    /// Resolve API key with automatic OAuth refresh.
583    pub async fn resolve_or_refresh<F, Fut>(
584        &mut self,
585        provider: &str,
586        refresh_fn: F,
587    ) -> Result<ApiKey>
588    where
589        F: FnOnce(&str) -> Fut,
590        Fut: std::future::Future<Output = Result<OAuthCredential>>,
591    {
592        if let Some(StoredCredential::OAuth(oauth)) = self.stored.get(provider) {
593            if oauth.is_expired() {
594                let refresh_token = oauth.refresh_token.clone();
595                let new_cred = refresh_fn(&refresh_token).await?;
596                let access_token = new_cred.access_token.clone();
597                self.store(provider, StoredCredential::OAuth(new_cred))?;
598                return Ok(access_token);
599            }
600        }
601        self.resolve(provider)
602    }
603
604    /// Remove a stored credential (logout).
605    pub fn remove(&mut self, provider: &str) -> Result<()> {
606        if let Some(StoredCredential::SecretFields { fields }) = self.stored.remove(provider) {
607            for field in fields {
608                self.backend.delete(provider, &field)?;
609            }
610        }
611        self.save()
612    }
613
614    fn save(&self) -> Result<()> {
615        if let Some(parent) = self.path.parent() {
616            std::fs::create_dir_all(parent)?;
617        }
618        let data = serde_json::to_string_pretty(&self.stored)?;
619        let temp_path = self.path.with_extension("json.tmp");
620        std::fs::write(&temp_path, data)?;
621        #[cfg(unix)]
622        {
623            use std::os::unix::fs::PermissionsExt;
624            let perms = std::fs::Permissions::from_mode(0o600);
625            let _ = std::fs::set_permissions(&temp_path, perms);
626        }
627        std::fs::rename(&temp_path, &self.path)?;
628        #[cfg(unix)]
629        {
630            use std::os::unix::fs::PermissionsExt;
631            let perms = std::fs::Permissions::from_mode(0o600);
632            let _ = std::fs::set_permissions(&self.path, perms);
633        }
634        Ok(())
635    }
636}
637
638fn resolve_env_secret(provider: &str, field: &str) -> Option<String> {
639    if field == "api_key" {
640        let registry = crate::model::ProviderRegistry::with_builtins();
641        if let Some(meta) = registry.find(provider) {
642            for env_var in meta.env_vars {
643                if let Ok(value) = std::env::var(env_var) {
644                    if !value.trim().is_empty() {
645                        return Some(value);
646                    }
647                }
648            }
649        }
650    }
651
652    let env_var = env_var_name(provider, field);
653    std::env::var(&env_var)
654        .ok()
655        .filter(|value| !value.trim().is_empty())
656}
657
658fn env_var_name(provider: &str, field: &str) -> String {
659    let provider = provider.to_uppercase().replace('-', "_");
660    let field = field.to_uppercase().replace('-', "_");
661    format!("{provider}_{field}")
662}
663
664fn missing_secret_error(provider: &str, field: &str) -> crate::error::Error {
665    crate::error::Error::Auth(format!(
666        "No readable secret field '{field}' found for {provider}. Set {} or run `imp secrets {provider}` to save it again.",
667        env_var_name(provider, field)
668    ))
669}
670
671pub fn oauth_display_info_for_credential(
672    provider: &str,
673    credential: &OAuthCredential,
674) -> Option<OAuthDisplayInfo> {
675    match provider {
676        "anthropic" => Some(OAuthDisplayInfo {
677            account_id: None,
678            plan: Some("Claude Max/Pro".into()),
679            using_subscription: true,
680        }),
681        "openai" | "openai-codex" => decode_openai_oauth_display_info(&credential.access_token),
682        "kimi-code" => Some(OAuthDisplayInfo {
683            account_id: None,
684            plan: Some("Kimi Code".into()),
685            using_subscription: true,
686        }),
687        _ => None,
688    }
689}
690
691fn decode_openai_oauth_display_info(access_token: &str) -> Option<OAuthDisplayInfo> {
692    let payload = access_token.split('.').nth(1)?;
693    let decoded = URL_SAFE_NO_PAD.decode(payload).ok()?;
694    let claims: Value = serde_json::from_slice(&decoded).ok()?;
695    let auth = claims.get("https://api.openai.com/auth")?;
696
697    Some(OAuthDisplayInfo {
698        account_id: auth
699            .get("chatgpt_account_id")
700            .and_then(Value::as_str)
701            .map(str::to_string),
702        plan: auth
703            .get("chatgpt_plan_type")
704            .and_then(Value::as_str)
705            .map(str::to_string),
706        using_subscription: true,
707    })
708}
709
710#[cfg(test)]
711mod tests {
712    use super::*;
713    use serde_json::json;
714    use std::sync::Mutex;
715
716    #[derive(Default)]
717    struct MockSecretBackend {
718        values: Mutex<HashMap<(String, String), String>>,
719    }
720
721    impl SecretBackend for MockSecretBackend {
722        fn get(&self, provider: &str, field: &str) -> Result<Option<String>> {
723            Ok(self
724                .values
725                .lock()
726                .unwrap()
727                .get(&(provider.to_string(), field.to_string()))
728                .cloned())
729        }
730
731        fn set(&self, provider: &str, field: &str, value: &str) -> Result<()> {
732            self.values
733                .lock()
734                .unwrap()
735                .insert((provider.to_string(), field.to_string()), value.to_string());
736            Ok(())
737        }
738
739        fn delete(&self, provider: &str, field: &str) -> Result<()> {
740            self.values
741                .lock()
742                .unwrap()
743                .remove(&(provider.to_string(), field.to_string()));
744            Ok(())
745        }
746    }
747
748    struct FailingSetBackend;
749
750    impl SecretBackend for FailingSetBackend {
751        fn get(&self, _provider: &str, _field: &str) -> Result<Option<String>> {
752            Ok(None)
753        }
754
755        fn set(&self, provider: &str, field: &str, _value: &str) -> Result<()> {
756            Err(crate::error::Error::Auth(format!(
757                "test secure storage write failed for {provider}.{field}"
758            )))
759        }
760
761        fn delete(&self, _provider: &str, _field: &str) -> Result<()> {
762            Ok(())
763        }
764    }
765
766    fn test_store(path: std::path::PathBuf) -> AuthStore {
767        AuthStore::new_with_backend(path, Arc::new(MockSecretBackend::default()))
768    }
769
770    fn test_store_with_backend(
771        path: std::path::PathBuf,
772        backend: Arc<dyn SecretBackend>,
773    ) -> AuthStore {
774        AuthStore::new_with_backend(path, backend)
775    }
776
777    fn test_load_with_backend(
778        path: &std::path::Path,
779        backend: Arc<dyn SecretBackend>,
780    ) -> AuthStore {
781        AuthStore::load_with_backend(path, backend).unwrap()
782    }
783
784    fn jwt_with_openai_auth(plan: &str, account_id: &str) -> String {
785        let header = URL_SAFE_NO_PAD.encode(r#"{"alg":"RS256","typ":"JWT"}"#);
786        let payload = URL_SAFE_NO_PAD.encode(
787            json!({
788                "https://api.openai.com/auth": {
789                    "chatgpt_account_id": account_id,
790                    "chatgpt_plan_type": plan,
791                }
792            })
793            .to_string(),
794        );
795        format!("{header}.{payload}.signature")
796    }
797
798    #[test]
799    fn test_oauth_credential_not_expired() {
800        let cred = OAuthCredential {
801            access_token: "token".into(),
802            refresh_token: "refresh".into(),
803            expires_at: crate::now() + 3600,
804        };
805        assert!(!cred.is_expired());
806    }
807
808    #[test]
809    fn test_oauth_credential_expired() {
810        let cred = OAuthCredential {
811            access_token: "token".into(),
812            refresh_token: "refresh".into(),
813            expires_at: crate::now().saturating_sub(100),
814        };
815        assert!(cred.is_expired());
816    }
817
818    #[test]
819    fn test_oauth_store_and_resolve() {
820        let dir = tempfile::tempdir().unwrap();
821        let path = dir.path().join("auth.json");
822        let mut store = test_store(path);
823
824        let cred = OAuthCredential {
825            access_token: "sk-ant-access".into(),
826            refresh_token: "rt-refresh".into(),
827            expires_at: crate::now() + 3600,
828        };
829        store
830            .store("anthropic", StoredCredential::OAuth(cred))
831            .unwrap();
832
833        let key = store.resolve("anthropic").unwrap();
834        assert_eq!(key, "sk-ant-access");
835    }
836
837    #[test]
838    fn test_secure_secret_fields_store_and_resolve() {
839        let dir = tempfile::tempdir().unwrap();
840        let path = dir.path().join("auth.json");
841        let mut store = test_store(path.clone());
842        let mut fields = HashMap::new();
843        fields.insert("api_key".to_string(), "test-api".to_string());
844        fields.insert("secret_key".to_string(), "test-secret".to_string());
845        store.store_secret_fields("test-service", fields).unwrap();
846
847        let data = std::fs::read_to_string(&path).unwrap();
848        assert!(!data.contains("test-api"));
849        assert!(!data.contains("test-secret"));
850        assert_eq!(
851            store
852                .resolve_secret_field("test-service", "api_key")
853                .unwrap(),
854            "test-api"
855        );
856        assert_eq!(
857            store
858                .resolve_secret_field("test-service", "secret_key")
859                .unwrap(),
860            "test-secret"
861        );
862    }
863
864    #[test]
865    fn store_secret_fields_does_not_save_metadata_when_secure_write_fails() {
866        let dir = tempfile::tempdir().unwrap();
867        let path = dir.path().join("auth.json");
868        let mut store = test_store_with_backend(path.clone(), Arc::new(FailingSetBackend));
869
870        let result = store.store_secret_fields(
871            "google",
872            HashMap::from([("api_key".to_string(), "test-api".to_string())]),
873        );
874
875        assert!(result.is_err());
876        assert!(!store.stored.contains_key("google"));
877        assert!(!path.exists());
878    }
879
880    #[test]
881    fn test_secure_secret_fields_persist_and_load() {
882        let dir = tempfile::tempdir().unwrap();
883        let path = dir.path().join("auth.json");
884        let backend: Arc<dyn SecretBackend> = Arc::new(MockSecretBackend::default());
885        let mut store = test_store_with_backend(path.clone(), Arc::clone(&backend));
886        store
887            .store_secret_fields(
888                "test-service",
889                HashMap::from([
890                    ("api_key".to_string(), "test-api".to_string()),
891                    ("secret_key".to_string(), "test-secret".to_string()),
892                ]),
893            )
894            .unwrap();
895
896        let loaded = test_load_with_backend(&path, backend);
897        let resolved = loaded.resolve_secret_fields("test-service").unwrap();
898        assert_eq!(
899            resolved.get("api_key").map(String::as_str),
900            Some("test-api")
901        );
902        assert_eq!(
903            resolved.get("secret_key").map(String::as_str),
904            Some("test-secret")
905        );
906    }
907
908    #[test]
909    fn test_secure_remove_deletes_secret_fields() {
910        let dir = tempfile::tempdir().unwrap();
911        let path = dir.path().join("auth.json");
912        let backend: Arc<dyn SecretBackend> = Arc::new(MockSecretBackend::default());
913        let mut store = test_store_with_backend(path, Arc::clone(&backend));
914        store
915            .store_secret_fields(
916                "test-service",
917                HashMap::from([
918                    ("api_key".to_string(), "test-api".to_string()),
919                    ("secret_key".to_string(), "test-secret".to_string()),
920                ]),
921            )
922            .unwrap();
923
924        store.remove("test-service").unwrap();
925        assert!(store
926            .resolve_secret_field("test-service", "api_key")
927            .is_err());
928        assert!(backend.get("test-service", "api_key").unwrap().is_none());
929    }
930
931    #[test]
932    fn test_oauth_detect_expiry() {
933        let dir = tempfile::tempdir().unwrap();
934        let path = dir.path().join("auth.json");
935        let mut store = test_store(path);
936
937        let fresh = OAuthCredential {
938            access_token: "fresh".into(),
939            refresh_token: "rt".into(),
940            expires_at: crate::now() + 3600,
941        };
942        store
943            .store("anthropic", StoredCredential::OAuth(fresh))
944            .unwrap();
945        assert!(!store.is_oauth_expired("anthropic"));
946
947        let expired = OAuthCredential {
948            access_token: "expired".into(),
949            refresh_token: "rt".into(),
950            expires_at: 0,
951        };
952        store
953            .store("anthropic", StoredCredential::OAuth(expired))
954            .unwrap();
955        assert!(store.is_oauth_expired("anthropic"));
956    }
957
958    #[tokio::test]
959    async fn test_oauth_resolve_or_refresh() {
960        let dir = tempfile::tempdir().unwrap();
961        let path = dir.path().join("auth.json");
962        let mut store = test_store(path);
963
964        let expired = OAuthCredential {
965            access_token: "old-access".into(),
966            refresh_token: "rt-for-refresh".into(),
967            expires_at: 0,
968        };
969        store
970            .store("anthropic", StoredCredential::OAuth(expired))
971            .unwrap();
972
973        let key = store
974            .resolve_or_refresh("anthropic", |refresh_tok| {
975                let refresh_tok = refresh_tok.to_string();
976                async move {
977                    assert_eq!(refresh_tok, "rt-for-refresh");
978                    Ok(OAuthCredential {
979                        access_token: "new-access".into(),
980                        refresh_token: "new-rt".into(),
981                        expires_at: crate::now() + 3600,
982                    })
983                }
984            })
985            .await
986            .unwrap();
987
988        assert_eq!(key, "new-access");
989        let resolved = store.resolve("anthropic").unwrap();
990        assert_eq!(resolved, "new-access");
991    }
992
993    #[tokio::test]
994    async fn test_oauth_resolve_or_refresh_not_expired() {
995        let dir = tempfile::tempdir().unwrap();
996        let path = dir.path().join("auth.json");
997        let mut store = test_store(path);
998
999        let fresh = OAuthCredential {
1000            access_token: "still-valid".into(),
1001            refresh_token: "rt".into(),
1002            expires_at: crate::now() + 3600,
1003        };
1004        store
1005            .store("anthropic", StoredCredential::OAuth(fresh))
1006            .unwrap();
1007
1008        let key = store
1009            .resolve_or_refresh("anthropic", |_| async {
1010                panic!("refresh should not be called for non-expired token");
1011            })
1012            .await
1013            .unwrap();
1014
1015        assert_eq!(key, "still-valid");
1016    }
1017
1018    #[test]
1019    fn test_load_invalid_auth_metadata_returns_error() {
1020        let dir = tempfile::tempdir().unwrap();
1021        let path = dir.path().join("auth.json");
1022        std::fs::write(&path, "{not valid json").unwrap();
1023
1024        let backend: Arc<dyn SecretBackend> = Arc::new(MockSecretBackend::default());
1025        let err = match AuthStore::load_with_backend(&path, backend) {
1026            Ok(_) => panic!("invalid auth metadata should error"),
1027            Err(err) => err,
1028        };
1029        let msg = err.to_string();
1030        assert!(msg.contains("Failed to parse auth metadata"));
1031        assert!(msg.contains("auth.json"));
1032    }
1033
1034    #[test]
1035    fn test_save_writes_atomically_without_leaving_temp_file() {
1036        let dir = tempfile::tempdir().unwrap();
1037        let path = dir.path().join("auth.json");
1038        let mut store = test_store(path.clone());
1039
1040        store
1041            .store(
1042                "openai",
1043                StoredCredential::ApiKey {
1044                    key: "sk-atomic".into(),
1045                },
1046            )
1047            .unwrap();
1048
1049        assert!(path.exists());
1050        assert!(!path.with_extension("json.tmp").exists());
1051        let loaded = test_load_with_backend(&path, Arc::new(MockSecretBackend::default()));
1052        assert_eq!(loaded.resolve("openai").unwrap(), "sk-atomic");
1053    }
1054
1055    #[test]
1056    fn test_oauth_store_persist_and_load() {
1057        let dir = tempfile::tempdir().unwrap();
1058        let path = dir.path().join("auth.json");
1059        let backend: Arc<dyn SecretBackend> = Arc::new(MockSecretBackend::default());
1060
1061        {
1062            let mut store = test_store_with_backend(path.clone(), Arc::clone(&backend));
1063            let cred = OAuthCredential {
1064                access_token: "persisted-token".into(),
1065                refresh_token: "persisted-rt".into(),
1066                expires_at: crate::now() + 3600,
1067            };
1068            store
1069                .store("anthropic", StoredCredential::OAuth(cred))
1070                .unwrap();
1071        }
1072
1073        let store = test_load_with_backend(&path, backend);
1074        let key = store.resolve("anthropic").unwrap();
1075        assert_eq!(key, "persisted-token");
1076    }
1077
1078    #[test]
1079    fn test_oauth_remove_credential() {
1080        let dir = tempfile::tempdir().unwrap();
1081        let path = dir.path().join("auth.json");
1082        let mut store = test_store(path);
1083
1084        let cred = OAuthCredential {
1085            access_token: "to-remove".into(),
1086            refresh_token: "rt".into(),
1087            expires_at: crate::now() + 3600,
1088        };
1089        store
1090            .store("anthropic", StoredCredential::OAuth(cred))
1091            .unwrap();
1092        assert!(store.resolve("anthropic").is_ok());
1093
1094        store.remove("anthropic").unwrap();
1095        std::env::remove_var("ANTHROPIC_API_KEY");
1096        assert!(store.resolve("anthropic").is_err());
1097    }
1098
1099    #[test]
1100    fn test_resolve_order_runtime_over_stored() {
1101        let dir = tempfile::tempdir().unwrap();
1102        let path = dir.path().join("auth.json");
1103        let mut store = test_store(path);
1104
1105        store
1106            .store(
1107                "anthropic",
1108                StoredCredential::ApiKey {
1109                    key: "stored-key".into(),
1110                },
1111            )
1112            .unwrap();
1113
1114        store.set_runtime_key("anthropic", "runtime-key".into());
1115        let key = store.resolve("anthropic").unwrap();
1116        assert_eq!(key, "runtime-key");
1117    }
1118
1119    #[test]
1120    fn test_set_runtime_key_ignores_empty_or_whitespace_values() {
1121        let dir = tempfile::tempdir().unwrap();
1122        let path = dir.path().join("auth.json");
1123        let mut store = test_store(path);
1124
1125        store.set_runtime_key("openai", "runtime-key".into());
1126        assert_eq!(store.resolve("openai").unwrap(), "runtime-key");
1127
1128        store.set_runtime_key("openai", "   ".into());
1129        assert!(store.resolve("openai").is_err());
1130    }
1131
1132    #[test]
1133    fn test_resolve_stored_api_key() {
1134        let dir = tempfile::tempdir().unwrap();
1135        let path = dir.path().join("auth.json");
1136        let mut store = test_store(path);
1137
1138        store
1139            .store(
1140                "openai",
1141                StoredCredential::ApiKey {
1142                    key: "sk-stored".into(),
1143                },
1144            )
1145            .unwrap();
1146
1147        let key = store.resolve("openai").unwrap();
1148        assert_eq!(key, "sk-stored");
1149    }
1150
1151    #[test]
1152    fn test_resolve_env_secret_uses_moonshot_env_vars() {
1153        let dir = tempfile::tempdir().unwrap();
1154        let path = dir.path().join("auth.json");
1155        let store = test_store(path);
1156
1157        std::env::remove_var("KIMI_API_KEY");
1158        std::env::set_var("MOONSHOT_API_KEY", "moonshot-env-key");
1159        let key = store.resolve("moonshot").unwrap();
1160        assert_eq!(key, "moonshot-env-key");
1161        std::env::remove_var("MOONSHOT_API_KEY");
1162
1163        std::env::set_var("KIMI_API_KEY", "kimi-env-key");
1164        let key = store.resolve("moonshot").unwrap();
1165        assert_eq!(key, "kimi-env-key");
1166        std::env::remove_var("KIMI_API_KEY");
1167    }
1168
1169    #[test]
1170    fn test_resolve_api_key_only_ignores_oauth_credentials() {
1171        let dir = tempfile::tempdir().unwrap();
1172        let path = dir.path().join("auth.json");
1173        let mut store = test_store(path);
1174
1175        store
1176            .store(
1177                "openai",
1178                StoredCredential::OAuth(OAuthCredential {
1179                    access_token: "oauth-token".into(),
1180                    refresh_token: "refresh-token".into(),
1181                    expires_at: crate::now() + 3600,
1182                }),
1183            )
1184            .unwrap();
1185
1186        assert!(store.resolve_api_key_only("openai").is_err());
1187    }
1188
1189    #[tokio::test]
1190    async fn test_resolve_chatgpt_oauth_prefers_openai_codex() {
1191        let dir = tempfile::tempdir().unwrap();
1192        let path = dir.path().join("auth.json");
1193        let mut store = test_store(path);
1194
1195        store
1196            .store(
1197                "openai",
1198                StoredCredential::OAuth(OAuthCredential {
1199                    access_token: "openai-oauth".into(),
1200                    refresh_token: "openai-refresh".into(),
1201                    expires_at: crate::now() + 3600,
1202                }),
1203            )
1204            .unwrap();
1205        store
1206            .store(
1207                "openai-codex",
1208                StoredCredential::OAuth(OAuthCredential {
1209                    access_token: "codex-oauth".into(),
1210                    refresh_token: "codex-refresh".into(),
1211                    expires_at: crate::now() + 3600,
1212                }),
1213            )
1214            .unwrap();
1215
1216        let key = store.resolve_chatgpt_oauth().await.unwrap();
1217        assert_eq!(key, "codex-oauth");
1218    }
1219
1220    #[tokio::test]
1221    async fn test_resolve_chatgpt_oauth_falls_back_to_openai() {
1222        let dir = tempfile::tempdir().unwrap();
1223        let path = dir.path().join("auth.json");
1224        let mut store = test_store(path);
1225
1226        store
1227            .store(
1228                "openai",
1229                StoredCredential::OAuth(OAuthCredential {
1230                    access_token: "openai-oauth".into(),
1231                    refresh_token: "openai-refresh".into(),
1232                    expires_at: crate::now() + 3600,
1233                }),
1234            )
1235            .unwrap();
1236
1237        let key = store.resolve_chatgpt_oauth().await.unwrap();
1238        assert_eq!(key, "openai-oauth");
1239    }
1240
1241    #[test]
1242    fn test_oauth_display_info_for_openai_credential() {
1243        let credential = OAuthCredential {
1244            access_token: jwt_with_openai_auth("pro", "acct-12345678"),
1245            refresh_token: "refresh".into(),
1246            expires_at: crate::now() + 3600,
1247        };
1248
1249        let info = oauth_display_info_for_credential("openai", &credential).unwrap();
1250        assert_eq!(info.account_id.as_deref(), Some("acct-12345678"));
1251        assert_eq!(info.plan.as_deref(), Some("pro"));
1252        assert_eq!(info.short_account_id().as_deref(), Some("acct-123…"));
1253    }
1254
1255    #[test]
1256    fn test_oauth_display_info_for_anthropic_credential() {
1257        let credential = OAuthCredential {
1258            access_token: "sk-ant-oat01-example".into(),
1259            refresh_token: "refresh".into(),
1260            expires_at: crate::now() + 3600,
1261        };
1262
1263        let info = oauth_display_info_for_credential("anthropic", &credential).unwrap();
1264        assert_eq!(info.plan.as_deref(), Some("Claude Max/Pro"));
1265        assert!(info.account_id.is_none());
1266        assert_eq!(
1267            info.login_message("anthropic"),
1268            "Logged in to Anthropic with Claude Max/Pro subscription credentials."
1269        );
1270    }
1271
1272    #[test]
1273    fn test_remove_then_resolve_falls_through() {
1274        let dir = tempfile::tempdir().unwrap();
1275        let path = dir.path().join("auth.json");
1276        let mut store = test_store(path);
1277
1278        store
1279            .store(
1280                "google",
1281                StoredCredential::ApiKey {
1282                    key: "google-key".into(),
1283                },
1284            )
1285            .unwrap();
1286        assert!(store.resolve("google").is_ok());
1287
1288        store.remove("google").unwrap();
1289        std::env::remove_var("GOOGLE_API_KEY");
1290        let result = store.resolve("google");
1291        assert!(result.is_err());
1292    }
1293
1294    #[test]
1295    fn provider_lookup_candidates_include_legacy_render_casing() {
1296        assert_eq!(
1297            provider_lookup_candidates("render"),
1298            vec!["render".to_string(), "Render".to_string()]
1299        );
1300        assert_eq!(
1301            provider_lookup_candidates("Render"),
1302            vec!["Render".to_string(), "render".to_string()]
1303        );
1304    }
1305
1306    #[test]
1307    fn field_lookup_candidates_support_porkbun_secret_key_typo() {
1308        assert_eq!(
1309            field_lookup_candidates("secrets_key"),
1310            vec!["secrets_key".to_string(), "secret_key".to_string()]
1311        );
1312    }
1313
1314    #[test]
1315    fn resolve_secret_fields_uses_provider_alias_candidates() {
1316        let dir = tempfile::tempdir().unwrap();
1317        let path = dir.path().join("auth.json");
1318        let mut store = test_store(path);
1319
1320        store
1321            .store_secret_fields(
1322                "Render",
1323                HashMap::from([("api_key".to_string(), "render-secret".to_string())]),
1324            )
1325            .unwrap();
1326
1327        let fields = store.resolve_secret_fields("render").unwrap();
1328        assert_eq!(
1329            fields.get("api_key").map(String::as_str),
1330            Some("render-secret")
1331        );
1332    }
1333
1334    #[test]
1335    fn test_secret_status_reports_missing_keychain_values() {
1336        let dir = tempfile::tempdir().unwrap();
1337        let path = dir.path().join("auth.json");
1338        let mut store = test_store(path);
1339        store.stored.insert(
1340            "google".into(),
1341            StoredCredential::SecretFields {
1342                fields: vec!["api_key".into()],
1343            },
1344        );
1345
1346        let status = store.secret_status("google").unwrap();
1347        assert_eq!(status.provider, "google");
1348        assert_eq!(
1349            status.fields,
1350            vec![("api_key".to_string(), SecretFieldStatus::Missing)]
1351        );
1352        assert!(!status.is_usable());
1353        assert!(!store.has_credentials("google"));
1354    }
1355
1356    #[test]
1357    fn test_has_credentials_detects_kimi_env_var() {
1358        let dir = tempfile::tempdir().unwrap();
1359        let path = dir.path().join("auth.json");
1360        let store = test_store(path);
1361
1362        std::env::remove_var("MOONSHOT_API_KEY");
1363        std::env::set_var("KIMI_API_KEY", "kimi-env-key");
1364        assert!(store.has_credentials("moonshot"));
1365        std::env::remove_var("KIMI_API_KEY");
1366    }
1367
1368    #[test]
1369    fn test_unknown_provider_returns_auth_error() {
1370        let dir = tempfile::tempdir().unwrap();
1371        let path = dir.path().join("auth.json");
1372        let store = test_store(path);
1373        let result = store.resolve("unknown_provider");
1374        assert!(result.is_err());
1375        let err = result.unwrap_err();
1376        assert!(matches!(err, crate::error::Error::Auth(_)));
1377    }
1378}
imp_llm/auth.rs

imp_llm/
auth.rs
