diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint.py b/python/packages/core/agent_framework/_workflows/_checkpoint.py index d9b30d8391..270e6c138d 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint.py +++ b/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) diff --git a/python/packages/core/tests/workflow/test_checkpoint.py b/python/packages/core/tests/workflow/test_checkpoint.py index be8de7c13b..16329ef260 100644 --- a/python/packages/core/tests/workflow/test_checkpoint.py +++ b/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)