From 7bd0818ad171c893f8e8355274e6f83bcd54dc30 Mon Sep 17 00:00:00 2001
From: Jeremy Brown <jeremy@tenfourty.com>
Date: Wed, 24 Jun 2026 15:05:37 +0200
Subject: [PATCH] feat(writes): all-or-nothing batch corrections (refs #1)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

apply_corrections now runs in two phases: compute every change in memory
first, then write them; if any write fails, restore the files already
written from their in-memory originals and re-raise. A bulk `correct
--apply` across many files can no longer leave a half-applied state — it's
all-or-nothing (issue #1's COW slice 2). Single-file writes were already
atomic, so this is the only batch that needed it.
---
 src/kb/correct.py     | 29 +++++++++++++++++++++--------
 tests/test_correct.py | 37 +++++++++++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+), 8 deletions(-)

diff --git a/src/kb/correct.py b/src/kb/correct.py
index aea92df..db9e657 100644
--- a/src/kb/correct.py
+++ b/src/kb/correct.py
@@ -290,10 +290,9 @@ def apply_corrections(
             changed_paths=[],
         )
 
-    files_changed = 0
+    # Phase 1 — prepare: read + compute every change in memory; write nothing yet.
+    planned: list[tuple[Path, str, str, str]] = []  # (full_path, rel_path, original, new)
     total_replaced = 0
-    changed_paths: list[str] = []
-
     for match in matches:
         full_path = memory_root / match.rel_path
         try:
@@ -309,18 +308,32 @@ def apply_corrections(
 
         new_content, count = pattern.subn(lambda _: new_term, content)
         if count > 0:
-            _atomic_write(full_path, new_content)
-            files_changed += 1
+            planned.append((full_path, match.rel_path, content, new_content))
             total_replaced += count
-            changed_paths.append(match.rel_path)
+
+    # Phase 2 — commit: write all planned changes; on any write failure, restore the
+    # files already written so the batch is all-or-nothing (#1 COW atomicity).
+    written: list[tuple[Path, str]] = []  # (full_path, original) for rollback
+    changed_paths: list[str] = []
+    try:
+        for full_path, rel_path, original, new_content in planned:
+            _atomic_write(full_path, new_content)
+            written.append((full_path, original))
+            changed_paths.append(rel_path)
+    except OSError:
+        for fp, original in reversed(written):
+            # Best-effort restore; originals are held in memory.
+            with contextlib.suppress(OSError):
+                _atomic_write(fp, original)
+        raise
 
     result = CorrectionResult(
-        files_changed=files_changed,
+        files_changed=len(written),
         occurrences_replaced=total_replaced,
         changed_paths=changed_paths,
     )
 
-    if log_path is not None and files_changed > 0:
+    if log_path is not None and result.files_changed > 0:
         _write_audit_log(log_path, matches[0].search_term, new_term, result)
 
     return result
diff --git a/tests/test_correct.py b/tests/test_correct.py
index c75521a..d160c0d 100644
--- a/tests/test_correct.py
+++ b/tests/test_correct.py
@@ -690,6 +690,43 @@ def test_scan_nfc_normalizes_rel_paths(self, memory_tree: Path):
             )
 
 
+class TestApplyCorrectionsAtomicity:
+    """apply_corrections is all-or-nothing: a mid-batch failure rolls back (#1 COW)."""
+
+    def test_rolls_back_on_write_failure(self, memory_tree: Path, monkeypatch):
+        from kb import correct as correct_mod
+
+        matches = scan(memory_tree, "Quartz Indexer")
+        assert len(matches) >= 2, "need >=2 matched files to exercise rollback"
+        paths = [memory_tree / m.rel_path for m in matches]
+        originals = {p: p.read_text(encoding="utf-8") for p in paths}
+
+        real_write = correct_mod._atomic_write
+        calls = {"n": 0}
+
+        def flaky(path, content):  # type: ignore[no-untyped-def]
+            calls["n"] += 1
+            if calls["n"] == 2:
+                raise OSError("simulated disk failure on the 2nd write")
+            real_write(path, content)
+
+        monkeypatch.setattr(correct_mod, "_atomic_write", flaky)
+        with pytest.raises(OSError):
+            apply_corrections(memory_tree, matches, "Datalux")
+
+        # No partial application: every file is back to its original content.
+        for p in paths:
+            assert p.read_text(encoding="utf-8") == originals[p], f"{p} not rolled back"
+
+    def test_happy_path_changes_all_files(self, memory_tree: Path):
+        matches = scan(memory_tree, "Quartz Indexer")
+        assert len(matches) >= 2
+        result = apply_corrections(memory_tree, matches, "Datalux")
+        assert result.files_changed == len(matches)
+        for m in matches:
+            assert "Quartz Indexer" not in (memory_tree / m.rel_path).read_text(encoding="utf-8")
+
+
 class TestCorrectAutoCommit:
     """`kbx correct` wires auto-commit on --apply only (issue #1).