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///
18/// Retained for backward compatibility. New code should use [`get_api_key`] directly —
19/// the struct is no longer consumed by the key resolution logic.
20#[derive(Debug, Clone, Default)]
21pub struct ApiKeySources {
22    pub gemini_env: String,
23    pub anthropic_env: String,
24    pub openai_env: String,
25    pub openrouter_env: String,
26    pub deepseek_env: String,
27    pub zai_env: String,
28    pub ollama_env: String,
29    pub lmstudio_env: String,
30    pub gemini_config: Option<String>,
31    pub anthropic_config: Option<String>,
32    pub openai_config: Option<String>,
33    pub openrouter_config: Option<String>,
34    pub deepseek_config: Option<String>,
35    pub zai_config: Option<String>,
36    pub ollama_config: Option<String>,
37    pub lmstudio_config: Option<String>,
38}
39
40pub fn api_key_env_var(provider: &str) -> String {
41    let trimmed = provider.trim();
42    if trimmed.is_empty() {
43        return defaults::DEFAULT_API_KEY_ENV.to_owned();
44    }
45
46    if trimmed.eq_ignore_ascii_case("codex") {
47        return String::new();
48    }
49
50    if let Ok(resolved) = Provider::from_str(trimmed)
51        && resolved.uses_managed_auth()
52    {
53        return String::new();
54    }
55
56    Provider::from_str(trimmed)
57        .map(|resolved| resolved.default_api_key_env().to_owned())
58        .unwrap_or_else(|_| format!("{}_API_KEY", trimmed.to_ascii_uppercase()))
59}
60
61pub fn resolve_api_key_env(provider: &str, configured_env: &str) -> String {
62    let trimmed = configured_env.trim();
63    if trimmed.is_empty() || trimmed.eq_ignore_ascii_case(defaults::DEFAULT_API_KEY_ENV) {
64        api_key_env_var(provider)
65    } else {
66        trimmed.to_owned()
67    }
68}
69
70#[cfg(test)]
71mod test_env_overrides {
72    use hashbrown::HashMap;
73    use std::sync::{LazyLock, Mutex};
74
75    static OVERRIDES: LazyLock<Mutex<HashMap<String, Option<String>>>> =
76        LazyLock::new(|| Mutex::new(HashMap::new()));
77
78    pub(super) fn get(key: &str) -> Option<Option<String>> {
79        OVERRIDES.lock().ok().and_then(|map| map.get(key).cloned())
80    }
81
82    pub(super) fn set(key: &str, value: Option<&str>) {
83        if let Ok(mut map) = OVERRIDES.lock() {
84            map.insert(key.to_string(), value.map(ToString::to_string));
85        }
86    }
87
88    pub(super) fn restore(key: &str, previous: Option<Option<String>>) {
89        if let Ok(mut map) = OVERRIDES.lock() {
90            match previous {
91                Some(value) => {
92                    map.insert(key.to_string(), value);
93                }
94                None => {
95                    map.remove(key);
96                }
97            }
98        }
99    }
100}
101
102fn read_env_var(key: &str) -> Option<String> {
103    #[cfg(test)]
104    if let Some(override_value) = test_env_overrides::get(key) {
105        return override_value;
106    }
107
108    env::var(key).ok()
109}
110
111/// Load environment variables from .env file
112///
113/// This function attempts to load environment variables from a .env file
114/// in the current directory. It logs a warning if the file exists but cannot
115/// be loaded, but doesn't fail if the file doesn't exist.
116pub fn load_dotenv() -> Result<()> {
117    match dotenvy::dotenv() {
118        Ok(path) => {
119            // Only print in verbose mode to avoid polluting stdout/stderr in scripts
120            if read_env_var("VTCODE_VERBOSE").is_some() || read_env_var("RUST_LOG").is_some() {
121                tracing::info!("Loaded environment variables from: {}", path.display());
122            }
123            Ok(())
124        }
125        Err(dotenvy::Error::Io(e)) if e.kind() == std::io::ErrorKind::NotFound => {
126            // .env file doesn't exist, which is fine
127            Ok(())
128        }
129        Err(e) => {
130            tracing::warn!("Failed to load .env file: {}", e);
131            Ok(())
132        }
133    }
134}
135
136/// Get API key for a specific provider.
137///
138/// Resolution order:
139/// 1. Environment variable inferred from the provider name (e.g. `POOLSIDE_API_KEY`)
140/// 2. Provider-specific fallbacks (OAuth tokens, alternate env vars, etc.)
141/// 3. OS keyring / encrypted file storage
142///
143/// Adding a new built-in provider only requires:
144/// - A `Provider` variant with `default_api_key_env()` returning the env var name
145/// - (Optional) a match arm here only if the provider needs special fallback logic
146pub fn get_api_key(provider: &str, _sources: &ApiKeySources) -> Result<String> {
147    let normalized_provider = provider.to_lowercase();
148    let inferred_env = api_key_env_var(&normalized_provider);
149
150    // Generic path: read the inferred env var for any provider.
151    if let Some(key) = read_env_var(&inferred_env)
152        && !key.is_empty()
153    {
154        return Ok(key);
155    }
156
157    // Provider-specific fallback logic. Most providers are handled by the generic
158    // env-var lookup above. Only providers with special behavior (alternate env vars,
159    // OAuth tokens, optional keys, or managed-auth error messages) need a match arm.
160    let provider_result = match normalized_provider.as_str() {
161        // Gemini falls back to GOOGLE_API_KEY for backward compatibility
162        "gemini" => {
163            if let Some(key) = read_env_var("GOOGLE_API_KEY").filter(|k| !k.is_empty()) {
164                return Ok(key);
165            }
166            Err(anyhow::anyhow!("GEMINI_API_KEY or GOOGLE_API_KEY not set"))
167        }
168        // OpenRouter tries OAuth token from encrypted storage first
169        "openrouter" => {
170            if let Ok(Some(token)) = crate::auth::load_oauth_token() {
171                tracing::debug!("Using OAuth token for OpenRouter authentication");
172                return Ok(token.api_key);
173            }
174            Err(anyhow::anyhow!("OPENROUTER_API_KEY not set"))
175        }
176        // Qwen has an alternate env var name
177        "qwen" => {
178            if let Some(key) = read_env_var("DASHSCOPE_API_KEY").filter(|k| !k.is_empty()) {
179                return Ok(key);
180            }
181            Err(anyhow::anyhow!("QWEN_API_KEY or DASHSCOPE_API_KEY not set"))
182        }
183        // Ollama and LM Studio allow empty keys (local deployment)
184        "ollama" | "lmstudio" | "llamacpp" | "llama.cpp" | "llama-cpp" => Ok(String::new()),
185        // Managed-auth providers show a specific error message
186        "copilot" => Err(anyhow::anyhow!(
187            "GitHub Copilot authentication is managed by the official `copilot` CLI. Run `vtcode login copilot`."
188        )),
189        "codex" => Err(anyhow::anyhow!(
190            "Codex authentication is managed by the official `codex app-server`. Run `vtcode login codex`."
191        )),
192        // All other providers: env var was already checked above, nothing more to do
193        _ => {
194            return Err(anyhow::anyhow!(
195                "{} API key not found. Set {} environment variable or add to .env file.",
196                normalized_provider,
197                inferred_env,
198            ));
199        }
200    };
201
202    if provider_result.is_ok() {
203        return provider_result;
204    }
205
206    // Try secure storage (keyring) only after env/config lookup fails.
207    if let Ok(Some(key)) = get_custom_api_key_from_secure_storage(&normalized_provider) {
208        return Ok(key);
209    }
210
211    provider_result
212}
213
214/// Get a custom API key from secure storage.
215///
216/// This function retrieves API keys that were stored securely via the model picker
217/// or interactive configuration flows. When the OS keyring is unavailable, the
218/// auth layer falls back to encrypted file storage automatically.
219///
220/// # Arguments
221/// * `provider` - The provider name
222///
223/// # Returns
224/// * `Ok(Some(String))` - The API key if found in secure storage
225/// * `Ok(None)` - If no key is stored for this provider
226/// * `Err` - If there was an error accessing secure storage
227fn get_custom_api_key_from_secure_storage(provider: &str) -> Result<Option<String>> {
228    let storage = CustomApiKeyStorage::new(provider);
229    // The auth layer handles keyring-to-file fallback internally.
230    let mode = crate::auth::AuthCredentialsStoreMode::default();
231    storage.load(mode)
232}
233
234#[cfg(test)]
235mod tests {
236    use super::*;
237
238    struct EnvOverrideGuard {
239        key: &'static str,
240        previous: Option<Option<String>>,
241    }
242
243    impl EnvOverrideGuard {
244        fn set(key: &'static str, value: Option<&str>) -> Self {
245            let previous = test_env_overrides::get(key);
246            test_env_overrides::set(key, value);
247            Self { key, previous }
248        }
249    }
250
251    impl Drop for EnvOverrideGuard {
252        fn drop(&mut self) {
253            test_env_overrides::restore(self.key, self.previous.clone());
254        }
255    }
256
257    fn with_override<F>(key: &'static str, value: Option<&str>, f: F)
258    where
259        F: FnOnce(),
260    {
261        let _guard = EnvOverrideGuard::set(key, value);
262        f();
263    }
264
265    fn with_overrides<F>(overrides: &[(&'static str, Option<&str>)], f: F)
266    where
267        F: FnOnce(),
268    {
269        let _guards: Vec<_> = overrides
270            .iter()
271            .map(|(key, value)| EnvOverrideGuard::set(key, *value))
272            .collect();
273        f();
274    }
275
276    fn default_sources() -> ApiKeySources {
277        ApiKeySources::default()
278    }
279
280    #[test]
281    fn gemini_reads_env_var() {
282        with_override("GEMINI_API_KEY", Some("test-gemini-key"), || {
283            let result = get_api_key("gemini", &default_sources());
284            assert_eq!(result.unwrap(), "test-gemini-key");
285        });
286    }
287
288    #[test]
289    fn gemini_falls_back_to_google_api_key() {
290        // Clear both GEMINI_API_KEY and set GOOGLE_API_KEY to verify fallback
291        with_overrides(
292            &[
293                ("GEMINI_API_KEY", Some("gemini-primary")),
294                ("GOOGLE_API_KEY", Some("google-fallback")),
295            ],
296            || {
297                // With GEMINI_API_KEY set, it should be preferred
298                let result = get_api_key("gemini", &default_sources());
299                assert_eq!(result.unwrap(), "gemini-primary");
300            },
301        );
302        with_overrides(
303            &[
304                ("GEMINI_API_KEY", None),
305                ("GOOGLE_API_KEY", Some("google-fallback")),
306            ],
307            || {
308                // Without GEMINI_API_KEY, it should fall back to GOOGLE_API_KEY
309                let result = get_api_key("gemini", &default_sources());
310                assert_eq!(result.unwrap(), "google-fallback");
311            },
312        );
313    }
314
315    #[test]
316    fn anthropic_reads_env_var() {
317        with_override("ANTHROPIC_API_KEY", Some("test-anthropic-key"), || {
318            let result = get_api_key("anthropic", &default_sources());
319            assert_eq!(result.unwrap(), "test-anthropic-key");
320        });
321    }
322
323    #[test]
324    fn openai_reads_env_var() {
325        with_override("OPENAI_API_KEY", Some("test-openai-key"), || {
326            let result = get_api_key("openai", &default_sources());
327            assert_eq!(result.unwrap(), "test-openai-key");
328        });
329    }
330
331    #[test]
332    fn deepseek_reads_env_var() {
333        with_override("DEEPSEEK_API_KEY", Some("test-deepseek-key"), || {
334            let result = get_api_key("deepseek", &default_sources());
335            assert_eq!(result.unwrap(), "test-deepseek-key");
336        });
337    }
338
339    #[test]
340    fn qwen_falls_back_to_dashscope() {
341        with_overrides(
342            &[
343                ("QWEN_API_KEY", None),
344                ("DASHSCOPE_API_KEY", Some("dashscope-key")),
345            ],
346            || {
347                let result = get_api_key("qwen", &default_sources());
348                assert_eq!(result.unwrap(), "dashscope-key");
349            },
350        );
351    }
352
353    #[test]
354    fn ollama_allows_empty_key() {
355        with_override("OLLAMA_API_KEY", None, || {
356            let result = get_api_key("ollama", &default_sources());
357            assert!(result.is_ok());
358            assert!(result.unwrap().is_empty());
359        });
360    }
361
362    #[test]
363    fn lmstudio_allows_empty_key() {
364        with_override("LMSTUDIO_API_KEY", None, || {
365            let result = get_api_key("lmstudio", &default_sources());
366            assert!(result.is_ok());
367            assert!(result.unwrap().is_empty());
368        });
369    }
370
371    #[test]
372    fn ollama_reads_env_var_when_set() {
373        with_override("OLLAMA_API_KEY", Some("test-ollama-key"), || {
374            let result = get_api_key("ollama", &default_sources());
375            assert_eq!(result.unwrap(), "test-ollama-key");
376        });
377    }
378
379    #[test]
380    fn copilot_returns_managed_auth_error() {
381        let result = get_api_key("copilot", &default_sources());
382        assert!(result.is_err());
383        assert!(result.unwrap_err().to_string().contains("copilot"));
384    }
385
386    #[test]
387    fn codex_returns_managed_auth_error() {
388        let result = get_api_key("codex", &default_sources());
389        assert!(result.is_err());
390        assert!(result.unwrap_err().to_string().contains("codex"));
391    }
392
393    #[test]
394    fn unknown_provider_returns_error_with_env_hint() {
395        let result = get_api_key("someunknown", &default_sources());
396        assert!(result.is_err());
397        let msg = result.unwrap_err().to_string();
398        assert!(msg.contains("SOMEUNKNOWN_API_KEY"));
399    }
400
401    #[test]
402    fn poolside_reads_env_var() {
403        with_override("POOLSIDE_API_KEY", Some("test-poolside-key"), || {
404            let result = get_api_key("poolside", &default_sources());
405            assert_eq!(result.unwrap(), "test-poolside-key");
406        });
407    }
408
409    #[test]
410    fn poolside_returns_error_when_missing() {
411        with_override("POOLSIDE_API_KEY", None, || {
412            let result = get_api_key("poolside", &default_sources());
413            assert!(result.is_err());
414            assert!(result.unwrap_err().to_string().contains("POOLSIDE_API_KEY"));
415        });
416    }
417
418    #[test]
419    fn api_key_env_var_uses_provider_defaults() {
420        assert_eq!(api_key_env_var("codex"), "");
421        assert_eq!(api_key_env_var("minimax"), "MINIMAX_API_KEY");
422        assert_eq!(api_key_env_var("huggingface"), "HF_TOKEN");
423        assert_eq!(api_key_env_var("poolside"), "POOLSIDE_API_KEY");
424    }
425
426    #[test]
427    fn resolve_api_key_env_uses_provider_default_for_placeholder() {
428        assert_eq!(
429            resolve_api_key_env("minimax", defaults::DEFAULT_API_KEY_ENV),
430            "MINIMAX_API_KEY"
431        );
432    }
433
434    #[test]
435    fn resolve_api_key_env_preserves_explicit_override() {
436        assert_eq!(
437            resolve_api_key_env("openai", "CUSTOM_OPENAI_KEY"),
438            "CUSTOM_OPENAI_KEY"
439        );
440    }
441}
vtcode_config/api_keys.rs

vtcode_config/
api_keys.rs
