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 682103a24d6..871d158f4bd 100644 --- a/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py +++ b/python/packages/foundry_hosting/agent_framework_foundry_hosting/_responses.py @@ -550,6 +550,13 @@ def __init__( self._agent_init_lock = asyncio.Lock() self.shutdown_handler(self._cleanup_agent) + # Give the cached backing session stores an explicit async owner: closing + # the server closes each cached store's pooled pipeline + owned credential + # rather than leaving that to loop collection. Registered only when the + # provider exposes ``aclose`` so custom providers stay opt-in. + session_provider_aclose = getattr(self._session_storage_provider, "aclose", None) + if callable(session_provider_aclose): + self.shutdown_handler(session_provider_aclose) self.response_handler(self._handle_response) mark_feature_used(FeatureIndex.FOUNDRY_HOSTING) 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..f03418d87d8 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,9 +1,11 @@ # Copyright (c) Microsoft. All rights reserved. +import asyncio +import logging from abc import ABC, abstractmethod from datetime import datetime -from typing import Generic, Protocol, TypeVar +from typing import ClassVar, Generic, Protocol, TypeVar from agent_framework import ( AgentSession, @@ -17,6 +19,8 @@ from azure.ai.agentserver.core import AgentConfig, FoundryAgentRequestContext from azure.ai.agentserver.core.storage import FoundryStateStore, FoundryStorageConflictError +logger = logging.getLogger(__name__) + StoreT = TypeVar("StoreT") @@ -295,37 +299,120 @@ def get_store(self, *, config: AgentConfig, platform_context: FoundryAgentReques # region Agent session persistence +class _SessionStoreLoopCache: + """Per-event-loop cache of backing ``FoundryStateStore`` instances, by scope. + + Anchored on the event loop (see ``FoundryAgentSessionStore._loop_cache``) so + that it -- along with the stores' pooled pipelines and credentials -- is + reclaimed together with the loop, rather than surviving in a process-global + map after the loop closes. + """ + + __slots__ = ("lock", "stores") + + def __init__(self) -> None: + self.stores: dict[str, FoundryStateStore] = {} + self.lock = asyncio.Lock() + + class FoundryAgentSessionStore(SessionStore): """Agent session store backed by the `FoundryStateStore`.""" DEFAULT_ROOT_SCOPE = "agent_sessions" + # Name of the attribute under which each event loop carries its own + # ``_SessionStoreLoopCache`` (see ``_loop_cache``). + _LOOP_CACHE_ATTR: ClassVar[str] = "_agent_framework_foundry_session_store_cache" + def __init__(self, platform_context: FoundryAgentRequestContext) -> None: self.platform_context = platform_context + @classmethod + def _loop_cache(cls, loop: "asyncio.AbstractEventLoop") -> "_SessionStoreLoopCache | None": + # Store the cache ON the loop rather than in a process-global map keyed + # by the loop. A backing ``FoundryStateStore`` (and the ``asyncio.Lock`` + # guarding its creation) strongly references the loop it was bound to, so + # holding either in a module-level ``WeakKeyDictionary`` would keep that + # "weak" key -- and the store's open pipeline + credential -- alive + # forever, leaking one entry per closed loop (e.g. every ``asyncio.run``). + # Anchored to the loop, the cache is collected together with the loop. + cache: _SessionStoreLoopCache | None = getattr(loop, cls._LOOP_CACHE_ATTR, None) + if cache is not None: + return cache + cache = _SessionStoreLoopCache() + try: + setattr(loop, cls._LOOP_CACHE_ATTR, cache) + except (AttributeError, TypeError): + # A C-level loop that forbids attribute assignment: skip caching + # rather than leak. Correctness is unaffected, only the reuse. + return None + return cache + async def _get_store(self) -> FoundryStateStore: - return await FoundryStateStore.get_or_create( - f"{self.DEFAULT_ROOT_SCOPE}", - user_isolation=True, - ) + loop = asyncio.get_running_loop() + scope = self.DEFAULT_ROOT_SCOPE + cache = self._loop_cache(loop) + if cache is None: + # Loop cannot hold the cache; resolve without reuse (and no leak). + return await FoundryStateStore.get_or_create(scope, user_isolation=True) + # Fast path: already resolved on this loop -> no lock, no round-trip. + store = cache.stores.get(scope) + if store is not None: + return store + async with cache.lock: + store = cache.stores.get(scope) + if store is None: + store = await FoundryStateStore.get_or_create(scope, user_isolation=True) + cache.stores[scope] = store + return store async def get(self, session_id: str) -> AgentSession | None: + # The shared store is intentionally NOT entered as an ``async with`` + # context manager: its ``__aexit__`` calls ``aclose()``, which would + # close the pooled pipeline + owned credential and defeat the cache. It + # stays open for reuse and is closed explicitly at server shutdown via + # ``aclose_loop_cache`` (or reclaimed with its loop if never registered). store = await self._get_store() - async with store: - item = await store.get_item(session_id, call_id=self.platform_context.call_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) 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(), call_id=self.platform_context.call_id) + 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, call_id=self.platform_context.call_id) + await store.delete_item(session_id, call_id=self.platform_context.call_id) + + @classmethod + async def aclose_loop_cache(cls) -> None: + """Close every cached backing store on the current event loop. + + The per-loop cache keeps each ``FoundryStateStore`` open for reuse and + otherwise relies on the loop being collected to reclaim it. Registered as + a host-server shutdown handler (see ``AgentSessionStoreProvider.aclose``) + so the owning server explicitly awaits closure of every store's pooled + pipeline and owned credential at shutdown instead of leaving it to GC. + Idempotent: the cache is drained under its lock, so a second call is a + no-op. + """ + try: + loop = asyncio.get_running_loop() + except RuntimeError: + return + cache: _SessionStoreLoopCache | None = getattr(loop, cls._LOOP_CACHE_ATTR, None) + if cache is None: + return + async with cache.lock: + stores = list(cache.stores.values()) + cache.stores.clear() + for store in stores: + try: + await store.aclose() + except Exception: # pragma: no cover - best-effort cleanup at shutdown + logger.debug("Failed to close a cached session state store during shutdown.", exc_info=True) class AgentSessionStoreProvider(StoreProvider[SessionStore]): @@ -338,5 +425,15 @@ def get_store(self, *, config: AgentConfig, platform_context: FoundryAgentReques """Get agent session store for the requested hosting environment.""" return FoundryAgentSessionStore(platform_context) + async def aclose(self) -> None: + """Release the backing stores cached by ``FoundryAgentSessionStore``. + + Registered by the host server as a shutdown handler so server shutdown + closes each cached store's pooled pipeline and owned credential. The + per-loop cache otherwise stays open for reuse and relies on loop + collection; this gives those resources an explicit async owner. + """ + await FoundryAgentSessionStore.aclose_loop_cache() + # endregion Agent session persistence diff --git a/python/packages/foundry_hosting/tests/test_state_store.py b/python/packages/foundry_hosting/tests/test_state_store.py index fc36d00373c..b36ad0d4a3e 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 from collections.abc import Callable from dataclasses import dataclass from types import SimpleNamespace @@ -72,6 +73,18 @@ def _platform_context(call_id: str = "call-1", user_id: str = "user-1") -> Found return FoundryAgentRequestContext(call_id=call_id, user_id=user_id) +@pytest.fixture(autouse=True) +def _reset_agent_session_store_cache() -> None: + """Guard the FoundryAgentSessionStore per-(loop, scope) state-store cache. + + The backing store is cached on the running event loop itself, so a store (and + its mocked ``get_or_create``) cannot outlive the loop that created it. Every + test runs on its own function-scoped event loop, so the cache is inherently + isolated per test -- no explicit teardown is required. + """ + return None + + def test_storage_providers_use_public_abstraction() -> None: assert issubclass(CheckpointStoreProvider, ContextScopedStoreProvider) assert not issubclass(CheckpointStoreProvider, StoreProvider) @@ -508,3 +521,98 @@ def test_agent_session_storage_provider_creates_request_scoped_storage() -> None assert storage_type.call_args_list[0].args == (first_context,) assert storage_type.call_args_list[1].args == (second_context,) + + +async def test_agent_session_store_is_cached_across_operations() -> None: + store = _store() + store.get_item = AsyncMock(return_value=None) + session_store = 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 session_store.set("s1", AgentSession(session_id="agent-session-1")) + await session_store.get("s1") + await session_store.delete("s1") + + # The backing state store is resolved once via get_or_create and reused for + # every subsequent operation instead of being rebuilt per call. + get_or_create.assert_awaited_once_with("agent_sessions", user_isolation=True) + store.set_item.assert_awaited_once() + store.get_item.assert_awaited_once() + store.delete_item.assert_awaited_once() + + +async def test_agent_session_store_concurrent_init_resolves_once() -> None: + store = _store() + store.get_item = AsyncMock(return_value=None) + session_store = FoundryAgentSessionStore(_platform_context()) + + async def _slow_get_or_create(scope: str, *, user_isolation: bool) -> MagicMock: + # Suspend before returning so every gathered task reaches ``_get_store`` + # and blocks on the creation lock while the first resolve is in flight. + # A non-suspending mock would let the first task populate the cache + # synchronously, so even a lock-less implementation would pass -- this + # forces genuine contention that exercises the lock + second cache check. + await asyncio.sleep(0) + return store + + with patch( + "agent_framework_foundry_hosting._state_store.FoundryStateStore.get_or_create", + new=AsyncMock(side_effect=_slow_get_or_create), + ) as get_or_create: + await asyncio.gather(*(session_store.get(f"s{i}") for i in range(25))) + + # Concurrent first-use must resolve the backing store exactly once. + get_or_create.assert_awaited_once_with("agent_sessions", user_isolation=True) + assert store.get_item.await_count == 25 + + +async def test_agent_session_store_subclass_scope_is_isolated() -> None: + class OtherScopeSessionStore(FoundryAgentSessionStore): + DEFAULT_ROOT_SCOPE = "other_sessions" + + base_store = _store() + base_store.get_item = AsyncMock(return_value=None) + other_store = _store() + other_store.get_item = AsyncMock(return_value=None) + + async def _fake_get_or_create(scope: str, *, user_isolation: bool) -> MagicMock: + return other_store if scope == "other_sessions" else base_store + + with patch( + "agent_framework_foundry_hosting._state_store.FoundryStateStore.get_or_create", + new=AsyncMock(side_effect=_fake_get_or_create), + ) as get_or_create: + await FoundryAgentSessionStore(_platform_context()).get("s1") + await OtherScopeSessionStore(_platform_context()).get("s1") + + # Each scope resolves and caches its own backing store -- no cross-routing. + assert get_or_create.await_count == 2 + base_store.get_item.assert_awaited_once() + other_store.get_item.assert_awaited_once() + + +async def test_provider_aclose_closes_and_drops_cached_stores() -> None: + store = _store() + store.get_item = AsyncMock(return_value=None) + store.aclose = AsyncMock() + provider = AgentSessionStoreProvider() + session_store = provider.get_store(config=_config(is_hosted=True), platform_context=_platform_context()) + + with patch( + "agent_framework_foundry_hosting._state_store.FoundryStateStore.get_or_create", + new=AsyncMock(return_value=store), + ) as get_or_create: + await session_store.get("s1") + # Shutdown must close the cached store's pooled pipeline + owned credential. + await provider.aclose() + store.aclose.assert_awaited_once() + # A second shutdown is a no-op (cache already drained). + await provider.aclose() + store.aclose.assert_awaited_once() + # The drained cache resolves a fresh backing store on the next request. + await session_store.get("s2") + + assert get_or_create.await_count == 2