feat: add virtualKeysByID secondary index and cache signing key + VK lookups on /mcp JWT auth path - #4783
Merged
Pratham-Mishra04 merged 1 commit intoJun 30, 2026
Conversation
This was referenced Jun 29, 2026
|
|
18 tasks
Contributor
|
Caution Review failedThe pull request is closed. ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (10)
📝 WalkthroughSummary by CodeRabbit
WalkthroughAdds ID-based virtual-key lookup and OAuth2 signing-key caching across the governance store and MCP HTTP handlers, updates JWT verification to accept a pre-fetched signing key, and changes three Anthropic stream tests to use exported stream-state helpers. ChangesMCP JWT and virtual-key caching
Anthropic stream-state test updates
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Suggested reviewers
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Warning There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure. 🔧 golangci-lint (2.12.2)level=error msg="[linters_context] typechecking error: pattern ./...: directory prefix . does not contain main module or its selected dependencies" Comment |
18 tasks
Collaborator
Author
18 tasks
Contributor
Confidence Score: 5/5This looks safe to merge.
Important Files Changed
Reviews (6): Last reviewed commit: "fix: removes db ops from mcp oauth paths" | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Pratham-Mishra04
force-pushed
the
06-29-fix_removes_db_ops_from_mcp_oauth_paths
branch
from
June 30, 2026 07:53
b281729 to
fa29eb3
Compare
Pratham-Mishra04
force-pushed
the
06-29-feat_adds_pagination_to_mcp_oauth_grants_table
branch
2 times, most recently
from
June 30, 2026 07:56
df6e27c to
6ea5663
Compare
Pratham-Mishra04
force-pushed
the
06-29-fix_removes_db_ops_from_mcp_oauth_paths
branch
from
June 30, 2026 07:56
fa29eb3 to
0609b71
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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 `@plugins/governance/store.go`:
- Around line 894-898: The virtual-key mutation path in
LocalGovernanceStore.storeVirtualKey is not fail-closed because the by-ID cache
can briefly retain stale auth data while the primary virtualKeys map changes.
Update the mutation flow to clear or evict the virtualKeysByID entry first
during updates/deletes/rebuilds, then write the new state so GetVirtualKeyByID
cannot serve the old active key during concurrent JWT auth. Review the other
virtual-key mutation helpers in the governance store that feed this index and
apply the same ordering consistently.
In `@transports/bifrost-http/handlers/mcpoauth2jwt.go`:
- Around line 31-36: The JWT key caching in mcpoauth2jwt.go is storing an
interior pointer from the parsed private key via parseRSAPrivateKeyPEM, which
unnecessarily keeps the private key allocation alive. Update the key-loading
logic around the public key cache to parse the standalone rsa.PublicKey from
signingKey.PublicKeyPEM instead of using &privKey.PublicKey, then store that
parsed public key in mcpJWTPublicKeys so verification only retains the public
material.
🪄 Autofix (Beta)
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: 02bcd1d2-84d0-4f78-bad2-4f28393bd020
📒 Files selected for processing (6)
plugins/governance/store.gotransports/bifrost-http/handlers/mcpoauth2issuance_test.gotransports/bifrost-http/handlers/mcpoauth2jwt.gotransports/bifrost-http/handlers/mcpoauth2jwt_test.gotransports/bifrost-http/handlers/mcpserver.gotransports/bifrost-http/server/server.go
danpiths
reviewed
Jun 30, 2026
Pratham-Mishra04
force-pushed
the
06-29-feat_adds_pagination_to_mcp_oauth_grants_table
branch
from
June 30, 2026 11:43
6ea5663 to
afb130d
Compare
Pratham-Mishra04
force-pushed
the
06-29-fix_removes_db_ops_from_mcp_oauth_paths
branch
from
June 30, 2026 11:43
0609b71 to
5f80820
Compare
Pratham-Mishra04
force-pushed
the
06-29-fix_removes_db_ops_from_mcp_oauth_paths
branch
from
June 30, 2026 13:44
5f80820 to
fabed5e
Compare
Pratham-Mishra04
force-pushed
the
06-29-feat_adds_pagination_to_mcp_oauth_grants_table
branch
from
June 30, 2026 13:44
afb130d to
375e4e3
Compare
Collaborator
Author
Merge activity
|
Pratham-Mishra04
changed the base branch from
06-29-feat_adds_pagination_to_mcp_oauth_grants_table
to
graphite-base/4783
June 30, 2026 14:43
Pratham-Mishra04
dismissed
coderabbitai[bot]’s stale review
June 30, 2026 14:46
The base branch was changed.
Pratham-Mishra04
force-pushed
the
06-29-fix_removes_db_ops_from_mcp_oauth_paths
branch
from
June 30, 2026 14:46
fabed5e to
01f6a6a
Compare
Pratham-Mishra04
deleted the
06-29-fix_removes_db_ops_from_mcp_oauth_paths
branch
June 30, 2026 14:48
18 tasks
akshaydeo
pushed a commit
that referenced
this pull request
Jul 1, 2026
…K lookups on `/mcp` JWT auth path (#4783) ## Summary The `/mcp` JWT authentication path previously performed a database read on every request to fetch the OAuth2 signing key and resolve virtual keys by ID. This PR eliminates those per-request DB reads by introducing a process-lifetime signing key cache on `MCPServerHandler` and a secondary in-memory index on `LocalGovernanceStore` that enables O(1) virtual key lookups by row ID. ## Changes - **Secondary VK index by ID**: Added `virtualKeysByID sync.Map` to `LocalGovernanceStore` as a secondary index over `virtualKeys`, keyed by row ID instead of VK value. Introduced `storeVirtualKey`, `deleteVirtualKeyByValue`, and `GetVirtualKeyByID` to keep both maps in lock-step. All existing write paths (`CreateVirtualKeyInMemory`, `UpdateVirtualKeyInMemory`, `DeleteVirtualKeyInMemory`, `rebuildInMemoryStructures`, and reference-update helpers) now route through these helpers instead of writing `virtualKeys` directly. - **Process-lifetime signing key cache**: Added an `atomic.Pointer[tables.OAuth2SigningKey]` field (`signingKey`) to `MCPServerHandler` with a `cachedSigningKey` loader that reads from the config store once and reuses the result. The signing key is created via an idempotent insert and never rotated, making a process-lifetime cache safe. The cache is warmed at handler construction when OAuth discovery is enabled, replacing the equivalent warm-up that previously lived in `Bootstrap`. - **RSA public key cache**: Added a package-level `mcpJWTPublicKeys sync.Map` that caches parsed `*rsa.PublicKey` values keyed by public-key PEM content, so JWT verification skips PEM parsing on every request. `verifyMCPJWT` now accepts the signing key as a parameter rather than fetching it internally, separating key loading from verification. - **`VirtualKeyCache` interface and `getVirtualKeyByID` helper**: Introduced a `VirtualKeyCache` interface (`GetVirtualKeyByID`) satisfied by `LocalGovernanceStore`. `MCPServerHandler` holds an optional `vkCache` field; `getVirtualKeyByID` checks the cache first and falls back to the config store on a miss or when no cache is wired. Both the vk-mode and user-mode JWT auth paths now use this helper instead of calling the config store directly. - **`NewMCPServerHandler` signature**: Added a `vkCache VirtualKeyCache` parameter. The server wires the governance store's in-memory cache at startup if available, with a type-assertion guard so the dependency remains optional. - **Test updates**: Updated `verifyMCPJWT` call sites to pass the signing key explicitly. Split the config-fault test into `TestVerifyMCPJWT_NilSigningKeyNotLabeledInvalidToken` (nil key → `verifyMCPJWT`) and `TestCachedSigningKey_ConfigFaults` (nil store / load error → `cachedSigningKey`). Added `nbf` claims to rejection test cases that were missing them. ## Type of change - [ ] Bug fix - [x] Feature - [x] Refactor - [ ] Documentation - [ ] Chore/CI ## Affected areas - [ ] Core (Go) - [x] Transports (HTTP) - [ ] Providers/Integrations - [x] Plugins - [ ] UI (React) - [ ] Docs ## How to test ```sh go test ./plugins/governance/... ./transports/bifrost-http/... ``` Validate that: - `/mcp` requests authenticated via JWT succeed without triggering DB reads for the signing key or virtual key on subsequent calls. - A virtual key created after the governance store last refreshed is still resolved correctly (config store fallback path). - Inactive virtual keys are still rejected. - JWT rejection cases (wrong algorithm, wrong audience, nil signing key) still return appropriate errors and are never labeled as `invalid token`. ## Breaking changes - [x] Yes `NewMCPServerHandler` now requires a `vkCache VirtualKeyCache` parameter (pass `nil` to preserve previous behaviour with config-store-only lookups). ## Security considerations - The signing key cache is keyed by public-key PEM content, not `kid`, preventing a rotated key with a reused `kid` from aliasing the wrong cached key. - `verifyMCPJWT` no longer reads the signing key internally; the caller (`getMCPServerForRequest`) owns loading it and logs infrastructure faults distinctly so they are never surfaced to clients as token-validation errors. - The `VirtualKeyCache` interface is read-only; no write path is exposed to the transport layer. ## Checklist - [x] I read `docs/contributing/README.md` and followed the guidelines - [x] I added/updated tests where appropriate - [ ] I updated documentation where needed - [x] I verified builds succeed (Go and UI) - [ ] I verified the CI pipeline passes locally if applicable
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
3 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
The
/mcpJWT authentication path previously performed a database read on every request to fetch the OAuth2 signing key and resolve virtual keys by ID. This PR eliminates those per-request DB reads by introducing a process-lifetime signing key cache onMCPServerHandlerand a secondary in-memory index onLocalGovernanceStorethat enables O(1) virtual key lookups by row ID.Changes
Secondary VK index by ID: Added
virtualKeysByID sync.MaptoLocalGovernanceStoreas a secondary index overvirtualKeys, keyed by row ID instead of VK value. IntroducedstoreVirtualKey,deleteVirtualKeyByValue, andGetVirtualKeyByIDto keep both maps in lock-step. All existing write paths (CreateVirtualKeyInMemory,UpdateVirtualKeyInMemory,DeleteVirtualKeyInMemory,rebuildInMemoryStructures, and reference-update helpers) now route through these helpers instead of writingvirtualKeysdirectly.Process-lifetime signing key cache: Added an
atomic.Pointer[tables.OAuth2SigningKey]field (signingKey) toMCPServerHandlerwith acachedSigningKeyloader that reads from the config store once and reuses the result. The signing key is created via an idempotent insert and never rotated, making a process-lifetime cache safe. The cache is warmed at handler construction when OAuth discovery is enabled, replacing the equivalent warm-up that previously lived inBootstrap.RSA public key cache: Added a package-level
mcpJWTPublicKeys sync.Mapthat caches parsed*rsa.PublicKeyvalues keyed by public-key PEM content, so JWT verification skips PEM parsing on every request.verifyMCPJWTnow accepts the signing key as a parameter rather than fetching it internally, separating key loading from verification.VirtualKeyCacheinterface andgetVirtualKeyByIDhelper: Introduced aVirtualKeyCacheinterface (GetVirtualKeyByID) satisfied byLocalGovernanceStore.MCPServerHandlerholds an optionalvkCachefield;getVirtualKeyByIDchecks the cache first and falls back to the config store on a miss or when no cache is wired. Both the vk-mode and user-mode JWT auth paths now use this helper instead of calling the config store directly.NewMCPServerHandlersignature: Added avkCache VirtualKeyCacheparameter. The server wires the governance store's in-memory cache at startup if available, with a type-assertion guard so the dependency remains optional.Test updates: Updated
verifyMCPJWTcall sites to pass the signing key explicitly. Split the config-fault test intoTestVerifyMCPJWT_NilSigningKeyNotLabeledInvalidToken(nil key →verifyMCPJWT) andTestCachedSigningKey_ConfigFaults(nil store / load error →cachedSigningKey). Addednbfclaims to rejection test cases that were missing them.Type of change
Affected areas
How to test
go test ./plugins/governance/... ./transports/bifrost-http/...Validate that:
/mcprequests authenticated via JWT succeed without triggering DB reads for the signing key or virtual key on subsequent calls.invalid token.Breaking changes
NewMCPServerHandlernow requires avkCache VirtualKeyCacheparameter (passnilto preserve previous behaviour with config-store-only lookups).Security considerations
kid, preventing a rotated key with a reusedkidfrom aliasing the wrong cached key.verifyMCPJWTno longer reads the signing key internally; the caller (getMCPServerForRequest) owns loading it and logs infrastructure faults distinctly so they are never surfaced to clients as token-validation errors.VirtualKeyCacheinterface is read-only; no write path is exposed to the transport layer.Checklist
docs/contributing/README.mdand followed the guidelines