Python: fix: raise WorkflowCheckpointException when loading a corrupted checkpoint file (#8181, item 4)

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

@@ -358,7 +358,16 @@ async def load(self, checkpoint_id: CheckpointID) -> WorkflowCheckpoint:
 
         def _read() -> dict[str, Any]:
             with open(file_path) as f:
-                return json.load(f)
+                try:
+                    return json.load(f)
+                except (json.JSONDecodeError, UnicodeDecodeError) as exc:
+                    # `load` is documented to raise WorkflowCheckpointException
+                    # when checkpoint decoding fails; a truncated file or one
+                    # with invalid utf-8 should surface as that, not a raw
+                    # json/unicode error (#8181).
+                    raise WorkflowCheckpointException(
+                        f"Checkpoint file for ID {checkpoint_id} is corrupted: {exc}"
+                    ) from exc
 
         encoded_checkpoint = await asyncio.to_thread(_read)
 

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

@@ -1275,6 +1275,41 @@ async def test_file_checkpoint_storage_directory_creation():
         assert file_path.exists()
 
 
+async def test_file_checkpoint_storage_load_corrupted_raises_checkpoint_exception():
+    """`load` on a corrupted file raises the documented exception (#8181, item 4).
+
+    Distinct from the graceful-list behavior pinned by
+    `test_file_checkpoint_storage_corrupted_file`: loading a truncated
+    checkpoint must surface as `WorkflowCheckpointException`, per the
+    `CheckpointStorage.load` contract, not as a raw `json.JSONDecodeError`.
+    """
+    with tempfile.TemporaryDirectory() as temp_dir:
+        storage = FileCheckpointStorage(temp_dir)
+        checkpoint = WorkflowCheckpoint(workflow_name="wf", graph_signature_hash="sig", state={"x": 1})
+        await storage.save(checkpoint)
+
+        file_path = Path(temp_dir) / f"{checkpoint.checkpoint_id}.json"
+        raw = file_path.read_text()
+        file_path.write_text(raw[: len(raw) // 2])  # truncate mid-JSON
+
+        with pytest.raises(WorkflowCheckpointException, match="corrupted"):
+            await storage.load(checkpoint.checkpoint_id)
+
+
+async def test_file_checkpoint_storage_load_invalid_utf8_raises_checkpoint_exception():
+    """`load` on a file with invalid utf-8 raises the documented exception (PR review)."""
+    with tempfile.TemporaryDirectory() as temp_dir:
+        storage = FileCheckpointStorage(temp_dir)
+        checkpoint = WorkflowCheckpoint(workflow_name="wf", graph_signature_hash="sig", state={"x": 1})
+        await storage.save(checkpoint)
+
+        file_path = Path(temp_dir) / f"{checkpoint.checkpoint_id}.json"
+        file_path.write_bytes(b'{"workflow_name": "' + b"\xff\xfe" + b'"}')  # invalid utf-8
+
+        with pytest.raises(WorkflowCheckpointException, match="corrupted"):
+            await storage.load(checkpoint.checkpoint_id)
+
+
 async def test_file_checkpoint_storage_corrupted_file():
     with tempfile.TemporaryDirectory() as temp_dir:
         storage = FileCheckpointStorage(temp_dir)