feat: add in-memory LRU cache for per-user MCP OAuth access tokens with targeted eviction

[Pratham-Mishra04]

Summary

Per-user MCP OAuth access tokens were fetched from the database on every request. This PR introduces a bounded, LRU in-memory cache for those lookups, keyed by (auth mode, identity, mcp client), so repeated requests for the same binding skip the database round-trip. Expired entries are treated as misses and fall through to the full database path, which retains ownership of all expiry and refresh decisions. The cache is kept consistent through targeted evictions wired into every write path that can invalidate a cached token: refresh, revocation, force-refresh, OAuth flow completion, credential rotation, access reconciliation, VK reload/delete, and MCP client update/delete.

Changes

• Added userTokenCache in framework/oauth2/usertokencache.go: a thin wrapper around the existing LRU cache that owns the (mode, identity, mcp client) key scheme, registers each entry under its token row ID for targeted eviction, validates entries against their expiry on read, and deduplicates concurrent fills for the same key so a refresh burst triggers at most one upstream call.
• Split GetUserAccessTokenByMode into a cache-aware outer function and a loadUserAccessTokenByMode cache-miss path. Cache hits bypass the database entirely; misses run the full lookup including the lazy pre-flight refresh.
• Added eviction methods on OAuth2Provider: EvictUserToken, EvictUserTokenByID, EvictUserTokensByMCPClient, EvictUserTokensByVirtualKey, EvictUserTokensByUser, and FlushUserTokenCache. All are nil-safe.
• Wired evictions into every write path in main.go: after RefreshAccessToken (both success and permanent rejection), ForceRefreshAccessToken (pre-eviction before the refresh call), RevokeToken, CompleteOAuthFlow, and CompleteUserOAuthFlow.
• Added MCPOauthTokenCacheManager interface in mcpsessions.go with EvictOauthTokenCacheByID and EvictOauthTokenCacheByMCPClient. Implemented on BifrostHTTPServer and wired into MCPSessionsHandler (session revoke) and MCPHandler (credential rotation, reconciliation after MCP client update).
• Added EvictOauthTokenCacheByID, EvictOauthTokenCacheByMCPClient, EvictOauthTokenCacheByVirtualKey, and FlushOauthTokenCache to ServerCallbacks and implemented them on BifrostHTTPServer, delegating to the provider's eviction methods.
• Evictions added to RemoveMCPClient, ReloadVirtualKey, and RemoveVirtualKey in server.go to cover token rows cascaded by database deletes.
• Reordered updateVirtualKey in governance.go so ReconcileOauthAfterVKChange runs before ReloadVirtualKey: the reload evicts VK-scoped tokens, and an eviction landing before the reconcile writes could be refilled from pre-reconcile rows and never dropped again.
• Added comprehensive tests in usertokencache_test.go covering: hit/miss, LRU capacity eviction, expired-entry miss-and-remove, exact-key eviction, token-ID eviction, MCP-client bulk eviction, virtual-key bulk eviction, user bulk eviction, flush, in-flight deduplication, error sharing without caching, generation-guard discarding a stale fill that raced an eviction, token-ID index rebinding after upsert, nil safety, and integration tests through GetUserAccessTokenByMode for cache hits, concurrent refresh deduplication, post-delete eviction, post-needs-reauth flush, negative-result non-caching, and force-refresh eviction.

Type of change

• Bug fix
• Feature
• Refactor
• Documentation
• Chore/CI

Affected areas

• Core (Go)
• Transports (HTTP)
• Providers/Integrations
• Plugins
• UI (React)
• Docs

How to test

go test ./framework/oauth2/... ./transports/bifrost-http/...

The new usertokencache_test.go file covers the cache unit behaviour and integration through GetUserAccessTokenByMode. Key scenarios to validate manually:

1. Make two identical per-user token requests back-to-back and confirm only one database query is issued.
2. Revoke a session token through the sessions API and confirm the next request re-fetches from the database rather than returning the cached token.
3. Rotate OAuth credentials on an MCP client and confirm subsequent token lookups reflect the new credentials.
4. Delete a virtual key and confirm its cached tokens are no longer served.

Breaking changes

• Yes
• No

NewMCPHandler and NewMCPSessionsHandler have new required parameters (MCPOauthTokenCacheManager). Any caller constructing these handlers directly outside the server must pass an implementation of MCPOauthTokenCacheManager. The ServerCallbacks interface has two new methods (EvictOauthTokenCacheByID, EvictOauthTokenCacheByMCPClient); any custom ServerCallbacks implementation must add them.

Security considerations

Cached access tokens are held in process memory and are never persisted. The cache is bounded to 4096 entries by default to prevent unbounded growth from caller-asserted session identities. Evictions are wired to every write path that can invalidate a token, including revocation and needs-reauth transitions, so a revoked or expired token cannot be served from cache beyond the current request's lifetime.

Checklist

• I read docs/contributing/README.md and followed the guidelines
• I added/updated tests where appropriate
• I updated documentation where needed
• I verified builds succeed (Go and UI)
• I verified the CI pipeline passes locally if applicable

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[Pratham-Mishra04]

• feat: add stack-absorb Claude skill for manual Graphite stack change distribution #5786 : 2 dependent PRs (#5787 , #5812 )
• docs+feat: expand MCP connection states into reference table and add tooltip to State column header #5785
• feat: add token-exchange demo MCP server for testing against real identity providers #5730
• docs: add token exchange (on-behalf-of) MCP auth documentation #5729
• docs: document make-before-break reconnection and shared-client retry behavior #5728
• feat: make-before-break reconnect for HTTP/SSE MCP clients with generation-guarded stale write protection #5727
• feat: retry shared-connection MCP auth failures after bounded reconnect wait with concurrent dedup #5726
• feat: add ui for token_exchange (On-Behalf-Of) auth type MCP clients #5725
• feat: add token_exchange MCP auth type with delegated identity-provider token exchange #5724
• docs: add needs_reauth state, reauthorize endpoint, per-client tool sync/timeout, allowed_extra_headers, and OAuth config rotation to MCP OpenAPI spec #5723
• docs: add OAuth reauthorization, admin discovery credentials, auth failure recovery, and credential rotation docs #5722
• docs: add env./vault. reference support for OAuth client_id/client_secret and update OpenAPI EnvVar schema #4860
• docs: add pending_verification state and initiate-verification / verify-headers MCP client endpoints to OpenAPI spec #4858
• docs: add config.json bootstrap support for auth-aware MCP clients (oauth, per_user_oauth, per_user_headers) #3876
• feat: use *SecretVar for OAuth2Config credentials to preserve env./vault. #4859
• feat: add LRU credential cache to MCP headers provider with scoped eviction #5721
• feat: add in-memory LRU cache for per-user MCP OAuth access tokens with targeted eviction #5720 👈 (View in Graphite)
• feat: add generation-guarded LRU cache with single-flight fills and secondary index #5719
• feat: add per-user OAuth/headers admin discovery credential repair flow with registry needs_reauth projection #5718
• feat: retain admin bootstrap credential for periodic per-user MCP tool discovery refresh #5717
• feat: add MCPAuthModeAdmin + reauthorize endpoint for shared OAuth MCP clients #5716
• feat: replace OAuth config pin-and-warn with in-place rotation and needs_reauth cascade #5715
• feat: add needs_reauth state and replace per-user connection badge with sessions link #5714
• feat: add auth-failure retry with forced credential refresh for MCP tool calls #5713
• feat: add NeedsReauth MCP connection state for dead OAuth2 credentials #5712
• refactor: make GetExpiringOauthTokens auth modes caller-configurable via AuthModes field #5711
• refactor: unify shared and per-user OAuth refresh under a single RefreshAccessToken(tokenID) path, dropping oauth_configs.token_id FK shortcut #5710
• refactor: migrate OAuth flow state/PKCE fields from oauth_configs to new mcp_oauth_flows table #5709
• refactor: merge oauth_tokens/oauth_user_tokens into unified mcp_oauth_tokens table #5708
• feat: add per_user_headers bootstrap verification flow via submitHandler abstraction #3875
• feat: add verify-headers endpoint and pending_verification state for per-user-headers MCP clients #3874
• fix: extend pending_verification bootstrap path to cover per_user_oauth clients #3873
• feat: add OAuth authorization flow for config.json-bootstrapped pending_verification MCP clients #3872
• feat: add pending_verification state and initiate-verification endpoint for config.json shared-OAuth MCP clients #3871
• docs: document flow mode access rules, temp token toggle, and per-server tool manager overrides #3870 : 1 other dependent PR (#3882 )
• dev

This stack of pull requests is managed by Graphite. Learn more about stacking.

[coderabbitai]

📝 Walkthrough

Summary by CodeRabbit

• New Features
  • Added efficient per-user OAuth token caching with expiry-aware refresh handling.
  • Added controls to remove individual tokens or clear cached tokens by client, virtual key, user, or all entries.
• Bug Fixes
  • Prevented stale credentials after updates, rotations, deletions, revocations, and reauthentication.
  • Improved concurrent token retrieval and refresh behavior.
  • Failed token lookups are no longer retained in cache.
• Documentation
  • Clarified that removing virtual keys or MCP clients also removes related cached OAuth tokens.

Walkthrough

The PR adds a bounded per-user OAuth token cache with expiry-aware loading, single-flight fills, sanitized cached values, and scoped eviction. MCP handlers and server callbacks evict entries after credential, token, client, and virtual-key changes.

Changes

OAuth token cache

Layer / File(s)	Summary
Cache storage and eviction primitives framework/oauth2/usertokencache.go, framework/oauth2/usertokencache_test.go	Adds bounded LRU storage, expiry checks, single-flight fills, token indexing, scoped eviction, flushing, and lifecycle tests.
Provider cache-backed token flow framework/oauth2/main.go, framework/oauth2/usertokencache_test.go	Routes per-user token reads through the cache, refreshes expired credentials, sanitizes cached access tokens, and evicts entries after token mutations.
MCP and server invalidation wiring transports/bifrost-http/handlers/*.go, transports/bifrost-http/server/server.go	Wires cache-manager callbacks into MCP handlers and evicts tokens after MCP-client, session-token, virtual-key, and OAuth configuration changes.

Estimated code review effort: 4 (Complex) | ~45 minutes

Sequence Diagram(s)

sequenceDiagram
  participant MCPHandler
  participant BifrostHTTPServer
  participant OAuth2Provider
  participant UserTokenCache
  MCPHandler->>BifrostHTTPServer: reconcile MCP credentials
  BifrostHTTPServer->>OAuth2Provider: evict tokens by MCP client
  OAuth2Provider->>UserTokenCache: remove matching entries
  MCPHandler-->>BifrostHTTPServer: complete configuration update

Loading

Possibly related PRs

• maximhq/bifrost#5713: Adds forced refresh behavior supported by this PR’s cache eviction.
• maximhq/bifrost#5715: Coordinates OAuth refresh and MCP credential rotation with token invalidation.
• maximhq/bifrost#5721: Uses related per-user MCP credential cache and handler eviction integration.

Suggested reviewers: akshaydeo, bearts, danpiths

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 20.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly and concisely summarizes the primary change: adding a bounded per-user MCP OAuth access-token cache with targeted eviction.
Description check	✅ Passed	The description covers the purpose, implementation, testing, affected areas, breaking changes, security considerations, and checklist, with only minor omissions.

✨ Finishing Touches 💡 1

📝 Generate docstrings 💡

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch 07-30-feat_cache_per-user_mcp_oauth_token_lookups_in_memory

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

[Pratham-Mishra04]

Merge activity

• Aug 8, 8:47 AM UTC: A user started a stack merge that includes this pull request via Graphite.
• Aug 8, 9:36 AM UTC: Graphite rebased this pull request as part of a merge.
• Aug 8, 9:37 AM UTC: @Pratham-Mishra04 merged this pull request with Graphite.

The base branch was changed.