Export ContextScopedStoreProvider

python/PACKAGE_STATUS.md

@@ -140,8 +140,8 @@ listed below.
 
 - `agent-framework-core`: `SessionStore` and `FileSessionStore` from
   `agent_framework/_sessions.py`
-- `agent-framework-foundry-hosting`: `FoundrySessionStore` from
-  `agent_framework_foundry_hosting/_session_store.py`
+- `agent-framework-foundry-hosting`: `FoundryAgentSessionStore` from
+  `agent_framework_foundry_hosting/_state_store.py`
 
 #### `TO_PROMPT_AGENT`
 

python/packages/core/AGENTS.md

@@ -213,8 +213,8 @@ agent_framework/
 ### Foundry (`foundry/`)
 
 - **`FoundryChatClient`** - Chat client for Microsoft Foundry project endpoints
-- **`FoundrySessionStore`** - Experimental Foundry-hosting session store, lazily re-exported from
-  `agent-framework-foundry-hosting`; currently file-backed and scoped by Agent Server request context
+- **`FoundryAgentSessionStore`** - Experimental Foundry-hosting session store, lazily re-exported from
+  `agent-framework-foundry-hosting`; currently backed by `azure.ai.agentserver.core.storage.FoundryStateStore`
 
 ## Key Patterns
 

python/packages/core/agent_framework/foundry/__init__.py

@@ -32,7 +32,15 @@ _IMPORTS: dict[str, tuple[str, str]] = {
     "FoundryEmbeddingSettings": ("agent_framework_foundry", "agent-framework-foundry"),
     "FoundryEvals": ("agent_framework_foundry", "agent-framework-foundry"),
     "FoundryMemoryProvider": ("agent_framework_foundry", "agent-framework-foundry"),
-    "FoundrySessionStore": ("agent_framework_foundry_hosting", "agent-framework-foundry-hosting"),
+    "FoundryAgentSessionStore": ("agent_framework_foundry_hosting", "agent-framework-foundry-hosting"),
+    "StoreProvider": ("agent_framework_foundry_hosting", "agent-framework-foundry-hosting"),
+    "AgentSessionStoreProvider": ("agent_framework_foundry_hosting", "agent-framework-foundry-hosting"),
+    "CheckpointStoreProvider": ("agent_framework_foundry_hosting", "agent-framework-foundry-hosting"),
+    "FoundryCheckpointStore": ("agent_framework_foundry_hosting", "agent-framework-foundry-hosting"),
+    "ContextScopedStoreProvider": ("agent_framework_foundry_hosting", "agent-framework-foundry-hosting"),
+    "FoundryFunctionApprovalStore": ("agent_framework_foundry_hosting", "agent-framework-foundry-hosting"),
+    "FunctionApprovalStore": ("agent_framework_foundry_hosting", "agent-framework-foundry-hosting"),
+    "FunctionApprovalStoreProvider": ("agent_framework_foundry_hosting", "agent-framework-foundry-hosting"),
     "FoundryToolbox": ("agent_framework_foundry_hosting", "agent-framework-foundry-hosting"),
     "FoundryLocalChatOptions": ("agent_framework_foundry_local", "agent-framework-foundry-local"),
     "FoundryLocalClient": ("agent_framework_foundry_local", "agent-framework-foundry-local"),

python/packages/core/agent_framework/foundry/__init__.pyi

@@ -30,10 +30,18 @@ from agent_framework_foundry import (
     to_prompt_agent,
 )
 from agent_framework_foundry_hosting import (
-    FoundrySessionStore,
+    AgentSessionStoreProvider,
+    CheckpointStoreProvider,
+    ContextScopedStoreProvider,
+    FoundryAgentSessionStore,
+    FoundryCheckpointStore,
+    FoundryFunctionApprovalStore,
     FoundryToolbox,
+    FunctionApprovalStore,
+    FunctionApprovalStoreProvider,
     InvocationsHostServer,
     ResponsesHostServer,
+    StoreProvider,
 )
 from agent_framework_foundry_local import (
     FoundryLocalChatOptions,
@@ -42,25 +50,32 @@ from agent_framework_foundry_local import (
 )
 
 __all__ = [
+    "AgentSessionStoreProvider",
     "AnalysisSection",
     "AnthropicFoundryClient",
+    "CheckpointStoreProvider",
     "ContentUnderstandingContextProvider",
+    "ContextScopedStoreProvider",
     "DocumentStatus",
     "FileSearchBackend",
     "FileSearchConfig",
     "FoundryAgent",
+    "FoundryAgentSessionStore",
     "FoundryChatClient",
     "FoundryChatOptions",
+    "FoundryCheckpointStore",
     "FoundryEmbeddingClient",
     "FoundryEmbeddingOptions",
     "FoundryEmbeddingSettings",
     "FoundryEvals",
+    "FoundryFunctionApprovalStore",
     "FoundryLocalChatOptions",
     "FoundryLocalClient",
     "FoundryLocalSettings",
     "FoundryMemoryProvider",
-    "FoundrySessionStore",
     "FoundryToolbox",
+    "FunctionApprovalStore",
+    "FunctionApprovalStoreProvider",
     "GeneratedEvaluatorRef",
     "InvocationsHostServer",
     "RawAnthropicFoundryClient",
@@ -69,6 +84,7 @@ __all__ = [
     "RawFoundryChatClient",
     "RawFoundryEmbeddingClient",
     "ResponsesHostServer",
+    "StoreProvider",
     "evaluate_foundry_target",
     "evaluate_traces",
     "to_prompt_agent",

python/packages/core/tests/core/test_foundry_namespace.py

@@ -11,7 +11,7 @@ _foundry_local = pytest.importorskip("agent_framework_foundry_local")
 
 FoundryChatClient = _foundry.FoundryChatClient
 FoundryMemoryProvider = _foundry.FoundryMemoryProvider
-FoundrySessionStore = _foundry_hosting.FoundrySessionStore
+FoundryAgentSessionStore = _foundry_hosting.FoundryAgentSessionStore
 ResponsesHostServer = _foundry_hosting.ResponsesHostServer
 FoundryLocalClient = _foundry_local.FoundryLocalClient
 
@@ -19,12 +19,12 @@ FoundryLocalClient = _foundry_local.FoundryLocalClient
 def test_foundry_namespace_exposes_cloud_and_local_symbols() -> None:
     assert foundry.FoundryChatClient is FoundryChatClient
     assert foundry.FoundryMemoryProvider is FoundryMemoryProvider
-    assert foundry.FoundrySessionStore is FoundrySessionStore
+    assert foundry.FoundryAgentSessionStore is FoundryAgentSessionStore
     assert foundry.ResponsesHostServer is ResponsesHostServer
     assert foundry.FoundryLocalClient is FoundryLocalClient
     assert "FoundryChatClient" in dir(foundry)
     assert "FoundryLocalClient" in dir(foundry)
-    assert "FoundrySessionStore" in dir(foundry)
+    assert "FoundryAgentSessionStore" in dir(foundry)
     assert "ResponsesHostServer" in dir(foundry)
 
 

python/packages/foundry_hosting/README.md

@@ -2,67 +2,23 @@
 
 This package provides the integration of Agent Framework agents and workflows with the Foundry Agent Server, which can be hosted on Foundry infrastructure.
 
-`ResponsesHostServer` persists the Agent Framework `AgentSession` used by regular
-agents in addition to the Responses provider's message history. By default it
-uses the experimental `FoundrySessionStore` under `/.sessions` when hosted and
-an in-memory `SessionStore` locally. Hosted snapshots are partitioned by the
-Agent Server request context's platform user ID. Snapshot filenames use the
-Responses `conversation_id` or `response_id`, depending on the continuation
-mode.
+## State store
 
-Foundry's session file API exposes the hosted `$HOME` directory as `/`, so the
-API path `/.sessions` is stored on disk at `$HOME/.sessions`.
+### Agent Sessions
 
-Workflow agents continue to use their existing checkpoint storage layout.
+`ResponsesHostServer` persists the Agent Framework `AgentSession` durably. By default it
+uses the `FoundryAgentSessionStore` when hosted and an in-memory `SessionStore` locally.
+When hosted, the stored sessions will be isolated by the platform user ID and scoped
+under `agent_sessions`.
 
-## Foundry session isolation
+### Workflow checkpoints
 
-`FoundrySessionStore` currently subclasses core's `FileSessionStore`. It reads
-the active request through `azure.ai.agentserver.core.get_request_context()` and
-validates the platform `user_id` (the same `x-agent-user-id` value exposed as
-`ResponseContext.platform_context.user_id_key`) before selecting its on-disk
-directory.
+`ResponsesHostServer` persists workflow checkpoints durably. By default, it uses the
+`FoundryCheckpointStore` when hosted and an in-memory `InMemoryCheckpointStorage` locally.
+When hosted, the stored checkpoints will be isolated by the platform user ID and scoped
+under `checkpoints`.
 
-Regular-agent session snapshots use the platform user ID and a Responses key:
+### Function approvals
 
-```text
-/.sessions/<user-id>/<conversation-id-or-response-id>.json
-```
-
-A Foundry session controls hosted compute and filesystem lifetime and may host
-multiple users and Responses conversations. The Foundry session ID is not used
-as the MAF session identifier.
-
-When `conversation_id` is used, the host reads and writes the same snapshot
-under that ID. When `previous_response_id` is used, the host reads that response
-snapshot, runs the loaded MAF session, and writes the updated snapshot under the
-current response's `response_id`. Multiple responses can therefore branch from
-one prior response without overwriting its snapshot.
-
-Foundry does not infer the hosted `agent_session_id` from
-`previous_response_id`. Callers using response chains must also reuse the
-`agent_session_id` returned by the previous response so the request reaches the
-same sandbox and `$HOME/.sessions` filesystem. Conversation objects bind to a
-stable hosted session automatically.
-
-Workflow checkpoints and function approvals preserve the existing Foundry
-Hosting layout. Hosted paths insert the validated raw platform user ID:
-
-```text
-/.checkpoints/<user-id>/<context-id>/
-/.function_approvals/<user-id>/approval_requests.json
-```
-
-Local workflow checkpoints use `{cwd}/.checkpoints/<context-id>/`, and local
-function approvals remain in memory.
-
-Hosted requests require container protocol `2.0.0`. The v2-only request
-`call_id` is checked before session, checkpoint, or approval storage is used,
-and a missing platform user ID fails closed. Regular agents also require the
-normal Responses continuation ID for restoration. Local requests may remain
-unscoped.
-
-The Foundry-specific store type intentionally hides the current filesystem
-implementation from `ResponsesHostServer` setup. A future version may move
-`FoundrySessionStore` to a Foundry storage API without changing the host's
-default configuration.
+`ResponsesHostServer` persists function approvals durably. By default, it uses the
+`FoundryFunctionApprovalStore` when hosted and an in-memory `InMemoryFunctionApprovalStore` locally. When hosted, the stored approvals will be isolated by the platform user ID and scoped under `function_approvals`.

python/packages/foundry_hosting/agent_framework_foundry_hosting/__init__.py

@@ -4,7 +4,17 @@ import importlib.metadata
 
 from ._invocations import InvocationsHostServer
 from ._responses import ResponsesHostServer
-from ._session_store import FoundrySessionStore
+from ._state_store import (
+    AgentSessionStoreProvider,
+    CheckpointStoreProvider,
+    ContextScopedStoreProvider,
+    FoundryAgentSessionStore,
+    FoundryCheckpointStore,
+    FoundryFunctionApprovalStore,
+    FunctionApprovalStore,
+    FunctionApprovalStoreProvider,
+    StoreProvider,
+)
 from ._toolbox import FoundryToolbox
 
 try:
@@ -13,8 +23,16 @@ except importlib.metadata.PackageNotFoundError:
     __version__ = "0.0.0"
 
 __all__ = [
-    "FoundrySessionStore",
+    "AgentSessionStoreProvider",
+    "CheckpointStoreProvider",
+    "ContextScopedStoreProvider",
+    "FoundryAgentSessionStore",
+    "FoundryCheckpointStore",
+    "FoundryFunctionApprovalStore",
     "FoundryToolbox",
+    "FunctionApprovalStore",
+    "FunctionApprovalStoreProvider",
     "InvocationsHostServer",
     "ResponsesHostServer",
+    "StoreProvider",
 ]

python/packages/foundry_hosting/agent_framework_foundry_hosting/_session_store.py

@@ -1,61 +0,0 @@
-# Copyright (c) Microsoft. All rights reserved.
-
-from __future__ import annotations
-
-from pathlib import Path
-from typing import Literal
-
-from agent_framework import ExperimentalFeature, FileSessionStore
-from agent_framework._feature_stage import experimental
-from azure.ai.agentserver.core import get_request_context
-
-from ._request_context import validate_path_segment
-
-
-@experimental(feature_id=ExperimentalFeature.SESSION_STORE)
-class FoundrySessionStore(FileSessionStore):
-    """Persist MAF AgentSession snapshots within a Foundry hosted session.
-
-    A Foundry hosted session controls platform compute and filesystem lifetime
-    and may host multiple users and Responses conversations. A MAF
-    :class:`AgentSession` contains framework context state. Snapshots are keyed
-    by ``conversation_id`` for stored conversations or by Responses
-    ``response_id`` for response chains; these storage keys are independent of
-    the MAF session's own identifier.
-
-    This implementation currently persists through :class:`FileSessionStore`,
-    with each validated platform user ID as a child directory. The
-    Foundry-specific type leaves room to use a platform storage API later
-    without changing :class:`ResponsesHostServer` configuration.
-    """
-
-    def __init__(
-        self,
-        storage_path: str | Path,
-        *,
-        serialization_format: Literal["json", "msgpack"] = "json",
-    ) -> None:
-        """Initialize a Foundry-scoped file store rooted at ``storage_path``."""
-        super().__init__(storage_path, serialization_format=serialization_format)
-
-    def _session_file_path(self, session_id: str) -> Path:
-        """Resolve a snapshot path within the active Foundry user's directory."""
-        user_directory = self._storage_root
-        if user_id := get_request_context().user_id:
-            validate_path_segment(user_id, kind="user id")
-            candidate_user_directory = self._storage_root / user_id
-            user_directory = candidate_user_directory.resolve()
-            if user_directory != candidate_user_directory or not user_directory.is_relative_to(self._storage_root):
-                raise ValueError(f"User directory escaped storage directory: '{user_directory}'.")
-        user_directory.mkdir(parents=True, exist_ok=True)
-
-        session_file_name = self._session_file_name(session_id)
-        candidate_session_file_path = user_directory / session_file_name
-        session_file_path = candidate_session_file_path.resolve()
-        if (
-            session_file_path != candidate_session_file_path
-            or not session_file_path.is_relative_to(user_directory)
-            or not session_file_path.is_relative_to(self._storage_root)
-        ):
-            raise ValueError(f"Session file path escaped user directory: {session_id!r}")
-        return session_file_path

python/packages/foundry_hosting/agent_framework_foundry_hosting/_state_store.py

@@ -0,0 +1,338 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+
+from abc import ABC, abstractmethod
+from datetime import datetime
+from typing import Generic, Protocol, TypeVar
+
+from agent_framework import (
+    AgentSession,
+    CheckpointID,
+    CheckpointStorage,
+    Content,
+    InMemoryCheckpointStorage,
+    SessionStore,
+    WorkflowCheckpoint,
+    WorkflowCheckpointException,
+)
+from azure.ai.agentserver.core import AgentConfig
+from azure.ai.agentserver.core.storage import FoundryStateStore, FoundryStorageConflictError
+
+StoreT = TypeVar("StoreT")
+
+
+class StoreProvider(ABC, Generic[StoreT]):
+    """Provide store for a hosting environment."""
+
+    @abstractmethod
+    def get_store(self, *, config: AgentConfig) -> StoreT:
+        """Get store for a hosting environment.
+
+        Args:
+            config: The resolved agent server configuration.
+
+        Returns:
+            The store instance for the given hosting environment.
+        """
+
+
+class ContextScopedStoreProvider(ABC, Generic[StoreT]):
+    """Provide a context-scoped store for a hosting environment.
+
+    Use this when state must be queried or managed as a collection belonging to
+    one context, rather than accessed only by an individual item ID. The context
+    can be a conversation ID or another identifier that groups related state.
+    """
+
+    @abstractmethod
+    def get_store(self, *, config: AgentConfig, context_id: str) -> StoreT:
+        """Get a context-scoped store for a hosting environment.
+
+        Args:
+            config: The resolved agent server configuration.
+            context_id: A string that uniquely identifies the context for which the store is scoped.
+
+        Returns:
+            The store instance for the given hosting environment and context ID.
+        """
+
+
+# region Checkpoint persistence
+
+
+class FoundryCheckpointStore:
+    """Checkpoint store backed by the `FoundryStateStore`."""
+
+    DEFAULT_ROOT_SCOPE = "checkpoints"
+
+    def __init__(self, context_id: str) -> 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.
+        """
+        if not context_id:
+            raise ValueError("context_id must be provided to initialize a FoundryCheckpointStore.")
+
+        self.context_id = context_id
+
+    async def _get_store(self) -> FoundryStateStore:
+        return await FoundryStateStore.get_or_create(
+            f"{self.DEFAULT_ROOT_SCOPE}/{self.context_id}", user_isolation=True
+        )
+
+    async def save(self, checkpoint: WorkflowCheckpoint) -> CheckpointID:
+        """Save a workflow checkpoint to the store.
+
+        Args:
+            checkpoint: The workflow checkpoint to save.
+
+        The checkpoint will be serialized to a dictionary and then encoded before being stored.
+
+        Returns:
+            The ID of the saved checkpoint.
+        """
+        from agent_framework._workflows._checkpoint_encoding import encode_checkpoint_value
+
+        encoded_checkpoint = encode_checkpoint_value(checkpoint.to_dict())
+
+        store = await self._get_store()
+        async with store:
+            await store.set_item(checkpoint.checkpoint_id, encoded_checkpoint)
+            return checkpoint.checkpoint_id
+
+    async def load(self, checkpoint_id: CheckpointID) -> WorkflowCheckpoint:
+        """Load a workflow checkpoint from the store.
+
+        Args:
+            checkpoint_id: The ID of the checkpoint to load.
+
+        Returns:
+            The loaded workflow checkpoint.
+
+        Raises:
+            WorkflowCheckpointException: If no checkpoint is found with the given ID.
+        """
+        from agent_framework._workflows._checkpoint_encoding import decode_checkpoint_value
+
+        store = await self._get_store()
+        async with store:
+            item = await store.get_item(checkpoint_id)
+        if item is None:
+            raise WorkflowCheckpointException(f"No checkpoint found with ID {checkpoint_id}")
+        return WorkflowCheckpoint.from_dict(decode_checkpoint_value(item.value))
+
+    async def list_checkpoints(self, *, workflow_name: str) -> list[WorkflowCheckpoint]:
+        """List all workflow checkpoints for a given workflow name."""
+        from agent_framework._workflows._checkpoint_encoding import decode_checkpoint_value
+
+        store = await self._get_store()
+        checkpoints: list[WorkflowCheckpoint] = []
+        after: str | None = None
+        async with store:
+            while True:
+                page = await store.list_keys(after=after)
+                for item_key in page.keys:
+                    item = await store.get_item(item_key.key)
+                    if item is None:
+                        continue
+                    checkpoint = WorkflowCheckpoint.from_dict(decode_checkpoint_value(item.value))
+                    if checkpoint.workflow_name == workflow_name:
+                        checkpoints.append(checkpoint)
+                if not page.has_more or page.last_id is None:
+                    break
+                after = page.last_id
+        return checkpoints
+
+    async def delete(self, checkpoint_id: CheckpointID) -> bool:
+        store = await self._get_store()
+        async with store:
+            deleted_item = await store.delete_item(checkpoint_id)
+        return deleted_item.id is not None
+
+    async def get_latest(self, *, workflow_name: str) -> WorkflowCheckpoint | None:
+        checkpoints = await self.list_checkpoints(workflow_name=workflow_name)
+        if not checkpoints:
+            return None
+        return max(checkpoints, key=lambda checkpoint: datetime.fromisoformat(checkpoint.timestamp))
+
+    async def list_checkpoint_ids(self, *, workflow_name: str) -> list[CheckpointID]:
+        checkpoints = await self.list_checkpoints(workflow_name=workflow_name)
+        return [checkpoint.checkpoint_id for checkpoint in checkpoints]
+
+
+class CheckpointStoreProvider(ContextScopedStoreProvider[CheckpointStorage]):
+    """Provide workflow checkpoint store scoped to a context.
+
+    A workflow context can contain multiple checkpoints. Scoping establishes the
+    collection boundary used to list checkpoints, restore the latest checkpoint,
+    and clean up older checkpoints without affecting another workflow context.
+
+    This will default to using the `FoundryCheckpointStore` when hosted in Foundry,
+    and an in-memory store otherwise.
+    """
+
+    def __init__(self) -> None:
+        self._foundry_storages: dict[str, CheckpointStorage] = {}
+        self._in_memory_storages: dict[str, CheckpointStorage] = {}
+
+    def get_store(
+        self,
+        *,
+        config: AgentConfig,
+        context_id: str,
+    ) -> CheckpointStorage:
+        """Get checkpoint store for the requested hosting environment."""
+        stores = self._foundry_storages if config.is_hosted else self._in_memory_storages
+
+        if not context_id:
+            raise ValueError("context_id must be provided to get a checkpoint store.")
+
+        if context_id not in stores:
+            stores[context_id] = FoundryCheckpointStore(context_id) if config.is_hosted else InMemoryCheckpointStorage()
+        return stores[context_id]
+
+
+# endregion Checkpoint persistence
+
+# region Function approval persistence
+
+
+class FunctionApprovalStore(Protocol):
+    """Store for saving function approval requests."""
+
+    async def save_approval_request(self, approval_request_id: str, request: Content) -> None:
+        """Save a function approval request under the given ID."""
+        ...
+
+    async def load_approval_request(self, approval_request_id: str) -> Content:
+        """Load a function approval request by its ID."""
+        ...
+
+
+class FoundryFunctionApprovalStore:
+    """Function approval store backed by the `FoundryStateStore`.
+
+    This storage implements a hybrid approach where the checkpoint metadata and structure are
+    stored in JSON format, while the actual state data (which may contain complex Python objects)
+    is serialized using pickle and embedded as base64-encoded strings within the JSON. This allows
+    for human-readable checkpoint files while preserving the ability to store complex Python objects.
+    """
+
+    DEFAULT_ROOT_SCOPE = "function_approvals"
+
+    async def _get_store(self) -> FoundryStateStore:
+        return await FoundryStateStore.get_or_create(self.DEFAULT_ROOT_SCOPE, user_isolation=True)
+
+    async def save_approval_request(self, approval_request_id: str, request: Content) -> None:
+        store = await self._get_store()
+        async with store:
+            try:
+                await store.create_item(approval_request_id, request.to_dict())
+            except FoundryStorageConflictError as ex:
+                raise ValueError(f"Approval request with ID '{approval_request_id}' already exists.") from ex
+
+    async def load_approval_request(self, approval_request_id: str) -> Content:
+        store = await self._get_store()
+        async with store:
+            item = await store.get_item(approval_request_id)
+        if item is None:
+            raise KeyError(f"Approval request with ID '{approval_request_id}' does not exist.")
+        return Content.from_dict(item.value)
+
+
+class InMemoryFunctionApprovalStore:
+    """An in-memory store for function approval requests."""
+
+    def __init__(self) -> None:
+        self._store: dict[str, Content] = {}
+
+    async def save_approval_request(self, approval_request_id: str, request: Content) -> None:
+        if approval_request_id in self._store:
+            raise ValueError(f"Approval request with ID '{approval_request_id}' already exists.")
+        self._store[approval_request_id] = request
+
+    async def load_approval_request(self, approval_request_id: str) -> Content:
+        if approval_request_id not in self._store:
+            raise KeyError(f"Approval request with ID '{approval_request_id}' does not exist.")
+        return self._store[approval_request_id]
+
+
+class FunctionApprovalStoreProvider(StoreProvider[FunctionApprovalStore]):
+    """Provide function approval store for the active hosting environment.
+
+    This will default to using the `FoundryFunctionApprovalStore` when hosted in Foundry,
+    and an in-memory store otherwise.
+    """
+
+    def __init__(self) -> None:
+        self._foundry_storage: FunctionApprovalStore | None = None
+        self._in_memory_storage: FunctionApprovalStore | None = None
+
+    def get_store(self, *, config: AgentConfig) -> FunctionApprovalStore:
+        """Get function approval store for the requested hosting environment."""
+        if config.is_hosted:
+            if self._foundry_storage is None:
+                self._foundry_storage = FoundryFunctionApprovalStore()
+            return self._foundry_storage
+        if self._in_memory_storage is None:
+            self._in_memory_storage = InMemoryFunctionApprovalStore()
+        return self._in_memory_storage
+
+
+# endregion Function approval persistence
+
+# region Agent session persistence
+
+
+class FoundryAgentSessionStore(SessionStore):
+    """Agent session store backed by the `FoundryStateStore`."""
+
+    DEFAULT_ROOT_SCOPE = "agent_sessions"
+
+    async def _get_store(self) -> FoundryStateStore:
+        return await FoundryStateStore.get_or_create(f"{self.DEFAULT_ROOT_SCOPE}", user_isolation=True)
+
+    async def get(self, session_id: str) -> AgentSession | None:
+        store = await self._get_store()
+        async with store:
+            item = await store.get_item(session_id)
+        if item is None:
+            return None
+        return AgentSession.from_dict(item.value)
+
+    async def set(self, session_id: str, session: AgentSession) -> None:
+        store = await self._get_store()
+        async with store:
+            await store.set_item(session_id, session.to_dict())
+
+    async def delete(self, session_id: str) -> None:
+        store = await self._get_store()
+        async with store:
+            await store.delete_item(session_id)
+
+
+class AgentSessionStoreProvider(StoreProvider[SessionStore]):
+    """Provide agent session store for the active hosting environment.
+
+    This will default to using the `FoundryAgentSessionStore` when hosted in Foundry,
+    and an in-memory store otherwise.
+    """
+
+    def __init__(self) -> None:
+        self._foundry_storage: SessionStore | None = None
+        self._in_memory_storage: SessionStore | None = None
+
+    def get_store(self, *, config: AgentConfig) -> SessionStore:
+        """Get agent session store for the requested hosting environment."""
+        if config.is_hosted:
+            if self._foundry_storage is None:
+                self._foundry_storage = FoundryAgentSessionStore()
+            return self._foundry_storage
+        if self._in_memory_storage is None:
+            self._in_memory_storage = SessionStore()
+        return self._in_memory_storage
+
+
+# endregion Agent session persistence

python/packages/foundry_hosting/pyproject.toml

@@ -24,9 +24,9 @@ classifiers = [
 ]
 dependencies = [
     "agent-framework-core>=1.13.0,<2",
-    "azure-ai-agentserver-core>=2.0.0b7,<3",
-    "azure-ai-agentserver-responses>=1.0.0b8,<2",
-    "azure-ai-agentserver-invocations>=1.0.0b6,<2",
+    "azure-ai-agentserver-core>=2.0.0b11,<3",
+    "azure-ai-agentserver-responses>=2.0.0b1,<3",
+    "azure-ai-agentserver-invocations>=1.0.0b8,<2",
     "httpx>=0.28,<1",
     "mcp>=1.24.0,<2",
 ]

python/packages/foundry_hosting/tests/test_responses_int.py

@@ -532,35 +532,58 @@ class TestReasoningHostedMcpReplay:
                 "status": "completed",
             }
 
+        def _streaming_response(response_id: str, output: list[dict[str, Any]]) -> httpx.Response:
+            response = _response(response_id, output)
+            events: list[dict[str, Any]] = []
+            for output_index, item in enumerate(output):
+                events.extend([
+                    {
+                        "type": "response.output_item.added",
+                        "output_index": output_index,
+                        "item": {**item, "status": "in_progress"},
+                        "sequence_number": len(events),
+                    },
+                    {
+                        "type": "response.output_item.done",
+                        "output_index": output_index,
+                        "item": item,
+                        "sequence_number": len(events) + 1,
+                    },
+                ])
+            events.append({
+                "type": "response.completed",
+                "response": response,
+                "sequence_number": len(events),
+            })
+            body = "".join(f"data: {json.dumps(event)}\n\n" for event in events) + "data: [DONE]\n\n"
+            return httpx.Response(200, text=body, headers={"content-type": "text/event-stream"})
+
         async def foundry_responses_boundary(request: httpx.Request) -> httpx.Response:
             nonlocal call_count
             call_count += 1
             payload = json.loads(request.content)
             provider_payloads.append(payload)
             if call_count == 1:
-                return httpx.Response(
-                    200,
-                    json=_response(
-                        "resp_first",
-                        [
-                            {
-                                "encrypted_content": "encrypted-reasoning",
-                                "id": reasoning_id,
-                                "summary": [{"text": "The MCP server has the answer.", "type": "summary_text"}],
-                                "type": "reasoning",
-                            },
-                            {
-                                "id": "mcp_paired",
-                                "arguments": '{"query":"Agent Framework overview"}',
-                                "name": "microsoft_docs_search",
-                                "server_label": "Microsoft_Learn",
-                                "type": "mcp_call",
-                                "output": "Microsoft Agent Framework",
-                                "status": "completed",
-                            },
-                            _message("msg_first"),
-                        ],
-                    ),
+                return _streaming_response(
+                    "resp_first",
+                    [
+                        {
+                            "encrypted_content": "encrypted-reasoning",
+                            "id": reasoning_id,
+                            "summary": [{"text": "The MCP server has the answer.", "type": "summary_text"}],
+                            "type": "reasoning",
+                        },
+                        {
+                            "id": "mcp_paired",
+                            "arguments": '{"query":"Agent Framework overview"}',
+                            "name": "microsoft_docs_search",
+                            "server_label": "Microsoft_Learn",
+                            "type": "mcp_call",
+                            "output": "Microsoft Agent Framework",
+                            "status": "completed",
+                        },
+                        _message("msg_first"),
+                    ],
                 )
 
             input_items = payload["input"]
@@ -586,13 +609,7 @@ class TestReasoningHostedMcpReplay:
                     },
                 )
 
-            return httpx.Response(
-                200,
-                json=_response(
-                    "resp_second",
-                    [_message("msg_second")],
-                ),
-            )
+            return _streaming_response("resp_second", [_message("msg_second")])
 
         transport = httpx.MockTransport(foundry_responses_boundary)
         responses_client = AsyncOpenAI(

python/packages/foundry_hosting/tests/test_state_store.py

@@ -0,0 +1,384 @@
+# Copyright (c) Microsoft. All rights reserved.
+from collections.abc import Callable
+from types import SimpleNamespace
+from typing import Any
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+from agent_framework import AgentSession, Content, SessionStore, WorkflowCheckpoint, WorkflowCheckpointException
+from azure.ai.agentserver.core import AgentConfig
+from azure.ai.agentserver.core.storage import FoundryStorageConflictError
+
+from agent_framework_foundry_hosting import ContextScopedStoreProvider, StoreProvider
+from agent_framework_foundry_hosting._state_store import (
+    AgentSessionStoreProvider,
+    CheckpointStoreProvider,
+    FoundryAgentSessionStore,
+    FoundryCheckpointStore,
+    FoundryFunctionApprovalStore,
+    FunctionApprovalStoreProvider,
+    InMemoryFunctionApprovalStore,
+)
+
+
+def _checkpoint(
+    checkpoint_id: str, *, workflow_name: str = "workflow", timestamp: str = "2026-01-01T00:00:00+00:00"
+) -> WorkflowCheckpoint:
+    return WorkflowCheckpoint(
+        workflow_name=workflow_name,
+        graph_signature_hash="graph-hash",
+        checkpoint_id=checkpoint_id,
+        timestamp=timestamp,
+    )
+
+
+def _store() -> MagicMock:
+    store = MagicMock()
+    store.__aenter__ = AsyncMock(return_value=store)
+    store.__aexit__ = AsyncMock(return_value=None)
+    store.create_item = AsyncMock()
+    store.set_item = AsyncMock()
+    store.get_item = AsyncMock()
+    store.list_keys = AsyncMock()
+    store.delete_item = AsyncMock()
+    return store
+
+
+def _config(*, is_hosted: bool) -> AgentConfig:
+    return AgentConfig(
+        agent_name="",
+        agent_version="",
+        agent_id="",
+        is_hosted=is_hosted,
+        project_endpoint="",
+        project_id="",
+        session_id="",
+        port=8088,
+        appinsights_connection_string="",
+        otlp_endpoint="",
+        sse_keepalive_interval=0,
+    )
+
+
+def test_storage_providers_use_public_abstraction() -> None:
+    assert issubclass(CheckpointStoreProvider, ContextScopedStoreProvider)
+    assert not issubclass(CheckpointStoreProvider, StoreProvider)
+    assert issubclass(FunctionApprovalStoreProvider, StoreProvider)
+    assert issubclass(AgentSessionStoreProvider, StoreProvider)
+
+
+async def test_save_uses_context_scoped_store() -> None:
+    store = _store()
+    checkpoint = _checkpoint("checkpoint-1")
+
+    with patch(
+        "agent_framework_foundry_hosting._state_store.FoundryStateStore.get_or_create",
+        new=AsyncMock(return_value=store),
+    ) as get_or_create:
+        result = await FoundryCheckpointStore("context-1").save(checkpoint)
+
+    assert result == "checkpoint-1"
+    get_or_create.assert_awaited_once_with("checkpoints/context-1", user_isolation=True)
+    store.set_item.assert_awaited_once_with("checkpoint-1", checkpoint.to_dict())
+
+
+async def test_load_returns_checkpoint() -> None:
+    store = _store()
+    checkpoint = _checkpoint("checkpoint-1")
+    store.get_item = AsyncMock(return_value=SimpleNamespace(value=checkpoint.to_dict()))
+
+    with patch(
+        "agent_framework_foundry_hosting._state_store.FoundryStateStore.get_or_create",
+        new=AsyncMock(return_value=store),
+    ):
+        result = await FoundryCheckpointStore("context-1").load("checkpoint-1")
+
+    assert result == checkpoint
+
+
+async def test_load_raises_for_missing_checkpoint() -> None:
+    store = _store()
+    store.get_item = AsyncMock(return_value=None)
+
+    with (
+        patch(
+            "agent_framework_foundry_hosting._state_store.FoundryStateStore.get_or_create",
+            new=AsyncMock(return_value=store),
+        ),
+        pytest.raises(WorkflowCheckpointException, match="No checkpoint found with ID missing"),
+    ):
+        await FoundryCheckpointStore("context-1").load("missing")
+
+
+async def test_list_checkpoints_paginates_and_filters_by_workflow() -> None:
+    store = _store()
+    matching = _checkpoint("checkpoint-1")
+    other = _checkpoint("checkpoint-2", workflow_name="other")
+    store.list_keys = AsyncMock(
+        side_effect=[
+            SimpleNamespace(keys=[SimpleNamespace(key="checkpoint-1")], has_more=True, last_id="cursor-1"),
+            SimpleNamespace(
+                keys=[SimpleNamespace(key="deleted"), SimpleNamespace(key="checkpoint-2")], has_more=False, last_id=None
+            ),
+        ]
+    )
+    store.get_item = AsyncMock(
+        side_effect=[
+            SimpleNamespace(value=matching.to_dict()),
+            None,
+            SimpleNamespace(value=other.to_dict()),
+        ]
+    )
+
+    with patch(
+        "agent_framework_foundry_hosting._state_store.FoundryStateStore.get_or_create",
+        new=AsyncMock(return_value=store),
+    ):
+        result = await FoundryCheckpointStore("context-1").list_checkpoints(workflow_name="workflow")
+
+    assert result == [matching]
+    assert store.list_keys.await_args_list[0].kwargs == {"after": None}
+    assert store.list_keys.await_args_list[1].kwargs == {"after": "cursor-1"}
+
+
+@pytest.mark.parametrize(("deleted_id", "expected"), [("item-id", True), (None, False)])
+async def test_delete_reports_whether_checkpoint_existed(deleted_id: str | None, expected: bool) -> None:
+    store = _store()
+    store.delete_item = AsyncMock(return_value=SimpleNamespace(id=deleted_id))
+
+    with patch(
+        "agent_framework_foundry_hosting._state_store.FoundryStateStore.get_or_create",
+        new=AsyncMock(return_value=store),
+    ):
+        result = await FoundryCheckpointStore("context-1").delete("checkpoint-1")
+
+    assert result is expected
+
+
+async def test_get_latest_uses_timestamp_and_list_ids_filters() -> None:
+    storage = FoundryCheckpointStore("context-1")
+    older = _checkpoint("older", timestamp="2026-01-01T00:00:00+00:00")
+    newer = _checkpoint("newer", timestamp="2026-01-02T00:00:00+00:00")
+    storage.list_checkpoints = AsyncMock(return_value=[newer, older])  # zuban:ignore
+
+    assert await storage.get_latest(workflow_name="workflow") == newer
+    assert await storage.list_checkpoint_ids(workflow_name="workflow") == ["newer", "older"]
+
+
+async def test_get_latest_returns_none_when_no_checkpoints_exist() -> None:
+    storage = FoundryCheckpointStore("context-1")
+    storage.list_checkpoints = AsyncMock(return_value=[])  # zuban:ignore
+
+    assert await storage.get_latest(workflow_name="workflow") is None
+
+
+@pytest.mark.parametrize("is_hosted", [True, False])
+def test_checkpoint_storage_provider_caches_storage_by_context(is_hosted: bool) -> None:
+    provider = CheckpointStoreProvider()
+    config = _config(is_hosted=is_hosted)
+
+    first = provider.get_store(config=config, context_id="context-1")
+    second = provider.get_store(config=config, context_id="context-2")
+
+    if is_hosted:
+        assert type(first) is FoundryCheckpointStore
+        assert type(second) is FoundryCheckpointStore
+
+    assert provider.get_store(config=config, context_id="context-1") is first
+    assert second is not first
+
+
+@pytest.mark.parametrize(
+    "create_store",
+    [
+        lambda: FoundryCheckpointStore(""),
+        lambda: CheckpointStoreProvider().get_store(config=_config(is_hosted=True), context_id=""),
+    ],
+)
+def test_checkpoint_stores_require_context_id(create_store: Callable[[], Any]) -> None:
+    with pytest.raises(ValueError, match="context_id must be provided"):
+        create_store()
+
+
+def _approval_request(approval_request_id: str) -> Content:
+    function_call = Content.from_function_call(
+        "call-1",
+        "delete_file",
+        arguments='{"path": "/foo"}',
+        additional_properties={"server_label": "my_server"},
+    )
+    return Content.from_function_approval_request(approval_request_id, function_call)
+
+
+async def test_save_and_load_function_approval_request() -> None:
+    store = _store()
+    request = _approval_request("approval-1")
+    store.get_item = AsyncMock(return_value=SimpleNamespace(value=request.to_dict()))
+
+    with patch(
+        "agent_framework_foundry_hosting._state_store.FoundryStateStore.get_or_create",
+        new=AsyncMock(return_value=store),
+    ) as get_or_create:
+        storage = FoundryFunctionApprovalStore()
+        await storage.save_approval_request("approval-1", request)
+        loaded = await storage.load_approval_request("approval-1")
+
+    assert get_or_create.await_count == 2
+    get_or_create.assert_awaited_with("function_approvals", user_isolation=True)
+    store.create_item.assert_awaited_once_with("approval-1", request.to_dict())
+    assert loaded == request
+    function_call = loaded.function_call
+    assert function_call is not None
+    assert function_call.name == "delete_file"
+    assert function_call.additional_properties["server_label"] == "my_server"
+
+
+async def test_save_duplicate_function_approval_request_raises() -> None:
+    store = _store()
+    store.create_item = AsyncMock(side_effect=FoundryStorageConflictError("already exists"))
+
+    with (
+        patch(
+            "agent_framework_foundry_hosting._state_store.FoundryStateStore.get_or_create",
+            new=AsyncMock(return_value=store),
+        ),
+        pytest.raises(ValueError, match="Approval request with ID 'approval-1' already exists"),
+    ):
+        await FoundryFunctionApprovalStore().save_approval_request("approval-1", _approval_request("approval-1"))
+
+
+async def test_load_missing_function_approval_request_raises() -> None:
+    store = _store()
+    store.get_item = AsyncMock(return_value=None)
+
+    with (
+        patch(
+            "agent_framework_foundry_hosting._state_store.FoundryStateStore.get_or_create",
+            new=AsyncMock(return_value=store),
+        ),
+        pytest.raises(KeyError, match="Approval request with ID 'missing' does not exist"),
+    ):
+        await FoundryFunctionApprovalStore().load_approval_request("missing")
+
+
+@pytest.mark.parametrize(
+    ("is_hosted", "expected_type"),
+    [(True, FoundryFunctionApprovalStore), (False, InMemoryFunctionApprovalStore)],
+)
+def test_function_approval_storage_provider_selects_backend(
+    is_hosted: bool,
+    expected_type: type[FoundryFunctionApprovalStore] | type[InMemoryFunctionApprovalStore],
+) -> None:
+    provider = FunctionApprovalStoreProvider()
+    config = _config(is_hosted=is_hosted)
+
+    storage = provider.get_store(config=config)
+
+    assert type(storage) is expected_type
+    assert provider.get_store(config=config) is storage
+
+
+def test_function_approval_storage_provider_creates_only_requested_backend() -> None:
+    with (
+        patch("agent_framework_foundry_hosting._state_store.FoundryFunctionApprovalStore") as foundry_storage_type,
+        patch("agent_framework_foundry_hosting._state_store.InMemoryFunctionApprovalStore") as in_memory_storage_type,
+    ):
+        provider = FunctionApprovalStoreProvider()
+        config = _config(is_hosted=True)
+        storage = provider.get_store(config=config)
+
+        assert provider.get_store(config=config) is storage
+
+    foundry_storage_type.assert_called_once_with()
+    in_memory_storage_type.assert_not_called()
+
+
+async def test_set_agent_session_uses_scoped_store() -> None:
+    store = _store()
+    session = AgentSession(session_id="agent-session-1")
+    session.state["turn_count"] = 2
+
+    with patch(
+        "agent_framework_foundry_hosting._state_store.FoundryStateStore.get_or_create",
+        new=AsyncMock(return_value=store),
+    ) as get_or_create:
+        await FoundryAgentSessionStore().set("storage-session-1", session)
+
+    get_or_create.assert_awaited_once_with("agent_sessions", user_isolation=True)
+    store.set_item.assert_awaited_once_with("storage-session-1", session.to_dict())
+
+
+async def test_get_agent_session_returns_deserialized_session() -> None:
+    store = _store()
+    session = AgentSession(session_id="agent-session-1")
+    session.state["turn_count"] = 2
+    store.get_item = AsyncMock(return_value=SimpleNamespace(value=session.to_dict()))
+
+    with patch(
+        "agent_framework_foundry_hosting._state_store.FoundryStateStore.get_or_create",
+        new=AsyncMock(return_value=store),
+    ) as get_or_create:
+        result = await FoundryAgentSessionStore().get("storage-session-1")
+
+    assert result is not None
+    assert result.to_dict() == session.to_dict()
+    assert result is not session
+    get_or_create.assert_awaited_once_with("agent_sessions", user_isolation=True)
+    store.get_item.assert_awaited_once_with("storage-session-1")
+
+
+async def test_get_missing_agent_session_returns_none() -> None:
+    store = _store()
+    store.get_item = AsyncMock(return_value=None)
+
+    with patch(
+        "agent_framework_foundry_hosting._state_store.FoundryStateStore.get_or_create",
+        new=AsyncMock(return_value=store),
+    ):
+        result = await FoundryAgentSessionStore().get("missing")
+
+    assert result is None
+
+
+async def test_delete_agent_session_is_idempotent() -> None:
+    store = _store()
+
+    with patch(
+        "agent_framework_foundry_hosting._state_store.FoundryStateStore.get_or_create",
+        new=AsyncMock(return_value=store),
+    ):
+        await FoundryAgentSessionStore().delete("storage-session-1")
+
+    store.delete_item.assert_awaited_once_with("storage-session-1")
+
+
+@pytest.mark.parametrize(
+    ("is_hosted", "expected_type"),
+    [(True, FoundryAgentSessionStore), (False, SessionStore)],
+)
+def test_agent_session_storage_provider_selects_backend(
+    is_hosted: bool,
+    expected_type: type[FoundryAgentSessionStore] | type[SessionStore],
+) -> None:
+    provider = AgentSessionStoreProvider()
+    config = _config(is_hosted=is_hosted)
+
+    storage = provider.get_store(config=config)
+
+    assert type(storage) is expected_type
+    assert provider.get_store(config=config) is storage
+
+
+def test_agent_session_storage_provider_creates_only_requested_backend() -> None:
+    with (
+        patch("agent_framework_foundry_hosting._state_store.FoundryAgentSessionStore") as foundry_storage_type,
+        patch("agent_framework_foundry_hosting._state_store.SessionStore") as in_memory_storage_type,
+    ):
+        provider = AgentSessionStoreProvider()
+        config = _config(is_hosted=False)
+        storage = provider.get_store(config=config)
+
+        assert provider.get_store(config=config) is storage
+
+    foundry_storage_type.assert_not_called()
+    in_memory_storage_type.assert_called_once_with()

python/pyproject.toml

@@ -63,7 +63,7 @@ prerelease = "if-necessary-or-explicit"
 # Security floors for transitive deps; overrides bypass litellm[proxy]'s strict pins.
 constraint-dependencies = ["litellm>=1.83.7", "fastapi-sso>=0.19.0"]
 # python-multipart>=0.0.31 overrides litellm[proxy]'s exact pin of <=0.0.27 for security.
-override-dependencies = ["mcp[ws]>=1.27.0", "uvicorn[standard]>=0.34.0", "python-multipart>=0.0.31"]
+override-dependencies = ["mcp[ws]>=1.27.0,<2", "uvicorn[standard]>=0.34.0", "python-multipart>=0.0.31"]
 environments = [
     "sys_platform == 'darwin'",
     "sys_platform == 'linux'",