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..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 @@ -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 diff --git a/python/packages/foundry_hosting/tests/test_state_store.py b/python/packages/foundry_hosting/tests/test_state_store.py index 23a488c03d..fc36d00373 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,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)