From 724c147b2c672e7fd8e33ca1958033b22f5aad5c Mon Sep 17 00:00:00 2001 From: sricursion Date: Fri, 4 Sep 2026 00:04:37 +0530 Subject: [PATCH 1/2] Python: restrict checkpoint deserialization in FoundryCheckpointStore `FileCheckpointStorage` and the Cosmos checkpoint storage both take an `allowed_checkpoint_types` argument, hold it as `self._allowed_types`, and pass it to `decode_checkpoint_value`. `FoundryCheckpointStore` had neither, and called the decoder with the argument omitted in `load` and again in `list_checkpoints`. An omitted `allowed_types` means no restriction and falls through to plain `pickle.loads`, while the empty frozenset the other two stores pass by default selects the restricted unpickler. So the Foundry store was the only one of the three not following the module's own guidance that the argument be specified whenever possible. Give it the same argument and pass it on, so a workflow that restores on one store restores on the others. This is a behaviour change for applications on this store whose checkpoints hold their own types and which never registered them, since those are relying on the store being more permissive than the other two. They register the types with `register_checkpoint_type` or pass `allowed_checkpoint_types`, which an application on the file or Cosmos store already has to do. --- .../_state_store.py | 19 ++++- .../foundry_hosting/tests/test_state_store.py | 72 +++++++++++++++++++ 2 files changed, 88 insertions(+), 3 deletions(-) diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_state_store.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_state_store.py index d73b582f9b..e68e41be8e 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_state_store.py +++ b/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: diff --git a/python/packages/foundry_hosting/tests/test_state_store.py b/python/packages/foundry_hosting/tests/test_state_store.py index 23a488c03d..750bbe3947 100644 --- a/python/packages/foundry_hosting/tests/test_state_store.py +++ b/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,69 @@ 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_load_raises_for_missing_checkpoint() -> None: store = _store() store.get_item = AsyncMock(return_value=None) From 9fad23473eaa4cb4453ff628b31b529b4ccd9bc8 Mon Sep 17 00:00:00 2001 From: sricursion Date: Fri, 4 Sep 2026 19:07:58 +0530 Subject: [PATCH 2/2] Forward allowed_checkpoint_types through CheckpointStoreProvider `ResponsesHostServer` builds a `CheckpointStoreProvider` itself on the default path, so an option settable only on `FoundryCheckpointStore` was out of reach for a hosted app: it would have had to register types process-wide or replace the whole provider. Take the list on the provider and pass it to each store it creates. The default is unchanged, so a provider built with no arguments still restricts exactly as before. --- .../_state_store.py | 18 ++++++- .../foundry_hosting/tests/test_state_store.py | 54 +++++++++++++++++++ 2 files changed, 71 insertions(+), 1 deletion(-) diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_state_store.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_state_store.py index e68e41be8e..a88c984ead 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_state_store.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_state_store.py @@ -195,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, *, @@ -206,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 diff --git a/python/packages/foundry_hosting/tests/test_state_store.py b/python/packages/foundry_hosting/tests/test_state_store.py index 750bbe3947..fc36d00373 100644 --- a/python/packages/foundry_hosting/tests/test_state_store.py +++ b/python/packages/foundry_hosting/tests/test_state_store.py @@ -172,6 +172,60 @@ async def test_list_checkpoints_restricts_checkpoint_deserialization() -> None: 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)