Python: Harden FoundryCheckpointStore: apply RestrictedUnpickler allowlist on checkpoint load

python/packages/foundry_hosting/agent_framework_foundry_hosting/_state_store.py

@@ -68,23 +68,43 @@ def get_store(
 
 
 class FoundryCheckpointStore:
-    """Checkpoint store backed by the `FoundryStateStore`."""
+    """Checkpoint store backed by the `FoundryStateStore`.
+
+    By default, checkpoint deserialization is restricted to a built-in set of safe Python types
+    (primitives, datetime, uuid, ...), all ``agent_framework`` internal types, and OpenAI SDK types
+    (``openai.types``). To allow additional application-specific types, pass them via the
+    ``allowed_checkpoint_types`` parameter using ``"module:qualname"`` format. This mirrors the
+    behavior of ``FileCheckpointStorage`` and ``CosmosCheckpointStorage``, ensuring the
+    ``RestrictedUnpickler`` allowlist is applied on every hosted checkpoint load rather than
+    falling back to unrestricted ``pickle.loads``.
+    """
 
     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(
@@ -132,7 +152,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."""
@@ -148,7 +168,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:
@@ -181,8 +203,23 @@ class CheckpointStoreProvider(ContextScopedStoreProvider[CheckpointStorage]):
     and clean up older checkpoints without affecting another workflow context.
 
     This defaults to using the `FoundryCheckpointStore` in all environments.
+
+    Checkpoint deserialization is restricted to the built-in safe type set plus framework and
+    OpenAI SDK types by default. Applications that persist custom state types can widen the
+    allowlist by passing ``allowed_checkpoint_types`` (``"module:qualname"`` strings), which is
+    threaded to every ``FoundryCheckpointStore`` this provider creates.
     """
 
+    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) permitted during checkpoint deserialization,
+                as ``"module:qualname"`` strings.
+        """
+        self._allowed_checkpoint_types = allowed_checkpoint_types
+
     def get_store(
         self,
         *,
@@ -194,7 +231,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

@@ -114,6 +114,51 @@ async def test_load_raises_for_missing_checkpoint() -> None:
         await FoundryCheckpointStore("context-1", _platform_context()).load("missing")
 
 
+async def test_load_defaults_to_restricted_unpickler_and_blocks_unlisted_types() -> None:
+    # A raw store item that was NOT produced by encode_checkpoint_value (e.g. planted
+    # directly into the backing store) carries an arbitrary pickle payload. Under the
+    # default (restricted) unpickler this must be rejected rather than executed.
+    import base64
+    import os
+    import pickle
+
+    class _Gadget:
+        def __reduce__(self) -> Any:
+            return (os.system, ("echo should-not-run",))
+
+    malicious_value = {
+        "__pickled__": base64.b64encode(pickle.dumps(_Gadget())).decode("ascii"),
+        "__type__": "builtins:int",
+    }
+
+    store = _store()
+    store.get_item = AsyncMock(return_value=SimpleNamespace(value=malicious_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")
+
+
+def test_checkpoint_store_provider_threads_allowed_types() -> None:
+    provider = CheckpointStoreProvider(allowed_checkpoint_types=["my_app.models:MyState"])
+    store = provider.get_store(
+        config=_config(is_hosted=True), context_id="context-1", platform_context=_platform_context()
+    )
+
+    assert isinstance(store, FoundryCheckpointStore)
+    assert store._allowed_types == frozenset({"my_app.models:MyState"})
+
+
+def test_checkpoint_store_defaults_to_empty_allowlist() -> None:
+    store = FoundryCheckpointStore("context-1", _platform_context())
+    assert store._allowed_types == frozenset()
+
+
 async def test_list_checkpoints_paginates_and_filters_by_workflow() -> None:
     store = _store()
     matching = _checkpoint("checkpoint-1")