vtcode_config/
workspace_env.rs

1use anyhow::{Context, Result, anyhow};
2use std::fs::{self, File};
3use std::io::{BufWriter, Write};
4use std::path::{Path, PathBuf};
5use tempfile::Builder;
6
7use vtcode_auth::AuthCredentialsStoreMode;
8use vtcode_commons::provider::Provider;
9
10use crate::api_keys::store_credential_with_mode;
11
12/// Returns the workspace `.env` file path.
13pub fn workspace_env_path(workspace: &Path) -> PathBuf {
14    workspace.join(".env")
15}
16
17/// Display-friendly representation of the workspace `.env` path.
18pub fn workspace_env_path_display(workspace: &Path) -> String {
19    workspace_env_path(workspace).display().to_string()
20}
21
22pub fn read_workspace_env_value(workspace: &Path, env_key: &str) -> Result<Option<String>> {
23    let env_path = workspace.join(".env");
24    let iter = match dotenvy::from_path_iter(&env_path) {
25        Ok(iter) => iter,
26        Err(dotenvy::Error::Io(err)) if err.kind() == std::io::ErrorKind::NotFound => {
27            return Ok(None);
28        }
29        Err(err) => {
30            return Err(anyhow!(err).context(format!("Failed to read {}", env_path.display())));
31        }
32    };
33
34    for item in iter {
35        let (key, value) = item
36            .map_err(|err: dotenvy::Error| anyhow!(err))
37            .with_context(|| format!("Failed to parse {}", env_path.display()))?;
38        if key == env_key {
39            if value.trim().is_empty() {
40                return Ok(None);
41            }
42            return Ok(Some(value));
43        }
44    }
45
46    Ok(None)
47}
48
49/// Read multiple env keys from workspace `.env` in a single pass.
50///
51/// Returns a map of key -> value for all found keys. Missing keys are
52/// omitted from the map.
53pub fn read_workspace_env_values(
54    workspace: &Path,
55    env_keys: &[&str],
56) -> Result<std::collections::HashMap<String, String>> {
57    let env_path = workspace.join(".env");
58    let iter = match dotenvy::from_path_iter(&env_path) {
59        Ok(iter) => iter,
60        Err(dotenvy::Error::Io(err)) if err.kind() == std::io::ErrorKind::NotFound => {
61            return Ok(std::collections::HashMap::new());
62        }
63        Err(err) => {
64            return Err(anyhow!(err).context(format!("Failed to read {}", env_path.display())));
65        }
66    };
67
68    let wanted: std::collections::HashSet<String> = env_keys.iter().map(|&s| s.to_string()).collect();
69    let mut found: std::collections::HashMap<String, String> = std::collections::HashMap::new();
70
71    for item in iter {
72        let (key, value) = item
73            .map_err(|err: dotenvy::Error| anyhow!(err))
74            .with_context(|| format!("Failed to parse {}", env_path.display()))?;
75        if wanted.contains(&key) && !value.trim().is_empty() {
76            found.insert(key, value);
77        }
78    }
79
80    Ok(found)
81}
82
83pub fn remove_workspace_env_value(workspace: &Path, key: &str) -> Result<()> {
84    let env_path = workspace.join(".env");
85    let mut lines = read_existing_lines(&env_path)?;
86    if !lines
87        .iter()
88        .any(|line| line.split_once('=').map(|(k, _)| k.trim()) == Some(key))
89    {
90        return Ok(());
91    }
92    lines.retain(|line| line.split_once('=').map(|(k, _)| k.trim()) != Some(key));
93
94    let parent = env_path.parent().unwrap_or(workspace);
95    fs::create_dir_all(parent).with_context(|| format!("Failed to create directory {}", parent.display()))?;
96
97    let temp = Builder::new()
98        .prefix(".env.")
99        .suffix(".tmp")
100        .tempfile_in(parent)
101        .with_context(|| format!("Failed to create temporary file in {}", parent.display()))?;
102
103    set_private_permissions(temp.as_file(), temp.path())?;
104
105    {
106        let mut writer = BufWriter::new(temp.as_file());
107        for line in &lines {
108            writeln!(writer, "{line}").with_context(|| format!("Failed to write .env entry for {key}"))?;
109        }
110        writer
111            .flush()
112            .with_context(|| format!("Failed to flush temporary .env for {key}"))?;
113    }
114
115    temp.as_file()
116        .sync_all()
117        .with_context(|| format!("Failed to sync temporary .env for {key}"))?;
118
119    let _persisted = temp
120        .persist(&env_path)
121        .with_context(|| format!("Failed to persist {}", env_path.display()))?;
122
123    set_private_path_permissions(&env_path)?;
124    Ok(())
125}
126
127pub fn write_workspace_env_value(workspace: &Path, key: &str, value: &str) -> Result<()> {
128    let env_path = workspace.join(".env");
129    let mut lines = read_existing_lines(&env_path)?;
130    upsert_env_line(&mut lines, key, value);
131
132    let parent = env_path.parent().unwrap_or(workspace);
133    fs::create_dir_all(parent).with_context(|| format!("Failed to create directory {}", parent.display()))?;
134
135    let temp = Builder::new()
136        .prefix(".env.")
137        .suffix(".tmp")
138        .tempfile_in(parent)
139        .with_context(|| format!("Failed to create temporary file in {}", parent.display()))?;
140
141    set_private_permissions(temp.as_file(), temp.path())?;
142
143    {
144        let mut writer = BufWriter::new(temp.as_file());
145        for line in &lines {
146            writeln!(writer, "{line}").with_context(|| format!("Failed to write .env entry for {key}"))?;
147        }
148        writer
149            .flush()
150            .with_context(|| format!("Failed to flush temporary .env for {key}"))?;
151    }
152
153    temp.as_file()
154        .sync_all()
155        .with_context(|| format!("Failed to sync temporary .env for {key}"))?;
156
157    let _persisted = temp
158        .persist(&env_path)
159        .with_context(|| format!("Failed to persist {}", env_path.display()))?;
160
161    set_private_path_permissions(&env_path)?;
162    Ok(())
163}
164
165fn read_existing_lines(env_path: &Path) -> Result<Vec<String>> {
166    if !env_path.exists() {
167        return Ok(Vec::new());
168    }
169
170    let contents = fs::read_to_string(env_path).with_context(|| format!("Failed to read {}", env_path.display()))?;
171    Ok(contents.lines().map(|line| line.to_string()).collect())
172}
173
174fn upsert_env_line(lines: &mut Vec<String>, key: &str, value: &str) {
175    let mut replaced = false;
176    for line in lines.iter_mut() {
177        if let Some((existing_key, _)) = line.split_once('=')
178            && existing_key.trim() == key
179        {
180            *line = format!("{key}={value}");
181            replaced = true;
182        }
183    }
184
185    if !replaced {
186        lines.push(format!("{key}={value}"));
187    }
188}
189
190#[cfg(unix)]
191fn set_private_permissions(file: &File, path: &Path) -> Result<()> {
192    use std::os::unix::fs::PermissionsExt;
193
194    file.set_permissions(fs::Permissions::from_mode(0o600))
195        .with_context(|| format!("Failed to set permissions on {}", path.display()))
196}
197
198#[cfg(not(unix))]
199fn set_private_permissions(_file: &File, _path: &Path) -> Result<()> {
200    Ok(())
201}
202
203#[cfg(unix)]
204fn set_private_path_permissions(path: &Path) -> Result<()> {
205    use std::os::unix::fs::PermissionsExt;
206
207    fs::set_permissions(path, fs::Permissions::from_mode(0o600))
208        .with_context(|| format!("Failed to set permissions on {}", path.display()))
209}
210
211#[cfg(not(unix))]
212fn set_private_path_permissions(_path: &Path) -> Result<()> {
213    Ok(())
214}
215
216#[derive(Debug, Default)]
217pub struct MigrationSummary {
218    pub migrated: u32,
219    pub skipped: u32,
220    pub failed: u32,
221}
222
223#[derive(Debug, Clone, Copy)]
224pub enum MigrationOutcome {
225    Migrated,
226    Skipped,
227    Failed,
228}
229
230#[derive(Debug, Clone, Copy)]
231pub enum SkipReason {
232    LocalProvider,
233    ManagedAuth,
234    NoEnvKeyDefined,
235    NotFoundInEnv,
236    EmptyValue,
237}
238
239/// Migrate a single provider's API key from workspace `.env` into secure storage.
240///
241/// Returns the outcome so callers can render progress in their own UI. The
242/// provider is skipped if it is local, uses managed auth, has no env key
243/// defined, or has no value in `.env`.
244pub fn migrate_single_env_key(
245    workspace: &Path,
246    provider: Provider,
247    store_mode: AuthCredentialsStoreMode,
248    value: Option<&str>,
249) -> Result<MigrationOutcome> {
250    if provider.is_local() || provider.uses_managed_auth() {
251        return Ok(MigrationOutcome::Skipped);
252    }
253
254    let env_key = provider.default_api_key_env();
255    if env_key.is_empty() {
256        return Ok(MigrationOutcome::Skipped);
257    }
258
259    let raw: String = match value {
260        Some(v) => v.to_owned(),
261        None => match read_workspace_env_value(workspace, env_key) {
262            Ok(Some(v)) => v,
263            Ok(None) => return Ok(MigrationOutcome::Skipped),
264            Err(err) => {
265                tracing::warn!("Failed to read {} from {}: {}", env_key, workspace.join(".env").display(), err);
266                return Ok(MigrationOutcome::Failed);
267            }
268        },
269    };
270    let trimmed = raw.trim();
271    if trimmed.is_empty() {
272        return Ok(MigrationOutcome::Skipped);
273    }
274
275    match store_credential_with_mode(provider.as_ref(), env_key, trimmed, store_mode) {
276        Ok(Some(_)) => match remove_workspace_env_value(workspace, env_key) {
277            Ok(()) => Ok(MigrationOutcome::Migrated),
278            Err(err) => {
279                tracing::warn!("Stored {} in keyring but failed to remove from .env: {}", env_key, err);
280                Ok(MigrationOutcome::Failed)
281            }
282        },
283        Ok(None) => Ok(MigrationOutcome::Skipped),
284        Err(err) => {
285            tracing::warn!("Failed to store API key for {}: {}", provider.label(), err);
286            Ok(MigrationOutcome::Failed)
287        }
288    }
289}
290
291/// Migrate API keys for multiple providers from workspace `.env` into secure
292/// storage.
293///
294/// This is the batch-optimized path: the `.env` file is read once, all
295/// providers are processed in memory, and the file is rewritten at most once.
296pub fn migrate_workspace_env_keys(
297    workspace: &Path,
298    providers: &[Provider],
299    store_mode: AuthCredentialsStoreMode,
300) -> Result<(MigrationSummary, Vec<(Provider, MigrationOutcome)>)> {
301    let env_path = workspace.join(".env");
302    if !env_path.exists() {
303        return Ok((MigrationSummary::default(), Vec::new()));
304    }
305
306    let mut lines = read_existing_lines(&env_path)?;
307    let mut line_map: std::collections::HashMap<&str, usize> = std::collections::HashMap::with_capacity(lines.len());
308    for (idx, line) in lines.iter().enumerate() {
309        if let Some((key, _)) = line.split_once('=') {
310            line_map.insert(key.trim(), idx);
311        }
312    }
313
314    let mut summary = MigrationSummary::default();
315    let mut outcomes = Vec::with_capacity(providers.len());
316    let mut removed = std::collections::HashSet::new();
317
318    for &provider in providers {
319        if provider.is_local() || provider.uses_managed_auth() {
320            summary.skipped += 1;
321            outcomes.push((provider, MigrationOutcome::Skipped));
322            continue;
323        }
324
325        let env_key = provider.default_api_key_env();
326        if env_key.is_empty() {
327            summary.skipped += 1;
328            outcomes.push((provider, MigrationOutcome::Skipped));
329            continue;
330        }
331
332        let value = match line_map
333            .get(env_key)
334            .and_then(|&idx| lines.get(idx))
335            .and_then(|line| line.split_once('='))
336            .map(|(_, v)| v.trim())
337        {
338            Some(v) if !v.is_empty() => v,
339            _ => {
340                summary.skipped += 1;
341                outcomes.push((provider, MigrationOutcome::Skipped));
342                continue;
343            }
344        };
345
346        match store_credential_with_mode(provider.as_ref(), env_key, value, store_mode) {
347            Ok(Some(_)) => {
348                if let Some(&idx) = line_map.get(env_key) {
349                    removed.insert(idx);
350                }
351                summary.migrated += 1;
352                outcomes.push((provider, MigrationOutcome::Migrated));
353            }
354            Ok(None) => {
355                summary.skipped += 1;
356                outcomes.push((provider, MigrationOutcome::Skipped));
357            }
358            Err(err) => {
359                tracing::warn!("Failed to store API key for {}: {}", provider.label(), err);
360                summary.failed += 1;
361                outcomes.push((provider, MigrationOutcome::Failed));
362            }
363        }
364    }
365
366    if !removed.is_empty() {
367        let mut new_lines = Vec::with_capacity(lines.len() - removed.len());
368        for (idx, line) in lines.into_iter().enumerate() {
369            if !removed.contains(&idx) {
370                new_lines.push(line);
371            }
372        }
373        lines = new_lines;
374
375        let parent = env_path.parent().unwrap_or(workspace);
376        fs::create_dir_all(parent).with_context(|| format!("Failed to create directory {}", parent.display()))?;
377
378        let temp = Builder::new()
379            .prefix(".env.")
380            .suffix(".tmp")
381            .tempfile_in(parent)
382            .with_context(|| format!("Failed to create temporary file in {}", parent.display()))?;
383
384        set_private_permissions(temp.as_file(), temp.path())?;
385
386        {
387            let mut writer = BufWriter::new(temp.as_file());
388            for line in &lines {
389                writeln!(writer, "{line}").with_context(|| "Failed to write .env entry for migration")?;
390            }
391            writer.flush().with_context(|| "Failed to flush temporary .env for migration")?;
392        }
393
394        temp.as_file()
395            .sync_all()
396            .with_context(|| "Failed to sync temporary .env for migration")?;
397
398        let _persisted = temp
399            .persist(&env_path)
400            .with_context(|| format!("Failed to persist {}", env_path.display()))?;
401
402        set_private_path_permissions(&env_path)?;
403    }
404
405    Ok((summary, outcomes))
406}
407
408#[cfg(test)]
409mod tests {
410    use super::{read_workspace_env_value, remove_workspace_env_value, write_workspace_env_value};
411    use anyhow::Result;
412    use std::fs;
413    use tempfile::tempdir;
414
415    #[test]
416    fn read_returns_value_when_present() -> Result<()> {
417        let dir = tempdir()?;
418        fs::write(dir.path().join(".env"), "OPENAI_API_KEY=sk-test\n")?;
419
420        let value = read_workspace_env_value(dir.path(), "OPENAI_API_KEY")?;
421
422        assert_eq!(value, Some("sk-test".to_string()));
423        Ok(())
424    }
425
426    #[test]
427    fn read_returns_none_when_missing() -> Result<()> {
428        let dir = tempdir()?;
429
430        let value = read_workspace_env_value(dir.path(), "OPENAI_API_KEY")?;
431
432        assert_eq!(value, None);
433        Ok(())
434    }
435
436    #[test]
437    fn write_adds_new_key() -> Result<()> {
438        let dir = tempdir()?;
439
440        write_workspace_env_value(dir.path(), "OPENAI_API_KEY", "sk-test")?;
441
442        let contents = fs::read_to_string(dir.path().join(".env"))?;
443        assert_eq!(contents, "OPENAI_API_KEY=sk-test\n");
444        Ok(())
445    }
446
447    #[test]
448    fn write_replaces_existing_key() -> Result<()> {
449        let dir = tempdir()?;
450        fs::write(dir.path().join(".env"), "OPENAI_API_KEY=old-value\nOTHER_KEY=value\n")?;
451
452        write_workspace_env_value(dir.path(), "OPENAI_API_KEY", "new-value")?;
453
454        let contents = fs::read_to_string(dir.path().join(".env"))?;
455        assert_eq!(contents, "OPENAI_API_KEY=new-value\nOTHER_KEY=value\n");
456        Ok(())
457    }
458
459    #[test]
460    fn remove_deletes_existing_key() -> Result<()> {
461        let dir = tempdir()?;
462        fs::write(dir.path().join(".env"), "OPENAI_API_KEY=sk-test\nOTHER_KEY=value\n")?;
463
464        remove_workspace_env_value(dir.path(), "OPENAI_API_KEY")?;
465
466        let contents = fs::read_to_string(dir.path().join(".env"))?;
467        assert_eq!(contents, "OTHER_KEY=value\n");
468        Ok(())
469    }
470
471    #[test]
472    fn remove_is_idempotent_when_key_absent() -> Result<()> {
473        let dir = tempdir()?;
474        fs::write(dir.path().join(".env"), "OTHER_KEY=value\n")?;
475
476        remove_workspace_env_value(dir.path(), "OPENAI_API_KEY")?;
477
478        let contents = fs::read_to_string(dir.path().join(".env"))?;
479        assert_eq!(contents, "OTHER_KEY=value\n");
480        Ok(())
481    }
482
483    #[test]
484    fn remove_cleans_up_empty_file() -> Result<()> {
485        let dir = tempdir()?;
486        fs::write(dir.path().join(".env"), "OPENAI_API_KEY=sk-test\n")?;
487
488        remove_workspace_env_value(dir.path(), "OPENAI_API_KEY")?;
489
490        let contents = fs::read_to_string(dir.path().join(".env"))?;
491        assert_eq!(contents, "");
492        Ok(())
493    }
494}
vtcode_config/workspace_env.rs

vtcode_config/
workspace_env.rs
