mirror of
https://github.com/danny-avila/LibreChat.git
synced 2026-10-21 01:22:54 +00:00
🔒 fix: Single-Flight MCP OAuth Token Refresh per User/Server (#14596)
* 🔒 fix: Single-Flight MCP OAuth Token Refresh per User/Server Concurrent refresh-token redemptions (tool-call 401, ping, reconnect retries, expired-token reads) each replayed the same stored refresh token at the OAuth token endpoint. RFC 9700 reuse detection treats the replay as theft and revokes the entire grant family, forcing manual re-consent every access-token expiry. MCPTokenStorage.forceRefreshTokens is the choke point every refresh path converges on; it now single-flights redemptions per (tenantId, userId, serverName) so concurrent callers share one wire call and receive the same rotated result. The refresh token is re-read from storage inside the locked execution — never from a caller snapshot — so a redemption starting after another refresh completed uses the rotated token instead of replaying the consumed one. Fixes #14583 * 🧪 test: Isolate Single-Flight Keys per Test via Unique Server Names * 🔒 fix: Evict Stalled Refresh Slots, Decouple Waiter Aborts from Shared Redemption Codex review round 1: - A redemption that never settles no longer wedges the single-flight slot until process restart: a stale-entry timer evicts the map entry so later refreshes start fresh, while existing waiters keep their promise. - Caller AbortSignals no longer thread into the shared redemption. An impatient waiter (silent refresh's short timeout) resolves its own wait with null via a per-waiter race; the shared wire call proceeds for everyone else, bounded by transport timeouts plus eviction. * 🔒 fix: Abort Stalled Refreshes Before Slot Release, Hook Cache Invalidation to Redemption Codex review round 2: - The stale timer now aborts the wedged execution instead of deleting its slot; the slot frees only once the execution has settled, and an abort guard before the token-endpoint call stops a woken pre-wire stall from replaying a refresh token a successor already rotated. - New onRefreshSuccess hook runs inside the shared redemption after rotated tokens persist, so the silent-refresh path's mcp_get_tokens cache invalidation fires even when the initiating waiter timed out before the redemption completed. * 📝 docs: Record Post-Dispatch Abort Recovery Rationale on Stale-Refresh Valve
This commit is contained in:
Danny Avila
GitHub
parent
7e74f8eb8c
commit
928b14f5bc
5 changed files with 649 additions and 43 deletions
|
|
@ -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<string, Promise<MCPOAuthTokens | null>>();
|
||||
|
||||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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<string, unknown>) => Promise<void>;
|
||||
|
|
@ -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({
|
||||
|
|
|
|||
|
|
@ -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<typeof MCPTokenStorage.forceRefreshTokens>[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<MCPOAuthTokens>((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<MCPOAuthTokens>((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<InMemoryTokenStore['findToken']>[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<InMemoryTokenStore['findToken']>[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<MCPOAuthTokens>((resolve) => {
|
||||
resolveFirst = resolve;
|
||||
}),
|
||||
);
|
||||
const refreshSecond = jest.fn(
|
||||
() =>
|
||||
new Promise<MCPOAuthTokens>((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<MCPOAuthTokens>((_, 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<MCPOAuthTokens>((_, 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<void>((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<InMemoryTokenStore['findToken']>[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<MCPOAuthTokens>((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<MCPOAuthTokens>((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<MCPOAuthTokens>((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<MCPOAuthTokens>((_, 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({
|
||||
|
|
|
|||
|
|
@ -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:/, '')),
|
||||
}));
|
||||
|
|
|
|||
|
|
@ -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<void>;
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -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<string, Promise<MCPOAuthTokens | null>>();
|
||||
|
||||
/**
|
||||
* 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<MCPOAuthTokens | null> {
|
||||
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<MCPOAuthTokens | null>,
|
||||
signal?: AbortSignal,
|
||||
): Promise<MCPOAuthTokens | null> {
|
||||
if (!signal) {
|
||||
return promise;
|
||||
}
|
||||
if (signal.aborted) {
|
||||
return Promise.resolve(null);
|
||||
}
|
||||
return new Promise<MCPOAuthTokens | null>((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<GetTokensParams['refreshTokens']>;
|
||||
createToken: NonNullable<GetTokensParams['createToken']>;
|
||||
/** Internal stale-abort signal owned by `forceRefreshTokens` — never a caller's. */
|
||||
signal: AbortSignal;
|
||||
}): Promise<MCPOAuthTokens | null> {
|
||||
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,
|
||||
});
|
||||
}
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue