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

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]):

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()