diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint.py b/python/packages/core/agent_framework/_workflows/_checkpoint.py index d9b30d8391..c6a6a9dc89 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint.py +++ b/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] diff --git a/python/packages/core/tests/workflow/test_checkpoint.py b/python/packages/core/tests/workflow/test_checkpoint.py index be8de7c13b..5c06b617c8 100644 --- a/python/packages/core/tests/workflow/test_checkpoint.py +++ b/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) 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 4e89dd1b2c..1f78f4e6ad 100644 --- a/python/packages/core/tests/workflow/test_checkpoint_unrestricted_pickle.py +++ b/python/packages/core/tests/workflow/test_checkpoint_unrestricted_pickle.py @@ -11,6 +11,7 @@ - Built-in safe types and framework types are always allowed """ +import asyncio import base64 import enum import os @@ -18,6 +19,7 @@ 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