From b03e996cdca3475648eb5b0095c62bb57b47ffd3 Mon Sep 17 00:00:00 2001 From: dev404ai Date: Sun, 23 Aug 2026 14:05:07 +0300 Subject: [PATCH 1/3] Python: raise instead of silently returning a stale checkpoint from get_latest --- .../agent_framework/_workflows/_checkpoint.py | 36 ++++++++++-- .../core/tests/workflow/test_checkpoint.py | 58 +++++++++++++++++++ 2 files changed, 89 insertions(+), 5 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint.py b/python/packages/core/agent_framework/_workflows/_checkpoint.py index 3de9460c86d..dc98381f61d 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint.py @@ -423,18 +423,44 @@ def _delete() -> bool: async def get_latest(self, *, workflow_name: str) -> WorkflowCheckpoint | None: """Get the latest checkpoint for a given workflow name. + The latest checkpoint is identified from stored metadata, which does not require + decoding any checkpoint payload, and only that checkpoint is then loaded. A checkpoint + that cannot be decoded therefore surfaces as an error rather than being skipped in + favour of an older one: silently resuming from earlier state is worse than failing. + 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. + + Raises: + WorkflowCheckpointException: If the latest checkpoint exists but cannot be loaded. """ - checkpoints = await self.list_checkpoints(workflow_name=workflow_name) - if not checkpoints: + + def _latest_checkpoint_id() -> CheckpointID | None: + latest: tuple[datetime, CheckpointID] | None = None + for file_path in self.storage_path.glob("*.json"): + try: + with open(file_path) as f: + stored = json.load(f) + if stored.get("workflow_name") != workflow_name: + continue + timestamp = datetime.fromisoformat(stored["timestamp"]) + checkpoint_id = stored["checkpoint_id"] + except Exception as e: + logger.warning(f"Failed to read checkpoint metadata from {file_path}: {e}") + continue + if latest is None or timestamp > latest[0]: + latest = (timestamp, checkpoint_id) + return latest[1] if latest else None + + latest_checkpoint_id = await asyncio.to_thread(_latest_checkpoint_id) + if latest_checkpoint_id is None: 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 + + logger.debug(f"Latest checkpoint for workflow {workflow_name} is {latest_checkpoint_id}") + return await self.load(latest_checkpoint_id) 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 5f3da78cd1d..894b6ba5c6f 100644 --- a/python/packages/core/tests/workflow/test_checkpoint.py +++ b/python/packages/core/tests/workflow/test_checkpoint.py @@ -1765,4 +1765,62 @@ 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): + 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 + + +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 From b5d9ec7498bbbda055e23e0b7e88723a02afb777 Mon Sep 17 00:00:00 2001 From: dev404ai Date: Sun, 6 Sep 2026 10:32:51 +0300 Subject: [PATCH 2/3] Python: harden checkpoint recovery validation --- .../agent_framework/_workflows/_checkpoint.py | 134 +++-- .../core/tests/workflow/test_checkpoint.py | 468 +++++++++++++++++- 2 files changed, 573 insertions(+), 29 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint.py b/python/packages/core/agent_framework/_workflows/_checkpoint.py index dc98381f61d..76d12f374f6 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,44 +425,120 @@ def _delete() -> bool: async def get_latest(self, *, workflow_name: str) -> WorkflowCheckpoint | None: """Get the latest checkpoint for a given workflow name. - The latest checkpoint is identified from stored metadata, which does not require - decoding any checkpoint payload, and only that checkpoint is then loaded. A checkpoint - that cannot be decoded therefore surfaces as an error rather than being skipped in - favour of an older one: silently resuming from earlier state is worse than failing. + 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. + The directory must be flat: nested directories are rejected because their checkpoints + would otherwise be invisible to this scan. Use load with an explicit ID for a nested path. 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 latest checkpoint exists but cannot be loaded. + 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. """ - def _latest_checkpoint_id() -> CheckpointID | None: - latest: tuple[datetime, CheckpointID] | None = None - for file_path in self.storage_path.glob("*.json"): + 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() -> list[Path]: + # Path.glob suppresses some filesystem errors, including permission failures. + with os.scandir(self.storage_path) as entries: + paths: list[Path] = [] + for entry in entries: + if entry.is_dir(): + raise WorkflowCheckpointException( + f"Checkpoint directory must be flat; found directory {entry.path}" + ) + if os.path.normcase(entry.name).endswith(".json"): + paths.append(Path(entry.path)) + return paths + + 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() -> dict[str, Any] | None: + directory_version = _version(self.storage_path.stat()) + paths = _checkpoint_paths() + 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(file_path.name[:-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: - stored = json.load(f) - if stored.get("workflow_name") != workflow_name: + 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: - logger.warning(f"Failed to read checkpoint metadata from {file_path}: {e}") - continue - if latest is None or timestamp > latest[0]: - latest = (timestamp, checkpoint_id) + 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" + ) + if set(_checkpoint_paths()) != set(paths) or _version(self.storage_path.stat()) != directory_version: + raise WorkflowCheckpointException( + f"Checkpoint directory {self.storage_path} changed during the scan; retry get_latest" + ) return latest[1] if latest else None - latest_checkpoint_id = await asyncio.to_thread(_latest_checkpoint_id) - if latest_checkpoint_id is None: - return None - - logger.debug(f"Latest checkpoint for workflow {workflow_name} is {latest_checkpoint_id}") - return await self.load(latest_checkpoint_id) + try: + encoded_checkpoint = await asyncio.to_thread(_read_latest) + if encoded_checkpoint is None: + return None + # Do not reopen through load(): it could decode a replacement rather than the selected data. + checkpoint = self._decode_checkpoint(encoded_checkpoint) + 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 894b6ba5c6f..cc467ec865b 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 @@ -1800,7 +1803,7 @@ async def test_file_checkpoint_storage_get_latest_does_not_silently_return_stale 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): + with pytest.raises(WorkflowCheckpointException, match="allowed_checkpoint_types"): await storage.get_latest(workflow_name="recovery-workflow") @@ -1811,6 +1814,469 @@ async def test_file_checkpoint_storage_get_latest_without_checkpoints_is_none(): 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"]) +async def test_file_checkpoint_storage_get_latest_rejects_changes_during_read( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, change: str +) -> None: + 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_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") + (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 + + +async def test_file_checkpoint_storage_get_latest_rejects_hidden_nested_checkpoint(tmp_path: Path) -> 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" + await storage.save(newer) + + with pytest.raises(WorkflowCheckpointException, match="flat"): + await storage.get_latest(workflow_name="recovery-workflow") + assert await storage.load(newer.checkpoint_id) == newer + + +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: From 3574c9b6c57b536ce580a560d5786cc4813e4e05 Mon Sep 17 00:00:00 2001 From: dev404ai Date: Sun, 6 Sep 2026 11:27:27 +0300 Subject: [PATCH 3/3] Python: preserve nested checkpoint recovery and async decoding --- .../agent_framework/_workflows/_checkpoint.py | 60 ++-- .../core/tests/workflow/test_checkpoint.py | 281 +++++++++++++++++- 2 files changed, 314 insertions(+), 27 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint.py b/python/packages/core/agent_framework/_workflows/_checkpoint.py index 76d12f374f6..09a6944bd42 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint.py @@ -435,8 +435,8 @@ async def get_latest(self, *, workflow_name: str) -> WorkflowCheckpoint | None: 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. - The directory must be flat: nested directories are rejected because their checkpoints - would otherwise be invisible to this scan. Use load with an explicit ID for a nested path. + 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. @@ -454,18 +454,27 @@ 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() -> list[Path]: + 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. - with os.scandir(self.storage_path) as entries: - paths: list[Path] = [] - for entry in entries: - if entry.is_dir(): - raise WorkflowCheckpointException( - f"Checkpoint directory must be flat; found directory {entry.path}" - ) - if os.path.normcase(entry.name).endswith(".json"): - paths.append(Path(entry.path)) - return paths + 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] = {} @@ -475,14 +484,14 @@ def _unique_object(pairs: list[tuple[str, Any]]) -> dict[str, Any]: result[key] = value return result - def _read_latest() -> dict[str, Any] | None: - directory_version = _version(self.storage_path.stat()) - paths = _checkpoint_paths() + 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(file_path.name[:-5]) + 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") @@ -519,18 +528,23 @@ def _read_latest() -> dict[str, Any] | None: raise WorkflowCheckpointException( f"Checkpoint {file_path} changed during the scan; retry get_latest" ) - if set(_checkpoint_paths()) != set(paths) or _version(self.storage_path.stat()) != directory_version: + 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" ) - return latest[1] if latest else None + # 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: - encoded_checkpoint = await asyncio.to_thread(_read_latest) - if encoded_checkpoint is None: + checkpoint = await asyncio.to_thread(_read_latest) + if checkpoint is None: return None - # Do not reopen through load(): it could decode a replacement rather than the selected data. - checkpoint = self._decode_checkpoint(encoded_checkpoint) except WorkflowCheckpointException: raise except Exception as e: diff --git a/python/packages/core/tests/workflow/test_checkpoint.py b/python/packages/core/tests/workflow/test_checkpoint.py index cc467ec865b..6ed07f889ec 100644 --- a/python/packages/core/tests/workflow/test_checkpoint.py +++ b/python/packages/core/tests/workflow/test_checkpoint.py @@ -1981,11 +1981,15 @@ def deny(*args: Any, **kwargs: Any) -> Any: @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 + 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 @@ -1998,6 +2002,7 @@ def change_after_read(stream: Any, **kwargs: Any) -> Any: 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() @@ -2120,17 +2125,285 @@ def pause_after_read(stream: Any, **kwargs: Any) -> Any: await reading -async def test_file_checkpoint_storage_get_latest_rejects_hidden_nested_checkpoint(tmp_path: Path) -> None: +@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, match="flat"): + with pytest.raises(WorkflowCheckpointException): await storage.get_latest(workflow_name="recovery-workflow") - assert await storage.load(newer.checkpoint_id) == newer + + +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: