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

vtcode_config/
api_keys.rs
