Skip to content

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

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

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

Conversation

@harsheet-shah

@harsheet-shah Harsheet Shah (harsheet-shah) commented Sep 11, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Motivation & Context

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.

Description & Review Guide

  • What are the major changes?
    Cache one per-event-loop, per-scope FoundryStateStore for the agent-session scope. The cache is a WeakKeyDictionary keyed by the running event loop (a closed loop's store is GC'd, so a stale store is never reused on another loop), further keyed by scope (a subclass overriding DEFAULT_ROOT_SCOPE gets its own store instead of clobbering the base collection). Concurrent first-use is guarded by a per-loop lock with double-checked init. The shared store is intentionally not entered via async with (its __aexit__ would aclose() the pooled pipeline + credential and defeat the cache); it is opened once and kept open for the process lifetime.

  • What is the impact of these changes?
    Only the real item GET/PUT/DELETE remains on the hot path. Item get/set/delete semantics are byte-for-byte identical; get_or_create still creates-on-first-use exactly once. Checkpoint / function-approval stores are left untouched (not on the per-request hot path). No public API change.

    Measured (500 cold + 500 warm streaming req/agent, concurrency 50, private/VNet Foundry project, gpt-4o-mini, 0 errors): WARM TTLB p50 7,736 → 7,464 ms (−272 ms, −3.5%); WARM TTFB p50 flat (~5.68 s model first-token floor, unaffected). The win is elimination of the per-request credential + metadata round-trip.

  • What do you want reviewers to focus on?
    Event-loop scoping / GC of the cache, and the subclass-scope isolation.

Benchmark — in-region A/B (baseline vs this PR)

Head-to-head warm-path measurement, both arms deployed in the same private/VNet + shared-NC Foundry project and driven from an in-region VMSS (client↔service RTT removed). Streaming responses, concurrency 25, n=1,999 warm req/arm, 0 errors, 0 throttles. Warm TTLB is the primary metric (the private flush-floor masks TTFB).

Warm metric baseline #8281 Δ
TTLB p50 2,812.2 ms 2,567.7 ms −244.5 ms (−8.7%)
TTLB p95 4,427.5 ms 4,116.3 ms −311.2 ms (−7.0%)
TTLB p99 6,469.7 ms 4,912.1 ms −1,557.6 ms (−24.1%)
TTFB p50 874.8 ms 907.0 ms flat (flush-floor, non-diagnostic)

This PR improves the typical warm request (median TTLB −8.7%) and substantially tightens the tail (p99 −24.1%), with no regression. Consistent with removing the per-request credential + state_stores metadata round-trip from the hot path.

Related Issue

Closes #8280

Contribution Checklist

  • The code builds clean without any errors or warnings
  • All unit tests pass, and I have added new tests where possible
  • The PR follows the Contribution Guidelines
  • This PR is linked to an issue and there is no other open PR for this issue (see Related Issue above).
  • This is not a breaking change.

@harsheet-shah
perf(foundry-hosting): cache FoundryStateStore in FoundryAgentSession…
97d85e8 
…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.
@harsheet-shah
test(foundry-hosting): isolate + cover cached FoundryAgentSessionStore
b8babbe 
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.
@harsheet-shah
perf(foundry-hosting): scope FoundryAgentSessionStore cache per event…
c34026b 
… 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.
@harsheet-shah
test(foundry-hosting): cover per-loop cache reset, concurrent init an…
5ba3127 
…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.
Copilot AI balanced review requested due to automatic review settings September 11, 2026 04:58
@harsheet-shah Harsheet Shah (harsheet-shah) mentioned this pull request Sep 11, 2026
Closed
@agent-framework-automation agent-framework-automation Bot added the python Usage: [Issues, PRs], Target: Python label Sep 11, 2026
Copilot AI reviewed Sep 11, 2026
View reviewed changes

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

Contended locks retain closed event loops, preventing cached stores and resources from being reclaimed.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

Caches Foundry agent-session state stores to remove repeated credential creation and metadata requests.

Changes:

  • Adds per-loop, per-scope caching with initialization locking.
  • Adds cache reuse, concurrency, and scope-isolation tests.
File summaries
File Description
_state_store.py Implements cached store reuse.
test_state_store.py Tests caching and initialization.
Review details
  • Files reviewed: 2/2 changed files
  • Comments generated: 2
  • Review effort level: Balanced

💡 Add a code-review agent skill for context-aware, tailored reviews. Learn more in the docs.

Comment thread python/packages/foundry_hosting/agent_framework_foundry_hosting/_state_store.py Resolved
Comment thread python/packages/foundry_hosting/tests/test_state_store.py Resolved
Harsheet Shah
perf(foundry-hosting): anchor session-store cache on the loop to fix …
10dd90f 
…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
@eavanvalkenburg

Eduard van Valkenburg (eavanvalkenburg) commented Sep 14, 2026

Copy link
Copy Markdown
Member

Thanks for the update. Before this is ready, could you please:

Once those are addressed, please re-request review. Thanks!

Harsheet Shah
perf(foundry-hosting): give cached session stores an explicit shutdow…
031f161 
…n owner

Address review: the per-loop FoundryStateStore cache deliberately skips
aclose(), relying on loop collection to reclaim each store's pooled pipeline
and owned credential. Add FoundryAgentSessionStore.aclose_loop_cache() and
AgentSessionStoreProvider.aclose(), registered as a ResponsesHostServer
shutdown handler, so server shutdown explicitly closes those resources while
preserving the reuse fast path.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Copilot-Session: 3d900395-13d0-4698-bf9d-f6670f9e545c
Closed
@harsheet-shah Harsheet Shah (harsheet-shah) mentioned this pull request Sep 28, 2026
Closed
4 of 5 tasks
@eavanvalkenburg

Eduard van Valkenburg (eavanvalkenburg) commented Sep 28, 2026

Copy link
Copy Markdown
Member

Harsheet Shah (@harsheet-shah) This branch currently conflicts with main. Please merge or rebase current main and resolve the conflicts before requesting another review.

@eavanvalkenburg

Eduard van Valkenburg (eavanvalkenburg) commented Sep 28, 2026

Copy link
Copy Markdown
Member

Harsheet Shah (@harsheet-shah) Thank you for the work and benchmark data here. We are going to continue in #8785 because its request-scoped client ownership fits the current sandbox-scoped, ETag-guarded session storage and avoids introducing a cross-request cache lifecycle. Please carry any useful performance evidence into #8785. I am closing this alternative so there is one implementation path.

This branch was previously deployed

1 inactive deployment
github-app-auth — 031f1613 Deployed Sep 17, 2026 by harsheet-shah via add_label #23065
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@moonbox3 Evan Mattson (moonbox3) moonbox3 left review comments
Copilot code review Copilot
Copilot review effort
Applies to this pull request for everyone.Learn more about Copilot code review.
Copilot left review comments
@chetantoshniwal chetantoshniwal Awaiting requested review from chetantoshniwal
@TaoChenOSU Tao Chen (TaoChenOSU) Awaiting requested review from TaoChenOSU
@jpalvarezl Jose Alvarez (jpalvarezl) Awaiting requested review from jpalvarezl jpalvarezl is a code owner
@eavanvalkenburg Eduard van Valkenburg (eavanvalkenburg) Awaiting requested review from eavanvalkenburg eavanvalkenburg is a code owner
@SergeyMenshykh SergeyMenshykh Awaiting requested review from SergeyMenshykh SergeyMenshykh is a code owner
@rogerbarreto Roger Barreto (rogerbarreto) Awaiting requested review from rogerbarreto rogerbarreto is a code owner
@westey-m westey (westey-m) Awaiting requested review from westey-m westey-m is a code owner
@giles17 Giles Odigwe (giles17) Awaiting requested review from giles17

Assignees

No one assigned

Labels

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

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

Python: perf(foundry-hosting): FoundryAgentSessionStore rebuilds FoundryStateStore (credential + metadata round-trip) on every request

4 participants

@harsheet-shah @eavanvalkenburg @moonbox3