From 6824afc8b1aefa0981ab1284a491084fbfb59b5a Mon Sep 17 00:00:00 2001 From: NNIIOVO Date: Tue, 29 Sep 2026 18:10:34 -0500 Subject: [PATCH 1/2] Python: Refuse to save Cosmos DB checkpoints that cannot be restored CosmosCheckpointStorage.save only encoded the checkpoint, while load decodes it under the storage's allowed_checkpoint_types. A checkpoint holding an unregistered application type was written successfully and only failed when a workflow tried to resume from it. FileCheckpointStorage has refused such checkpoints at save time since #8214. Validate the encoded checkpoint the same way before writing, wrap encoding failures in WorkflowCheckpointException, and skip creating the database or container when the save is refused. Register ComputeTask in the Cosmos checkpointing sample, whose checkpoints could not be restored. --- .../_checkpoint_storage.py | 19 ++++++- .../tests/test_cosmos_checkpoint_storage.py | 56 +++++++++++++++++++ .../cosmos_workflow_checkpointing.py | 5 ++ 3 files changed, 78 insertions(+), 2 deletions(-) diff --git a/python/packages/azure-cosmos/agent_framework_azure_cosmos/_checkpoint_storage.py b/python/packages/azure-cosmos/agent_framework_azure_cosmos/_checkpoint_storage.py index 3ddd0e0c73f..419d538fe96 100644 --- a/python/packages/azure-cosmos/agent_framework_azure_cosmos/_checkpoint_storage.py +++ b/python/packages/azure-cosmos/agent_framework_azure_cosmos/_checkpoint_storage.py @@ -218,12 +218,27 @@ 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``. """ mark_feature_used(FeatureIndex.AZURE_COSMOS) - await self._ensure_container_proxy() checkpoint_dict = checkpoint.to_dict() - encoded = encode_checkpoint_value(checkpoint_dict) + # Fail at save time if encoding or restore validation fails, matching FileCheckpointStorage. + try: + encoded = encode_checkpoint_value(checkpoint_dict) + decode_checkpoint_value(encoded, 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 + + await self._ensure_container_proxy() document: dict[str, Any] = { "id": self._make_document_id(checkpoint.workflow_name, checkpoint.checkpoint_id), diff --git a/python/packages/azure-cosmos/tests/test_cosmos_checkpoint_storage.py b/python/packages/azure-cosmos/tests/test_cosmos_checkpoint_storage.py index e4fd305171c..552a9e854c4 100644 --- a/python/packages/azure-cosmos/tests/test_cosmos_checkpoint_storage.py +++ b/python/packages/azure-cosmos/tests/test_cosmos_checkpoint_storage.py @@ -3,6 +3,7 @@ from __future__ import annotations import os +import threading import uuid from collections.abc import AsyncIterator from contextlib import suppress @@ -203,6 +204,18 @@ async def test_save_returns_checkpoint_id(mock_container: MagicMock) -> None: assert result == "cp-123" +async def test_save_unencodable_state_raises_checkpoint_exception(mock_container: MagicMock) -> None: + """save wraps encoding failures in WorkflowCheckpointException and writes nothing.""" + checkpoint = _make_checkpoint() + checkpoint.state = {"lock": threading.Lock()} + storage = CosmosCheckpointStorage(container_client=mock_container) + + with pytest.raises(WorkflowCheckpointException, match="cannot be encoded"): + await storage.save(checkpoint) + + mock_container.upsert_item.assert_not_awaited() + + # --- Tests for load --- @@ -763,3 +776,46 @@ async def test_get_latest_allows_listed_app_type(mock_container: MagicMock) -> N assert result is not None assert isinstance(result.state["data"], _AppState) assert result.state["data"].label == "latest" + + +async def test_save_rejects_unlisted_app_type(mock_container: MagicMock) -> None: + """save refuses a checkpoint that the same storage could not load back.""" + checkpoint = _make_checkpoint_with_state({"data": _AppState(label="x", count=1)}) + storage = CosmosCheckpointStorage(container_client=mock_container) + + with pytest.raises(WorkflowCheckpointException, match="deserialization blocked"): + await storage.save(checkpoint) + + mock_container.upsert_item.assert_not_awaited() + + +async def test_save_rejected_checkpoint_does_not_create_database( + mock_cosmos_client: MagicMock, mock_container: MagicMock +) -> None: + """A refused save does not create the database or container.""" + checkpoint = _make_checkpoint_with_state({"data": _AppState(label="x", count=1)}) + storage = CosmosCheckpointStorage( + cosmos_client=mock_cosmos_client, + database_name="db1", + container_name="checkpoints", + ) + + with pytest.raises(WorkflowCheckpointException): + await storage.save(checkpoint) + + mock_cosmos_client.create_database_if_not_exists.assert_not_awaited() + mock_container.upsert_item.assert_not_awaited() + + +async def test_save_allows_listed_app_type(mock_container: MagicMock) -> None: + """save accepts application types listed in allowed_checkpoint_types.""" + checkpoint = _make_checkpoint_with_state({"data": _AppState(label="ok", count=7)}) + storage = CosmosCheckpointStorage( + container_client=mock_container, + allowed_checkpoint_types=[_APP_STATE_TYPE_KEY], + ) + + result = await storage.save(checkpoint) + + assert result == checkpoint.checkpoint_id + mock_container.upsert_item.assert_awaited_once() diff --git a/python/samples/03-workflows/checkpoint/cosmos_workflow_checkpointing.py b/python/samples/03-workflows/checkpoint/cosmos_workflow_checkpointing.py index fd4608db03d..44247d5de2e 100644 --- a/python/samples/03-workflows/checkpoint/cosmos_workflow_checkpointing.py +++ b/python/samples/03-workflows/checkpoint/cosmos_workflow_checkpointing.py @@ -38,6 +38,7 @@ WorkflowCheckpoint, WorkflowContext, handler, + register_checkpoint_type, ) if sys.version_info >= (3, 12): @@ -114,6 +115,10 @@ async def main() -> None: print("Please set AZURE_COSMOS_ENDPOINT, AZURE_COSMOS_DATABASE_NAME, and AZURE_COSMOS_CONTAINER_NAME.") return + # Register the application-defined message type so checkpoints that contain it can be restored. + # Alternatively, pass allowed_checkpoint_types to CosmosCheckpointStorage to scope it to that storage. + register_checkpoint_type(ComputeTask) + # Authentication: supports both managed identity/RBAC and key-based auth. # When AZURE_COSMOS_KEY is set, key-based auth is used. # Otherwise, falls back to DefaultAzureCredential (properly closed via async with). From 6416f649eb8e78e59363cfe623838d135de16dce Mon Sep 17 00:00:00 2001 From: NNIIOVO Date: Wed, 30 Sep 2026 01:00:08 -0500 Subject: [PATCH 2/2] Python: Share save-time checkpoint validation between file and Cosmos storage Move the encode-and-restore check into a private core helper, _encode_checkpoint_for_storage(checkpoint, allowed_types), which returns the validated encoded payload. FileCheckpointStorage.save and CosmosCheckpointStorage.save both call it and keep their own I/O, so the two stores can no longer drift apart. Neither store's behavior changes. Add a FileCheckpointStorage test for the wrapped encoding failure, a path that only the Cosmos tests reached before. --- .../_checkpoint_storage.py | 23 +++------ .../agent_framework/_workflows/_checkpoint.py | 51 ++++++++++++++----- .../test_checkpoint_unrestricted_pickle.py | 18 +++++++ 3 files changed, 63 insertions(+), 29 deletions(-) diff --git a/python/packages/azure-cosmos/agent_framework_azure_cosmos/_checkpoint_storage.py b/python/packages/azure-cosmos/agent_framework_azure_cosmos/_checkpoint_storage.py index 419d538fe96..052172a95b6 100644 --- a/python/packages/azure-cosmos/agent_framework_azure_cosmos/_checkpoint_storage.py +++ b/python/packages/azure-cosmos/agent_framework_azure_cosmos/_checkpoint_storage.py @@ -9,8 +9,12 @@ from agent_framework._settings import SecretString, load_settings from agent_framework._telemetry import get_user_agent, mark_feature_used -from agent_framework._workflows._checkpoint import CheckpointID, WorkflowCheckpoint -from agent_framework._workflows._checkpoint_encoding import decode_checkpoint_value, encode_checkpoint_value +from agent_framework._workflows._checkpoint import ( + CheckpointID, + WorkflowCheckpoint, + _encode_checkpoint_for_storage, # pyright: ignore[reportPrivateUsage] +) +from agent_framework._workflows._checkpoint_encoding import decode_checkpoint_value from agent_framework.exceptions import WorkflowCheckpointException from azure.core.credentials import TokenCredential from azure.core.credentials_async import AsyncTokenCredential @@ -225,19 +229,8 @@ async def save(self, checkpoint: WorkflowCheckpoint) -> CheckpointID: """ mark_feature_used(FeatureIndex.AZURE_COSMOS) - checkpoint_dict = checkpoint.to_dict() - # Fail at save time if encoding or restore validation fails, matching FileCheckpointStorage. - try: - encoded = encode_checkpoint_value(checkpoint_dict) - decode_checkpoint_value(encoded, 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 - + # Validate before resolving the container, so a refused save creates no Cosmos resources. + encoded = _encode_checkpoint_for_storage(checkpoint, self._allowed_types) await self._ensure_container_proxy() document: dict[str, Any] = { diff --git a/python/packages/core/agent_framework/_workflows/_checkpoint.py b/python/packages/core/agent_framework/_workflows/_checkpoint.py index cd78725d6a1..c06497bd2b5 100644 --- a/python/packages/core/agent_framework/_workflows/_checkpoint.py +++ b/python/packages/core/agent_framework/_workflows/_checkpoint.py @@ -252,6 +252,42 @@ async def list_checkpoint_ids(self, *, workflow_name: str) -> list[CheckpointID] return [cp.checkpoint_id for cp in self._checkpoints.values() if cp.workflow_name == workflow_name] +def _encode_checkpoint_for_storage(checkpoint: WorkflowCheckpoint, allowed_types: frozenset[str]) -> dict[str, Any]: + """Encode a checkpoint for persistent storage, refusing one the storage could not restore. + + Storages that persist encoded checkpoints decode them under their ``allowed_checkpoint_types`` + on load, so the encoded payload is decoded here under the same allowlist before it is written. + A checkpoint that ``load`` would reject then fails at save time instead of when a workflow + tries to resume from it (#8181). + + Args: + checkpoint: The checkpoint to encode. + allowed_types: The storage's ``allowed_checkpoint_types``, as passed to + ``decode_checkpoint_value`` on load. + + Returns: + The encoded checkpoint, ready to be written. + + Raises: + WorkflowCheckpointException: If the checkpoint cannot be encoded or would + fail to decode under ``allowed_types``. + """ + from ._checkpoint_encoding import decode_checkpoint_value, encode_checkpoint_value + + checkpoint_dict = checkpoint.to_dict() + try: + encoded_checkpoint: dict[str, Any] = encode_checkpoint_value(checkpoint_dict) + decode_checkpoint_value(encoded_checkpoint, allowed_types=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 + return encoded_checkpoint + + # Process-wide serialization of writes per destination file. # # asyncio.Lock is loop-bound, so a per-(loop, checkpoint-id) registry cannot @@ -589,21 +625,8 @@ async def save(self, checkpoint: WorkflowCheckpoint) -> CheckpointID: WorkflowCheckpointException: If the checkpoint cannot be encoded or would fail to decode under this storage's ``allowed_checkpoint_types``. """ - from ._checkpoint_encoding import decode_checkpoint_value, encode_checkpoint_value - file_path = self._validate_file_path(checkpoint.checkpoint_id) - checkpoint_dict = checkpoint.to_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 + encoded_checkpoint = _encode_checkpoint_for_storage(checkpoint, self._allowed_types) def _replace_with_retry(tmp_path: Path) -> None: # On Windows, os.replace can transiently fail with PermissionError when a 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 77f0ab84235..9fdb0fe2e4a 100644 --- a/python/packages/core/tests/workflow/test_checkpoint_unrestricted_pickle.py +++ b/python/packages/core/tests/workflow/test_checkpoint_unrestricted_pickle.py @@ -17,6 +17,7 @@ import os import pickle import tempfile +import threading from dataclasses import dataclass from datetime import datetime, timezone from pathlib import Path @@ -308,6 +309,23 @@ async def test_file_storage_rejects_unlisted_user_type_at_save(): assert not await asyncio.to_thread(lambda: list(Path(tmpdir).glob("*.json"))) +async def test_file_storage_wraps_unencodable_state_at_save(): + """save wraps an encoding failure in WorkflowCheckpointException and writes nothing.""" + from agent_framework import WorkflowCheckpoint + + with tempfile.TemporaryDirectory() as tmpdir: + storage = FileCheckpointStorage(tmpdir) + checkpoint = WorkflowCheckpoint( + workflow_name="test", + graph_signature_hash="hash", + state={"lock": threading.Lock()}, + ) + with pytest.raises(WorkflowCheckpointException, match="cannot be encoded") as exc_info: + await storage.save(checkpoint) + assert isinstance(exc_info.value.__cause__, TypeError) + assert not await asyncio.to_thread(lambda: list(Path(tmpdir).iterdir())) + + async def test_file_storage_allows_listed_user_type(): """FileCheckpointStorage allows user types listed in allowed_checkpoint_types.""" from agent_framework import WorkflowCheckpoint