diff --git a/python/packages/foundry_hosting/README.md b/python/packages/foundry_hosting/README.md index a7a59803df4..8a2dba22c03 100644 --- a/python/packages/foundry_hosting/README.md +++ b/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. diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py index 325bf264cd1..49c84b00e3e 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py +++ b/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 diff --git a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_state_store.py b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_state_store.py index a88c984ead5..050c93a8a1e 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_state_store.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_state_store.py @@ -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,6 +305,22 @@ 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( @@ -309,22 +328,34 @@ async def _get_store(self) -> FoundryStateStore: 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) diff --git a/python/packages/foundry_hosting/tests/test_responses.py b/python/packages/foundry_hosting/tests/test_responses.py index f781bfaa2f4..13da91e8d29 100644 --- a/python/packages/foundry_hosting/tests/test_responses.py +++ b/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] = [] diff --git a/python/packages/foundry_hosting/tests/test_state_store.py b/python/packages/foundry_hosting/tests/test_state_store.py index 7ae5088a4d9..85fe93c7878 100644 --- a/python/packages/foundry_hosting/tests/test_state_store.py +++ b/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()