Python: Update agentserver to x.1.0b1 (#7621)

* Update agentserver to 2.1.0

* Update agentserver responses and invocations to x.1.0b1

* Pass platform context to state store provider

* Pass user id

* Correct requirements.txt

* Fix unit tests

* Fix unit tests

python/packages/azure-contentunderstanding/agent_framework_azure_contentunderstanding/_context_provider.py

@@ -762,7 +762,7 @@ class ContentUnderstandingContextProvider(ContextProvider):
             result,
             include_markdown="markdown" in self.output_sections,
             include_fields="fields" in self.output_sections,
-            metadata={"source": filename},
+            custom_metadata={"source": filename},
         )
 
     # ------------------------------------------------------------------

python/packages/azure-contentunderstanding/tests/cu/test_context_provider.py

@@ -1663,13 +1663,13 @@ class TestAnalyzerAutoDetectionE2E:
 class TestWarningsExtraction:
     """Verify that CU RAI warnings are surfaced via ``to_llm_input`` rendering.
 
-    The SDK serializes ``result.warnings`` under the reserved ``rai_warnings``
+    The SDK serializes ``result.warnings`` under the reserved ``warnings``
     YAML front-matter key. Telemetry filtering of stray ``LLMStats:`` lines is
     handled by the SDK helper (azure-ai-contentunderstanding >= 1.2.0b2).
     """
 
     def test_warnings_included_when_present(self) -> None:
-        """Non-empty warnings should appear under ``rai_warnings`` front-matter key."""
+        """Non-empty warnings should appear under ``warnings`` front-matter key."""
         provider = _make_provider()
         fixture = {
             "contents": [
@@ -1694,16 +1694,16 @@ class TestWarningsExtraction:
         result_obj = AnalysisResult(fixture)
         rendered = provider._render_for_llm(result_obj, "doc.pdf")
 
-        assert "rai_warnings:" in rendered
+        assert "warnings:" in rendered
         assert "ContentFiltered" in rendered
         assert "Content was filtered due to Responsible AI policy." in rendered
         assert "Violence content detected and filtered." in rendered
 
     def test_warnings_omitted_when_empty(self, pdf_analysis_result: AnalysisResult) -> None:
-        """The PDF fixture has no warnings, so ``rai_warnings:`` should not appear."""
+        """The PDF fixture has no warnings, so ``warnings:`` should not appear."""
         provider = _make_provider()
         rendered = provider._render_for_llm(pdf_analysis_result, "report.pdf")
-        assert "rai_warnings:" not in rendered
+        assert "warnings:" not in rendered
 
 
 class TestCategoryExtraction:

python/packages/foundry_hosting/README.md

@@ -4,21 +4,46 @@ This package provides the integration of Agent Framework agents and workflows wi
 
 ## State store
 
+### Local persistence
+
+Outside the Foundry hosting environment, state is persisted as JSON files under
+`~/.agentserver/state_stores` by default. Set `AGENTSERVER_STATE_ROOT` to use a
+different root directory; the files will be written to its `state_stores`
+subdirectory instead.
+
+Each logical store is saved as one JSON file whose name is a URL-safe Base64
+encoding of the store name. For example:
+
+- Agent sessions: `YWdlbnRfc2Vzc2lvbnM.json`
+- Function approvals: `ZnVuY3Rpb25fYXBwcm92YWxz.json`
+- Workflow checkpoints: one file per context, encoded from `checkpoints/<context_id>`
+
+> Read more about the Foundry durable state store in the [developer guide](https://github.com/Azure/azure-sdk-for-python/blob/main/sdk/agentserver/azure-ai-agentserver-core/docs/state-store-guide.md).
+
+### User isolation
+
+When hosted on Foundry, the default state stores automatically isolate data by the
+platform user ID supplied with each request. Sessions, workflow checkpoints, and
+function approvals written for one user cannot be read or modified by another user.
+No additional partitioning configuration is required when using the default stores.
+
 ### Agent Sessions
 
 `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`.
+uses the `FoundryAgentSessionStore`, backed by Foundry storage when hosted and file-based
+storage locally. Stored sessions are scoped under `agent_sessions`.
+
+See the [custom storage provider sample](../../samples/04-hosting/foundry-hosted-agents/responses/custom_storage/)
+for an example that uses an in-memory session store locally and Azure Cosmos DB when hosted.
 
 ### Workflow checkpoints
 
 `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`.
+`FoundryCheckpointStore`, backed by Foundry storage when hosted and file-based storage
+locally. Stored checkpoints are scoped under `checkpoints`.
 
 ### Function approvals
 
 `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`.
+`FoundryFunctionApprovalStore`, backed by Foundry storage when hosted and file-based
+storage locally. Stored approvals are scoped under `function_approvals`.

python/packages/foundry_hosting/agent_framework_foundry_hosting/_invocations.py

@@ -9,7 +9,6 @@ from starlette.responses import Response, StreamingResponse
 from typing_extensions import Any, AsyncGenerator
 
 from ._feature_usage import FeatureIndex
-from ._request_context import validate_foundry_request_context
 
 
 class InvocationsHostServer(InvocationAgentServerHost):
@@ -55,7 +54,6 @@ class InvocationsHostServer(InvocationAgentServerHost):
             RuntimeError: If the context doesn't contain the expected IDs.
         """
         context = get_request_context()
-        validate_foundry_request_context(context, is_hosted=self.config.is_hosted)
 
         if self.config.is_hosted:
             if not context.session_id or not context.user_id:

python/packages/foundry_hosting/agent_framework_foundry_hosting/_request_context.py

@@ -1,53 +0,0 @@
-# Copyright (c) Microsoft. All rights reserved.
-
-from __future__ import annotations
-
-import os
-from typing import Literal
-
-from azure.ai.agentserver.core import FoundryAgentRequestContext
-
-_PROTOCOL_V2_REQUIRED_MESSAGE = (
-    "The hosted environment is running on protocol 1.0.0, but the agent requires protocol 2.0.0. "
-    "Please upgrade your agent protocol to 2.0.0 in `agent.manifest.yaml` or `agent.yaml`, or "
-    "downgrade the `agent-framework-foundry-hosting` package to `1.0.0a260625` or before to use 1.0.0."
-)
-
-
-def validate_path_segment(
-    segment: str,
-    *,
-    kind: Literal["context id", "user id"],
-) -> None:
-    """Validate that ``segment`` is a single safe path component (CWE-22).
-
-    Request context values are untrusted when used as path segments. Reject
-    separators, drive letters, parent references, and similar values rather
-    than attempting to sanitize them and risk collisions.
-    """
-    if not isinstance(segment, str) or not segment:
-        raise RuntimeError(f"Invalid {kind}: must be a non-empty string.")
-    if (
-        "/" in segment
-        or "\\" in segment
-        or "\x00" in segment
-        or segment.strip(".") == ""
-        or os.path.isabs(segment)
-        or os.path.splitdrive(segment)[0]
-    ):
-        raise RuntimeError(f"Invalid {kind}: {segment!r}")
-
-
-def validate_foundry_request_context(
-    context: FoundryAgentRequestContext,
-    *,
-    is_hosted: bool,
-) -> None:
-    """Validate that a hosted request contains protocol-v2 user identity."""
-    if is_hosted and context.call_id is None:
-        raise RuntimeError(_PROTOCOL_V2_REQUIRED_MESSAGE)
-    if is_hosted and not context.user_id:
-        raise RuntimeError(
-            "The hosted environment is missing the platform user ID in the request context. "
-            "Please ensure that the request is coming from a valid Foundry platform service."
-        )

python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py

@@ -6,6 +6,7 @@ import asyncio
 import base64
 import json
 import logging
+import os
 from collections.abc import AsyncIterable, AsyncIterator, Generator, Mapping, Sequence
 from contextlib import AbstractAsyncContextManager, AsyncExitStack
 from dataclasses import asdict, dataclass, is_dataclass
@@ -62,10 +63,6 @@ from mcp import McpError
 from typing_extensions import Any
 
 from ._feature_usage import FeatureIndex
-from ._request_context import (
-    validate_foundry_request_context,
-    validate_path_segment,
-)
 from ._state_store import (
     AgentSessionStoreProvider,
     CheckpointStoreProvider,
@@ -80,6 +77,20 @@ logger = logging.getLogger(__name__)
 _HOSTED_RESPONSES_HISTORY_SOURCE_ID = "_foundry_responses_history"
 
 
+def _validate_checkpoint_context_id(context_id: str) -> None:
+    """Validate that a checkpoint context ID is a single safe path component in case file-based storage is used."""
+    if (
+        not context_id
+        or "/" in context_id
+        or "\\" in context_id
+        or "\x00" in context_id
+        or context_id.strip(".") == ""
+        or os.path.isabs(context_id)
+        or os.path.splitdrive(context_id)[0]
+    ):
+        raise RuntimeError(f"Invalid context id: {context_id!r}")
+
+
 def _is_hosted_responses_history_sentinel(provider: ContextProvider) -> bool:
     """Return whether ``provider`` is the host's transient history buffer."""
     return (
@@ -311,9 +322,6 @@ class ResponsesHostServer(ResponsesAgentServerHost):
         cancellation_signal: asyncio.Event,
     ) -> AsyncIterable[ResponseStreamEvent | dict[str, Any]]:
         """Handle the creation of a response."""
-        request_context = get_request_context()
-        validate_foundry_request_context(request_context, is_hosted=self.config.is_hosted)
-
         if self._is_workflow_agent:
             # Workflow agents are handled differently because they require checkpoint restoration
             return self._handle_inner_workflow(request, context)
@@ -370,32 +378,24 @@ class ResponsesHostServer(ResponsesAgentServerHost):
             return
 
         try:
-            approval_storage = self._function_approval_storage_provider.get_store(config=self.config)
-            session_storage = self._session_storage_provider.get_store(config=self.config)
-            # Agent sessions are either tied to the conversation_id (for multi-turn conversation mode)
-            # or the previous_response_id (for response chaining). If neither is present, a new session
-            # is created for this request and stored under the current response_id. The current response_id
-            # will become the previous_response_id for the next request in a response chain, allowing the
-            # session to be retrieved.
-            if (previous_response_id := request.get("previous_response_id")) is not None:
-                session = await session_storage.get(previous_response_id)
-                if session is None:
+            request_context = get_request_context()
+            approval_storage = self._function_approval_storage_provider.get_store(
+                config=self.config, platform_context=request_context
+            )
+            session_storage = self._session_storage_provider.get_store(
+                config=self.config, platform_context=request_context
+            )
+
+            previous_response_id = request.get("previous_response_id")
+            session_load_id = context.conversation_id or previous_response_id
+            session = await session_storage.get(session_load_id) if session_load_id is not None else None
+            if session is None:
+                if previous_response_id is not None and context.conversation_id is None:
                     raise RuntimeError(
-                        f"Cannot find an existing agent session for previous_response_id={previous_response_id}. "
-                        "Ensure that the previous response was created successfully and that the ID is correct."
+                        f"Cannot find an existing agent session for previous_response_id={previous_response_id}."
                     )
-            elif (conversation_id := context.conversation_id) is not None:
-                session = await session_storage.get(conversation_id)
-                if session is None:
-                    # Note that we cannot determine if the session was deleted or never existed,
-                    # so we log a warning and create a new session.
-                    logger.info(
-                        "Cannot find an existing agent session for id=%s. Creating a new session.",
-                        conversation_id,
-                    )
-                    session = self._agent.create_session()
-            else:
                 session = self._agent.create_session()
+            session_save_id = context.conversation_id or context.response_id
         except Exception as ex:
             logger.error("Failed to prepare state storage: %s", ex, exc_info=(type(ex), ex, ex.__traceback__))
             for event in self._emit_failure(response_event_stream, None, ex):
@@ -456,7 +456,7 @@ class ResponsesHostServer(ResponsesAgentServerHost):
             if self._uses_hosted_responses_history:
                 session.state.pop(_HOSTED_RESPONSES_HISTORY_SOURCE_ID, None)
             try:
-                await session_storage.set(context.conversation_id or context.response_id, session)
+                await session_storage.set(session_save_id, session)
             except Exception as save_error:
                 save_failure = save_error
                 if request_interrupted:
@@ -498,7 +498,10 @@ class ResponsesHostServer(ResponsesAgentServerHost):
         tracker: _OutputItemTracker | None = None
 
         try:
-            approval_storage = self._function_approval_storage_provider.get_store(config=self.config)
+            request_context = get_request_context()
+            approval_storage = self._function_approval_storage_provider.get_store(
+                config=self.config, platform_context=request_context
+            )
             input_items = await context.get_input_items()
             input_messages = await _items_to_messages(input_items, approval_storage=approval_storage)
 
@@ -506,10 +509,6 @@ class ResponsesHostServer(ResponsesAgentServerHost):
             if are_options_set:
                 logger.warning("Workflow agent doesn't support runtime options. They will be ignored.")
 
-            if request.get("previous_response_id") is not None and context.conversation_id is not None:
-                raise RuntimeError("Previous response ID cannot be used in conjunction with conversation ID.")
-            context_id = request.get("previous_response_id") or context.conversation_id
-
             if not isinstance(self._agent, WorkflowAgent):
                 raise RuntimeError("Agent is not a workflow agent.")
 
@@ -518,42 +517,42 @@ class ResponsesHostServer(ResponsesAgentServerHost):
             # any future async resources owned by the workflow are entered here.
             await self._ensure_agent_ready()
 
+            checkpoint_save_id = context.conversation_id or context.response_id
+            _validate_checkpoint_context_id(checkpoint_save_id)
+            checkpoint_storage = self._checkpoint_storage_provider.get_store(
+                config=self.config,
+                context_id=checkpoint_save_id,
+                platform_context=request_context,
+            )
+
             # Determine the latest checkpoint (if any) so we can resume the
             # workflow's prior state for this turn. The directory is keyed by
-            # the inbound context id (conversation_id when set, otherwise
-            # previous_response_id). Multi-turn declarative workflows need the
-            # workflow's internal state (e.g. Conversation.messages,
+            # the platform derived context_id. Multi-turn declarative workflows
+            # need the workflow's internal state (e.g. Conversation.messages,
             # intermediate Local.* variables) to survive across user turns;
             # the only place that state lives is the workflow checkpoint, so
             # on every turn we restore the latest checkpoint and feed the new
             # input back into the start executor as a continuation rather than
             # a fresh run.
-            latest_checkpoint_id: str | None = None
-            restore_storage: CheckpointStorage | None = None
-            if context_id is not None:
-                validate_path_segment(context_id, kind="context id")
-                restore_storage = self._checkpoint_storage_provider.get_store(
-                    config=self.config,
-                    context_id=context_id,
-                )
-                latest_checkpoint = await restore_storage.get_latest(workflow_name=self._agent.workflow.name)
-                if latest_checkpoint is not None:
-                    latest_checkpoint_id = latest_checkpoint.checkpoint_id
-
-            # Storage that will receive checkpoints written during this turn.
-            # When the caller chains with previous_response_id, the next turn
-            # will reference the current response_id as its previous_response_id,
-            # so new checkpoints must land under the current response_id (or the
-            # conversation_id when set). When conversation_id is set, this
-            # matches restore_storage; when only previous_response_id was
-            # supplied, restore_storage points at the *prior* response's
-            # directory and write_storage points at the *current* response's.
-            write_context_id = context.conversation_id or context.response_id
-            validate_path_segment(write_context_id, kind="context id")
-            write_storage = self._checkpoint_storage_provider.get_store(
-                config=self.config,
-                context_id=write_context_id,
-            )
+            if request.get("previous_response_id") is not None and context.conversation_id is not None:
+                raise RuntimeError("Previous response ID cannot be used in conjunction with conversation ID.")
+            previous_response_id = request.get("previous_response_id")
+            checkpoint_load_id = context.conversation_id or previous_response_id
+            latest_checkpoint = None
+            restore_checkpoint_storage = checkpoint_storage
+            if checkpoint_load_id is not None:
+                _validate_checkpoint_context_id(checkpoint_load_id)
+                if checkpoint_load_id != checkpoint_save_id:
+                    restore_checkpoint_storage = self._checkpoint_storage_provider.get_store(
+                        config=self.config,
+                        context_id=checkpoint_load_id,
+                        platform_context=request_context,
+                    )
+                latest_checkpoint = await restore_checkpoint_storage.get_latest(workflow_name=self._agent.workflow.name)
+                if latest_checkpoint is None and previous_response_id is not None:
+                    raise RuntimeError(
+                        f"Cannot find an existing workflow checkpoint for previous_response_id={previous_response_id}."
+                    )
 
             # Multi-turn pattern: when we have a prior checkpoint, restore it
             # first (drive the workflow back to idle with prior state intact),
@@ -571,11 +570,11 @@ class ResponsesHostServer(ResponsesAgentServerHost):
             # ``run(input_messages, ...)`` call may contain ``function_call_output``
             # items (carried as FunctionResult/FunctionApprovalResponse content)
             # that fulfill them via :meth:`WorkflowAgent._process_pending_requests`.
-            if latest_checkpoint_id is not None:
+            if latest_checkpoint is not None:
                 async for _ in self._agent.run(
                     stream=True,
-                    checkpoint_id=latest_checkpoint_id,
-                    checkpoint_storage=restore_storage,
+                    checkpoint_id=latest_checkpoint.checkpoint_id,
+                    checkpoint_storage=restore_checkpoint_storage,
                 ):
                     pass
 
@@ -585,7 +584,7 @@ class ResponsesHostServer(ResponsesAgentServerHost):
             async for update in self._agent.run(
                 input_messages,
                 stream=True,
-                checkpoint_storage=write_storage,
+                checkpoint_storage=checkpoint_storage,
             ):
                 for content in update.contents:
                     for event in tracker.handle(content):
@@ -600,27 +599,12 @@ class ResponsesHostServer(ResponsesAgentServerHost):
             # Close any remaining active builder
             for event in tracker.close():
                 yield event
-
-            await self._delete_not_latest_checkpoints(write_storage, self._agent.workflow.name)
             yield response_event_stream.emit_completed()
         except Exception as ex:
             logger.exception("Failed to produce response for workflow agent")
             for event in self._emit_failure(response_event_stream, tracker, ex):
                 yield event
 
-    @staticmethod
-    async def _delete_not_latest_checkpoints(checkpoint_storage: CheckpointStorage, workflow_name: str) -> None:
-        """Delete all checkpoints except the latest one.
-
-        We only need the last checkpoint for each invocation.
-        """
-        latest_checkpoint = await checkpoint_storage.get_latest(workflow_name=workflow_name)
-        if latest_checkpoint is not None:
-            all_checkpoints = await checkpoint_storage.list_checkpoints(workflow_name=workflow_name)
-            for checkpoint in all_checkpoints:
-                if checkpoint.checkpoint_id != latest_checkpoint.checkpoint_id:
-                    await checkpoint_storage.delete(checkpoint.checkpoint_id)
-
     @staticmethod
     def _emit_failure(
         response_event_stream: ResponseEventStream,

python/packages/foundry_hosting/agent_framework_foundry_hosting/_state_store.py

@@ -10,12 +10,11 @@ from agent_framework import (
     CheckpointID,
     CheckpointStorage,
     Content,
-    InMemoryCheckpointStorage,
     SessionStore,
     WorkflowCheckpoint,
     WorkflowCheckpointException,
 )
-from azure.ai.agentserver.core import AgentConfig
+from azure.ai.agentserver.core import AgentConfig, FoundryAgentRequestContext
 from azure.ai.agentserver.core.storage import FoundryStateStore, FoundryStorageConflictError
 
 StoreT = TypeVar("StoreT")
@@ -25,11 +24,12 @@ class StoreProvider(ABC, Generic[StoreT]):
     """Provide store for a hosting environment."""
 
     @abstractmethod
-    def get_store(self, *, config: AgentConfig) -> StoreT:
+    def get_store(self, *, config: AgentConfig, platform_context: FoundryAgentRequestContext) -> StoreT:
         """Get store for a hosting environment.
 
         Args:
             config: The resolved agent server configuration.
+            platform_context: The request-scoped platform context for the current request.
 
         Returns:
             The store instance for the given hosting environment.
@@ -45,12 +45,19 @@ class ContextScopedStoreProvider(ABC, Generic[StoreT]):
     """
 
     @abstractmethod
-    def get_store(self, *, config: AgentConfig, context_id: str) -> StoreT:
+    def get_store(
+        self,
+        *,
+        config: AgentConfig,
+        context_id: str,
+        platform_context: FoundryAgentRequestContext,
+    ) -> 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.
+            platform_context: The request-scoped platform context for the current request.
 
         Returns:
             The store instance for the given hosting environment and context ID.
@@ -65,21 +72,25 @@ class FoundryCheckpointStore:
 
     DEFAULT_ROOT_SCOPE = "checkpoints"
 
-    def __init__(self, context_id: str) -> None:
+    def __init__(self, context_id: str, platform_context: FoundryAgentRequestContext) -> 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.
         """
         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
 
     async def _get_store(self) -> FoundryStateStore:
         return await FoundryStateStore.get_or_create(
-            f"{self.DEFAULT_ROOT_SCOPE}/{self.context_id}", user_isolation=True
+            f"{self.DEFAULT_ROOT_SCOPE}/{self.context_id}",
+            user_isolation=True,
+            user_id=self.platform_context.user_id,
         )
 
     async def save(self, checkpoint: WorkflowCheckpoint) -> CheckpointID:
@@ -99,7 +110,7 @@ class FoundryCheckpointStore:
 
         store = await self._get_store()
         async with store:
-            await store.set_item(checkpoint.checkpoint_id, encoded_checkpoint)
+            await store.set_item(checkpoint.checkpoint_id, encoded_checkpoint, call_id=self.platform_context.call_id)
             return checkpoint.checkpoint_id
 
     async def load(self, checkpoint_id: CheckpointID) -> WorkflowCheckpoint:
@@ -118,7 +129,7 @@ class FoundryCheckpointStore:
 
         store = await self._get_store()
         async with store:
-            item = await store.get_item(checkpoint_id)
+            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))
@@ -132,9 +143,9 @@ class FoundryCheckpointStore:
         after: str | None = None
         async with store:
             while True:
-                page = await store.list_keys(after=after)
+                page = await store.list_keys(after=after, call_id=self.platform_context.call_id)
                 for item_key in page.keys:
-                    item = await store.get_item(item_key.key)
+                    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))
@@ -148,7 +159,7 @@ class FoundryCheckpointStore:
     async def delete(self, checkpoint_id: CheckpointID) -> bool:
         store = await self._get_store()
         async with store:
-            deleted_item = await store.delete_item(checkpoint_id)
+            deleted_item = await store.delete_item(checkpoint_id, call_id=self.platform_context.call_id)
         return deleted_item.id is not None
 
     async def get_latest(self, *, workflow_name: str) -> WorkflowCheckpoint | None:
@@ -169,29 +180,21 @@ class CheckpointStoreProvider(ContextScopedStoreProvider[CheckpointStorage]):
     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.
+    This defaults to using the `FoundryCheckpointStore` in all environments.
     """
 
-    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,
+        platform_context: FoundryAgentRequestContext,
     ) -> 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]
+        return FoundryCheckpointStore(context_id, platform_context)
 
 
 # endregion Checkpoint persistence
@@ -222,63 +225,42 @@ class FoundryFunctionApprovalStore:
 
     DEFAULT_ROOT_SCOPE = "function_approvals"
 
+    def __init__(self, platform_context: FoundryAgentRequestContext) -> None:
+        self.platform_context = platform_context
+
     async def _get_store(self) -> FoundryStateStore:
-        return await FoundryStateStore.get_or_create(self.DEFAULT_ROOT_SCOPE, user_isolation=True)
+        return await FoundryStateStore.get_or_create(
+            self.DEFAULT_ROOT_SCOPE,
+            user_isolation=True,
+            user_id=self.platform_context.user_id,
+        )
 
     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())
+                await store.create_item(approval_request_id, request.to_dict(), call_id=self.platform_context.call_id)
             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)
+            item = await store.get_item(approval_request_id, call_id=self.platform_context.call_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.
+    This defaults to using the `FoundryFunctionApprovalStore` in all environments.
     """
 
-    def __init__(self) -> None:
-        self._foundry_storage: FunctionApprovalStore | None = None
-        self._in_memory_storage: FunctionApprovalStore | None = None
-
-    def get_store(self, *, config: AgentConfig) -> FunctionApprovalStore:
+    def get_store(self, *, config: AgentConfig, platform_context: FoundryAgentRequestContext) -> 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
+        return FoundryFunctionApprovalStore(platform_context)
 
 
 # endregion Function approval persistence
@@ -291,13 +273,20 @@ class FoundryAgentSessionStore(SessionStore):
 
     DEFAULT_ROOT_SCOPE = "agent_sessions"
 
+    def __init__(self, platform_context: FoundryAgentRequestContext) -> None:
+        self.platform_context = platform_context
+
     async def _get_store(self) -> FoundryStateStore:
-        return await FoundryStateStore.get_or_create(f"{self.DEFAULT_ROOT_SCOPE}", user_isolation=True)
+        return await FoundryStateStore.get_or_create(
+            f"{self.DEFAULT_ROOT_SCOPE}",
+            user_isolation=True,
+            user_id=self.platform_context.user_id,
+        )
 
     async def get(self, session_id: str) -> AgentSession | None:
         store = await self._get_store()
         async with store:
-            item = await store.get_item(session_id)
+            item = await store.get_item(session_id, call_id=self.platform_context.call_id)
         if item is None:
             return None
         return AgentSession.from_dict(item.value)
@@ -305,34 +294,23 @@ class FoundryAgentSessionStore(SessionStore):
     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())
+            await store.set_item(session_id, session.to_dict(), call_id=self.platform_context.call_id)
 
     async def delete(self, session_id: str) -> None:
         store = await self._get_store()
         async with store:
-            await store.delete_item(session_id)
+            await store.delete_item(session_id, call_id=self.platform_context.call_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.
+    This defaults to using the `FoundryAgentSessionStore` in all environments.
     """
 
-    def __init__(self) -> None:
-        self._foundry_storage: SessionStore | None = None
-        self._in_memory_storage: SessionStore | None = None
-
-    def get_store(self, *, config: AgentConfig) -> SessionStore:
+    def get_store(self, *, config: AgentConfig, platform_context: FoundryAgentRequestContext) -> 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
+        return FoundryAgentSessionStore(platform_context)
 
 
 # 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.0b11,<3",
-    "azure-ai-agentserver-responses>=2.0.0b1,<3",
-    "azure-ai-agentserver-invocations>=1.0.0b8,<2",
+    "azure-ai-agentserver-core>=2.1.0b1,<3",
+    "azure-ai-agentserver-responses>=2.1.0b1,<3",
+    "azure-ai-agentserver-invocations>=1.1.0b1,<2",
     "httpx>=0.28,<1",
     "mcp>=1.24.0,<2",
 ]

python/packages/foundry_hosting/tests/test_invocations.py

@@ -169,21 +169,12 @@ class TestPartitionKey:
         with _request_context(), pytest.raises(RuntimeError, match="missing session_id"):
             server._partition_key()  # pyright: ignore[reportPrivateUsage]
 
-    def test_hosted_without_call_id_raises_protocol_error(self) -> None:
-        server = InvocationsHostServer(_make_agent(response_text="hi"))
-        server.config.is_hosted = True
-        with (
-            _request_context(session_id="sess-1", user_id="user-1"),
-            pytest.raises(RuntimeError, match="protocol 2.0.0"),
-        ):
-            server._partition_key()  # pyright: ignore[reportPrivateUsage]
-
     def test_hosted_missing_user_id_raises(self) -> None:
         server = InvocationsHostServer(_make_agent(response_text="hi"))
         server.config.is_hosted = True
         with (
             _request_context(call_id="call-1", session_id="sess-1"),
-            pytest.raises(RuntimeError, match="missing the platform user ID"),
+            pytest.raises(RuntimeError, match="missing session_id or user_id"),
         ):
             server._partition_key()  # pyright: ignore[reportPrivateUsage]
 

python/packages/foundry_hosting/tests/test_responses.py

@@ -13,8 +13,7 @@ from __future__ import annotations
 import asyncio
 import json
 import uuid
-from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Generator, Mapping, Sequence
-from contextlib import contextmanager
+from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Mapping, Sequence
 from dataclasses import dataclass
 from typing import Literal, cast, overload
 from unittest.mock import AsyncMock, MagicMock, patch
@@ -47,11 +46,7 @@ from agent_framework import (
     executor,
     tool,
 )
-from azure.ai.agentserver.core import (
-    FoundryAgentRequestContext,
-    reset_request_context,
-    set_request_context,
-)
+from azure.ai.agentserver.core import get_request_context
 from azure.ai.agentserver.responses import InMemoryResponseProvider, ResponseContext
 from azure.ai.agentserver.responses.models import CreateResponse, Item, OutputItem
 from mcp import McpError
@@ -70,10 +65,15 @@ from agent_framework_foundry_hosting._state_store import (
     AgentSessionStoreProvider,
     CheckpointStoreProvider,
     FunctionApprovalStoreProvider,
-    InMemoryFunctionApprovalStore,
 )
 
 
+def _function_approval_store(request: Content) -> MagicMock:
+    storage = MagicMock()
+    storage.load_approval_request = AsyncMock(return_value=request)
+    return storage
+
+
 def _make_function_approval_request_content(
     *,
     request_id: str = "apr_test",
@@ -89,21 +89,6 @@ def _make_function_approval_request_content(
     return Content.from_function_approval_request(request_id, function_call)
 
 
-@contextmanager
-def _request_context(
-    *,
-    call_id: str | None = None,
-    user_id: str | None = None,
-    session_id: str | None = None,
-) -> Generator[None]:
-    """Install a Foundry request context for the duration of the block."""
-    token = set_request_context(FoundryAgentRequestContext(call_id=call_id, user_id=user_id, session_id=session_id))
-    try:
-        yield
-    finally:
-        reset_request_context(token)
-
-
 # region Helpers
 
 
@@ -419,44 +404,6 @@ class TestResponsesHostServerInit:
         with pytest.raises(RuntimeError, match="history provider"):
             ResponsesHostServer(agent)
 
-    async def test_hosted_request_requires_user_partition_key(self) -> None:
-        agent = _make_agent(
-            response=AgentResponse(messages=[Message(role="assistant", contents=[Content.from_text("hi")])])
-        )
-        server = _make_server(agent)
-        request = CreateResponse(model="m", input="hi")
-        context = ResponseContext(
-            response_id="caresp_aaaaaaaaaaaaaaaa00" + "1" * 32,
-            mode_flags=MagicMock(),
-        )
-
-        with (
-            patch.object(server.config, "is_hosted", True),
-            _request_context(call_id="call-1"),
-            pytest.raises(RuntimeError, match="platform user ID"),
-        ):
-            await server._handle_response(  # pyright: ignore[reportPrivateUsage]
-                request,
-                context,
-                asyncio.Event(),
-            )
-
-    async def test_hosted_request_requires_protocol_v2(self) -> None:
-        server = _make_server(_make_agent())
-        request = CreateResponse(model="m", input="hi")
-        context = ResponseContext(response_id="response-1", mode_flags=MagicMock())
-
-        with (
-            patch.object(server.config, "is_hosted", True),
-            _request_context(user_id="user-1"),
-            pytest.raises(RuntimeError, match="protocol 2.0.0"),
-        ):
-            await server._handle_response(  # pyright: ignore[reportPrivateUsage]
-                request,
-                context,
-                asyncio.Event(),
-            )
-
     async def test_previous_response_requires_existing_agent_session(self) -> None:
         agent = _make_agent()
         server = _make_server(agent, session_store=SessionStore())
@@ -517,7 +464,7 @@ class TestAgentSessionPersistence:
 
         provider = server._session_storage_provider  # pyright: ignore[reportPrivateUsage]
         assert provider is not None
-        session_store = provider.get_store(config=server.config)
+        session_store = provider.get_store(config=server.config, platform_context=get_request_context())
         assert session_store is not None
         first_session = await session_store.get(first.json()["id"])
         second_session = await session_store.get(second.json()["id"])
@@ -1483,9 +1430,8 @@ class TestOutputItemToMessage:
     async def test_mcp_approval_request(self) -> None:
         from azure.ai.agentserver.responses.models import OutputItemMcpApprovalRequest
 
-        storage = InMemoryFunctionApprovalStore()
         saved = _make_function_approval_request_content(request_id="apr-1")
-        await storage.save_approval_request("apr-1", saved)
+        storage = _function_approval_store(saved)
 
         item = OutputItemMcpApprovalRequest({
             "type": "mcp_approval_request",
@@ -1501,9 +1447,8 @@ class TestOutputItemToMessage:
     async def test_mcp_approval_response(self) -> None:
         from azure.ai.agentserver.responses.models import OutputItemMcpApprovalResponseResource
 
-        storage = InMemoryFunctionApprovalStore()
         saved = _make_function_approval_request_content(request_id="apr-1")
-        await storage.save_approval_request("apr-1", saved)
+        storage = _function_approval_store(saved)
 
         item = OutputItemMcpApprovalResponseResource({
             "type": "mcp_approval_response",
@@ -1993,9 +1938,8 @@ class TestItemToMessage:
     async def test_mcp_approval_request(self) -> None:
         from azure.ai.agentserver.responses.models import ItemMcpApprovalRequest
 
-        storage = InMemoryFunctionApprovalStore()
         saved = _make_function_approval_request_content(request_id="apr-1")
-        await storage.save_approval_request("apr-1", saved)
+        storage = _function_approval_store(saved)
 
         item = ItemMcpApprovalRequest({
             "type": "mcp_approval_request",
@@ -2012,9 +1956,8 @@ class TestItemToMessage:
     async def test_mcp_approval_response(self) -> None:
         from azure.ai.agentserver.responses.models import MCPApprovalResponse
 
-        storage = InMemoryFunctionApprovalStore()
         saved = _make_function_approval_request_content(request_id="apr-1")
-        await storage.save_approval_request("apr-1", saved)
+        storage = _function_approval_store(saved)
 
         item = MCPApprovalResponse({
             "type": "mcp_approval_response",
@@ -3244,39 +3187,14 @@ class TestMultiTurnMixedContent:
 # region Function approval round-trip
 
 
-class TestFunctionApprovalStore:
-    """Unit tests for the function approval storage classes."""
-
-    async def test_in_memory_save_and_load(self) -> None:
-        storage = InMemoryFunctionApprovalStore()
-        request = _make_function_approval_request_content(request_id="apr_1")
-        await storage.save_approval_request("apr_1", request)
-        loaded = await storage.load_approval_request("apr_1")
-        assert loaded.type == "function_approval_request"
-        assert loaded.id == "apr_1"
-
-    async def test_in_memory_duplicate_save_raises(self) -> None:
-        storage = InMemoryFunctionApprovalStore()
-        request = _make_function_approval_request_content(request_id="apr_1")
-        await storage.save_approval_request("apr_1", request)
-        with pytest.raises(ValueError, match="already exists"):
-            await storage.save_approval_request("apr_1", request)
-
-    async def test_in_memory_missing_load_raises(self) -> None:
-        storage = InMemoryFunctionApprovalStore()
-        with pytest.raises(KeyError):
-            await storage.load_approval_request("missing")
-
-
 class TestFunctionApprovalConversion:
     """Tests for the approval-aware paths in `_item_to_message` / `_output_item_to_message`."""
 
     async def test_output_item_mcp_approval_request_loads_from_storage(self) -> None:
         from azure.ai.agentserver.responses.models import OutputItemMcpApprovalRequest
 
-        storage = InMemoryFunctionApprovalStore()
         saved = _make_function_approval_request_content(request_id="apr-1")
-        await storage.save_approval_request("apr-1", saved)
+        storage = _function_approval_store(saved)
 
         item = OutputItemMcpApprovalRequest({
             "type": "mcp_approval_request",
@@ -3311,9 +3229,8 @@ class TestFunctionApprovalConversion:
     async def test_output_item_mcp_approval_response_resolves_to_approval_response(self) -> None:
         from azure.ai.agentserver.responses.models import OutputItemMcpApprovalResponseResource
 
-        storage = InMemoryFunctionApprovalStore()
         saved = _make_function_approval_request_content(request_id="apr-1")
-        await storage.save_approval_request("apr-1", saved)
+        storage = _function_approval_store(saved)
 
         item = OutputItemMcpApprovalResponseResource({
             "type": "mcp_approval_response",
@@ -3346,9 +3263,8 @@ class TestFunctionApprovalConversion:
     async def test_input_item_mcp_approval_request_loads_from_storage(self) -> None:
         from azure.ai.agentserver.responses.models import ItemMcpApprovalRequest
 
-        storage = InMemoryFunctionApprovalStore()
         saved = _make_function_approval_request_content(request_id="apr-1")
-        await storage.save_approval_request("apr-1", saved)
+        storage = _function_approval_store(saved)
 
         item = ItemMcpApprovalRequest({
             "type": "mcp_approval_request",
@@ -3365,9 +3281,8 @@ class TestFunctionApprovalConversion:
     async def test_input_item_mcp_approval_response_resolves_to_approval_response(self) -> None:
         from azure.ai.agentserver.responses.models import MCPApprovalResponse
 
-        storage = InMemoryFunctionApprovalStore()
         saved = _make_function_approval_request_content(request_id="apr-1")
-        await storage.save_approval_request("apr-1", saved)
+        storage = _function_approval_store(saved)
 
         item = MCPApprovalResponse({
             "type": "mcp_approval_response",
@@ -3410,7 +3325,7 @@ class TestFunctionApprovalRoundTrip:
 
         # Storage must contain a saved entry under the emitted request id.
         loaded = await server._function_approval_storage_provider.get_store(  # pyright: ignore[reportPrivateUsage]
-            config=server.config
+            config=server.config, platform_context=get_request_context()
         ).load_approval_request(approval_request_id)
         assert loaded.type == "function_approval_request"
         assert loaded.function_call is not None
@@ -3440,7 +3355,7 @@ class TestFunctionApprovalRoundTrip:
         assert approval_request_id is not None
 
         loaded = await server._function_approval_storage_provider.get_store(  # pyright: ignore[reportPrivateUsage]
-            config=server.config
+            config=server.config, platform_context=get_request_context()
         ).load_approval_request(approval_request_id)
         assert loaded.type == "function_approval_request"
 
@@ -4283,7 +4198,7 @@ class TestWorkflowAgentHosting:
         # ``function_call``) must be persisted under that id so the next
         # turn can reconstruct it.
         loaded = await server._function_approval_storage_provider.get_store(  # pyright: ignore[reportPrivateUsage]
-            config=server.config
+            config=server.config, platform_context=get_request_context()
         ).load_approval_request(approval_request_id)
         assert loaded.type == "function_approval_request"
         assert loaded.function_call is not None
@@ -4313,7 +4228,7 @@ class TestWorkflowAgentHosting:
         assert approval_request_id is not None
 
         loaded = await server._function_approval_storage_provider.get_store(  # pyright: ignore[reportPrivateUsage]
-            config=server.config
+            config=server.config, platform_context=get_request_context()
         ).load_approval_request(approval_request_id)
         assert loaded.type == "function_approval_request"
         assert mock_agent.run_count == 1
@@ -4333,32 +4248,39 @@ class TestWorkflowAgentHosting:
             final_text="done with approval",
         )
         server = _make_server(workflow_agent)
+        checkpoint_provider = server._checkpoint_storage_provider  # pyright: ignore[reportPrivateUsage]
 
-        first = await _post(server, stream=False)
-        assert first.status_code == 200
-        first_body = first.json()
-        first_response_id = first_body["id"]
-        approval_items = [it for it in first_body["output"] if it["type"] == "mcp_approval_request"]
-        assert len(approval_items) == 1
-        approval_request_id = approval_items[0]["id"]
-        assert mock_agent.run_count == 1
+        with patch.object(checkpoint_provider, "get_store", wraps=checkpoint_provider.get_store) as get_store:
+            first = await _post(server, stream=False)
+            assert first.status_code == 200
+            first_body = first.json()
+            first_response_id = first_body["id"]
+            approval_items = [it for it in first_body["output"] if it["type"] == "mcp_approval_request"]
+            assert len(approval_items) == 1
+            approval_request_id = approval_items[0]["id"]
+            assert mock_agent.run_count == 1
 
-        second_payload: dict[str, Any] = {
-            "model": "test-model",
-            "input": [
-                {
-                    "type": "mcp_approval_response",
-                    "approval_request_id": approval_request_id,
-                    "approve": True,
-                }
-            ],
-            "stream": False,
-            "previous_response_id": first_response_id,
-        }
-        second = await _post_json(server, second_payload)
+            second_payload: dict[str, Any] = {
+                "model": "test-model",
+                "input": [
+                    {
+                        "type": "mcp_approval_response",
+                        "approval_request_id": approval_request_id,
+                        "approve": True,
+                    }
+                ],
+                "stream": False,
+                "previous_response_id": first_response_id,
+            }
+            second = await _post_json(server, second_payload)
         assert second.status_code == 200
         second_body = second.json()
         assert second_body["status"] == "completed"
+        assert [call.kwargs["context_id"] for call in get_store.call_args_list] == [
+            first_response_id,
+            second_body["id"],
+            first_response_id,
+        ]
 
         # The inner agent must have been resumed (restore replay + new turn).
         # Restore call is a no-op for the mock (no input); the new-turn call

python/packages/foundry_hosting/tests/test_state_store.py

@@ -5,8 +5,8 @@ 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 agent_framework import AgentSession, Content, WorkflowCheckpoint, WorkflowCheckpointException
+from azure.ai.agentserver.core import AgentConfig, FoundryAgentRequestContext
 from azure.ai.agentserver.core.storage import FoundryStorageConflictError
 
 from agent_framework_foundry_hosting import ContextScopedStoreProvider, StoreProvider
@@ -17,7 +17,6 @@ from agent_framework_foundry_hosting._state_store import (
     FoundryCheckpointStore,
     FoundryFunctionApprovalStore,
     FunctionApprovalStoreProvider,
-    InMemoryFunctionApprovalStore,
 )
 
 
@@ -60,6 +59,10 @@ def _config(*, is_hosted: bool) -> AgentConfig:
     )
 
 
+def _platform_context(call_id: str = "call-1", user_id: str = "user-1") -> FoundryAgentRequestContext:
+    return FoundryAgentRequestContext(call_id=call_id, user_id=user_id)
+
+
 def test_storage_providers_use_public_abstraction() -> None:
     assert issubclass(CheckpointStoreProvider, ContextScopedStoreProvider)
     assert not issubclass(CheckpointStoreProvider, StoreProvider)
@@ -75,11 +78,11 @@ async def test_save_uses_context_scoped_store() -> None:
         "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)
+        result = await FoundryCheckpointStore("context-1", _platform_context()).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())
+    get_or_create.assert_awaited_once_with("checkpoints/context-1", user_isolation=True, user_id="user-1")
+    store.set_item.assert_awaited_once_with("checkpoint-1", checkpoint.to_dict(), call_id="call-1")
 
 
 async def test_load_returns_checkpoint() -> None:
@@ -91,9 +94,10 @@ async def test_load_returns_checkpoint() -> None:
         "agent_framework_foundry_hosting._state_store.FoundryStateStore.get_or_create",
         new=AsyncMock(return_value=store),
     ):
-        result = await FoundryCheckpointStore("context-1").load("checkpoint-1")
+        result = await FoundryCheckpointStore("context-1", _platform_context()).load("checkpoint-1")
 
     assert result == checkpoint
+    store.get_item.assert_awaited_once_with("checkpoint-1", call_id="call-1")
 
 
 async def test_load_raises_for_missing_checkpoint() -> None:
@@ -107,7 +111,7 @@ async def test_load_raises_for_missing_checkpoint() -> None:
         ),
         pytest.raises(WorkflowCheckpointException, match="No checkpoint found with ID missing"),
     ):
-        await FoundryCheckpointStore("context-1").load("missing")
+        await FoundryCheckpointStore("context-1", _platform_context()).load("missing")
 
 
 async def test_list_checkpoints_paginates_and_filters_by_workflow() -> None:
@@ -134,11 +138,14 @@ async def test_list_checkpoints_paginates_and_filters_by_workflow() -> None:
         "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")
+        result = await FoundryCheckpointStore("context-1", _platform_context()).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"}
+    assert store.list_keys.await_args_list[0].kwargs == {"after": None, "call_id": "call-1"}
+    assert store.list_keys.await_args_list[1].kwargs == {"after": "cursor-1", "call_id": "call-1"}
+    assert all(call.kwargs == {"call_id": "call-1"} for call in store.get_item.await_args_list)
 
 
 @pytest.mark.parametrize(("deleted_id", "expected"), [("item-id", True), (None, False)])
@@ -150,13 +157,14 @@ async def test_delete_reports_whether_checkpoint_existed(deleted_id: str | None,
         "agent_framework_foundry_hosting._state_store.FoundryStateStore.get_or_create",
         new=AsyncMock(return_value=store),
     ):
-        result = await FoundryCheckpointStore("context-1").delete("checkpoint-1")
+        result = await FoundryCheckpointStore("context-1", _platform_context()).delete("checkpoint-1")
 
     assert result is expected
+    store.delete_item.assert_awaited_once_with("checkpoint-1", call_id="call-1")
 
 
 async def test_get_latest_uses_timestamp_and_list_ids_filters() -> None:
-    storage = FoundryCheckpointStore("context-1")
+    storage = FoundryCheckpointStore("context-1", _platform_context())
     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
@@ -166,33 +174,36 @@ async def test_get_latest_uses_timestamp_and_list_ids_filters() -> None:
 
 
 async def test_get_latest_returns_none_when_no_checkpoints_exist() -> None:
-    storage = FoundryCheckpointStore("context-1")
+    storage = FoundryCheckpointStore("context-1", _platform_context())
     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:
+def test_checkpoint_storage_provider_creates_request_scoped_storage(is_hosted: bool) -> None:
     provider = CheckpointStoreProvider()
     config = _config(is_hosted=is_hosted)
+    first_context = _platform_context("call-1")
+    second_context = _platform_context("call-2")
 
-    first = provider.get_store(config=config, context_id="context-1")
-    second = provider.get_store(config=config, context_id="context-2")
+    first = provider.get_store(config=config, context_id="context-1", platform_context=first_context)
+    second = provider.get_store(config=config, context_id="context-1", platform_context=second_context)
 
-    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 type(first) is FoundryCheckpointStore
+    assert type(second) is FoundryCheckpointStore
     assert second is not first
+    assert first.platform_context is first_context
+    assert second.platform_context is second_context
 
 
 @pytest.mark.parametrize(
     "create_store",
     [
-        lambda: FoundryCheckpointStore(""),
-        lambda: CheckpointStoreProvider().get_store(config=_config(is_hosted=True), context_id=""),
+        lambda: FoundryCheckpointStore("", _platform_context()),
+        lambda: CheckpointStoreProvider().get_store(
+            config=_config(is_hosted=True), context_id="", platform_context=_platform_context()
+        ),
     ],
 )
 def test_checkpoint_stores_require_context_id(create_store: Callable[[], Any]) -> None:
@@ -219,13 +230,14 @@ async def test_save_and_load_function_approval_request() -> None:
         "agent_framework_foundry_hosting._state_store.FoundryStateStore.get_or_create",
         new=AsyncMock(return_value=store),
     ) as get_or_create:
-        storage = FoundryFunctionApprovalStore()
+        storage = FoundryFunctionApprovalStore(_platform_context())
         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())
+    get_or_create.assert_awaited_with("function_approvals", user_isolation=True, user_id="user-1")
+    store.create_item.assert_awaited_once_with("approval-1", request.to_dict(), call_id="call-1")
+    store.get_item.assert_awaited_once_with("approval-1", call_id="call-1")
     assert loaded == request
     function_call = loaded.function_call
     assert function_call is not None
@@ -244,7 +256,9 @@ async def test_save_duplicate_function_approval_request_raises() -> None:
         ),
         pytest.raises(ValueError, match="Approval request with ID 'approval-1' already exists"),
     ):
-        await FoundryFunctionApprovalStore().save_approval_request("approval-1", _approval_request("approval-1"))
+        await FoundryFunctionApprovalStore(_platform_context()).save_approval_request(
+            "approval-1", _approval_request("approval-1")
+        )
 
 
 async def test_load_missing_function_approval_request_raises() -> None:
@@ -258,39 +272,32 @@ async def test_load_missing_function_approval_request_raises() -> None:
         ),
         pytest.raises(KeyError, match="Approval request with ID 'missing' does not exist"),
     ):
-        await FoundryFunctionApprovalStore().load_approval_request("missing")
+        await FoundryFunctionApprovalStore(_platform_context()).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:
+@pytest.mark.parametrize("is_hosted", [True, False])
+def test_function_approval_storage_provider_uses_foundry_store(is_hosted: bool) -> None:
     provider = FunctionApprovalStoreProvider()
     config = _config(is_hosted=is_hosted)
+    platform_context = _platform_context()
 
-    storage = provider.get_store(config=config)
+    storage = provider.get_store(config=config, platform_context=platform_context)
 
-    assert type(storage) is expected_type
-    assert provider.get_store(config=config) is storage
+    assert type(storage) is FoundryFunctionApprovalStore
+    assert storage.platform_context is platform_context
 
 
-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,
-    ):
+def test_function_approval_storage_provider_creates_request_scoped_storage() -> None:
+    with patch("agent_framework_foundry_hosting._state_store.FoundryFunctionApprovalStore") as storage_type:
         provider = FunctionApprovalStoreProvider()
-        config = _config(is_hosted=True)
-        storage = provider.get_store(config=config)
+        config = _config(is_hosted=False)
+        first_context = _platform_context("call-1")
+        second_context = _platform_context("call-2")
+        provider.get_store(config=config, platform_context=first_context)
+        provider.get_store(config=config, platform_context=second_context)
 
-        assert provider.get_store(config=config) is storage
-
-    foundry_storage_type.assert_called_once_with()
-    in_memory_storage_type.assert_not_called()
+    assert storage_type.call_args_list[0].args == (first_context,)
+    assert storage_type.call_args_list[1].args == (second_context,)
 
 
 async def test_set_agent_session_uses_scoped_store() -> None:
@@ -302,10 +309,10 @@ async def test_set_agent_session_uses_scoped_store() -> None:
         "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)
+        await FoundryAgentSessionStore(_platform_context()).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())
+    get_or_create.assert_awaited_once_with("agent_sessions", user_isolation=True, user_id="user-1")
+    store.set_item.assert_awaited_once_with("storage-session-1", session.to_dict(), call_id="call-1")
 
 
 async def test_get_agent_session_returns_deserialized_session() -> None:
@@ -318,13 +325,13 @@ async def test_get_agent_session_returns_deserialized_session() -> None:
         "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")
+        result = await FoundryAgentSessionStore(_platform_context()).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")
+    get_or_create.assert_awaited_once_with("agent_sessions", user_isolation=True, user_id="user-1")
+    store.get_item.assert_awaited_once_with("storage-session-1", call_id="call-1")
 
 
 async def test_get_missing_agent_session_returns_none() -> None:
@@ -335,7 +342,7 @@ async def test_get_missing_agent_session_returns_none() -> None:
         "agent_framework_foundry_hosting._state_store.FoundryStateStore.get_or_create",
         new=AsyncMock(return_value=store),
     ):
-        result = await FoundryAgentSessionStore().get("missing")
+        result = await FoundryAgentSessionStore(_platform_context()).get("missing")
 
     assert result is None
 
@@ -347,38 +354,31 @@ async def test_delete_agent_session_is_idempotent() -> None:
         "agent_framework_foundry_hosting._state_store.FoundryStateStore.get_or_create",
         new=AsyncMock(return_value=store),
     ):
-        await FoundryAgentSessionStore().delete("storage-session-1")
+        await FoundryAgentSessionStore(_platform_context()).delete("storage-session-1")
 
-    store.delete_item.assert_awaited_once_with("storage-session-1")
+    store.delete_item.assert_awaited_once_with("storage-session-1", call_id="call-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:
+@pytest.mark.parametrize("is_hosted", [True, False])
+def test_agent_session_storage_provider_uses_foundry_store(is_hosted: bool) -> None:
     provider = AgentSessionStoreProvider()
     config = _config(is_hosted=is_hosted)
+    platform_context = _platform_context()
 
-    storage = provider.get_store(config=config)
+    storage = provider.get_store(config=config, platform_context=platform_context)
 
-    assert type(storage) is expected_type
-    assert provider.get_store(config=config) is storage
+    assert type(storage) is FoundryAgentSessionStore
+    assert storage.platform_context is platform_context
 
 
-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,
-    ):
+def test_agent_session_storage_provider_creates_request_scoped_storage() -> None:
+    with patch("agent_framework_foundry_hosting._state_store.FoundryAgentSessionStore") as storage_type:
         provider = AgentSessionStoreProvider()
         config = _config(is_hosted=False)
-        storage = provider.get_store(config=config)
+        first_context = _platform_context("call-1")
+        second_context = _platform_context("call-2")
+        provider.get_store(config=config, platform_context=first_context)
+        provider.get_store(config=config, platform_context=second_context)
 
-        assert provider.get_store(config=config) is storage
-
-    foundry_storage_type.assert_not_called()
-    in_memory_storage_type.assert_called_once_with()
+    assert storage_type.call_args_list[0].args == (first_context,)
+    assert storage_type.call_args_list[1].args == (second_context,)