diff --git a/src/git/repo_storage.rs b/src/git/repo_storage.rs index a72d851..038b314 100644 --- a/src/git/repo_storage.rs +++ b/src/git/repo_storage.rs @@ -11,6 +11,7 @@ use std::collections::{HashMap, HashSet}; use std::fs; use std::io::{BufRead, Write}; use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{Duration, SystemTime}; /// Initial attributions data structure stored in the INITIAL file @@ -393,22 +394,20 @@ impl PersistedWorkingLog { /* append checkpoint */ pub fn append_checkpoint(&self, checkpoint: &Checkpoint) -> Result<(), AutterError> { - // Read existing checkpoints - let mut checkpoints = self.read_all_checkpoints().unwrap_or_default(); - - // Create a copy, potentially without transcript to reduce storage size. - // - // Tools that DON'T support refetch (transcript must be kept): - // - "mock_ai" - test preset, transcript not stored externally - // - Any other agent-v1 custom tools (detected by lack of tool-specific metadata) - checkpoints.push(checkpoint.clone()); - - // Prune char-level attributions from older checkpoints for the same files - // Only the most recent checkpoint per file needs char-level precision - self.prune_old_char_attributions(&mut checkpoints); - - // Write all checkpoints back - self.write_all_checkpoints(&checkpoints) + self.mutate_all_checkpoints(|checkpoints| { + // Create a copy, potentially without transcript to reduce storage size. + // + // Tools that DON'T support refetch (transcript must be kept): + // - "mock_ai" - test preset, transcript not stored externally + // - Any other agent-v1 custom tools (detected by lack of tool-specific metadata) + checkpoints.push(checkpoint.clone()); + + // Prune char-level attributions from older checkpoints for the same files + // Only the most recent checkpoint per file needs char-level precision + self.prune_old_char_attributions(checkpoints); + Ok(()) + })?; + Ok(()) } pub fn read_all_checkpoints(&self) -> Result, AutterError> { @@ -425,15 +424,27 @@ impl PersistedWorkingLog { let reader = std::io::BufReader::new(file); let mut checkpoints = Vec::new(); - // Parse JSONL file - each line is a separate JSON object - for line in reader.lines() { + // Parse JSONL file - each line is a separate JSON object. Read raw bytes + // so that a corrupt (e.g. NUL-padded or non-UTF-8) line is skipped + // instead of failing the whole read and every checkpoint after it. + for (index, line) in reader.split(b'\n').enumerate() { let line = line?; - if line.trim().is_empty() { + if line.trim_ascii().is_empty() { continue; } - let checkpoint: Checkpoint = serde_json::from_str(&line) - .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?; + let checkpoint: Checkpoint = match serde_json::from_slice(&line) { + Ok(checkpoint) => checkpoint, + Err(error) => { + tracing::warn!( + %error, + path = %checkpoints_file.display(), + line = index + 1, + "skipping corrupt checkpoint line" + ); + continue; + } + }; if checkpoint.api_version != CHECKPOINT_API_VERSION { tracing::debug!( @@ -535,14 +546,36 @@ impl PersistedWorkingLog { // collecting every line into a Vec and join()-ing them into one // giant String. On large histories that intermediate buffer (plus the // joined copy) was a major peak-memory amplifier. - let file = fs::File::create(&checkpoints_file)?; - let mut writer = std::io::BufWriter::new(file); + // + // Write to a uniquely named temp file and rename it over the target so + // that a crash or a concurrent writer can never leave a partially + // written checkpoints.jsonl behind. + static TMP_COUNTER: AtomicU64 = AtomicU64::new(0); + let tmp_file = self.dir.join(format!( + "checkpoints.jsonl.{}.{}.tmp", + std::process::id(), + TMP_COUNTER.fetch_add(1, Ordering::Relaxed) + )); + let result = Self::write_checkpoints_atomically(&tmp_file, &checkpoints_file, checkpoints); + if result.is_err() { + let _ = fs::remove_file(&tmp_file); + } + result + } + + fn write_checkpoints_atomically( + tmp_file: &Path, + checkpoints_file: &Path, + checkpoints: &[Checkpoint], + ) -> Result<(), AutterError> { + let mut writer = std::io::BufWriter::new(fs::File::create(tmp_file)?); for checkpoint in checkpoints { serde_json::to_writer(&mut writer, checkpoint)?; writer.write_all(b"\n")?; } writer.flush()?; - + writer.get_ref().sync_all()?; + fs::rename(tmp_file, checkpoints_file)?; Ok(()) } diff --git a/tests/integration/repo_storage_unit.rs b/tests/integration/repo_storage_unit.rs index c1bb57f..c8cb1e3 100644 --- a/tests/integration/repo_storage_unit.rs +++ b/tests/integration/repo_storage_unit.rs @@ -1,3 +1,4 @@ +use crate::repos::test_file::ExpectedLineExt; use crate::repos::test_repo::TestRepo; use autter::authorship::attribution_tracker::LineAttribution; use autter::authorship::working_log::{ @@ -605,3 +606,123 @@ fn test_prune_does_not_touch_active_working_logs() { "Active working logs should not be pruned" ); } + +// --------------------------------------------------------------------------- +// 16. test_read_all_checkpoints_skips_corrupt_lines +// --------------------------------------------------------------------------- + +#[test] +fn test_read_all_checkpoints_skips_corrupt_lines() { + let repo = TestRepo::new(); + let repo_storage = storage_for(&repo); + let working_log = repo_storage + .working_log_for_base_commit("test-commit-sha") + .unwrap(); + + let first = Checkpoint::new( + CheckpointKind::Human, + "diff-1".to_string(), + "first-author".to_string(), + vec![], + ); + let mut second = first.clone(); + second.author = "second-author".to_string(); + + let mut content = Vec::new(); + content.extend_from_slice(serde_json::to_string(&first).unwrap().as_bytes()); + content.push(b'\n'); + // NUL padding, as left behind by an interrupted in-place rewrite + content.extend_from_slice(&[0u8; 64]); + content.push(b'\n'); + // Invalid UTF-8 + content.extend_from_slice(&[0xff, 0xfe, b'{', b'\n']); + // Truncated JSON + content.extend_from_slice(b"{\"kind\":\"Hum\n"); + content.extend_from_slice(serde_json::to_string(&second).unwrap().as_bytes()); + content.push(b'\n'); + fs::write(working_log.dir.join("checkpoints.jsonl"), content).unwrap(); + + let checkpoints = working_log + .read_all_checkpoints() + .expect("corrupt lines should not fail the read"); + let authors: Vec<&str> = checkpoints.iter().map(|c| c.author.as_str()).collect(); + assert_eq!(authors, vec!["first-author", "second-author"]); +} + +// --------------------------------------------------------------------------- +// 17. test_append_checkpoint_keeps_earlier_checkpoints_past_corrupt_line +// --------------------------------------------------------------------------- + +#[test] +fn test_append_checkpoint_keeps_earlier_checkpoints_past_corrupt_line() { + let repo = TestRepo::new(); + let repo_storage = storage_for(&repo); + let working_log = repo_storage + .working_log_for_base_commit("test-commit-sha") + .unwrap(); + + let earlier = Checkpoint::new( + CheckpointKind::Human, + "diff-1".to_string(), + "earlier-author".to_string(), + vec![], + ); + working_log.append_checkpoint(&earlier).unwrap(); + + let checkpoints_file = working_log.dir.join("checkpoints.jsonl"); + let mut content = fs::read(&checkpoints_file).unwrap(); + content.extend_from_slice(&[0u8; 32]); + content.push(b'\n'); + fs::write(&checkpoints_file, content).unwrap(); + + let later = Checkpoint::new( + CheckpointKind::AiAgent, + "diff-2".to_string(), + "later-author".to_string(), + vec![], + ); + working_log + .append_checkpoint(&later) + .expect("append should succeed with a corrupt line present"); + + let checkpoints = working_log.read_all_checkpoints().unwrap(); + let authors: Vec<&str> = checkpoints.iter().map(|c| c.author.as_str()).collect(); + assert_eq!(authors, vec!["earlier-author", "later-author"]); + + let rewritten = fs::read(&checkpoints_file).unwrap(); + assert!( + !rewritten.contains(&0u8), + "the rewrite should drop the corrupt line" + ); +} + +// --------------------------------------------------------------------------- +// 18. test_ai_checkpoint_succeeds_with_corrupt_checkpoint_line +// --------------------------------------------------------------------------- + +#[test] +fn test_ai_checkpoint_succeeds_with_corrupt_checkpoint_line() { + let repo = TestRepo::new(); + let file_path = repo.path().join("example.txt"); + + fs::write(&file_path, "Untracked line\n").unwrap(); + repo.stage_all_and_commit("Initial commit").unwrap(); + let mut file = repo.filename("example.txt"); + file.assert_committed_lines(lines!["Untracked line".unattributed_human()]); + + let checkpoints_file = repo.current_working_logs().dir.join("checkpoints.jsonl"); + let mut content = fs::read(&checkpoints_file).unwrap_or_default(); + content.extend_from_slice(&[0u8; 64]); + content.push(b'\n'); + fs::write(&checkpoints_file, content).unwrap(); + + fs::write(&file_path, "Untracked line\nAI line\n").unwrap(); + repo.autter(&["checkpoint", "mock_ai", "example.txt"]) + .expect("AI checkpoint should succeed despite a corrupt checkpoint line"); + + repo.stage_all_and_commit("AI commit").unwrap(); + file.assert_committed_lines(lines![ + "Untracked line".unattributed_human(), + "AI line".ai(), + ]); +}