refactor: unify shared and per-user OAuth refresh under a single RefreshAccessToken(tokenID) path, dropping oauth_configs.token_id FK shortcut - #5710
Merged
Pratham-Mishra04 merged 1 commit intoAug 8, 2026
Conversation
|
|
Collaborator
Author
This was referenced Jul 30, 2026
refactor: migrate OAuth flow state/PKCE fields from
oauth_configs to new mcp_oauth_flows table
#5709
Merged
Merged
Merged
Merged
Merged
Contributor
📝 WalkthroughSummary by CodeRabbit
WalkthroughOAuth credential ownership moves from ChangesUnified OAuth token lifecycle
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant OAuth2Service
participant ConfigStore
participant OAuthProvider
OAuth2Service->>ConfigStore: Load token by token ID
OAuth2Service->>ConfigStore: Load template OAuth config
OAuth2Service->>OAuthProvider: Refresh access token
OAuthProvider-->>OAuth2Service: Return refreshed token or error
OAuth2Service->>ConfigStore: Mark token needs_reauth on permanent rejection
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 Jul 30, 2026
This was referenced Aug 6, 2026
![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: 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/configstore/tables/mcpoauth2.go`:
- Around line 22-25: Expand the token lookup documentation near the existing
holder-token comment to specify shared lookups use (oauth_config_id, auth_mode =
"shared") but are not unique, while per-identity lookups additionally require
mcp_client_id and the matching user_id, virtual_key_id, or session_id, plus
status = "active". Document that refresh operations use the token row ID and
FlowMode = "admin" produces AuthMode = "shared".
🪄 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: 04f5c7fb-454d-4e26-9cc8-f71802e0a681
📒 Files selected for processing (14)
core/mcp/credstore/per_user_oauth_test.gocore/schemas/oauth.goframework/configstore/migrations.goframework/configstore/migrations_perf_test.goframework/configstore/rdb.goframework/configstore/rdb_mcp_sessions_test.goframework/configstore/rdb_oauth2_test.goframework/configstore/store.goframework/configstore/tables/mcpoauth2.goframework/oauth2/main.goframework/oauth2/sync.goframework/oauth2/sync_test.gotransports/bifrost-http/handlers/mcpoauth2.gotransports/bifrost-http/lib/config_test.go
🚧 Files skipped from review as they are similar to previous changes (13)
- transports/bifrost-http/handlers/mcpoauth2.go
- framework/configstore/rdb_mcp_sessions_test.go
- framework/configstore/migrations_perf_test.go
- core/mcp/credstore/per_user_oauth_test.go
- framework/oauth2/sync.go
- framework/configstore/store.go
- framework/configstore/rdb_oauth2_test.go
- transports/bifrost-http/lib/config_test.go
- framework/oauth2/sync_test.go
- framework/oauth2/main.go
- framework/configstore/migrations.go
- framework/configstore/rdb.go
- core/schemas/oauth.go
Pratham-Mishra04
force-pushed
the
07-29-feat_unify_oauth_token_refresh_into_refreshaccesstoken
branch
from
August 6, 2026 21:53
53a270a to
c57aa1e
Compare
Pratham-Mishra04
force-pushed
the
07-29-feat_unify_oauth_flow_storage_into_tablemcpoauthflow
branch
from
August 6, 2026 21:53
404348f to
83be929
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. |
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (2)
framework/oauth2/main.go (2)
225-229: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRename
refreshAccessTokenLockedto match its new synchronization model.The
Lockedsuffix normally means the caller holds a mutex. This function now runs insidesingleflight.Group.DoChan, and no lock is held. ConsiderrefreshAccessTokenOnceordoRefreshAccessTokenso future readers do not assume mutual exclusion beyond per-token deduplication.🤖 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/oauth2/main.go` around lines 225 - 229, Rename OAuth2Provider.refreshAccessTokenLocked to reflect singleflight-based per-token deduplication rather than mutex protection, using a name such as refreshAccessTokenOnce or doRefreshAccessToken. Update every declaration, invocation, and related comment consistently while preserving the existing DoChan synchronization behavior.
325-333: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winDistinguish a store failure from an invalid token.
Line 326 collapses a
GetSharedOauthTokenByConfigIDerror into(false, nil). A transient database failure then reports the credential as invalid, and callers cannot retry or alert. Return the error for the failure case, and keep(false, nil)for the missing-token case.As per coding guidelines: "explicit error handling and wrapping".
♻️ Proposed change
token, err := p.configStore.GetSharedOauthTokenByConfigID(ctx, oauthConfigID) - if err != nil || token == nil { + if err != nil { + return false, fmt.Errorf("failed to load shared oauth token: %w", err) + } + if token == nil { return false, nil }🤖 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/oauth2/main.go` around lines 325 - 333, Update the token lookup handling in the surrounding OAuth validation method: return a wrapped error when GetSharedOauthTokenByConfigID fails, while continuing to return (false, nil) when the token is absent. Preserve the existing status-based validity checks for present tokens.Source: Coding guidelines
🤖 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/oauth2/main.go`:
- Around line 225-229: Rename OAuth2Provider.refreshAccessTokenLocked to reflect
singleflight-based per-token deduplication rather than mutex protection, using a
name such as refreshAccessTokenOnce or doRefreshAccessToken. Update every
declaration, invocation, and related comment consistently while preserving the
existing DoChan synchronization behavior.
- Around line 325-333: Update the token lookup handling in the surrounding OAuth
validation method: return a wrapped error when GetSharedOauthTokenByConfigID
fails, while continuing to return (false, nil) when the token is absent.
Preserve the existing status-based validity checks for present tokens.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: d5344fd8-4d89-4784-9f00-3522241f193d
📒 Files selected for processing (14)
core/mcp/credstore/per_user_oauth_test.gocore/schemas/oauth.goframework/configstore/migrations.goframework/configstore/migrations_perf_test.goframework/configstore/rdb.goframework/configstore/rdb_mcp_sessions_test.goframework/configstore/rdb_oauth2_test.goframework/configstore/store.goframework/configstore/tables/mcpoauth2.goframework/oauth2/main.goframework/oauth2/sync.goframework/oauth2/sync_test.gotransports/bifrost-http/handlers/mcpoauth2.gotransports/bifrost-http/lib/config_test.go
🚧 Files skipped from review as they are similar to previous changes (12)
- core/schemas/oauth.go
- framework/configstore/rdb_mcp_sessions_test.go
- framework/oauth2/sync.go
- core/mcp/credstore/per_user_oauth_test.go
- transports/bifrost-http/handlers/mcpoauth2.go
- framework/configstore/rdb_oauth2_test.go
- framework/configstore/migrations_perf_test.go
- framework/configstore/rdb.go
- framework/oauth2/sync_test.go
- transports/bifrost-http/lib/config_test.go
- framework/configstore/store.go
- framework/configstore/migrations.go
Pratham-Mishra04
force-pushed
the
07-29-feat_unify_oauth_flow_storage_into_tablemcpoauthflow
branch
from
August 8, 2026 08:43
83be929 to
f2c9723
Compare
Pratham-Mishra04
force-pushed
the
07-29-feat_unify_oauth_token_refresh_into_refreshaccesstoken
branch
from
August 8, 2026 08:43
c57aa1e to
c50fb76
Compare
Collaborator
Author
Merge activity
|
Pratham-Mishra04
changed the base branch from
07-29-feat_unify_oauth_flow_storage_into_tablemcpoauthflow
to
graphite-base/5710
August 8, 2026 09:03
Pratham-Mishra04
dismissed
coderabbitai[bot]’s stale review
August 8, 2026 09:05
The base branch was changed.
…ntity oauth token lookup contract
Pratham-Mishra04
force-pushed
the
07-29-feat_unify_oauth_token_refresh_into_refreshaccesstoken
branch
from
August 8, 2026 09:05
c50fb76 to
f30b500
Compare
Pratham-Mishra04
deleted the
07-29-feat_unify_oauth_token_refresh_into_refreshaccesstoken
branch
August 8, 2026 09:07
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
This PR removes the
token_idFK shortcut column fromoauth_configsand unifies the two separate OAuth token refresh paths (one keyed byoauth_config_idfor shared tokens, one keyed by token ID for per-user tokens) into a singleRefreshAccessToken(ctx, tokenID)method. Credential health tracking moves entirely onto the token row's ownStatusfield (active/needs_reauth), replacing the previous pattern of flippingoauth_configs.statustoexpiredorrevokedon permanent refresh failures.Changes
TableOauthConfig.TokenID: The FK shortcut that pointed from an oauth config to its single shared-mode token row is dropped. All callers that previously readoauthConfig.TokenIDnow resolve the token viaGetSharedOauthTokenByConfigID(ctx, oauthConfigID), which queriesmcp_oauth_tokenson(oauth_config_id, auth_mode='shared').RefreshAccessToken:RefreshUserAccessToken(per-identity, keyed by token ID) and the oldRefreshAccessToken(shared, keyed byoauth_config_id) are merged into one method keyed by token ID. BothGetAccessTokenandGetUserAccessTokenByModefunnel their lazy pre-flight refresh through this single path.invalid_grant,unauthorized_client) now flipmcp_oauth_tokens.statustoneeds_reauthviaMarkOauthUserTokenNeedsReauthByID, which is no longer scoped away fromauth_mode='shared'. Theoauth_configs.statuscolumn is now a one-time bootstrap lifecycle field only (pending/authorized/failed) and is never written by the refresh path.GetExpiringOauthTokensfilter change: The query previously excluded tokens whose owningoauth_confighad a terminal status (expired/revoked). It now filters directly onmcp_oauth_tokens.status = 'active', and the join tooauth_configsfor the enabled-client check usesoauth_configs.id = mcp_oauth_tokens.oauth_config_idinstead of the retiredtoken_idcolumn.TokenRefreshWorkersimplification: The worker no longer looks up the owningoauth_configfor each expiring token before calling refresh. It callsRefreshAccessToken(ctx, token.ID)directly; permanent-failure handling is entirely insideRefreshAccessToken.GetOauthUserSessionByIDflow-mode filter: Addedflow_mode IN (perUserOauthFlowModes)to prevent an admin-mode flow row from being reachable through the per-user-facing ID lookup endpoint.migrationDropOauthConfigTokenIDColumnmigration dropsoauth_configs.token_id. The existingmigrationMergeOauthTokenTablesbackfill is guarded with aHasColumncheck so it skips thetoken_id-referencing SQL on fresh installs where the column never existed.CompleteOAuthFlowexpiry branch: Changed the terminal bootstrap status written on flow expiry from the since-retired"expired"to"failed", consistent with the other bootstrap-failure branch.GetOauthConfigByTokenIDremoved: No longer needed; replaced byGetSharedOauthTokenByConfigID.Type of change
Affected areas
How to test
go test ./core/mcp/credstore/... ./framework/configstore/... ./framework/oauth2/... ./transports/bifrost-http/...status = needs_reauthon themcp_oauth_tokensrow and thatoauth_configs.statusremainsauthorized.TokenRefreshWorkerdoes not retry a token already markedneeds_reauth.token_idcolumn) runs all migrations without error.token_idpopulated correctly backfillsmcp_oauth_tokens.oauth_config_idand then drops the column.Breaking changes
GetOauthConfigByTokenIDis removed from theConfigStoreinterface. Any external implementation ofConfigStoremust addGetSharedOauthTokenByConfigIDand removeGetOauthConfigByTokenID. TheOAuth2Providerinterface losesRefreshAccessToken(ctx, oauthConfigID)andRefreshUserAccessToken(ctx, tokenID)and gains a singleRefreshAccessToken(ctx, tokenID). Theoauth_configs.token_idcolumn is dropped by migration; any raw SQL or tooling that references it will need updating.Related issues
N/A
Security considerations
GetOauthUserSessionByIDnow filters byflow_mode IN (perUserOauthFlowModes), preventing an admin-mode flow row from being fetched through a per-user-facing endpoint by ID.MarkOauthUserTokenNeedsReauthByIDis intentionally not scoped byauth_modebecause the token ID it receives always comes from a trusted internal lookup, never an arbitrary caller-supplied value.Checklist
docs/contributing/README.mdand followed the guidelines