diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint.py b/python/packages/core/agent_framework/_workflows/_checkpoint.py index 3de9460c86..09a6944bd4 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint.py @@ -7,6 +7,7 @@ import json import logging import os +import stat import uuid from collections.abc import Mapping from dataclasses import dataclass, field, fields @@ -360,16 +361,17 @@ def _read() -> dict[str, Any]: encoded_checkpoint = await asyncio.to_thread(_read) - from ._checkpoint_encoding import decode_checkpoint_value - - try: - decoded_checkpoint_dict = decode_checkpoint_value(encoded_checkpoint, allowed_types=self._allowed_types) - except WorkflowCheckpointException: - raise - checkpoint = WorkflowCheckpoint.from_dict(decoded_checkpoint_dict) + checkpoint = self._decode_checkpoint(encoded_checkpoint) logger.info(f"Loaded checkpoint {checkpoint_id} from {file_path}") return checkpoint + def _decode_checkpoint(self, encoded_checkpoint: dict[str, Any]) -> WorkflowCheckpoint: + """Decode an already-read checkpoint using this reader's deserialization policy.""" + from ._checkpoint_encoding import decode_checkpoint_value + + decoded = decode_checkpoint_value(encoded_checkpoint, allowed_types=self._allowed_types) + return WorkflowCheckpoint.from_dict(decoded) + async def list_checkpoints(self, *, workflow_name: str) -> list[WorkflowCheckpoint]: """List checkpoint objects for a given workflow name. @@ -423,18 +425,134 @@ def _delete() -> bool: async def get_latest(self, *, workflow_name: str) -> WorkflowCheckpoint | None: """Get the latest checkpoint for a given workflow name. + Select from stored metadata and decode only the selected checkpoint, using the same + data that was read during selection. Unreadable files and files without a valid workflow + name raise because they might hide a newer checkpoint for the requested workflow. + Invalid metadata for a known different workflow is ignored. + + File versions and directory membership are checked before accepting the scan. Detected + concurrent changes raise rather than returning a potentially stale checkpoint; the caller + can retry. Writes after validation belong to a subsequent snapshot. This relies on the + filesystem reporting file changes and writers using atomic replacement, as save does. + Checkpoints with equal timestamps retain directory iteration order. + Nested directories are included, with in-tree directory aliases visited only once. + Reading and decoding run in a worker thread so decoding does not block the event loop. + Args: workflow_name: The name of the workflow to get the latest checkpoint for. Returns: - The latest WorkflowCheckpoint object for the specified workflow name, or None if no checkpoints exist. + The latest WorkflowCheckpoint from the validated scan, or None if no matching checkpoints exist. + + Raises: + WorkflowCheckpointException: If the directory or a checkpoint cannot be read, a file's + workflow cannot be identified, matching metadata is invalid, the scan changes, + or the selected checkpoint cannot be decoded. """ - checkpoints = await self.list_checkpoints(workflow_name=workflow_name) - if not checkpoints: - return None - latest_checkpoint = max(checkpoints, key=lambda cp: datetime.fromisoformat(cp.timestamp)) - logger.debug(f"Latest checkpoint for workflow {workflow_name} is {latest_checkpoint.checkpoint_id}") - return latest_checkpoint + + def _version(info: os.stat_result) -> tuple[int, int, int, int, int]: + # Reads may change atime; it is not a content/version signal. + return info.st_dev, info.st_ino, info.st_size, info.st_mtime_ns, info.st_ctime_ns + + def _checkpoint_paths(root: Path) -> tuple[list[Path], dict[Path, tuple[int, int, int, int, int]]]: + # Path.glob suppresses some filesystem errors, including permission failures. + paths: list[Path] = [] + directories: dict[Path, tuple[int, int, int, int, int]] = {} + pending = [root] + while pending: + directory = pending.pop().resolve() + if not directory.is_relative_to(root): + raise WorkflowCheckpointException( + f"Checkpoint directory {directory} is outside storage root {root}" + ) + if directory in directories: + continue # An in-tree symlink may alias an already visited directory, including an ancestor. + directories[directory] = _version(directory.stat()) + with os.scandir(directory) as entries: + for entry in entries: + if entry.is_dir(): + pending.append(Path(entry.path)) + elif os.path.normcase(entry.name).endswith(".json"): + paths.append(Path(entry.path)) + return paths, directories + + def _unique_object(pairs: list[tuple[str, Any]]) -> dict[str, Any]: + result: dict[str, Any] = {} + for key, value in pairs: + if key in result: + raise ValueError(f"Duplicate checkpoint JSON key: {key}") + result[key] = value + return result + + def _read_latest() -> WorkflowCheckpoint | None: + root = self.storage_path.resolve() + paths, directories = _checkpoint_paths(root) + versions: dict[Path, tuple[int, int, int, int, int]] = {} + latest: tuple[datetime, dict[str, Any]] | None = None + for file_path in paths: + try: + self._validate_file_path(str(file_path.relative_to(root))[:-5]) + info = file_path.stat() + if not stat.S_ISREG(info.st_mode): + raise ValueError("checkpoint path is not a regular file") + version = _version(info) + with open(file_path) as f: + if _version(os.fstat(f.fileno())) != version: + raise ValueError("checkpoint changed before reading") + stored: dict[str, Any] = json.load(f, object_pairs_hook=_unique_object) + if _version(os.fstat(f.fileno())) != version: + raise ValueError("checkpoint changed while reading") + versions[file_path] = version + if not isinstance(stored, dict) or not isinstance(stored.get("workflow_name"), str): + raise ValueError("checkpoint workflow_name must be a string") + # save() writes plain metadata at the root; an envelope could hide different metadata. + if "__pickled__" in stored or "__type__" in stored: + raise ValueError("Top-level pickle markers are not checkpoint metadata") + if stored["workflow_name"] != workflow_name: + continue + timestamp = datetime.fromisoformat(stored["timestamp"]) + checkpoint_id = stored["checkpoint_id"] + if not isinstance(checkpoint_id, str): + raise ValueError("checkpoint_id must be a string") + if self._validate_file_path(checkpoint_id) != file_path.resolve(): + raise ValueError("checkpoint_id does not match its checkpoint file") + if latest is None or timestamp > latest[0]: + latest = (timestamp, stored) + except Exception as e: + raise WorkflowCheckpointException( + f"Failed to read checkpoint metadata for workflow {workflow_name} from {file_path}: {e}" + ) from e + + for file_path, version in versions.items(): + if _version(file_path.stat()) != version: + raise WorkflowCheckpointException( + f"Checkpoint {file_path} changed during the scan; retry get_latest" + ) + current_paths, current_directories = _checkpoint_paths(root) + if ( + set(current_paths) != set(paths) + or current_directories != directories + or any(_version(path.stat()) != version for path, version in directories.items()) + or self.storage_path.resolve() != root + ): + raise WorkflowCheckpointException( + f"Checkpoint directory {self.storage_path} changed during the scan; retry get_latest" + ) + # Do not reopen through load(): it could decode a replacement rather than the selected data. + return self._decode_checkpoint(latest[1]) if latest else None + + try: + checkpoint = await asyncio.to_thread(_read_latest) + if checkpoint is None: + return None + except WorkflowCheckpointException: + raise + except Exception as e: + raise WorkflowCheckpointException( + f"Failed to get latest checkpoint for workflow {workflow_name} from {self.storage_path}: {e}" + ) from e + logger.debug(f"Latest checkpoint for workflow {workflow_name} is {checkpoint.checkpoint_id}") + return checkpoint async def list_checkpoint_ids(self, *, workflow_name: str) -> list[CheckpointID]: """List checkpoint IDs for a given workflow name. diff --git a/python/packages/core/tests/workflow/test_checkpoint.py b/python/packages/core/tests/workflow/test_checkpoint.py index 5f3da78cd1..6ed07f889e 100644 --- a/python/packages/core/tests/workflow/test_checkpoint.py +++ b/python/packages/core/tests/workflow/test_checkpoint.py @@ -1,7 +1,10 @@ # Copyright (c) Microsoft. All rights reserved. +import asyncio import json +import os import tempfile +import threading from dataclasses import dataclass from datetime import datetime, timezone from pathlib import Path @@ -1765,4 +1768,798 @@ async def test_file_checkpoint_storage_roundtrip_empty_collections(): assert loaded.pending_request_info_events == {} +@dataclass +class _UnlistedState: + """A state type outside the deserialization allow list.""" + + stage: str + + +def _stamped_checkpoint(state: dict[str, Any], timestamp: str) -> WorkflowCheckpoint: + return WorkflowCheckpoint( + workflow_name="recovery-workflow", + graph_signature_hash="test-hash", + state=state, + timestamp=timestamp, + ) + + +async def test_file_checkpoint_storage_get_latest_returns_the_newest(): + with tempfile.TemporaryDirectory() as temp_dir: + storage = FileCheckpointStorage(temp_dir) + await storage.save(_stamped_checkpoint({"stage": "first"}, "2026-01-01T10:00:00+00:00")) + newest_id = await storage.save(_stamped_checkpoint({"stage": "second"}, "2026-01-01T11:00:00+00:00")) + + latest = await storage.get_latest(workflow_name="recovery-workflow") + + assert latest is not None + assert latest.checkpoint_id == newest_id + + +async def test_file_checkpoint_storage_get_latest_does_not_silently_return_stale_state(): + """A newest checkpoint that cannot be decoded must raise, not fall back to an older one.""" + with tempfile.TemporaryDirectory() as temp_dir: + storage = FileCheckpointStorage(temp_dir) + await storage.save(_stamped_checkpoint({"stage": "first"}, "2026-01-01T10:00:00+00:00")) + await storage.save(_stamped_checkpoint({"app": _UnlistedState(stage="second")}, "2026-01-01T11:00:00+00:00")) + + with pytest.raises(WorkflowCheckpointException, match="allowed_checkpoint_types"): + await storage.get_latest(workflow_name="recovery-workflow") + + +async def test_file_checkpoint_storage_get_latest_without_checkpoints_is_none(): + with tempfile.TemporaryDirectory() as temp_dir: + storage = FileCheckpointStorage(temp_dir) + + assert await storage.get_latest(workflow_name="recovery-workflow") is None + + +@pytest.mark.parametrize("has_older_checkpoint", [False, True]) +async def test_file_checkpoint_storage_get_latest_rejects_saved_invalid_timestamp( + tmp_path: Path, has_older_checkpoint: bool +) -> None: + storage = FileCheckpointStorage(tmp_path) + if has_older_checkpoint: + await storage.save(_stamped_checkpoint({"stage": "first"}, "2026-01-01T10:00:00+00:00")) + checkpoint = _stamped_checkpoint({"stage": "second"}, "not-a-timestamp") + assert await storage.save(checkpoint) == checkpoint.checkpoint_id + + with pytest.raises(WorkflowCheckpointException, match=checkpoint.checkpoint_id) as exc: + await storage.get_latest(workflow_name="recovery-workflow") + + assert isinstance(exc.value.__cause__, ValueError) + + +@pytest.mark.parametrize( + ("field_name", "value", "remove"), + [ + ("timestamp", None, True), + ("timestamp", None, False), + ("timestamp", 42, False), + ("timestamp", "", False), + ("checkpoint_id", None, True), + ("checkpoint_id", None, False), + ("checkpoint_id", 42, False), + ("checkpoint_id", "../outside", False), + ], +) +async def test_file_checkpoint_storage_get_latest_rejects_invalid_metadata( + tmp_path: Path, field_name: str, value: Any, remove: bool +) -> None: + storage = FileCheckpointStorage(tmp_path) + await storage.save(_stamped_checkpoint({"stage": "first"}, "2026-01-01T10:00:00+00:00")) + newest_id = await storage.save(_stamped_checkpoint({"stage": "second"}, "2026-01-01T11:00:00+00:00")) + file_path = tmp_path / f"{newest_id}.json" + data = json.loads(file_path.read_text()) + if remove: + del data[field_name] + else: + data[field_name] = value + file_path.write_text(json.dumps(data)) + + with pytest.raises(WorkflowCheckpointException, match=newest_id) as exc: + await storage.get_latest(workflow_name="recovery-workflow") + + assert exc.value.__cause__ is not None + + +@pytest.mark.parametrize("other_workflow", [False, True]) +async def test_file_checkpoint_storage_get_latest_rejects_redirected_id(tmp_path: Path, other_workflow: bool) -> None: + storage = FileCheckpointStorage(tmp_path) + older = _stamped_checkpoint({"stage": "first"}, "2026-01-01T10:00:00+00:00") + if other_workflow: + older.workflow_name = "other-workflow" + older_id = await storage.save(older) + newest_id = await storage.save(_stamped_checkpoint({"stage": "second"}, "2026-01-01T11:00:00+00:00")) + file_path = tmp_path / f"{newest_id}.json" + data = json.loads(file_path.read_text()) + data["checkpoint_id"] = older_id + file_path.write_text(json.dumps(data)) + + with pytest.raises(WorkflowCheckpointException, match=newest_id): + await storage.get_latest(workflow_name="recovery-workflow") + + +@pytest.mark.parametrize("reverse_order", [False, True]) +async def test_file_checkpoint_storage_get_latest_wraps_incomparable_timestamps( + tmp_path: Path, reverse_order: bool +) -> None: + storage = FileCheckpointStorage(tmp_path) + timestamps = ["2026-01-01T10:00:00+00:00", "2026-01-01T11:00:00"] + if reverse_order: + timestamps.reverse() + for timestamp in timestamps: + await storage.save(_stamped_checkpoint({}, timestamp)) + + with pytest.raises(WorkflowCheckpointException) as exc: + await storage.get_latest(workflow_name="recovery-workflow") + + assert isinstance(exc.value.__cause__, TypeError) + + +async def test_file_checkpoint_storage_get_latest_ignores_unrelated_invalid_metadata(tmp_path: Path) -> None: + storage = FileCheckpointStorage(tmp_path) + unrelated = _stamped_checkpoint({}, "not-a-timestamp") + unrelated.workflow_name = "other-workflow" + await storage.save(unrelated) + newest_id = await storage.save(_stamped_checkpoint({}, "2026-01-01T11:00:00+00:00")) + + latest = await storage.get_latest(workflow_name="recovery-workflow") + + assert latest is not None + assert latest.checkpoint_id == newest_id + + +@pytest.mark.parametrize( + ("older_timestamp", "newer_timestamp"), + [ + ("2026-01-01T12:00:00+03:00", "2026-01-01T10:00:00+00:00"), + ("2026-01-01T09:00:00", "2026-01-01T10:00:00"), + ], +) +async def test_file_checkpoint_storage_get_latest_only_decodes_selected_checkpoint( + tmp_path: Path, older_timestamp: str, newer_timestamp: str +) -> None: + storage = FileCheckpointStorage(tmp_path) + await storage.save(_stamped_checkpoint({"app": _UnlistedState(stage="first")}, older_timestamp)) + newest_id = await storage.save(_stamped_checkpoint({"stage": "second"}, newer_timestamp)) + + latest = await storage.get_latest(workflow_name="recovery-workflow") + + assert latest is not None + assert latest.checkpoint_id == newest_id + assert latest.state == {"stage": "second"} + + +@pytest.mark.parametrize( + "contents", ["{ invalid json }", "[]", "null", "{}", '{"workflow_name": null}', '{"workflow_name": 1}'] +) +@pytest.mark.parametrize("has_older_checkpoint", [False, True]) +async def test_file_checkpoint_storage_get_latest_rejects_unidentifiable_files( + tmp_path: Path, contents: str, has_older_checkpoint: bool +) -> None: + storage = FileCheckpointStorage(tmp_path) + if has_older_checkpoint: + await storage.save(_stamped_checkpoint({"stage": "first"}, "2026-01-01T10:00:00+00:00")) + corrupt = tmp_path / "unidentifiable.json" + corrupt.write_text(contents) + + with pytest.raises(WorkflowCheckpointException, match="unidentifiable.json"): + await storage.get_latest(workflow_name="recovery-workflow") + + assert corrupt.read_text() == contents + + +async def test_file_checkpoint_storage_get_latest_rejects_missing_directory(tmp_path: Path) -> None: + storage = FileCheckpointStorage(tmp_path / "checkpoints") + storage.storage_path.rmdir() + + with pytest.raises(WorkflowCheckpointException) as exc: + await storage.get_latest(workflow_name="recovery-workflow") + + assert isinstance(exc.value.__cause__, FileNotFoundError) + + +@pytest.mark.parametrize("deny_directory", [False, True]) +async def test_file_checkpoint_storage_get_latest_surfaces_permission_errors( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, deny_directory: bool +) -> None: + import builtins + + storage = FileCheckpointStorage(tmp_path) + await storage.save(_stamped_checkpoint({}, "2026-01-01T11:00:00+00:00")) + + def deny(*args: Any, **kwargs: Any) -> Any: + raise PermissionError("checkpoint storage read denied") + + monkeypatch.setattr(os if deny_directory else builtins, "scandir" if deny_directory else "open", deny) + + with pytest.raises(WorkflowCheckpointException) as exc: + await storage.get_latest(workflow_name="recovery-workflow") + + assert isinstance(exc.value.__cause__, PermissionError) + + +@pytest.mark.parametrize("change", ["create", "replace", "delete", "overwrite"]) +@pytest.mark.parametrize("nested", [False, True]) +async def test_file_checkpoint_storage_get_latest_rejects_changes_during_read( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, change: str, nested: bool +) -> None: + storage = FileCheckpointStorage(tmp_path) + checkpoint = _stamped_checkpoint({"stage": "first"}, "2026-01-01T10:00:00+00:00") + if nested: + (tmp_path / "nested").mkdir() + checkpoint.checkpoint_id = f"nested/{checkpoint.checkpoint_id}" + checkpoint_id = await storage.save(checkpoint) + path = tmp_path / f"{checkpoint_id}.json" + original_load = json.load + changed = False + + def change_after_read(stream: Any, **kwargs: Any) -> Any: + nonlocal changed + data = original_load(stream, **kwargs) + if not changed: + changed = True + if change == "create": + newer = _stamped_checkpoint({}, "2026-01-01T11:00:00+00:00") + newer.checkpoint_id = f"nested/{newer.checkpoint_id}" if nested else newer.checkpoint_id + (tmp_path / f"{newer.checkpoint_id}.json").write_text(json.dumps(newer.to_dict())) + elif change == "delete": + path.unlink() + else: + replacement = checkpoint.to_dict() + replacement["state"] = {"stage": "other"} # Same identity, timestamp and payload length. + if change == "replace": + temporary = tmp_path / "replacement.tmp" + temporary.write_text(json.dumps(replacement, indent=2, ensure_ascii=False)) + os.replace(temporary, path) + else: + path.write_text(json.dumps(replacement, indent=2, ensure_ascii=False)) + return data + + monkeypatch.setattr(json, "load", change_after_read) + + with pytest.raises(WorkflowCheckpointException): + await storage.get_latest(workflow_name="recovery-workflow") + + +async def test_file_checkpoint_storage_get_latest_revalidates_previously_read_files( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + storage = FileCheckpointStorage(tmp_path) + for hour in (10, 11): + await storage.save(_stamped_checkpoint({"stage": "first"}, f"2026-01-01T{hour}:00:00+00:00")) + original_load = json.load + seen: list[Path] = [] + + def change_previous_file(stream: Any, **kwargs: Any) -> Any: + data = original_load(stream, **kwargs) + seen.append(Path(stream.name)) + if len(seen) == 2: + previous = seen[0] + previous.write_text(previous.read_text().replace('"first"', '"other"')) + return data + + monkeypatch.setattr(json, "load", change_previous_file) + + with pytest.raises(WorkflowCheckpointException): + await storage.get_latest(workflow_name="recovery-workflow") + + +async def test_file_checkpoint_storage_get_latest_reads_selected_file_once( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + storage = FileCheckpointStorage(tmp_path) + checkpoint = _stamped_checkpoint({"stage": "first"}, "2026-01-01T10:00:00+00:00") + await storage.save(checkpoint) + original_load = json.load + reads = 0 + + def count_reads(stream: Any, **kwargs: Any) -> Any: + nonlocal reads + reads += 1 + return original_load(stream, **kwargs) + + monkeypatch.setattr(json, "load", count_reads) + + assert await storage.get_latest(workflow_name="recovery-workflow") == checkpoint + assert reads == 1 + + +@pytest.mark.parametrize("change", ["create", "replace", "delete"]) +async def test_file_checkpoint_storage_get_latest_decodes_the_validated_snapshot( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, change: str +) -> None: + from agent_framework._workflows import _checkpoint_encoding + + storage = FileCheckpointStorage(tmp_path) + checkpoint = _stamped_checkpoint({"stage": "first"}, "2026-01-01T10:00:00+00:00") + checkpoint_id = await storage.save(checkpoint) + path = tmp_path / f"{checkpoint_id}.json" + original_decode = _checkpoint_encoding.decode_checkpoint_value + + def change_after_validation(value: Any, **kwargs: Any) -> Any: + if change == "delete": + path.unlink() + else: + newer = _stamped_checkpoint({"stage": "second"}, "2026-01-01T11:00:00+00:00") + target = tmp_path / f"{newer.checkpoint_id}.json" if change == "create" else path + target.write_text(json.dumps(newer.to_dict())) + return original_decode(value, **kwargs) + + monkeypatch.setattr(_checkpoint_encoding, "decode_checkpoint_value", change_after_validation) + + assert await storage.get_latest(workflow_name="recovery-workflow") == checkpoint + + +@pytest.mark.parametrize("overwrite", [False, True]) +async def test_file_checkpoint_storage_get_latest_detects_concurrent_save( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, overwrite: bool +) -> None: + reader = FileCheckpointStorage(tmp_path) + writer = FileCheckpointStorage(tmp_path) + original = _stamped_checkpoint({"stage": "first"}, "2026-01-01T10:00:00+00:00") + await writer.save(original) + read_started = threading.Event() + write_finished = threading.Event() + original_load = json.load + + def pause_after_read(stream: Any, **kwargs: Any) -> Any: + data = original_load(stream, **kwargs) + read_started.set() + if not write_finished.wait(timeout=5): + raise TimeoutError("concurrent writer did not finish") + return data + + monkeypatch.setattr(json, "load", pause_after_read) + reading = asyncio.create_task(reader.get_latest(workflow_name="recovery-workflow")) + try: + assert await asyncio.to_thread(read_started.wait, 5) + newer = _stamped_checkpoint({"stage": "second"}, "2026-01-01T11:00:00+00:00") + if overwrite: + newer.checkpoint_id = original.checkpoint_id + await writer.save(newer) + finally: + write_finished.set() + with pytest.raises(WorkflowCheckpointException): + await reading + + +@pytest.mark.parametrize( + "checkpoint_id", ["nested/latest", "nested/deeper/latest", ".hidden/latest", "nested/./latest", "nested/../latest"] +) +async def test_file_checkpoint_storage_get_latest_finds_nested_checkpoint(tmp_path: Path, checkpoint_id: str) -> None: + storage = FileCheckpointStorage(tmp_path) + await storage.save(_stamped_checkpoint({}, "2026-01-01T10:00:00+00:00")) + (tmp_path / "nested").mkdir() + (tmp_path / checkpoint_id).parent.mkdir(parents=True, exist_ok=True) + newer = _stamped_checkpoint({}, "2026-01-01T11:00:00+00:00") + newer.checkpoint_id = checkpoint_id + await storage.save(newer) + + assert await storage.get_latest(workflow_name="recovery-workflow") == newer + assert await storage.load(newer.checkpoint_id) == newer + + +@pytest.mark.parametrize("newest_location", ["", "nested", "sibling"]) +async def test_file_checkpoint_storage_get_latest_compares_across_directories( + tmp_path: Path, newest_location: str +) -> None: + storage = FileCheckpointStorage(tmp_path) + newest = None + for location in ("", "nested", "sibling"): + (tmp_path / location).mkdir(exist_ok=True) + checkpoint = _stamped_checkpoint({}, "2026-01-01T10:00:00+00:00") + checkpoint.checkpoint_id = f"{location}/latest" if location else "latest" + if location == newest_location: + checkpoint.timestamp = "2026-01-01T11:00:00+00:00" + newest = checkpoint + await storage.save(checkpoint) + + assert newest is not None + assert await storage.get_latest(workflow_name="recovery-workflow") == newest + + +async def test_file_checkpoint_storage_get_latest_supports_relative_storage_root( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.chdir(tmp_path) + storage = FileCheckpointStorage("storage") + (storage.storage_path / "nested").mkdir() + checkpoint = _stamped_checkpoint({}, "2026-01-01T10:00:00+00:00") + checkpoint.checkpoint_id = "nested/latest" + await storage.save(checkpoint) + + assert await storage.get_latest(workflow_name="recovery-workflow") == checkpoint + + +@pytest.mark.parametrize("defect", ["timestamp", "checkpoint_id", "json", "payload"]) +async def test_file_checkpoint_storage_get_latest_rejects_invalid_nested_checkpoint( + tmp_path: Path, defect: str +) -> None: + storage = FileCheckpointStorage(tmp_path) + await storage.save(_stamped_checkpoint({}, "2026-01-01T10:00:00+00:00")) + (tmp_path / "nested").mkdir() + newer = _stamped_checkpoint({}, "2026-01-01T11:00:00+00:00") + newer.checkpoint_id = "nested/latest" + if defect == "payload": + newer.state = {"app": _UnlistedState(stage="second")} + await storage.save(newer) + path = tmp_path / "nested" / "latest.json" + if defect in ("timestamp", "checkpoint_id"): + data = json.loads(path.read_text()) + data[defect] = "latest" # Invalid time, or an ID pointing to the root instead of this file. + path.write_text(json.dumps(data)) + elif defect == "json": + path.write_text("{ incomplete }") + + with pytest.raises(WorkflowCheckpointException): + await storage.get_latest(workflow_name="recovery-workflow") + + +async def test_file_checkpoint_storage_get_latest_surfaces_nested_directory_permission_error( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + storage = FileCheckpointStorage(tmp_path) + await storage.save(_stamped_checkpoint({}, "2026-01-01T10:00:00+00:00")) + nested = tmp_path / "nested" + nested.mkdir() + original_scandir = os.scandir + + def deny_nested(path: Any) -> Any: + if Path(path) == nested: + raise PermissionError("nested checkpoint directory read denied") + return original_scandir(path) + + monkeypatch.setattr(os, "scandir", deny_nested) + with pytest.raises(WorkflowCheckpointException) as exc: + await storage.get_latest(workflow_name="recovery-workflow") + assert isinstance(exc.value.__cause__, PermissionError) + + +@pytest.mark.parametrize("change", ["create_directory", "delete_directory", "transient_file"]) +async def test_file_checkpoint_storage_get_latest_detects_nested_directory_changes( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, change: str +) -> None: + storage = FileCheckpointStorage(tmp_path) + checkpoint = _stamped_checkpoint({}, "2026-01-01T10:00:00+00:00") + await storage.save(checkpoint) + nested = tmp_path / "nested" + (nested / "empty").mkdir(parents=True) + original_load = json.load + + def change_after_read(stream: Any, **kwargs: Any) -> Any: + data = original_load(stream, **kwargs) + if change == "create_directory": + (nested / "new").mkdir() + elif change == "delete_directory": + (nested / "empty").rmdir() + else: + transient = nested / "transient.json" + transient.write_text(json.dumps(checkpoint.to_dict())) + transient.unlink() + return data + + monkeypatch.setattr(json, "load", change_after_read) + with pytest.raises(WorkflowCheckpointException, match="changed during the scan"): + await storage.get_latest(workflow_name="recovery-workflow") + + +@pytest.mark.skipif(os.name == "nt", reason="directory symlink creation requires Windows privileges") +async def test_file_checkpoint_storage_get_latest_handles_directory_aliases_and_cycles( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + storage = FileCheckpointStorage(tmp_path) + nested = tmp_path / "nested" + nested.mkdir() + (tmp_path / "alias").symlink_to(nested, target_is_directory=True) + (nested / "back").symlink_to(tmp_path, target_is_directory=True) + checkpoint = _stamped_checkpoint({}, "2026-01-01T10:00:00+00:00") + checkpoint.checkpoint_id = "alias/latest" + await storage.save(checkpoint) + original_load = json.load + reads = 0 + + def count_reads(stream: Any, **kwargs: Any) -> Any: + nonlocal reads + reads += 1 + return original_load(stream, **kwargs) + + monkeypatch.setattr(json, "load", count_reads) + + assert await storage.get_latest(workflow_name="recovery-workflow") == checkpoint + assert reads == 1 + + +async def test_file_checkpoint_storage_get_latest_revalidates_directories_after_final_scan( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + from collections.abc import Generator, Iterator + from contextlib import contextmanager + + storage = FileCheckpointStorage(tmp_path) + checkpoint = _stamped_checkpoint({}, "2026-01-01T10:00:00+00:00") + await storage.save(checkpoint) + nested = tmp_path / "nested" + nested.mkdir() + original_scandir = os.scandir + nested_scans = 0 + + @contextmanager + def change_after_scan(path: Any) -> Generator[Iterator[os.DirEntry[str]], None, None]: + nonlocal nested_scans + with original_scandir(path) as entries: + yield entries + if Path(path) == nested: + nested_scans += 1 + if nested_scans == 2: + transient = nested / "transient.json" + transient.write_text(json.dumps(checkpoint.to_dict())) + transient.unlink() + + monkeypatch.setattr(os, "scandir", change_after_scan) + with pytest.raises(WorkflowCheckpointException, match="changed during the scan"): + await storage.get_latest(workflow_name="recovery-workflow") + + +@pytest.mark.skipif(os.name == "nt", reason="directory symlink creation requires Windows privileges") +@pytest.mark.parametrize("retarget", [False, True]) +async def test_file_checkpoint_storage_get_latest_validates_symlinked_storage_root( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, retarget: bool +) -> None: + root = tmp_path / "root" + root.mkdir() + alias = tmp_path / "alias" + alias.symlink_to(root, target_is_directory=True) + other = tmp_path / "other" + other.mkdir() + storage = FileCheckpointStorage(alias) + checkpoint = _stamped_checkpoint({}, "2026-01-01T10:00:00+00:00") + await storage.save(checkpoint) + original_scandir = os.scandir + scans = 0 + + def retarget_after_read(path: Any) -> Any: + nonlocal scans + scans += 1 + if retarget and scans == 2: + alias.unlink() + alias.symlink_to(other, target_is_directory=True) + return original_scandir(path) + + monkeypatch.setattr(os, "scandir", retarget_after_read) + if retarget: + with pytest.raises(WorkflowCheckpointException, match="changed during the scan"): + await storage.get_latest(workflow_name="recovery-workflow") + else: + assert await storage.get_latest(workflow_name="recovery-workflow") == checkpoint + + +@pytest.mark.skipif(os.name == "nt", reason="symlink creation requires Windows privileges") +@pytest.mark.parametrize("directory_link", [False, True]) +async def test_file_checkpoint_storage_get_latest_rejects_links_outside_storage( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, directory_link: bool +) -> None: + import builtins + + storage = FileCheckpointStorage(tmp_path / "storage") + outside = tmp_path / "outside" + outside.mkdir() + (outside / "latest.json").write_text("{ untrusted }") + if directory_link: + (storage.storage_path / "link").symlink_to(outside, target_is_directory=True) + else: + (storage.storage_path / "link.json").symlink_to(outside / "latest.json") + + def unexpected_read(*args: Any, **kwargs: Any) -> Any: + pytest.fail("must reject the escaped path before reading JSON") + + monkeypatch.setattr(builtins, "open", unexpected_read) + with pytest.raises(WorkflowCheckpointException, match="outside storage root|Invalid checkpoint ID"): + await storage.get_latest(workflow_name="recovery-workflow") + + +@pytest.mark.parametrize("has_checkpoint", [False, True]) +async def test_file_checkpoint_storage_get_latest_ignores_unrelated_directories( + tmp_path: Path, has_checkpoint: bool +) -> None: + storage = FileCheckpointStorage(tmp_path) + checkpoint = _stamped_checkpoint({}, "2026-01-01T10:00:00+00:00") if has_checkpoint else None + if checkpoint is not None: + await storage.save(checkpoint) + (tmp_path / "unrelated" / "empty.json").mkdir(parents=True) + (tmp_path / "unrelated" / "notes.txt").write_text("not a checkpoint") + + assert await storage.get_latest(workflow_name="recovery-workflow") == checkpoint + + +async def test_file_checkpoint_storage_get_latest_keeps_event_loop_responsive( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + from agent_framework._workflows import _checkpoint_encoding + + storage = FileCheckpointStorage(tmp_path) + checkpoint = _stamped_checkpoint({}, "2026-01-01T10:00:00+00:00") + await storage.save(checkpoint) + loop = asyncio.get_running_loop() + decode_started = asyncio.Event() + release_decode = threading.Event() + original_decode = _checkpoint_encoding.decode_checkpoint_value + + def pause_decode(value: Any, **kwargs: Any) -> Any: + loop.call_soon_threadsafe(decode_started.set) + if not release_decode.wait(timeout=5): + raise TimeoutError("event loop could not resume while decoding") + return original_decode(value, **kwargs) + + async def resume_decode() -> None: + await decode_started.wait() + release_decode.set() + + monkeypatch.setattr(_checkpoint_encoding, "decode_checkpoint_value", pause_decode) + resuming = asyncio.create_task(resume_decode()) + try: + assert await storage.get_latest(workflow_name="recovery-workflow") == checkpoint + finally: + release_decode.set() + resuming.cancel() + await asyncio.gather(resuming, return_exceptions=True) + + +async def test_file_checkpoint_storage_get_latest_rejects_duplicate_json_keys(tmp_path: Path) -> None: + storage = FileCheckpointStorage(tmp_path) + await storage.save(_stamped_checkpoint({}, "2026-01-01T10:00:00+00:00")) + (tmp_path / "ambiguous.json").write_text( + '{"workflow_name": "recovery-workflow", "workflow_name": "other-workflow"}' + ) + + with pytest.raises(WorkflowCheckpointException, match="Duplicate checkpoint JSON key"): + await storage.get_latest(workflow_name="recovery-workflow") + + +@pytest.mark.parametrize("header_workflow", ["recovery-workflow", "other-workflow"]) +async def test_file_checkpoint_storage_get_latest_rejects_root_pickle_envelopes( + tmp_path: Path, header_workflow: str +) -> None: + from agent_framework._workflows._checkpoint_encoding import _pickle_to_base64 + + storage = FileCheckpointStorage(tmp_path) + older = _stamped_checkpoint({"stage": "first"}, "2026-01-01T10:00:00+00:00") + await storage.save(older) + newer = _stamped_checkpoint({"stage": "second"}, "2026-01-01T11:00:00+00:00") + await storage.save(newer) + data = newer.to_dict() + data.update({ + "workflow_name": header_workflow, + "__pickled__": _pickle_to_base64(older.to_dict()), + "__type__": "builtins:dict", + }) + (tmp_path / f"{newer.checkpoint_id}.json").write_text(json.dumps(data)) + + with pytest.raises(WorkflowCheckpointException, match="pickle"): + await storage.get_latest(workflow_name="recovery-workflow") + + +async def test_file_checkpoint_storage_get_latest_rejects_replacement_before_open( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + import builtins + + storage = FileCheckpointStorage(tmp_path) + checkpoint = _stamped_checkpoint({}, "2026-01-01T10:00:00+00:00") + checkpoint_id = await storage.save(checkpoint) + path = tmp_path / f"{checkpoint_id}.json" + original_open = builtins.open + + def replace_before_open(file: Any, *args: Any, **kwargs: Any) -> Any: + if Path(file) == path: + replacement = tmp_path / "replacement.tmp" + replacement.write_text(path.read_text()) + os.replace(replacement, path) + return original_open(file, *args, **kwargs) + + monkeypatch.setattr(builtins, "open", replace_before_open) + + with pytest.raises(WorkflowCheckpointException, match="changed before reading"): + await storage.get_latest(workflow_name="recovery-workflow") + + +async def test_file_checkpoint_storage_get_latest_revalidates_empty_scan( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + from collections.abc import Generator, Iterator + from contextlib import contextmanager + + reader = FileCheckpointStorage(tmp_path) + writer = FileCheckpointStorage(tmp_path) + enumerated = threading.Event() + write_finished = threading.Event() + original_scandir = os.scandir + first_scan = True + + @contextmanager + def pause_empty_scan(path: Any) -> Generator[Iterator[os.DirEntry[str]], None, None]: + nonlocal first_scan + with original_scandir(path) as entries: + snapshot = list(entries) + if first_scan: + first_scan = False + enumerated.set() + if not write_finished.wait(timeout=5): + raise TimeoutError("concurrent writer did not finish") + yield iter(snapshot) + + monkeypatch.setattr(os, "scandir", pause_empty_scan) + reading = asyncio.create_task(reader.get_latest(workflow_name="recovery-workflow")) + try: + assert await asyncio.to_thread(enumerated.wait, 5) + await writer.save(_stamped_checkpoint({}, "2026-01-01T10:00:00+00:00")) + finally: + write_finished.set() + with pytest.raises(WorkflowCheckpointException, match="changed during the scan"): + await reading + + +async def test_file_checkpoint_storage_get_latest_ignores_uncommitted_temporary_files(tmp_path: Path) -> None: + storage = FileCheckpointStorage(tmp_path) + checkpoint = _stamped_checkpoint({}, "2026-01-01T10:00:00+00:00") + await storage.save(checkpoint) + (tmp_path / "unfinished.json.tmp").write_text("{ incomplete }") + + assert await storage.get_latest(workflow_name="recovery-workflow") == checkpoint + + +@pytest.mark.skipif(not hasattr(os, "mkfifo"), reason="requires POSIX FIFO support") +async def test_file_checkpoint_storage_get_latest_rejects_non_regular_files(tmp_path: Path) -> None: + storage = FileCheckpointStorage(tmp_path) + os.mkfifo(tmp_path / "pipe.json") + + with pytest.raises(WorkflowCheckpointException, match="regular file"): + await storage.get_latest(workflow_name="recovery-workflow") + + +async def test_file_checkpoint_storage_get_latest_rejects_undecodable_text(tmp_path: Path) -> None: + storage = FileCheckpointStorage(tmp_path) + (tmp_path / "invalid-encoding.json").write_bytes(b"\xff\xfe\x00") + + with pytest.raises(WorkflowCheckpointException) as exc: + await storage.get_latest(workflow_name="recovery-workflow") + assert isinstance(exc.value.__cause__, UnicodeError) + + +@pytest.mark.parametrize("checkpoint_id", ["", "./custom-id"]) +async def test_file_checkpoint_storage_get_latest_preserves_valid_path_aliases( + tmp_path: Path, checkpoint_id: str +) -> None: + storage = FileCheckpointStorage(tmp_path) + checkpoint = _stamped_checkpoint({}, "2026-01-01T11:00:00+00:00") + checkpoint.checkpoint_id = checkpoint_id + await storage.save(checkpoint) + + assert await storage.get_latest(workflow_name="recovery-workflow") == checkpoint + + +async def test_file_checkpoint_storage_get_latest_honors_reader_allowed_types(tmp_path: Path) -> None: + writer = FileCheckpointStorage(tmp_path) + checkpoint = _stamped_checkpoint({"app": _UnlistedState(stage="second")}, "2026-01-01T11:00:00+00:00") + await writer.save(checkpoint) + reader = FileCheckpointStorage( + tmp_path, allowed_checkpoint_types=[f"{_UnlistedState.__module__}:{_UnlistedState.__qualname__}"] + ) + + assert await reader.get_latest(workflow_name="recovery-workflow") == checkpoint + + +async def test_file_checkpoint_storage_list_checkpoints_still_skips_undecodable_entries(): + """list_checkpoints keeps returning what it can read, so one bad file cannot break listing.""" + with tempfile.TemporaryDirectory() as temp_dir: + storage = FileCheckpointStorage(temp_dir) + readable_id = await storage.save(_stamped_checkpoint({"stage": "first"}, "2026-01-01T10:00:00+00:00")) + await storage.save(_stamped_checkpoint({"app": _UnlistedState(stage="second")}, "2026-01-01T11:00:00+00:00")) + + listed = await storage.list_checkpoints(workflow_name="recovery-workflow") + + assert [cp.checkpoint_id for cp in listed] == [readable_id] + + # endregion