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