From 7e5cb69916af44a58332c112ceb7efba553996af Mon Sep 17 00:00:00 2001 From: Pratham-Mishra04 Date: Thu, 30 Jul 2026 13:24:31 +0530 Subject: [PATCH] fix: propagate ctx through userTokenCache.Fill so a canceled request unblocks instead of waiting on an unrelated leader --- framework/oauth2/main.go | 133 +++- framework/oauth2/usertokencache.go | 197 ++++++ framework/oauth2/usertokencache_test.go | 594 ++++++++++++++++++ .../bifrost-http/handlers/governance.go | 20 +- transports/bifrost-http/handlers/mcp.go | 60 +- .../bifrost-http/handlers/mcpsessions.go | 21 +- transports/bifrost-http/server/server.go | 57 +- 7 files changed, 1046 insertions(+), 36 deletions(-) create mode 100644 framework/oauth2/usertokencache.go create mode 100644 framework/oauth2/usertokencache_test.go diff --git a/framework/oauth2/main.go b/framework/oauth2/main.go index 3be0d9e028b..f51204eb944 100644 --- a/framework/oauth2/main.go +++ b/framework/oauth2/main.go @@ -68,6 +68,15 @@ type OAuth2Provider struct { // RevokeToken's database I/O. Sharing p.mu would stall flow init/cleanup // reads behind unrelated revoke traffic. tempTokens atomic.Pointer[temptoken.Service] + + // userTokens caches per-user access-token lookups keyed by + // (auth mode, identity, mcp client). It owns its own locking and must + // never be guarded by p.mu, for the same reason tempTokens is not: + // p.mu is write-locked across token-endpoint network I/O, and a cached + // read that had to wait on it would lose the entire point of the cache. + // Write paths in this file evict inline; handler-level database writes + // that bypass the provider evict through the exported eviction methods. + userTokens *userTokenCache } // NewOAuth2Provider creates a new OAuth provider instance @@ -79,6 +88,7 @@ func NewOAuth2Provider(configStore configstore.ConfigStore, logger schemas.Logge return &OAuth2Provider{ configStore: configStore, retryBaseDelay: time.Second, + userTokens: newUserTokenCache(defaultUserTokenCacheCapacity), } } @@ -309,6 +319,9 @@ func (p *OAuth2Provider) refreshAccessTokenLocked(ctx context.Context, tokenID s return fmt.Errorf("oauth refresh permanently rejected but status update failed (mcp_client=%s auth_mode=%s upstream_status=%d): %w", token.MCPClientID, token.AuthMode, permErr.StatusCode, markErr) } + // The row is no longer 'active'; drop any cached copy so lookups + // surface the re-auth requirement instead of a dead access token. + p.EvictUserTokenByID(token.ID) logger.Debug("OAuth refresh permanently rejected; token marked needs_reauth: mcp_client=%s auth_mode=%s upstream_status=%d", token.MCPClientID, token.AuthMode, permErr.StatusCode) return fmt.Errorf("refresh token rejected by upstream OAuth server, re-authentication required: %w", schemas.ErrOAuth2TokenExpired) @@ -344,6 +357,10 @@ func (p *OAuth2Provider) refreshAccessTokenLocked(ctx context.Context, tokenID s return fmt.Errorf("token was reauthorized or rotated during refresh, discarding this refresh: %w", schemas.ErrOAuth2TokenExpired) } + // Drop any cached copy of the old access token; the next lookup reads + // the freshly written row. + p.EvictUserTokenByID(token.ID) + logger.Debug("OAuth token refreshed successfully: token_id=%s auth_mode=%s", token.ID, token.AuthMode) return nil @@ -394,6 +411,12 @@ func (p *OAuth2Provider) ForceRefreshAccessToken(ctx *schemas.BifrostContext, co if token == nil { return schemas.ErrOAuth2TokenNotFound } + // Force-refresh means the caller just saw the current access token + // rejected upstream. Drop the cached copy for this binding up front, + // in addition to RefreshAccessToken's own post-write eviction, so + // the follow-up lookup re-reads the database even if the refresh + // call itself fails. + p.EvictUserToken(mode, identity, config.ID) return p.RefreshAccessToken(ctx, token.ID) default: return fmt.Errorf("force-refresh is not supported for MCP auth type %q", config.AuthType) @@ -459,6 +482,7 @@ func (p *OAuth2Provider) RevokeToken(ctx context.Context, oauthConfigID string) if err := p.configStore.DeleteSharedOauthTokensByConfigID(ctx, oauthConfigID); err != nil { return fmt.Errorf("failed to delete token: %w", err) } + p.EvictUserTokenByID(token.ID) logger.Debug("OAuth token revoked", "oauth_config_id", oauthConfigID) @@ -922,6 +946,10 @@ func (p *OAuth2Provider) CompleteOAuthFlow(ctx context.Context, state, code stri p.cleanupFlow(ctx, flow.ID) return err } + // CreateOauthToken upserts by binding and rewrites tokenRecord.ID to the + // reused row's ID when one existed, so this evicts exactly the row that + // was written. A fresh row has nothing cached and the evict is a no-op. + p.EvictUserTokenByID(tokenRecord.ID) // Flow row's purpose ends here — same reasoning as CompleteUserOAuthFlow // (see cleanupFlow): the token row is the durable record now. Harmless @@ -1521,6 +1549,10 @@ func (p *OAuth2Provider) CompleteUserOAuthFlow(ctx context.Context, state string _ = p.configStore.DeleteOauthUserSession(ctx, session.ID) return "", fmt.Errorf("failed to create per-user oauth token: %w", err) } + // CreateOauthToken upserts by binding and rewrites tokenRecord.ID to the + // reused row's ID when one existed, so this evicts exactly the row that + // was written; the next lookup for this binding reads the new credential. + p.EvictUserTokenByID(tokenRecord.ID) // Token row is written; the flow row's purpose ends here. cleanupFlow // deletes both the flow row (transient PKCE/state carrier — the token @@ -1538,32 +1570,119 @@ func (p *OAuth2Provider) CompleteUserOAuthFlow(ctx context.Context, state string // GetUserAccessTokenByMode retrieves the upstream access token using exactly // one identity column determined by mode. No fallback chain. Filters // status='active' so orphaned rows never satisfy a lookup. +// +// Reads are served from an in-memory cache when possible: a cached entry +// whose expiry has passed is dropped and the lookup falls through to the +// database path, which owns every expiry and refresh decision. Failed +// lookups are never cached, so a caller completing OAuth (or a token being +// reactivated) becomes visible on the very next call. func (p *OAuth2Provider) GetUserAccessTokenByMode(ctx context.Context, mode schemas.MCPAuthMode, identity, mcpClientID string) (string, error) { + key := userTokenCacheKey(mode, identity, mcpClientID) + if cached, ok := p.userTokens.Get(key); ok { + return cached.accessToken, nil + } + cached, err := p.userTokens.Fill(ctx, key, func() (cachedUserToken, error) { + return p.loadUserAccessTokenByMode(ctx, mode, identity, mcpClientID) + }) + if err != nil { + return "", err + } + return cached.accessToken, nil +} + +// loadUserAccessTokenByMode is GetUserAccessTokenByMode's cache-miss path: +// the full database lookup, including the lazy pre-flight refresh for a +// known-expired row with a refresh token available. +func (p *OAuth2Provider) loadUserAccessTokenByMode(ctx context.Context, mode schemas.MCPAuthMode, identity, mcpClientID string) (cachedUserToken, error) { token, err := p.configStore.GetOauthUserTokenByMode(ctx, mode, identity, mcpClientID) if err != nil { - return "", fmt.Errorf("failed to load per-user oauth token (mode=%s): %w", mode, err) + return cachedUserToken{}, fmt.Errorf("failed to load per-user oauth token (mode=%s): %w", mode, err) } if token == nil { - return "", schemas.ErrOAuth2TokenNotFound + return cachedUserToken{}, schemas.ErrOAuth2TokenNotFound } // Refresh only when known-expired and refresh token exists. if token.ExpiresAt != nil && time.Now().After(*token.ExpiresAt) && strings.TrimSpace(token.RefreshToken) != "" { if err := p.RefreshAccessToken(ctx, token.ID); err != nil { - return "", fmt.Errorf("per-user token expired and refresh failed: %w", err) + return cachedUserToken{}, fmt.Errorf("per-user token expired and refresh failed: %w", err) } token, err = p.configStore.GetOauthUserTokenByMode(ctx, mode, identity, mcpClientID) if err != nil || token == nil { - return "", fmt.Errorf("failed to reload per-user token after refresh") + return cachedUserToken{}, fmt.Errorf("failed to reload per-user token after refresh") } } if token.ExpiresAt != nil && time.Now().After(*token.ExpiresAt) { - return "", fmt.Errorf("per-user token expired and no refresh token is available; re-authorization required: %w", schemas.ErrOAuth2TokenExpired) + return cachedUserToken{}, fmt.Errorf("per-user token expired and no refresh token is available; re-authorization required: %w", schemas.ErrOAuth2TokenExpired) } accessToken := strings.TrimSpace(token.AccessToken) if accessToken == "" { - return "", fmt.Errorf("per-user access token is empty after sanitization") + return cachedUserToken{}, fmt.Errorf("per-user access token is empty after sanitization") } - return accessToken, nil + return cachedUserToken{tokenID: token.ID, accessToken: accessToken, expiresAt: token.ExpiresAt}, nil +} + +// EvictUserToken drops the cached access token for one (mode, identity, +// mcp client) binding. Side-effect only and safe to call when the cache is +// absent; the next lookup for the binding reads the database. +func (p *OAuth2Provider) EvictUserToken(mode schemas.MCPAuthMode, identity, mcpClientID string) { + if p == nil { + return + } + p.userTokens.Evict(userTokenCacheKey(mode, identity, mcpClientID)) +} + +// EvictUserTokenByID drops the cached access token backed by the given token +// row ID, if any binding currently holds it. Side-effect only and safe to +// call when the cache is absent or the ID is not cached. +func (p *OAuth2Provider) EvictUserTokenByID(tokenID string) { + if p == nil { + return + } + p.userTokens.EvictByTokenID(tokenID) +} + +// EvictUserTokensByMCPClient drops every cached access token bound to the +// given MCP client, across all auth modes and identities. Used after +// client-level mutations that invalidate its token rows as a set, such as +// credential rotation, access reconciliation, or client deletion. +// Side-effect only and safe to call when the cache is absent. +func (p *OAuth2Provider) EvictUserTokensByMCPClient(mcpClientID string) { + if p == nil { + return + } + p.userTokens.EvictByMCPClient(mcpClientID) +} + +// EvictUserTokensByVirtualKey drops every cached vk-mode access token bound +// to the given virtual key, across all MCP clients. Used after virtual key +// mutations that orphan or delete its token rows as a set. Side-effect only +// and safe to call when the cache is absent. +func (p *OAuth2Provider) EvictUserTokensByVirtualKey(virtualKeyID string) { + if p == nil { + return + } + p.userTokens.EvictByVirtualKey(virtualKeyID) +} + +// EvictUserTokensByUser drops every cached user-mode access token bound to +// the given user, across all MCP clients. Used after user-level mutations +// that orphan or delete the user's token rows as a set. Side-effect only +// and safe to call when the cache is absent. +func (p *OAuth2Provider) EvictUserTokensByUser(userID string) { + if p == nil { + return + } + p.userTokens.EvictByUser(userID) +} + +// FlushUserTokenCache drops every cached per-user access token. The coarse +// fallback for mutations whose blast radius cannot be scoped to one client +// or virtual key. Side-effect only. +func (p *OAuth2Provider) FlushUserTokenCache() { + if p == nil { + return + } + p.userTokens.Flush() } diff --git a/framework/oauth2/usertokencache.go b/framework/oauth2/usertokencache.go new file mode 100644 index 00000000000..d3f383ab8a0 --- /dev/null +++ b/framework/oauth2/usertokencache.go @@ -0,0 +1,197 @@ +package oauth2 + +import ( + "context" + "fmt" + "strconv" + "strings" + "time" + + "github.com/maximhq/bifrost/core/schemas" + "github.com/maximhq/bifrost/framework/lrucache" +) + +// defaultUserTokenCacheCapacity bounds the per-user token cache. Session-mode +// identities are caller-asserted strings, so the cache must stay bounded no +// matter what identities callers present; least-recently-used entries are +// dropped once the capacity is reached. +const defaultUserTokenCacheCapacity = 4096 + +// cachedUserToken is the trimmed value cached per (auth mode, identity, +// mcp client) binding. Only the fields the request path needs are kept: +// the row ID (for targeted eviction), the sanitized access token, and the +// expiry used to treat a stale hit as a miss. Refresh tokens are never +// cached; anything needing one goes back to the database row. +type cachedUserToken struct { + tokenID string + accessToken string + expiresAt *time.Time +} + +// userTokenCache adapts lrucache.Cache to per-user MCP OAuth access-token +// lookups: it owns the binding-key scheme, registers each entry under its +// token row ID for targeted eviction (refresh, revoke, needs_reauth), treats +// expired entries as misses via the validator so the database path keeps +// every expiry and refresh decision, and provides the scoped bulk evictions +// the credential lifecycle needs (by MCP client, virtual key, and user). +// +// Locking: the cache owns its consistency and must never be guarded by the +// provider's own lock. The provider's lock is held across token-endpoint +// network I/O during refresh and revocation, so sharing it would stall +// every cached read behind unrelated upstream traffic. +type userTokenCache struct { + cache *lrucache.Cache[cachedUserToken] +} + +// userTokenCacheKey builds the cache key for a (mode, identity, mcp client) +// binding. Length-prefixed rather than NUL-separated: identity is a +// caller-asserted string (see the doc comment above), so a NUL-separated +// scheme would let an identity value containing "\x00" forge a component +// boundary — e.g. (mode, "a\x00b", "c") and (mode, "a", "b\x00c") would build +// the identical key, letting one binding's cached access token be served to +// another. Prefixing each component with its own byte length makes that +// forgery impossible regardless of what bytes a component contains. +func userTokenCacheKey(mode schemas.MCPAuthMode, identity, mcpClientID string) string { + return fmt.Sprintf("%d:%s%d:%s%d:%s", len(mode), mode, len(identity), identity, len(mcpClientID), mcpClientID) +} + +// splitUserTokenCacheKey is userTokenCacheKey's inverse, for the scoped +// eviction predicates. ok is false for a key that isn't in the +// length-prefixed form userTokenCacheKey builds (impossible for keys this +// package produces). +func splitUserTokenCacheKey(key string) (mode, identity, clientID string, ok bool) { + rest := key + parts := make([]string, 0, 3) + for range 3 { + i := strings.IndexByte(rest, ':') + if i < 0 { + return "", "", "", false + } + n, err := strconv.Atoi(rest[:i]) + if err != nil || n < 0 { + return "", "", "", false + } + rest = rest[i+1:] + if n > len(rest) { + return "", "", "", false + } + parts = append(parts, rest[:n]) + rest = rest[n:] + } + if rest != "" { + return "", "", "", false + } + return parts[0], parts[1], parts[2], true +} + +func newUserTokenCache(capacity int) *userTokenCache { + if capacity <= 0 { + capacity = defaultUserTokenCacheCapacity + } + return &userTokenCache{ + cache: lrucache.New(capacity, lrucache.WithValidator(func(v cachedUserToken) bool { + // An expired entry is a miss: the caller falls through to the + // full database path, which owns all expiry and refresh + // decisions. nil means non-expiring. + return v.expiresAt == nil || time.Now().Before(*v.expiresAt) + })), + } +} + +// Get returns the cached value for key. A hit whose expiry has passed is +// removed and reported as a miss. +func (c *userTokenCache) Get(key string) (cachedUserToken, bool) { + if c == nil { + return cachedUserToken{}, false + } + return c.cache.Get(key) +} + +// Fill runs fill for key with single-flight deduplication: concurrent +// callers for the same key wait for one leader and share its result and +// error, so a refresh burst for one identity performs a single database +// read (and at most one upstream refresh) instead of a stampede. A +// successful result is cached under both the binding key and its token row +// ID; errors are propagated but never cached. +func (c *userTokenCache) Fill(ctx context.Context, key string, fill func() (cachedUserToken, error)) (cachedUserToken, error) { + if c == nil { + return fill() + } + return c.cache.Fill(ctx, key, func() (cachedUserToken, string, error) { + value, err := fill() + return value, value.tokenID, err + }) +} + +// Evict removes the entry for an exact binding key, if present. +func (c *userTokenCache) Evict(key string) { + if c == nil { + return + } + c.cache.Evict(key) +} + +// EvictByTokenID removes the entry holding the given token row ID, if any. +func (c *userTokenCache) EvictByTokenID(tokenID string) { + if c == nil { + return + } + c.cache.EvictByIndex(tokenID) +} + +// EvictByMCPClient removes every cached entry bound to the given MCP client, +// across all auth modes and identities. Used when a client-level change +// invalidates its tokens as a set, such as credential rotation or client +// deletion. A linear sweep is fine here: these are rare admin operations and +// the cache is bounded. +func (c *userTokenCache) EvictByMCPClient(mcpClientID string) { + if c == nil || mcpClientID == "" { + return + } + c.cache.EvictWhere(func(key string) bool { + _, _, clientID, ok := splitUserTokenCacheKey(key) + return ok && clientID == mcpClientID + }) +} + +// EvictByVirtualKey removes every cached vk-mode entry bound to the given +// virtual key, across all MCP clients. Used when a virtual key change +// orphans or deletes its token rows as a set. +func (c *userTokenCache) EvictByVirtualKey(virtualKeyID string) { + if c == nil || virtualKeyID == "" { + return + } + c.cache.EvictWhere(func(key string) bool { + mode, identity, _, ok := splitUserTokenCacheKey(key) + return ok && mode == string(schemas.MCPAuthModeVK) && identity == virtualKeyID + }) +} + +// EvictByUser removes every cached user-mode entry bound to the given user, +// across all MCP clients. Used when a user-level change orphans or deletes +// the user's token rows as a set. +func (c *userTokenCache) EvictByUser(userID string) { + if c == nil || userID == "" { + return + } + c.cache.EvictWhere(func(key string) bool { + mode, identity, _, ok := splitUserTokenCacheKey(key) + return ok && mode == string(schemas.MCPAuthModeUser) && identity == userID + }) +} + +// Flush drops every cached entry. +func (c *userTokenCache) Flush() { + if c == nil { + return + } + c.cache.Flush() +} + +// Len reports the number of cached entries. +func (c *userTokenCache) Len() int { + if c == nil { + return 0 + } + return c.cache.Len() +} diff --git a/framework/oauth2/usertokencache_test.go b/framework/oauth2/usertokencache_test.go new file mode 100644 index 00000000000..381c073a13c --- /dev/null +++ b/framework/oauth2/usertokencache_test.go @@ -0,0 +1,594 @@ +package oauth2 + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "net/http/httptest" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + bifrost "github.com/maximhq/bifrost/core" + "github.com/maximhq/bifrost/core/schemas" + "github.com/maximhq/bifrost/framework/configstore/tables" +) + +func testToken(id, access string, expiresAt *time.Time) cachedUserToken { + return cachedUserToken{tokenID: id, accessToken: access, expiresAt: expiresAt} +} + +func fillWith(v cachedUserToken) func() (cachedUserToken, error) { + return func() (cachedUserToken, error) { return v, nil } +} + +func TestUserTokenCache_HitAndMiss(t *testing.T) { + c := newUserTokenCache(4) + + _, ok := c.Get("missing") + assert.False(t, ok, "empty cache must miss") + + v, err := c.Fill(context.Background(), "k1", fillWith(testToken("t1", "access-1", nil))) + require.NoError(t, err) + assert.Equal(t, "access-1", v.accessToken) + + got, ok := c.Get("k1") + require.True(t, ok, "filled key must hit") + assert.Equal(t, "access-1", got.accessToken) + assert.Equal(t, "t1", got.tokenID) + + _, ok = c.Get("k2") + assert.False(t, ok, "unrelated key must miss") +} + +func TestUserTokenCache_CapacityEviction(t *testing.T) { + c := newUserTokenCache(2) + + for i := 1; i <= 3; i++ { + _, err := c.Fill(context.Background(), + fmt.Sprintf("k%d", i), + fillWith(testToken(fmt.Sprintf("t%d", i), fmt.Sprintf("access-%d", i), nil)), + ) + require.NoError(t, err) + } + + assert.Equal(t, 2, c.Len(), "cache must stay at capacity") + _, ok := c.Get("k1") + assert.False(t, ok, "least recently used entry must be evicted") + _, ok = c.Get("k2") + assert.True(t, ok) + _, ok = c.Get("k3") + assert.True(t, ok) + + // The evicted entry's token-ID index mapping must be gone: evicting by + // its row ID must not disturb the surviving entries. + c.EvictByTokenID("t1") + assert.Equal(t, 2, c.Len()) +} + +func TestUserTokenCache_ExpiredEntryIsMissAndRemoved(t *testing.T) { + c := newUserTokenCache(4) + past := time.Now().Add(-1 * time.Minute) + + _, err := c.Fill(context.Background(), "k1", fillWith(testToken("t1", "stale-access", &past))) + require.NoError(t, err) + require.Equal(t, 1, c.Len()) + + _, ok := c.Get("k1") + assert.False(t, ok, "expired entry must read as a miss") + assert.Equal(t, 0, c.Len(), "expired entry must be removed on read") +} + +func TestUserTokenCache_EvictExactKey(t *testing.T) { + c := newUserTokenCache(4) + _, err := c.Fill(context.Background(), "k1", fillWith(testToken("t1", "access-1", nil))) + require.NoError(t, err) + + c.Evict("k1") + _, ok := c.Get("k1") + assert.False(t, ok) + assert.Equal(t, 0, c.Len()) + + // Evicting an absent key is a no-op, not a panic. + c.Evict("never-set") +} + +func TestUserTokenCache_EvictByTokenID(t *testing.T) { + c := newUserTokenCache(4) + _, err := c.Fill(context.Background(), "k1", fillWith(testToken("t1", "access-1", nil))) + require.NoError(t, err) + _, err = c.Fill(context.Background(), "k2", fillWith(testToken("t2", "access-2", nil))) + require.NoError(t, err) + + c.EvictByTokenID("t1") + _, ok := c.Get("k1") + assert.False(t, ok, "entry holding the evicted token ID must be gone") + _, ok = c.Get("k2") + assert.True(t, ok, "unrelated entry must survive") + + // Unknown token IDs are a no-op. + c.EvictByTokenID("unknown") + assert.Equal(t, 1, c.Len()) +} + +func TestUserTokenCache_EvictByMCPClient(t *testing.T) { + c := newUserTokenCache(8) + keyA1 := userTokenCacheKey(schemas.MCPAuthModeUser, "u1", "client-a") + keyA2 := userTokenCacheKey(schemas.MCPAuthModeVK, "vk1", "client-a") + keyB := userTokenCacheKey(schemas.MCPAuthModeUser, "u1", "client-b") + for i, k := range []string{keyA1, keyA2, keyB} { + _, err := c.Fill(context.Background(), k, fillWith(testToken(fmt.Sprintf("t%d", i), "access", nil))) + require.NoError(t, err) + } + + c.EvictByMCPClient("client-a") + _, ok := c.Get(keyA1) + assert.False(t, ok, "user-mode entry for the client must be gone") + _, ok = c.Get(keyA2) + assert.False(t, ok, "vk-mode entry for the client must be gone") + _, ok = c.Get(keyB) + assert.True(t, ok, "entry for another client must survive") + assert.Equal(t, 1, c.Len()) + + // Empty client ID and unknown client IDs are no-ops. + c.EvictByMCPClient("") + c.EvictByMCPClient("client-x") + assert.Equal(t, 1, c.Len()) +} + +func TestUserTokenCache_EvictByVirtualKey(t *testing.T) { + c := newUserTokenCache(8) + keyVK1A := userTokenCacheKey(schemas.MCPAuthModeVK, "vk1", "client-a") + keyVK1B := userTokenCacheKey(schemas.MCPAuthModeVK, "vk1", "client-b") + keyVK2 := userTokenCacheKey(schemas.MCPAuthModeVK, "vk2", "client-a") + // A user-mode identity that happens to equal the VK ID must survive: + // the eviction is scoped to vk-mode entries only. + keyUser := userTokenCacheKey(schemas.MCPAuthModeUser, "vk1", "client-a") + for i, k := range []string{keyVK1A, keyVK1B, keyVK2, keyUser} { + _, err := c.Fill(context.Background(), k, fillWith(testToken(fmt.Sprintf("vt%d", i), "access", nil))) + require.NoError(t, err) + } + + c.EvictByVirtualKey("vk1") + _, ok := c.Get(keyVK1A) + assert.False(t, ok) + _, ok = c.Get(keyVK1B) + assert.False(t, ok, "the VK's entries must be evicted across clients") + _, ok = c.Get(keyVK2) + assert.True(t, ok, "another VK's entry must survive") + _, ok = c.Get(keyUser) + assert.True(t, ok, "a user-mode identity equal to the VK ID must survive") + + // Empty and unknown VK IDs are no-ops. + c.EvictByVirtualKey("") + c.EvictByVirtualKey("vk-x") + assert.Equal(t, 2, c.Len()) +} + +func TestUserTokenCache_EvictByUser(t *testing.T) { + c := newUserTokenCache(8) + keyU1A := userTokenCacheKey(schemas.MCPAuthModeUser, "u1", "client-a") + keyU1B := userTokenCacheKey(schemas.MCPAuthModeUser, "u1", "client-b") + keyU2 := userTokenCacheKey(schemas.MCPAuthModeUser, "u2", "client-a") + // A vk-mode identity that happens to equal the user ID must survive: + // the eviction is scoped to user-mode entries only. + keyVK := userTokenCacheKey(schemas.MCPAuthModeVK, "u1", "client-a") + for i, k := range []string{keyU1A, keyU1B, keyU2, keyVK} { + _, err := c.Fill(context.Background(), k, fillWith(testToken(fmt.Sprintf("ut%d", i), "access", nil))) + require.NoError(t, err) + } + + c.EvictByUser("u1") + _, ok := c.Get(keyU1A) + assert.False(t, ok) + _, ok = c.Get(keyU1B) + assert.False(t, ok, "the user's entries must be evicted across clients") + _, ok = c.Get(keyU2) + assert.True(t, ok, "another user's entry must survive") + _, ok = c.Get(keyVK) + assert.True(t, ok, "a vk-mode identity equal to the user ID must survive") + + // Empty and unknown user IDs are no-ops. + c.EvictByUser("") + c.EvictByUser("u-x") + assert.Equal(t, 2, c.Len()) +} + +func TestUserTokenCache_Flush(t *testing.T) { + c := newUserTokenCache(4) + for i := 1; i <= 3; i++ { + _, err := c.Fill(context.Background(), fmt.Sprintf("k%d", i), fillWith(testToken(fmt.Sprintf("t%d", i), "access", nil))) + require.NoError(t, err) + } + + c.Flush() + assert.Equal(t, 0, c.Len()) + for i := 1; i <= 3; i++ { + _, ok := c.Get(fmt.Sprintf("k%d", i)) + assert.False(t, ok) + } +} + +func TestUserTokenCache_InflightDedup(t *testing.T) { + c := newUserTokenCache(4) + + var calls atomic.Int64 + release := make(chan struct{}) + started := make(chan struct{}) + + fill := func() (cachedUserToken, error) { + calls.Add(1) + close(started) + <-release + return testToken("t1", "shared-access", nil), nil + } + + var wg sync.WaitGroup + results := make([]cachedUserToken, 2) + wg.Add(1) + fillErrs := make([]error, 2) + go func() { + defer wg.Done() + v, err := c.Fill(context.Background(), "k1", fill) + results[0], fillErrs[0] = v, err + }() + <-started + + wg.Add(1) + go func() { + defer wg.Done() + // This fill func must never run: the leader is already in flight. + v, err := c.Fill(context.Background(), "k1", func() (cachedUserToken, error) { + calls.Add(1) + return testToken("t-other", "other", nil), nil + }) + results[1], fillErrs[1] = v, err + }() + + // Give the second goroutine a moment to register as a waiter, then let + // the leader finish. + time.Sleep(20 * time.Millisecond) + close(release) + wg.Wait() + + // require, not assert, calls t.FailNow(), which per the testing + // package's own contract must only be called from the goroutine running + // the test — hence collecting errors above and asserting them here. + require.NoError(t, fillErrs[0]) + require.NoError(t, fillErrs[1]) + + assert.Equal(t, int64(1), calls.Load(), "concurrent fills for one key must run the handler once") + assert.Equal(t, "shared-access", results[0].accessToken) + assert.Equal(t, "shared-access", results[1].accessToken, "waiter must share the leader's result") +} + +func TestUserTokenCache_ErrorSharedNotCached(t *testing.T) { + c := newUserTokenCache(4) + + var calls atomic.Int64 + release := make(chan struct{}) + started := make(chan struct{}) + fillErr := errors.New("db unavailable") + + var wg sync.WaitGroup + errs := make([]error, 2) + wg.Add(1) + go func() { + defer wg.Done() + _, err := c.Fill(context.Background(), "k1", func() (cachedUserToken, error) { + calls.Add(1) + close(started) + <-release + return cachedUserToken{}, fillErr + }) + errs[0] = err + }() + <-started + + wg.Add(1) + go func() { + defer wg.Done() + _, err := c.Fill(context.Background(), "k1", func() (cachedUserToken, error) { + calls.Add(1) + return cachedUserToken{}, fillErr + }) + errs[1] = err + }() + + time.Sleep(20 * time.Millisecond) + close(release) + wg.Wait() + + assert.Equal(t, int64(1), calls.Load(), "waiter must share the leader's error, not re-run the fill") + assert.ErrorIs(t, errs[0], fillErr) + assert.ErrorIs(t, errs[1], fillErr, "error must be shared with waiters") + assert.Equal(t, 0, c.Len(), "errors must never be cached") + + // A later fill runs the handler again: the failure was not cached. + _, err := c.Fill(context.Background(), "k1", func() (cachedUserToken, error) { + calls.Add(1) + return testToken("t1", "recovered", nil), nil + }) + require.NoError(t, err) + assert.Equal(t, int64(2), calls.Load()) + got, ok := c.Get("k1") + require.True(t, ok) + assert.Equal(t, "recovered", got.accessToken) +} + +func TestUserTokenCache_GenerationGuardDiscardsStaleFill(t *testing.T) { + c := newUserTokenCache(4) + + inFill := make(chan struct{}) + release := make(chan struct{}) + + done := make(chan struct{}) + go func() { + defer close(done) + v, err := c.Fill(context.Background(), "k1", func() (cachedUserToken, error) { + close(inFill) + <-release + return testToken("t1", "stale-value", nil), nil + }) + // The caller still receives the value it read; only the cache + // install is discarded. assert, not require: require.NoError calls + // t.FailNow(), which per the testing package's own contract must + // only be called from the goroutine running the test. + assert.NoError(t, err) + assert.Equal(t, "stale-value", v.accessToken) + }() + + <-inFill + // An eviction lands while the fill is mid-read: whatever the fill read + // is now suspect and must not be installed. + c.Evict("k1") + close(release) + <-done + + _, ok := c.Get("k1") + assert.False(t, ok, "a fill that raced an eviction must not install its result") + assert.Equal(t, 0, c.Len()) +} + +func TestUserTokenCache_UpsertRebindsTokenIDIndex(t *testing.T) { + c := newUserTokenCache(4) + _, err := c.Fill(context.Background(), "k1", fillWith(testToken("t-old", "old", nil))) + require.NoError(t, err) + + // Same key, new backing row (the binding re-authenticated onto a fresh + // row): the index must follow the new row ID and drop the old one. + c.Evict("k1") + _, err = c.Fill(context.Background(), "k1", fillWith(testToken("t-new", "new", nil))) + require.NoError(t, err) + + c.EvictByTokenID("t-old") + _, ok := c.Get("k1") + assert.True(t, ok, "stale token-ID mapping must not evict the rebound entry") + + c.EvictByTokenID("t-new") + _, ok = c.Get("k1") + assert.False(t, ok, "current token-ID mapping must evict the entry") +} + +func TestUserTokenCache_NilSafety(t *testing.T) { + var c *userTokenCache + _, ok := c.Get("k") + assert.False(t, ok) + c.Evict("k") + c.EvictByTokenID("t") + c.Flush() + assert.Equal(t, 0, c.Len()) + v, err := c.Fill(context.Background(), "k", fillWith(testToken("t", "pass-through", nil))) + require.NoError(t, err) + assert.Equal(t, "pass-through", v.accessToken) +} + +// ---------- Integration through GetUserAccessTokenByMode ---------- + +// countingConfigStore wraps testConfigStore with per-method call counters so +// integration tests can assert which lookups actually reached the store. +type countingConfigStore struct { + *testConfigStore + getUserTokenByModeCalls atomic.Int64 +} + +func (s *countingConfigStore) GetOauthUserTokenByMode(ctx context.Context, mode schemas.MCPAuthMode, identity, mcpClientID string) (*tables.TableMCPOauthToken, error) { + s.getUserTokenByModeCalls.Add(1) + return s.testConfigStore.GetOauthUserTokenByMode(ctx, mode, identity, mcpClientID) +} + +func newCountingConfigStore() *countingConfigStore { + return &countingConfigStore{testConfigStore: newTestConfigStore()} +} + +// seedUserToken inserts an active session-mode per-user token row. +func seedUserToken(store *testConfigStore, tokenID, oauthConfigID, mcpClientID, sessionID, access string, expiresAt *time.Time) { + store.oauthTokens[tokenID] = &tables.TableMCPOauthToken{ + ID: tokenID, + AuthMode: "session", + MCPClientID: mcpClientID, + OauthConfigID: oauthConfigID, + SessionID: sessionID, + Status: "active", + AccessToken: access, + RefreshToken: "refresh-token", + TokenType: "bearer", + ExpiresAt: expiresAt, + Scopes: "[]", + } +} + +func TestGetUserAccessTokenByMode_SecondCallServedFromCache(t *testing.T) { + store := newCountingConfigStore() + seedUserToken(store.testConfigStore, "tok-1", "cfg-1", "mcp-1", "sess-1", "cached-access", bifrost.Ptr(time.Now().Add(1*time.Hour))) + + provider := NewOAuth2Provider(store, bifrost.NewDefaultLogger(schemas.LogLevelError)) + ctx := context.Background() + + access, err := provider.GetUserAccessTokenByMode(ctx, schemas.MCPAuthModeSession, "sess-1", "mcp-1") + require.NoError(t, err) + assert.Equal(t, "cached-access", access) + assert.Equal(t, int64(1), store.getUserTokenByModeCalls.Load()) + + access, err = provider.GetUserAccessTokenByMode(ctx, schemas.MCPAuthModeSession, "sess-1", "mcp-1") + require.NoError(t, err) + assert.Equal(t, "cached-access", access) + assert.Equal(t, int64(1), store.getUserTokenByModeCalls.Load(), "second call must be served from cache, not the store") +} + +func TestGetUserAccessTokenByMode_ExpiredRefreshesOnceUnderConcurrency(t *testing.T) { + var tokenEndpointCalls atomic.Int64 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + tokenEndpointCalls.Add(1) + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(map[string]any{ + "access_token": "refreshed-access", + "refresh_token": "new-refresh-token", + "token_type": "bearer", + "expires_in": 3600, + }) + })) + defer server.Close() + + store := newCountingConfigStore() + store.oauthConfigs["cfg-1"] = &tables.TableOauthConfig{ + ID: "cfg-1", + ClientID: schemas.NewSecretVar("client-id"), + TokenURL: server.URL + "/token", + RedirectURI: "http://localhost/callback", + Scopes: `["read"]`, + Status: "authorized", + } + seedUserToken(store.testConfigStore, "tok-1", "cfg-1", "mcp-1", "sess-1", "expired-access", bifrost.Ptr(time.Now().Add(-1*time.Minute))) + + provider := NewOAuth2Provider(store, bifrost.NewDefaultLogger(schemas.LogLevelError)) + ctx := context.Background() + + const callers = 8 + var wg sync.WaitGroup + results := make([]string, callers) + errs := make([]error, callers) + for i := range callers { + wg.Add(1) + go func() { + defer wg.Done() + results[i], errs[i] = provider.GetUserAccessTokenByMode(ctx, schemas.MCPAuthModeSession, "sess-1", "mcp-1") + }() + } + wg.Wait() + + for i := range callers { + require.NoError(t, errs[i]) + assert.Equal(t, "refreshed-access", results[i]) + } + assert.Equal(t, int64(1), tokenEndpointCalls.Load(), "concurrent callers must trigger exactly one upstream refresh") +} + +func TestGetUserAccessTokenByMode_EvictByIDAfterDelete(t *testing.T) { + store := newCountingConfigStore() + seedUserToken(store.testConfigStore, "tok-1", "cfg-1", "mcp-1", "sess-1", "cached-access", bifrost.Ptr(time.Now().Add(1*time.Hour))) + + provider := NewOAuth2Provider(store, bifrost.NewDefaultLogger(schemas.LogLevelError)) + ctx := context.Background() + + _, err := provider.GetUserAccessTokenByMode(ctx, schemas.MCPAuthModeSession, "sess-1", "mcp-1") + require.NoError(t, err) + + // Delete the row directly through the store (as the sessions revoke + // handler does) and evict by ID (as its cache callback does). + store.mu.Lock() + delete(store.oauthTokens, "tok-1") + store.mu.Unlock() + provider.EvictUserTokenByID("tok-1") + + _, err = provider.GetUserAccessTokenByMode(ctx, schemas.MCPAuthModeSession, "sess-1", "mcp-1") + require.Error(t, err) + assert.ErrorIs(t, err, schemas.ErrOAuth2TokenNotFound, "post-eviction lookup must see the delete, not the cached token") +} + +func TestGetUserAccessTokenByMode_FlushAfterNeedsReauth(t *testing.T) { + store := newCountingConfigStore() + seedUserToken(store.testConfigStore, "tok-1", "cfg-1", "mcp-1", "sess-1", "cached-access", bifrost.Ptr(time.Now().Add(1*time.Hour))) + + provider := NewOAuth2Provider(store, bifrost.NewDefaultLogger(schemas.LogLevelError)) + ctx := context.Background() + + _, err := provider.GetUserAccessTokenByMode(ctx, schemas.MCPAuthModeSession, "sess-1", "mcp-1") + require.NoError(t, err) + + // Credential rotation marks every token row for the config + // needs_reauth, then the handler-level callback flushes the cache. + require.NoError(t, store.MarkTokensNeedsReauthByConfigID(ctx, "cfg-1")) + provider.FlushUserTokenCache() + + // The active-only lookup no longer matches the row, so the caller sees + // the re-auth requirement instead of the cached token. + _, err = provider.GetUserAccessTokenByMode(ctx, schemas.MCPAuthModeSession, "sess-1", "mcp-1") + require.Error(t, err) + assert.ErrorIs(t, err, schemas.ErrOAuth2TokenNotFound) +} + +func TestGetUserAccessTokenByMode_NegativeNotCached(t *testing.T) { + store := newCountingConfigStore() + provider := NewOAuth2Provider(store, bifrost.NewDefaultLogger(schemas.LogLevelError)) + ctx := context.Background() + + _, err := provider.GetUserAccessTokenByMode(ctx, schemas.MCPAuthModeSession, "sess-1", "mcp-1") + require.Error(t, err) + assert.ErrorIs(t, err, schemas.ErrOAuth2TokenNotFound) + + // The user completes OAuth: a fresh row appears. No eviction happens + // (there is nothing to evict) and the very next call must see it. + store.mu.Lock() + seedUserToken(store.testConfigStore, "tok-1", "cfg-1", "mcp-1", "sess-1", "fresh-access", bifrost.Ptr(time.Now().Add(1*time.Hour))) + store.mu.Unlock() + + access, err := provider.GetUserAccessTokenByMode(ctx, schemas.MCPAuthModeSession, "sess-1", "mcp-1") + require.NoError(t, err) + assert.Equal(t, "fresh-access", access, "a failed lookup must never be cached") +} + +func TestGetUserAccessTokenByMode_ForceRefreshEvictsAndServesNewToken(t *testing.T) { + server := tokenRefreshServer(t, "force-refreshed-access") + + store := newCountingConfigStore() + store.oauthConfigs["cfg-1"] = &tables.TableOauthConfig{ + ID: "cfg-1", + ClientID: schemas.NewSecretVar("client-id"), + TokenURL: server.URL + "/token", + RedirectURI: "http://localhost/callback", + Scopes: `["read"]`, + Status: "authorized", + } + // Not expired: the cached copy would keep serving without the forced + // refresh's eviction. + seedUserToken(store.testConfigStore, "tok-1", "cfg-1", "mcp-1", "sess-1", "rejected-upstream", bifrost.Ptr(time.Now().Add(1*time.Hour))) + + provider := NewOAuth2Provider(store, bifrost.NewDefaultLogger(schemas.LogLevelError)) + ctx := context.Background() + + access, err := provider.GetUserAccessTokenByMode(ctx, schemas.MCPAuthModeSession, "sess-1", "mcp-1") + require.NoError(t, err) + require.Equal(t, "rejected-upstream", access) + + oauthConfigID := "cfg-1" + config := &schemas.MCPClientConfig{ + ID: "mcp-1", + Name: "Test Client", + AuthType: schemas.MCPAuthTypePerUserOauth, + OauthConfigID: &oauthConfigID, + } + bfCtx := schemas.NewBifrostContext(context.Background(), schemas.NoDeadline) + bfCtx.SetValue(schemas.BifrostContextKeyMCPSessionID, "sess-1") + require.NoError(t, provider.ForceRefreshAccessToken(bfCtx, config)) + + access, err = provider.GetUserAccessTokenByMode(ctx, schemas.MCPAuthModeSession, "sess-1", "mcp-1") + require.NoError(t, err) + assert.Equal(t, "force-refreshed-access", access, "force refresh must evict the cached copy so the next read serves the new token") +} diff --git a/transports/bifrost-http/handlers/governance.go b/transports/bifrost-http/handlers/governance.go index f89fca5ec19..a85a9e1a1df 100644 --- a/transports/bifrost-http/handlers/governance.go +++ b/transports/bifrost-http/handlers/governance.go @@ -2091,18 +2091,15 @@ func (h *GovernanceHandler) updateVirtualKey(ctx *fasthttp.RequestCtx) { } // Reverse-map governance from VK-scoped model configs for display. h.hydrateVKGovernance(ctx, preloadedVk) - if _, err := h.governanceManager.ReloadVirtualKey(ctx, vk.ID); err != nil { - // Should never happen but just in case - logger.Error("failed to reload virtual key after update: %v", err) - SendError(ctx, 500, "Virtual key updated in database but failed to reload in-memory state") - return - } // Per-user credential reconciliation when the VK's MCP allowlist // changed. Mirrors the AP-propagation path: enterprise orphans / // reactivates credentials keyed to this VK (vk-keyed creds) and to the // VK's owner (user-keyed creds) against the new effective allowlist // (explicit rows ∪ MCPs with AllowOnAllVirtualKeys=true). OSS no-ops. + // Must run before ReloadVirtualKey: the reload also evicts this VK's + // cached OAuth tokens, and an eviction that lands before these writes + // could be refilled from the pre-reconcile rows and then never dropped. if req.MCPConfigs != nil && h.configStore != nil { if err := h.configStore.ReconcileOauthAfterVKChange(ctx, vk.ID); err != nil { logger.Error("reconcile OAuth credentials after VK %s update failed: %v", vk.ID, err) @@ -2112,6 +2109,13 @@ func (h *GovernanceHandler) updateVirtualKey(ctx *fasthttp.RequestCtx) { } } + if _, err := h.governanceManager.ReloadVirtualKey(ctx, vk.ID); err != nil { + // Should never happen but just in case + logger.Error("failed to reload virtual key after update: %v", err) + SendError(ctx, 500, "Virtual key updated in database but failed to reload in-memory state") + return + } + SendJSON(ctx, map[string]interface{}{ "message": "Virtual key updated successfully", "virtual_key": preloadedVk, @@ -2239,7 +2243,9 @@ func (h *GovernanceHandler) deleteVirtualKey(ctx *fasthttp.RequestCtx) { SendError(ctx, 500, "Failed to delete virtual key") return } - // Removing key from in-memory store + // Removing key from in-memory store. RemoveVirtualKey also evicts the + // VK's cached OAuth access tokens internally, covering the token rows + // the database delete above cascaded over. err = h.governanceManager.RemoveVirtualKey(ctx, vk.ID) if err != nil { // But we ignore this error because its not diff --git a/transports/bifrost-http/handlers/mcp.go b/transports/bifrost-http/handlers/mcp.go index faea1d52fd0..952404b06c7 100644 --- a/transports/bifrost-http/handlers/mcp.go +++ b/transports/bifrost-http/handlers/mcp.go @@ -60,16 +60,29 @@ type MCPHandler struct { mcpManager MCPManager governanceManager GovernanceManager oauthHandler *OAuthHandler + // mcpOauthTokenCacheManager invalidates cached OAuth access tokens after + // mutations that rewrite or delete token rows through the configstore + // (credential rotation, access reconciliation). Always wired by the + // server; see the interface doc for the non-nil requirement. + mcpOauthTokenCacheManager MCPOauthTokenCacheManager } // NewMCPHandler creates a new MCP handler instance -func NewMCPHandler(mcpManager MCPManager, governanceManager GovernanceManager, client *bifrost.Bifrost, store *lib.Config, oauthHandler *OAuthHandler) *MCPHandler { +func NewMCPHandler( + mcpManager MCPManager, + governanceManager GovernanceManager, + client *bifrost.Bifrost, + store *lib.Config, + oauthHandler *OAuthHandler, + mcpOauthTokenCacheManager MCPOauthTokenCacheManager, +) *MCPHandler { return &MCPHandler{ - client: client, - store: store, - mcpManager: mcpManager, - governanceManager: governanceManager, - oauthHandler: oauthHandler, + client: client, + store: store, + mcpManager: mcpManager, + governanceManager: governanceManager, + oauthHandler: oauthHandler, + mcpOauthTokenCacheManager: mcpOauthTokenCacheManager, } } @@ -1972,7 +1985,8 @@ func (h *MCPHandler) updateMCPClient(ctx *fasthttp.RequestCtx) { // config update above have succeeded — see the comment where // shouldRotateOAuthConfig was computed for why this can't run earlier. if shouldRotateOAuthConfig { - if _, err := h.store.ConfigStore.RotateMCPOAuthConfig(ctx, existingOauthConfig, resolvedOauthFields); err != nil { + rotated, err := h.store.ConfigStore.RotateMCPOAuthConfig(ctx, existingOauthConfig, resolvedOauthFields) + if err != nil { // The rest of the update already committed; only credential // rotation failed. Report it as a partial success rather than a // full failure so the caller doesn't retry the whole request @@ -1981,17 +1995,21 @@ func (h *MCPHandler) updateMCPClient(ctx *fasthttp.RequestCtx) { SendError(ctx, fasthttp.StatusInternalServerError, fmt.Sprintf("MCP client updated but rotating its OAuth credentials failed: %v", err)) return } - // Rotation just cascaded every token bound to this oauth_config to - // needs_reauth in the DB, but the in-memory client above only had its - // ExecutionConfig replaced — its live connection, if any, is still - // open on the now-invalidated Authorization header. Close it and - // flip the in-memory state to match, rather than leaving a stale - // connection serving calls until the health monitor eventually - // notices. Not a hard failure: the DB rotation is what actually - // matters for correctness, and the health monitor's next cycle would - // eventually catch a connection this call failed to close. - if err := h.mcpManager.CloseAndMarkNeedsReauth(ctx, id); err != nil && !errors.Is(err, schemas.ErrMCPReconnectNotApplicable) { - logger.Error(fmt.Sprintf("Failed to close MCP client %s's connection after OAuth credential rotation: %v", id, err)) + if rotated { + h.mcpOauthTokenCacheManager.EvictOauthTokenCacheByMCPClient(ctx, id) + // Rotation just cascaded every token bound to this oauth_config to + // needs_reauth in the DB, but the in-memory client above only had + // its ExecutionConfig replaced — its live connection, if any, is + // still open on the now-invalidated Authorization header. Close + // it and flip the in-memory state to match, rather than leaving + // a stale connection serving calls until the health monitor + // eventually notices. Not a hard failure: the DB rotation is + // what actually matters for correctness, and the health + // monitor's next cycle would eventually catch a connection this + // call failed to close. + if err := h.mcpManager.CloseAndMarkNeedsReauth(ctx, id); err != nil && !errors.Is(err, schemas.ErrMCPReconnectNotApplicable) { + logger.Error(fmt.Sprintf("Failed to close MCP client %s's connection after OAuth credential rotation: %v", id, err)) + } } } @@ -2147,6 +2165,9 @@ func (h *MCPHandler) updateMCPClient(ctx *fasthttp.RequestCtx) { if err := h.store.ConfigStore.ReconcileMCPHeadersAfterMCPChange(ctx, id); err != nil { logger.Error(fmt.Sprintf("reconcile per-user-headers credentials after MCP %s update failed: %v", id, err)) } + // Reconciliation may have orphaned or reactivated token rows; + // cached copies no longer reflect the database. + h.mcpOauthTokenCacheManager.EvictOauthTokenCacheByMCPClient(ctx, id) } } @@ -2175,6 +2196,9 @@ func (h *MCPHandler) deleteMCPClient(ctx *fasthttp.RequestCtx) { return } } + // RemoveMCPClient also evicts the client's cached OAuth access tokens + // internally, covering the token rows the database delete above + // cascaded over. if err := h.mcpManager.RemoveMCPClient(ctx, id); err != nil { SendError(ctx, fasthttp.StatusInternalServerError, fmt.Sprintf("failed to remove MCP client: %v", err)) return diff --git a/transports/bifrost-http/handlers/mcpsessions.go b/transports/bifrost-http/handlers/mcpsessions.go index c82da73a99d..3af954a637e 100644 --- a/transports/bifrost-http/handlers/mcpsessions.go +++ b/transports/bifrost-http/handlers/mcpsessions.go @@ -9,6 +9,7 @@ package handlers import ( + "context" "errors" "sort" "strconv" @@ -22,14 +23,29 @@ import ( "github.com/valyala/fasthttp" ) +// MCPOauthTokenCacheManager invalidates cached per-user MCP OAuth access +// tokens after a database write that bypasses the OAuth provider's own write +// paths. The base server implementation evicts from local memory; a +// clustered deployment overrides it to also notify peers so their caches +// stay current. Handlers call it directly after successful writes, so the +// server must always wire it (it is nil-safe internally when no provider is +// configured, but the interface value itself must be non-nil). +type MCPOauthTokenCacheManager interface { + EvictOauthTokenCacheByID(ctx context.Context, tokenID string) + EvictOauthTokenCacheByMCPClient(ctx context.Context, mcpClientID string) +} + // MCPSessionsHandler serves the sessions tab API. type MCPSessionsHandler struct { store *lib.Config + // mcpOauthTokenCacheManager invalidates cached OAuth access tokens after + // this handler writes token rows directly through the configstore. + mcpOauthTokenCacheManager MCPOauthTokenCacheManager } // NewMCPSessionsHandler creates the handler. -func NewMCPSessionsHandler(store *lib.Config) *MCPSessionsHandler { - return &MCPSessionsHandler{store: store} +func NewMCPSessionsHandler(store *lib.Config, mcpOauthTokenCacheManager MCPOauthTokenCacheManager) *MCPSessionsHandler { + return &MCPSessionsHandler{store: store, mcpOauthTokenCacheManager: mcpOauthTokenCacheManager} } // RegisterRoutes registers the sessions tab routes. @@ -666,6 +682,7 @@ func (h *MCPSessionsHandler) revoke(ctx *fasthttp.RequestCtx) { SendError(ctx, fasthttp.StatusInternalServerError, "Failed to delete MCP session") return } + h.mcpOauthTokenCacheManager.EvictOauthTokenCacheByID(ctx, tok.ID) logger.Debug("[mcp/sessions] revoked: token=%s mcp_client=%s mode=%s", rowID, tok.MCPClientID, tok.AuthMode) ctx.SetStatusCode(fasthttp.StatusNoContent) } diff --git a/transports/bifrost-http/server/server.go b/transports/bifrost-http/server/server.go index 6572b94ea21..2c09b417106 100644 --- a/transports/bifrost-http/server/server.go +++ b/transports/bifrost-http/server/server.go @@ -140,6 +140,15 @@ type ServerCallbacks interface { CloseAndMarkNeedsReauth(ctx context.Context, id string) error DisableMCPClient(ctx context.Context, id string) error EnableMCPClient(ctx context.Context, id string) error + // EvictOauthTokenCacheByID drops the cached per-user MCP OAuth access + // token backed by the given token row ID after a database write that + // bypassed the OAuth provider's own write paths. + EvictOauthTokenCacheByID(ctx context.Context, tokenID string) + // EvictOauthTokenCacheByMCPClient drops every cached per-user MCP OAuth + // access token bound to the given MCP client after a client-level + // mutation that invalidates its token rows as a set (credential + // rotation, access reconciliation, client deletion). + EvictOauthTokenCacheByMCPClient(ctx context.Context, mcpClientID string) } // LogRedactionMappingResolverProvider is implemented by servers that can attach reveal data to log-detail responses. @@ -323,6 +332,7 @@ func (s *BifrostHTTPServer) RemoveMCPClient(ctx context.Context, id string) erro if err := s.MCPServerHandler.SyncAllMCPServers(ctx); err != nil { logger.Warn("failed to sync MCP servers after removing client: %v", err) } + s.Config.OAuthProvider.EvictUserTokensByMCPClient(id) return nil } @@ -482,6 +492,7 @@ func (s *BifrostHTTPServer) ReloadVirtualKey(ctx context.Context, id string) (*t store.DeleteModelConfigInMemory(ctx, mcID) } s.MCPServerHandler.SyncVKMCPServer(virtualKey) + s.Config.OAuthProvider.EvictUserTokensByVirtualKey(id) return virtualKey, nil } @@ -504,6 +515,7 @@ func (s *BifrostHTTPServer) RemoveVirtualKey(ctx context.Context, id string) err } governancePlugin.GetGovernanceStore().DeleteVirtualKeyInMemory(ctx, id) s.MCPServerHandler.DeleteVKMCPServer(preloadedVk.Value.GetValue()) + s.Config.OAuthProvider.EvictUserTokensByVirtualKey(id) return nil } @@ -946,6 +958,47 @@ func (s *BifrostHTTPServer) RemoveWebhookEndpoint(ctx context.Context, id string return nil } +// EvictOauthTokenCacheByID drops the cached per-user MCP OAuth access token +// backed by the given token row ID from the in-memory cache after a database +// write. A clustered deployment overrides this to also notify peers. +func (s *BifrostHTTPServer) EvictOauthTokenCacheByID(ctx context.Context, tokenID string) { + if s.Config == nil || s.Config.OAuthProvider == nil { + return + } + s.Config.OAuthProvider.EvictUserTokenByID(tokenID) +} + +// EvictOauthTokenCacheByMCPClient drops every cached per-user MCP OAuth +// access token bound to the given MCP client from the in-memory cache. A +// clustered deployment overrides this to also notify peers. +func (s *BifrostHTTPServer) EvictOauthTokenCacheByMCPClient(ctx context.Context, mcpClientID string) { + if s.Config == nil || s.Config.OAuthProvider == nil { + return + } + s.Config.OAuthProvider.EvictUserTokensByMCPClient(mcpClientID) +} + +// EvictOauthTokenCacheByVirtualKey drops every cached vk-mode MCP OAuth +// access token bound to the given virtual key from the in-memory cache. A +// clustered deployment overrides this to also notify peers. +func (s *BifrostHTTPServer) EvictOauthTokenCacheByVirtualKey(ctx context.Context, virtualKeyID string) { + if s.Config == nil || s.Config.OAuthProvider == nil { + return + } + s.Config.OAuthProvider.EvictUserTokensByVirtualKey(virtualKeyID) +} + +// FlushOauthTokenCache drops every cached per-user MCP OAuth access token +// from the in-memory cache. The coarse fallback for mutations whose blast +// radius cannot be scoped to one client or virtual key. A clustered +// deployment overrides this to also notify peers. +func (s *BifrostHTTPServer) FlushOauthTokenCache(ctx context.Context) { + if s.Config == nil || s.Config.OAuthProvider == nil { + return + } + s.Config.OAuthProvider.FlushUserTokenCache() +} + // ReloadClientConfigFromConfigStore reloads the client config from config store func (s *BifrostHTTPServer) ReloadClientConfigFromConfigStore(ctx context.Context) error { if s.Config == nil || s.Config.ConfigStore == nil { @@ -1874,9 +1927,9 @@ func (s *BifrostHTTPServer) RegisterAPIRoutes(ctx context.Context, callbacks Ser healthHandler := handlers.NewHealthHandler(s.Config) providerHandler := handlers.NewProviderHandler(callbacks, s.Config, s.Client) oauthHandler := handlers.NewOAuthHandler(s.Config.OAuthProvider, s.Client, s.Config) - mcpHandler := handlers.NewMCPHandler(callbacks, callbacks, s.Client, s.Config, oauthHandler) + mcpHandler := handlers.NewMCPHandler(callbacks, callbacks, s.Client, s.Config, oauthHandler, callbacks) mcpPerUserHeadersHandler := handlers.NewMCPPerUserHeadersHandler(callbacks, s.Config, s.TempTokens) - mcpSessionsHandler := handlers.NewMCPSessionsHandler(s.Config) + mcpSessionsHandler := handlers.NewMCPSessionsHandler(s.Config, callbacks) configHandler := handlers.NewConfigHandler(callbacks, s.Config) pluginsHandler := handlers.NewPluginsHandler(callbacks, s.Config.ConfigStore) sessionHandler := handlers.NewSessionHandler(s.Config.ConfigStore, s.WSTicketStore)