feat: mcp per user oauth flow refactor - #3565
Merged
Merged
Conversation
Contributor
📝 WalkthroughSummary by CodeRabbit
WalkthroughRefactors per-user OAuth to a mode-driven model (user/vk/session), migrates DB schema and store APIs, updates provider initiation/completion/refresh to bind identity by mode, adds a sweep worker, removes legacy per-user OAuth handlers, and adds MCP sessions HTTP API and UI. ChangesMode-driven Per-User OAuth and Sessions Management
Estimated code review effort 🎯 4 (Complex) | ⏱️ ~60 minutes Suggested reviewers
Poem
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
|
Collaborator
Author
This stack of pull requests is managed by Graphite. Learn more about stacking. |
|
|
Pratham-Mishra04
force-pushed
the
05-18-refactor_mcp_per_user_oauth_flow_refactor
branch
2 times, most recently
from
May 18, 2026 13:14
cc1d3ec to
05cfa3e
Compare
Pratham-Mishra04
marked this pull request as ready for review
May 18, 2026 13:15
Contributor
Confidence Score: 3/5The OAuth refactor and new sessions API are well-structured, but the tool-key re-prefixing logic in the MCP client update handler corrupts persisted tool maps for clients with hyphenated names on every update. The new handler code in mcp.go uses strings.Cut(oldKey, "-") splitting at the first hyphen rather than at the full old client-name boundary. For a client named github-mcp, every updateMCPClient call rewrites github-mcp-list_repos as github-mcp-mcp-list_repos, silently corrupting the DB. The same block copies DiscoveredToolNameMapping verbatim, leaving stale keys after a rename that break tool dispatch on restart. transports/bifrost-http/handlers/mcp.go lines 939-956 (tool-key re-prefixing block) and framework/configstore/migrations.go (BeforeSave AuthMode guard promised in comment but absent in tables/oauth.go). Important Files Changed
Reviews (18): Last reviewed commit: "refactor: mcp per user oauth flow refact..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
![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: 13
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
core/mcp/utils/utils.go (1)
21-50:⚠️ Potential issue | 🟠 Major | ⚡ Quick winDon't fail user-mode OAuth before creating the pending flow.
InitiateUserOAuthFlownow explicitly supportsMCPAuthModeUserwith noUserIDin context yet, but thisidentity == ""path returns before that branch can run. That breaks externally initiated/user-not-yet-bound MCP auth by turning it into an immediate error instead of creating the pending flow row that gets stamped on callback. Onlyvk/sessionmodes need a non-empty identity up front; user mode should skip lookup and fall through to flow initiation.🤖 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 `@core/mcp/utils/utils.go` around lines 21 - 50, The current early return when identity == "" blocks user-mode flows; change the logic in the block that calls identityForMCPAuthMode(ctx, mode) so it only rejects missing identity for modes that require it (e.g., VK/session modes) and does not return for MCPAuthModeUser; i.e., only error when identity == "" && mode != MCPAuthModeUser (or use an explicit list of modes that require identity). Keep the rest of the flow intact so InitiateUserOAuthFlow can be called for user-mode even when UserID/identity is not yet present; continue to call oauth2Provider.GetUserAccessTokenByMode, BuildRedirectURIFromContext, and InitiateUserOAuthFlow as before.transports/bifrost-http/handlers/mcpserver.go (1)
72-76:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winDuplicate tool filter registration.
The same filter is registered twice on
globalMCPServer(lines 72 and 74-75). This appears to be a copy-paste error that could cause the filter to execute twice pertools/listrequest.Proposed fix
// Register per-request tool filter so x-bf-mcp-include-clients and x-bf-mcp-include-tools are respected on tools/list server.WithToolFilter(handler.makeIncludeClientsFilter())(handler.globalMCPServer) - // Register per-request tool filter so x-bf-mcp-include-clients and x-bf-mcp-include-tools are respected on tools/list - server.WithToolFilter(handler.makeIncludeClientsFilter())(handler.globalMCPServer) - if err := handler.SyncAllMCPServers(ctx); err != 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 `@transports/bifrost-http/handlers/mcpserver.go` around lines 72 - 76, The same tool filter is registered twice on globalMCPServer via server.WithToolFilter(handler.makeIncludeClientsFilter()), causing it to run twice; remove the duplicate call so only a single registration of server.WithToolFilter(handler.makeIncludeClientsFilter()) on handler.globalMCPServer remains (or replace the second call with the correct filter if it was meant to be a different one), ensuring the per-request tool filter is registered exactly once.framework/configstore/rdb.go (1)
4711-4750:⚠️ Potential issue | 🟠 Major | ⚡ Quick winHandle concurrent first-time token writes as a real upsert.
This still falls back to
SELECT-then-CREATEon the miss path. Two same-identity callbacks can both missexisting, and one will fail on the unique index instead of replacing the row, which makes re-auth flaky under retries/concurrent callbacks. Use a conflict-aware insert or retry the create path by re-reading andSaveing on unique-violation.🤖 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/configstore/rdb.go` around lines 4711 - 4750, The current Transaction in s.DB().WithContext(ctx).Transaction does a SELECT-then-CREATE and can race: when lookupErr is gorm.ErrRecordNotFound two concurrent callers may both try tx.Create(token) and one will hit a unique-constraint error; change the miss path to do a conflict-aware upsert or retry-on-unique-violation. Concretely, replace the tx.Create(token) branch with either a GORM OnConflict clause (e.g., use tx.Clauses(clause.OnConflict{UpdateAll: true}).Create(token)) to perform an atomic upsert, or catch the unique-constraint error from tx.Create, re-query the existing row via dbForUpdate(tx) (same WHERE logic using token.UserID/VirtualKeyID/SessionID) and then set token.ID/CreatedAt/LastRefreshedAt and call tx.Save(token) so the second writer merges instead of failing.
🧹 Nitpick comments (4)
ui/app/_fallbacks/enterprise/components/mcp-sessions/dacFilterPill.tsx (1)
5-7: ⚡ Quick winRename this file to PascalCase to match UI component conventions.
The component export is correctly PascalCase, but the filename
dacFilterPill.tsxis not. Please rename it toDACFilterPill.tsx(and update imports accordingly) for consistency with the TSX naming rule.As per coding guidelines:
ui/**/*.tsx: React component files must use PascalCase for component exports and filenames.🤖 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 `@ui/app/_fallbacks/enterprise/components/mcp-sessions/dacFilterPill.tsx` around lines 5 - 7, The file name needs to be changed to PascalCase to match the component export; rename the file from dacFilterPill.tsx to DACFilterPill.tsx and update all imports that reference this file to the new filename (search for imports referencing dacFilterPill and update them). Ensure the exported component function DACFilterPill remains unchanged and that build/test imports compile after the rename.ui/app/_fallbacks/enterprise/components/mcp-sessions/grantedViaVKChips.tsx (1)
11-13: ⚡ Quick winUse PascalCase for this TSX filename as well.
grantedViaVKChips.tsxshould be renamed toGrantedViaVKChips.tsxto stay consistent with the component file naming convention.As per coding guidelines:
ui/**/*.tsx: React component files must use PascalCase for component exports and filenames.🤖 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 `@ui/app/_fallbacks/enterprise/components/mcp-sessions/grantedViaVKChips.tsx` around lines 11 - 13, Rename the TSX file from grantedViaVKChips.tsx to GrantedViaVKChips.tsx to match the component export GrantedViaVKChips and the project's PascalCase convention; perform a git mv to preserve history, then update any imports that reference the old filename (search for "grantedViaVKChips" usages) so they point to "GrantedViaVKChips" and run the build/typecheck to ensure no broken imports remain.ui/lib/store/apis/mcpSessionsApi.ts (1)
22-25: ⚡ Quick winConsider optimistic cache patching for
revokeMCPSessionto match patterns in this directory, but notereauthMCPSessioncannot be optimistically patched.The pattern of using
onQueryStarted + updateQueryDatafor mutations inui/lib/store/apis/exists (seedeleteRoutingRule), but is inconsistently applied across the codebase—deleteLogsanddeleteMCPLogsalso use onlyinvalidatesTags.For
revokeMCPSession, implementing optimistic cache patching would improve UX by removing the deleted session fromgetMCPSessionscache immediately, avoiding a full refetch.However,
reauthMCPSessionreturnsMCPSessionReauthResponse(an OAuth URL), not updated session data, so it cannot be optimistically patched. TheinvalidatesTagsapproach is appropriate for this mutation.Also applies to: lines 28–31 (revokeMCPSession only; reauthMCPSession is unsuitable for optimistic patching)
🤖 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 `@ui/lib/store/apis/mcpSessionsApi.ts` around lines 22 - 25, The review notes that revokeMCPSession should follow the optimistic cache-patching pattern used elsewhere (use onQueryStarted + updateQueryData) to remove the deleted session from the getMCPSessions cache immediately, while reauthMCPSession should remain using invalidatesTags because it returns an OAuth URL rather than updated session data; therefore, modify the revokeMCPSession mutation to implement onQueryStarted that patches getMCPSessions (remove the session by id and roll back on error), referencing the revokeMCPSession mutation and getMCPSessions cache-updater, and leave reauthMCPSession as-is (invalidatesTags only).transports/bifrost-http/lib/config.go (1)
3080-3085: 💤 Low valueConsider making the OAuth sweep retention period configurable.
The 30-day retention period is currently hardcoded. While this is a reasonable default and clearly documented, consider allowing configuration via
FrameworkConfigor an environment variable in a future iteration to give operators control over orphan cleanup policy.Current implementation is correct and safe for this refactor.
🤖 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 `@transports/bifrost-http/lib/config.go` around lines 3080 - 3085, The hardcoded 30-day retention in the call to oauth2.NewPerUserOAuthSweepWorker should be made configurable: add a retention field (e.g., OAuthSweepRetention time.Duration) to the FrameworkConfig (or read from an ENV var with a 30*24*time.Hour default), plumb that value into config.OAuthSweepWorker creation instead of the literal 30*24*time.Hour, and ensure any config initialization/validation sets the default when unset; update references to NewPerUserOAuthSweepWorker and config.OAuthSweepWorker initialization to use the new configurable retention.
🤖 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/migrations.go`:
- Around line 7646-7651: The partial unique indexes
idx_oauth_user_tokens_user_mcp and idx_oauth_user_tokens_vk_mcp on table
oauth_user_tokens currently only exclude NULLs, but the backfill treats empty
string ('') as unset so rows with user_id = '' or virtual_key_id = '' will still
be indexed and can cause uniqueness failures; update the CREATE UNIQUE INDEX
predicates to also exclude empty strings (e.g. add "AND user_id <> ''" for
idx_oauth_user_tokens_user_mcp and "AND virtual_key_id <> ''" for
idx_oauth_user_tokens_vk_mcp) so only meaningful identities participate in the
unique index.
- Around line 7683-7690: The CASE updating oauth_user_sessions.flow_mode
incorrectly maps legacy session-backed rows to 'none'; modify the SQL in the
migration (the tx.Exec(...) block) to include a WHEN branch that sets 'session'
when session_token_hash IS NOT NULL AND session_token_hash != '' (place this
check before the ELSE), so existing session-mode records are backfilled
correctly prior to migrationReplaceOauthSessionTokenWithSessionID.
- Around line 7762-7784: After adding the SessionID column on
TableOauthUserSession, create the corresponding index by calling
mg.CreateIndex(&tables.TableOauthUserSession{}, "SessionID") (and check/return
any error) so the new column gets the same indexed lookup as before; place this
call immediately after the AddColumn block for SessionID and before finishing
the migration (mirroring how the oauth_user_tokens migration recreates its
index).
In `@framework/configstore/rdb.go`:
- Around line 4631-4652: The lookup
GetOauthUserSessionByModeIdentityAndMCPClient currently returns rows stuck in
status "claiming" and must ignore/expire stale claiming rows; update the query
in GetOauthUserSessionByModeIdentityAndMCPClient to exclude sessions whose
status == "claiming" unless their last-update timestamp is older than the
claiming timeout (use a shared ClaimExpiryDuration or config constant and
compare against updated_at/claimed_at), so only genuinely expired claiming rows
are returned; apply the identical change to the other analogous lookup (the
block at the other occurrence around lines 5038-5046) and ensure any error
messages remain unchanged.
In `@framework/oauth2/sync.go`:
- Around line 182-184: The Stop method on PerUserOAuthSweepWorker closes
w.stopCh unguarded which can panic if Stop is called multiple times; make Stop
idempotent by adding and using a sync.Once (e.g., a field like stopOnce
sync.Once on PerUserOAuthSweepWorker) and call stopOnce.Do(func(){
close(w.stopCh) }) inside the Stop method (retain the existing logger nil-check
after the guarded close). This mirrors the pattern used in
framework/tracing/store.go where channel closes are protected by sync.Once.
In `@transports/bifrost-http/handlers/mcp_sessions.go`:
- Around line 376-380: BuildUpstreamAuthorizeURL can return errors for
stale/expired/not-pending flows but the handler currently maps all failures to
500; update the error handling in the mcp_sessions.go block that calls
h.store.OAuthProvider.BuildUpstreamAuthorizeURL(ctx, flowID) so that it checks
the returned error for the flow-stale conditions (e.g., compare against the
provider's sentinel errors like ErrFlowExpired, ErrFlowNotPending or inspect an
exported error code/type) and when matched return a 4xx response (410 Gone or
400 Bad Request) via SendError instead of 500; keep logging (logger.Error) but
include the error details and still return 500 only for unexpected/internal
errors.
- Around line 77-99: The code currently treats identity == "" as a signal to use
the admin (ListAll...) branch regardless of auth mode; change the guard so that
the admin branch is only taken when no auth mode is configured (mode == ""), and
when mode != "" but identity == "" the handler fails closed (return an
unauthorized/identity resolution error or empty result) instead of calling
ListAllOauthUserTokens/ListAllPendingOauthUserSessions; update the branch around
callerModeAndIdentity(bfCtx) (checking mode and identity) to call
ListOauthUserTokensByMode/ListOauthUserSessionsByMode when mode != "" and
identity present, only call the ListAll* functions when mode == "" and identity
== "", and return an error when mode != "" && identity == "".
In `@ui/app/workspace/config/views/mcpView.tsx`:
- Around line 279-283: The new AccordionTrigger element for the "Advanced
Settings" section is missing a stable test selector; add a data-testid attribute
to the AccordionTrigger (the JSX element with <AccordionTrigger> that is inside
the AccordionItem value="advanced-settings") using the 3-part convention, e.g.
data-testid="mcp-settings-advanced-trigger", so E2E tests can reliably target
this interactive control.
In `@ui/app/workspace/mcp-sessions/auth/page.tsx`:
- Around line 135-140: Add stable data-testid attributes to the new interactive
auth controls so E2E tests can select them: add data-testid="auth-button" on the
primary Button that calls handleAuthenticate (handles both Authenticate and
Re-authenticate states), add data-testid="sessions-tab-link" on the
SessionsTabLink component, and add data-testid="sign-in-button" on the sign-in
control (the button/link that triggers sign-in in the same file — also apply
identical data-testid additions for the duplicate auth/sign-in controls around
lines 244-271). Ensure the attributes are added directly to the React elements
(Button, SessionsTabLink, sign-in button) without changing their behavior.
- Around line 222-233: The formatExpiry helper can produce "in NaN minutes"
because new Date(iso).getTime() returns NaN for invalid strings; update
formatExpiry to validate the parsed time (e.g., const t = d.getTime() and if
(Number.isNaN(t) ) return iso) before computing diffMs and minutes so invalid
ISO inputs fall back to the original iso string; keep the existing
negative/difference logic intact and reference the formatExpiry function when
making the change.
In `@ui/app/workspace/mcp-sessions/oauth-callback/page.tsx`:
- Around line 52-54: The new interactive "Back to MCP sessions" control (the
Button component using asChild wrapping the Link with
to="/workspace/mcp-sessions") needs a data-testid added to satisfy E2E
conventions; add a data-testid prop (for example
data-testid="back-to-mcp-sessions") to the interactive element (either on the
Link or passed through Button when using asChild) so the selector targets this
new control (refer to the Button and Link instances in the
oauth-callback/page.tsx snippet).
In `@ui/app/workspace/mcp-sessions/views/sessionsTable.tsx`:
- Around line 71-82: Add stable data-testid attributes to the new MCP session
interactive elements so E2E can target them: add data-testid to the dialog
buttons rendered by AlertDialogCancel and AlertDialogAction (e.g.
"mcp-session-revoke-cancel" and "mcp-session-revoke-confirm") while keeping the
existing onClick handler confirmRevoke; add a stable data-testid on the row
actions trigger component used to open the menu (e.g.
"mcp-session-row-actions-<sessionId>" or similar unique id bound to the row) and
add data-testid attributes to the menu items for Revoke and Re-auth flows (e.g.
"mcp-session-revoke-menu-item" and "mcp-session-reauth-menu-item"); ensure the
ids include the session identifier where appropriate for uniqueness and attach
them directly to the rendered interactive elements (AlertDialogCancel,
AlertDialogAction, the row-actions trigger, and the corresponding menu item
components).
- Around line 301-334: Both formatRelativePast and formatAccessExpiry need to
guard against invalid timestamps because new Date(...).getTime() can be NaN (it
doesn’t throw); update formatRelativePast(iso: string) to compute const t = new
Date(iso).getTime() and if Number.isNaN(t) return iso, then use diffMs =
Date.now() - t; similarly in formatAccessExpiry(row: MCPSessionRow) compute
const t = new Date(row.expires_at).getTime() and if Number.isNaN(t) return
row.expires_at, then use diffMs = t - Date.now() and continue existing
logic—this prevents "NaNd" or "in NaN min" outputs.
---
Outside diff comments:
In `@core/mcp/utils/utils.go`:
- Around line 21-50: The current early return when identity == "" blocks
user-mode flows; change the logic in the block that calls
identityForMCPAuthMode(ctx, mode) so it only rejects missing identity for modes
that require it (e.g., VK/session modes) and does not return for
MCPAuthModeUser; i.e., only error when identity == "" && mode != MCPAuthModeUser
(or use an explicit list of modes that require identity). Keep the rest of the
flow intact so InitiateUserOAuthFlow can be called for user-mode even when
UserID/identity is not yet present; continue to call
oauth2Provider.GetUserAccessTokenByMode, BuildRedirectURIFromContext, and
InitiateUserOAuthFlow as before.
In `@framework/configstore/rdb.go`:
- Around line 4711-4750: The current Transaction in
s.DB().WithContext(ctx).Transaction does a SELECT-then-CREATE and can race: when
lookupErr is gorm.ErrRecordNotFound two concurrent callers may both try
tx.Create(token) and one will hit a unique-constraint error; change the miss
path to do a conflict-aware upsert or retry-on-unique-violation. Concretely,
replace the tx.Create(token) branch with either a GORM OnConflict clause (e.g.,
use tx.Clauses(clause.OnConflict{UpdateAll: true}).Create(token)) to perform an
atomic upsert, or catch the unique-constraint error from tx.Create, re-query the
existing row via dbForUpdate(tx) (same WHERE logic using
token.UserID/VirtualKeyID/SessionID) and then set
token.ID/CreatedAt/LastRefreshedAt and call tx.Save(token) so the second writer
merges instead of failing.
In `@transports/bifrost-http/handlers/mcpserver.go`:
- Around line 72-76: The same tool filter is registered twice on globalMCPServer
via server.WithToolFilter(handler.makeIncludeClientsFilter()), causing it to run
twice; remove the duplicate call so only a single registration of
server.WithToolFilter(handler.makeIncludeClientsFilter()) on
handler.globalMCPServer remains (or replace the second call with the correct
filter if it was meant to be a different one), ensuring the per-request tool
filter is registered exactly once.
---
Nitpick comments:
In `@transports/bifrost-http/lib/config.go`:
- Around line 3080-3085: The hardcoded 30-day retention in the call to
oauth2.NewPerUserOAuthSweepWorker should be made configurable: add a retention
field (e.g., OAuthSweepRetention time.Duration) to the FrameworkConfig (or read
from an ENV var with a 30*24*time.Hour default), plumb that value into
config.OAuthSweepWorker creation instead of the literal 30*24*time.Hour, and
ensure any config initialization/validation sets the default when unset; update
references to NewPerUserOAuthSweepWorker and config.OAuthSweepWorker
initialization to use the new configurable retention.
In `@ui/app/_fallbacks/enterprise/components/mcp-sessions/dacFilterPill.tsx`:
- Around line 5-7: The file name needs to be changed to PascalCase to match the
component export; rename the file from dacFilterPill.tsx to DACFilterPill.tsx
and update all imports that reference this file to the new filename (search for
imports referencing dacFilterPill and update them). Ensure the exported
component function DACFilterPill remains unchanged and that build/test imports
compile after the rename.
In `@ui/app/_fallbacks/enterprise/components/mcp-sessions/grantedViaVKChips.tsx`:
- Around line 11-13: Rename the TSX file from grantedViaVKChips.tsx to
GrantedViaVKChips.tsx to match the component export GrantedViaVKChips and the
project's PascalCase convention; perform a git mv to preserve history, then
update any imports that reference the old filename (search for
"grantedViaVKChips" usages) so they point to "GrantedViaVKChips" and run the
build/typecheck to ensure no broken imports remain.
In `@ui/lib/store/apis/mcpSessionsApi.ts`:
- Around line 22-25: The review notes that revokeMCPSession should follow the
optimistic cache-patching pattern used elsewhere (use onQueryStarted +
updateQueryData) to remove the deleted session from the getMCPSessions cache
immediately, while reauthMCPSession should remain using invalidatesTags because
it returns an OAuth URL rather than updated session data; therefore, modify the
revokeMCPSession mutation to implement onQueryStarted that patches
getMCPSessions (remove the session by id and roll back on error), referencing
the revokeMCPSession mutation and getMCPSessions cache-updater, and leave
reauthMCPSession as-is (invalidatesTags only).
🪄 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: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 8a9d4ac5-f8a4-4862-8b1d-426f7d7a8ce0
📒 Files selected for processing (42)
core/mcp/utils/utils.gocore/schemas/bifrost.gocore/schemas/context.gocore/schemas/oauth.goframework/configstore/clientconfig.goframework/configstore/migrations.goframework/configstore/rdb.goframework/configstore/rdb_test.goframework/configstore/store.goframework/configstore/tables/clientconfig.goframework/configstore/tables/oauth.goframework/oauth2/main.goframework/oauth2/sync.gotransports/bifrost-http/handlers/config.gotransports/bifrost-http/handlers/mcp_sessions.gotransports/bifrost-http/handlers/mcpserver.gotransports/bifrost-http/handlers/oauth2.gotransports/bifrost-http/handlers/oauth2_consent.gotransports/bifrost-http/handlers/oauth2_metadata.gotransports/bifrost-http/handlers/oauth2_per_user.gotransports/bifrost-http/lib/config.gotransports/bifrost-http/lib/config_test.gotransports/bifrost-http/lib/ctx.gotransports/bifrost-http/server/server.goui/app/_fallbacks/enterprise/components/mcp-sessions/dacFilterPill.tsxui/app/_fallbacks/enterprise/components/mcp-sessions/grantedViaVKChips.tsxui/app/_fallbacks/enterprise/components/mcp-sessions/userDisplay.tsxui/app/workspace/config/views/mcpView.tsxui/app/workspace/mcp-sessions/auth/layout.tsxui/app/workspace/mcp-sessions/auth/page.tsxui/app/workspace/mcp-sessions/layout.tsxui/app/workspace/mcp-sessions/oauth-callback/layout.tsxui/app/workspace/mcp-sessions/oauth-callback/page.tsxui/app/workspace/mcp-sessions/page.tsxui/app/workspace/mcp-sessions/views/sessionsTable.tsxui/components/sidebar.tsxui/components/ui/alert.tsxui/lib/store/apis/baseApi.tsui/lib/store/apis/index.tsui/lib/store/apis/mcpSessionsApi.tsui/lib/types/config.tsui/lib/types/mcpSessions.ts
💤 Files with no reviewable changes (8)
- transports/bifrost-http/handlers/oauth2_consent.go
- ui/lib/types/config.ts
- framework/configstore/clientconfig.go
- transports/bifrost-http/handlers/oauth2_per_user.go
- framework/configstore/rdb_test.go
- framework/configstore/tables/clientconfig.go
- transports/bifrost-http/handlers/oauth2_metadata.go
- transports/bifrost-http/handlers/config.go
Pratham-Mishra04
force-pushed
the
05-18-refactor_mcp_per_user_oauth_flow_refactor
branch
from
May 19, 2026 14:23
05cfa3e to
9bab60b
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 5
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
framework/configstore/rdb.go (1)
4780-4819:⚠️ Potential issue | 🟠 Major | 🏗️ Heavy liftThe new "upsert" path still races on first insert.
When no row exists yet, the
SELECT ... FOR UPDATElocks nothing. Two concurrent callbacks/reauths for the same(identity, mcp_client_id)can both miss, both hitCreate, and one will fail on the unique constraint instead of converging on a single row. Use a realON CONFLICT DO UPDATEpath here, or retry the locked read/update path after a duplicate-key insert error.🤖 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/configstore/rdb.go` around lines 4780 - 4819, The upsert currently does a SELECT ... FOR UPDATE via dbForUpdate inside the Transaction and falls back to tx.Create which races on first insert; change this to a true DB-level upsert or retry-on-duplicate-key: either replace the lookup+Create path with a single ON CONFLICT DO UPDATE upsert using GORM's Clauses(clause.OnConflict{UpdateAll: true} or a specific SET list) against tables.TableOauthUserToken, or keep the SELECT/LOCK flow but catch a unique-constraint error from tx.Create and in that case re-run the locked lookup (dbForUpdate(...).First(&existing)) and then tx.Save(token) to update the existing row; adjust tx.Save/tx.Create usage accordingly inside s.DB().WithContext(ctx).Transaction to ensure only one row is created for (user_id|virtual_key_id|session_id, mcp_client_id).
♻️ Duplicate comments (1)
framework/configstore/migrations.go (1)
7772-7779:⚠️ Potential issue | 🟠 Major | ⚡ Quick winBackfill
flow_modefor legacy session rows too.At this point
session_token_hash/session_idis still the only way to recognize session-backed rows. Falling through to'none'means the immediately followingmigrationDropNonVKOauthUserRowswill delete valid session-mode records instead of just the bindings you intend to retire.🤖 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/configstore/migrations.go` around lines 7772 - 7779, The UPDATE that sets flow_mode in migrations.go currently falls back to 'none' and thus will misclassify legacy session-backed rows; modify the CASE in the tx.Exec(...) SQL (the block that updates oauth_user_sessions) to include a WHEN that checks for session_token_hash IS NOT NULL OR session_id IS NOT NULL (and not empty if applicable) and sets flow_mode = 'session' for those rows before the ELSE branch so legacy session records are backfilled correctly.
🧹 Nitpick comments (3)
ui/app/workspace/mcp-sessions/page.tsx (1)
3-3: ⚡ Quick winUse
@/...alias import instead of a relative path.Please switch this import to the UI alias pattern for consistency with the rest of the codebase.
Suggested change
-import SessionsTable from "./views/sessionsTable"; +import SessionsTable from "`@/app/workspace/mcp-sessions/views/sessionsTable`";Based on learnings, in the UI codebase prefer alias imports using
@/...over relative imports in TS/TSX files underui/.🤖 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 `@ui/app/workspace/mcp-sessions/page.tsx` at line 3, Replace the relative import of SessionsTable in page.tsx with the UI alias import pattern: update the import statement that currently reads import SessionsTable from "./views/sessionsTable" to use the '`@/`...' alias (import SessionsTable from "`@/`..." pointing to the same views/sessionsTable module) so the file uses the project-wide UI alias convention for TSX imports.transports/bifrost-http/handlers/governance.go (2)
1431-1445: 🏗️ Heavy liftConsider aligning deletion order across all entity handlers.
The new cleanup order (DB deletion → in-memory removal) in
deleteVirtualKeyfollows the safer pattern also used indeleteRoutingRule(line 3554): persist the DB change first, then update in-memory state as a non-fatal operation. However,deleteTeam(line 1926) anddeleteCustomer(line 2307) still use the old order (in-memory first, then DB).The new pattern is preferable because:
- DB becomes the definitive source of truth before in-memory state changes
- If DB deletion fails, in-memory state remains untouched
- In-memory cleanup failures after successful DB deletion are safely logged without blocking the operation
Recommend updating
deleteTeamanddeleteCustomerto follow the same pattern for consistency and to ensure all deletion handlers treat the database as the authoritative source.🤖 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 `@transports/bifrost-http/handlers/governance.go` around lines 1431 - 1445, Standardize deletion order across handlers: follow the pattern used in deleteVirtualKey and deleteRoutingRule—perform the persistent deletion via configStore (e.g., configStore.DeleteTeam / configStore.DeleteCustomer) first, handle configstore.ErrNotFound appropriately, then perform in-memory cleanup (e.g., governanceManager.RemoveTeam / governanceManager.RemoveCustomer) as a best-effort non-fatal step and log any removal errors without returning failure; update deleteTeam and deleteCustomer to mirror deleteVirtualKey's flow and logging behavior so the DB remains the source of truth and in-memory failures do not block the successful API response.
1443-1443: ⚡ Quick winComplete the error-handling comment.
The comment on line 1443 is incomplete:
"But we ignore this error because its not"ends mid-sentence. The same incomplete comment appears at lines 1929 (deleteTeam) and 2309 (deleteCustomer).Complete the explanation or adopt the clearer pattern from
deleteRoutingRule(line 3564):"non-fatal: DB already updated".📝 Suggested fix for all three locations
- // But we ignore this error because its not + // In-memory removal is non-fatal: DB deletion already succeeded logger.Error("failed to remove virtual key: %v", err)Apply the same fix at lines 1929 and 2309.
🤖 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 `@transports/bifrost-http/handlers/governance.go` at line 1443, Replace the incomplete comment "But we ignore this error because its not" with a clear non-fatal explanation matching the pattern used in deleteRoutingRule: e.g., "non-fatal: DB already updated". Do this for the three occurrences in governance.go: the comment at line ~1443 and the ones in the deleteTeam and deleteCustomer functions (the existing deleteRoutingRule comment can be used as the reference wording).
🤖 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/migrations.go`:
- Around line 7993-8009: The Rollback implementation for the migration (the
Rollback func that uses tx.Exec to ALTER TABLE and recreate oauth_per_user_*
tables) is unsafe because it fakes table restoration; replace that logic with an
explicit irreversible error: remove the best-effort ALTER/CREATE statements and
have Rollback immediately return fmt.Errorf("irreversible migration: legacy
oauth_per_user schema not restorable") (or similar) so downgrades fail fast
unless the real legacy schema and data restoration is implemented; ensure the
function signature and returned error type remain compatible with the current
migration framework.
In `@framework/configstore/rdb.go`:
- Around line 4721-4727: Update the outdated comment on
GetOauthUserSessionByModeIdentityAndMCPClient to remove the mention of hashing
the session token for lookups and clearly state that for AuthModeSession the
store now queries session_id directly (no hashing), and make the same
clarification in the corresponding helper comments around the later block (the
other session-mode helper in the same file). Ensure the per-mode identity
mapping lines (AuthModeUser, AuthModeVK, AuthModeSession) accurately reflect
current behavior and mention that session lookups/deletes use the raw/session_id
column rather than a hashed value.
- Around line 4954-4968: ListAllOauthUserTokens currently filters statuses to
only "active" (and "orphaned") so tokens set to "needs_reauth" (by
MarkOauthUserTokenNeedsReauthByID) are excluded; update the statuses slice and
the Where("status IN ?", statuses) usage in ListAllOauthUserTokens to include
"needs_reauth" (or conditionally include it if desired) so those rows appear in
the sessions listing, ensuring Preload and Order calls remain unchanged.
In `@transports/bifrost-http/handlers/mcp_sessions.go`:
- Around line 83-98: The list handler currently returns all sessions because no
DAC QueryScope or caller identity is applied; fix by extracting the caller
identity from the request context (the auth info injected by middleware) and set
a QueryScope in ctx before calling h.store.ConfigStore.ListAllOauthUserTokens
and ListAllPendingOauthUserSessions so the store returns only visible rows, or
if setting QueryScope is not possible, post-filter the returned slices
(flows/tokens) by comparing each row's owner/team fields to the caller identity
(use the same visibility rules used elsewhere), then continue to call tokenRow()
/ flowRow() only on the filtered set; update list() to perform one of these two
actions and ensure any context helper used to set/query scope (the ScopedDB(ctx)
path) sees the new QueryScope.
---
Outside diff comments:
In `@framework/configstore/rdb.go`:
- Around line 4780-4819: The upsert currently does a SELECT ... FOR UPDATE via
dbForUpdate inside the Transaction and falls back to tx.Create which races on
first insert; change this to a true DB-level upsert or retry-on-duplicate-key:
either replace the lookup+Create path with a single ON CONFLICT DO UPDATE upsert
using GORM's Clauses(clause.OnConflict{UpdateAll: true} or a specific SET list)
against tables.TableOauthUserToken, or keep the SELECT/LOCK flow but catch a
unique-constraint error from tx.Create and in that case re-run the locked lookup
(dbForUpdate(...).First(&existing)) and then tx.Save(token) to update the
existing row; adjust tx.Save/tx.Create usage accordingly inside
s.DB().WithContext(ctx).Transaction to ensure only one row is created for
(user_id|virtual_key_id|session_id, mcp_client_id).
---
Duplicate comments:
In `@framework/configstore/migrations.go`:
- Around line 7772-7779: The UPDATE that sets flow_mode in migrations.go
currently falls back to 'none' and thus will misclassify legacy session-backed
rows; modify the CASE in the tx.Exec(...) SQL (the block that updates
oauth_user_sessions) to include a WHEN that checks for session_token_hash IS NOT
NULL OR session_id IS NOT NULL (and not empty if applicable) and sets flow_mode
= 'session' for those rows before the ELSE branch so legacy session records are
backfilled correctly.
---
Nitpick comments:
In `@transports/bifrost-http/handlers/governance.go`:
- Around line 1431-1445: Standardize deletion order across handlers: follow the
pattern used in deleteVirtualKey and deleteRoutingRule—perform the persistent
deletion via configStore (e.g., configStore.DeleteTeam /
configStore.DeleteCustomer) first, handle configstore.ErrNotFound appropriately,
then perform in-memory cleanup (e.g., governanceManager.RemoveTeam /
governanceManager.RemoveCustomer) as a best-effort non-fatal step and log any
removal errors without returning failure; update deleteTeam and deleteCustomer
to mirror deleteVirtualKey's flow and logging behavior so the DB remains the
source of truth and in-memory failures do not block the successful API response.
- Line 1443: Replace the incomplete comment "But we ignore this error because
its not" with a clear non-fatal explanation matching the pattern used in
deleteRoutingRule: e.g., "non-fatal: DB already updated". Do this for the three
occurrences in governance.go: the comment at line ~1443 and the ones in the
deleteTeam and deleteCustomer functions (the existing deleteRoutingRule comment
can be used as the reference wording).
In `@ui/app/workspace/mcp-sessions/page.tsx`:
- Line 3: Replace the relative import of SessionsTable in page.tsx with the UI
alias import pattern: update the import statement that currently reads import
SessionsTable from "./views/sessionsTable" to use the '`@/`...' alias (import
SessionsTable from "`@/`..." pointing to the same views/sessionsTable module) so
the file uses the project-wide UI alias convention for TSX imports.
🪄 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: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 844e4b7c-3bff-4563-8d4c-65c5ac3a86a6
📒 Files selected for processing (41)
core/mcp/utils/utils.gocore/schemas/bifrost.gocore/schemas/context.gocore/schemas/mcp.gocore/schemas/oauth.goframework/configstore/clientconfig.goframework/configstore/migrations.goframework/configstore/rdb.goframework/configstore/rdb_test.goframework/configstore/store.goframework/configstore/tables/clientconfig.goframework/configstore/tables/oauth.goframework/oauth2/main.goframework/oauth2/sync.gotransports/bifrost-http/handlers/config.gotransports/bifrost-http/handlers/governance.gotransports/bifrost-http/handlers/mcp_sessions.gotransports/bifrost-http/handlers/mcpserver.gotransports/bifrost-http/handlers/oauth2.gotransports/bifrost-http/handlers/oauth2_consent.gotransports/bifrost-http/handlers/oauth2_metadata.gotransports/bifrost-http/handlers/oauth2_per_user.gotransports/bifrost-http/lib/config.gotransports/bifrost-http/lib/config_test.gotransports/bifrost-http/lib/ctx.gotransports/bifrost-http/server/server.goui/app/workspace/config/views/mcpView.tsxui/app/workspace/mcp-sessions/auth/layout.tsxui/app/workspace/mcp-sessions/auth/page.tsxui/app/workspace/mcp-sessions/layout.tsxui/app/workspace/mcp-sessions/oauth-callback/layout.tsxui/app/workspace/mcp-sessions/oauth-callback/page.tsxui/app/workspace/mcp-sessions/page.tsxui/app/workspace/mcp-sessions/views/sessionsTable.tsxui/components/sidebar.tsxui/components/ui/alert.tsxui/lib/store/apis/baseApi.tsui/lib/store/apis/index.tsui/lib/store/apis/mcpSessionsApi.tsui/lib/types/config.tsui/lib/types/mcpSessions.ts
💤 Files with no reviewable changes (8)
- ui/lib/types/config.ts
- transports/bifrost-http/handlers/oauth2_metadata.go
- framework/configstore/rdb_test.go
- framework/configstore/clientconfig.go
- framework/configstore/tables/clientconfig.go
- transports/bifrost-http/handlers/oauth2_per_user.go
- transports/bifrost-http/handlers/oauth2_consent.go
- transports/bifrost-http/handlers/config.go
✅ Files skipped from review due to trivial changes (1)
- ui/lib/store/apis/baseApi.ts
🚧 Files skipped from review as they are similar to previous changes (18)
- ui/app/workspace/mcp-sessions/layout.tsx
- core/schemas/context.go
- ui/app/workspace/mcp-sessions/auth/layout.tsx
- core/schemas/bifrost.go
- transports/bifrost-http/lib/ctx.go
- transports/bifrost-http/handlers/mcpserver.go
- framework/oauth2/sync.go
- ui/app/workspace/config/views/mcpView.tsx
- ui/lib/store/apis/index.ts
- ui/app/workspace/mcp-sessions/oauth-callback/page.tsx
- ui/components/sidebar.tsx
- ui/app/workspace/mcp-sessions/views/sessionsTable.tsx
- ui/components/ui/alert.tsx
- ui/app/workspace/mcp-sessions/oauth-callback/layout.tsx
- core/mcp/utils/utils.go
- ui/lib/store/apis/mcpSessionsApi.ts
- transports/bifrost-http/lib/config.go
- framework/configstore/tables/oauth.go
Pratham-Mishra04
force-pushed
the
05-18-refactor_mcp_per_user_oauth_flow_refactor
branch
from
May 19, 2026 15:57
9bab60b to
3b6be1e
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 5
🤖 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/migrations.go`:
- Around line 8056-8120: This migration mixes a heavy SELECT/UPDATE backfill of
config_hash with an ALTER TABLE DROP COLUMN in the same migration (the Migrate
func), which can lock config_client; split this into two migrations like
migrationDropAllowDirectKeysColumn / migrationDropAllowDirectKeysColumnDDL: keep
the backfill logic (the tx.Find loop that regenerates and updates config_hash
using ClientConfig.GenerateClientConfigHash and tx.Model(&cc).Update) in one
migration that performs only the row updates, and move the ALTER TABLE DROP
COLUMN mcp_external_server_url into a separate migration that runs the DDL only
(non-transactional/DDL migration variant) so the expensive row updates and the
DROP COLUMN do not run in the same transaction or migration step.
- Around line 7725-7753: The migration currently constructs partialUniques and
will fail if duplicate (identity, mcp_client_id) rows exist; before executing
the CREATE UNIQUE INDEX statements (idx_oauth_user_tokens_user_mcp,
idx_oauth_user_tokens_vk_mcp and the session-based index), add a deduplication
step that finds duplicate groups in oauth_user_tokens grouped by (user_id,
mcp_client_id) and (virtual_key_id, mcp_client_id) (and the sessionCol variant
when sessionCol != ""), and collapse them (e.g. keep the single canonical row
per group by newest/highest id or archive duplicates into a separate table and
delete the extras) so no more than one row per unique key remains; perform this
dedupe inside the same tx before iterating partialUniques so the subsequent
tx.Exec(...) CREATE UNIQUE INDEX calls succeed.
In `@framework/configstore/tables/oauth.go`:
- Around line 166-170: Validate and enforce the FlowMode invariants inside
TableOauthUserSession.BeforeSave by (1) ensuring each flow_mode value only has
its allowed identity field set—'session' requires SessionID and disallows
VirtualKeyID/UserID, 'vk' requires VirtualKeyID and disallows SessionID/UserID,
'user' requires UserID and disallows SessionID/VirtualKeyID—and return an error
if the combination is invalid, and (2) prevent changes to FlowMode on updates by
comparing the existing DB row's FlowMode to the current struct and rejecting any
modification; apply the same validation logic to the related
BeforeSave/BeforeUpdate hooks referenced near the other oauth user session
handling block so invariants are enforced on both create and update paths.
- Around line 229-234: The BeforeSave hook must enforce the one-identity-per-row
contract for TableOauthUserToken: validate that AuthMode is one of "session",
"vk", or "user" and that exactly the matching identity column is present (for
"session" SessionID is non-empty and VirtualKeyID/UserID are nil/empty; for "vk"
VirtualKeyID is non-nil/non-empty and SessionID/UserID are nil/empty; for "user"
UserID is non-nil/non-empty and SessionID/VirtualKeyID are nil/empty). Update
TableOauthUserToken.BeforeSave to perform this check and return an error if the
shape is invalid; include the field names SessionID, VirtualKeyID, UserID and
AuthMode in the validation logic. Apply the same validation to the other
lifecycle hook(s) that handle persistence for TableOauthUserToken (the similar
block referenced around lines 258-271) so the invariant is enforced on
create/update paths.
In `@transports/bifrost-http/handlers/oauth2.go`:
- Around line 114-116: The code sets oauthConfig.Status = "failed" and calls
h.store.ConfigStore.UpdateOauthConfig(context.Background(), oauthConfig) but
ignores the returned error; modify the call to capture the error from
UpdateOauthConfig and log it (including the error value and relevant context
such as oauthConfig.ID or client name) via the existing logger (e.g., h.logger
or processLogger) so failures to persist the status are visible for debugging
when h.store.ConfigStore.UpdateOauthConfig(...) returns an error.
🪄 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: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 6608334f-5dc4-440b-8343-ab0dde1c7725
📒 Files selected for processing (29)
core/mcp/utils/utils.gocore/schemas/bifrost.gocore/schemas/context.gocore/schemas/mcp.gocore/schemas/oauth.goframework/configstore/clientconfig.goframework/configstore/migrations.goframework/configstore/rdb.goframework/configstore/rdb_test.goframework/configstore/store.goframework/configstore/tables/clientconfig.goframework/configstore/tables/oauth.goframework/oauth2/main.goframework/oauth2/sync.gotransports/bifrost-http/handlers/config.gotransports/bifrost-http/handlers/governance.gotransports/bifrost-http/handlers/mcp_sessions.gotransports/bifrost-http/handlers/mcpserver.gotransports/bifrost-http/handlers/oauth2.gotransports/bifrost-http/lib/config.gotransports/bifrost-http/lib/config_test.gotransports/bifrost-http/lib/ctx.goui/app/workspace/config/views/mcpView.tsxui/app/workspace/mcp-sessions/auth/page.tsxui/app/workspace/mcp-sessions/oauth-callback/page.tsxui/app/workspace/mcp-sessions/views/sessionsTable.tsxui/components/ui/alert.tsxui/lib/types/config.tsui/lib/types/mcpSessions.ts
💤 Files with no reviewable changes (5)
- ui/lib/types/config.ts
- framework/configstore/tables/clientconfig.go
- framework/configstore/clientconfig.go
- framework/configstore/rdb_test.go
- transports/bifrost-http/handlers/config.go
🚧 Files skipped from review as they are similar to previous changes (18)
- transports/bifrost-http/handlers/governance.go
- core/schemas/mcp.go
- core/schemas/bifrost.go
- transports/bifrost-http/lib/ctx.go
- ui/components/ui/alert.tsx
- ui/app/workspace/config/views/mcpView.tsx
- ui/app/workspace/mcp-sessions/oauth-callback/page.tsx
- core/mcp/utils/utils.go
- transports/bifrost-http/handlers/mcpserver.go
- transports/bifrost-http/lib/config.go
- framework/oauth2/sync.go
- core/schemas/context.go
- ui/app/workspace/mcp-sessions/views/sessionsTable.tsx
- transports/bifrost-http/handlers/mcp_sessions.go
- transports/bifrost-http/lib/config_test.go
- framework/oauth2/main.go
- core/schemas/oauth.go
- framework/configstore/rdb.go
Pratham-Mishra04
force-pushed
the
05-18-refactor_mcp_per_user_oauth_flow_refactor
branch
from
May 19, 2026 16:18
3b6be1e to
8b3deb7
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 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/migrations.go`:
- Around line 7671-7821: The dedupe SQL (dedupeStmts) and the partialUniques
index statements must be constrained by auth_mode so rows that contain multiple
identity columns only participate in the intended uniqueness/dedupe domain;
update each dedupe DELETE and each CREATE UNIQUE INDEX predicate to include the
matching auth_mode filter (e.g. for the user partition/index add "AND auth_mode
= 'user'" / WHERE auth_mode = 'user', for vk use 'vk', for session use
'session'), using the existing variables sessionCol, dedupeStmts and
partialUniques and leaving the backfillSQL unchanged; this ensures dedupe keeps
the correct winner and CREATE UNIQUE INDEX succeeds prior to
migrationDropNonVKOauthUserRows.
In `@transports/bifrost-http/handlers/oauth2.go`:
- Around line 79-80: The logger.Error calls in
transports/bifrost-http/handlers/oauth2.go are emitting raw OAuth `state` values
(e.g., the call that currently logs "[oauth] per-user callback completion
failed: state=%s err=%v" along with similar logger.Error uses later), which must
be removed; update those logger.Error invocations to omit the raw state and
instead log only non-sensitive context and the error (perUserErr), or log a
deterministic safe surrogate (e.g., a short hash or the literal "<redacted>") if
you need to correlate attempts. Apply the same change for all occurrences of
logger.Error that include `state` and ensure redirects (ctx.Redirect calls)
remain unchanged except for not relying on logged state.
- Line 77: Replace usage of context.Background() with the incoming fasthttp
request context variable (ctx) when calling request-scoped methods: pass ctx
into h.oauthProvider.CompleteUserOAuthFlow(...) instead of context.Background(),
and do the same for subsequent store/OAuth calls in this handler (calls on
h.store and any other methods invoked around the CompleteUserOAuthFlow call and
in the blocks around lines 89 and 111-116). Ensure you propagate the same ctx so
cancellation, tracing and request-scoped values from *fasthttp.RequestCtx flow
through CompleteUserOAuthFlow and the store method calls.
🪄 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: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 1ed179c8-e44b-41c7-849c-6e3c2f069cf9
📒 Files selected for processing (29)
core/mcp/utils/utils.gocore/schemas/bifrost.gocore/schemas/context.gocore/schemas/mcp.gocore/schemas/oauth.goframework/configstore/clientconfig.goframework/configstore/migrations.goframework/configstore/rdb.goframework/configstore/rdb_test.goframework/configstore/store.goframework/configstore/tables/clientconfig.goframework/configstore/tables/oauth.goframework/oauth2/main.goframework/oauth2/sync.gotransports/bifrost-http/handlers/config.gotransports/bifrost-http/handlers/governance.gotransports/bifrost-http/handlers/mcp_sessions.gotransports/bifrost-http/handlers/mcpserver.gotransports/bifrost-http/handlers/oauth2.gotransports/bifrost-http/lib/config.gotransports/bifrost-http/lib/config_test.gotransports/bifrost-http/lib/ctx.goui/app/workspace/config/views/mcpView.tsxui/app/workspace/mcp-sessions/auth/page.tsxui/app/workspace/mcp-sessions/oauth-callback/page.tsxui/app/workspace/mcp-sessions/views/sessionsTable.tsxui/components/ui/alert.tsxui/lib/types/config.tsui/lib/types/mcpSessions.ts
💤 Files with no reviewable changes (5)
- ui/lib/types/config.ts
- framework/configstore/tables/clientconfig.go
- framework/configstore/clientconfig.go
- framework/configstore/rdb_test.go
- transports/bifrost-http/handlers/config.go
🚧 Files skipped from review as they are similar to previous changes (20)
- core/schemas/mcp.go
- transports/bifrost-http/lib/ctx.go
- core/schemas/context.go
- core/schemas/bifrost.go
- ui/lib/types/mcpSessions.ts
- transports/bifrost-http/handlers/mcpserver.go
- core/mcp/utils/utils.go
- ui/app/workspace/config/views/mcpView.tsx
- framework/oauth2/sync.go
- core/schemas/oauth.go
- ui/app/workspace/mcp-sessions/oauth-callback/page.tsx
- ui/app/workspace/mcp-sessions/auth/page.tsx
- transports/bifrost-http/handlers/governance.go
- transports/bifrost-http/handlers/mcp_sessions.go
- transports/bifrost-http/lib/config_test.go
- framework/configstore/tables/oauth.go
- ui/app/workspace/mcp-sessions/views/sessionsTable.tsx
- framework/configstore/store.go
- framework/oauth2/main.go
- framework/configstore/rdb.go
Pratham-Mishra04
force-pushed
the
05-18-refactor_mcp_per_user_oauth_flow_refactor
branch
from
May 19, 2026 16:36
8b3deb7 to
0c200e7
Compare
Contributor
There was a problem hiding this comment.
♻️ Duplicate comments (1)
framework/configstore/migrations.go (1)
7972-7976:⚠️ Potential issue | 🟠 Major | ⚡ Quick winKeep the recreated session index scoped to session-mode rows.
This reintroduces the cross-mode collision the previous migration just fixed: any row with a populated
session_idnow participates inidx_oauth_user_tokens_session_mcp, even when its winning identity isuserorvk.Suggested fix
if err := tx.Exec(` CREATE UNIQUE INDEX IF NOT EXISTS idx_oauth_user_tokens_session_mcp ON oauth_user_tokens (session_id, mcp_client_id) - WHERE session_id IS NOT NULL AND session_id != '' + WHERE auth_mode = 'session' AND session_id IS NOT NULL AND session_id != '' `).Error; err != nil { return fmt.Errorf("create idx_oauth_user_tokens_session_mcp on session_id: %w", err) }🤖 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/configstore/migrations.go` around lines 7972 - 7976, The recreated unique index idx_oauth_user_tokens_session_mcp on oauth_user_tokens is too broad and must be limited to session-mode rows; update the WHERE clause in the tx.Exec that creates idx_oauth_user_tokens_session_mcp to include the session-mode predicate (e.g. AND winning_identity = 'session') so only rows where winning_identity == 'session' participate (keep the existing session_id IS NOT NULL AND session_id != '' checks and add the winning_identity = 'session' condition).
🧹 Nitpick comments (1)
transports/bifrost-http/lib/config.go (1)
3163-3168: 💤 Low valueConsider extracting the orphan retention duration to a named constant.
The
30*24*time.Hourmagic number represents a policy decision (30-day orphan retention). Extracting it to a package-level constant would improve readability and make future adjustments easier to locate.♻️ Suggested refactor
+// DefaultOAuthOrphanRetention is the default duration after which orphaned +// per-user OAuth token rows are reaped by the sweep worker. +const DefaultOAuthOrphanRetention = 30 * 24 * time.Hour + // Start per-user OAuth sweep worker: expires stale pending flows and reaps // long-orphaned token rows. Orphan retention defaults to 30 days. -config.OAuthSweepWorker = oauth2.NewPerUserOAuthSweepWorker(config.OAuthProvider, 30*24*time.Hour, logger) +config.OAuthSweepWorker = oauth2.NewPerUserOAuthSweepWorker(config.OAuthProvider, DefaultOAuthOrphanRetention, logger)🤖 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 `@transports/bifrost-http/lib/config.go` around lines 3163 - 3168, Extract the 30*24*time.Hour literal into a package-level constant (e.g., const OAuthOrphanRetention = 30*24*time.Hour) and use that constant when calling oauth2.NewPerUserOAuthSweepWorker in the config initialization (replace the literal in the config.OAuthSweepWorker assignment); keep the constant near the top of the file so the orphan retention policy is easy to find and adjust in the future.
🤖 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.
Duplicate comments:
In `@framework/configstore/migrations.go`:
- Around line 7972-7976: The recreated unique index
idx_oauth_user_tokens_session_mcp on oauth_user_tokens is too broad and must be
limited to session-mode rows; update the WHERE clause in the tx.Exec that
creates idx_oauth_user_tokens_session_mcp to include the session-mode predicate
(e.g. AND winning_identity = 'session') so only rows where winning_identity ==
'session' participate (keep the existing session_id IS NOT NULL AND session_id
!= '' checks and add the winning_identity = 'session' condition).
---
Nitpick comments:
In `@transports/bifrost-http/lib/config.go`:
- Around line 3163-3168: Extract the 30*24*time.Hour literal into a
package-level constant (e.g., const OAuthOrphanRetention = 30*24*time.Hour) and
use that constant when calling oauth2.NewPerUserOAuthSweepWorker in the config
initialization (replace the literal in the config.OAuthSweepWorker assignment);
keep the constant near the top of the file so the orphan retention policy is
easy to find and adjust in the future.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: ecdf56db-2ad9-4554-9829-52e3c2e2f0f7
📒 Files selected for processing (29)
core/mcp/utils/utils.gocore/schemas/bifrost.gocore/schemas/context.gocore/schemas/mcp.gocore/schemas/oauth.goframework/configstore/clientconfig.goframework/configstore/migrations.goframework/configstore/rdb.goframework/configstore/rdb_test.goframework/configstore/store.goframework/configstore/tables/clientconfig.goframework/configstore/tables/oauth.goframework/oauth2/main.goframework/oauth2/sync.gotransports/bifrost-http/handlers/config.gotransports/bifrost-http/handlers/governance.gotransports/bifrost-http/handlers/mcp_sessions.gotransports/bifrost-http/handlers/mcpserver.gotransports/bifrost-http/handlers/oauth2.gotransports/bifrost-http/lib/config.gotransports/bifrost-http/lib/config_test.gotransports/bifrost-http/lib/ctx.goui/app/workspace/config/views/mcpView.tsxui/app/workspace/mcp-sessions/auth/page.tsxui/app/workspace/mcp-sessions/oauth-callback/page.tsxui/app/workspace/mcp-sessions/views/sessionsTable.tsxui/components/ui/alert.tsxui/lib/types/config.tsui/lib/types/mcpSessions.ts
💤 Files with no reviewable changes (5)
- ui/lib/types/config.ts
- framework/configstore/tables/clientconfig.go
- framework/configstore/rdb_test.go
- transports/bifrost-http/handlers/config.go
- framework/configstore/clientconfig.go
🚧 Files skipped from review as they are similar to previous changes (20)
- core/schemas/mcp.go
- ui/components/ui/alert.tsx
- core/schemas/context.go
- transports/bifrost-http/lib/ctx.go
- transports/bifrost-http/handlers/governance.go
- ui/app/workspace/mcp-sessions/auth/page.tsx
- core/schemas/bifrost.go
- ui/lib/types/mcpSessions.ts
- ui/app/workspace/mcp-sessions/oauth-callback/page.tsx
- transports/bifrost-http/handlers/oauth2.go
- ui/app/workspace/config/views/mcpView.tsx
- transports/bifrost-http/handlers/mcpserver.go
- core/mcp/utils/utils.go
- framework/oauth2/sync.go
- ui/app/workspace/mcp-sessions/views/sessionsTable.tsx
- core/schemas/oauth.go
- framework/configstore/tables/oauth.go
- framework/configstore/store.go
- framework/oauth2/main.go
- framework/configstore/rdb.go
roroghost17
force-pushed
the
05-18-refactor_mcp_per_user_oauth_flow_refactor
branch
from
May 19, 2026 18:46
0c200e7 to
f649709
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (2)
framework/configstore/migrations.go (2)
7972-7976:⚠️ Potential issue | 🟠 Major | ⚡ Quick winKeep the recreated session index scoped to
auth_mode = 'session'.The previous migration deliberately partitions every identity-domain unique index by
auth_mode. Recreatingidx_oauth_user_tokens_session_mcpwithout that predicate reopens cross-mode collisions for any row that carries a non-emptysession_idoutside session mode.Suggested fix
if err := tx.Exec(` CREATE UNIQUE INDEX IF NOT EXISTS idx_oauth_user_tokens_session_mcp ON oauth_user_tokens (session_id, mcp_client_id) - WHERE session_id IS NOT NULL AND session_id != '' + WHERE auth_mode = 'session' AND session_id IS NOT NULL AND session_id != '' `).Error; err != nil { return fmt.Errorf("create idx_oauth_user_tokens_session_mcp on session_id: %w", err) }🤖 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/configstore/migrations.go` around lines 7972 - 7976, Recreate the unique index idx_oauth_user_tokens_session_mcp on table oauth_user_tokens so it is still scoped to session-mode rows; update the tx.Exec call that creates idx_oauth_user_tokens_session_mcp to include the predicate "WHERE auth_mode = 'session' AND session_id IS NOT NULL AND session_id != ''" (keeping mcp_client_id and session_id as the indexed columns) so cross-mode collisions are not reintroduced.
7920-7942:⚠️ Potential issue | 🟠 Major | ⚡ Quick winRecreate the
oauth_user_sessions.session_idindex after the column swap.After dropping
idx_oauth_user_sessions_session_token_hash, this migration never creates the replacement index forsession_id.AddColumnonly adds the column, so lookups on the new identifier fall back to a full scan.Suggested fix
if !mg.HasColumn(&tables.TableOauthUserSession{}, "session_id") { if err := mg.AddColumn(&tables.TableOauthUserSession{}, "SessionID"); err != nil { return fmt.Errorf("add session_id to oauth_user_sessions: %w", err) } } + if !mg.HasIndex(&tables.TableOauthUserSession{}, "idx_oauth_user_sessions_session_id") { + if err := mg.CreateIndex(&tables.TableOauthUserSession{}, "SessionID"); err != nil { + return fmt.Errorf("create session_id index on oauth_user_sessions: %w", err) + } + } // The legacy session_token_hash column had a uniqueIndex declared // via gorm tag; GORM names it something like // "idx_oauth_user_sessions_session_token_hash". DROP COLUMN dropsExpected result:
framework/configstore/tables/oauth.goshowsSessionIDis indexed, while this migration currently has no matchingCreateIndexcall.#!/bin/bash set -euo pipefail printf '%s\n' '--- oauth session schema ---' rg -n -C2 'SessionID|session_id|session_token_hash' framework/configstore/tables/oauth.go printf '\n%s\n' '--- migration index handling ---' rg -n -C3 'replace_oauth_session_token_with_session_id|CreateIndex|idx_oauth_user_sessions_session_token_hash|session_id' framework/configstore/migrations.go🤖 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/configstore/migrations.go` around lines 7920 - 7942, The migration adds SessionID to tables.TableOauthUserSession but never creates the corresponding index, so after dropping idx_oauth_user_sessions_session_token_hash you should create a new index on the session_id column; update the migration (the block using mg.HasColumn(&tables.TableOauthUserSession{}, "session_id"), mg.AddColumn, and tx.Exec(...drop legacy index...)) to call the appropriate GORM index creation (e.g., mg.CreateIndex or tx.Exec to CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_oauth_user_sessions_session_id ON oauth_user_sessions(session_id)) for SessionID so lookups use the new index.
🤖 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 `@ui/lib/types/mcpSessions.ts`:
- Around line 8-16: Update the comment above MCPSessionStatus to document flow
statuses as "pending" | "authorized" | "failed" | "expired" (not
"needs_reauth"), and change the MCPFlowDetail.status union type to "pending" |
"authorized" | "failed" | "expired" so flow statuses match the backend
TableOauthUserSession; leave token-only statuses (e.g., "needs_reauth" if used
by TableOauthUserToken) out of MCPFlowDetail.status and ensure
MCPSessionStatus/comment clearly separates token vs flow status sets.
---
Duplicate comments:
In `@framework/configstore/migrations.go`:
- Around line 7972-7976: Recreate the unique index
idx_oauth_user_tokens_session_mcp on table oauth_user_tokens so it is still
scoped to session-mode rows; update the tx.Exec call that creates
idx_oauth_user_tokens_session_mcp to include the predicate "WHERE auth_mode =
'session' AND session_id IS NOT NULL AND session_id != ''" (keeping
mcp_client_id and session_id as the indexed columns) so cross-mode collisions
are not reintroduced.
- Around line 7920-7942: The migration adds SessionID to
tables.TableOauthUserSession but never creates the corresponding index, so after
dropping idx_oauth_user_sessions_session_token_hash you should create a new
index on the session_id column; update the migration (the block using
mg.HasColumn(&tables.TableOauthUserSession{}, "session_id"), mg.AddColumn, and
tx.Exec(...drop legacy index...)) to call the appropriate GORM index creation
(e.g., mg.CreateIndex or tx.Exec to CREATE INDEX CONCURRENTLY IF NOT EXISTS
idx_oauth_user_sessions_session_id ON oauth_user_sessions(session_id)) for
SessionID so lookups use the new index.
🪄 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: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: f05b1fa2-ec9a-49b9-9d4b-adc17c36cfa1
📒 Files selected for processing (41)
core/mcp/utils/utils.gocore/schemas/bifrost.gocore/schemas/context.gocore/schemas/mcp.gocore/schemas/oauth.goframework/configstore/clientconfig.goframework/configstore/migrations.goframework/configstore/rdb.goframework/configstore/rdb_test.goframework/configstore/store.goframework/configstore/tables/clientconfig.goframework/configstore/tables/oauth.goframework/oauth2/main.goframework/oauth2/sync.gotransports/bifrost-http/handlers/config.gotransports/bifrost-http/handlers/governance.gotransports/bifrost-http/handlers/mcp_sessions.gotransports/bifrost-http/handlers/mcpserver.gotransports/bifrost-http/handlers/oauth2.gotransports/bifrost-http/handlers/oauth2_consent.gotransports/bifrost-http/handlers/oauth2_metadata.gotransports/bifrost-http/handlers/oauth2_per_user.gotransports/bifrost-http/lib/config.gotransports/bifrost-http/lib/config_test.gotransports/bifrost-http/lib/ctx.gotransports/bifrost-http/server/server.goui/app/workspace/config/views/mcpView.tsxui/app/workspace/mcp-sessions/auth/layout.tsxui/app/workspace/mcp-sessions/auth/page.tsxui/app/workspace/mcp-sessions/layout.tsxui/app/workspace/mcp-sessions/oauth-callback/layout.tsxui/app/workspace/mcp-sessions/oauth-callback/page.tsxui/app/workspace/mcp-sessions/page.tsxui/app/workspace/mcp-sessions/views/sessionsTable.tsxui/components/sidebar.tsxui/components/ui/alert.tsxui/lib/store/apis/baseApi.tsui/lib/store/apis/index.tsui/lib/store/apis/mcpSessionsApi.tsui/lib/types/config.tsui/lib/types/mcpSessions.ts
💤 Files with no reviewable changes (8)
- framework/configstore/tables/clientconfig.go
- ui/lib/types/config.ts
- transports/bifrost-http/handlers/oauth2_per_user.go
- framework/configstore/clientconfig.go
- transports/bifrost-http/handlers/oauth2_metadata.go
- framework/configstore/rdb_test.go
- transports/bifrost-http/handlers/config.go
- transports/bifrost-http/handlers/oauth2_consent.go
🚧 Files skipped from review as they are similar to previous changes (30)
- ui/lib/store/apis/index.ts
- core/schemas/mcp.go
- ui/components/ui/alert.tsx
- ui/app/workspace/mcp-sessions/oauth-callback/page.tsx
- ui/app/workspace/mcp-sessions/auth/layout.tsx
- transports/bifrost-http/server/server.go
- ui/app/workspace/mcp-sessions/layout.tsx
- ui/app/workspace/mcp-sessions/oauth-callback/layout.tsx
- ui/app/workspace/config/views/mcpView.tsx
- core/schemas/bifrost.go
- core/mcp/utils/utils.go
- ui/lib/store/apis/baseApi.ts
- transports/bifrost-http/lib/ctx.go
- transports/bifrost-http/handlers/governance.go
- framework/oauth2/sync.go
- core/schemas/context.go
- ui/components/sidebar.tsx
- ui/app/workspace/mcp-sessions/auth/page.tsx
- ui/app/workspace/mcp-sessions/views/sessionsTable.tsx
- transports/bifrost-http/handlers/mcpserver.go
- framework/configstore/tables/oauth.go
- core/schemas/oauth.go
- ui/lib/store/apis/mcpSessionsApi.ts
- transports/bifrost-http/handlers/oauth2.go
- framework/oauth2/main.go
- transports/bifrost-http/lib/config_test.go
- transports/bifrost-http/lib/config.go
- transports/bifrost-http/handlers/mcp_sessions.go
- framework/configstore/rdb.go
- framework/configstore/store.go
Pratham-Mishra04
force-pushed
the
05-18-refactor_mcp_per_user_oauth_flow_refactor
branch
from
May 19, 2026 19:24
f649709 to
da09599
Compare
Contributor
There was a problem hiding this comment.
♻️ Duplicate comments (1)
framework/configstore/migrations.go (1)
7920-7933:⚠️ Potential issue | 🟠 Major | ⚡ Quick winRecreate the dropped lookup index on
oauth_user_sessions.session_id.Line 7920 adds
session_id, and Line 7930 drops the oldsession_token_hashindex, but this migration never creates the replacement index on the new lookup column.AddColumndoes not materialize struct-tag indexes, so upgraded databases lose indexedsession_idlookups here.#!/bin/bash set -euo pipefail printf '--- TableOauthUserSession schema ---\n' rg -n -C3 'type TableOauthUserSession struct|SessionID' framework/configstore/tables/oauth.go printf '\n--- MigrationReplaceOauthSessionTokenWithSessionID ---\n' sed -n '7918,7980p' framework/configstore/migrations.goExpected result:
TableOauthUserSession.SessionIDis indexed in the table schema, while this migration only adds the column and drops the legacysession_token_hashindex without a matchingCreateIndex(&tables.TableOauthUserSession{}, "SessionID").🤖 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/configstore/migrations.go` around lines 7920 - 7933, The migration adds the session_id column on TableOauthUserSession but never recreates the lookup index after dropping the old session_token_hash index; update the migration (the block using mg.HasColumn(&tables.TableOauthUserSession{}, "session_id") / mg.AddColumn) to call mg.CreateIndex(&tables.TableOauthUserSession{}, "SessionID") (or the equivalent migrator.CreateIndex) after the DROP INDEX step and after confirming the column exists so that SessionID is indexed for lookups; ensure you use the same struct/field name TableOauthUserSession.SessionID and handle/return any error from CreateIndex similarly to the existing error handling.
🤖 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.
Duplicate comments:
In `@framework/configstore/migrations.go`:
- Around line 7920-7933: The migration adds the session_id column on
TableOauthUserSession but never recreates the lookup index after dropping the
old session_token_hash index; update the migration (the block using
mg.HasColumn(&tables.TableOauthUserSession{}, "session_id") / mg.AddColumn) to
call mg.CreateIndex(&tables.TableOauthUserSession{}, "SessionID") (or the
equivalent migrator.CreateIndex) after the DROP INDEX step and after confirming
the column exists so that SessionID is indexed for lookups; ensure you use the
same struct/field name TableOauthUserSession.SessionID and handle/return any
error from CreateIndex similarly to the existing error handling.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 95cf3fe6-782c-4508-a805-416c044f87de
📒 Files selected for processing (41)
core/mcp/utils/utils.gocore/schemas/bifrost.gocore/schemas/context.gocore/schemas/mcp.gocore/schemas/oauth.goframework/configstore/clientconfig.goframework/configstore/migrations.goframework/configstore/rdb.goframework/configstore/rdb_test.goframework/configstore/store.goframework/configstore/tables/clientconfig.goframework/configstore/tables/oauth.goframework/oauth2/main.goframework/oauth2/sync.gotransports/bifrost-http/handlers/config.gotransports/bifrost-http/handlers/governance.gotransports/bifrost-http/handlers/mcp_sessions.gotransports/bifrost-http/handlers/mcpserver.gotransports/bifrost-http/handlers/oauth2.gotransports/bifrost-http/handlers/oauth2_consent.gotransports/bifrost-http/handlers/oauth2_metadata.gotransports/bifrost-http/handlers/oauth2_per_user.gotransports/bifrost-http/lib/config.gotransports/bifrost-http/lib/config_test.gotransports/bifrost-http/lib/ctx.gotransports/bifrost-http/server/server.goui/app/workspace/config/views/mcpView.tsxui/app/workspace/mcp-sessions/auth/layout.tsxui/app/workspace/mcp-sessions/auth/page.tsxui/app/workspace/mcp-sessions/layout.tsxui/app/workspace/mcp-sessions/oauth-callback/layout.tsxui/app/workspace/mcp-sessions/oauth-callback/page.tsxui/app/workspace/mcp-sessions/page.tsxui/app/workspace/mcp-sessions/views/sessionsTable.tsxui/components/sidebar.tsxui/components/ui/alert.tsxui/lib/store/apis/baseApi.tsui/lib/store/apis/index.tsui/lib/store/apis/mcpSessionsApi.tsui/lib/types/config.tsui/lib/types/mcpSessions.ts
💤 Files with no reviewable changes (8)
- transports/bifrost-http/handlers/oauth2_consent.go
- framework/configstore/tables/clientconfig.go
- ui/lib/types/config.ts
- framework/configstore/clientconfig.go
- transports/bifrost-http/handlers/oauth2_per_user.go
- transports/bifrost-http/handlers/config.go
- transports/bifrost-http/handlers/oauth2_metadata.go
- framework/configstore/rdb_test.go
✅ Files skipped from review due to trivial changes (1)
- ui/components/ui/alert.tsx
🚧 Files skipped from review as they are similar to previous changes (31)
- ui/app/workspace/mcp-sessions/layout.tsx
- ui/lib/store/apis/index.ts
- ui/app/workspace/mcp-sessions/oauth-callback/layout.tsx
- transports/bifrost-http/server/server.go
- ui/components/sidebar.tsx
- core/schemas/mcp.go
- core/schemas/context.go
- transports/bifrost-http/lib/ctx.go
- core/mcp/utils/utils.go
- transports/bifrost-http/handlers/governance.go
- ui/app/workspace/mcp-sessions/auth/layout.tsx
- ui/lib/types/mcpSessions.ts
- ui/lib/store/apis/mcpSessionsApi.ts
- transports/bifrost-http/handlers/mcpserver.go
- core/schemas/bifrost.go
- ui/app/workspace/mcp-sessions/oauth-callback/page.tsx
- ui/app/workspace/mcp-sessions/page.tsx
- ui/app/workspace/mcp-sessions/views/sessionsTable.tsx
- framework/configstore/tables/oauth.go
- ui/app/workspace/mcp-sessions/auth/page.tsx
- ui/app/workspace/config/views/mcpView.tsx
- framework/oauth2/sync.go
- framework/configstore/store.go
- transports/bifrost-http/handlers/mcp_sessions.go
- ui/lib/store/apis/baseApi.ts
- core/schemas/oauth.go
- transports/bifrost-http/handlers/oauth2.go
- transports/bifrost-http/lib/config.go
- transports/bifrost-http/lib/config_test.go
- framework/oauth2/main.go
- framework/configstore/rdb.go
roroghost17
force-pushed
the
05-18-refactor_mcp_per_user_oauth_flow_refactor
branch
from
May 19, 2026 19:47
da09599 to
791502b
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (1)
framework/configstore/migrations.go (1)
7918-7943:⚠️ Potential issue | 🟠 Major | ⚡ Quick winRecreate the
oauth_user_sessions.session_idindex during the column swap.This drops the old
session_token_hashindex but never adds the replacement onsession_id, so upgraded databases lose indexed session lookups.💡 Minimal fix
if mg.HasTable(&tables.TableOauthUserSession{}) { if !mg.HasColumn(&tables.TableOauthUserSession{}, "session_id") { if err := mg.AddColumn(&tables.TableOauthUserSession{}, "SessionID"); err != nil { return fmt.Errorf("add session_id to oauth_user_sessions: %w", err) } } + if err := tx.Exec(` + CREATE INDEX IF NOT EXISTS idx_oauth_user_sessions_session_id + ON oauth_user_sessions (session_id) + `).Error; err != nil { + return fmt.Errorf("create session_id index on oauth_user_sessions: %w", err) + } // The legacy session_token_hash column had a uniqueIndex declared // via gorm tag; GORM names it something like // "idx_oauth_user_sessions_session_token_hash". DROP COLUMN drops🤖 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/configstore/migrations.go` around lines 7918 - 7943, After adding SessionID in the oauth_user_sessions migration block (the code that checks mg.HasTable(&tables.TableOauthUserSession{}), mg.AddColumn, and drops the old session_token_hash/index), recreate the replacement index on the new SessionID column so session lookups remain indexed: add a step after adding the column (and after dropping the old index/columns) to create an index on session_id (use the same uniqueness semantics as the original uniqueIndex on session_token_hash), e.g. execute a CREATE INDEX IF NOT EXISTS (or use the migration helper equivalent) named similarly (e.g. idx_oauth_user_sessions_session_id) via tx.Exec or the mg helper so upgraded databases regain the index. Ensure you reference TableOauthUserSession/SessionID when adding the index and handle errors consistently with the surrounding fmt.Errorf style.
🧹 Nitpick comments (1)
ui/lib/store/apis/mcpSessionsApi.ts (1)
10-10: ⚡ Quick winUse the
@/lib/*alias forbaseApiimport consistency.Line 10 should follow the UI alias convention instead of a relative path.
Based on learnings: "In the Bifrost codebase, prefer using the `@/lib` path alias for imports instead of relative paths from within the ui/lib directory."Suggested diff
-import { baseApi } from "./baseApi"; +import { baseApi } from "`@/lib/store/apis/baseApi`";🤖 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 `@ui/lib/store/apis/mcpSessionsApi.ts` at line 10, The import in mcpSessionsApi uses a relative path; update the import of baseApi in ui/lib/store/apis/mcpSessionsApi.ts to use the project alias (change the import that currently reads import { baseApi } from "./baseApi" to use the "`@/lib`" alias, e.g. import { baseApi } from "`@/lib/baseApi`") so it follows the codebase convention and stays consistent with other modules that reference baseApi.
🤖 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/migrations.go`:
- Around line 7684-7703: The backfill SQL in migrations.go incorrectly sets ELSE
'vk' which classifies identity-less rows as vk; update both backfillSQL branches
(the ones referencing sessionCol and the else branch) to use ELSE NULL (unquoted
NULL) for auth_mode instead of 'vk' so rows with no user_id, virtual_key_id, or
session identifier are not treated as vk and can be removed by
migrationDropNonVKOauthUserRows; locate the string construction around
sessionCol, backfillSQL, and oauth_user_tokens and replace the final ELSE 'vk'
with ELSE NULL in both templated SQL blocks.
In `@ui/lib/store/apis/mcpSessionsApi.ts`:
- Around line 22-25: Replace the invalidatesTags approach in the MCP mutations
with an optimistic cache patch using onQueryStarted + dispatch(updateQueryData)
for the MCPSessions list: in the reauthMCPSession mutation (and the other
mutation at lines 28-31), implement onQueryStarted that calls patchResult =
dispatch(api.util.updateQueryData('getMCPSessions'/* or the actual query
endpoint name */, undefined, draft => { apply the same state change you expect
from the mutation to the draft })), then await the mutation, call
patchResult.undo() on error (or finalize commit on success). Ensure you
reference the MCPSessions query name used by the slice, perform rollback on
exception, and remove invalidatesTags from these mutations.
---
Duplicate comments:
In `@framework/configstore/migrations.go`:
- Around line 7918-7943: After adding SessionID in the oauth_user_sessions
migration block (the code that checks
mg.HasTable(&tables.TableOauthUserSession{}), mg.AddColumn, and drops the old
session_token_hash/index), recreate the replacement index on the new SessionID
column so session lookups remain indexed: add a step after adding the column
(and after dropping the old index/columns) to create an index on session_id (use
the same uniqueness semantics as the original uniqueIndex on
session_token_hash), e.g. execute a CREATE INDEX IF NOT EXISTS (or use the
migration helper equivalent) named similarly (e.g.
idx_oauth_user_sessions_session_id) via tx.Exec or the mg helper so upgraded
databases regain the index. Ensure you reference TableOauthUserSession/SessionID
when adding the index and handle errors consistently with the surrounding
fmt.Errorf style.
---
Nitpick comments:
In `@ui/lib/store/apis/mcpSessionsApi.ts`:
- Line 10: The import in mcpSessionsApi uses a relative path; update the import
of baseApi in ui/lib/store/apis/mcpSessionsApi.ts to use the project alias
(change the import that currently reads import { baseApi } from "./baseApi" to
use the "`@/lib`" alias, e.g. import { baseApi } from "`@/lib/baseApi`") so it
follows the codebase convention and stays consistent with other modules that
reference baseApi.
🪄 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: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: c6709b71-96ae-4508-93f3-336283c691db
📒 Files selected for processing (41)
core/mcp/utils/utils.gocore/schemas/bifrost.gocore/schemas/context.gocore/schemas/mcp.gocore/schemas/oauth.goframework/configstore/clientconfig.goframework/configstore/migrations.goframework/configstore/rdb.goframework/configstore/rdb_test.goframework/configstore/store.goframework/configstore/tables/clientconfig.goframework/configstore/tables/oauth.goframework/oauth2/main.goframework/oauth2/sync.gotransports/bifrost-http/handlers/config.gotransports/bifrost-http/handlers/governance.gotransports/bifrost-http/handlers/mcp_sessions.gotransports/bifrost-http/handlers/mcpserver.gotransports/bifrost-http/handlers/oauth2.gotransports/bifrost-http/handlers/oauth2_consent.gotransports/bifrost-http/handlers/oauth2_metadata.gotransports/bifrost-http/handlers/oauth2_per_user.gotransports/bifrost-http/lib/config.gotransports/bifrost-http/lib/config_test.gotransports/bifrost-http/lib/ctx.gotransports/bifrost-http/server/server.goui/app/workspace/config/views/mcpView.tsxui/app/workspace/mcp-sessions/auth/layout.tsxui/app/workspace/mcp-sessions/auth/page.tsxui/app/workspace/mcp-sessions/layout.tsxui/app/workspace/mcp-sessions/oauth-callback/layout.tsxui/app/workspace/mcp-sessions/oauth-callback/page.tsxui/app/workspace/mcp-sessions/page.tsxui/app/workspace/mcp-sessions/views/sessionsTable.tsxui/components/sidebar.tsxui/components/ui/alert.tsxui/lib/store/apis/baseApi.tsui/lib/store/apis/index.tsui/lib/store/apis/mcpSessionsApi.tsui/lib/types/config.tsui/lib/types/mcpSessions.ts
💤 Files with no reviewable changes (8)
- framework/configstore/tables/clientconfig.go
- transports/bifrost-http/handlers/oauth2_consent.go
- framework/configstore/rdb_test.go
- ui/lib/types/config.ts
- transports/bifrost-http/handlers/oauth2_metadata.go
- transports/bifrost-http/handlers/config.go
- framework/configstore/clientconfig.go
- transports/bifrost-http/handlers/oauth2_per_user.go
🚧 Files skipped from review as they are similar to previous changes (28)
- ui/lib/store/apis/index.ts
- ui/app/workspace/mcp-sessions/oauth-callback/layout.tsx
- ui/app/workspace/mcp-sessions/auth/layout.tsx
- core/schemas/mcp.go
- transports/bifrost-http/server/server.go
- ui/components/sidebar.tsx
- transports/bifrost-http/handlers/governance.go
- core/schemas/context.go
- ui/lib/types/mcpSessions.ts
- core/schemas/bifrost.go
- core/schemas/oauth.go
- ui/components/ui/alert.tsx
- transports/bifrost-http/lib/ctx.go
- ui/lib/store/apis/baseApi.ts
- core/mcp/utils/utils.go
- ui/app/workspace/mcp-sessions/oauth-callback/page.tsx
- ui/app/workspace/mcp-sessions/views/sessionsTable.tsx
- framework/oauth2/sync.go
- framework/configstore/tables/oauth.go
- ui/app/workspace/mcp-sessions/auth/page.tsx
- transports/bifrost-http/lib/config_test.go
- transports/bifrost-http/lib/config.go
- transports/bifrost-http/handlers/mcp_sessions.go
- transports/bifrost-http/handlers/mcpserver.go
- transports/bifrost-http/handlers/oauth2.go
- framework/configstore/store.go
- framework/configstore/rdb.go
- framework/oauth2/main.go
Pratham-Mishra04
force-pushed
the
05-18-refactor_mcp_per_user_oauth_flow_refactor
branch
from
May 19, 2026 20:10
791502b to
bb4e5d4
Compare
Pratham-Mishra04
force-pushed
the
05-18-refactor_mcp_per_user_oauth_flow_refactor
branch
from
May 20, 2026 10:18
6d1fd81 to
fc6475d
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
core/schemas/context.go (1)
302-325:⚠️ Potential issue | 🟠 Major | ⚡ Quick winProtect MCP auth identity source keys from restricted-write overrides.
Line 315, Line 318, and Line 321 read identity from context for OAuth lookup mode, but those keys are still writable when restricted writes are blocked because they are not in
reservedKeys. This allows identity override before token resolution.🔒 Suggested fix
var reservedKeys = []any{ BifrostContextKeyVirtualKey, + BifrostContextKeyGovernanceVirtualKeyID, + BifrostContextKeyUserID, + BifrostContextKeyMCPSessionID, BifrostContextKeyAPIKeyName, BifrostContextKeyAPIKeyID, BifrostContextKeyRequestID,🤖 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 `@core/schemas/context.go` around lines 302 - 325, The MCPAuthMode() method currently reads identity keys that can be overwritten by restricted-write code paths; add BifrostContextKeyUserID, BifrostContextKeyGovernanceVirtualKeyID, and BifrostContextKeyMCPSessionID to the set of reserved keys used by BifrostContext (or the existing reservedKeys collection) so these context entries cannot be modified by restricted writes, then ensure BifrostContext's Set/With methods respect that reservedKeys collection so MCPAuthMode() always reads protected, non-overridable identity values from the context.
🤖 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 `@core/schemas/bifrost.go`:
- Around line 178-186: Update the MCPAuthModeSession documentation and any
related help text to require server-minted or cryptographically high-entropy
session IDs rather than “client-issued opaque” values: change the comment for
MCPAuthModeSession to state that x-bf-mcp-session-id must be a server-generated,
unpredictable token (e.g., >=128 bits from a CSPRNG) and update any downstream
explanatory text that references x-bf-mcp-session-id to reflect this
server-minted/high-entropy requirement and guidance for generation/rotation.
---
Duplicate comments:
In `@core/schemas/context.go`:
- Around line 302-325: The MCPAuthMode() method currently reads identity keys
that can be overwritten by restricted-write code paths; add
BifrostContextKeyUserID, BifrostContextKeyGovernanceVirtualKeyID, and
BifrostContextKeyMCPSessionID to the set of reserved keys used by BifrostContext
(or the existing reservedKeys collection) so these context entries cannot be
modified by restricted writes, then ensure BifrostContext's Set/With methods
respect that reservedKeys collection so MCPAuthMode() always reads protected,
non-overridable identity values from the context.
🪄 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: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: b3b014c9-1f51-4f82-93d7-1c7c4f91ea38
📒 Files selected for processing (24)
core/mcp/clientmanager.gocore/mcp/codemode/starlark/listfiles.gocore/mcp/utils/utils.gocore/schemas/bifrost.gocore/schemas/context.gocore/schemas/mcp.gocore/schemas/oauth.goframework/configstore/clientconfig.goframework/configstore/migrations.goframework/configstore/rdb.goframework/configstore/rdb_test.goframework/configstore/store.goframework/configstore/tables/clientconfig.goframework/configstore/tables/oauth.goframework/oauth2/main.goframework/oauth2/sync.gotransports/bifrost-http/handlers/config.gotransports/bifrost-http/handlers/governance.gotransports/bifrost-http/handlers/mcp.gotransports/bifrost-http/handlers/mcp_sessions.gotransports/bifrost-http/handlers/mcpserver.gotransports/bifrost-http/handlers/oauth2.gotransports/bifrost-http/handlers/oauth2_consent.gotransports/bifrost-http/handlers/oauth2_metadata.go
💤 Files with no reviewable changes (16)
- transports/bifrost-http/handlers/mcp.go
- framework/configstore/tables/clientconfig.go
- transports/bifrost-http/handlers/oauth2_metadata.go
- framework/configstore/rdb_test.go
- transports/bifrost-http/handlers/governance.go
- transports/bifrost-http/handlers/oauth2_consent.go
- framework/oauth2/sync.go
- transports/bifrost-http/handlers/oauth2.go
- framework/configstore/store.go
- transports/bifrost-http/handlers/mcpserver.go
- framework/configstore/tables/oauth.go
- transports/bifrost-http/handlers/mcp_sessions.go
- framework/oauth2/main.go
- transports/bifrost-http/handlers/config.go
- framework/configstore/migrations.go
- framework/configstore/rdb.go
✅ Files skipped from review due to trivial changes (1)
- core/mcp/codemode/starlark/listfiles.go
akshaydeo
approved these changes
May 20, 2026
Contributor
Merge activity
|
akshaydeo
added a commit
that referenced
this pull request
May 20, 2026
## Summary This PR cuts the `v1.5.11` / `v1.3.11` release across core, framework, and all plugins, and introduces a new Claude skill (`release-checklist`) for pre-release migration safety auditing. ## Changes - **`release-checklist` skill** — Adds `.claude/skills/release-checklist/SKILL.md`, a read-only pre-release audit tool that scans Go-defined database migrations changed in a release for high-scale deadlock/lock-contention risks and boot-time-blocking operations. It produces a structured `PASS`/`WARN`/`FAIL` report with a concrete remediation plan per finding. The skill is designed to grow via an extensible Checks Registry. - **Version bumps** — `core` → `1.5.11`, `framework` → `1.3.11`, `transports` → `1.5.3`, `plugins/governance` → `1.5.11`, `plugins/logging` → `1.5.11`, `plugins/semanticcache` → `1.5.11`, `plugins/otel` → `1.2.11`, `plugins/maxim` → `1.6.11`, `plugins/prompts` → `1.0.11`, and remaining plugins bumped accordingly. - **Changelogs populated** — All per-package changelogs updated with the full set of features and fixes shipping in this release. Key highlights in this release: - Temporary access tokens for scoped, time-limited API access - MCP per-user OAuth flow refactor - Bedrock Mantle inference engine support - Azure Realtime provider with enriched session tracking - Direct access control (DAC) and virtual key rotation - Cluster-aware log metadata and per-node usage aggregation - Feature flag framework - Config-hash-based file value override of DB on restart - Semantic cache plugin rewrite - Numerous streaming stability, Bedrock, Anthropic, and Gemini fixes - AWS SDK and dependency security updates ## Type of change - [ ] Bug fix - [ ] Feature - [ ] Refactor - [ ] Documentation - [x] Chore/CI ## Affected areas - [x] Core (Go) - [x] Transports (HTTP) - [x] Providers/Integrations - [x] Plugins - [x] UI (React) - [ ] Docs ## How to test ```sh # Verify version files reflect the new release cat core/version # expect 1.5.11 cat framework/version # expect 1.3.11 cat transports/version # expect 1.5.3 # Core/Transports go test ./... ``` To exercise the new `release-checklist` skill, invoke it via Claude with: ``` /release-checklist origin/dev...HEAD ``` Expected output: a structured report with `PASS`/`WARN`/`FAIL` per check and a Remediation Plan table for any findings. ## Screenshots/Recordings N/A ## Breaking changes - [ ] Yes - [x] No ## Related issues #3603, #3565, #3489, #3334, #3335, #3435, #3554, #3590, #3444, #3198, #3581, #3610, #3599, #3567, #3382, #3461 and others listed in the changelogs. ## Security considerations - AWS SDK and dependency security updates are included (#3461). - `FullyRedacted()` for proxy passwords and `MarshalForStorage()` for `ProxyConfig` prevent partial secret leakage in API responses (#3445). - The `release-checklist` skill is strictly read-only and never modifies files. ## Checklist - [x] I read `docs/contributing/README.md` and followed the guidelines - [x] I added/updated tests where appropriate - [x] I updated documentation where needed - [x] I verified builds succeed (Go and UI) - [ ] I verified the CI pipeline passes locally if applicable
Merged
This was referenced May 21, 2026
18 tasks
18 tasks
This was referenced Jun 17, 2026
Merged
This was referenced Jul 2, 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
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
Briefly explain the purpose of this PR and the problem it solves.
Changes
Type of change
Affected areas
How to test
Describe the steps to validate this change. Include commands and expected outcomes.
If adding new configs or environment variables, document them here.
Screenshots/Recordings
If UI changes, add before/after screenshots or short clips.
Breaking changes
If yes, describe impact and migration instructions.
Related issues
Link related issues and discussions. Example: Closes #123
Security considerations
Note any security implications (auth, secrets, PII, sandboxing, etc.).
Checklist
docs/contributing/README.mdand followed the guidelines