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

vtcode_config/
api_keys.rs
