Skip to content

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

Closed
Harsheet Shah (harsheet-shah) wants to merge 4 commits into
microsoft:mainfrom
harsheet-shah:harsheetshah/maf-caching-improvements
Closed

Python: perf(foundry-hosting): cache FoundryStateStore in FoundryAgentSessionStore#8178
Harsheet Shah (harsheet-shah) wants to merge 4 commits into
microsoft:mainfrom
harsheet-shah:harsheetshah/maf-caching-improvements

Conversation

@harsheet-shah

Copy link
Copy Markdown

Summary

FoundryAgentSessionStore (the default agent-session SessionStore for hosted MAF agents) rebuilds its backing FoundryStateStore on every get/set/delete. Each _get_store() call goes through FoundryStateStore.get_or_create("agent_sessions", user_isolation=True), which:

  1. constructs a fresh credential (empty token cache → a new managed-identity token fetch), and
  2. issues an agent_sessions metadata round-trip (GET/POST state_stores) before the actual item operation.

Because the hosting infra calls set() in the finally of every Responses request, this redundant credential + metadata work lands on the critical path of every request.

This PR caches one process-wide FoundryStateStore for the agent-session scope, so only the real item GET/PUT/DELETE remains on the hot path.

Why it's behaviour-preserving

  • The agent-session scope ("agent_sessions", user_isolation=True) is identical for every request, so one shared store == a per-request store.
  • Per-request user isolation is enforced through the per-operation call_id argument (unchanged), not through the store instance.
  • Item get/set/delete semantics are byte-for-byte identical; get_or_create still creates-on-first-use exactly once.
  • The one intentional change: the shared store is not entered via async with — its __aexit__ calls aclose(), 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 an asyncio.Lock with 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):

Metric Baseline Cached Δ
WARM TTLB p50 7,736 7,464 −272 ms (−3.5%)
WARM TTFB p50 5,680 5,683 ~flat

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, FoundryAgentSessionStore only. No public API change.

…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.
Copilot AI balanced review requested due to automatic review settings September 9, 2026 08:57
@agent-framework-automation agent-framework-automation Bot added the python Usage: [Issues, PRs], Target: Python label Sep 9, 2026
@github-actions github-actions Bot changed the title perf(foundry-hosting): cache FoundryStateStore in FoundryAgentSessionStore Python: perf(foundry-hosting): cache FoundryStateStore in FoundryAgentSessionStore Sep 9, 2026

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 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 overridable self.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.

Comment on lines +315 to +316
_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.
@harsheet-shah

Copy link
Copy Markdown
Author

Thanks for the review — both points addressed in the latest commits:

1. Event-loop scoping (lines 316/327). The FoundryStateStore owns a loop-bound async pipeline + credential, so a process-wide singleton could be reused from a different loop (across asyncio.run() calls or between loop-scoped tests). The cache is now a WeakKeyDictionary keyed by the running event loop with a per-loop lock; when a loop closes its store is garbage-collected with it (verified: entry count drops to 0 after loop close), so a stale store is never reused on another loop. The long-running server still runs one loop, so the latency win is preserved. Added test_agent_session_store_concurrent_init_resolves_once to cover concurrent first-use resolving the store exactly once.

2. Subclass scope routing (line 328). The cache is now also keyed by scope, so a subclass overriding DEFAULT_ROOT_SCOPE gets its own cached store instead of sharing/clobbering the base collection. Added test_agent_session_store_subclass_scope_is_isolated.

The per-test fixture now clears the per-(loop, scope) cache. All Python test / typing / coverage gates are green.

@eavanvalkenburg

Copy link
Copy Markdown
Member

Harsheet Shah (@harsheet-shah) please use the defined PR template

@eavanvalkenburg

Copy link
Copy Markdown
Member

Closing, please create a issue first Harsheet Shah (@harsheet-shah)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

python Usage: [Issues, PRs], Target: Python

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants