Python: perf(foundry-hosting): cache FoundryStateStore in FoundryAgentSessionStore

python/packages/foundry_hosting/agent_framework_foundry_hosting/_state_store.py

@@ -1,9 +1,10 @@
 # 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 agent_framework import (
     AgentSession,
@@ -295,37 +296,91 @@ 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 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]):

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,74 @@ 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()