code_kb_core/
edit.rs

1use rusqlite::Connection;
2use serde::{Deserialize, Serialize};
3use sha2::{Digest, Sha256};
4use std::fs;
5use std::io::Write;
6use std::path::Path;
7use thiserror::Error;
8
9use crate::queries;
10use crate::slicer;
11use crate::sync;
12use crate::syntax::{self, SyntaxError};
13use crate::workspace::Workspace;
14
15#[derive(Debug, Error)]
16pub enum EditError {
17    #[error("Symbol '{0}' not found in '{1}'")]
18    SymbolNotFound(String, String),
19    #[error("Symbol has no body defined (e.g. trait declaration without default implementation)")]
20    NoBodyDefined,
21    #[error("Optimistic lock failed: expected body hash '{0}', found '{1}'")]
22    HashMismatch(String, String),
23    #[error(
24        "File offsets out of bounds: range [{0}..{1}], but file length is {2} bytes (file may have shrunk or changed)"
25    )]
26    InvalidOffsetRange(usize, usize, usize),
27    #[error("Pre-flight syntax validation failed: {0}")]
28    Syntax(#[from] SyntaxError),
29    #[error("Replacement content is not valid UTF-8: {0}")]
30    InvalidUtf8(String),
31    #[error("Failed to read/write file '{0}': {1}")]
32    Io(String, #[source] std::io::Error),
33    #[error("Synchronization failed and file was rolled back: {0}")]
34    SyncWithRollback(String),
35    #[error(
36        "Synchronization failed after edit ({sync_error}) and rollback also failed: {rollback_error}"
37    )]
38    SyncRollbackFailed {
39        sync_error: String,
40        rollback_error: String,
41    },
42    #[error("File '{0}' was concurrently modified; edit aborted")]
43    ConcurrentModification(String),
44    #[error("Synchronization failed after edit: {0}")]
45    Sync(#[from] sync::SyncError),
46    #[error("Workspace path error: {0}")]
47    Workspace(#[from] crate::workspace::WorkspaceError),
48    #[error("Query error: {0}")]
49    Query(#[from] queries::QueryError),
50}
51
52/// Result of an atomic symbol body replacement.
53#[derive(Debug, Clone, Serialize, Deserialize)]
54pub struct EditResult {
55    pub symbol_name: String,
56    pub file_path: String,
57    pub old_body_hash: String,
58    pub new_body_hash: String,
59    pub bytes_written: usize,
60    pub syntax_checked: bool,
61}
62
63/// Computes SHA256 of a string content.
64pub fn hash_content(content: &str) -> String {
65    let mut hasher = Sha256::new();
66    hasher.update(content.as_bytes());
67    hex::encode(hasher.finalize())
68}
69
70fn is_transient_lock_error(err: &std::io::Error) -> bool {
71    #[cfg(windows)]
72    {
73        if let Some(code) = err.raw_os_error() {
74            // ERROR_ACCESS_DENIED = 5, ERROR_SHARING_VIOLATION = 32, ERROR_LOCK_VIOLATION = 33
75            if code == 5 || code == 32 || code == 33 {
76                return true;
77            }
78        }
79    }
80    matches!(err.kind(), std::io::ErrorKind::PermissionDenied)
81}
82
83fn persist_with_retry(
84    mut temp_file: tempfile::NamedTempFile,
85    dest: &Path,
86    expected_dest_bytes: Option<&[u8]>,
87) -> Result<(), std::io::Error> {
88    for attempt in 0..5 {
89        if attempt > 0
90            && let Some(expected) = expected_dest_bytes
91            && let Ok(current) = fs::read(dest)
92            && current != expected
93        {
94            return Err(std::io::Error::other(
95                "Destination file was concurrently modified during retry",
96            ));
97        }
98        match temp_file.persist(dest) {
99            Ok(_) => return Ok(()),
100            Err(e) => {
101                let is_transient = is_transient_lock_error(&e.error);
102                temp_file = e.file;
103                if is_transient && attempt < 4 {
104                    std::thread::sleep(std::time::Duration::from_millis(10 * (1 << attempt)));
105                    continue;
106                }
107                return Err(e.error);
108            }
109        }
110    }
111    unreachable!()
112}
113
114/// Atomically replaces the implementation body of a symbol by name.
115pub fn replace_symbol_body(
116    workspace: &Workspace,
117    db_path: &Path,
118    conn: &Connection,
119    symbol_name: &str,
120    file_path: &str,
121    new_body: &str,
122    expected_body_hash: Option<&str>,
123) -> Result<EditResult, EditError> {
124    let (abs_path, rel_path) = workspace.resolve_path(Path::new(file_path))?;
125
126    // Tier 2: Refresh file in index before querying symbol offsets, propagating any sync errors
127    sync::ensure_fresh_file(workspace, db_path, conn, &rel_path)?;
128
129    // Find symbol in database with exact path
130    let symbol = queries::get_symbol_by_name_exact(conn, symbol_name, &rel_path)?
131        .ok_or_else(|| EditError::SymbolNotFound(symbol_name.to_string(), rel_path.clone()))?;
132
133    let body_start = symbol.body_start_byte.ok_or(EditError::NoBodyDefined)?;
134    let body_end = symbol.body_end_byte.ok_or(EditError::NoBodyDefined)?;
135
136    // Read existing file metadata (permissions) and content
137    let existing_metadata =
138        fs::metadata(&abs_path).map_err(|e| EditError::Io(abs_path.display().to_string(), e))?;
139    let existing_permissions = existing_metadata.permissions();
140    let existing_bytes =
141        fs::read(&abs_path).map_err(|e| EditError::Io(abs_path.display().to_string(), e))?;
142
143    // Bounds check to prevent out-of-bounds panics
144    if body_start > body_end || body_end > existing_bytes.len() {
145        return Err(EditError::InvalidOffsetRange(
146            body_start,
147            body_end,
148            existing_bytes.len(),
149        ));
150    }
151
152    let existing_body =
153        slicer::slice_bytes_safe(&existing_bytes, body_start, body_end).map_err(|e| {
154            EditError::Io(
155                abs_path.display().to_string(),
156                std::io::Error::new(std::io::ErrorKind::InvalidData, e.to_string()),
157            )
158        })?;
159
160    let current_sha256 = hash_content(existing_body);
161
162    // Verify optimistic lock if caller specified expected_body_hash
163    if let Some(expected) = expected_body_hash
164        && expected != current_sha256
165    {
166        return Err(EditError::HashMismatch(
167            expected.to_string(),
168            current_sha256,
169        ));
170    }
171
172    // Match file line endings (CRLF vs LF)
173    let is_crlf = existing_bytes.windows(2).any(|w| w == b"\r\n");
174    let normalized_body = if is_crlf {
175        // Uniformly normalize all line endings to CRLF, including mixed inputs
176        let lf_only = new_body.replace("\r\n", "\n");
177        lf_only.replace('\n', "\r\n")
178    } else {
179        // Uniformly normalize all line endings to LF, including mixed inputs
180        new_body.replace("\r\n", "\n")
181    };
182
183    // Construct new file content with replaced byte span
184    let mut new_file_bytes = Vec::with_capacity(existing_bytes.len() + normalized_body.len());
185    new_file_bytes.extend_from_slice(&existing_bytes[..body_start]);
186    new_file_bytes.extend_from_slice(normalized_body.as_bytes());
187    new_file_bytes.extend_from_slice(&existing_bytes[body_end..]);
188
189    // Pre-flight syntax validation before touching disk
190    let new_file_str =
191        std::str::from_utf8(&new_file_bytes).map_err(|e| EditError::InvalidUtf8(e.to_string()))?;
192    let syntax_checked = syntax::validate_syntax(&rel_path, new_file_str)?;
193
194    // Backup original bytes for rollback if re-indexing fails
195    let backup_bytes = existing_bytes.clone();
196
197    // Final pre-commit disk check: ensure file was not concurrently modified between read and write
198    let current_disk =
199        fs::read(&abs_path).map_err(|e| EditError::Io(abs_path.display().to_string(), e))?;
200    if current_disk != existing_bytes {
201        return Err(EditError::ConcurrentModification(rel_path));
202    }
203
204    // Write file atomically: write to a temporary file in the target directory, then persist
205    let target_dir = abs_path.parent().unwrap_or(Path::new("."));
206    let mut temp_file = tempfile::Builder::new()
207        .prefix(".code-kb-edit-")
208        .suffix(".tmp")
209        .tempfile_in(target_dir)
210        .map_err(|e| EditError::Io(abs_path.display().to_string(), e))?;
211
212    temp_file
213        .write_all(&new_file_bytes)
214        .map_err(|e| EditError::Io(abs_path.display().to_string(), e))?;
215    temp_file
216        .flush()
217        .map_err(|e| EditError::Io(abs_path.display().to_string(), e))?;
218
219    // Preserve existing file permissions (e.g. +x 0755)
220    let _ = temp_file
221        .as_file()
222        .set_permissions(existing_permissions.clone());
223
224    persist_with_retry(temp_file, &abs_path, Some(&existing_bytes)).map_err(|e| {
225        if e.to_string().contains("concurrently modified") {
226            EditError::ConcurrentModification(rel_path.clone())
227        } else {
228            EditError::Io(abs_path.display().to_string(), e)
229        }
230    })?;
231
232    // Tier 1: Immediately re-index the file so catalog is 100% fresh.
233    // If indexing fails, roll back to original content safely and atomically.
234    if let Err(err) = sync::update_file(workspace, db_path, &rel_path) {
235        // First check if the file on disk is still our newly written file
236        let disk_post_write = fs::read(&abs_path);
237        if disk_post_write.as_deref().ok() != Some(new_file_bytes.as_slice()) {
238            return Err(EditError::ConcurrentModification(format!(
239                "File was concurrently modified during re-indexing; rollback aborted: {err}"
240            )));
241        }
242
243        // Perform rollback atomically via temporary file
244        let rollback_res = (|| -> Result<(), std::io::Error> {
245            let mut rollback_tmp = tempfile::Builder::new()
246                .prefix(".code-kb-rollback-")
247                .suffix(".tmp")
248                .tempfile_in(target_dir)?;
249            rollback_tmp.write_all(&backup_bytes)?;
250            rollback_tmp.flush()?;
251            let _ = rollback_tmp
252                .as_file()
253                .set_permissions(existing_permissions.clone());
254            persist_with_retry(rollback_tmp, &abs_path, Some(&new_file_bytes))?;
255            Ok(())
256        })();
257
258        match rollback_res {
259            Ok(()) => return Err(EditError::SyncWithRollback(err.to_string())),
260            Err(rollback_err) => {
261                return Err(EditError::SyncRollbackFailed {
262                    sync_error: err.to_string(),
263                    rollback_error: rollback_err.to_string(),
264                });
265            }
266        }
267    }
268
269    let new_body_hash = hash_content(&normalized_body);
270
271    Ok(EditResult {
272        symbol_name: symbol_name.to_string(),
273        file_path: rel_path,
274        old_body_hash: current_sha256,
275        new_body_hash,
276        bytes_written: new_file_bytes.len(),
277        syntax_checked,
278    })
279}
280
281#[cfg(test)]
282mod tests {
283    use super::*;
284
285    #[test]
286    fn test_mixed_crlf_lf_normalization() {
287        let mixed_input = "line1\r\nline2\nline3\r\nline4\n";
288        // When file is CRLF:
289        let lf_only = mixed_input.replace("\r\n", "\n");
290        let crlf_normalized = lf_only.replace('\n', "\r\n");
291        assert_eq!(crlf_normalized, "line1\r\nline2\r\nline3\r\nline4\r\n");
292
293        // When file is LF:
294        let lf_normalized = mixed_input.replace("\r\n", "\n");
295        assert_eq!(lf_normalized, "line1\nline2\nline3\nline4\n");
296    }
297
298    #[test]
299    fn test_persist_with_retry_succeeds() {
300        let dir = crate::safe_tempdir();
301        let target_file = dir.path().join("test_persist.txt");
302        fs::write(&target_file, "initial").unwrap();
303
304        let mut temp_file = tempfile::Builder::new()
305            .prefix(".test-persist-")
306            .suffix(".tmp")
307            .tempfile_in(dir.path())
308            .unwrap();
309        temp_file.write_all(b"updated").unwrap();
310        temp_file.flush().unwrap();
311
312        persist_with_retry(temp_file, &target_file, Some(b"initial"))
313            .expect("persist_with_retry must succeed");
314        assert_eq!(fs::read_to_string(&target_file).unwrap(), "updated");
315    }
316}
code_kb_core/edit.rs

code_kb_core/
edit.rs
