Skip to main content

vtcode_core/terminal_setup/
config_writer.rs

1//! Safe atomic configuration file writing with smart merging.
2//!
3//! Provides atomic writes and marker-based config merging to preserve user customizations.
4
5use crate::utils::file_utils::ensure_dir_exists_sync;
6use anyhow::{Context, Result};
7use std::io::Write;
8use std::path::Path;
9use tempfile::NamedTempFile;
10
11/// Configuration file format
12#[derive(Debug, Clone, Copy, PartialEq, Eq)]
13pub enum ConfigFormat {
14    /// Plain text or shell-style config (# comments)
15    PlainText,
16    /// TOML format
17    Toml,
18    /// JSON format
19    Json,
20    /// YAML format
21    Yaml,
22    /// JavaScript format (for .hyper.js, etc.)
23    JavaScript,
24}
25
26/// Markers for identifying VT Code-managed configuration sections
27pub const VTCODE_BEGIN_MARKER: &str = "BEGIN VTCODE CONFIGURATION";
28pub const VTCODE_END_MARKER: &str = "END VTCODE CONFIGURATION";
29
30/// Safe atomic configuration file writer
31pub struct ConfigWriter;
32
33impl ConfigWriter {
34    /// Write content to a file atomically using a temp file and rename
35    ///
36    /// This ensures that the file is never left in a partially-written state
37    pub fn write_atomic(path: &Path, content: &str) -> Result<()> {
38        // Create parent directory if it doesn't exist
39        if let Some(parent) = path.parent() {
40            ensure_dir_exists_sync(parent)
41                .with_context(|| format!("Failed to create directory: {}", parent.display()))?;
42        }
43
44        // Create temp file in the same directory for atomic rename
45        let temp_file = NamedTempFile::new_in(path.parent().unwrap_or_else(|| Path::new(".")))
46            .with_context(|| {
47                format!(
48                    "Failed to create temp file in directory: {}",
49                    path.display()
50                )
51            })?;
52
53        // Write content to temp file
54        temp_file
55            .as_file()
56            .write_all(content.as_bytes())
57            .with_context(|| "Failed to write content to temp file")?;
58
59        // Sync to ensure data is written to disk
60        temp_file
61            .as_file()
62            .sync_all()
63            .with_context(|| "Failed to sync temp file to disk")?;
64
65        // Atomically rename temp file to target path
66        temp_file
67            .persist(path)
68            .with_context(|| format!("Failed to write config file: {}", path.display()))?;
69
70        Ok(())
71    }
72
73    /// Merge new VT Code configuration section with existing config
74    ///
75    /// Removes old VT Code sections and adds the new section, preserving user customizations
76    pub fn merge_with_markers(
77        existing: &str,
78        new_section: &str,
79        format: ConfigFormat,
80    ) -> Result<String> {
81        // Remove any existing VT Code section
82        let cleaned = Self::remove_vtcode_section(existing);
83
84        // Add new VT Code section with markers
85        let vtcode_section = Self::wrap_with_markers(new_section, format);
86
87        // Determine where to insert the new section
88        let merged = if cleaned.trim().is_empty() {
89            // File is empty, just use the VT Code section
90            vtcode_section
91        } else {
92            // Prepend VT Code section so user-defined settings stay authoritative.
93            // Most terminal configs resolve duplicate keys as "last value wins".
94            format!("{}\n\n{}", vtcode_section, cleaned.trim_start())
95        };
96
97        Ok(merged)
98    }
99
100    /// Remove existing VT Code configuration section from content
101    fn remove_vtcode_section(content: &str) -> String {
102        let mut result = Vec::new();
103        let mut in_vtcode_section = false;
104
105        for line in content.lines() {
106            if line.contains(VTCODE_BEGIN_MARKER) {
107                in_vtcode_section = true;
108                continue;
109            }
110
111            if line.contains(VTCODE_END_MARKER) {
112                in_vtcode_section = false;
113                continue;
114            }
115
116            if !in_vtcode_section {
117                result.push(line);
118            }
119        }
120
121        result.join("\n")
122    }
123
124    /// Wrap configuration content with VT Code markers
125    fn wrap_with_markers(content: &str, format: ConfigFormat) -> String {
126        let comment_prefix = match format {
127            ConfigFormat::PlainText | ConfigFormat::Toml | ConfigFormat::Yaml => "#",
128            ConfigFormat::Json => "//", // JSON doesn't support comments, but some parsers allow them
129            ConfigFormat::JavaScript => "//",
130        };
131
132        let header = format!(
133            "{} {}\n{} VT Code-managed section - auto-generated\n{} Do not edit manually",
134            comment_prefix, VTCODE_BEGIN_MARKER, comment_prefix, comment_prefix
135        );
136
137        let footer = format!("{} {}", comment_prefix, VTCODE_END_MARKER);
138
139        format!("{}\n{}\n{}", header, content.trim(), footer)
140    }
141
142    /// Detect config file format from extension
143    pub fn detect_format(path: &Path) -> ConfigFormat {
144        if let Some(ext) = path.extension().and_then(|e| e.to_str()) {
145            match ext {
146                "toml" => ConfigFormat::Toml,
147                "json" => ConfigFormat::Json,
148                "yaml" | "yml" => ConfigFormat::Yaml,
149                "js" => ConfigFormat::JavaScript,
150                _ => ConfigFormat::PlainText,
151            }
152        } else {
153            ConfigFormat::PlainText
154        }
155    }
156}
157
158#[cfg(test)]
159mod tests {
160    use super::*;
161    use tempfile::TempDir;
162
163    #[test]
164    fn test_atomic_write() {
165        let temp_dir = TempDir::new().unwrap();
166        let path = temp_dir.path().join("test.conf");
167
168        let content = "test content";
169        ConfigWriter::write_atomic(&path, content).unwrap();
170
171        assert!(path.exists());
172        assert_eq!(std::fs::read_to_string(&path).unwrap(), content);
173    }
174
175    #[test]
176    fn test_remove_vtcode_section() {
177        let content = r#"# User config
178user_setting = 1
179
180# BEGIN VTCODE CONFIGURATION
181# VT Code-managed section - auto-generated
182vtcode_setting = 2
183# END VTCODE CONFIGURATION
184
185# More user config
186another_setting = 3
187"#;
188
189        let result = ConfigWriter::remove_vtcode_section(content);
190
191        assert!(!result.contains("vtcode_setting"));
192        assert!(result.contains("user_setting"));
193        assert!(result.contains("another_setting"));
194        assert!(!result.contains("VTCODE CONFIGURATION"));
195    }
196
197    #[test]
198    fn test_merge_with_markers() {
199        let existing = r#"# User config
200user_setting = 1
201"#;
202
203        let new_section = "vtcode_setting = 2";
204
205        let result =
206            ConfigWriter::merge_with_markers(existing, new_section, ConfigFormat::PlainText)
207                .unwrap();
208
209        assert!(result.contains("user_setting"));
210        assert!(result.contains("vtcode_setting"));
211        assert!(result.contains(VTCODE_BEGIN_MARKER));
212        assert!(result.contains(VTCODE_END_MARKER));
213        assert!(
214            result
215                .find("vtcode_setting")
216                .expect("missing VT Code setting")
217                < result.find("user_setting").expect("missing user setting")
218        );
219    }
220
221    #[test]
222    fn test_merge_empty_file() {
223        let new_section = "vtcode_setting = 1";
224
225        let result =
226            ConfigWriter::merge_with_markers("", new_section, ConfigFormat::PlainText).unwrap();
227
228        assert!(result.contains("vtcode_setting"));
229        assert!(result.contains(VTCODE_BEGIN_MARKER));
230    }
231
232    #[test]
233    fn test_detect_format() {
234        assert_eq!(
235            ConfigWriter::detect_format(Path::new("test.toml")),
236            ConfigFormat::Toml
237        );
238        assert_eq!(
239            ConfigWriter::detect_format(Path::new("test.json")),
240            ConfigFormat::Json
241        );
242        assert_eq!(
243            ConfigWriter::detect_format(Path::new("test.yaml")),
244            ConfigFormat::Yaml
245        );
246        assert_eq!(
247            ConfigWriter::detect_format(Path::new("test.js")),
248            ConfigFormat::JavaScript
249        );
250        assert_eq!(
251            ConfigWriter::detect_format(Path::new("test.conf")),
252            ConfigFormat::PlainText
253        );
254    }
255
256    #[test]
257    fn test_wrap_with_markers_toml() {
258        let content = "setting = 1";
259        let result = ConfigWriter::wrap_with_markers(content, ConfigFormat::Toml);
260
261        assert!(result.starts_with("# BEGIN VTCODE CONFIGURATION"));
262        assert!(result.ends_with("# END VTCODE CONFIGURATION"));
263        assert!(result.contains("setting = 1"));
264    }
265
266    #[test]
267    fn test_wrap_with_markers_javascript() {
268        let content = "const setting = 1;";
269        let result = ConfigWriter::wrap_with_markers(content, ConfigFormat::JavaScript);
270
271        assert!(result.starts_with("// BEGIN VTCODE CONFIGURATION"));
272        assert!(result.contains("const setting = 1;"));
273    }
274}