feat: add LRU credential cache to MCP headers provider with scoped eviction - #5721
Merged
Pratham-Mishra04 merged 1 commit intoAug 8, 2026
Conversation
This was referenced Jul 30, 2026
Merged
refactor: migrate OAuth flow state/PKCE fields from
oauth_configs to new mcp_oauth_flows table
#5709
Merged
|
|
This was referenced Jul 30, 2026
Merged
Merged
Collaborator
Author
Contributor
📝 WalkthroughSummary by CodeRabbit
WalkthroughAdds an indexed LRU cache for MCP header credentials. Provider lookups use cached copies and invalidate entries after mutations. MCP handlers and server callbacks evict cached credentials during MCP lifecycle operations. ChangesMCP credential cache
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant MCPHandler
participant MCPSessionsHandler
participant BifrostHTTPServer
participant Provider
MCPHandler->>Provider: update or reconcile credentials
MCPHandler->>BifrostHTTPServer: evict header credentials
MCPSessionsHandler->>BifrostHTTPServer: evict revoked credential or client
BifrostHTTPServer->>Provider: invoke scoped cache eviction
Provider-->>BifrostHTTPServer: invalidate matching cache entries
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Comment |
This was referenced Aug 6, 2026
Pratham-Mishra04
force-pushed
the
07-30-feat_cache_per-user_mcp_header_credential_lookups_in_memory
branch
from
August 6, 2026 21:53
f04d622 to
a5c418f
Compare
Pratham-Mishra04
force-pushed
the
07-30-feat_cache_per-user_mcp_oauth_token_lookups_in_memory
branch
from
August 6, 2026 21:53
9004cee to
2c5a041
Compare
Contributor
|
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. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
framework/mcp_headers/credentialcache_test.go (1)
58-72: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winTest access-based LRU eviction.
This test only verifies insertion-order eviction. A FIFO cache also passes it.
Read
k1after fillingk1andk2. Then fillk3and assert thatk2was evicted.Proposed test update
- for i := 1; i <= 3; i++ { + for i := 1; i <= 2; i++ { _, err := c.Fill(context.Background(), fmt.Sprintf("k%d", i), fillWithCredential(testCredential(fmt.Sprintf("c%d", i), nil)), ) require.NoError(t, err) } + _, ok := c.Get("k1") + require.True(t, ok) + _, err := c.Fill(context.Background(), "k3", fillWithCredential(testCredential("c3", nil))) + require.NoError(t, err) + assert.Equal(t, 2, c.Len(), "cache must stay at capacity") - _, ok := c.Get("k1") + _, ok = c.Get("k2") assert.False(t, ok, "least recently used entry must be evicted") - _, ok = c.Get("k2") + _, ok = c.Get("k1") assert.True(t, ok)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@framework/mcp_headers/credentialcache_test.go` around lines 58 - 72, Update the cache test around the Fill loop to access k1 after inserting k1 and k2, making it the most recently used entry before inserting k3. Keep the capacity assertion, then assert k2 is absent while k1 and k3 remain present, using the existing c.Get checks.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@framework/mcp_headers/credentialcache_test.go`:
- Around line 58-72: Update the cache test around the Fill loop to access k1
after inserting k1 and k2, making it the most recently used entry before
inserting k3. Keep the capacity assertion, then assert k2 is absent while k1 and
k3 remain present, using the existing c.Get checks.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 165b8cfb-31a9-4c8b-b473-dae06ed620d7
📒 Files selected for processing (11)
framework/lrucache/lrucache.goframework/lrucache/lrucache_test.goframework/mcp_headers/credentialcache.goframework/mcp_headers/credentialcache_test.goframework/mcp_headers/main.goframework/mcp_headers/sweep.goframework/oauth2/usertokencache.gotransports/bifrost-http/handlers/governance.gotransports/bifrost-http/handlers/mcp.gotransports/bifrost-http/handlers/mcpsessions.gotransports/bifrost-http/server/server.go
🚧 Files skipped from review as they are similar to previous changes (9)
- framework/mcp_headers/sweep.go
- framework/lrucache/lrucache_test.go
- framework/oauth2/usertokencache.go
- transports/bifrost-http/handlers/governance.go
- transports/bifrost-http/server/server.go
- transports/bifrost-http/handlers/mcpsessions.go
- framework/lrucache/lrucache.go
- framework/mcp_headers/main.go
- transports/bifrost-http/handlers/mcp.go
Pratham-Mishra04
force-pushed
the
07-30-feat_cache_per-user_mcp_oauth_token_lookups_in_memory
branch
from
August 8, 2026 08:43
2c5a041 to
c50dff4
Compare
Pratham-Mishra04
force-pushed
the
07-30-feat_cache_per-user_mcp_header_credential_lookups_in_memory
branch
from
August 8, 2026 08:43
a5c418f to
8157cc2
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@framework/lrucache/lrucache.go`:
- Around line 427-432: Update DecodeKey to reject any n greater than len(key)/2
before allocating the parts slice, while preserving the existing negative-count
rejection and valid decoding behavior. Add a regression test covering an
oversized count, such as DecodeKey("", 1<<30), and verify it returns false
without attempting a large allocation.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: f11bf2e1-149a-41b9-85cf-e7ca80063d8a
📒 Files selected for processing (2)
framework/lrucache/lrucache.goframework/lrucache/lrucache_test.go
Comment on lines
+427
to
+432
| func DecodeKey(key string, n int) (parts []string, ok bool) { | ||
| if n < 0 { | ||
| return nil, false | ||
| } | ||
| rest := key | ||
| parts = make([]string, 0, n) |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
Reject infeasible part counts before allocation.
Every encoded part consumes at least two bytes. A valid key requires n <= len(key)/2.
DecodeKey("", 1<<30) currently attempts a large allocation before it returns false. This can panic or exhaust memory if a caller passes an untrusted count. Reject the count before make, and add a regression test.
Proposed fix
func DecodeKey(key string, n int) (parts []string, ok bool) {
- if n < 0 {
+ if n < 0 || n > len(key)/2 {
return nil, false
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| func DecodeKey(key string, n int) (parts []string, ok bool) { | |
| if n < 0 { | |
| return nil, false | |
| } | |
| rest := key | |
| parts = make([]string, 0, n) | |
| func DecodeKey(key string, n int) (parts []string, ok bool) { | |
| if n < 0 || n > len(key)/2 { | |
| return nil, false | |
| } | |
| rest := key | |
| parts = make([]string, 0, n) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@framework/lrucache/lrucache.go` around lines 427 - 432, Update DecodeKey to
reject any n greater than len(key)/2 before allocating the parts slice, while
preserving the existing negative-count rejection and valid decoding behavior.
Add a regression test covering an oversized count, such as DecodeKey("", 1<<30),
and verify it returns false without attempting a large allocation.
Collaborator
Author
Merge activity
|
Pratham-Mishra04
changed the base branch from
07-30-feat_cache_per-user_mcp_oauth_token_lookups_in_memory
to
graphite-base/5721
August 8, 2026 09:35
…equest unblocks instead of waiting on an unrelated leader
Pratham-Mishra04
force-pushed
the
07-30-feat_cache_per-user_mcp_header_credential_lookups_in_memory
branch
from
August 8, 2026 09:38
8157cc2 to
91ea2dd
Compare
Pratham-Mishra04
deleted the
07-30-feat_cache_per-user_mcp_header_credential_lookups_in_memory
branch
August 8, 2026 09:40
This was referenced Aug 8, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds an in-memory LRU cache for per-user MCP header credentials in the
mcp_headersprovider, eliminating redundant database reads on every tool call for the same (auth mode, identity, MCP client) binding. The cache is bounded at 4096 entries, uses single-flight deduplication to prevent stampedes on concurrent misses, and is invalidated through targeted eviction methods wired into every write path that can change or delete credential rows.Changes
headerCredentialCacheinframework/mcp_headers/credentialcache.go, wrapping the existinglrucache.Cachewith a NUL-delimited composite key scheme (mode\x00identity\x00mcpClientID) and a secondary index on credential row ID for targeted eviction without a full scan.GetCredentialByModenow checks the cache before hitting the database. Admin-mode identity is normalized to an empty string before key construction so that stray non-empty identities cannot alias the same row under multiple keys. Bothactiveandneeds_updaterows are cached with their status preserved; failed lookups (including not-found) are never cached. Every cache hit returns a deep copy of the stored credential so caller mutations cannot corrupt the cached value.UpsertCredentialandDeleteCredentialevict by credential row ID immediately after their database writes complete.Provider(EvictCredential,EvictCredentialByID,EvictCredentialsByMCPClient,EvictCredentialsByVirtualKey,EvictCredentialsByUser) plusFlushCredentialCacheas a coarse fallback.MCPOauthTokenCacheManageris renamed toMCPCredentialCacheManagerand extended withEvictMCPHeaderCredentialCacheByIDandEvictMCPHeaderCredentialCacheByMCPClient. All handler construction sites and call sites are updated accordingly.BifrostHTTPServerimplements the two newMCPCredentialCacheManagermethods and callsEvictCredentialsByMCPClient/EvictCredentialsByVirtualKeyfromRemoveMCPClient,ReloadVirtualKey, andRemoveVirtualKeyso cascaded database deletes are reflected in the cache.updateMCPClienthandler evicts by MCP client after aMarkMCPPerUserHeaderCredentialsNeedsUpdateflip and afterReconcileMCPHeadersAfterMCPChange, so cached copies carrying pre-flip status or pre-reconcile rows are not served.revokesessions handler evicts by credential ID after deleting a header credential row directly through the configstore.sweepOrphanedCredentialsdocuments why it does not evict: the cache can never hold an orphaned row because theByModeSQL query filters them out, and the reconcile paths that flip rows toorphanedalready evict at flip time.Type of change
Affected areas
How to test
Key scenarios covered by the new test suite:
GetCredentialByModecovering: second call served from cache, isolated deep copy returned to callers, upsert evicts so the next read sees new values, delete evicts so the next read returns not-found,needs_updaterows cached with status, negative results not cached, and client-scoped eviction covering admin-mode bindings.Breaking changes
MCPOauthTokenCacheManageris renamed toMCPCredentialCacheManagerand gains two new methods (EvictMCPHeaderCredentialCacheByID,EvictMCPHeaderCredentialCacheByMCPClient). Any server implementation that embeds or implementsMCPOauthTokenCacheManagermust be updated to implement the fullMCPCredentialCacheManagerinterface and use the new name.Related issues
Security considerations
The cache is bounded (default 4096 entries) and uses LRU eviction, so a caller presenting arbitrary session-mode identity strings cannot cause unbounded memory growth. Cached credential values are never shared between callers: every hit returns a deep copy, preventing one caller from reading or mutating another caller's credential headers.
Checklist
docs/contributing/README.mdand followed the guidelines