[BREAKING] Python: restrict checkpoint deserialization in FoundryCheckpointStore

python/packages/foundry_hosting/agent_framework_foundry_hosting/_state_store.py

@@ -72,19 +72,30 @@ class FoundryCheckpointStore:
 
     DEFAULT_ROOT_SCOPE = "checkpoints"
 
-    def __init__(self, context_id: str, platform_context: FoundryAgentRequestContext) -> None:
+    def __init__(
+        self,
+        context_id: str,
+        platform_context: FoundryAgentRequestContext,
+        *,
+        allowed_checkpoint_types: list[str] | None = None,
+    ) -> None:
         """Initialize a Foundry-scoped checkpoint store for the given context ID.
 
         Args:
             context_id: A string that uniquely identifies the context for which the checkpoint store is scoped.
                         This can be used to isolate checkpoints for different workflow runs.
             platform_context: The request-scoped platform context for the current request.
+            allowed_checkpoint_types: Additional types (beyond the built-in safe set
+                and framework types) that are permitted during checkpoint
+                deserialization.  Each entry should be a ``"module:qualname"``
+                string (e.g., ``"my_app.models:MyState"``).
         """
         if not context_id:
             raise ValueError("context_id must be provided to initialize a FoundryCheckpointStore.")
 
         self.context_id = context_id
         self.platform_context = platform_context
+        self._allowed_types: frozenset[str] = frozenset(allowed_checkpoint_types or [])
 
     async def _get_store(self) -> FoundryStateStore:
         return await FoundryStateStore.get_or_create(
@@ -131,7 +142,7 @@ async def load(self, checkpoint_id: CheckpointID) -> WorkflowCheckpoint:
             item = await store.get_item(checkpoint_id, call_id=self.platform_context.call_id)
         if item is None:
             raise WorkflowCheckpointException(f"No checkpoint found with ID {checkpoint_id}")
-        return WorkflowCheckpoint.from_dict(decode_checkpoint_value(item.value))
+        return WorkflowCheckpoint.from_dict(decode_checkpoint_value(item.value, allowed_types=self._allowed_types))
 
     async def list_checkpoints(self, *, workflow_name: str) -> list[WorkflowCheckpoint]:
         """List all workflow checkpoints for a given workflow name."""
@@ -147,7 +158,9 @@ async def list_checkpoints(self, *, workflow_name: str) -> list[WorkflowCheckpoi
                     item = await store.get_item(item_key.key, call_id=self.platform_context.call_id)
                     if item is None:
                         continue
-                    checkpoint = WorkflowCheckpoint.from_dict(decode_checkpoint_value(item.value))
+                    checkpoint = WorkflowCheckpoint.from_dict(
+                        decode_checkpoint_value(item.value, allowed_types=self._allowed_types)
+                    )
                     if checkpoint.workflow_name == workflow_name:
                         checkpoints.append(checkpoint)
                 if not page.has_more or page.last_id is None:
@@ -182,6 +195,18 @@ class CheckpointStoreProvider(ContextScopedStoreProvider[CheckpointStorage]):
     This defaults to using the `FoundryCheckpointStore` in all environments.
     """
 
+    def __init__(self, *, allowed_checkpoint_types: list[str] | None = None) -> None:
+        """Initialize the provider.
+
+        Args:
+            allowed_checkpoint_types: Additional types (beyond the built-in safe set
+                and framework types) that are permitted during checkpoint
+                deserialization, forwarded to every store this provider creates.
+                Each entry should be a ``"module:qualname"`` string
+                (e.g., ``"my_app.models:MyState"``).
+        """
+        self._allowed_checkpoint_types = allowed_checkpoint_types
+
     def get_store(
         self,
         *,
@@ -193,7 +218,11 @@ def get_store(
         if not context_id:
             raise ValueError("context_id must be provided to get a checkpoint store.")
 
-        return FoundryCheckpointStore(context_id, platform_context)
+        return FoundryCheckpointStore(
+            context_id,
+            platform_context,
+            allowed_checkpoint_types=self._allowed_checkpoint_types,
+        )
 
 
 # endregion Checkpoint persistence

python/packages/foundry_hosting/tests/test_state_store.py

@@ -1,11 +1,13 @@
 # Copyright (c) Microsoft. All rights reserved.
 from collections.abc import Callable
+from dataclasses import dataclass
 from types import SimpleNamespace
 from typing import Any
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 from agent_framework import AgentSession, Content, WorkflowCheckpoint, WorkflowCheckpointException
+from agent_framework._workflows._checkpoint_encoding import encode_checkpoint_value
 from azure.ai.agentserver.core import AgentConfig, FoundryAgentRequestContext
 from azure.ai.agentserver.core.storage import FoundryStorageConflictError
 
@@ -20,6 +22,13 @@
 )
 
 
+@dataclass
+class _NotAllowed:
+    """A type outside the built-in safe set, standing in for an application type."""
+
+    value: int
+
+
 def _checkpoint(
     checkpoint_id: str, *, workflow_name: str = "workflow", timestamp: str = "2026-01-01T00:00:00+00:00"
 ) -> WorkflowCheckpoint:
@@ -100,6 +109,123 @@ async def test_load_returns_checkpoint() -> None:
     store.get_item.assert_awaited_once_with("checkpoint-1", call_id="call-1")
 
 
+async def test_load_restricts_checkpoint_deserialization() -> None:
+    """A checkpoint value naming a type outside the allow set is refused.
+
+    The file and Cosmos checkpoint stores both restrict deserialization this
+    way; this store reaches the same decoder, so it restricts it too.
+    """
+    store = _store()
+    checkpoint = _checkpoint("checkpoint-1")
+    value = checkpoint.to_dict()
+    value["state"] = encode_checkpoint_value({"payload": _NotAllowed(7)})
+    store.get_item = AsyncMock(return_value=SimpleNamespace(value=value))
+
+    with (
+        patch(
+            "agent_framework_foundry_hosting._state_store.FoundryStateStore.get_or_create",
+            new=AsyncMock(return_value=store),
+        ),
+        pytest.raises(WorkflowCheckpointException),
+    ):
+        await FoundryCheckpointStore("context-1", _platform_context()).load("checkpoint-1")
+
+
+async def test_load_accepts_a_declared_checkpoint_type() -> None:
+    """A caller can still name the types its checkpoints carry."""
+    store = _store()
+    checkpoint = _checkpoint("checkpoint-1")
+    value = checkpoint.to_dict()
+    value["state"] = encode_checkpoint_value({"payload": _NotAllowed(7)})
+    store.get_item = AsyncMock(return_value=SimpleNamespace(value=value))
+
+    with patch(
+        "agent_framework_foundry_hosting._state_store.FoundryStateStore.get_or_create",
+        new=AsyncMock(return_value=store),
+    ):
+        result = await FoundryCheckpointStore(
+            "context-1",
+            _platform_context(),
+            allowed_checkpoint_types=[f"{_NotAllowed.__module__}:{_NotAllowed.__qualname__}"],
+        ).load("checkpoint-1")
+
+    assert result.state["payload"].value == 7
+
+
+async def test_list_checkpoints_restricts_checkpoint_deserialization() -> None:
+    store = _store()
+    checkpoint = _checkpoint("checkpoint-1")
+    value = checkpoint.to_dict()
+    value["state"] = encode_checkpoint_value({"payload": _NotAllowed(7)})
+    store.list_keys = AsyncMock(
+        return_value=SimpleNamespace(keys=[SimpleNamespace(key="checkpoint-1")], has_more=False, last_id=None)
+    )
+    store.get_item = AsyncMock(return_value=SimpleNamespace(value=value))
+
+    with (
+        patch(
+            "agent_framework_foundry_hosting._state_store.FoundryStateStore.get_or_create",
+            new=AsyncMock(return_value=store),
+        ),
+        pytest.raises(WorkflowCheckpointException),
+    ):
+        await FoundryCheckpointStore("context-1", _platform_context()).list_checkpoints(workflow_name="workflow")
+
+
+async def test_provider_forwards_allowed_checkpoint_types() -> None:
+    """A hosted app reaches the option through the provider it actually gets.
+
+    `ResponsesHostServer` builds a `CheckpointStoreProvider` itself on the default
+    path, so an option only settable on the store would be out of reach there.
+    """
+    store = _store()
+    checkpoint = _checkpoint("checkpoint-1")
+    value = checkpoint.to_dict()
+    value["state"] = encode_checkpoint_value({"payload": _NotAllowed(7)})
+    store.get_item = AsyncMock(return_value=SimpleNamespace(value=value))
+
+    provider = CheckpointStoreProvider(
+        allowed_checkpoint_types=[f"{_NotAllowed.__module__}:{_NotAllowed.__qualname__}"]
+    )
+    storage = provider.get_store(
+        config=MagicMock(),
+        context_id="context-1",
+        platform_context=_platform_context(),
+    )
+
+    with patch(
+        "agent_framework_foundry_hosting._state_store.FoundryStateStore.get_or_create",
+        new=AsyncMock(return_value=store),
+    ):
+        result = await storage.load("checkpoint-1")
+
+    assert result.state["payload"].value == 7
+
+
+async def test_provider_restricts_by_default() -> None:
+    """Without the option the provider's stores restrict, as before."""
+    store = _store()
+    checkpoint = _checkpoint("checkpoint-1")
+    value = checkpoint.to_dict()
+    value["state"] = encode_checkpoint_value({"payload": _NotAllowed(7)})
+    store.get_item = AsyncMock(return_value=SimpleNamespace(value=value))
+
+    storage = CheckpointStoreProvider().get_store(
+        config=MagicMock(),
+        context_id="context-1",
+        platform_context=_platform_context(),
+    )
+
+    with (
+        patch(
+            "agent_framework_foundry_hosting._state_store.FoundryStateStore.get_or_create",
+            new=AsyncMock(return_value=store),
+        ),
+        pytest.raises(WorkflowCheckpointException),
+    ):
+        await storage.load("checkpoint-1")
+
+
 async def test_load_raises_for_missing_checkpoint() -> None:
     store = _store()
     store.get_item = AsyncMock(return_value=None)