diff --git a/packages/api/src/mcp/MCPConnectionFactory.ts b/packages/api/src/mcp/MCPConnectionFactory.ts index 01f20cb1bf..c5f5180035 100644 --- a/packages/api/src/mcp/MCPConnectionFactory.ts +++ b/packages/api/src/mcp/MCPConnectionFactory.ts @@ -87,9 +87,10 @@ export class MCPConnectionFactory { * The map only holds in-flight promises (no result caching), so each new 401 * after the previous attempt resolves triggers a fresh redemption. * - * NOTE: this is a single-process lock. Across multiple worker processes the - * race with `getOAuthTokens()`'s `mcp_get_tokens` flow remains; that's an - * inherent limitation of process-local coalescing and tracked separately. + * NOTE: `MCPTokenStorage.forceRefreshTokens` additionally single-flights the + * refresh-token redemption itself, covering the `mcp_get_tokens` and + * reconnect paths that bypass this map. Both locks are process-local; the + * cross-worker race is an inherent limitation and tracked separately. */ private static inflightSilentRefreshes = new Map>(); @@ -717,16 +718,19 @@ export class MCPConnectionFactory { deleteTokens: this.tokenMethods!.deleteTokens, refreshTokens: this.createRefreshTokensFunction(), signal, + /** + * Drop any previously cached `mcp_get_tokens` result so the next + * `getOAuthTokens` reads the freshly persisted tokens rather than the + * now-stale flow-cached value. Attached to the shared redemption + * (not this waiter) so a refresh that completes after this caller's + * timeout still invalidates the cache. + */ + onRefreshSuccess: async (refreshed) => { + await this.invalidateGetTokensFlow(refreshed); + }, }), ); - if (tokens) { - // Drop any previously cached `mcp_get_tokens` result so the next call - // to `getOAuthTokens` reads the freshly persisted tokens from the - // token store rather than the now-stale flow-cached value. - await this.invalidateGetTokensFlow(tokens); - } - if (tokens) { logger.info(`${this.logPrefix} Silent token refresh succeeded`); } else { diff --git a/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts b/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts index bd67e68995..01339baaea 100644 --- a/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts +++ b/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts @@ -1392,7 +1392,12 @@ describe('MCPConnectionFactory', () => { completedAt: Date.now(), result: staleCachedTokens, }); - mockMCPTokenStorage.forceRefreshTokens.mockResolvedValueOnce(freshlyRefreshedTokens); + // Mirror the real storage contract: `onRefreshSuccess` runs inside the + // shared redemption after the rotated tokens are persisted. + mockMCPTokenStorage.forceRefreshTokens.mockImplementationOnce(async (params) => { + await params.onRefreshSuccess?.(freshlyRefreshedTokens); + return freshlyRefreshedTokens; + }); mockConnectionInstance.isConnected.mockResolvedValue(false); let oauthRequiredHandler: (data: Record) => Promise; @@ -2674,7 +2679,12 @@ describe('MCPConnectionFactory', () => { (getTenantId as jest.Mock).mockReturnValue('tenant-a'); mockProcessMCPEnv.mockReturnValue(sseConfig); mockMCPOAuthHandler.generateFlowId.mockReturnValue('flow123'); - mockMCPTokenStorage.forceRefreshTokens.mockResolvedValueOnce(refreshedTokens); + // Mirror the real storage contract: `onRefreshSuccess` runs inside the + // shared redemption after the rotated tokens are persisted. + mockMCPTokenStorage.forceRefreshTokens.mockImplementationOnce(async (params) => { + await params.onRefreshSuccess?.(refreshedTokens); + return refreshedTokens; + }); // A concurrent `getOAuthTokens` for this user/server is currently // PENDING — joiners are monitoring it via `monitorFlow`. mockFlowManager.getFlowState.mockResolvedValue({ diff --git a/packages/api/src/mcp/__tests__/MCPOAuthTokenStorage.test.ts b/packages/api/src/mcp/__tests__/MCPOAuthTokenStorage.test.ts index 674350f2b0..628ca202b4 100644 --- a/packages/api/src/mcp/__tests__/MCPOAuthTokenStorage.test.ts +++ b/packages/api/src/mcp/__tests__/MCPOAuthTokenStorage.test.ts @@ -5,6 +5,8 @@ * refresh callback wiring, and ReauthenticationRequiredError paths. */ +import type { TokenMethods } from '@librechat/data-schemas'; +import type { MCPOAuthTokens } from '~/mcp/oauth'; import { MCPTokenStorage, ReauthenticationRequiredError } from '~/mcp/oauth'; import { InMemoryTokenStore } from './helpers/oauthTestServer'; @@ -15,6 +17,7 @@ jest.mock('@librechat/data-schemas', () => ({ error: jest.fn(), debug: jest.fn(), }, + getTenantId: jest.fn(), encryptV2: jest.fn(async (val: string) => `enc:${val}`), decryptV2: jest.fn(async (val: string) => val.replace(/^enc:/, '')), })); @@ -348,7 +351,7 @@ describe('MCPTokenStorage', () => { expect(refreshTokens).toHaveBeenCalledWith( 'rt', expect.objectContaining({ userId: 'u1', serverName: 'srv1' }), - undefined, + expect.any(AbortSignal), ); }); @@ -476,7 +479,7 @@ describe('MCPTokenStorage', () => { expect.objectContaining({ clientInfo: expect.objectContaining({ client_id: 'cid' }), }), - undefined, + expect.any(AbortSignal), ); }); @@ -765,7 +768,7 @@ describe('MCPTokenStorage', () => { serverName: 'srv1', identifier: 'mcp:srv1', }), - undefined, + expect.any(AbortSignal), ); // The new access token is persisted, replacing the stale one. @@ -865,6 +868,439 @@ describe('MCPTokenStorage', () => { }); }); + describe('forceRefreshTokens concurrent coalescing (single-flight)', () => { + /** Seeds an expired access token plus a refresh token for `serverName`. */ + const seedRefreshableTokens = async (serverName = 'srv1', refreshToken = 'rt-1') => { + await store.createToken({ + userId: 'u1', + type: 'mcp_oauth', + identifier: `mcp:${serverName}`, + token: 'enc:expired-access-token', + expiresIn: -1, + }); + await store.createToken({ + userId: 'u1', + type: 'mcp_oauth_refresh', + identifier: `mcp:${serverName}:refresh`, + token: `enc:${refreshToken}`, + expiresIn: 86400, + }); + }; + + const rotatedTokens = (generation: number): MCPOAuthTokens => ({ + access_token: `at-${generation}`, + refresh_token: `rt-${generation}`, + token_type: 'Bearer', + expires_in: 3600, + obtained_at: Date.now(), + }); + + const refreshParams = (refreshTokens: GetTokensRefresh, serverName = 'srv1') => ({ + userId: 'u1', + serverName, + findToken: store.findToken, + createToken: store.createToken, + updateToken: store.updateToken, + refreshTokens, + }); + + type GetTokensRefresh = NonNullable< + Parameters[0]['refreshTokens'] + >; + + /** Flushes the task queue until `predicate` holds (bounded to avoid hangs). */ + const waitFor = async (predicate: () => boolean) => { + for (let i = 0; i < 50 && !predicate(); i++) { + await new Promise((resolve) => setImmediate(resolve)); + } + expect(predicate()).toBe(true); + }; + + it('coalesces concurrent refresh calls into a single token-endpoint redemption', async () => { + // Unique server names per test keep single-flight keys isolated, so a + // failed test can't leave an unsettled in-flight entry that later tests join. + await seedRefreshableTokens('coalesce-srv'); + + let resolveRefresh!: (tokens: MCPOAuthTokens) => void; + const refreshTokens = jest.fn( + () => + new Promise((resolve) => { + resolveRefresh = resolve; + }), + ); + + const params = refreshParams(refreshTokens, 'coalesce-srv'); + const first = MCPTokenStorage.forceRefreshTokens(params); + const second = MCPTokenStorage.forceRefreshTokens(params); + const third = MCPTokenStorage.forceRefreshTokens(params); + + await waitFor(() => refreshTokens.mock.calls.length > 0); + resolveRefresh(rotatedTokens(2)); + + const results = await Promise.all([first, second, third]); + + expect(refreshTokens).toHaveBeenCalledTimes(1); + for (const result of results) { + expect(result!.access_token).toBe('at-2'); + } + + const storedRefresh = await store.findToken({ + userId: 'u1', + type: 'mcp_oauth_refresh', + identifier: 'mcp:coalesce-srv:refresh', + }); + expect(storedRefresh!.token).toBe('enc:rt-2'); + }); + + it('getTokens joins an in-flight refresh instead of replaying the consumed refresh token', async () => { + // Mirrors issue #14583: a silent refresh (401-triggered) is mid-redemption + // when an expired-token read fires its own refresh. Without single-flight, + // the second caller replays rt-1 and RFC 9700 servers revoke the grant family. + await seedRefreshableTokens('join-srv'); + + let resolveRefresh!: (tokens: MCPOAuthTokens) => void; + const refreshTokens = jest.fn( + (_refreshToken: string) => + new Promise((resolve) => { + resolveRefresh = resolve; + }), + ); + + let refreshReads = 0; + const findToken = (async (filter: { type?: string }) => { + if (filter.type === 'mcp_oauth_refresh') { + refreshReads++; + } + return store.findToken(filter as Parameters[0]); + }) as TokenMethods['findToken']; + + const silentRefresh = MCPTokenStorage.forceRefreshTokens({ + ...refreshParams(refreshTokens, 'join-srv'), + findToken, + }); + await waitFor(() => refreshTokens.mock.calls.length === 1); + + const expiredRead = MCPTokenStorage.getTokens({ + ...refreshParams(refreshTokens, 'join-srv'), + findToken, + }); + /** getTokens probes the refresh token (2nd refresh-identifier read) and then + * synchronously joins the in-flight redemption; flush until the probe lands. */ + await waitFor(() => refreshReads >= 2); + await new Promise((resolve) => setImmediate(resolve)); + + resolveRefresh(rotatedTokens(2)); + const [silentResult, readResult] = await Promise.all([silentRefresh, expiredRead]); + + expect(refreshTokens).toHaveBeenCalledTimes(1); + expect(refreshTokens.mock.calls[0][0]).toBe('rt-1'); + expect(silentResult!.access_token).toBe('at-2'); + expect(readResult!.access_token).toBe('at-2'); + }); + + it('redeems the refresh token as stored at redemption time, not the pre-read snapshot', async () => { + // Simulates a refresh completing between getTokens' probe and the + // redemption (e.g. another replica rotated rt-1 → rt-2). The redemption + // must use the freshest stored token, never the probe snapshot. + await seedRefreshableTokens('snapshot-srv', 'rt-1'); + + let probeSeen = false; + const findToken = (async (filter: { type?: string }) => { + const result = await store.findToken( + filter as Parameters[0], + ); + if (filter.type === 'mcp_oauth_refresh' && !probeSeen) { + probeSeen = true; + await store.updateToken( + { userId: 'u1', type: 'mcp_oauth_refresh', identifier: 'mcp:snapshot-srv:refresh' }, + { token: 'enc:rt-2' }, + ); + } + return result; + }) as TokenMethods['findToken']; + + const refreshTokens = jest.fn().mockResolvedValue(rotatedTokens(3)); + + const result = await MCPTokenStorage.getTokens({ + ...refreshParams(refreshTokens, 'snapshot-srv'), + findToken, + }); + + expect(refreshTokens).toHaveBeenCalledTimes(1); + expect(refreshTokens.mock.calls[0][0]).toBe('rt-2'); + expect(result!.access_token).toBe('at-3'); + }); + + it('does not cache results: a refresh after settlement triggers a fresh redemption', async () => { + await seedRefreshableTokens('sequential-srv'); + + const refreshTokens = jest + .fn() + .mockResolvedValueOnce(rotatedTokens(2)) + .mockResolvedValueOnce(rotatedTokens(3)); + + const params = refreshParams(refreshTokens, 'sequential-srv'); + const first = await MCPTokenStorage.forceRefreshTokens(params); + expect(first!.access_token).toBe('at-2'); + + const second = await MCPTokenStorage.forceRefreshTokens(params); + expect(second!.access_token).toBe('at-3'); + + expect(refreshTokens).toHaveBeenCalledTimes(2); + expect(refreshTokens.mock.calls[1][0]).toBe('rt-2'); + }); + + it('does not coalesce refreshes for different servers', async () => { + await seedRefreshableTokens('multi-a'); + await seedRefreshableTokens('multi-b'); + + let resolveFirst!: (tokens: MCPOAuthTokens) => void; + let resolveSecond!: (tokens: MCPOAuthTokens) => void; + const refreshFirst = jest.fn( + () => + new Promise((resolve) => { + resolveFirst = resolve; + }), + ); + const refreshSecond = jest.fn( + () => + new Promise((resolve) => { + resolveSecond = resolve; + }), + ); + + const first = MCPTokenStorage.forceRefreshTokens(refreshParams(refreshFirst, 'multi-a')); + const second = MCPTokenStorage.forceRefreshTokens(refreshParams(refreshSecond, 'multi-b')); + + await waitFor( + () => refreshFirst.mock.calls.length === 1 && refreshSecond.mock.calls.length === 1, + ); + resolveFirst(rotatedTokens(2)); + resolveSecond(rotatedTokens(5)); + + const [firstResult, secondResult] = await Promise.all([first, second]); + expect(firstResult!.access_token).toBe('at-2'); + expect(secondResult!.access_token).toBe('at-5'); + }); + + it('propagates refresh failure to all joined callers and releases the single-flight slot', async () => { + await seedRefreshableTokens('fail-srv'); + + let rejectRefresh!: (error: Error) => void; + const refreshTokens = jest.fn( + () => + new Promise((_, reject) => { + rejectRefresh = reject; + }), + ); + + const params = refreshParams(refreshTokens, 'fail-srv'); + const first = MCPTokenStorage.forceRefreshTokens(params); + const second = MCPTokenStorage.forceRefreshTokens(params); + + await waitFor(() => refreshTokens.mock.calls.length === 1); + rejectRefresh(new Error('network blew up')); + + await expect(first).resolves.toBeNull(); + await expect(second).resolves.toBeNull(); + + refreshTokens.mockResolvedValueOnce(rotatedTokens(2)); + const third = await MCPTokenStorage.forceRefreshTokens(params); + expect(third!.access_token).toBe('at-2'); + expect(refreshTokens).toHaveBeenCalledTimes(2); + }); + + it('aborts a stalled redemption and frees the slot only after it settles', async () => { + jest.useFakeTimers({ doNotFake: ['setImmediate', 'nextTick'] }); + try { + await seedRefreshableTokens('stall-srv'); + + const refreshTokens = jest + .fn() + .mockImplementationOnce( + (_refreshToken: string, _metadata: unknown, signal: AbortSignal) => + new Promise((_, reject) => { + signal.addEventListener('abort', () => reject(new Error('aborted')), { + once: true, + }); + }), + ) + .mockResolvedValueOnce(rotatedTokens(2)); + + const params = refreshParams(refreshTokens, 'stall-srv'); + const stalled = MCPTokenStorage.forceRefreshTokens(params); + await waitFor(() => refreshTokens.mock.calls.length === 1); + + jest.advanceTimersByTime(MCPTokenStorage.INFLIGHT_REFRESH_STALE_MS + 1); + /** The stale abort settles the stalled execution, which frees the slot. */ + await expect(stalled).resolves.toBeNull(); + + const fresh = await MCPTokenStorage.forceRefreshTokens(params); + expect(fresh!.access_token).toBe('at-2'); + expect(refreshTokens).toHaveBeenCalledTimes(2); + } finally { + jest.useRealTimers(); + } + }); + + it('never reaches the token endpoint when a stalled execution wakes after abort', async () => { + jest.useFakeTimers({ doNotFake: ['setImmediate', 'nextTick'] }); + try { + await seedRefreshableTokens('wake-srv'); + + let releaseRead!: () => void; + const readGate = new Promise((resolve) => { + releaseRead = resolve; + }); + let refreshReadStarted = false; + const findToken = (async (filter: { type?: string }) => { + if (filter.type === 'mcp_oauth_refresh') { + refreshReadStarted = true; + await readGate; + } + return store.findToken(filter as Parameters[0]); + }) as TokenMethods['findToken']; + + const refreshTokens = jest.fn(); + const stalled = MCPTokenStorage.forceRefreshTokens({ + ...refreshParams(refreshTokens, 'wake-srv'), + findToken, + }); + await waitFor(() => refreshReadStarted); + + jest.advanceTimersByTime(MCPTokenStorage.INFLIGHT_REFRESH_STALE_MS + 1); + releaseRead(); + + await expect(stalled).resolves.toBeNull(); + expect(refreshTokens).not.toHaveBeenCalled(); + } finally { + jest.useRealTimers(); + } + }); + + it('runs onRefreshSuccess on the shared redemption even after the initiating waiter aborted', async () => { + await seedRefreshableTokens('hook-srv'); + + let resolveRefresh!: (tokens: MCPOAuthTokens) => void; + const refreshTokens = jest.fn( + () => + new Promise((resolve) => { + resolveRefresh = resolve; + }), + ); + const onRefreshSuccess = jest.fn().mockResolvedValue(undefined); + + const controller = new AbortController(); + const initiator = MCPTokenStorage.forceRefreshTokens({ + ...refreshParams(refreshTokens, 'hook-srv'), + signal: controller.signal, + onRefreshSuccess, + }); + await waitFor(() => refreshTokens.mock.calls.length === 1); + + controller.abort(); + await expect(initiator).resolves.toBeNull(); + expect(onRefreshSuccess).not.toHaveBeenCalled(); + + resolveRefresh(rotatedTokens(2)); + await waitFor(() => onRefreshSuccess.mock.calls.length === 1); + expect(onRefreshSuccess).toHaveBeenCalledWith( + expect.objectContaining({ access_token: 'at-2' }), + ); + }); + + it("an initiator's abort resolves only its own wait, not the shared redemption", async () => { + await seedRefreshableTokens('abort-srv'); + + let resolveRefresh!: (tokens: MCPOAuthTokens) => void; + const refreshTokens = jest.fn( + () => + new Promise((resolve) => { + resolveRefresh = resolve; + }), + ); + + const params = refreshParams(refreshTokens, 'abort-srv'); + const controller = new AbortController(); + const initiator = MCPTokenStorage.forceRefreshTokens({ + ...params, + signal: controller.signal, + }); + await waitFor(() => refreshTokens.mock.calls.length === 1); + const joiner = MCPTokenStorage.forceRefreshTokens(params); + + controller.abort(); + await expect(initiator).resolves.toBeNull(); + + resolveRefresh(rotatedTokens(2)); + const joined = await joiner; + expect(joined!.access_token).toBe('at-2'); + expect(refreshTokens).toHaveBeenCalledTimes(1); + }); + + it("a joiner's abort resolves only its own wait, not the shared redemption", async () => { + await seedRefreshableTokens('joiner-abort-srv'); + + let resolveRefresh!: (tokens: MCPOAuthTokens) => void; + const refreshTokens = jest.fn( + () => + new Promise((resolve) => { + resolveRefresh = resolve; + }), + ); + + const params = refreshParams(refreshTokens, 'joiner-abort-srv'); + const initiator = MCPTokenStorage.forceRefreshTokens(params); + await waitFor(() => refreshTokens.mock.calls.length === 1); + + const controller = new AbortController(); + const joiner = MCPTokenStorage.forceRefreshTokens({ + ...params, + signal: controller.signal, + }); + + controller.abort(); + await expect(joiner).resolves.toBeNull(); + + resolveRefresh(rotatedTokens(2)); + const initiated = await initiator; + expect(initiated!.access_token).toBe('at-2'); + expect(refreshTokens).toHaveBeenCalledTimes(1); + }); + + it('propagates ReauthenticationRequiredError to concurrent callers', async () => { + await seedRefreshableTokens('reauth-srv'); + await store.createToken({ + userId: 'u1', + type: 'mcp_oauth_client', + identifier: 'mcp:reauth-srv:client', + token: 'enc:{"client_id":"cid"}', + expiresIn: 86400, + }); + + let rejectRefresh!: (error: Error) => void; + const refreshTokens = jest.fn( + () => + new Promise((_, reject) => { + rejectRefresh = reject; + }), + ); + + const params = { + ...refreshParams(refreshTokens, 'reauth-srv'), + deleteTokens: store.deleteTokens, + }; + const first = MCPTokenStorage.forceRefreshTokens(params); + const second = MCPTokenStorage.forceRefreshTokens(params); + + await waitFor(() => refreshTokens.mock.calls.length === 1); + rejectRefresh(new Error('invalid_client')); + + await expect(first).rejects.toThrow(ReauthenticationRequiredError); + await expect(second).rejects.toThrow(ReauthenticationRequiredError); + }); + }); + describe('storeTokens + getTokens round-trip', () => { it('should store and retrieve tokens with full encrypt/decrypt cycle', async () => { await MCPTokenStorage.storeTokens({ diff --git a/packages/api/src/mcp/oauth/tokens.test.ts b/packages/api/src/mcp/oauth/tokens.test.ts index f22c7f69c8..4a5944c866 100644 --- a/packages/api/src/mcp/oauth/tokens.test.ts +++ b/packages/api/src/mcp/oauth/tokens.test.ts @@ -10,6 +10,7 @@ jest.mock('@librechat/data-schemas', () => ({ error: jest.fn(), debug: jest.fn(), }, + getTenantId: jest.fn(), encryptV2: jest.fn(async (value: string) => `encrypted:${value}`), decryptV2: jest.fn(async (value: string) => value.replace(/^encrypted:/, '')), })); diff --git a/packages/api/src/mcp/oauth/tokens.ts b/packages/api/src/mcp/oauth/tokens.ts index 65c686e7ef..b242072bc2 100644 --- a/packages/api/src/mcp/oauth/tokens.ts +++ b/packages/api/src/mcp/oauth/tokens.ts @@ -1,5 +1,5 @@ import jwt from 'jsonwebtoken'; -import { logger, encryptV2, decryptV2 } from '@librechat/data-schemas'; +import { logger, encryptV2, decryptV2, getTenantId } from '@librechat/data-schemas'; import type { OAuthTokens, OAuthClientInformation } from '@modelcontextprotocol/sdk/shared/auth.js'; import type { TokenMethods, IToken } from '@librechat/data-schemas'; import type { MCPOAuthTokens, ExtendedOAuthTokens, OAuthStoredClientMetadata } from './types'; @@ -55,7 +55,14 @@ interface GetTokensParams { updateToken?: TokenMethods['updateToken']; /** Enables cleanup of stale client registration and refresh token on invalid_client errors during refresh. */ deleteTokens?: TokenMethods['deleteTokens']; + /** Waiter-specific: aborting resolves this caller's wait with `null` without cancelling the shared redemption. */ signal?: AbortSignal; + /** + * Invoked inside the shared redemption after rotated tokens are persisted. + * Runs even when the initiating waiter has already aborted its wait, so + * cache invalidation tied to the fresh tokens cannot be skipped by a timeout. + */ + onRefreshSuccess?: (tokens: MCPOAuthTokens) => Promise; } /** @@ -86,6 +93,29 @@ function getJwtAccessTokenExpiry(accessToken?: string): number | null { } export class MCPTokenStorage { + /** + * Process-local in-flight refresh-token redemptions, keyed by + * `tenantId:userId:serverName`. Every code path that redeems a refresh token + * (expired-token refresh via `getTokens`, silent refresh on 401, reconnect + * retries) converges on `forceRefreshTokens`, so coalescing here guarantees + * at most one wire call to the token endpoint per user/server at a time. + * RFC 9700 servers treat a replayed (already-consumed) refresh token as + * theft and revoke the entire grant family, so concurrent redemptions are + * not merely wasteful — they destroy the freshly issued tokens. Only + * in-flight promises are held (no result caching): each new refresh request + * after settlement triggers a fresh redemption. + */ + private static inflightRefreshes = new Map>(); + + /** + * How long an in-flight redemption may run before it is aborted. Generous + * relative to a healthy refresh round trip (well under a minute) so the + * abort only fires for genuinely wedged executions. The single-flight slot + * is freed when the aborted execution settles — never while it might still + * reach the token endpoint with the old refresh token. + */ + static readonly INFLIGHT_REFRESH_STALE_MS = 60_000; + static getLogPrefix(userId: string, serverName: string): string { return isSystemUserId(userId) ? `[MCP][${serverName}]` @@ -301,11 +331,131 @@ export class MCPTokenStorage { * server has signaled token invalidity (e.g. a 401 mid-session) — the 401 is * the authoritative signal, not the local `expires_at`. * + * Single-flighted per `(tenantId, userId, serverName)`: concurrent callers + * (tool-call 401s, pings, reconnect retries, expired-token reads) share one + * redemption and receive the same rotated result instead of each replaying + * the refresh token at the token endpoint. + * * Returns the new tokens, or `null` when refresh is not possible (no refresh * token stored, no refresh callback, etc.). Throws `ReauthenticationRequiredError` * when the refresh server response indicates the client registration is stale. */ - static async forceRefreshTokens({ + static async forceRefreshTokens( + params: GetTokensParams & { + existingAccessToken?: IToken | null; + }, + ): Promise { + const { userId, serverName, refreshTokens, createToken, signal } = params; + const logPrefix = this.getLogPrefix(userId, serverName); + + const refreshKey = `${getTenantId() ?? ''}:${userId}:${serverName}`; + const inflight = this.inflightRefreshes.get(refreshKey); + if (inflight) { + logger.debug(`${logPrefix} Joining in-flight token refresh`); + return this.raceWithAbort(inflight, signal); + } + + if (!refreshTokens) { + logger.warn(`${logPrefix} Cannot refresh tokens: no \`refreshTokens\` callback provided`); + return null; + } + + if (!createToken) { + logger.warn(`${logPrefix} Cannot refresh tokens: no \`createToken\` function provided`); + return null; + } + + /** + * The shared redemption is owner-neutral: no caller's `AbortSignal` is + * threaded into the execution, so an impatient waiter (e.g. the silent + * refresh path's short timeout) cannot cancel the wire call for everyone + * who joined. Cancellation is waiter-specific via `raceWithAbort`; the + * execution itself is bounded by the internal stale-abort controller below. + */ + const executionController = new AbortController(); + const refreshPromise = this.executeTokenRefresh({ + ...params, + refreshTokens, + createToken, + signal: executionController.signal, + }).finally(() => { + clearTimeout(staleTimer); + if (this.inflightRefreshes.get(refreshKey) === refreshPromise) { + this.inflightRefreshes.delete(refreshKey); + } + }); + /** + * Safety valve for wedged executions: after the stale window the execution + * is aborted so it can never reach the token endpoint with a refresh token + * that a successor is about to redeem. The slot itself is freed only by the + * `.finally` above, that is, once the aborted execution has actually + * settled. Deleting the entry while the redemption might still consume the + * stored refresh token would re-open the concurrent-replay window this + * single-flight exists to close. + * + * If the abort lands after the endpoint already processed the request + * (response lost in transit), the rotated tokens are unrecoverable and the + * stored refresh token is deliberately left in place rather than deleted. + * A later redemption then either succeeds (request never actually + * processed, or the server grants rotation leeway) or trips reuse + * detection on a family whose fresh tokens were never received and whose + * access token was already expired or rejected. That failure ends in the + * same re-authentication the proactive deletion would force on every + * stall, while deletion would also foreclose the silent recovery paths. + */ + const staleTimer = setTimeout(() => { + if (this.inflightRefreshes.get(refreshKey) === refreshPromise) { + logger.warn( + `${logPrefix} Aborting stalled in-flight token refresh after ${MCPTokenStorage.INFLIGHT_REFRESH_STALE_MS}ms`, + ); + executionController.abort(); + } + }, MCPTokenStorage.INFLIGHT_REFRESH_STALE_MS); + staleTimer.unref?.(); + this.inflightRefreshes.set(refreshKey, refreshPromise); + return this.raceWithAbort(refreshPromise, signal); + } + + /** + * Wraps the shared redemption promise for a single waiter: if the waiter's + * `signal` aborts first, that waiter receives `null` (matching the prior + * abort contract) while the shared execution continues for other callers. + */ + private static raceWithAbort( + promise: Promise, + signal?: AbortSignal, + ): Promise { + if (!signal) { + return promise; + } + if (signal.aborted) { + return Promise.resolve(null); + } + return new Promise((resolve, reject) => { + const onAbort = () => resolve(null); + signal.addEventListener('abort', onAbort, { once: true }); + promise.then( + (value) => { + signal.removeEventListener('abort', onAbort); + resolve(value); + }, + (error: unknown) => { + signal.removeEventListener('abort', onAbort); + reject(error); + }, + ); + }); + } + + /** + * Runs a single refresh-token redemption under the `inflightRefreshes` + * single-flight lock. The refresh token is read from storage at execution + * time — never from a caller-provided snapshot — so a redemption that starts + * after another refresh completed uses the rotated token instead of + * replaying the consumed one (which RFC 9700 reuse detection punishes by + * revoking the whole grant family). + */ + private static async executeTokenRefresh({ userId, serverName, findToken, @@ -314,39 +464,29 @@ export class MCPTokenStorage { deleteTokens, refreshTokens, existingAccessToken, - existingRefreshToken, + onRefreshSuccess, signal, }: GetTokensParams & { existingAccessToken?: IToken | null; - existingRefreshToken?: IToken | null; + refreshTokens: NonNullable; + createToken: NonNullable; + /** Internal stale-abort signal owned by `forceRefreshTokens` — never a caller's. */ + signal: AbortSignal; }): Promise { const logPrefix = this.getLogPrefix(userId, serverName); const identifier = `mcp:${serverName}`; - const refreshTokenData = - existingRefreshToken !== undefined - ? existingRefreshToken - : await findToken({ - userId, - type: 'mcp_oauth_refresh', - identifier: `${identifier}:refresh`, - }); + const refreshTokenData = await findToken({ + userId, + type: 'mcp_oauth_refresh', + identifier: `${identifier}:refresh`, + }); if (!refreshTokenData) { logger.debug(`${logPrefix} No refresh token in storage`); return null; } - if (!refreshTokens) { - logger.warn(`${logPrefix} Refresh token available but no \`refreshTokens\` provided`); - return null; - } - - if (!createToken) { - logger.warn(`${logPrefix} Refresh token available but no \`createToken\` function provided`); - return null; - } - try { logger.info(`${logPrefix} Attempting to refresh token`); const decryptedRefreshToken = await decryptV2(refreshTokenData.token); @@ -402,12 +542,17 @@ export class MCPTokenStorage { resource, }; - const newTokens = await refreshTokens(decryptedRefreshToken, metadata, signal); - - if (signal?.aborted) { - throw new Error('Token refresh aborted'); + /** + * A stalled execution may wake here after the stale-abort fired and a + * successor already redeemed (and rotated) the refresh token — reaching + * the endpoint with the old token would trip RFC 9700 reuse detection. + */ + if (signal.aborted) { + throw new Error('Token refresh aborted before reaching the token endpoint'); } + const newTokens = await refreshTokens(decryptedRefreshToken, metadata, signal); + logger.debug(`${logPrefix} Refresh completed`, { has_new_access_token: !!newTokens.access_token, has_new_refresh_token: !!newTokens.refresh_token, @@ -433,6 +578,14 @@ export class MCPTokenStorage { metadata: storedClientMetadata, }); + if (onRefreshSuccess) { + try { + await onRefreshSuccess(newTokens); + } catch (hookError) { + logger.warn(`${logPrefix} onRefreshSuccess callback failed`, hookError); + } + } + logger.info(`${logPrefix} Successfully refreshed and stored OAuth tokens`); return newTokens; } catch (refreshError) { @@ -504,7 +657,10 @@ export class MCPTokenStorage { logger.info(`${logPrefix} Access token ${isMissing ? 'missing' : 'expired'}`); /** Probe for a refresh token first so we can throw `ReauthenticationRequiredError` - * when none exists, matching the prior contract. */ + * when none exists, matching the prior contract. The probe result is + * intentionally not passed to `forceRefreshTokens` — the redemption + * re-reads storage under its single-flight lock so it never replays a + * refresh token that a concurrent refresh already consumed. */ const refreshTokenData = await findToken({ userId, type: 'mcp_oauth_refresh', @@ -528,7 +684,6 @@ export class MCPTokenStorage { deleteTokens, refreshTokens, existingAccessToken: accessTokenData, - existingRefreshToken: refreshTokenData, }); }