diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint.py b/python/packages/core/agent_framework/_workflows/_checkpoint.py index cd78725d6a1..ddcd24f482f 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,156 @@ 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 _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() + 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) + 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=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): + 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..887447edd74 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,281 @@ 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) + # 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": + 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 + + +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