Skip to content

feat: add OAuth2 session listing, revocation, family-revocation, VK liveness check, and sweep worker#4509

Open
Pratham-Mishra04 wants to merge 1 commit into
06-17-feat_wires_jwt_authentication_on__mcp_endpointfrom
06-17-feat_adds_mcp_oauth2_server_lifecycle_token_rotation_liveness_checks_sessions_api
Open

feat: add OAuth2 session listing, revocation, family-revocation, VK liveness check, and sweep worker#4509
Pratham-Mishra04 wants to merge 1 commit into
06-17-feat_wires_jwt_authentication_on__mcp_endpointfrom
06-17-feat_adds_mcp_oauth2_server_lifecycle_token_rotation_liveness_checks_sessions_api

Conversation

@Pratham-Mishra04

Copy link
Copy Markdown
Collaborator

Summary

This PR adds OAuth2 session management capabilities including stolen-token detection with automatic family revocation, a Connected Clients API for listing and revoking active downstream grants, VK liveness checks on token refresh, and a background sweep worker to clean up revoked refresh tokens.

Changes

  • Added GetOAuth2RefreshTokenByHashAny to look up refresh tokens including revoked ones, enabling detection of token reuse attacks where a previously rotated token is re-presented
  • When a revoked token is re-presented during refresh, all active tokens in the same family are now revoked per RFC 9700 §2.2.2 to limit damage from a potentially stolen token
  • Added RevokeOAuth2RefreshTokensByFamilyID and RevokeOAuth2RefreshTokensByMode for bulk revocation of token families and mode-scoped grants respectively
  • Added SweepOAuth2RefreshTokens to delete revoked tokens older than a configurable retention window (default 30 days), swept every 10 minutes via a new oauth2SweepWorker
  • Added ListOAuth2Sessions which joins refresh token rows with client names and VK names for human-readable display in the Connected Clients UI
  • Added GetOAuth2SessionByID and RevokeOAuth2Session for single-session lookup and revocation
  • Added OAuth2SessionsHandler exposing GET /api/oauth2/sessions and DELETE /api/oauth2/sessions/{id}; user-mode sessions enforce that the caller's identity matches bf_sub before allowing revocation
  • Added a VK liveness check during token refresh: if the virtual key referenced by a VK-mode token has been deleted or disabled, the refresh is rejected with invalid_grant
  • The oauth2SweepWorker is started during server bootstrap and stopped cleanly on all error and shutdown paths alongside the existing temp-token sweep worker

Type of change

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

Affected areas

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

How to test

go test ./...
  1. Issue a refresh token and rotate it once. Re-present the original (now revoked) token to POST /token — expect invalid_grant and all sibling tokens in the family to be revoked.
  2. Disable or delete a virtual key that has an active VK-mode refresh token. Attempt a token refresh — expect invalid_grant with "virtual key is no longer active".
  3. GET /api/oauth2/sessions — expect a JSON list of active grants with client names and VK display names populated.
  4. DELETE /api/oauth2/sessions/{id} with a user-mode session ID using a mismatched caller identity — expect 403 Forbidden.
  5. DELETE /api/oauth2/sessions/{id} with the correct caller identity — expect 204 No Content and the session absent from subsequent list calls.
  6. Wait for or trigger the sweep worker; confirm revoked tokens older than the retention window are removed from the database.

Breaking changes

  • Yes
  • No

The ConfigStore interface has new required methods. Any custom implementations of ConfigStore must add GetOAuth2RefreshTokenByHashAny, RevokeOAuth2RefreshTokensByFamilyID, RevokeOAuth2RefreshTokensByMode, SweepOAuth2RefreshTokens, ListOAuth2Sessions, GetOAuth2SessionByID, and RevokeOAuth2Session.

Related issues

Security considerations

  • Implements RFC 9700 §2.2.2 refresh token family revocation: re-use of a revoked token triggers revocation of all active tokens in the family, limiting the blast radius of a stolen token.
  • The session revoke endpoint enforces identity matching for user-mode grants, preventing one user from revoking another user's session even if they share visibility of the row.
  • VK liveness checks prevent deleted or disabled virtual keys from silently continuing to obtain access tokens via refresh.

Checklist

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

@CLAassistant

Copy link
Copy Markdown

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

Pratham-Mishra04 commented Jun 17, 2026

Copy link
Copy Markdown
Collaborator Author

@coderabbitai

coderabbitai Bot commented Jun 17, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 73439e47-7b85-4a93-a7e8-e29215151dcf

📥 Commits

Reviewing files that changed from the base of the PR and between 1c19dfb and 0ecdc90.

📒 Files selected for processing (10)
  • framework/configstore/rdb.go
  • framework/configstore/store.go
  • framework/mcp_headers/sweep.go
  • framework/oauth2/sync.go
  • framework/temptoken/sweeper.go
  • transports/bifrost-http/handlers/oauth2_issuance.go
  • transports/bifrost-http/handlers/oauth2_sessions.go
  • transports/bifrost-http/lib/config_test.go
  • transports/bifrost-http/server/oauth2.go
  • transports/bifrost-http/server/server.go
🚧 Files skipped from review as they are similar to previous changes (10)
  • framework/mcp_headers/sweep.go
  • framework/oauth2/sync.go
  • framework/temptoken/sweeper.go
  • framework/configstore/store.go
  • transports/bifrost-http/lib/config_test.go
  • transports/bifrost-http/handlers/oauth2_issuance.go
  • transports/bifrost-http/handlers/oauth2_sessions.go
  • transports/bifrost-http/server/server.go
  • transports/bifrost-http/server/oauth2.go
  • framework/configstore/rdb.go

📝 Walkthrough

Summary by CodeRabbit

  • New Features
    • Added Connected Clients session APIs: GET /api/oauth2/sessions to list active sessions and DELETE /api/oauth2/sessions/{id} to revoke a session.
  • Improvements
    • Enhanced refresh-token refresh flow to detect revoked tokens (returns invalid_grant when appropriate) and revoke related active sessions by family or mode.
    • Added VK-mode liveness checks to prevent minting when the referenced virtual key is no longer active.
  • Bug Fixes
    • Improved error handling for session revocation and lookups (more precise 404/403/400/500 responses).
  • Chores
    • Improved shutdown behavior for OAuth2-related background sweeps.

Walkthrough

Adds OAuth2 connected-clients session management with refresh-token breach detection and session HTTP endpoints. Extends ConfigStore with refresh-token lookup (including revoked state), family/mode revocation, session listing with client joins, and session fetch/revoke operations. Adds breach detection in token refresh to revoke token families when a revoked token is replayed, and VK-mode liveness checking. Introduces OAuth2SessionsHandler for listing and revoking sessions via HTTP, wires a background oauth2SweepWorker into the server lifecycle with 30-day revoked-token retention. Also applies graceful shutdown context-cancellation pattern to sweep workers across multiple framework components.

Changes

OAuth2 Session Management & Breach Detection

Layer / File(s) Summary
ConfigStore interface OAuth2 extensions
framework/configstore/store.go
ConfigStore interface gains eight new OAuth2 methods: GetOAuth2RefreshTokenByHashAny, RevokeOAuth2RefreshTokensByFamilyID, RevokeOAuth2RefreshTokensByMode, SweepOAuth2RefreshTokens, SweepOrphanedOAuth2Clients, ListOAuth2Sessions, GetOAuth2SessionByID, and RevokeOAuth2Session.
RDB refresh-token lookup and revocation operations
framework/configstore/rdb.go
GetOAuth2RefreshTokenByHashAny looks up tokens across any revocation state; RevokeOAuth2RefreshTokensByFamilyID and RevokeOAuth2RefreshTokensByMode bulk-revoke active tokens; SweepOAuth2RefreshTokens hard-deletes revoked tokens older than a supplied duration.
OAuth2SessionRow type and session DB operations
framework/configstore/rdb.go
Defines OAuth2SessionRow with client and VK display names; ListOAuth2Sessions joins refresh-token rows with oauth2_clients and governance_virtual_keys; RevokeOAuth2Session and GetOAuth2SessionByID operate on non-revoked rows; SweepOrphanedOAuth2Clients deletes dynamically-registered clients with no active tokens.
Token refresh breach detection and VK liveness check
transports/bifrost-http/handlers/oauth2_issuance.go
When a refresh token is not found in active state, probes whether it was previously revoked and revokes the entire family if so; for VK-mode refreshes, validates that the referenced virtual key exists and is active before minting new tokens.
OAuth2 sessions HTTP handler and routes
transports/bifrost-http/handlers/oauth2_sessions.go
OAuth2SessionsHandler exposes GET /api/oauth2/sessions to list active sessions and DELETE /api/oauth2/sessions/{id} to revoke a session; delete enforces caller-identity matching for user-mode sessions and maps not-found/already-revoked to 404.
OAuth2 sweep worker implementation
transports/bifrost-http/server/oauth2.go
oauth2SweepWorker runs every 10 minutes sweeping authorize-request expiration, revoked-refresh-token deletion (30-day retention), and orphaned-client cleanup (1-hour grace); implements graceful start/stop with context cancellation.
Server lifecycle wiring for sweep worker and sessions routes
transports/bifrost-http/server/server.go
BifrostHTTPServer adds OAuth2SweepWorker field; Bootstrap constructs and starts the worker with explicit cleanup on init failures; graceful shutdown stops and clears the worker; RegisterAPIRoutes registers OAuth2SessionsHandler routes.
Mock store OAuth2 method stubs
transports/bifrost-http/lib/config_test.go
Extends MockConfigStore with no-op implementations of all new OAuth2 methods for testing.

Graceful Shutdown Pattern for Sweep Workers

Layer / File(s) Summary
MCP CredentialSweepWorker context cancellation
framework/mcp_headers/sweep.go
CredentialSweepWorker derives a cancellable context in Start, stores the cancel function, and runs the sweep loop with that context; Stop invokes cancel before closing the stop channel.
OAuth2 sync workers context cancellation
framework/oauth2/sync.go
TokenRefreshWorker and PerUserOAuthSweepWorker derive cancellable contexts in Start, store cancel functions, and run with those contexts; Stop invokes cancel before closing the stop channel.
Temp-token SweepWorker context cancellation
framework/temptoken/sweeper.go
SweepWorker derives a cancellable context in Start, stores the cancel function, and runs the sweep loop with that context; Stop invokes cancel before closing the stop channel.

Sequence Diagrams

sequenceDiagram
  participant Client
  participant OAuth2SessionsHandler
  participant ConfigStore
  participant DB

  rect rgba(173, 216, 230, 0.5)
    note over Client,DB: List Active Sessions
    Client->>OAuth2SessionsHandler: GET /api/oauth2/sessions
    OAuth2SessionsHandler->>ConfigStore: ListOAuth2Sessions(ctx)
    ConfigStore->>DB: SELECT rt JOIN oauth2_clients JOIN virtual_keys
    DB-->>ConfigStore: []OAuth2SessionRow
    OAuth2SessionsHandler-->>Client: 200 {"sessions": [...]}
  end

  rect rgba(255, 200, 150, 0.5)
    note over Client,DB: Revoke Session
    Client->>OAuth2SessionsHandler: DELETE /api/oauth2/sessions/{id}
    OAuth2SessionsHandler->>ConfigStore: GetOAuth2SessionByID(ctx, id)
    ConfigStore-->>OAuth2SessionsHandler: session (non-revoked)
    OAuth2SessionsHandler->>ConfigStore: RevokeOAuth2Session(ctx, id)
    ConfigStore->>DB: UPDATE revoked_at=now WHERE id AND revoked_at IS NULL
    OAuth2SessionsHandler-->>Client: 204 No Content
  end

  rect rgba(200, 230, 180, 0.5)
    note over Client,DB: Token Refresh with Breach Detection
    Client->>handleTokenRefresh: POST /token (refresh_token)
    handleTokenRefresh->>ConfigStore: GetOAuth2RefreshTokenByHash(active)
    alt not found
      handleTokenRefresh->>ConfigStore: GetOAuth2RefreshTokenByHashAny
      alt was revoked
        handleTokenRefresh->>ConfigStore: RevokeOAuth2RefreshTokensByFamilyID
        handleTokenRefresh-->>Client: 400 invalid_grant (breach)
      else never issued
        handleTokenRefresh-->>Client: 400 invalid_grant
      end
    else found, check VK if mode=vk
      handleTokenRefresh->>ConfigStore: GetVirtualKey(bf_sub)
      alt missing or inactive
        handleTokenRefresh-->>Client: 400 invalid_grant
      else active
        handleTokenRefresh-->>Client: 200 new tokens
      end
    end
  end
Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~50 minutes

Suggested reviewers

  • akshaydeo
  • danpiths

Poem

🐇 Hopping through the token garden, I detect each breach,
A family revoked when stale tokens within reach.
Sessions list, sessions vanish — order maintained with care,
Ten-minute sweeps keep the database bright and fair.
🌿 Graceful shutdown cancels in-flight work with grace — the warren's clean once more!

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Title check ✅ Passed The title concisely and accurately summarizes all major changes in the PR: OAuth2 session listing, revocation, family-revocation, VK liveness check, and sweep worker.
Description check ✅ Passed The PR description is comprehensive and well-structured, covering summary, changes, type of change, affected areas, testing steps, breaking changes, and security considerations. All required template sections are present and filled out appropriately.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
📝 Generate docstrings
  • Create stacked PR
  • Commit on current branch
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch 06-17-feat_adds_mcp_oauth2_server_lifecycle_token_rotation_liveness_checks_sessions_api

Warning

There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure.

🔧 golangci-lint (2.12.2)

level=error msg="[linters_context] typechecking error: pattern ./...: directory prefix . does not contain main module or its selected dependencies"


Comment @coderabbitai help to get the list of available commands and usage tips.

@greptile-apps

greptile-apps Bot commented Jun 17, 2026

Copy link
Copy Markdown
Contributor

Confidence Score: 5/5

Safe to merge. The security-critical paths (family revocation on token replay, VK liveness check, user-mode identity gating) are logically correct and the worker goroutine lifecycle is properly wired.

All three concerns raised in earlier review threads are addressed in the current revision. The remaining findings are a missing zero-duration guard on SweepOrphanedOAuth2Clients (a cosmetic inconsistency with its sibling function) and absent unit tests for the new handler paths — neither affects runtime correctness.

framework/configstore/rdb.go (SweepOrphanedOAuth2Clients missing zero-guard) and transports/bifrost-http/lib/config_test.go (no tests for new session/revocation handlers).

Important Files Changed

Filename Overview
framework/configstore/rdb.go Adds seven new RDB methods: revoked-token lookup, family/mode revocation, token and orphaned-client sweeps, session list/get/revoke. Logic is correct; hard deletes are safe because the table has no soft-delete (DeletedAt) field. SweepOrphanedOAuth2Clients lacks the zero-duration guard present on SweepOAuth2RefreshTokens.
framework/configstore/store.go Extends ConfigStore interface with seven new methods. Breaking change is acknowledged in the PR description. Mock in config_test.go is updated with stub implementations.
transports/bifrost-http/handlers/oauth2_issuance.go Adds RFC 9700 §2.2.2 family revocation on revoked-token replay and VK liveness check on refresh. The transient-DB-error / ErrNotFound separation addressed in the previous thread is present in the new code.
transports/bifrost-http/handlers/oauth2_sessions.go New handler for GET/DELETE /api/oauth2/sessions. Identity gating for user-mode is correct. Marshal error is handled. listSessions returns 200 with no explicit SetStatusCode (fasthttp defaults to 200 — fine).
transports/bifrost-http/server/oauth2.go New oauth2SweepWorker mirrors the temptoken sweep pattern. Runs an immediate sweep on start then every 10 minutes. Shutdown paths (cancel + stopCh) are consistent with other workers.
transports/bifrost-http/server/server.go Integrates OAuth2SweepWorker into Bootstrap and Start. All three Bootstrap error paths and the graceful-shutdown path each have their own top-level nil guard for OAuth2SweepWorker, independent of TempTokenSweepWorker.
framework/oauth2/sync.go Adds cancel context propagation to TokenRefreshWorker and PerUserOAuthSweepWorker so in-flight DB calls are cancelled promptly on Stop().
framework/mcp_headers/sweep.go Same cancel-context pattern added to CredentialSweepWorker. Straightforward and correct.
framework/temptoken/sweeper.go Adds cancel context propagation to the temp-token SweepWorker. Consistent with all other workers updated in this PR.
transports/bifrost-http/lib/config_test.go MockConfigStore updated with stub implementations for all new interface methods. Stubs compile and satisfy the interface but do not exercise new logic.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant Client
    participant TokenEndpoint as POST /oauth2/token
    participant Store as ConfigStore
    participant SweepWorker as oauth2SweepWorker

    Client->>TokenEndpoint: refresh_token (possibly revoked)
    TokenEndpoint->>Store: GetOAuth2RefreshTokenByHash (active only)
    alt Token active
        Store-->>TokenEndpoint: token row
        TokenEndpoint->>Store: GetVirtualKey (VK mode only)
        alt VK deleted/disabled
            Store-->>TokenEndpoint: ErrNotFound / inactive
            TokenEndpoint-->>Client: invalid_grant
        else VK healthy
            Store-->>TokenEndpoint: VK row
            TokenEndpoint->>Store: RotateOAuth2RefreshToken
            TokenEndpoint-->>Client: new access + refresh token
        end
    else ErrNotFound (may be revoked)
        Store-->>TokenEndpoint: ErrNotFound
        TokenEndpoint->>Store: GetOAuth2RefreshTokenByHashAny
        alt Found (revoked token reuse)
            Store-->>TokenEndpoint: revoked token row
            TokenEndpoint->>Store: RevokeOAuth2RefreshTokensByFamilyID
            TokenEndpoint-->>Client: invalid_grant (family revoked)
        else Not found at all
            TokenEndpoint-->>Client: invalid_grant
        end
    end

    loop every 10 minutes
        SweepWorker->>Store: SweepExpiredOAuth2AuthorizeRequests
        SweepWorker->>Store: "SweepOAuth2RefreshTokens (revoked > 30d)"
        SweepWorker->>Store: "SweepOrphanedOAuth2Clients (no tokens, > 1h old)"
    end
Loading
%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant Client
    participant TokenEndpoint as POST /oauth2/token
    participant Store as ConfigStore
    participant SweepWorker as oauth2SweepWorker

    Client->>TokenEndpoint: refresh_token (possibly revoked)
    TokenEndpoint->>Store: GetOAuth2RefreshTokenByHash (active only)
    alt Token active
        Store-->>TokenEndpoint: token row
        TokenEndpoint->>Store: GetVirtualKey (VK mode only)
        alt VK deleted/disabled
            Store-->>TokenEndpoint: ErrNotFound / inactive
            TokenEndpoint-->>Client: invalid_grant
        else VK healthy
            Store-->>TokenEndpoint: VK row
            TokenEndpoint->>Store: RotateOAuth2RefreshToken
            TokenEndpoint-->>Client: new access + refresh token
        end
    else ErrNotFound (may be revoked)
        Store-->>TokenEndpoint: ErrNotFound
        TokenEndpoint->>Store: GetOAuth2RefreshTokenByHashAny
        alt Found (revoked token reuse)
            Store-->>TokenEndpoint: revoked token row
            TokenEndpoint->>Store: RevokeOAuth2RefreshTokensByFamilyID
            TokenEndpoint-->>Client: invalid_grant (family revoked)
        else Not found at all
            TokenEndpoint-->>Client: invalid_grant
        end
    end

    loop every 10 minutes
        SweepWorker->>Store: SweepExpiredOAuth2AuthorizeRequests
        SweepWorker->>Store: "SweepOAuth2RefreshTokens (revoked > 30d)"
        SweepWorker->>Store: "SweepOrphanedOAuth2Clients (no tokens, > 1h old)"
    end
Loading

Reviews (6): Last reviewed commit: "feat: adds mcp oauth2 server lifecycle —..." | Re-trigger Greptile

Comment thread transports/bifrost-http/handlers/oauth2_issuance.go Outdated
Comment thread transports/bifrost-http/server/server.go
Comment thread transports/bifrost-http/handlers/oauth2_sessions.go

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 6

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)

6882-6888: ⚠️ Potential issue | 🔴 Critical

Add single-use guard to refresh-token rotation under concurrency.

The rotation update on id alone allows two concurrent refresh requests to both read the old active token, both update it, and both insert active descendants. This bypasses the revoked-token reuse detection because neither transaction observes ErrNotFound.

The current error handling (line 374 in oauth2_issuance.go) also maps any rotation failure to "server_error" without distinguishing a rotation conflict from other errors, preventing proper family revocation when rotation races are detected.

Add AND revoked_at IS NULL to the rotation update's Where clause and check RowsAffected == 0 to return ErrNotFound when the old token is no longer active. The caller must then map this to "invalid_grant" and revoke the family using the already-loaded rt.FamilyID.

🔒 Required fix
 func (s *RDBConfigStore) RotateOAuth2RefreshToken(ctx context.Context, oldID string, newRT *tables.TableOAuth2RefreshToken) error {
 	now := time.Now()
 	return s.DB().WithContext(ctx).Transaction(func(tx *gorm.DB) error {
-		if err := tx.Model(&tables.TableOAuth2RefreshToken{}).
-			Where("id = ?", oldID).
-			Update("revoked_at", &now).Error; err != nil {
-			return fmt.Errorf("revoke old refresh token: %w", err)
+		result := tx.Model(&tables.TableOAuth2RefreshToken{}).
+			Where("id = ? AND revoked_at IS NULL", oldID).
+			Update("revoked_at", &now)
+		if result.Error != nil {
+			return fmt.Errorf("revoke old refresh token: %w", result.Error)
+		}
+		if result.RowsAffected == 0 {
+			return ErrNotFound
 		}
 		if err := tx.Create(newRT).Error; err != nil {
 			return fmt.Errorf("create new refresh token: %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/rdb.go` around lines 6882 - 6888, The Update operation
on tables.TableOAuth2RefreshToken only checks the token id in its Where clause,
allowing concurrent refresh requests to both read and update the same old token.
Add AND revoked_at IS NULL to the Where clause to ensure only active tokens are
updated, then check if the tx result RowsAffected equals 0 after the update and
return ErrNotFound in that case. This prevents concurrent rotation races from
both inserting new active descendants and allows the caller to properly
distinguish rotation conflicts from other errors for family revocation.
🧹 Nitpick comments (1)
framework/configstore/rdb.go (1)

6960-6973: 🏗️ Heavy lift

Bound the connected-clients session query.

ListOAuth2Sessions scans every active refresh token and performs two joins with no limit/offset. Since active grants have no expiry column in TableOAuth2RefreshToken, this new API path can grow unbounded; add pagination with a stable tiebreaker before exposing it broadly.

♻️ Suggested shape
-func (s *RDBConfigStore) ListOAuth2Sessions(ctx context.Context) ([]OAuth2SessionRow, error) {
+func (s *RDBConfigStore) ListOAuth2Sessions(ctx context.Context, limit, offset int) ([]OAuth2SessionRow, error) {
+	if limit <= 0 {
+		limit = 25
+	} else if limit > 100 {
+		limit = 100
+	}
+	if offset < 0 {
+		offset = 0
+	}
 	rows := []struct {
 		tables.TableOAuth2RefreshToken
 		ClientName string `gorm:"column:client_name"`
 		VKName     string `gorm:"column:vk_name"`
 	}{}
 	err := s.ScopedDB(ctx).
@@
 		Joins("LEFT JOIN oauth2_clients c ON c.client_id = rt.client_id").
 		Joins("LEFT JOIN governance_virtual_keys vk ON vk.id = rt.bf_sub AND rt.bf_mode = 'vk'").
 		Where("rt.revoked_at IS NULL").
-		Order("rt.created_at DESC").
+		Order("rt.created_at DESC, rt.id DESC").
+		Limit(limit).
+		Offset(offset).
 		Scan(&rows).Error

Apply the same contract through the interface and handler response. As per coding guidelines, framework/** changes should use bounded resource usage.

🤖 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 6960 - 6973, The
ListOAuth2Sessions function in RDBConfigStore queries all active refresh tokens
with no pagination limits, which can result in unbounded resource usage. Add
pagination support by accepting limit and offset parameters to the function,
apply these to the database query before the Scan call, and add a stable
tiebreaker to the Order clause (using rt.id in addition to rt.created_at DESC)
for consistent ordering across paginated requests. Update the interface contract
and handler response to accept and return pagination parameters consistently
with other bounded API endpoints.

Source: Coding guidelines

🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@framework/configstore/rdb.go`:
- Around line 6934-6939: The SweepOAuth2RefreshTokens method does not validate
that revokedOlderThan is a positive duration. When revokedOlderThan is zero or
negative, the cutoff time becomes now or in the future, causing the deletion
query to immediately remove nearly all revoked tokens instead of just old ones,
which breaks replay detection. Add a guard at the beginning of the
SweepOAuth2RefreshTokens method to check if revokedOlderThan is positive and
return an error if it is not.

In `@transports/bifrost-http/handlers/oauth2_issuance.go`:
- Around line 327-328: The code at the breach-detection block suppresses
critical store errors instead of handling them explicitly. The call to
GetOAuth2RefreshTokenByHashAny should be handled to address the lookupErr case
beyond just the success path, and the RevokeOAuth2RefreshTokensByFamilyID call
should not discard its error using the blank identifier. Instead, capture and
properly handle the error returned from RevokeOAuth2RefreshTokensByFamilyID to
ensure fail-closed behavior for this security-critical family token revocation
operation. According to coding guidelines, transport persistence and security
changes must have explicit error handling.
- Around line 346-353: In the GetVirtualKey error handling block in
oauth2_issuance.go, reorder the condition checks so that non-ErrNotFound lookup
errors are handled first by returning server_error, before checking if the
virtual key is nil or inactive. The current logic combines the nil check with
ErrNotFound in a single condition, causing transient store failures that return
(nil, err) to incorrectly send invalid_grant instead of server_error. Separate
the vkErr != nil check (excluding ErrNotFound) to execute first with
server_error response, then check vk == nil or !vk.IsActiveValue() for the
invalid_grant response.

In `@transports/bifrost-http/handlers/oauth2_sessions.go`:
- Around line 44-45: The sonic.Marshal call on line 44 ignores the error by
using underscore assignment, which means if JSON serialization fails, the code
will still set an empty or invalid body while returning a success status.
Capture the error returned by sonic.Marshal along with the data variable, check
if the error is not nil, and if so return a 500 status code using ctx methods
before proceeding to ctx.SetBody. Only call ctx.SetBody(data) when the marshal
operation succeeds to ensure the API fails closed on serialization errors.

In `@transports/bifrost-http/server/oauth2.go`:
- Around line 34-53: The oauth2SweepWorker does not cancel in-flight sweep
operations when stop() is called, only closing stopCh. If sweep(ctx) is blocked
on a DB operation, the goroutine remains alive indefinitely. Create a
cancellable context derived from the passed ctx parameter in the run method
using context.WithCancel, store the resulting cancel function as a field in the
worker, and modify the stop() method to call this cancel function in addition to
closing stopCh. This ensures any blocking sweep(ctx) call receives the
cancellation signal through ctx.Done() and can terminate promptly during
shutdown.

In `@transports/bifrost-http/server/server.go`:
- Around line 1764-1767: The OAuth2SweepWorker cleanup block is incorrectly
nested under a condition checking if TempTokenSweepWorker exists, which means
OAuth2 cleanup will be skipped if only OAuth2SweepWorker was created but
TempTokenSweepWorker is nil. Move the OAuth2SweepWorker cleanup code (the check
for s.OAuth2SweepWorker != nil, followed by stop() and nil assignment) to
execute independently without being gated by any TempTokenSweepWorker condition.
Apply this fix at all four locations: the block at lines 1764-1767 and the three
additional locations mentioned (1798-1801, 1836-1839, 1933-1936).

---

Outside diff comments:
In `@framework/configstore/rdb.go`:
- Around line 6882-6888: The Update operation on tables.TableOAuth2RefreshToken
only checks the token id in its Where clause, allowing concurrent refresh
requests to both read and update the same old token. Add AND revoked_at IS NULL
to the Where clause to ensure only active tokens are updated, then check if the
tx result RowsAffected equals 0 after the update and return ErrNotFound in that
case. This prevents concurrent rotation races from both inserting new active
descendants and allows the caller to properly distinguish rotation conflicts
from other errors for family revocation.

---

Nitpick comments:
In `@framework/configstore/rdb.go`:
- Around line 6960-6973: The ListOAuth2Sessions function in RDBConfigStore
queries all active refresh tokens with no pagination limits, which can result in
unbounded resource usage. Add pagination support by accepting limit and offset
parameters to the function, apply these to the database query before the Scan
call, and add a stable tiebreaker to the Order clause (using rt.id in addition
to rt.created_at DESC) for consistent ordering across paginated requests. Update
the interface contract and handler response to accept and return pagination
parameters consistently with other bounded API endpoints.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: ae0ec375-e9c9-4ecc-ab89-f1fe7617bb4d

📥 Commits

Reviewing files that changed from the base of the PR and between d6a46f9 and 98169a4.

📒 Files selected for processing (6)
  • framework/configstore/rdb.go
  • framework/configstore/store.go
  • transports/bifrost-http/handlers/oauth2_issuance.go
  • transports/bifrost-http/handlers/oauth2_sessions.go
  • transports/bifrost-http/server/oauth2.go
  • transports/bifrost-http/server/server.go

Comment thread framework/configstore/rdb.go
Comment thread transports/bifrost-http/handlers/oauth2_issuance.go Outdated
Comment thread transports/bifrost-http/handlers/oauth2_issuance.go
Comment thread transports/bifrost-http/handlers/oauth2_sessions.go Outdated
Comment thread transports/bifrost-http/server/oauth2.go
Comment thread transports/bifrost-http/server/server.go Outdated
@Pratham-Mishra04 Pratham-Mishra04 force-pushed the 06-17-feat_wires_jwt_authentication_on__mcp_endpoint branch from d6a46f9 to 5d1f810 Compare June 18, 2026 07:34
@Pratham-Mishra04 Pratham-Mishra04 force-pushed the 06-17-feat_adds_mcp_oauth2_server_lifecycle_token_rotation_liveness_checks_sessions_api branch from 98169a4 to be01ec8 Compare June 18, 2026 07:34
@Pratham-Mishra04 Pratham-Mishra04 force-pushed the 06-17-feat_wires_jwt_authentication_on__mcp_endpoint branch from 5d1f810 to bc90e18 Compare June 18, 2026 07:37
@Pratham-Mishra04 Pratham-Mishra04 force-pushed the 06-17-feat_adds_mcp_oauth2_server_lifecycle_token_rotation_liveness_checks_sessions_api branch from be01ec8 to 9282613 Compare June 18, 2026 07:37

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick comments (1)
framework/configstore/rdb.go (1)

7016-7025: 💤 Low value

Consider adding a guard for non-positive duration for consistency.

SweepOAuth2RefreshTokens guards against revokedOlderThan <= 0 to prevent unintended mass deletion. This method lacks an equivalent guard—if registeredOlderThan <= 0, the cutoff becomes now or later, deleting all orphaned clients immediately rather than respecting a grace window.

The risk is lower here (no security detection dependency), but adding a similar guard would maintain consistency across sweep methods.

♻️ Suggested guard for consistency
 func (s *RDBConfigStore) SweepOrphanedOAuth2Clients(ctx context.Context, registeredOlderThan time.Duration) (int64, error) {
+	if registeredOlderThan <= 0 {
+		return 0, nil
+	}
 	cutoff := time.Now().Add(-registeredOlderThan)
🤖 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 7016 - 7025, The
SweepOrphanedOAuth2Clients method lacks a guard against non-positive duration
values for the registeredOlderThan parameter, which could lead to unintended
mass deletion of all orphaned clients. Add a validation guard at the beginning
of the SweepOrphanedOAuth2Clients function (similar to what exists in
SweepOAuth2RefreshTokens) to check if registeredOlderThan is less than or equal
to zero and return an error in such cases to maintain consistency across sweep
methods and prevent accidental data loss.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@framework/configstore/rdb.go`:
- Around line 7016-7025: The SweepOrphanedOAuth2Clients method lacks a guard
against non-positive duration values for the registeredOlderThan parameter,
which could lead to unintended mass deletion of all orphaned clients. Add a
validation guard at the beginning of the SweepOrphanedOAuth2Clients function
(similar to what exists in SweepOAuth2RefreshTokens) to check if
registeredOlderThan is less than or equal to zero and return an error in such
cases to maintain consistency across sweep methods and prevent accidental data
loss.

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: ba94ffc9-a891-4373-b716-137921d834cb

📥 Commits

Reviewing files that changed from the base of the PR and between 98169a4 and be01ec8.

📒 Files selected for processing (10)
  • framework/configstore/rdb.go
  • framework/configstore/store.go
  • framework/mcp_headers/sweep.go
  • framework/oauth2/sync.go
  • framework/temptoken/sweeper.go
  • transports/bifrost-http/handlers/oauth2_issuance.go
  • transports/bifrost-http/handlers/oauth2_sessions.go
  • transports/bifrost-http/lib/config_test.go
  • transports/bifrost-http/server/oauth2.go
  • transports/bifrost-http/server/server.go
🚧 Files skipped from review as they are similar to previous changes (5)
  • transports/bifrost-http/server/oauth2.go
  • framework/configstore/store.go
  • transports/bifrost-http/handlers/oauth2_sessions.go
  • transports/bifrost-http/server/server.go
  • transports/bifrost-http/handlers/oauth2_issuance.go

@Pratham-Mishra04 Pratham-Mishra04 force-pushed the 06-17-feat_wires_jwt_authentication_on__mcp_endpoint branch from b156095 to 8a84174 Compare June 18, 2026 12:22
@Pratham-Mishra04 Pratham-Mishra04 force-pushed the 06-17-feat_adds_mcp_oauth2_server_lifecycle_token_rotation_liveness_checks_sessions_api branch from 0ecdc90 to 4178622 Compare June 18, 2026 12:22
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants