Python: FileCheckpointStorage save/load symmetry (#8181)

python/packages/core/agent_framework/_workflows/_checkpoint.py

@@ -320,16 +320,30 @@ 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 encoding or restore validation fails (#8181).
+        try:
+            encoded_checkpoint = encode_checkpoint_value(checkpoint_dict)
+            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 encoded or restored under "
+                "this storage's allowed types; refusing to save."
+            ) from ex
 
         def _write_atomic() -> None:
             tmp_path = file_path.with_suffix(".json.tmp")
-            with open(tmp_path, "w") as f:
+            with open(tmp_path, "w", encoding="utf-8") as f:
                 json.dump(encoded_checkpoint, f, indent=2, ensure_ascii=False)
             os.replace(tmp_path, file_path)
 
@@ -357,10 +371,19 @@ async def load(self, checkpoint_id: CheckpointID) -> WorkflowCheckpoint:
             raise WorkflowCheckpointException(f"No checkpoint found with ID {checkpoint_id}")
 
         def _read() -> dict[str, Any]:
-            with open(file_path) as f:
+            with open(file_path, encoding="utf-8") as f:
                 return json.load(f)
 
-        encoded_checkpoint = await asyncio.to_thread(_read)
+        try:
+            encoded_checkpoint = await asyncio.to_thread(_read)
+        except UnicodeDecodeError as ex:
+            raise WorkflowCheckpointException(
+                f"Checkpoint file for {checkpoint_id} is not valid UTF-8 and cannot be loaded."
+            ) from ex
+        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
 
@@ -386,7 +409,7 @@ def _list_checkpoints() -> list[WorkflowCheckpoint]:
             checkpoints: list[WorkflowCheckpoint] = []
             for file_path in self.storage_path.glob("*.json"):
                 try:
-                    with open(file_path) as f:
+                    with open(file_path, encoding="utf-8") as f:
                         encoded_checkpoint = json.load(f)
                         from ._checkpoint_encoding import decode_checkpoint_value
 
@@ -446,18 +469,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]:
-            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))
-                except Exception as e:
-                    logger.warning(f"Failed to read checkpoint file {file_path}: {e}")
-            return checkpoint_ids
-
-        return await asyncio.to_thread(_list_ids)
+        checkpoints = await self.list_checkpoints(workflow_name=workflow_name)
+        return [checkpoint.checkpoint_id for checkpoint in checkpoints]

python/packages/core/tests/workflow/test_checkpoint.py

@@ -1289,6 +1289,54 @@ 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_load_invalid_utf8_raises():
+    """Issue #8181: load wraps UnicodeDecodeError as WorkflowCheckpointException."""
+    with tempfile.TemporaryDirectory() as temp_dir:
+        storage = FileCheckpointStorage(temp_dir)
+        bad_id = "bad-utf8-checkpoint"
+        bad_file = Path(temp_dir) / f"{bad_id}.json"
+        bad_file.write_bytes(b'{"x": "\xff\xfe"}')
+        with pytest.raises(WorkflowCheckpointException, match="not valid UTF-8"):
+            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

@@ -11,13 +11,15 @@
 - Built-in safe types and framework types are always allowed
 """
 
+import asyncio
 import base64
 import enum
 import os
 import pickle
 import tempfile
 from dataclasses import dataclass
 from datetime import datetime, timezone
+from pathlib import Path
 from typing import Any
 
 import pytest
@@ -287,6 +289,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|cannot be encoded"):
+            await storage.save(checkpoint)
+        assert not await asyncio.to_thread(lambda: 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