diff --git a/tests/tools/test_checkpoint_manager.py b/tests/tools/test_checkpoint_manager.py index f36f97e112..b87bbcf970 100644 --- a/tests/tools/test_checkpoint_manager.py +++ b/tests/tools/test_checkpoint_manager.py @@ -338,6 +338,27 @@ class TestSafeRestore: assert "README.md" in result["skipped_user_edits"] assert "main.py" in result["restored_files"] + def test_safe_restore_finds_ledger_under_walked_key(self, mgr, work_dir, tmp_path): + base = self._checkpoint(mgr, work_dir) + + # A generic project marker in an ancestor dir (e.g. a stray + # package.json in /tmp) makes record_agent_write's marker walk key + # the ledger to the ancestor, while restore reads the exact dir's + # hash. Safe restore must still find the ledger, or it silently + # degrades to a full restore and overwrites user edits. + (tmp_path / "package.json").write_text("{}\n") + + (work_dir / "main.py").write_text("agent version\n") + (work_dir / "README.md").write_text("user hand edit\n") + mgr.record_agent_write(str(work_dir / "main.py")) + + result = mgr.restore(str(work_dir), base, safe=True) + assert result["success"] is True + assert (work_dir / "main.py").read_text() == "print('hello')\n" + assert (work_dir / "README.md").read_text() == "user hand edit\n" + assert result["restored_files"] == ["main.py"] + assert "README.md" in result["skipped_user_edits"] + def test_safe_restore_skips_file_user_edited_after_agent(self, mgr, work_dir): base = self._checkpoint(mgr, work_dir) diff --git a/tools/checkpoint_manager.py b/tools/checkpoint_manager.py index 790e5124b8..bee7ab758d 100644 --- a/tools/checkpoint_manager.py +++ b/tools/checkpoint_manager.py @@ -638,6 +638,19 @@ class CheckpointManager: return {"success": False, "error": f"Could not compute changed files: {err}"} ledger = _load_ledger(p.store, p.dir_hash) + if not ledger: + # record_agent_write keys the ledger by the marker-walk result + # (get_working_dir_for_path), which can diverge from the dir this + # restore names — e.g. a markerless project dir under an ancestor + # that carries a generic project marker, or a restore of a repo + # subdir while the write was recorded at the repo root. Fall back + # to the walked key so safe restore never silently degrades to a + # full restore (which overwrites user edits) just because the + # ledger landed under a sibling key. + walked = self.get_working_dir_for_path(working_dir) + walk_key = _project_hash(walked) + if walk_key != p.dir_hash: + ledger = _load_ledger(p.store, walk_key) if not ledger: return {"success": True, "restore": [], "skipped": [], "ledger_empty": True} out: Dict[str, List[str]] = {"restore": [], "skipped": []}