Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 21 additions & 8 deletions src/kb/correct.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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
Expand Down
37 changes: 37 additions & 0 deletions tests/test_correct.py
Original file line number Diff line number Diff line change
Expand Up @@ -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).

Expand Down