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

vtcode_config/
api_keys.rs
