Python: reuse Foundry session-store client within each response

python/packages/foundry_hosting/README.md

@@ -107,6 +107,12 @@ No additional partitioning configuration is required when using the default stor
 uses the `FoundryAgentSessionStore`, backed by Foundry storage when hosted and file-based
 storage locally. Stored sessions are scoped under `agent_sessions`.
 
+Within one Responses request, the default session store lazily opens and reuses
+one Foundry client for session load and save, then closes it before the terminal
+response event. This avoids a second store-metadata lookup without caching session
+contents or sharing clients across requests. Other session-store implementations
+retain their existing lifecycle.
+
 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.
 

python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py

@@ -90,6 +90,7 @@
     AgentSessionStoreProvider,
     CheckpointStoreProvider,
     ContextScopedStoreProvider,
+    FoundryAgentSessionStore,
     FunctionApprovalStore,
     FunctionApprovalStoreProvider,
     StoreProvider,
@@ -752,6 +753,10 @@ async def _handle_prepared_response(
                     session_storage = self._session_storage_provider.get_store(
                         config=self.config, platform_context=request_context
                     )
+                    if isinstance(session_storage, FoundryAgentSessionStore):
+                        await resources.enter_async_context(
+                            session_storage._request_scope()  # pyright: ignore[reportPrivateUsage]
+                        )
                     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
@@ -808,6 +813,7 @@ async def _handle_prepared_response(
                     cancellation_signal,
                     agent,
                     configuration,
+                    resources,
                 )
 
             try:
@@ -911,6 +917,7 @@ async def _handle_inner_agent(
         cancellation_signal: asyncio.Event,
         agent: SupportsAgentRun,
         configuration: _AgentConfiguration,
+        resources: AsyncExitStack,
     ) -> AsyncGenerator[ResponseStreamEvent]:
         """Handle a regular (non-workflow) agent.
 
@@ -934,6 +941,10 @@ async def _handle_inner_agent(
             session_storage = self._session_storage_provider.get_store(
                 config=self.config, platform_context=request_context
             )
+            if isinstance(session_storage, FoundryAgentSessionStore):
+                await resources.enter_async_context(
+                    session_storage._request_scope()  # pyright: ignore[reportPrivateUsage]
+                )
 
             # Load the caller's input items and prior conversation history concurrently with the
             # session load below. These are independent storage round-trips with no data dependency

python/packages/foundry_hosting/agent_framework_foundry_hosting/_state_store.py

@@ -1,7 +1,10 @@
 # Copyright (c) Microsoft. All rights reserved.
 
 
+import asyncio
 from abc import ABC, abstractmethod
+from collections.abc import AsyncGenerator
+from contextlib import AsyncExitStack, asynccontextmanager
 from datetime import datetime
 from typing import Generic, Protocol, TypeVar
 
@@ -302,29 +305,57 @@ class FoundryAgentSessionStore(SessionStore):
 
     def __init__(self, platform_context: FoundryAgentRequestContext) -> None:
         self.platform_context = platform_context
+        self._request_resources: AsyncExitStack | None = None
+        self._request_store: FoundryStateStore | None = None
+        self._store_lock = asyncio.Lock()
+
+    @asynccontextmanager
+    async def _request_scope(self) -> AsyncGenerator[None, None]:
+        """Reuse a lazily opened client until the hosting request finishes."""
+        if self._request_resources is not None:
+            raise RuntimeError("The session store already belongs to an active request scope.")
+        try:
+            async with AsyncExitStack() as resources:
+                self._request_resources = resources
+                yield
+        finally:
+            self._request_resources = None
+            self._request_store = None
 
     async def _get_store(self) -> FoundryStateStore:
         return await FoundryStateStore.get_or_create(
             f"{self.DEFAULT_ROOT_SCOPE}",
             user_isolation=True,
         )
 
+    @asynccontextmanager
+    async def _open_store(self) -> AsyncGenerator[FoundryStateStore, None]:
+        if self._request_resources is None:
+            store = await self._get_store()
+            async with store:
+                yield store
+            return
+        async with self._store_lock:
+            if self._request_store is None:
+                store = await self._get_store()
+                await self._request_resources.enter_async_context(store)
+                self._request_store = store
+            store = self._request_store
+        yield store
+
     async def get(self, session_id: str) -> AgentSession | None:
-        store = await self._get_store()
-        async with store:
+        async with self._open_store() as store:
             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)
 
     async def set(self, session_id: str, session: AgentSession) -> None:
-        store = await self._get_store()
-        async with store:
+        async with self._open_store() as store:
             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:
+        async with self._open_store() as store:
             await store.delete_item(session_id, call_id=self.platform_context.call_id)
 
 

python/packages/foundry_hosting/tests/test_responses.py

@@ -1086,6 +1086,36 @@ async def test_previous_response_requires_existing_agent_session(self) -> None:
 
 
 class TestAgentSessionPersistence:
+    @pytest.mark.parametrize("stream", [False, True])
+    async def test_default_store_reuses_one_client_per_response(self, stream: bool) -> None:
+        from azure.ai.agentserver.core.storage import FoundryStateStore
+
+        clients: list[FoundryStateStore] = []
+        get_or_create = FoundryStateStore.get_or_create
+
+        async def create(*args: Any, **kwargs: Any) -> FoundryStateStore:
+            client = await get_or_create(*args, **kwargs)
+            client.aclose = AsyncMock(wraps=client.aclose)
+            clients.append(client)
+            return client
+
+        server = _make_server(
+            _make_agent(stream_updates=[AgentResponseUpdate(contents=[Content.from_text("hello")], role="assistant")])
+        )
+        with patch(
+            "agent_framework_foundry_hosting._state_store.FoundryStateStore.get_or_create",
+            new=AsyncMock(side_effect=create),
+        ) as create_mock:
+            first = await _post(server)
+            second = await _post(server, previous_response_id=first.json()["id"], stream=stream)
+
+        second_body = _parse_sse_events(second.text)[-1]["data"]["response"] if stream else second.json()
+        assert first.json()["status"] == second_body["status"] == "completed"
+        assert create_mock.await_count == 2
+        assert clients[0] is not clients[1]
+        for client in clients:
+            cast(AsyncMock, client.aclose).assert_awaited_once()
+
     async def test_previous_response_chain_restores_session_state(self) -> None:
         seen_counts: list[int] = []
         seen_session_ids: list[str] = []

python/packages/foundry_hosting/tests/test_state_store.py

@@ -1,4 +1,5 @@
 # Copyright (c) Microsoft. All rights reserved.
+import asyncio
 import os
 from collections.abc import Callable
 from dataclasses import dataclass
@@ -494,6 +495,135 @@ async def test_delete_agent_session_is_idempotent() -> None:
     store.delete_item.assert_awaited_once_with("storage-session-1", call_id="call-1")
 
 
+async def test_request_scope_reuses_client_without_caching_session_contents() -> None:
+    store = _store()
+    session = AgentSession(session_id="session-1")
+    store.get_item.return_value = SimpleNamespace(value=session.to_dict())
+    storage = FoundryAgentSessionStore(_platform_context())
+
+    with patch(
+        "agent_framework_foundry_hosting._state_store.FoundryStateStore.get_or_create",
+        new=AsyncMock(return_value=store),
+    ) as get_or_create:
+        async with storage._request_scope():
+            get_or_create.assert_not_awaited()
+            await storage.get("session-1")
+            await storage.set("session-1", session)
+            await storage.get("session-1")
+            store.__aexit__.assert_not_awaited()
+        get_or_create.assert_awaited_once_with("agent_sessions", user_isolation=True)
+
+    assert store.get_item.await_count == 2
+    store.set_item.assert_awaited_once_with("session-1", session.to_dict(), call_id="call-1")
+    store.__aenter__.assert_awaited_once()
+    store.__aexit__.assert_awaited_once()
+
+
+async def test_request_scope_does_not_open_unused_client() -> None:
+    storage = FoundryAgentSessionStore(_platform_context())
+    with patch(
+        "agent_framework_foundry_hosting._state_store.FoundryStateStore.get_or_create",
+        new_callable=AsyncMock,
+    ) as get_or_create:
+        async with storage._request_scope():
+            pass
+    get_or_create.assert_not_awaited()
+
+
+@pytest.mark.parametrize("failure", [OSError("read failed"), asyncio.CancelledError(), GeneratorExit()])
+async def test_request_scope_closes_client_on_failure(failure: BaseException) -> None:
+    store = _store()
+    store.get_item.side_effect = failure
+    storage = FoundryAgentSessionStore(_platform_context())
+    with (
+        patch(
+            "agent_framework_foundry_hosting._state_store.FoundryStateStore.get_or_create",
+            new=AsyncMock(return_value=store),
+        ),
+        pytest.raises(type(failure)),
+    ):
+        async with storage._request_scope():
+            await storage.get("session-1")
+    store.__aexit__.assert_awaited_once()
+    assert storage._request_store is None
+    assert storage._request_resources is None
+
+
+async def test_request_scope_retries_failed_initialization() -> None:
+    store = _store()
+    store.get_item.return_value = None
+    storage = FoundryAgentSessionStore(_platform_context())
+    with patch(
+        "agent_framework_foundry_hosting._state_store.FoundryStateStore.get_or_create",
+        new=AsyncMock(side_effect=[OSError("initialization failed"), store]),
+    ) as get_or_create:
+        async with storage._request_scope():
+            with pytest.raises(OSError, match="initialization failed"):
+                await storage.get("session-1")
+            assert await storage.get("session-1") is None
+    assert get_or_create.await_count == 2
+    store.__aexit__.assert_awaited_once()
+
+
+async def test_request_scope_initializes_once_for_concurrent_reads() -> None:
+    store = _store()
+    store.get_item.return_value = None
+    started = asyncio.Event()
+    release = asyncio.Event()
+
+    async def create(*args: Any, **kwargs: Any) -> MagicMock:
+        started.set()
+        await release.wait()
+        return store
+
+    storage = FoundryAgentSessionStore(_platform_context())
+    with patch(
+        "agent_framework_foundry_hosting._state_store.FoundryStateStore.get_or_create",
+        new=AsyncMock(side_effect=create),
+    ) as get_or_create:
+        async with storage._request_scope():
+            first = asyncio.create_task(storage.get("first"))
+            await started.wait()
+            second = asyncio.create_task(storage.get("second"))
+            release.set()
+            await asyncio.gather(first, second)
+    get_or_create.assert_awaited_once()
+    assert store.get_item.await_count == 2
+    store.__aexit__.assert_awaited_once()
+
+
+async def test_request_scopes_keep_clients_and_call_ids_separate() -> None:
+    first, second = _store(), _store()
+    first.get_item.return_value = second.get_item.return_value = None
+    with patch(
+        "agent_framework_foundry_hosting._state_store.FoundryStateStore.get_or_create",
+        new=AsyncMock(side_effect=[first, second]),
+    ) as get_or_create:
+        for call_id in ("call-1", "call-2"):
+            storage = FoundryAgentSessionStore(_platform_context(call_id))
+            async with storage._request_scope():
+                await storage.get("session-1")
+    assert get_or_create.await_count == 2
+    first.get_item.assert_awaited_once_with("session-1", call_id="call-1")
+    second.get_item.assert_awaited_once_with("session-1", call_id="call-2")
+    first.__aexit__.assert_awaited_once()
+    second.__aexit__.assert_awaited_once()
+
+
+async def test_standalone_operations_still_close_each_client() -> None:
+    store = _store()
+    store.get_item.return_value = None
+    storage = FoundryAgentSessionStore(_platform_context())
+    with patch(
+        "agent_framework_foundry_hosting._state_store.FoundryStateStore.get_or_create",
+        new=AsyncMock(return_value=store),
+    ) as get_or_create:
+        await storage.get("session-1")
+        await storage.set("session-1", AgentSession())
+    assert get_or_create.await_count == 2
+    assert store.__aexit__.await_count == 2
+
+
 @pytest.mark.parametrize("is_hosted", [True, False])
 def test_agent_session_storage_provider_uses_foundry_store(is_hosted: bool) -> None:
     provider = AgentSessionStoreProvider()