From 243fa72d0ec2b5adb6907b02cdcda18266ec14dd Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Fri, 29 May 2026 12:03:22 -0700 Subject: [PATCH] fix: Release timed-out MCP refresh locks --- packages/api/src/mcp/MCPConnectionFactory.ts | 28 ++--- .../__tests__/MCPConnectionFactory.test.ts | 100 +++++++++++++++++- 2 files changed, 114 insertions(+), 14 deletions(-) diff --git a/packages/api/src/mcp/MCPConnectionFactory.ts b/packages/api/src/mcp/MCPConnectionFactory.ts index d331178443..c6a72737c8 100644 --- a/packages/api/src/mcp/MCPConnectionFactory.ts +++ b/packages/api/src/mcp/MCPConnectionFactory.ts @@ -86,6 +86,7 @@ export class MCPConnectionFactory { * connect budget for OAuth discovery, registration, and `oauthStart`. */ private static readonly SILENT_REFRESH_TIMEOUT_MS = 5_000; + private static readonly SILENT_REFRESH_ABORT_GRACE_MS = 1_000; /** Creates a new MCP connection with optional OAuth support */ static async create( @@ -534,10 +535,15 @@ export class MCPConnectionFactory { const timeoutMs = this.getSilentRefreshTimeoutMs(); const abortController = new AbortController(); let timeoutId: ReturnType | null = null; + let abortGraceTimeoutId: ReturnType | null = null; const refreshPromise = this.runSilentRefresh(abortController.signal); const promise = new Promise((resolve) => { timeoutId = setTimeout(() => { abortController.abort(); + abortGraceTimeoutId = setTimeout( + releaseLock, + MCPConnectionFactory.SILENT_REFRESH_ABORT_GRACE_MS, + ); logger.info( `${this.logPrefix} Silent token refresh timed out after ${timeoutMs}ms, falling back to interactive OAuth`, ); @@ -556,19 +562,17 @@ export class MCPConnectionFactory { clearTimeout(timeoutId); } }); + function releaseLock() { + if (abortGraceTimeoutId) { + clearTimeout(abortGraceTimeoutId); + abortGraceTimeoutId = null; + } + if (MCPConnectionFactory.inflightSilentRefreshes.get(lockKey) === promise) { + MCPConnectionFactory.inflightSilentRefreshes.delete(lockKey); + } + } MCPConnectionFactory.inflightSilentRefreshes.set(lockKey, promise); - void refreshPromise.then( - () => { - if (MCPConnectionFactory.inflightSilentRefreshes.get(lockKey) === promise) { - MCPConnectionFactory.inflightSilentRefreshes.delete(lockKey); - } - }, - () => { - if (MCPConnectionFactory.inflightSilentRefreshes.get(lockKey) === promise) { - MCPConnectionFactory.inflightSilentRefreshes.delete(lockKey); - } - }, - ); + void refreshPromise.then(releaseLock, releaseLock); return await promise; } diff --git a/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts b/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts index eb3440e708..80770ea8ff 100644 --- a/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts +++ b/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts @@ -1323,8 +1323,9 @@ describe('MCPConnectionFactory', () => { it('should keep the in-flight silent-refresh lock until the aborted refresh settles', async () => { // A timed-out caller should fall back to interactive OAuth immediately, - // but the shared lock must remain until the underlying abortable refresh - // settles so another 401 cannot redeem the same rotating refresh token. + // but the shared lock must remain briefly while the underlying abortable + // refresh settles so another 401 cannot immediately redeem the same + // rotating refresh token. const sseConfig = { ...mockServerConfig, url: 'https://api.example.com', @@ -1427,6 +1428,101 @@ describe('MCPConnectionFactory', () => { } }); + it('should drop timed-out silent-refresh locks after abort grace if refresh stays hung', async () => { + const sseConfig = { + ...mockServerConfig, + url: 'https://api.example.com', + type: 'sse' as const, + } as t.SSEOptions; + + const basicOptions = { + serverName: 'test-server', + serverConfig: sseConfig, + }; + + const oauthOptions = { + useOAuth: true as const, + user: mockUser, + flowManager: mockFlowManager, + oauthStart: jest.fn(), + oauthEnd: jest.fn(), + tokenMethods: { + findToken: jest.fn(), + createToken: jest.fn(), + updateToken: jest.fn(), + deleteTokens: jest.fn(), + }, + }; + + mockProcessMCPEnv.mockReturnValue(sseConfig); + mockMCPOAuthHandler.generateFlowId.mockReturnValue('flow123'); + let firstRefreshSignal: AbortSignal | undefined; + let secondRefreshSignal: AbortSignal | undefined; + mockMCPTokenStorage.forceRefreshTokens + .mockImplementationOnce((params) => { + firstRefreshSignal = params.signal; + return new Promise(() => {}); + }) + .mockImplementationOnce((params) => { + secondRefreshSignal = params.signal; + return Promise.resolve(null as unknown as MCPOAuthTokens); + }); + mockMCPOAuthHandler.initiateOAuthFlow.mockResolvedValue({ + authorizationUrl: 'https://auth.example.com', + flowId: 'flow123', + flowMetadata: { + serverName: 'test-server', + userId: 'user123', + serverUrl: 'https://api.example.com', + state: 'state-x', + }, + }); + mockFlowManager.getFlowState.mockResolvedValue(null); + mockFlowManager.createFlow.mockResolvedValue(null); + mockConnectionInstance.connect.mockRejectedValue(new Error('OAuth authentication required')); + mockConnectionInstance.isConnected.mockResolvedValue(false); + + let oauthRequiredHandler: (data: Record) => Promise; + mockConnectionInstance.on.mockImplementation((event, handler) => { + if (event === 'oauthRequired') { + oauthRequiredHandler = handler as (data: Record) => Promise; + } + return mockConnectionInstance; + }); + + try { + await MCPConnectionFactory.create(basicOptions, oauthOptions); + } catch { + // Expected + } + + jest.useFakeTimers(); + try { + const firstAttempt = oauthRequiredHandler!({ serverUrl: 'https://api.example.com' }); + await jest.advanceTimersByTimeAsync(2_001); + await firstAttempt; + + const inflightMap = ( + MCPConnectionFactory as unknown as { + inflightSilentRefreshes: Map; + } + ).inflightSilentRefreshes; + expect(inflightMap.size).toBe(1); + expect(firstRefreshSignal?.aborted).toBe(true); + + await jest.advanceTimersByTimeAsync(1_001); + expect(inflightMap.size).toBe(0); + + const secondAttempt = oauthRequiredHandler!({ serverUrl: 'https://api.example.com' }); + await secondAttempt; + + expect(mockMCPTokenStorage.forceRefreshTokens).toHaveBeenCalledTimes(2); + expect(secondRefreshSignal).toBeDefined(); + } finally { + jest.useRealTimers(); + } + }); + it('should reserve interactive OAuth fallback time when initTimeout is omitted', async () => { const sseConfig = { url: 'https://api.example.com',