From 33fe35407a62dbd80a31bd8075039ae13131a150 Mon Sep 17 00:00:00 2001 From: pratikwayase Date: Sun, 18 Oct 2026 11:11:47 +0530 Subject: [PATCH 1/3] fix(checkpoint): raise instead of silently returning stale checkpoint from get_latest --- .../agent_framework/_workflows/_checkpoint.py | 152 ++++++++++- .../core/tests/workflow/test_checkpoint.py | 247 ++++++++++++++++++ 2 files changed, 385 insertions(+), 14 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint.py b/python/packages/core/agent_framework/_workflows/_checkpoint.py index cd78725d6a1..d36c983a612 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint.py @@ -8,6 +8,7 @@ import json import logging import os +import stat import threading import time import uuid @@ -779,16 +780,23 @@ def _read() -> dict[str, Any]: f"Checkpoint file for {checkpoint_id} is not valid JSON and cannot be loaded." ) from ex - 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. + + Shared by :meth:`load` and :meth:`get_latest` so both apply the same + ``allowed_checkpoint_types`` policy. Callers must have already validated the + on-disk state if snapshot consistency matters; this method only handles + deserialization. + """ + 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. @@ -845,18 +853,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]: + # On Windows, Path.stat() and os.fstat() use different Win32 APIs and can + # return different values for st_dev (volume serial number) and st_ino + # (file ID). Only st_size + st_mtime_ns + st_ctime_ns are reliable across + # both APIs, which is sufficient to detect both content changes (mtime, + # size) and file replacement (ctime on Windows is creation time, which + # updates on os.replace). Symlink escapes are caught separately by + # _validate_file_path and the is_relative_to(root) check in _checkpoint_paths. + return 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]]]: + paths: list[Path] = [] + directories: dict[Path, tuple[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 + 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]] = {} + 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, encoding="utf-8") as f: + stored: dict[str, Any] = json.load(f, object_pairs_hook=_unique_object) + 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") + 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" + ) + 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 6e13a414f8b..1912ddfd4df 100644 --- a/python/packages/core/tests/workflow/test_checkpoint.py +++ b/python/packages/core/tests/workflow/test_checkpoint.py @@ -24,6 +24,7 @@ WorkflowCheckpointException, WorkflowEvent, ) +from agent_framework._workflows._checkpoint_encoding import encode_checkpoint_value from agent_framework._workflows._runner_context import WorkflowMessage @@ -62,6 +63,37 @@ class _TestCustomData: tags: list[str] +@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, + ) + + +def _write_checkpoint_file(storage_path: Path, checkpoint: WorkflowCheckpoint) -> None: + """Write a checkpoint JSON file directly, bypassing save()'s decode validation. + + Simulates the real-world scenario the bug is about: a file on disk that this + reader's allowed_checkpoint_types cannot decode (written by an older framework + version, a different storage config, or a type since unregistered). + """ + encoded = encode_checkpoint_value(checkpoint.to_dict()) + file_path = storage_path / f"{checkpoint.checkpoint_id}.json" + file_path.write_text( + json.dumps(encoded, indent=2, ensure_ascii=False), + encoding="utf-8", + ) + + # region test WorkflowCheckpoint @@ -3449,3 +3481,218 @@ async def test_await_signal_through_cancellation_returns_immediately_for_a_resol checkpoint_module._await_signal_through_cancellation(resolved), # pyright: ignore[reportPrivateUsage] timeout=5, ) + + +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")) + _write_checkpoint_file( + Path(temp_dir), + _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_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) + + +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_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( + ("older_timestamp", "newer_timestamp"), + [ + ("2026-01-01T12:00:00+03:00", "2026-01-01T10:00:00+00:00"), # tz-aware comparison + ("2026-01-01T09:00:00", "2026-01-01T10:00:00"), # naive comparison + ], +) +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) + _write_checkpoint_file( + tmp_path, + _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( + ("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( + "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 + + +@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") + encoded_newer = encode_checkpoint_value(newer.to_dict()) + (tmp_path / f"{newer.checkpoint_id}.json").write_text( + json.dumps(encoded_newer, indent=2, ensure_ascii=False), + encoding="utf-8", + ) + elif change == "delete": + path.unlink() + else: + replacement = checkpoint.to_dict() + replacement["state"] = {"stage": "other-different-content"} + encoded_replacement = encode_checkpoint_value(replacement) + if change == "replace": + temporary = tmp_path / "replacement.tmp" + temporary.write_text( + json.dumps(encoded_replacement, indent=2, ensure_ascii=False), + encoding="utf-8", + ) + os.replace(temporary, path) + else: + path.write_text( + json.dumps(encoded_replacement, indent=2, ensure_ascii=False), + encoding="utf-8", + ) + 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_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")) + _write_checkpoint_file( + Path(temp_dir), + _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] + + +async def test_file_checkpoint_storage_get_latest_honors_reader_allowed_types(tmp_path: Path) -> None: + checkpoint = _stamped_checkpoint({"app": _UnlistedState(stage="second")}, "2026-01-01T11:00:00+00:00") + _write_checkpoint_file(tmp_path, 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_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 From 0562e1735b305cc0cd95cbbecc47993a7e581ace Mon Sep 17 00:00:00 2001 From: pratikwayase Date: Sun, 18 Oct 2026 12:38:10 +0530 Subject: [PATCH 2/3] address copilot comments --- .../agent_framework/_workflows/_checkpoint.py | 38 ++++++++--- .../core/tests/workflow/test_checkpoint.py | 63 +++++++++++++++++++ 2 files changed, 93 insertions(+), 8 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint.py b/python/packages/core/agent_framework/_workflows/_checkpoint.py index d36c983a612..ddcd24f482f 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint.py @@ -911,13 +911,30 @@ def _checkpoint_paths(root: Path) -> tuple[list[Path], dict[Path, tuple[int, int 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 _make_duplicate_capturing_hook() -> tuple[ + Callable[[list[tuple[str, Any]]], dict[str, Any]], + list[str], + ]: + """Return a JSON ``object_pairs_hook`` that records duplicate keys instead of raising. + + The hook is paired with a list that collects duplicate key names as they are + encountered during parsing. The caller decides when to raise based on workflow + ownership, so a foreign-workflow file with duplicate keys does not block + recovery of healthy checkpoints for the requested workflow (per the documented + policy that invalid metadata for a known different workflow is ignored). + """ + duplicate_keys: list[str] = [] + + def hook(pairs: list[tuple[str, Any]]) -> dict[str, Any]: + result: dict[str, Any] = {} + for key, value in pairs: + if key in result: + duplicate_keys.append(key) + else: + result[key] = value + return result + + return hook, duplicate_keys def _read_latest() -> WorkflowCheckpoint | None: root = self.storage_path.resolve() @@ -931,15 +948,20 @@ def _read_latest() -> WorkflowCheckpoint | None: if not stat.S_ISREG(info.st_mode): raise ValueError("checkpoint path is not a regular file") version = _version(info) + hook, duplicate_keys = _make_duplicate_capturing_hook() with open(file_path, encoding="utf-8") as f: - stored: dict[str, Any] = json.load(f, object_pairs_hook=_unique_object) + stored: dict[str, Any] = json.load(f, object_pairs_hook=hook) versions[file_path] = version + if "workflow_name" in duplicate_keys: + raise ValueError("Duplicate checkpoint JSON key: workflow_name") if not isinstance(stored, dict) or not isinstance(stored.get("workflow_name"), str): raise ValueError("checkpoint workflow_name must be a string") 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 + if duplicate_keys: + raise ValueError(f"Duplicate checkpoint JSON key: {duplicate_keys[0]}") timestamp = datetime.fromisoformat(stored["timestamp"]) checkpoint_id = stored["checkpoint_id"] if not isinstance(checkpoint_id, str): diff --git a/python/packages/core/tests/workflow/test_checkpoint.py b/python/packages/core/tests/workflow/test_checkpoint.py index 1912ddfd4df..887447edd74 100644 --- a/python/packages/core/tests/workflow/test_checkpoint.py +++ b/python/packages/core/tests/workflow/test_checkpoint.py @@ -3623,6 +3623,11 @@ async def test_file_checkpoint_storage_get_latest_rejects_changes_during_read( def change_after_read(stream: Any, **kwargs: Any) -> Any: nonlocal changed data = original_load(stream, **kwargs) + # Release the read handle before deleting or replacing the file on Windows: + # an open file handle blocks unlink() and os.replace() with PermissionError, + # which get_latest() would wrap as WorkflowCheckpointException without + # actually exercising the mutation-detection path under test. + stream.close() if not changed: changed = True if change == "create": @@ -3696,3 +3701,61 @@ def count_reads(stream: Any, **kwargs: Any) -> Any: 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_rejects_duplicate_json_keys(tmp_path: Path) -> None: + """A matching-workflow file with a duplicate workflow_name key must raise. + + A duplicated workflow_name key means ownership is ambiguous: the file might + belong to the requested workflow even if the last-wins value says otherwise. + Reject before trusting the parsed workflow_name. + """ + 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") + + +async def test_file_checkpoint_storage_get_latest_rejects_duplicate_metadata_keys_in_matching_workflow( + tmp_path: Path, +) -> None: + """A matching-workflow file with duplicate non-ownership keys must raise. + + Once ownership is established and matches the requested workflow, key + uniqueness is enforced across the whole file. + """ + storage = FileCheckpointStorage(tmp_path) + await storage.save(_stamped_checkpoint({}, "2026-01-01T10:00:00+00:00")) + (tmp_path / "duplicate-timestamp.json").write_text( + '{"workflow_name": "recovery-workflow", ' + '"timestamp": "2026-01-01T10:00:00+00:00", ' + '"timestamp": "2026-01-01T10:00:00+00:00"}' + ) + with pytest.raises(WorkflowCheckpointException, match="Duplicate checkpoint JSON key"): + await storage.get_latest(workflow_name="recovery-workflow") + + +async def test_file_checkpoint_storage_get_latest_ignores_duplicate_keys_in_foreign_workflow( + tmp_path: Path, +) -> None: + """Duplicate keys in a file belonging to a different workflow must not block recovery. + + Per the documented policy, invalid metadata for a known different workflow is + ignored. A foreign file with duplicate keys must be skipped after ownership is + determined, not rejected during parsing. + """ + storage = FileCheckpointStorage(tmp_path) + newest_id = await storage.save( + _stamped_checkpoint({"stage": "first"}, "2026-01-01T11:00:00+00:00") + ) + (tmp_path / "foreign.json").write_text( + '{"workflow_name": "other-workflow", ' + '"timestamp": "2026-01-01T10:00:00+00:00", ' + '"timestamp": "2026-01-01T10:00:00+00:00"}' + ) + latest = await storage.get_latest(workflow_name="recovery-workflow") + assert latest is not None + assert latest.checkpoint_id == newest_id From 4d317e18ef99e9eb753a5dc1a408dcd08ec0cf4d Mon Sep 17 00:00:00 2001 From: pratikwayase Date: Sun, 25 Oct 2026 15:29:31 +0530 Subject: [PATCH 3/3] fix review comments --- .../agent_framework/_workflows/_checkpoint.py | 121 +++++++--- .../core/tests/workflow/test_checkpoint.py | 212 +++++++++++++++--- 2 files changed, 273 insertions(+), 60 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint.py b/python/packages/core/agent_framework/_workflows/_checkpoint.py index ddcd24f482f..79155df7531 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint.py @@ -30,6 +30,17 @@ # Type alias for checkpoint IDs in case we want to change the # underlying type in the future (e.g., to UUID or a custom class) CheckpointID: TypeAlias = str +_MAX_GET_LATEST_ATTEMPTS = 3 +_GET_LATEST_RETRY_DELAY_SECONDS = 0.05 + + +class _CheckpointScanInvalidated(WorkflowCheckpointException): + """Internal signal that an optimistic scan was invalidated by concurrent changes. + + Raised only inside the worker thread used by ``get_latest`` and always consumed + by its retry loop. Callers see a plain :class:`WorkflowCheckpointException` + only when every attempt is invalidated. + """ @dataclass(slots=True) @@ -854,14 +865,20 @@ 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. + data that was read during selection. Files whose owning workflow cannot be + identified (unparseable JSON, a missing or non-string ``workflow_name``, or + duplicate ``workflow_name`` keys) are quarantined: they are skipped with a + warning so a single malformed or legacy file cannot permanently break recovery + for every workflow sharing this storage root. Invalid content in a file that IS + identified as belonging to the requested workflow still raises, and 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. + Concurrent changes under the storage root, including an atomic save for a + different workflow, invalidate the optimistic scan; the scan is retried + internally a bounded number of times before raising, so ordinary concurrent + checkpointing does not fail unrelated recovery. 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 @@ -875,9 +892,9 @@ async def get_latest(self, *, workflow_name: str) -> WorkflowCheckpoint | None: 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. + WorkflowCheckpointException: If a checkpoint owned by ``workflow_name`` cannot + be read or decoded, the directory cannot be read, matching metadata is + invalid, or the scan changes on every retry attempt. """ def _version(info: os.stat_result) -> tuple[int, int, int]: @@ -952,14 +969,25 @@ def _read_latest() -> WorkflowCheckpoint | None: with open(file_path, encoding="utf-8") as f: stored: dict[str, Any] = json.load(f, object_pairs_hook=hook) versions[file_path] = version - if "workflow_name" in duplicate_keys: - raise ValueError("Duplicate checkpoint JSON key: workflow_name") - if not isinstance(stored, dict) or not isinstance(stored.get("workflow_name"), str): - raise ValueError("checkpoint workflow_name must be a string") + owner: str | None = None + if ( + isinstance(stored, dict) + and "workflow_name" not in duplicate_keys + and isinstance(stored.get("workflow_name"), str) + ): + owner = stored["workflow_name"] + if owner is None: + logger.warning( + "Quarantining unidentifiable checkpoint file %s while finding " + "the latest checkpoint for workflow '%s'", + file_path, + workflow_name, + ) + continue + if owner != workflow_name: + continue 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 if duplicate_keys: raise ValueError(f"Duplicate checkpoint JSON key: {duplicate_keys[0]}") timestamp = datetime.fromisoformat(stored["timestamp"]) @@ -970,13 +998,28 @@ def _read_latest() -> WorkflowCheckpoint | None: raise ValueError("checkpoint_id does not match its checkpoint file") if latest is None or timestamp > latest[0]: latest = (timestamp, stored) + except (json.JSONDecodeError, UnicodeDecodeError) as ex: + logger.warning( + "Quarantining checkpoint file %s that is not valid JSON while " + "finding the latest checkpoint for workflow '%s': %s", + file_path, + workflow_name, + ex, + ) + continue 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( + try: + current = _version(file_path.stat()) + except FileNotFoundError as ex: + raise _CheckpointScanInvalidated( + f"Checkpoint {file_path} was removed during the scan; retry get_latest" + ) from ex + if current != version: + raise _CheckpointScanInvalidated( f"Checkpoint {file_path} changed during the scan; retry get_latest" ) current_paths, current_directories = _checkpoint_paths(root) @@ -986,23 +1029,43 @@ def _read_latest() -> WorkflowCheckpoint | None: or any(_version(path.stat()) != version for path, version in directories.items()) or self.storage_path.resolve() != root ): - raise WorkflowCheckpointException( + raise _CheckpointScanInvalidated( f"Checkpoint directory {self.storage_path} changed during the scan; retry get_latest" ) return self._decode_checkpoint(latest[1]) if latest else None - try: - checkpoint = await asyncio.to_thread(_read_latest) + attempt = 0 + while True: + attempt += 1 + try: + checkpoint = await asyncio.to_thread(_read_latest) + except _CheckpointScanInvalidated as exc: + if attempt >= _MAX_GET_LATEST_ATTEMPTS: + raise WorkflowCheckpointException( + f"Checkpoint directory {self.storage_path} kept changing during " + f"get_latest for workflow '{workflow_name}' after " + f"{_MAX_GET_LATEST_ATTEMPTS} attempts; retry once concurrent " + "checkpoint writes settle" + ) from exc + logger.debug( + "Checkpoint directory changed while finding the latest checkpoint for " + "workflow '%s' (attempt %d of %d); retrying", + workflow_name, + attempt, + _MAX_GET_LATEST_ATTEMPTS, + ) + await asyncio.sleep(_GET_LATEST_RETRY_DELAY_SECONDS * attempt) + continue + 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 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 + 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 887447edd74..417c539cda5 100644 --- a/python/packages/core/tests/workflow/test_checkpoint.py +++ b/python/packages/core/tests/workflow/test_checkpoint.py @@ -3577,20 +3577,58 @@ async def test_file_checkpoint_storage_get_latest_rejects_invalid_metadata( "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 +@pytest.mark.parametrize("has_healthy_checkpoint", [False, True]) +async def test_file_checkpoint_storage_get_latest_quarantines_unidentifiable_files( + tmp_path: Path, contents: str, has_healthy_checkpoint: bool ) -> None: + """A file whose owner cannot be identified is skipped, not fatal. + + One malformed or legacy file must not poison get_latest() for every workflow + sharing the storage root, so it is quarantined (skipped with a warning) and + left in place for inspection. Healthy checkpoints for the requested workflow + are still returned. + """ storage = FileCheckpointStorage(tmp_path) - if has_older_checkpoint: - await storage.save(_stamped_checkpoint({"stage": "first"}, "2026-01-01T10:00:00+00:00")) + healthy_id: str | None = None + if has_healthy_checkpoint: + healthy_id = 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") + + latest = await storage.get_latest(workflow_name="recovery-workflow") + + if has_healthy_checkpoint: + assert latest is not None + assert latest.checkpoint_id == healthy_id + else: + assert latest is None assert corrupt.read_text() == contents +async def test_file_checkpoint_storage_get_latest_shared_root_is_not_poisoned_by_foreign_files( + tmp_path: Path, +) -> None: + """Regression: corruption owned by one workflow must not break another's recovery. + + Two workflows share one storage root. 'recovery-workflow' is healthy, while the + root also holds an unparseable file and a valid-but-foreign checkpoint with a + newer timestamp. get_latest() must still return the healthy newest checkpoint. + """ + storage = FileCheckpointStorage(tmp_path) + ours_id = await storage.save(_stamped_checkpoint({"stage": "ours"}, "2026-01-01T10:00:00+00:00")) + (tmp_path / "legacy-broken.json").write_text("{ not valid json") + foreign = _stamped_checkpoint({"stage": "theirs"}, "2026-01-01T11:00:00+00:00") + foreign_dict = foreign.to_dict() + foreign_dict["workflow_name"] = "other-workflow" + (tmp_path / f"{foreign.checkpoint_id}.json").write_text(json.dumps(encode_checkpoint_value(foreign_dict))) + + latest = await storage.get_latest(workflow_name="recovery-workflow") + + assert latest is not None + assert latest.checkpoint_id == ours_id + assert latest.state == {"stage": "ours"} + + @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 @@ -3609,27 +3647,31 @@ def deny(*args: Any, **kwargs: Any) -> Any: 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( +@pytest.mark.parametrize("change", ["create", "replace", "overwrite"]) +async def test_file_checkpoint_storage_get_latest_returns_fresh_snapshot_after_concurrent_change( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, change: str ) -> None: + """A concurrent change invalidates the scan once; the internal retry returns fresh state. + + The first attempt detects the mutation and retries instead of failing; the + second attempt observes the change and returns the new, consistent snapshot. + """ 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 + mutated = False - def change_after_read(stream: Any, **kwargs: Any) -> Any: - nonlocal changed + def mutate_once(stream: Any, **kwargs: Any) -> Any: + nonlocal mutated data = original_load(stream, **kwargs) # Release the read handle before deleting or replacing the file on Windows: # an open file handle blocks unlink() and os.replace() with PermissionError, - # which get_latest() would wrap as WorkflowCheckpointException without - # actually exercising the mutation-detection path under test. + # so the mutation below would not land while the scan holds the handle. stream.close() - if not changed: - changed = True + if not mutated: + mutated = True if change == "create": newer = _stamped_checkpoint({}, "2026-01-01T11:00:00+00:00") encoded_newer = encode_checkpoint_value(newer.to_dict()) @@ -3637,8 +3679,6 @@ def change_after_read(stream: Any, **kwargs: Any) -> Any: json.dumps(encoded_newer, indent=2, ensure_ascii=False), encoding="utf-8", ) - elif change == "delete": - path.unlink() else: replacement = checkpoint.to_dict() replacement["state"] = {"stage": "other-different-content"} @@ -3657,10 +3697,117 @@ def change_after_read(stream: Any, **kwargs: Any) -> Any: ) return data - monkeypatch.setattr(json, "load", change_after_read) - with pytest.raises(WorkflowCheckpointException): + monkeypatch.setattr(json, "load", mutate_once) + + latest = await storage.get_latest(workflow_name="recovery-workflow") + + assert latest is not None + if change == "create": + assert latest.state == {} + else: + assert latest.state == {"stage": "other-different-content"} + + +async def test_file_checkpoint_storage_get_latest_returns_none_when_file_removed_during_read( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """After a concurrent deletion the fresh snapshot honestly contains no checkpoints. + + The removal invalidates the first attempt; the retry rescans, finds no + checkpoint files, and returns None instead of stale data or an error. + """ + storage = FileCheckpointStorage(tmp_path) + checkpoint_id = await storage.save(_stamped_checkpoint({"stage": "first"}, "2026-01-01T10:00:00+00:00")) + path = tmp_path / f"{checkpoint_id}.json" + original_load = json.load + mutated = False + + def mutate_once(stream: Any, **kwargs: Any) -> Any: + nonlocal mutated + data = original_load(stream, **kwargs) + # Release the read handle before unlinking on Windows. + stream.close() + if not mutated: + mutated = True + path.unlink() + return data + + monkeypatch.setattr(json, "load", mutate_once) + + assert await storage.get_latest(workflow_name="recovery-workflow") is None + + +async def test_file_checkpoint_storage_get_latest_raises_after_exhausting_retries( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """When the directory changes on every attempt, get_latest raises a clear error. + + Each attempt is invalidated by a new foreign checkpoint file appearing during + the scan. After the bounded number of internal retries is exhausted, the error + tells the caller the storage root is under heavy concurrent writes. + """ + storage = FileCheckpointStorage(tmp_path) + checkpoint_id = await storage.save(_stamped_checkpoint({}, "2026-01-01T10:00:00+00:00")) + counter = 0 + original_load = json.load + + def churn_on_every_read(stream: Any, **kwargs: Any) -> Any: + nonlocal counter + data = original_load(stream, **kwargs) + stream.close() + churn = _stamped_checkpoint({}, f"2026-01-01T12:0{counter}:00+00:00") + churn_dict = churn.to_dict() + churn_dict["workflow_name"] = "churn-workflow" + (tmp_path / f"{churn.checkpoint_id}.json").write_text(json.dumps(encode_checkpoint_value(churn_dict))) + counter += 1 + return data + + monkeypatch.setattr(json, "load", churn_on_every_read) + + with pytest.raises(WorkflowCheckpointException, match="kept changing"): await storage.get_latest(workflow_name="recovery-workflow") + assert (tmp_path / f"{checkpoint_id}.json").exists() + + +async def test_file_checkpoint_storage_get_latest_survives_concurrent_foreign_save( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """An atomic save for a different workflow during the scan must not fail recovery. + + Regression for the shared-root invalidation comment: the optimistic scan is + retried internally, the second attempt skips the foreign checkpoint, and the + requested workflow's checkpoint is returned instead of a 'changed during the + scan' error. + """ + storage = FileCheckpointStorage(tmp_path) + checkpoint_id = await storage.save(_stamped_checkpoint({"stage": "ours"}, "2026-01-01T10:00:00+00:00")) + original_load = json.load + injected = False + + def inject_foreign_save(stream: Any, **kwargs: Any) -> Any: + nonlocal injected + data = original_load(stream, **kwargs) + stream.close() + if not injected: + injected = True + foreign = _stamped_checkpoint({"stage": "theirs"}, "2026-01-01T11:00:00+00:00") + foreign_dict = foreign.to_dict() + foreign_dict["workflow_name"] = "another-workflow" + (tmp_path / f"{foreign.checkpoint_id}.json").write_text( + json.dumps(encode_checkpoint_value(foreign_dict)), + encoding="utf-8", + ) + return data + + monkeypatch.setattr(json, "load", inject_foreign_save) + + latest = await storage.get_latest(workflow_name="recovery-workflow") + + assert latest is not None + assert latest.checkpoint_id == checkpoint_id + assert latest.state == {"stage": "ours"} + 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.""" @@ -3703,20 +3850,25 @@ def count_reads(stream: Any, **kwargs: Any) -> Any: assert reads == 1 -async def test_file_checkpoint_storage_get_latest_rejects_duplicate_json_keys(tmp_path: Path) -> None: - """A matching-workflow file with a duplicate workflow_name key must raise. +async def test_file_checkpoint_storage_get_latest_quarantines_duplicate_workflow_name_keys( + tmp_path: Path, +) -> None: + """A file with a duplicated workflow_name key has ambiguous ownership. - A duplicated workflow_name key means ownership is ambiguous: the file might - belong to the requested workflow even if the last-wins value says otherwise. - Reject before trusting the parsed workflow_name. + It might belong to any workflow sharing the root, so it is quarantined (skipped + with a warning) instead of poisoning get_latest() for all of them. Ownership is + never trusted from a duplicated key. """ storage = FileCheckpointStorage(tmp_path) - await storage.save(_stamped_checkpoint({}, "2026-01-01T10:00:00+00:00")) + saved_id = 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") + + latest = await storage.get_latest(workflow_name="recovery-workflow") + + assert latest is not None + assert latest.checkpoint_id == saved_id async def test_file_checkpoint_storage_get_latest_rejects_duplicate_metadata_keys_in_matching_workflow( @@ -3748,9 +3900,7 @@ async def test_file_checkpoint_storage_get_latest_ignores_duplicate_keys_in_fore determined, not rejected during parsing. """ storage = FileCheckpointStorage(tmp_path) - newest_id = await storage.save( - _stamped_checkpoint({"stage": "first"}, "2026-01-01T11:00:00+00:00") - ) + newest_id = await storage.save(_stamped_checkpoint({"stage": "first"}, "2026-01-01T11:00:00+00:00")) (tmp_path / "foreign.json").write_text( '{"workflow_name": "other-workflow", ' '"timestamp": "2026-01-01T10:00:00+00:00", '