diff --git a/tests/tools/test_checkpoint_manager.py b/tests/tools/test_checkpoint_manager.py index 5c6db10c0119..cb7af5b9aefd 100644 --- a/tests/tools/test_checkpoint_manager.py +++ b/tests/tools/test_checkpoint_manager.py @@ -353,6 +353,17 @@ def test_restore_to_previous(self, mgr, work_dir): assert result["success"] is True assert (work_dir / "main.py").read_text() == "original\n" + def test_restore_removes_file_added_after_checkpoint(self, mgr, work_dir): + mgr.ensure_checkpoint(str(work_dir), "original state") + added = work_dir / "added-after-checkpoint.txt" + added.write_text("temporary\n") + + cps = mgr.list_checkpoints(str(work_dir)) + result = mgr.restore(str(work_dir), cps[0]["hash"]) + + assert result["success"] is True + assert not added.exists() + def test_restore_invalid_hash(self, mgr, work_dir): mgr.ensure_checkpoint(str(work_dir), "initial") result = mgr.restore(str(work_dir), "deadbeef1234") @@ -375,6 +386,70 @@ def test_restore_creates_pre_rollback_snapshot(self, mgr, work_dir): assert len(all_cps) >= 2 assert "pre-rollback" in all_cps[0]["reason"] + def test_restore_preserves_excluded_files(self, mgr, work_dir): + (work_dir / "main.py").write_text("v1\n") + mgr.ensure_checkpoint(str(work_dir), "v1") + mgr.new_turn() + + excluded = work_dir / "node_modules" / "package" / "index.js" + excluded.parent.mkdir(parents=True) + excluded.write_text("keep me\n") + (work_dir / "main.py").write_text("v2\n") + + cps = mgr.list_checkpoints(str(work_dir)) + result = mgr.restore(str(work_dir), cps[0]["hash"]) + + assert result["success"] is True + assert (work_dir / "main.py").read_text() == "v1\n" + assert excluded.read_text() == "keep me\n" + + def test_restore_survives_retention_rewriting_target_hash( + self, work_dir, checkpoint_base, monkeypatch, + ): + monkeypatch.setattr("tools.checkpoint_manager.CHECKPOINT_BASE", checkpoint_base) + m = CheckpointManager(enabled=True, max_snapshots=2) + + (work_dir / "main.py").write_text("v1\n") + assert m.ensure_checkpoint(str(work_dir), "v1") is True + m.new_turn() + (work_dir / "main.py").write_text("v2\n") + assert m.ensure_checkpoint(str(work_dir), "v2") is True + m.new_turn() + + target_hash = m.list_checkpoints(str(work_dir))[-1]["hash"] + (work_dir / "main.py").write_text("v3\n") + + result = m.restore(str(work_dir), target_hash) + + assert result["success"] is True + assert (work_dir / "main.py").read_text() == "v1\n" + + def test_file_restore_survives_retention_rewriting_target_hash( + self, work_dir, checkpoint_base, monkeypatch, + ): + monkeypatch.setattr("tools.checkpoint_manager.CHECKPOINT_BASE", checkpoint_base) + m = CheckpointManager(enabled=True, max_snapshots=2) + + (work_dir / "main.py").write_text("v1\n") + (work_dir / "other.py").write_text("other-v1\n") + assert m.ensure_checkpoint(str(work_dir), "v1") is True + m.new_turn() + (work_dir / "main.py").write_text("v2\n") + (work_dir / "other.py").write_text("other-v2\n") + assert m.ensure_checkpoint(str(work_dir), "v2") is True + m.new_turn() + + target_hash = m.list_checkpoints(str(work_dir))[-1]["hash"] + (work_dir / "main.py").write_text("v3\n") + (work_dir / "other.py").write_text("other-v3\n") + + result = m.restore(str(work_dir), target_hash, file_path="main.py") + + assert result["success"] is True + assert result["file"] == "main.py" + assert (work_dir / "main.py").read_text() == "v1\n" + assert (work_dir / "other.py").read_text() == "other-v3\n" + def test_tilde_path_supports_diff_and_restore_flow( self, checkpoint_base, fake_home, monkeypatch, ): diff --git a/tools/checkpoint_manager.py b/tools/checkpoint_manager.py index b5ce04c6a64e..b9ecad9fa355 100644 --- a/tools/checkpoint_manager.py +++ b/tools/checkpoint_manager.py @@ -816,32 +816,61 @@ def restore(self, working_dir: str, commit_hash: str, file_path: str = None) -> return {"success": False, "error": f"Checkpoint '{commit_hash}' not found", "debug": err or None} - # Take a pre-rollback snapshot so you can undo the undo. - self._take(abs_dir, f"pre-rollback snapshot (restoring to {commit_hash[:8]})") + # Resolve the target tree and reason before taking the pre-rollback + # snapshot. That snapshot may cross max_snapshots and trigger _prune, + # which rebuilds the retained commit chain, expires reflogs, and runs + # gc --prune=now. Pin the tree with a temporary ref so it remains + # reachable until the restore finishes. + ok_tree, target_tree, tree_err = _run_git( + ["rev-parse", "--verify", f"{commit_hash}^{{tree}}"], store, abs_dir, + ) + if not ok_tree or not target_tree: + return {"success": False, "error": f"Checkpoint '{commit_hash}' has no readable tree", + "debug": tree_err or None} + ok_reason, reason_out, _ = _run_git( + ["log", "--format=%s", "-1", commit_hash], store, abs_dir, + ) + target_reason = reason_out if ok_reason else "unknown" dir_hash = _project_hash(abs_dir) index_file = _index_path(store, dir_hash) - - restore_target = file_path if file_path else "." - ok, stdout, err = _run_git( - ["checkout", commit_hash, "--", restore_target], - store, abs_dir, timeout=_GIT_TIMEOUT * 2, - index_file=index_file, + restore_ref = f"refs/hermes-restore/{dir_hash}-{os.getpid()}-{time.time_ns()}" + ok_pin, _, pin_err = _run_git( + ["update-ref", restore_ref, target_tree], store, abs_dir, ) + if not ok_pin: + return {"success": False, "error": "Could not pin checkpoint tree for restore", + "debug": pin_err or None} - if not ok: - return {"success": False, "error": f"Restore failed: {err}", - "debug": err or None} + try: + # Take a pre-rollback snapshot so you can undo the undo. + self._take(abs_dir, f"pre-rollback snapshot (restoring to {commit_hash[:8]})") - ok2, reason_out, _ = _run_git( - ["log", "--format=%s", "-1", commit_hash], store, abs_dir, - ) - reason = reason_out if ok2 else "unknown" + if file_path: + restore_args = ["checkout", restore_ref, "--", file_path] + else: + # ``git checkout -- .`` updates paths present in the + # target tree but leaves behind files added after the checkpoint. + # The per-project index currently represents the pre-rollback + # snapshot, so read-tree's reset+update mode can reconcile the + # complete tracked tree, including removing those added paths, + # without touching excluded files. + restore_args = ["read-tree", "--reset", "-u", restore_ref] + ok, stdout, err = _run_git( + restore_args, store, abs_dir, timeout=_GIT_TIMEOUT * 2, + index_file=index_file, + ) + + if not ok: + return {"success": False, "error": f"Restore failed: {err}", + "debug": err or None} + finally: + _run_git(["update-ref", "-d", restore_ref], store, abs_dir) result = { "success": True, "restored_to": commit_hash[:8], - "reason": reason, + "reason": target_reason, "directory": abs_dir, } if file_path: