vtcode_core/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;
10
11/// API key sources for different providers
12#[derive(Debug, Clone)]
13pub struct ApiKeySources {
14    /// Gemini API key environment variable name
15    pub gemini_env: String,
16    /// Anthropic API key environment variable name
17    pub anthropic_env: String,
18    /// OpenAI API key environment variable name
19    pub openai_env: String,
20    /// Gemini API key from configuration file
21    pub gemini_config: Option<String>,
22    /// Anthropic API key from configuration file
23    pub anthropic_config: Option<String>,
24    /// OpenAI API key from configuration file
25    pub openai_config: Option<String>,
26}
27
28impl Default for ApiKeySources {
29    fn default() -> Self {
30        Self {
31            gemini_env: "GEMINI_API_KEY".to_string(),
32            anthropic_env: "ANTHROPIC_API_KEY".to_string(),
33            openai_env: "OPENAI_API_KEY".to_string(),
34            gemini_config: None,
35            anthropic_config: None,
36            openai_config: None,
37        }
38    }
39}
40
41/// Load environment variables from .env file
42///
43/// This function attempts to load environment variables from a .env file
44/// in the current directory, ignoring errors if the file doesn't exist.
45pub fn load_dotenv() -> Result<()> {
46    // Try to load .env file, but don't fail if it doesn't exist
47    let _ = dotenvy::dotenv();
48    Ok(())
49}
50
51/// Get API key for a specific provider with secure fallback mechanism
52///
53/// This function implements a secure retrieval mechanism that:
54/// 1. First checks environment variables (highest priority for security)
55/// 2. Then checks .env file values
56/// 3. Falls back to configuration file values if neither above is set
57/// 4. Supports all major providers: Gemini, Anthropic, and OpenAI
58///
59/// # Arguments
60///
61/// * `provider` - The provider name ("gemini", "anthropic", or "openai")
62/// * `sources` - Configuration for where to look for API keys
63///
64/// # Returns
65///
66/// * `Ok(String)` - The API key if found
67/// * `Err` - If no API key could be found for the provider
68pub fn get_api_key(provider: &str, sources: &ApiKeySources) -> Result<String> {
69    match provider.to_lowercase().as_str() {
70        "gemini" => get_gemini_api_key(sources),
71        "anthropic" => get_anthropic_api_key(sources),
72        "openai" => get_openai_api_key(sources),
73        _ => Err(anyhow::anyhow!("Unsupported provider: {}", provider)),
74    }
75}
76
77/// Get API key for a specific environment variable with fallback
78fn get_api_key_with_fallback(
79    env_var: &str,
80    config_value: Option<&String>,
81    provider_name: &str,
82) -> Result<String> {
83    // First try environment variable (most secure)
84    if let Ok(key) = env::var(env_var) {
85        if !key.is_empty() {
86            return Ok(key);
87        }
88    }
89
90    // Then try configuration file value
91    if let Some(key) = config_value {
92        if !key.is_empty() {
93            return Ok(key.clone());
94        }
95    }
96
97    // If neither worked, return an error
98    Err(anyhow::anyhow!(
99        "No API key found for {} provider. Set {} environment variable or configure in vtcode.toml",
100        provider_name,
101        env_var
102    ))
103}
104
105/// Get Gemini API key with secure fallback
106fn get_gemini_api_key(sources: &ApiKeySources) -> Result<String> {
107    // Try primary Gemini environment variable
108    if let Ok(key) = env::var(&sources.gemini_env) {
109        if !key.is_empty() {
110            return Ok(key);
111        }
112    }
113
114    // Try Google API key as fallback (for backward compatibility)
115    if let Ok(key) = env::var("GOOGLE_API_KEY") {
116        if !key.is_empty() {
117            return Ok(key);
118        }
119    }
120
121    // Try configuration file value
122    if let Some(key) = &sources.gemini_config {
123        if !key.is_empty() {
124            return Ok(key.clone());
125        }
126    }
127
128    // If nothing worked, return an error
129    Err(anyhow::anyhow!(
130        "No API key found for Gemini provider. Set {} or GOOGLE_API_KEY environment variable or configure in vtcode.toml",
131        sources.gemini_env
132    ))
133}
134
135/// Get Anthropic API key with secure fallback
136fn get_anthropic_api_key(sources: &ApiKeySources) -> Result<String> {
137    get_api_key_with_fallback(
138        &sources.anthropic_env,
139        sources.anthropic_config.as_ref(),
140        "Anthropic",
141    )
142}
143
144/// Get OpenAI API key with secure fallback
145fn get_openai_api_key(sources: &ApiKeySources) -> Result<String> {
146    get_api_key_with_fallback(
147        &sources.openai_env,
148        sources.openai_config.as_ref(),
149        "OpenAI",
150    )
151}
152
153#[cfg(test)]
154mod tests {
155    use super::*;
156    use std::env;
157
158    #[test]
159    fn test_get_gemini_api_key_from_env() {
160        // Set environment variable
161        unsafe {
162            env::set_var("TEST_GEMINI_KEY", "test-gemini-key");
163        }
164
165        let sources = ApiKeySources {
166            gemini_env: "TEST_GEMINI_KEY".to_string(),
167            ..Default::default()
168        };
169
170        let result = get_gemini_api_key(&sources);
171        assert!(result.is_ok());
172        assert_eq!(result.unwrap(), "test-gemini-key");
173
174        // Clean up
175        unsafe {
176            env::remove_var("TEST_GEMINI_KEY");
177        }
178    }
179
180    #[test]
181    fn test_get_anthropic_api_key_from_env() {
182        // Set environment variable
183        unsafe {
184            env::set_var("TEST_ANTHROPIC_KEY", "test-anthropic-key");
185        }
186
187        let sources = ApiKeySources {
188            anthropic_env: "TEST_ANTHROPIC_KEY".to_string(),
189            ..Default::default()
190        };
191
192        let result = get_anthropic_api_key(&sources);
193        assert!(result.is_ok());
194        assert_eq!(result.unwrap(), "test-anthropic-key");
195
196        // Clean up
197        unsafe {
198            env::remove_var("TEST_ANTHROPIC_KEY");
199        }
200    }
201
202    #[test]
203    fn test_get_openai_api_key_from_env() {
204        // Set environment variable
205        unsafe {
206            env::set_var("TEST_OPENAI_KEY", "test-openai-key");
207        }
208
209        let sources = ApiKeySources {
210            openai_env: "TEST_OPENAI_KEY".to_string(),
211            ..Default::default()
212        };
213
214        let result = get_openai_api_key(&sources);
215        assert!(result.is_ok());
216        assert_eq!(result.unwrap(), "test-openai-key");
217
218        // Clean up
219        unsafe {
220            env::remove_var("TEST_OPENAI_KEY");
221        }
222    }
223
224    #[test]
225    fn test_get_gemini_api_key_from_config() {
226        let sources = ApiKeySources {
227            gemini_config: Some("config-gemini-key".to_string()),
228            ..Default::default()
229        };
230
231        let result = get_gemini_api_key(&sources);
232        assert!(result.is_ok());
233        assert_eq!(result.unwrap(), "config-gemini-key");
234    }
235
236    #[test]
237    fn test_get_api_key_with_fallback_prefers_env() {
238        // Set environment variable
239        unsafe {
240            env::set_var("TEST_FALLBACK_KEY", "env-key");
241        }
242
243        let sources = ApiKeySources {
244            openai_env: "TEST_FALLBACK_KEY".to_string(),
245            openai_config: Some("config-key".to_string()),
246            ..Default::default()
247        };
248
249        let result = get_openai_api_key(&sources);
250        assert!(result.is_ok());
251        assert_eq!(result.unwrap(), "env-key"); // Should prefer env var
252
253        // Clean up
254        unsafe {
255            env::remove_var("TEST_FALLBACK_KEY");
256        }
257    }
258
259    #[test]
260    fn test_get_api_key_fallback_to_config() {
261        let sources = ApiKeySources {
262            openai_env: "NONEXISTENT_ENV_VAR".to_string(),
263            openai_config: Some("config-key".to_string()),
264            ..Default::default()
265        };
266
267        let result = get_openai_api_key(&sources);
268        assert!(result.is_ok());
269        assert_eq!(result.unwrap(), "config-key");
270    }
271
272    #[test]
273    fn test_get_api_key_error_when_not_found() {
274        let sources = ApiKeySources {
275            openai_env: "NONEXISTENT_ENV_VAR".to_string(),
276            ..Default::default()
277        };
278
279        let result = get_openai_api_key(&sources);
280        assert!(result.is_err());
281    }
282}
vtcode_core/config/api_keys.rs

vtcode_core/config/
api_keys.rs
