From ef0682333fe66cb191f4c9d3c26462d5676fd484 Mon Sep 17 00:00:00 2001 From: minelhi <3417378192@qq.com> Date: Wed, 9 Sep 2026 21:58:35 +0800 Subject: [PATCH] Python: FileCheckpointStorage save/load symmetry (#8181) --- .../agent_framework/_workflows/_checkpoint.py | 37 ++++++++++++++++--- .../core/tests/workflow/test_checkpoint.py | 37 +++++++++++++++++++ .../test_checkpoint_unrestricted_pickle.py | 17 +++++++++ 3 files changed, 86 insertions(+), 5 deletions(-) diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint.py b/python/packages/core/agent_framework/_workflows/_checkpoint.py index d9b30d83919..7c113c36767 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint.py @@ -320,12 +320,26 @@ async def save(self, checkpoint: WorkflowCheckpoint) -> CheckpointID: Returns: The unique ID of the saved checkpoint. + + Raises: + WorkflowCheckpointException: If the checkpoint cannot be encoded or would + fail to decode under this storage's ``allowed_checkpoint_types``. """ - from ._checkpoint_encoding import encode_checkpoint_value + from ._checkpoint_encoding import decode_checkpoint_value, encode_checkpoint_value file_path = self._validate_file_path(checkpoint.checkpoint_id) checkpoint_dict = checkpoint.to_dict() encoded_checkpoint = encode_checkpoint_value(checkpoint_dict) + # Fail at save time if this storage could not restore the payload (#8181). + try: + decode_checkpoint_value(encoded_checkpoint, allowed_types=self._allowed_types) + except WorkflowCheckpointException: + raise + except Exception as ex: + raise WorkflowCheckpointException( + f"Checkpoint {checkpoint.checkpoint_id} cannot be restored under this " + "storage's allowed types; refusing to save." + ) from ex def _write_atomic() -> None: tmp_path = file_path.with_suffix(".json.tmp") @@ -360,7 +374,12 @@ def _read() -> dict[str, Any]: with open(file_path) as f: return json.load(f) - encoded_checkpoint = await asyncio.to_thread(_read) + try: + encoded_checkpoint = await asyncio.to_thread(_read) + except json.JSONDecodeError as ex: + raise WorkflowCheckpointException( + f"Checkpoint file for {checkpoint_id} is not valid JSON and cannot be loaded." + ) from ex from ._checkpoint_encoding import decode_checkpoint_value @@ -446,6 +465,8 @@ async def list_checkpoint_ids(self, *, workflow_name: str) -> list[CheckpointID] Returns: A list of checkpoint IDs for the specified workflow name. + Only includes checkpoints that can be decoded under this storage's + allowed types (aligned with :meth:`list_checkpoints`, #8181). """ def _list_ids() -> list[CheckpointID]: @@ -453,9 +474,15 @@ def _list_ids() -> list[CheckpointID]: for file_path in self.storage_path.glob("*.json"): try: with open(file_path) as f: - data = json.load(f) - if data.get("workflow_name") == workflow_name: - checkpoint_ids.append(data.get("checkpoint_id", file_path.stem)) + encoded_checkpoint = json.load(f) + from ._checkpoint_encoding import decode_checkpoint_value + + decoded_checkpoint_dict = decode_checkpoint_value( + encoded_checkpoint, allowed_types=self._allowed_types + ) + checkpoint = WorkflowCheckpoint.from_dict(decoded_checkpoint_dict) + if checkpoint.workflow_name == workflow_name: + checkpoint_ids.append(checkpoint.checkpoint_id) except Exception as e: logger.warning(f"Failed to read checkpoint file {file_path}: {e}") return checkpoint_ids diff --git a/python/packages/core/tests/workflow/test_checkpoint.py b/python/packages/core/tests/workflow/test_checkpoint.py index be8de7c13b6..1fe883abd3c 100644 --- a/python/packages/core/tests/workflow/test_checkpoint.py +++ b/python/packages/core/tests/workflow/test_checkpoint.py @@ -1289,6 +1289,43 @@ async def test_file_checkpoint_storage_corrupted_file(): assert checkpoints == [] +async def test_file_checkpoint_storage_load_invalid_json_raises(): + """Issue #8181: load wraps JSONDecodeError as WorkflowCheckpointException.""" + with tempfile.TemporaryDirectory() as temp_dir: + storage = FileCheckpointStorage(temp_dir) + bad_id = "bad-json-checkpoint" + bad_file = Path(temp_dir) / f"{bad_id}.json" + with open(bad_file, "w") as f: # noqa: ASYNC230 + f.write("{ not json") + with pytest.raises(WorkflowCheckpointException, match="not valid JSON"): + await storage.load(bad_id) + + +async def test_file_checkpoint_storage_list_ids_matches_list_decode_filter(): + """Issue #8181: list_checkpoint_ids skips undecodable files like list_checkpoints.""" + from tests.workflow.test_checkpoint_unrestricted_pickle import _AllowedTestState + + with tempfile.TemporaryDirectory() as temp_dir: + type_key = f"{_AllowedTestState.__module__}:{_AllowedTestState.__qualname__}" + writer = FileCheckpointStorage(temp_dir, allowed_checkpoint_types=[type_key]) + good = WorkflowCheckpoint(workflow_name="wf-a", graph_signature_hash="h") + await writer.save(good) + blocked = WorkflowCheckpoint( + workflow_name="wf-a", + graph_signature_hash="h", + checkpoint_id="orphan-blocked", + state={"x": _AllowedTestState(name="x", value=1)}, + ) + await writer.save(blocked) + + reader = FileCheckpointStorage(temp_dir) # no allow list + listed = await reader.list_checkpoints(workflow_name="wf-a") + ids = await reader.list_checkpoint_ids(workflow_name="wf-a") + assert [c.checkpoint_id for c in listed] == ids + assert "orphan-blocked" not in ids + assert good.checkpoint_id in ids + + async def test_file_checkpoint_storage_json_serialization(): with tempfile.TemporaryDirectory() as temp_dir: storage = FileCheckpointStorage(temp_dir) diff --git a/python/packages/core/tests/workflow/test_checkpoint_unrestricted_pickle.py b/python/packages/core/tests/workflow/test_checkpoint_unrestricted_pickle.py index 4e89dd1b2cc..02eb0601c5d 100644 --- a/python/packages/core/tests/workflow/test_checkpoint_unrestricted_pickle.py +++ b/python/packages/core/tests/workflow/test_checkpoint_unrestricted_pickle.py @@ -18,6 +18,7 @@ import tempfile from dataclasses import dataclass from datetime import datetime, timezone +from pathlib import Path from typing import Any import pytest @@ -287,6 +288,22 @@ async def test_file_storage_blocks_unlisted_user_type(): await load_storage.load(checkpoint.checkpoint_id) +async def test_file_storage_rejects_unlisted_user_type_at_save(): + """Issue #8181: same storage must refuse to save what it cannot restore.""" + from agent_framework import WorkflowCheckpoint + + with tempfile.TemporaryDirectory() as tmpdir: + storage = FileCheckpointStorage(tmpdir) + checkpoint = WorkflowCheckpoint( + workflow_name="test", + graph_signature_hash="hash", + state={"data": _AllowedTestState(name="test", value=1)}, + ) + with pytest.raises(WorkflowCheckpointException, match="deserialization blocked|Unable to save|cannot be restored"): + await storage.save(checkpoint) + assert list(Path(tmpdir).glob("*.json")) == [] + + async def test_file_storage_allows_listed_user_type(): """FileCheckpointStorage allows user types listed in allowed_checkpoint_types.""" from agent_framework import WorkflowCheckpoint