Python: FileCheckpointStorage save/load symmetry (#8181)

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,16 +465,24 @@ 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]:
             checkpoint_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

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)

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