diff --git a/tests/tools/test_checkpoint_manager.py b/tests/tools/test_checkpoint_manager.py index d0af8e9e58..038a8436f7 100644 --- a/tests/tools/test_checkpoint_manager.py +++ b/tests/tools/test_checkpoint_manager.py @@ -393,6 +393,102 @@ class TestSafeRestore: assert "README.md" in result["skipped_user_edits"] assert "agent.txt" in result["restored_files"] + # -- size-capped files ------------------------------------------------- + # + # ``max_file_size_mb`` keeps generated assets out of a checkpoint + # (test_max_file_size_mb_skips_large_files). Safe restore then sees such a + # file as "changed since the checkpoint and absent from it" — the same + # shape as a file Hermes created — and the delete branch treats absence as + # proof of authorship. For a capped file that is wrong: no checkpoint holds + # a copy, so deleting it destroys the only one. + + @staticmethod + def _capped_mgr(checkpoint_base, monkeypatch, cap_mb=1): + monkeypatch.setattr("tools.checkpoint_manager.CHECKPOINT_BASE", checkpoint_base) + return CheckpointManager(enabled=True, max_snapshots=50, max_file_size_mb=cap_mb) + + def test_safe_restore_keeps_an_oversize_file_the_cap_excluded( + self, work_dir, checkpoint_base, monkeypatch, + ): + """The file is in no checkpoint, so deleting it is unrecoverable.""" + m = self._capped_mgr(checkpoint_base, monkeypatch) + corpus = work_dir / "corpus.jsonl" + corpus.write_bytes(b"original\n" + b"x" * (2 * 1024 * 1024)) + base = self._checkpoint(m, work_dir) + + corpus.write_bytes(b"agent appended\n" + b"y" * (2 * 1024 * 1024)) + m.record_agent_write(str(corpus)) + + result = m.restore(str(work_dir), base, safe=True) + + assert result["success"] is True + assert corpus.exists(), "safe restore deleted a file no checkpoint holds" + assert corpus.read_bytes().startswith(b"agent appended") + assert "corpus.jsonl" in result.get("skipped_oversize", []) + + def test_safe_restore_does_not_report_a_kept_file_as_restored( + self, work_dir, checkpoint_base, monkeypatch, + ): + """Misreporting is what made the loss silent — the user was told + "Restored" for a path that had just been unlinked.""" + m = self._capped_mgr(checkpoint_base, monkeypatch) + corpus = work_dir / "corpus.jsonl" + corpus.write_bytes(b"o" * (2 * 1024 * 1024)) + base = self._checkpoint(m, work_dir) + + corpus.write_bytes(b"a" * (2 * 1024 * 1024)) + m.record_agent_write(str(corpus)) + (work_dir / "main.py").write_text("agent version\n") + m.record_agent_write(str(work_dir / "main.py")) + + result = m.restore(str(work_dir), base, safe=True) + + assert result["restored_files"] == ["main.py"] + assert (work_dir / "main.py").read_text() == "print('hello')\n" + + def test_safe_restore_still_reverts_a_file_that_grew_past_the_cap( + self, work_dir, checkpoint_base, monkeypatch, + ): + """Regression guard for the tempting wrong fix. + + Dropping every oversize path from the restore plan would also strand + this one — small enough to be checkpointed, then bloated by the agent. + It IS in the checkpoint, so a prior version exists and reverting to it + is exactly what the user asked for. The rule has to key on "absent from + the checkpoint", not on "large now". + """ + m = self._capped_mgr(checkpoint_base, monkeypatch) + grew = work_dir / "grew.txt" + grew.write_text("precious original\n") + base = self._checkpoint(m, work_dir) + + grew.write_bytes(b"agent bloated it\n" + b"b" * (2 * 1024 * 1024)) + m.record_agent_write(str(grew)) + + result = m.restore(str(work_dir), base, safe=True) + + assert result["success"] is True + assert grew.read_text() == "precious original\n" + assert "grew.txt" in result["restored_files"] + assert "grew.txt" not in result.get("skipped_oversize", []) + + def test_safe_restore_still_deletes_a_small_agent_created_file( + self, work_dir, checkpoint_base, monkeypatch, + ): + """The delete branch keeps working for the case it was written for.""" + m = self._capped_mgr(checkpoint_base, monkeypatch) + base = self._checkpoint(m, work_dir) + + scratch = work_dir / "scratch.txt" + scratch.write_text("agent scratch\n") + m.record_agent_write(str(scratch)) + + result = m.restore(str(work_dir), base, safe=True) + + assert not scratch.exists() + assert "scratch.txt" in result["restored_files"] + assert "skipped_oversize" not in result + def test_unsafe_restore_overwrites_everything(self, mgr, work_dir): base = self._checkpoint(mgr, work_dir) (work_dir / "main.py").write_text("agent version\n") diff --git a/tools/checkpoint_manager.py b/tools/checkpoint_manager.py index 7149eebc90..6b5657bfa2 100644 --- a/tools/checkpoint_manager.py +++ b/tools/checkpoint_manager.py @@ -1157,12 +1157,26 @@ class CheckpointManager: # Hermes-created files absent from it (delete to restore state). checkout_targets: List[str] = [] delete_targets: List[str] = [] + kept_oversize: List[str] = [] for rel in restore_paths: ok_in_commit, _, _ = _run_git( ["cat-file", "-e", f"{commit_hash}:{rel}"], store, abs_dir, allowed_returncodes={1, 128}, ) - (checkout_targets if ok_in_commit else delete_targets).append(rel) + if ok_in_commit: + checkout_targets.append(rel) + elif self._exceeds_size_cap(Path(abs_dir) / rel): + # Absent from the checkpoint because ``max_file_size_mb`` + # kept it out (_drop_oversize_from_index), not because + # Hermes created it. Deleting it would not restore a prior + # state — no checkpoint holds one — it would destroy the + # only copy. The ledger records a content hash, not whether + # a write created or modified the file, so an oversize path + # cannot be proven agent-created; leaving it costs a stale + # file, deleting it costs the file. + kept_oversize.append(rel) + else: + delete_targets.append(rel) for rel in delete_targets: try: target = Path(abs_dir) / rel @@ -1203,8 +1217,16 @@ class CheckpointManager: if file_path: result["file"] = file_path if restore_paths is not None: - result["restored_files"] = restore_paths + # Only what was actually acted on. A kept oversize path was not + # restored, and reporting it as such is how the data loss above + # stayed silent: the user was told "Restored" for a file that had + # just been unlinked. + result["restored_files"] = [ + rel for rel in restore_paths if rel not in kept_oversize + ] result["skipped_user_edits"] = skipped_user_edits + if kept_oversize: + result["skipped_oversize"] = kept_oversize return result def get_working_dir_for_path(self, file_path: str) -> str: @@ -1363,6 +1385,22 @@ class CheckpointManager: return True + def _exceeds_size_cap(self, path: Path) -> bool: + """Whether *path* is larger than ``max_file_size_mb``. + + The same test :meth:`_drop_oversize_from_index` applies when building a + checkpoint, so "excluded from the checkpoint" and "refused deletion at + restore" agree on one definition. A cap of 0 disables it, and an + unstattable path is not claimed to be oversize. + """ + cap = self.max_file_size_mb * 1024 * 1024 + if cap <= 0: + return False + try: + return path.stat().st_size > cap + except OSError: + return False + def _drop_oversize_from_index( self, store: Path, working_dir: str, index_file: Path, ) -> None: