vtcode_config/
api_keys.rs

1//! API key management module for secure retrieval from environment variables,
2//! .env files, and configuration files.
3//!
4//! This module provides a unified interface for retrieving API keys for different providers,
5//! prioritizing security by checking environment variables first, then .env files, and finally
6//! falling back to configuration file values.
7
8use anyhow::Result;
9use std::env;
10use std::str::FromStr;
11
12use crate::auth::CustomApiKeyStorage;
13use crate::constants::defaults;
14use crate::models::Provider;
15
16/// API key sources for different providers
17#[derive(Debug, Clone)]
18pub struct ApiKeySources {
19    /// Gemini API key environment variable name
20    pub gemini_env: String,
21    /// Anthropic API key environment variable name
22    pub anthropic_env: String,
23    /// OpenAI API key environment variable name
24    pub openai_env: String,
25    /// OpenRouter API key environment variable name
26    pub openrouter_env: String,
27    /// DeepSeek API key environment variable name
28    pub deepseek_env: String,
29    /// Z.AI API key environment variable name
30    pub zai_env: String,
31    /// Ollama API key environment variable name
32    pub ollama_env: String,
33    /// LM Studio API key environment variable name
34    pub lmstudio_env: String,
35    /// Gemini API key from configuration file
36    pub gemini_config: Option<String>,
37    /// Anthropic API key from configuration file
38    pub anthropic_config: Option<String>,
39    /// OpenAI API key from configuration file
40    pub openai_config: Option<String>,
41    /// OpenRouter API key from configuration file
42    pub openrouter_config: Option<String>,
43    /// DeepSeek API key from configuration file
44    pub deepseek_config: Option<String>,
45    /// Z.AI API key from configuration file
46    pub zai_config: Option<String>,
47    /// Ollama API key from configuration file
48    pub ollama_config: Option<String>,
49    /// LM Studio API key from configuration file
50    pub lmstudio_config: Option<String>,
51}
52
53impl Default for ApiKeySources {
54    fn default() -> Self {
55        Self {
56            gemini_env: "GEMINI_API_KEY".to_string(),
57            anthropic_env: "ANTHROPIC_API_KEY".to_string(),
58            openai_env: "OPENAI_API_KEY".to_string(),
59            openrouter_env: "OPENROUTER_API_KEY".to_string(),
60            deepseek_env: "DEEPSEEK_API_KEY".to_string(),
61            zai_env: "ZAI_API_KEY".to_string(),
62            ollama_env: "OLLAMA_API_KEY".to_string(),
63            lmstudio_env: "LMSTUDIO_API_KEY".to_string(),
64            gemini_config: None,
65            anthropic_config: None,
66            openai_config: None,
67            openrouter_config: None,
68            deepseek_config: None,
69            zai_config: None,
70            ollama_config: None,
71            lmstudio_config: None,
72        }
73    }
74}
75
76impl ApiKeySources {
77    /// Create API key sources for a specific provider with automatic environment variable inference
78    pub fn for_provider(_provider: &str) -> Self {
79        Self::default()
80    }
81}
82
83pub fn api_key_env_var(provider: &str) -> String {
84    let trimmed = provider.trim();
85    if trimmed.is_empty() {
86        return defaults::DEFAULT_API_KEY_ENV.to_owned();
87    }
88
89    if let Ok(resolved) = Provider::from_str(trimmed)
90        && resolved.uses_managed_auth()
91    {
92        return String::new();
93    }
94
95    Provider::from_str(trimmed)
96        .map(|resolved| resolved.default_api_key_env().to_owned())
97        .unwrap_or_else(|_| format!("{}_API_KEY", trimmed.to_ascii_uppercase()))
98}
99
100pub fn resolve_api_key_env(provider: &str, configured_env: &str) -> String {
101    let trimmed = configured_env.trim();
102    if trimmed.is_empty() || trimmed.eq_ignore_ascii_case(defaults::DEFAULT_API_KEY_ENV) {
103        api_key_env_var(provider)
104    } else {
105        trimmed.to_owned()
106    }
107}
108
109#[cfg(test)]
110mod test_env_overrides {
111    use hashbrown::HashMap;
112    use std::sync::{LazyLock, Mutex};
113
114    static OVERRIDES: LazyLock<Mutex<HashMap<String, Option<String>>>> =
115        LazyLock::new(|| Mutex::new(HashMap::new()));
116
117    pub(super) fn get(key: &str) -> Option<Option<String>> {
118        OVERRIDES.lock().ok().and_then(|map| map.get(key).cloned())
119    }
120
121    pub(super) fn set(key: &str, value: Option<&str>) {
122        if let Ok(mut map) = OVERRIDES.lock() {
123            map.insert(key.to_string(), value.map(ToString::to_string));
124        }
125    }
126
127    pub(super) fn restore(key: &str, previous: Option<Option<String>>) {
128        if let Ok(mut map) = OVERRIDES.lock() {
129            match previous {
130                Some(value) => {
131                    map.insert(key.to_string(), value);
132                }
133                None => {
134                    map.remove(key);
135                }
136            }
137        }
138    }
139}
140
141fn read_env_var(key: &str) -> Option<String> {
142    #[cfg(test)]
143    if let Some(override_value) = test_env_overrides::get(key) {
144        return override_value;
145    }
146
147    env::var(key).ok()
148}
149
150/// Load environment variables from .env file
151///
152/// This function attempts to load environment variables from a .env file
153/// in the current directory. It logs a warning if the file exists but cannot
154/// be loaded, but doesn't fail if the file doesn't exist.
155pub fn load_dotenv() -> Result<()> {
156    match dotenvy::dotenv() {
157        Ok(path) => {
158            // Only print in verbose mode to avoid polluting stdout/stderr in scripts
159            if read_env_var("VTCODE_VERBOSE").is_some() || read_env_var("RUST_LOG").is_some() {
160                tracing::info!("Loaded environment variables from: {}", path.display());
161            }
162            Ok(())
163        }
164        Err(dotenvy::Error::Io(e)) if e.kind() == std::io::ErrorKind::NotFound => {
165            // .env file doesn't exist, which is fine
166            Ok(())
167        }
168        Err(e) => {
169            tracing::warn!("Failed to load .env file: {}", e);
170            Ok(())
171        }
172    }
173}
174
175/// Get API key for a specific provider with secure fallback mechanism
176///
177/// This function implements a secure retrieval mechanism that:
178/// 1. First checks environment variables (highest priority for security)
179/// 2. Then checks .env file values
180/// 3. Falls back to configuration file values if neither above is set
181/// 4. Supports all major providers: Gemini, Anthropic, OpenAI, and OpenRouter
182/// 5. Automatically infers the correct environment variable based on provider
183///
184/// # Arguments
185///
186/// * `provider` - The provider name ("gemini", "anthropic", or "openai")
187/// * `sources` - Configuration for where to look for API keys
188///
189/// # Returns
190///
191/// * `Ok(String)` - The API key if found
192/// * `Err` - If no API key could be found for the provider
193pub fn get_api_key(provider: &str, sources: &ApiKeySources) -> Result<String> {
194    let normalized_provider = provider.to_lowercase();
195    // Automatically infer the correct environment variable based on provider
196    let inferred_env = api_key_env_var(&normalized_provider);
197
198    // Try the inferred environment variable first
199    if let Some(key) = read_env_var(&inferred_env)
200        && !key.is_empty()
201    {
202        return Ok(key);
203    }
204
205    // Try secure storage (keyring) for custom API keys
206    if let Ok(Some(key)) = get_custom_api_key_from_secure_storage(&normalized_provider) {
207        return Ok(key);
208    }
209
210    // Fall back to the provider-specific sources
211    match normalized_provider.as_str() {
212        "gemini" => get_gemini_api_key(sources),
213        "anthropic" => get_anthropic_api_key(sources),
214        "openai" => get_openai_api_key(sources),
215        "copilot" => Err(anyhow::anyhow!(
216            "GitHub Copilot authentication is managed by the official `copilot` CLI. Run `vtcode login copilot`."
217        )),
218        "deepseek" => get_deepseek_api_key(sources),
219        "openrouter" => get_openrouter_api_key(sources),
220        "zai" => get_zai_api_key(sources),
221        "ollama" => get_ollama_api_key(sources),
222        "lmstudio" => get_lmstudio_api_key(sources),
223        "huggingface" => {
224            read_env_var("HF_TOKEN").ok_or_else(|| anyhow::anyhow!("HF_TOKEN not set"))
225        }
226        _ => Err(anyhow::anyhow!("Unsupported provider: {}", provider)),
227    }
228}
229
230/// Get a custom API key from secure storage.
231///
232/// This function retrieves API keys that were stored securely via the model picker
233/// or interactive configuration flows. When the OS keyring is unavailable, the
234/// auth layer falls back to encrypted file storage automatically.
235///
236/// # Arguments
237/// * `provider` - The provider name
238///
239/// # Returns
240/// * `Ok(Some(String))` - The API key if found in secure storage
241/// * `Ok(None)` - If no key is stored for this provider
242/// * `Err` - If there was an error accessing secure storage
243fn get_custom_api_key_from_secure_storage(provider: &str) -> Result<Option<String>> {
244    let storage = CustomApiKeyStorage::new(provider);
245    // The auth layer handles keyring-to-file fallback internally.
246    let mode = crate::auth::AuthCredentialsStoreMode::default();
247    storage.load(mode)
248}
249
250/// Get API key for a specific environment variable with fallback
251fn get_api_key_with_fallback(
252    env_var: &str,
253    config_value: Option<&String>,
254    provider_name: &str,
255) -> Result<String> {
256    // First try environment variable (most secure)
257    if let Some(key) = read_env_var(env_var)
258        && !key.is_empty()
259    {
260        return Ok(key);
261    }
262
263    // Then try configuration file value
264    if let Some(key) = config_value
265        && !key.is_empty()
266    {
267        return Ok(key.clone());
268    }
269
270    // If neither worked, return an error
271    Err(anyhow::anyhow!(
272        "No API key found for {} provider. Set {} environment variable (or add to .env file) or configure in vtcode.toml",
273        provider_name,
274        env_var
275    ))
276}
277
278fn get_optional_api_key_with_fallback(env_var: &str, config_value: Option<&String>) -> String {
279    if let Some(key) = read_env_var(env_var)
280        && !key.is_empty()
281    {
282        return key;
283    }
284
285    if let Some(key) = config_value
286        && !key.is_empty()
287    {
288        return key.clone();
289    }
290
291    String::new()
292}
293
294/// Get Gemini API key with secure fallback
295fn get_gemini_api_key(sources: &ApiKeySources) -> Result<String> {
296    // Try primary Gemini environment variable
297    if let Some(key) = read_env_var(&sources.gemini_env)
298        && !key.is_empty()
299    {
300        return Ok(key);
301    }
302
303    // Try Google API key as fallback (for backward compatibility)
304    if let Some(key) = read_env_var("GOOGLE_API_KEY")
305        && !key.is_empty()
306    {
307        return Ok(key);
308    }
309
310    // Try configuration file value
311    if let Some(key) = &sources.gemini_config
312        && !key.is_empty()
313    {
314        return Ok(key.clone());
315    }
316
317    // If nothing worked, return an error
318    Err(anyhow::anyhow!(
319        "No API key found for Gemini provider. Set {} or GOOGLE_API_KEY environment variable (or add to .env file) or configure in vtcode.toml",
320        sources.gemini_env
321    ))
322}
323
324/// Get Anthropic API key with secure fallback
325fn get_anthropic_api_key(sources: &ApiKeySources) -> Result<String> {
326    get_api_key_with_fallback(
327        &sources.anthropic_env,
328        sources.anthropic_config.as_ref(),
329        "Anthropic",
330    )
331}
332
333/// Get OpenAI API key with secure fallback
334fn get_openai_api_key(sources: &ApiKeySources) -> Result<String> {
335    get_api_key_with_fallback(
336        &sources.openai_env,
337        sources.openai_config.as_ref(),
338        "OpenAI",
339    )
340}
341
342/// Get OpenRouter API key with secure fallback
343///
344/// This function checks for credentials in the following order:
345/// 1. OAuth token from encrypted storage (if OAuth is enabled)
346/// 2. Environment variable (OPENROUTER_API_KEY)
347/// 3. Configuration file value
348fn get_openrouter_api_key(sources: &ApiKeySources) -> Result<String> {
349    // First, try to load OAuth token from encrypted storage
350    if let Ok(Some(token)) = crate::auth::load_oauth_token() {
351        tracing::debug!("Using OAuth token for OpenRouter authentication");
352        return Ok(token.api_key);
353    }
354
355    // Fall back to standard API key retrieval
356    get_api_key_with_fallback(
357        &sources.openrouter_env,
358        sources.openrouter_config.as_ref(),
359        "OpenRouter",
360    )
361}
362
363/// Get DeepSeek API key with secure fallback
364fn get_deepseek_api_key(sources: &ApiKeySources) -> Result<String> {
365    get_api_key_with_fallback(
366        &sources.deepseek_env,
367        sources.deepseek_config.as_ref(),
368        "DeepSeek",
369    )
370}
371
372/// Get Z.AI API key with secure fallback
373fn get_zai_api_key(sources: &ApiKeySources) -> Result<String> {
374    get_api_key_with_fallback(&sources.zai_env, sources.zai_config.as_ref(), "Z.AI")
375}
376
377/// Get Ollama API key with secure fallback
378fn get_ollama_api_key(sources: &ApiKeySources) -> Result<String> {
379    // For Ollama we allow running without credentials when connecting to a local deployment.
380    // Cloud variants still rely on environment/config values when present.
381    Ok(get_optional_api_key_with_fallback(
382        &sources.ollama_env,
383        sources.ollama_config.as_ref(),
384    ))
385}
386
387/// Get LM Studio API key with secure fallback
388fn get_lmstudio_api_key(sources: &ApiKeySources) -> Result<String> {
389    Ok(get_optional_api_key_with_fallback(
390        &sources.lmstudio_env,
391        sources.lmstudio_config.as_ref(),
392    ))
393}
394
395#[cfg(test)]
396mod tests {
397    use super::*;
398
399    struct EnvOverrideGuard {
400        key: &'static str,
401        previous: Option<Option<String>>,
402    }
403
404    impl EnvOverrideGuard {
405        fn set(key: &'static str, value: Option<&str>) -> Self {
406            let previous = test_env_overrides::get(key);
407            test_env_overrides::set(key, value);
408            Self { key, previous }
409        }
410    }
411
412    impl Drop for EnvOverrideGuard {
413        fn drop(&mut self) {
414            test_env_overrides::restore(self.key, self.previous.clone());
415        }
416    }
417
418    fn with_override<F>(key: &'static str, value: Option<&str>, f: F)
419    where
420        F: FnOnce(),
421    {
422        let _guard = EnvOverrideGuard::set(key, value);
423        f();
424    }
425
426    fn with_overrides<F>(overrides: &[(&'static str, Option<&str>)], f: F)
427    where
428        F: FnOnce(),
429    {
430        let _guards: Vec<_> = overrides
431            .iter()
432            .map(|(key, value)| EnvOverrideGuard::set(key, *value))
433            .collect();
434        f();
435    }
436
437    #[test]
438    fn test_get_gemini_api_key_from_env() {
439        with_override("TEST_GEMINI_KEY", Some("test-gemini-key"), || {
440            let sources = ApiKeySources {
441                gemini_env: "TEST_GEMINI_KEY".to_string(),
442                ..Default::default()
443            };
444
445            let result = get_gemini_api_key(&sources);
446            assert!(result.is_ok());
447            assert_eq!(result.unwrap(), "test-gemini-key");
448        });
449    }
450
451    #[test]
452    fn test_get_anthropic_api_key_from_env() {
453        with_override("TEST_ANTHROPIC_KEY", Some("test-anthropic-key"), || {
454            let sources = ApiKeySources {
455                anthropic_env: "TEST_ANTHROPIC_KEY".to_string(),
456                ..Default::default()
457            };
458
459            let result = get_anthropic_api_key(&sources);
460            assert!(result.is_ok());
461            assert_eq!(result.unwrap(), "test-anthropic-key");
462        });
463    }
464
465    #[test]
466    fn test_get_openai_api_key_from_env() {
467        with_override("TEST_OPENAI_KEY", Some("test-openai-key"), || {
468            let sources = ApiKeySources {
469                openai_env: "TEST_OPENAI_KEY".to_string(),
470                ..Default::default()
471            };
472
473            let result = get_openai_api_key(&sources);
474            assert!(result.is_ok());
475            assert_eq!(result.unwrap(), "test-openai-key");
476        });
477    }
478
479    #[test]
480    fn test_get_deepseek_api_key_from_env() {
481        with_override("TEST_DEEPSEEK_KEY", Some("test-deepseek-key"), || {
482            let sources = ApiKeySources {
483                deepseek_env: "TEST_DEEPSEEK_KEY".to_string(),
484                ..Default::default()
485            };
486
487            let result = get_deepseek_api_key(&sources);
488            assert!(result.is_ok());
489            assert_eq!(result.unwrap(), "test-deepseek-key");
490        });
491    }
492
493    #[test]
494    fn test_get_gemini_api_key_from_config() {
495        with_overrides(
496            &[
497                ("TEST_GEMINI_CONFIG_KEY", None),
498                ("GOOGLE_API_KEY", None),
499                ("GEMINI_API_KEY", None),
500            ],
501            || {
502                let sources = ApiKeySources {
503                    gemini_env: "TEST_GEMINI_CONFIG_KEY".to_string(),
504                    gemini_config: Some("config-gemini-key".to_string()),
505                    ..Default::default()
506                };
507
508                let result = get_gemini_api_key(&sources);
509                assert!(result.is_ok());
510                assert_eq!(result.unwrap(), "config-gemini-key");
511            },
512        );
513    }
514
515    #[test]
516    fn test_get_api_key_with_fallback_prefers_env() {
517        with_override("TEST_FALLBACK_KEY", Some("env-key"), || {
518            let sources = ApiKeySources {
519                openai_env: "TEST_FALLBACK_KEY".to_string(),
520                openai_config: Some("config-key".to_string()),
521                ..Default::default()
522            };
523
524            let result = get_openai_api_key(&sources);
525            assert!(result.is_ok());
526            assert_eq!(result.unwrap(), "env-key"); // Should prefer env var
527        });
528    }
529
530    #[test]
531    fn test_get_api_key_fallback_to_config() {
532        let sources = ApiKeySources {
533            openai_env: "NONEXISTENT_ENV_VAR".to_string(),
534            openai_config: Some("config-key".to_string()),
535            ..Default::default()
536        };
537
538        let result = get_openai_api_key(&sources);
539        assert!(result.is_ok());
540        assert_eq!(result.unwrap(), "config-key");
541    }
542
543    #[test]
544    fn test_get_api_key_error_when_not_found() {
545        let sources = ApiKeySources {
546            openai_env: "NONEXISTENT_ENV_VAR".to_string(),
547            ..Default::default()
548        };
549
550        let result = get_openai_api_key(&sources);
551        assert!(result.is_err());
552    }
553
554    #[test]
555    fn test_get_ollama_api_key_missing_sources() {
556        let sources = ApiKeySources {
557            ollama_env: "NONEXISTENT_OLLAMA_ENV".to_string(),
558            ..Default::default()
559        };
560
561        let result = get_ollama_api_key(&sources).expect("Ollama key retrieval should succeed");
562        assert!(result.is_empty());
563    }
564
565    #[test]
566    fn test_get_ollama_api_key_from_env() {
567        with_override("TEST_OLLAMA_KEY", Some("test-ollama-key"), || {
568            let sources = ApiKeySources {
569                ollama_env: "TEST_OLLAMA_KEY".to_string(),
570                ..Default::default()
571            };
572
573            let result = get_ollama_api_key(&sources);
574            assert!(result.is_ok());
575            assert_eq!(result.unwrap(), "test-ollama-key");
576        });
577    }
578
579    #[test]
580    fn test_get_lmstudio_api_key_missing_sources() {
581        let sources = ApiKeySources {
582            lmstudio_env: "NONEXISTENT_LMSTUDIO_ENV".to_string(),
583            ..Default::default()
584        };
585
586        let result =
587            get_lmstudio_api_key(&sources).expect("LM Studio key retrieval should succeed");
588        assert!(result.is_empty());
589    }
590
591    #[test]
592    fn test_get_lmstudio_api_key_from_env() {
593        with_override("TEST_LMSTUDIO_KEY", Some("test-lmstudio-key"), || {
594            let sources = ApiKeySources {
595                lmstudio_env: "TEST_LMSTUDIO_KEY".to_string(),
596                ..Default::default()
597            };
598
599            let result = get_lmstudio_api_key(&sources);
600            assert!(result.is_ok());
601            assert_eq!(result.unwrap(), "test-lmstudio-key");
602        });
603    }
604
605    #[test]
606    fn test_get_api_key_ollama_provider() {
607        with_override(
608            "TEST_OLLAMA_PROVIDER_KEY",
609            Some("test-ollama-env-key"),
610            || {
611                let sources = ApiKeySources {
612                    ollama_env: "TEST_OLLAMA_PROVIDER_KEY".to_string(),
613                    ..Default::default()
614                };
615                let result = get_api_key("ollama", &sources);
616                assert!(result.is_ok());
617                assert_eq!(result.unwrap(), "test-ollama-env-key");
618            },
619        );
620    }
621
622    #[test]
623    fn test_get_api_key_lmstudio_provider() {
624        with_override(
625            "TEST_LMSTUDIO_PROVIDER_KEY",
626            Some("test-lmstudio-env-key"),
627            || {
628                let sources = ApiKeySources {
629                    lmstudio_env: "TEST_LMSTUDIO_PROVIDER_KEY".to_string(),
630                    ..Default::default()
631                };
632                let result = get_api_key("lmstudio", &sources);
633                assert!(result.is_ok());
634                assert_eq!(result.unwrap(), "test-lmstudio-env-key");
635            },
636        );
637    }
638
639    #[test]
640    fn api_key_env_var_uses_provider_defaults() {
641        assert_eq!(api_key_env_var("minimax"), "MINIMAX_API_KEY");
642        assert_eq!(api_key_env_var("huggingface"), "HF_TOKEN");
643    }
644
645    #[test]
646    fn resolve_api_key_env_uses_provider_default_for_placeholder() {
647        assert_eq!(
648            resolve_api_key_env("minimax", defaults::DEFAULT_API_KEY_ENV),
649            "MINIMAX_API_KEY"
650        );
651    }
652
653    #[test]
654    fn resolve_api_key_env_preserves_explicit_override() {
655        assert_eq!(
656            resolve_api_key_env("openai", "CUSTOM_OPENAI_KEY"),
657            "CUSTOM_OPENAI_KEY"
658        );
659    }
660}
vtcode_config/api_keys.rs

vtcode_config/
api_keys.rs
