From 97d85e81effe7afd0531ad85262303ed0ec66994 Mon Sep 17 00:00:00 2001 From: Harsheet Shah <50236780+harsheet-shah@users.noreply.github.com> Date: Wed, 9 Sep 2026 14:27:10 +0530 Subject: [PATCH 1/5] perf(foundry-hosting): cache FoundryStateStore in FoundryAgentSessionStore Reuse one process-wide FoundryStateStore for agent-session persistence instead of rebuilding it (new credential + agent_sessions metadata round-trip via get_or_create) on every get/set/delete. The session set() runs on the critical path of every Responses request, so the redundant work was pure per-request latency. Per-request user isolation is preserved via the per-operation call_id, so a shared store is equivalent; the store is no longer entered as an async-with context (its aclose() would defeat the cache) and is kept open for the process lifetime. --- .../_state_store.py | 45 ++++++++++++++----- 1 file changed, 34 insertions(+), 11 deletions(-) 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..5a3ab283c50 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,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, @@ -300,32 +301,54 @@ class FoundryAgentSessionStore(SessionStore): DEFAULT_ROOT_SCOPE = "agent_sessions" + # Process-wide cache of the backing state store. The agent-session scope + # ("agent_sessions", user_isolation=True) is identical for every request, + # and per-request user isolation is enforced through the per-operation + # ``call_id`` argument -- not through the store instance -- so a single + # shared store is equivalent to a per-request one. Caching it avoids + # rebuilding the store on every get/set/delete, where each rebuild creates a + # fresh credential (empty token cache -> a new managed-identity token fetch) + # and issues an ``agent_sessions`` metadata round-trip via ``get_or_create`` + # before the actual item operation. Because the session ``set`` runs on the + # critical path of every Responses request, that redundant work is pure + # per-request latency. + _shared_store: ClassVar[FoundryStateStore | None] = None + _shared_store_lock: ClassVar[asyncio.Lock] = asyncio.Lock() + def __init__(self, platform_context: FoundryAgentRequestContext) -> None: self.platform_context = platform_context async def _get_store(self) -> FoundryStateStore: - return await FoundryStateStore.get_or_create( - f"{self.DEFAULT_ROOT_SCOPE}", - user_isolation=True, - ) + # Fast path: already resolved -> no lock, no metadata round-trip. + if FoundryAgentSessionStore._shared_store is not None: + return FoundryAgentSessionStore._shared_store + async with FoundryAgentSessionStore._shared_store_lock: + if FoundryAgentSessionStore._shared_store is None: + FoundryAgentSessionStore._shared_store = await FoundryStateStore.get_or_create( + f"{self.DEFAULT_ROOT_SCOPE}", + user_isolation=True, + ) + return FoundryAgentSessionStore._shared_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 and owned credential and defeat the cache. + # The store is created once and kept open for the process lifetime; + # process exit reclaims 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]): From b8babbe6105a0177854623780e30f4a744e19eb4 Mon Sep 17 00:00:00 2001 From: Harsheet Shah <50236780+harsheet-shah@users.noreply.github.com> Date: Wed, 9 Sep 2026 14:37:46 +0530 Subject: [PATCH 2/5] test(foundry-hosting): isolate + cover cached FoundryAgentSessionStore Add an autouse fixture that resets the new process-wide FoundryStateStore cache between tests so each agent-session test observes its own patched get_or_create, and add a test asserting the store is resolved once and reused across set/get/delete. --- .../foundry_hosting/tests/test_state_store.py | 37 ++++++++++++++++++- 1 file changed, 36 insertions(+), 1 deletion(-) diff --git a/python/packages/foundry_hosting/tests/test_state_store.py b/python/packages/foundry_hosting/tests/test_state_store.py index fc36d00373c..fa72f16e5e9 100644 --- a/python/packages/foundry_hosting/tests/test_state_store.py +++ b/python/packages/foundry_hosting/tests/test_state_store.py @@ -1,5 +1,5 @@ # Copyright (c) Microsoft. All rights reserved. -from collections.abc import Callable +from collections.abc import Callable, Iterator from dataclasses import dataclass from types import SimpleNamespace from typing import Any @@ -72,6 +72,20 @@ 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 process-wide FoundryAgentSessionStore state-store cache. + + FoundryAgentSessionStore caches one FoundryStateStore for the whole process + (a latency optimisation), which would otherwise leak a test's mocked store + into later tests. Clear it before and after every test so each test observes + its own patched ``get_or_create``. + """ + FoundryAgentSessionStore._shared_store = None + yield + FoundryAgentSessionStore._shared_store = None + + def test_storage_providers_use_public_abstraction() -> None: assert issubclass(CheckpointStoreProvider, ContextScopedStoreProvider) assert not issubclass(CheckpointStoreProvider, StoreProvider) @@ -508,3 +522,24 @@ 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() From c34026ba5d6a1ff7447de1f2fadf553b35324183 Mon Sep 17 00:00:00 2001 From: Harsheet Shah <50236780+harsheet-shah@users.noreply.github.com> Date: Wed, 9 Sep 2026 15:10:44 +0530 Subject: [PATCH 3/5] perf(foundry-hosting): scope FoundryAgentSessionStore cache per event loop and scope Address review: the store owns a loop-bound async pipeline + credential, so a process-wide singleton could be reused from a different event loop (across asyncio.run() calls or loop-scoped tests). Cache the store in a WeakKeyDictionary keyed by the running loop (closed loops -> their stores are GC'd) and by scope, with a per-loop lock, so a subclass overriding DEFAULT_ROOT_SCOPE no longer shares or clobbers the base collection. The single-loop server still shares one store, preserving the latency win. --- .../_state_store.py | 64 +++++++++++-------- 1 file changed, 38 insertions(+), 26 deletions(-) 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 5a3ab283c50..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 @@ -5,6 +5,7 @@ from abc import ABC, abstractmethod from datetime import datetime from typing import ClassVar, Generic, Protocol, TypeVar +from weakref import WeakKeyDictionary from agent_framework import ( AgentSession, @@ -301,41 +302,52 @@ class FoundryAgentSessionStore(SessionStore): DEFAULT_ROOT_SCOPE = "agent_sessions" - # Process-wide cache of the backing state store. The agent-session scope - # ("agent_sessions", user_isolation=True) is identical for every request, - # and per-request user isolation is enforced through the per-operation - # ``call_id`` argument -- not through the store instance -- so a single - # shared store is equivalent to a per-request one. Caching it avoids - # rebuilding the store on every get/set/delete, where each rebuild creates a - # fresh credential (empty token cache -> a new managed-identity token fetch) - # and issues an ``agent_sessions`` metadata round-trip via ``get_or_create`` - # before the actual item operation. Because the session ``set`` runs on the - # critical path of every Responses request, that redundant work is pure - # per-request latency. - _shared_store: ClassVar[FoundryStateStore | None] = None - _shared_store_lock: ClassVar[asyncio.Lock] = asyncio.Lock() + # 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: - # Fast path: already resolved -> no lock, no metadata round-trip. - if FoundryAgentSessionStore._shared_store is not None: - return FoundryAgentSessionStore._shared_store - async with FoundryAgentSessionStore._shared_store_lock: - if FoundryAgentSessionStore._shared_store is None: - FoundryAgentSessionStore._shared_store = await FoundryStateStore.get_or_create( - f"{self.DEFAULT_ROOT_SCOPE}", - user_isolation=True, - ) - return FoundryAgentSessionStore._shared_store + 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 and owned credential and defeat the cache. - # The store is created once and kept open for the process lifetime; - # process exit reclaims it. + # 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() item = await store.get_item(session_id, call_id=self.platform_context.call_id) if item is None: From 5ba3127409a8de1d7c92c92b384c8d3238a7f8c3 Mon Sep 17 00:00:00 2001 From: Harsheet Shah <50236780+harsheet-shah@users.noreply.github.com> Date: Wed, 9 Sep 2026 15:10:47 +0530 Subject: [PATCH 4/5] test(foundry-hosting): cover per-loop cache reset, concurrent init and subclass-scope isolation Reset the per-(loop, scope) cache between tests, and add tests asserting the backing store is resolved once under concurrent first-use and that a subclass overriding DEFAULT_ROOT_SCOPE gets its own cached store. --- .../foundry_hosting/tests/test_state_store.py | 57 ++++++++++++++++--- 1 file changed, 50 insertions(+), 7 deletions(-) diff --git a/python/packages/foundry_hosting/tests/test_state_store.py b/python/packages/foundry_hosting/tests/test_state_store.py index fa72f16e5e9..5a18c41f016 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, Iterator from dataclasses import dataclass from types import SimpleNamespace @@ -74,16 +75,17 @@ def _platform_context(call_id: str = "call-1", user_id: str = "user-1") -> Found @pytest.fixture(autouse=True) def _reset_agent_session_store_cache() -> Iterator[None]: - """Isolate the process-wide FoundryAgentSessionStore state-store cache. + """Isolate the FoundryAgentSessionStore per-(loop, scope) state-store cache. - FoundryAgentSessionStore caches one FoundryStateStore for the whole process - (a latency optimisation), which would otherwise leak a test's mocked store - into later tests. Clear it before and after every test so each test observes - its own patched ``get_or_create``. + 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._shared_store = None + FoundryAgentSessionStore._store_cache.clear() + FoundryAgentSessionStore._cache_locks.clear() yield - FoundryAgentSessionStore._shared_store = None + FoundryAgentSessionStore._store_cache.clear() + FoundryAgentSessionStore._cache_locks.clear() def test_storage_providers_use_public_abstraction() -> None: @@ -543,3 +545,44 @@ async def test_agent_session_store_is_cached_across_operations() -> None: 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() From 10dd90f2a932d6ccb84bbc2d38020ab119cd36d7 Mon Sep 17 00:00:00 2001 From: Harsheet Shah Date: Fri, 11 Sep 2026 14:34:27 +0530 Subject: [PATCH 5/5] perf(foundry-hosting): anchor session-store cache on the loop to fix per-loop leak Address PR review feedback on the FoundryAgentSessionStore cache: - Store the per-loop FoundryStateStore cache (and its creation lock) on the running event loop itself instead of in process-global WeakKeyDictionaries. An asyncio.Lock (and the store's pooled pipeline/credential) strongly references its loop, so a module-global map keyed by the loop kept that "weak" key alive, leaking one cache + open pipeline/credential per closed loop (e.g. every asyncio.run). Anchored to the loop, the state is reclaimed with the loop. Falls back to an uncached resolve if a C-level loop forbids attribute assignment (no reuse, but no leak). - Make the concurrent-init test's mocked get_or_create actually suspend (await asyncio.sleep(0)) so all gathered tasks contend on the creation lock; the previous non-suspending mock let the first task populate the cache synchronously, so even a lock-less implementation would have passed. - Drop the now-unnecessary global cache-reset fixture: per-loop state is isolated automatically by the function-scoped event loop each test runs on. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 3d900395-13d0-4698-bf9d-f6670f9e545c --- .../_state_store.py | 82 ++++++++++++------- .../foundry_hosting/tests/test_state_store.py | 30 ++++--- 2 files changed, 69 insertions(+), 43 deletions(-) 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 f1d31976904..7bb3136c252 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 @@ -5,7 +5,6 @@ from abc import ABC, abstractmethod from datetime import datetime from typing import ClassVar, Generic, Protocol, TypeVar -from weakref import WeakKeyDictionary from agent_framework import ( AgentSession, @@ -297,51 +296,72 @@ 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" - # 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() + # 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_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 + 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: 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. - 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] + 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`` diff --git a/python/packages/foundry_hosting/tests/test_state_store.py b/python/packages/foundry_hosting/tests/test_state_store.py index 5a18c41f016..376fdb83962 100644 --- a/python/packages/foundry_hosting/tests/test_state_store.py +++ b/python/packages/foundry_hosting/tests/test_state_store.py @@ -1,6 +1,6 @@ # Copyright (c) Microsoft. All rights reserved. import asyncio -from collections.abc import Callable, Iterator +from collections.abc import Callable from dataclasses import dataclass from types import SimpleNamespace from typing import Any @@ -74,18 +74,15 @@ def _platform_context(call_id: str = "call-1", user_id: str = "user-1") -> Found @pytest.fixture(autouse=True) -def _reset_agent_session_store_cache() -> Iterator[None]: - """Isolate the FoundryAgentSessionStore per-(loop, scope) state-store cache. +def _reset_agent_session_store_cache() -> None: + """Guard 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. + 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. """ - FoundryAgentSessionStore._store_cache.clear() - FoundryAgentSessionStore._cache_locks.clear() - yield - FoundryAgentSessionStore._store_cache.clear() - FoundryAgentSessionStore._cache_locks.clear() + return None def test_storage_providers_use_public_abstraction() -> None: @@ -552,9 +549,18 @@ async def test_agent_session_store_concurrent_init_resolves_once() -> None: 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(return_value=store), + 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)))