vtcode_config/
workspace_env.rs

1use anyhow::{Context, Result, anyhow};
2use std::fs::{self, File};
3use std::io::{BufWriter, Write};
4use std::path::Path;
5
6use tempfile::Builder;
7
8pub fn read_workspace_env_value(workspace: &Path, env_key: &str) -> Result<Option<String>> {
9    let env_path = workspace.join(".env");
10    let iter = match dotenvy::from_path_iter(&env_path) {
11        Ok(iter) => iter,
12        Err(dotenvy::Error::Io(err)) if err.kind() == std::io::ErrorKind::NotFound => {
13            return Ok(None);
14        }
15        Err(err) => {
16            return Err(anyhow!(err).context(format!("Failed to read {}", env_path.display())));
17        }
18    };
19
20    for item in iter {
21        let (key, value) = item
22            .map_err(|err: dotenvy::Error| anyhow!(err))
23            .with_context(|| format!("Failed to parse {}", env_path.display()))?;
24        if key == env_key {
25            if value.trim().is_empty() {
26                return Ok(None);
27            }
28            return Ok(Some(value));
29        }
30    }
31
32    Ok(None)
33}
34
35pub fn write_workspace_env_value(workspace: &Path, key: &str, value: &str) -> Result<()> {
36    let env_path = workspace.join(".env");
37    let mut lines = read_existing_lines(&env_path)?;
38    upsert_env_line(&mut lines, key, value);
39
40    let parent = env_path.parent().unwrap_or(workspace);
41    fs::create_dir_all(parent)
42        .with_context(|| format!("Failed to create directory {}", parent.display()))?;
43
44    let temp = Builder::new()
45        .prefix(".env.")
46        .suffix(".tmp")
47        .tempfile_in(parent)
48        .with_context(|| format!("Failed to create temporary file in {}", parent.display()))?;
49
50    set_private_permissions(temp.as_file(), temp.path())?;
51
52    {
53        let mut writer = BufWriter::new(temp.as_file());
54        for line in &lines {
55            writeln!(writer, "{line}")
56                .with_context(|| format!("Failed to write .env entry for {key}"))?;
57        }
58        writer
59            .flush()
60            .with_context(|| format!("Failed to flush temporary .env for {}", key))?;
61    }
62
63    temp.as_file()
64        .sync_all()
65        .with_context(|| format!("Failed to sync temporary .env for {}", key))?;
66
67    let _persisted = temp
68        .persist(&env_path)
69        .with_context(|| format!("Failed to persist {}", env_path.display()))?;
70
71    set_private_path_permissions(&env_path)?;
72    Ok(())
73}
74
75fn read_existing_lines(env_path: &Path) -> Result<Vec<String>> {
76    if !env_path.exists() {
77        return Ok(Vec::new());
78    }
79
80    let contents = fs::read_to_string(env_path)
81        .with_context(|| format!("Failed to read {}", env_path.display()))?;
82    Ok(contents.lines().map(|line| line.to_string()).collect())
83}
84
85fn upsert_env_line(lines: &mut Vec<String>, key: &str, value: &str) {
86    let mut replaced = false;
87    for line in lines.iter_mut() {
88        if let Some((existing_key, _)) = line.split_once('=')
89            && existing_key.trim() == key
90        {
91            *line = format!("{key}={value}");
92            replaced = true;
93        }
94    }
95
96    if !replaced {
97        lines.push(format!("{key}={value}"));
98    }
99}
100
101#[cfg(unix)]
102fn set_private_permissions(file: &File, path: &Path) -> Result<()> {
103    use std::os::unix::fs::PermissionsExt;
104
105    file.set_permissions(fs::Permissions::from_mode(0o600))
106        .with_context(|| format!("Failed to set permissions on {}", path.display()))
107}
108
109#[cfg(not(unix))]
110fn set_private_permissions(_file: &File, _path: &Path) -> Result<()> {
111    Ok(())
112}
113
114#[cfg(unix)]
115fn set_private_path_permissions(path: &Path) -> Result<()> {
116    use std::os::unix::fs::PermissionsExt;
117
118    fs::set_permissions(path, fs::Permissions::from_mode(0o600))
119        .with_context(|| format!("Failed to set permissions on {}", path.display()))
120}
121
122#[cfg(not(unix))]
123fn set_private_path_permissions(_path: &Path) -> Result<()> {
124    Ok(())
125}
126
127#[cfg(test)]
128mod tests {
129    use super::{read_workspace_env_value, write_workspace_env_value};
130    use anyhow::Result;
131    use std::fs;
132    use tempfile::tempdir;
133
134    #[test]
135    fn read_returns_value_when_present() -> Result<()> {
136        let dir = tempdir()?;
137        fs::write(dir.path().join(".env"), "OPENAI_API_KEY=sk-test\n")?;
138
139        let value = read_workspace_env_value(dir.path(), "OPENAI_API_KEY")?;
140
141        assert_eq!(value, Some("sk-test".to_string()));
142        Ok(())
143    }
144
145    #[test]
146    fn read_returns_none_when_missing() -> Result<()> {
147        let dir = tempdir()?;
148
149        let value = read_workspace_env_value(dir.path(), "OPENAI_API_KEY")?;
150
151        assert_eq!(value, None);
152        Ok(())
153    }
154
155    #[test]
156    fn write_adds_new_key() -> Result<()> {
157        let dir = tempdir()?;
158
159        write_workspace_env_value(dir.path(), "OPENAI_API_KEY", "sk-test")?;
160
161        let contents = fs::read_to_string(dir.path().join(".env"))?;
162        assert_eq!(contents, "OPENAI_API_KEY=sk-test\n");
163        Ok(())
164    }
165
166    #[test]
167    fn write_replaces_existing_key() -> Result<()> {
168        let dir = tempdir()?;
169        fs::write(
170            dir.path().join(".env"),
171            "OPENAI_API_KEY=old-value\nOTHER_KEY=value\n",
172        )?;
173
174        write_workspace_env_value(dir.path(), "OPENAI_API_KEY", "new-value")?;
175
176        let contents = fs::read_to_string(dir.path().join(".env"))?;
177        assert_eq!(contents, "OPENAI_API_KEY=new-value\nOTHER_KEY=value\n");
178        Ok(())
179    }
180}
vtcode_config/workspace_env.rs

vtcode_config/
workspace_env.rs
