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

vtcode_config/
api_keys.rs
