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..f1d31976904 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 from abc import ABC, abstractmethod from datetime import datetime -from typing import Generic, Protocol, TypeVar +from typing import ClassVar, Generic, Protocol, TypeVar +from weakref import WeakKeyDictionary from agent_framework import ( AgentSession, @@ -300,32 +302,65 @@ class FoundryAgentSessionStore(SessionStore): DEFAULT_ROOT_SCOPE = "agent_sessions" + # Cache the backing ``FoundryStateStore`` per (event loop, scope). The store + # owns an async pipeline + credential bound to the loop it was created on, so + # it must never be reused from a different loop (e.g. across ``asyncio.run()`` + # calls, or between loop-scoped tests) -- keying by the running loop prevents + # that and lets a closed loop's store be garbage-collected along with it. + # Keying by scope keeps a subclass that overrides ``DEFAULT_ROOT_SCOPE`` + # isolated to its own collection rather than sharing (or clobbering) the base + # store. + # + # Within the long-running server (a single loop) every request shares one + # store, so the per-request credential rebuild + ``agent_sessions`` metadata + # round-trip that ``get_or_create`` would otherwise repeat is paid just once. + _store_cache: ClassVar[ + "WeakKeyDictionary[asyncio.AbstractEventLoop, dict[str, FoundryStateStore]]" + ] = WeakKeyDictionary() + _cache_locks: ClassVar["WeakKeyDictionary[asyncio.AbstractEventLoop, asyncio.Lock]"] = WeakKeyDictionary() + def __init__(self, platform_context: FoundryAgentRequestContext) -> None: self.platform_context = platform_context + @classmethod + def _loop_lock(cls, loop: "asyncio.AbstractEventLoop") -> asyncio.Lock: + lock = cls._cache_locks.get(loop) + if lock is None: + # ``setdefault`` collapses a concurrent first-use to a single lock. + lock = cls._cache_locks.setdefault(loop, asyncio.Lock()) + return lock + 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 + # Fast path: already resolved on this loop -> no lock, no round-trip. + by_scope = FoundryAgentSessionStore._store_cache.get(loop) + if by_scope is not None and scope in by_scope: + return by_scope[scope] + async with self._loop_lock(loop): + by_scope = FoundryAgentSessionStore._store_cache.setdefault(loop, {}) + if scope not in by_scope: + by_scope[scope] = await FoundryStateStore.get_or_create(scope, user_isolation=True) + return by_scope[scope] 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 the life of its event loop and is reclaimed with it. 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) class AgentSessionStoreProvider(StoreProvider[SessionStore]): diff --git a/python/packages/foundry_hosting/tests/test_state_store.py b/python/packages/foundry_hosting/tests/test_state_store.py index fc36d00373c..5a18c41f016 100644 --- a/python/packages/foundry_hosting/tests/test_state_store.py +++ b/python/packages/foundry_hosting/tests/test_state_store.py @@ -1,5 +1,6 @@ # Copyright (c) Microsoft. All rights reserved. -from collections.abc import Callable +import asyncio +from collections.abc import Callable, Iterator from dataclasses import dataclass from types import SimpleNamespace from typing import Any @@ -72,6 +73,21 @@ 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() -> Iterator[None]: + """Isolate the FoundryAgentSessionStore per-(loop, scope) state-store cache. + + The backing store is cached per event loop and scope; clear the caches + before and after every test so a test's mocked ``get_or_create`` never leaks + into another test that happens to share an event loop. + """ + FoundryAgentSessionStore._store_cache.clear() + FoundryAgentSessionStore._cache_locks.clear() + yield + FoundryAgentSessionStore._store_cache.clear() + FoundryAgentSessionStore._cache_locks.clear() + + def test_storage_providers_use_public_abstraction() -> None: assert issubclass(CheckpointStoreProvider, ContextScopedStoreProvider) assert not issubclass(CheckpointStoreProvider, StoreProvider) @@ -508,3 +524,65 @@ 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()) + + with patch( + "agent_framework_foundry_hosting._state_store.FoundryStateStore.get_or_create", + new=AsyncMock(return_value=store), + ) 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()