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`."
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        _ => Err(anyhow::anyhow!("Unsupported provider: {}", provider)),
234    }
235}
236
237/// Get a custom API key from secure storage.
238///
239/// This function retrieves API keys that were stored securely via the model picker
240/// or interactive configuration flows. When the OS keyring is unavailable, the
241/// auth layer falls back to encrypted file storage automatically.
242///
243/// # Arguments
244/// * `provider` - The provider name
245///
246/// # Returns
247/// * `Ok(Some(String))` - The API key if found in secure storage
248/// * `Ok(None)` - If no key is stored for this provider
249/// * `Err` - If there was an error accessing secure storage
250fn get_custom_api_key_from_secure_storage(provider: &str) -> Result<Option<String>> {
251    let storage = CustomApiKeyStorage::new(provider);
252    // The auth layer handles keyring-to-file fallback internally.
253    let mode = crate::auth::AuthCredentialsStoreMode::default();
254    storage.load(mode)
255}
256
257/// Get API key for a specific environment variable with fallback
258fn get_api_key_with_fallback(
259    env_var: &str,
260    config_value: Option<&String>,
261    provider_name: &str,
262) -> Result<String> {
263    // First try environment variable (most secure)
264    if let Some(key) = read_env_var(env_var)
265        && !key.is_empty()
266    {
267        return Ok(key);
268    }
269
270    // Then try configuration file value
271    if let Some(key) = config_value
272        && !key.is_empty()
273    {
274        return Ok(key.clone());
275    }
276
277    // If neither worked, return an error
278    Err(anyhow::anyhow!(
279        "No API key found for {} provider. Set {} environment variable (or add to .env file) or configure in vtcode.toml",
280        provider_name,
281        env_var
282    ))
283}
284
285fn get_optional_api_key_with_fallback(env_var: &str, config_value: Option<&String>) -> String {
286    if let Some(key) = read_env_var(env_var)
287        && !key.is_empty()
288    {
289        return key;
290    }
291
292    if let Some(key) = config_value
293        && !key.is_empty()
294    {
295        return key.clone();
296    }
297
298    String::new()
299}
300
301/// Get Gemini API key with secure fallback
302fn get_gemini_api_key(sources: &ApiKeySources) -> Result<String> {
303    // Try primary Gemini environment variable
304    if let Some(key) = read_env_var(&sources.gemini_env)
305        && !key.is_empty()
306    {
307        return Ok(key);
308    }
309
310    // Try Google API key as fallback (for backward compatibility)
311    if let Some(key) = read_env_var("GOOGLE_API_KEY")
312        && !key.is_empty()
313    {
314        return Ok(key);
315    }
316
317    // Try configuration file value
318    if let Some(key) = &sources.gemini_config
319        && !key.is_empty()
320    {
321        return Ok(key.clone());
322    }
323
324    // If nothing worked, return an error
325    Err(anyhow::anyhow!(
326        "No API key found for Gemini provider. Set {} or GOOGLE_API_KEY environment variable (or add to .env file) or configure in vtcode.toml",
327        sources.gemini_env
328    ))
329}
330
331/// Get Anthropic API key with secure fallback
332fn get_anthropic_api_key(sources: &ApiKeySources) -> Result<String> {
333    get_api_key_with_fallback(
334        &sources.anthropic_env,
335        sources.anthropic_config.as_ref(),
336        "Anthropic",
337    )
338}
339
340/// Get OpenAI API key with secure fallback
341fn get_openai_api_key(sources: &ApiKeySources) -> Result<String> {
342    get_api_key_with_fallback(
343        &sources.openai_env,
344        sources.openai_config.as_ref(),
345        "OpenAI",
346    )
347}
348
349/// Get OpenRouter API key with secure fallback
350///
351/// This function checks for credentials in the following order:
352/// 1. OAuth token from encrypted storage (if OAuth is enabled)
353/// 2. Environment variable (OPENROUTER_API_KEY)
354/// 3. Configuration file value
355fn get_openrouter_api_key(sources: &ApiKeySources) -> Result<String> {
356    // First, try to load OAuth token from encrypted storage
357    if let Ok(Some(token)) = crate::auth::load_oauth_token() {
358        tracing::debug!("Using OAuth token for OpenRouter authentication");
359        return Ok(token.api_key);
360    }
361
362    // Fall back to standard API key retrieval
363    get_api_key_with_fallback(
364        &sources.openrouter_env,
365        sources.openrouter_config.as_ref(),
366        "OpenRouter",
367    )
368}
369
370/// Get DeepSeek API key with secure fallback
371fn get_deepseek_api_key(sources: &ApiKeySources) -> Result<String> {
372    get_api_key_with_fallback(
373        &sources.deepseek_env,
374        sources.deepseek_config.as_ref(),
375        "DeepSeek",
376    )
377}
378
379/// Get Z.AI API key with secure fallback
380fn get_zai_api_key(sources: &ApiKeySources) -> Result<String> {
381    get_api_key_with_fallback(&sources.zai_env, sources.zai_config.as_ref(), "Z.AI")
382}
383
384/// Get Ollama API key with secure fallback
385fn get_ollama_api_key(sources: &ApiKeySources) -> Result<String> {
386    // For Ollama we allow running without credentials when connecting to a local deployment.
387    // Cloud variants still rely on environment/config values when present.
388    Ok(get_optional_api_key_with_fallback(
389        &sources.ollama_env,
390        sources.ollama_config.as_ref(),
391    ))
392}
393
394/// Get LM Studio API key with secure fallback
395fn get_lmstudio_api_key(sources: &ApiKeySources) -> Result<String> {
396    Ok(get_optional_api_key_with_fallback(
397        &sources.lmstudio_env,
398        sources.lmstudio_config.as_ref(),
399    ))
400}
401
402#[cfg(test)]
403mod tests {
404    use super::*;
405
406    struct EnvOverrideGuard {
407        key: &'static str,
408        previous: Option<Option<String>>,
409    }
410
411    impl EnvOverrideGuard {
412        fn set(key: &'static str, value: Option<&str>) -> Self {
413            let previous = test_env_overrides::get(key);
414            test_env_overrides::set(key, value);
415            Self { key, previous }
416        }
417    }
418
419    impl Drop for EnvOverrideGuard {
420        fn drop(&mut self) {
421            test_env_overrides::restore(self.key, self.previous.clone());
422        }
423    }
424
425    fn with_override<F>(key: &'static str, value: Option<&str>, f: F)
426    where
427        F: FnOnce(),
428    {
429        let _guard = EnvOverrideGuard::set(key, value);
430        f();
431    }
432
433    fn with_overrides<F>(overrides: &[(&'static str, Option<&str>)], f: F)
434    where
435        F: FnOnce(),
436    {
437        let _guards: Vec<_> = overrides
438            .iter()
439            .map(|(key, value)| EnvOverrideGuard::set(key, *value))
440            .collect();
441        f();
442    }
443
444    #[test]
445    fn test_get_gemini_api_key_from_env() {
446        with_override("TEST_GEMINI_KEY", Some("test-gemini-key"), || {
447            let sources = ApiKeySources {
448                gemini_env: "TEST_GEMINI_KEY".to_string(),
449                ..Default::default()
450            };
451
452            let result = get_gemini_api_key(&sources);
453            assert!(result.is_ok());
454            assert_eq!(result.unwrap(), "test-gemini-key");
455        });
456    }
457
458    #[test]
459    fn test_get_anthropic_api_key_from_env() {
460        with_override("TEST_ANTHROPIC_KEY", Some("test-anthropic-key"), || {
461            let sources = ApiKeySources {
462                anthropic_env: "TEST_ANTHROPIC_KEY".to_string(),
463                ..Default::default()
464            };
465
466            let result = get_anthropic_api_key(&sources);
467            assert!(result.is_ok());
468            assert_eq!(result.unwrap(), "test-anthropic-key");
469        });
470    }
471
472    #[test]
473    fn test_get_openai_api_key_from_env() {
474        with_override("TEST_OPENAI_KEY", Some("test-openai-key"), || {
475            let sources = ApiKeySources {
476                openai_env: "TEST_OPENAI_KEY".to_string(),
477                ..Default::default()
478            };
479
480            let result = get_openai_api_key(&sources);
481            assert!(result.is_ok());
482            assert_eq!(result.unwrap(), "test-openai-key");
483        });
484    }
485
486    #[test]
487    fn test_get_deepseek_api_key_from_env() {
488        with_override("TEST_DEEPSEEK_KEY", Some("test-deepseek-key"), || {
489            let sources = ApiKeySources {
490                deepseek_env: "TEST_DEEPSEEK_KEY".to_string(),
491                ..Default::default()
492            };
493
494            let result = get_deepseek_api_key(&sources);
495            assert!(result.is_ok());
496            assert_eq!(result.unwrap(), "test-deepseek-key");
497        });
498    }
499
500    #[test]
501    fn test_get_gemini_api_key_from_config() {
502        with_overrides(
503            &[
504                ("TEST_GEMINI_CONFIG_KEY", None),
505                ("GOOGLE_API_KEY", None),
506                ("GEMINI_API_KEY", None),
507            ],
508            || {
509                let sources = ApiKeySources {
510                    gemini_env: "TEST_GEMINI_CONFIG_KEY".to_string(),
511                    gemini_config: Some("config-gemini-key".to_string()),
512                    ..Default::default()
513                };
514
515                let result = get_gemini_api_key(&sources);
516                assert!(result.is_ok());
517                assert_eq!(result.unwrap(), "config-gemini-key");
518            },
519        );
520    }
521
522    #[test]
523    fn test_get_api_key_with_fallback_prefers_env() {
524        with_override("TEST_FALLBACK_KEY", Some("env-key"), || {
525            let sources = ApiKeySources {
526                openai_env: "TEST_FALLBACK_KEY".to_string(),
527                openai_config: Some("config-key".to_string()),
528                ..Default::default()
529            };
530
531            let result = get_openai_api_key(&sources);
532            assert!(result.is_ok());
533            assert_eq!(result.unwrap(), "env-key"); // Should prefer env var
534        });
535    }
536
537    #[test]
538    fn test_get_api_key_fallback_to_config() {
539        let sources = ApiKeySources {
540            openai_env: "NONEXISTENT_ENV_VAR".to_string(),
541            openai_config: Some("config-key".to_string()),
542            ..Default::default()
543        };
544
545        let result = get_openai_api_key(&sources);
546        assert!(result.is_ok());
547        assert_eq!(result.unwrap(), "config-key");
548    }
549
550    #[test]
551    fn test_get_api_key_error_when_not_found() {
552        let sources = ApiKeySources {
553            openai_env: "NONEXISTENT_ENV_VAR".to_string(),
554            ..Default::default()
555        };
556
557        let result = get_openai_api_key(&sources);
558        assert!(result.is_err());
559    }
560
561    #[test]
562    fn test_get_ollama_api_key_missing_sources() {
563        let sources = ApiKeySources {
564            ollama_env: "NONEXISTENT_OLLAMA_ENV".to_string(),
565            ..Default::default()
566        };
567
568        let result = get_ollama_api_key(&sources).expect("Ollama key retrieval should succeed");
569        assert!(result.is_empty());
570    }
571
572    #[test]
573    fn test_get_ollama_api_key_from_env() {
574        with_override("TEST_OLLAMA_KEY", Some("test-ollama-key"), || {
575            let sources = ApiKeySources {
576                ollama_env: "TEST_OLLAMA_KEY".to_string(),
577                ..Default::default()
578            };
579
580            let result = get_ollama_api_key(&sources);
581            assert!(result.is_ok());
582            assert_eq!(result.unwrap(), "test-ollama-key");
583        });
584    }
585
586    #[test]
587    fn test_get_lmstudio_api_key_missing_sources() {
588        let sources = ApiKeySources {
589            lmstudio_env: "NONEXISTENT_LMSTUDIO_ENV".to_string(),
590            ..Default::default()
591        };
592
593        let result =
594            get_lmstudio_api_key(&sources).expect("LM Studio key retrieval should succeed");
595        assert!(result.is_empty());
596    }
597
598    #[test]
599    fn test_get_lmstudio_api_key_from_env() {
600        with_override("TEST_LMSTUDIO_KEY", Some("test-lmstudio-key"), || {
601            let sources = ApiKeySources {
602                lmstudio_env: "TEST_LMSTUDIO_KEY".to_string(),
603                ..Default::default()
604            };
605
606            let result = get_lmstudio_api_key(&sources);
607            assert!(result.is_ok());
608            assert_eq!(result.unwrap(), "test-lmstudio-key");
609        });
610    }
611
612    #[test]
613    fn test_get_api_key_ollama_provider() {
614        with_override(
615            "TEST_OLLAMA_PROVIDER_KEY",
616            Some("test-ollama-env-key"),
617            || {
618                let sources = ApiKeySources {
619                    ollama_env: "TEST_OLLAMA_PROVIDER_KEY".to_string(),
620                    ..Default::default()
621                };
622                let result = get_api_key("ollama", &sources);
623                assert!(result.is_ok());
624                assert_eq!(result.unwrap(), "test-ollama-env-key");
625            },
626        );
627    }
628
629    #[test]
630    fn test_get_api_key_lmstudio_provider() {
631        with_override(
632            "TEST_LMSTUDIO_PROVIDER_KEY",
633            Some("test-lmstudio-env-key"),
634            || {
635                let sources = ApiKeySources {
636                    lmstudio_env: "TEST_LMSTUDIO_PROVIDER_KEY".to_string(),
637                    ..Default::default()
638                };
639                let result = get_api_key("lmstudio", &sources);
640                assert!(result.is_ok());
641                assert_eq!(result.unwrap(), "test-lmstudio-env-key");
642            },
643        );
644    }
645
646    #[test]
647    fn api_key_env_var_uses_provider_defaults() {
648        assert_eq!(api_key_env_var("codex"), "");
649        assert_eq!(api_key_env_var("minimax"), "MINIMAX_API_KEY");
650        assert_eq!(api_key_env_var("huggingface"), "HF_TOKEN");
651    }
652
653    #[test]
654    fn resolve_api_key_env_uses_provider_default_for_placeholder() {
655        assert_eq!(
656            resolve_api_key_env("minimax", defaults::DEFAULT_API_KEY_ENV),
657            "MINIMAX_API_KEY"
658        );
659    }
660
661    #[test]
662    fn resolve_api_key_env_preserves_explicit_override() {
663        assert_eq!(
664            resolve_api_key_env("openai", "CUSTOM_OPENAI_KEY"),
665            "CUSTOM_OPENAI_KEY"
666        );
667    }
668}
vtcode_config/api_keys.rs

vtcode_config/
api_keys.rs
