Python: perf(foundry-hosting): cache FoundryStateStore in FoundryAgentSessionStore - #8178
Conversation
…Store 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.
There was a problem hiding this comment.
🟡 Changes recommended
The shared cache is unsafe across event loops and can route subclassed stores to the wrong scope.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
Caches the Foundry-backed agent session store to remove repeated credential creation and metadata requests.
Changes:
- Adds locked, lazy shared-store initialization.
- Reuses the store without closing it after each operation.
File summaries
| File | Description |
|---|---|
_state_store.py |
Adds process-wide session-store caching. |
Review details
Suppressed comments (1)
python/packages/foundry_hosting/agent_framework_foundry_hosting/_state_store.py:328
- This cache is attached explicitly to
FoundryAgentSessionStore, while the lookup still uses overridableself.DEFAULT_ROOT_SCOPE. A subclass with a different scope will therefore either reuse the base class's store or cause the base class to reuse the subclass's store, routing sessions to the wrong collection. Key the cache by scope/type or consistently use a non-overridable scope.
FoundryAgentSessionStore._shared_store = await FoundryStateStore.get_or_create(
f"{self.DEFAULT_ROOT_SCOPE}",
- Files reviewed: 1/1 changed files
- Comments generated: 1
- Review effort level: Balanced
💡 Add a code-review agent skill for context-aware, tailored reviews. Learn more in the docs.
| _shared_store: ClassVar[FoundryStateStore | None] = None | ||
| _shared_store_lock: ClassVar[asyncio.Lock] = asyncio.Lock() |
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.
… 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.
…d 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.
|
Thanks for the review — both points addressed in the latest commits: 1. Event-loop scoping (lines 316/327). The 2. Subclass scope routing (line 328). The cache is now also keyed by scope, so a subclass overriding The per-test fixture now clears the per-(loop, scope) cache. All Python test / typing / coverage gates are green. |
|
Harsheet Shah (@harsheet-shah) please use the defined PR template |
|
Closing, please create a issue first Harsheet Shah (@harsheet-shah) |
Summary
FoundryAgentSessionStore(the default agent-sessionSessionStorefor hosted MAF agents) rebuilds its backingFoundryStateStoreon everyget/set/delete. Each_get_store()call goes throughFoundryStateStore.get_or_create("agent_sessions", user_isolation=True), which:agent_sessionsmetadata round-trip (GET/POST state_stores) before the actual item operation.Because the hosting infra calls
set()in thefinallyof every Responses request, this redundant credential + metadata work lands on the critical path of every request.This PR caches one process-wide
FoundryStateStorefor the agent-session scope, so only the real itemGET/PUT/DELETEremains on the hot path.Why it's behaviour-preserving
"agent_sessions",user_isolation=True) is identical for every request, so one shared store == a per-request store.call_idargument (unchanged), not through the store instance.get/set/deletesemantics are byte-for-byte identical;get_or_createstill creates-on-first-use exactly once.async with— its__aexit__callsaclose(), which would close the pooled pipeline + owned credential and defeat the cache. The store is opened once and kept open for the process lifetime (process exit reclaims it). Concurrent first-use is guarded by anasyncio.Lockwith double-checked init.Checkpoint / function-approval stores are intentionally left untouched — they are not on the per-request hot path.
Measured impact
500 cold + 500 warm streaming requests per agent, concurrency 50, private/VNet Foundry project,
gpt-4o-mini, 0 errors / 0 throttled. Same-conditions A/B, baseline vs cached Responses agent (ms, p50/p95):The remaining warm latency is the model first-token floor (~5.7 s TTFB), which is unaffected by this change. The win is the elimination of the per-request credential + metadata round-trip.
Scope
Single file,
FoundryAgentSessionStoreonly. No public API change.