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

vtcode_config/
api_keys.rs
