From 7fc62023eb4f1c65d4e5ab4580f8b46f0048eaab Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 10 Aug 2026 10:38:34 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=A7=B7=20fix:=20Safely=20Recover=20Runtim?= =?UTF-8?q?e=20MCP=20OAuth=20Rejections=20(#14684)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix runtime MCP OAuth recovery * style: sort LC-008 imports * fix: single-flight runtime OAuth handlers * fix: retain transport OAuth failures for recovery * fix(mcp): preserve OAuth recovery connections * test(mcp): type request-scoped config fixture * fix(mcp): harden shared OAuth recovery * fix(mcp): bound OAuth recovery escalation * style(mcp): sort OAuth integration imports * fix(mcp): harden OAuth recovery boundaries * fix(mcp): abort shared recovery waiters * fix(mcp): bound request OAuth recovery phases * fix(mcp): close OAuth recovery ownership gaps * fix(mcp): retry borrowers closed by OAuth recovery * fix(mcp): drain borrowers before OAuth reconnect * fix(mcp): preserve eviction across OAuth recovery * fix(mcp): unify OAuth recovery leases * fix(mcp): serialize cache reuse with recovery * fix(mcp): make recovery checkout atomic * test(mcp): use numeric config timestamp * fix(mcp): reacquire recovery checkouts * fix(mcp): retain shared recovery disposal * fix(mcp): restart checkout after recovery takeover * fix(mcp): close recovery lifecycle gaps * refactor(mcp): deepen OAuth recovery lifecycle * fix(mcp): harden OAuth lifecycle disposal * style(mcp): sort OAuth lifecycle imports * fix: lease MCP OAuth lifecycle edges * fix(mcp): isolate shared OAuth flows from aborts --------- Co-authored-by: Dennis Schenk --- api/server/routes/__tests__/mcp.spec.js | 136 +- api/server/routes/mcp.js | 52 +- api/server/services/MCP.js | 133 +- api/server/services/MCP.spec.js | 96 +- packages/api/src/mcp/MCPConnectionFactory.ts | 176 +- packages/api/src/mcp/MCPManager.ts | 922 ++++++++--- packages/api/src/mcp/UserConnectionManager.ts | 400 +++-- .../src/mcp/__tests__/MCPConnection.test.ts | 112 +- .../MCPConnectionAgentLifecycle.test.ts | 16 +- ...ectionFactory.oauthSdk.integration.test.ts | 270 +++- .../__tests__/MCPConnectionFactory.test.ts | 190 ++- .../api/src/mcp/__tests__/MCPManager.test.ts | 1420 ++++++++++++++++- .../__tests__/MCPOAuthRaceCondition.test.ts | 60 +- .../mcp/__tests__/helpers/oauthTestServer.ts | 20 +- .../api/src/mcp/__tests__/request.test.ts | 33 +- packages/api/src/mcp/connection.ts | 62 +- packages/api/src/mcp/errors.ts | 41 + packages/api/src/mcp/oauth/pending.ts | 139 ++ packages/api/src/mcp/request.ts | 28 +- packages/api/src/mcp/types/index.ts | 3 + 20 files changed, 3439 insertions(+), 870 deletions(-) diff --git a/api/server/routes/__tests__/mcp.spec.js b/api/server/routes/__tests__/mcp.spec.js index ed919ba703..7b2e702a67 100644 --- a/api/server/routes/__tests__/mcp.spec.js +++ b/api/server/routes/__tests__/mcp.spec.js @@ -187,6 +187,11 @@ const mockOAuthCompletion = (tokens) => { ); }; +const createLeasedMcpManager = (connection, overrides = {}) => ({ + ...overrides, + withUserConnectionLease: jest.fn((_options, useConnection) => useConnection(connection)), +}); + describe('MCP Routes', () => { let app; let mongoServer; @@ -853,11 +858,9 @@ describe('MCP Routes', () => { MCPTokenStorage.storeTokens.mockResolvedValue(); mockRegistryInstance.getServerConfig.mockResolvedValue({}); - const mockMcpManager = { - getUserConnection: jest.fn().mockResolvedValue({ - fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: [], complete: true }), - }), - }; + const mockMcpManager = createLeasedMcpManager({ + fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: [], complete: true }), + }); require('~/config').getMCPManager.mockReturnValue(mockMcpManager); require('~/config').getOAuthReconnectionManager.mockReturnValue({ clearReconnection: jest.fn(), @@ -915,11 +918,9 @@ describe('MCP Routes', () => { MCPTokenStorage.storeTokens.mockResolvedValue(); mockRegistryInstance.getServerConfig.mockResolvedValue({}); - const mockMcpManager = { - getUserConnection: jest.fn().mockResolvedValue({ - fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: [], complete: true }), - }), - }; + const mockMcpManager = createLeasedMcpManager({ + fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: [], complete: true }), + }); require('~/config').getMCPManager.mockReturnValue(mockMcpManager); require('~/config').getOAuthReconnectionManager.mockReturnValue({ clearReconnection: jest.fn(), @@ -973,11 +974,7 @@ describe('MCP Routes', () => { const fetchOrderedToolsSnapshot = jest .fn() .mockResolvedValue({ tools: fetchedTools, complete: true }); - const mockMcpManager = { - getUserConnection: jest.fn().mockResolvedValue({ - fetchOrderedToolsSnapshot, - }), - }; + const mockMcpManager = createLeasedMcpManager({ fetchOrderedToolsSnapshot }); require('~/config').getMCPManager.mockReturnValue(mockMcpManager); require('~/config').getOAuthReconnectionManager.mockReturnValue({ clearReconnection: jest.fn(), @@ -992,8 +989,9 @@ describe('MCP Routes', () => { expect(response.status).toBe(302); expect(mockResolveAllMcpConfigs).toHaveBeenCalledWith('test-user-id'); expect(fetchOrderedToolsSnapshot).toHaveBeenCalledTimes(1); - expect(mockMcpManager.getUserConnection).toHaveBeenCalledWith( + expect(mockMcpManager.withUserConnectionLease).toHaveBeenCalledWith( expect.objectContaining({ serverConfig: mergedServerConfig }), + expect.any(Function), ); expect(updateMCPServerTools).toHaveBeenCalledWith({ userId: 'test-user-id', @@ -1044,13 +1042,9 @@ describe('MCP Routes', () => { [`mcp_test-server`]: { LITELLM_KEY: 'sk-real-user-key' }, }); - const mockMcpManager = { - getUserConnection: jest.fn().mockResolvedValue({ - fetchToolsSnapshot: jest - .fn() - .mockResolvedValue({ tools: fetchedTools, complete: true }), - }), - }; + const mockMcpManager = createLeasedMcpManager({ + fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: fetchedTools, complete: true }), + }); require('~/config').getMCPManager.mockReturnValue(mockMcpManager); require('~/config').getOAuthReconnectionManager.mockReturnValue({ clearReconnection: jest.fn(), @@ -1068,8 +1062,9 @@ describe('MCP Routes', () => { servers: ['test-server'], findPluginAuthsByKeys: require('~/models').findPluginAuthsByKeys, }); - expect(mockMcpManager.getUserConnection).toHaveBeenCalledWith( + expect(mockMcpManager.withUserConnectionLease).toHaveBeenCalledWith( expect.objectContaining({ customUserVars: { LITELLM_KEY: 'sk-real-user-key' } }), + expect.any(Function), ); }); @@ -1109,13 +1104,9 @@ describe('MCP Routes', () => { mockResolveAllMcpConfigs.mockResolvedValueOnce({ 'test-server': mergedServerConfig }); require('@librechat/api').getUserMCPAuthMap.mockClear(); - const mockMcpManager = { - getUserConnection: jest.fn().mockResolvedValue({ - fetchToolsSnapshot: jest - .fn() - .mockResolvedValue({ tools: fetchedTools, complete: true }), - }), - }; + const mockMcpManager = createLeasedMcpManager({ + fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: fetchedTools, complete: true }), + }); require('~/config').getMCPManager.mockReturnValue(mockMcpManager); require('~/config').getOAuthReconnectionManager.mockReturnValue({ clearReconnection: jest.fn(), @@ -1129,8 +1120,9 @@ describe('MCP Routes', () => { expect(response.status).toBe(302); expect(require('@librechat/api').getUserMCPAuthMap).not.toHaveBeenCalled(); - expect(mockMcpManager.getUserConnection).toHaveBeenCalledWith( + expect(mockMcpManager.withUserConnectionLease).toHaveBeenCalledWith( expect.objectContaining({ customUserVars: undefined }), + expect.any(Function), ); }); @@ -1241,7 +1233,9 @@ describe('MCP Routes', () => { }), }; const mockMcpManager = { - getUserConnection: jest.fn().mockResolvedValue(mockUserConnection), + withUserConnectionLease: jest.fn((_options, useConnection) => + useConnection(mockUserConnection), + ), }; require('~/config').getMCPManager.mockReturnValue(mockMcpManager); @@ -1285,7 +1279,7 @@ describe('MCP Routes', () => { ); const storeInvocation = MCPTokenStorage.storeTokens.mock.invocationCallOrder[0]; const flowCompletionInvocation = mockFlowManager.completeFlow.mock.invocationCallOrder[0]; - const connectInvocation = mockMcpManager.getUserConnection.mock.invocationCallOrder[0]; + const connectInvocation = mockMcpManager.withUserConnectionLease.mock.invocationCallOrder[0]; expect(storeInvocation).toBeLessThan(flowCompletionInvocation); expect(storeInvocation).toBeLessThan(connectInvocation); expect(mockFlowManager.completeFlow).toHaveBeenCalledWith( @@ -1331,11 +1325,11 @@ describe('MCP Routes', () => { require('~/config').getOAuthReconnectionManager.mockReturnValue({ clearReconnection: jest.fn(), }); - require('~/config').getMCPManager.mockReturnValue({ - getUserConnection: jest.fn().mockResolvedValue({ + require('~/config').getMCPManager.mockReturnValue( + createLeasedMcpManager({ fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: [], complete: true }), }), - }); + ); const { getCachedTools, setCachedTools } = require('~/server/services/Config'); getCachedTools.mockResolvedValue({}); setCachedTools.mockResolvedValue(); @@ -1406,11 +1400,11 @@ describe('MCP Routes', () => { require('~/config').getOAuthReconnectionManager.mockReturnValue({ clearReconnection: jest.fn(), }); - require('~/config').getMCPManager.mockReturnValue({ - getUserConnection: jest.fn().mockResolvedValue({ + require('~/config').getMCPManager.mockReturnValue( + createLeasedMcpManager({ fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: [], complete: true }), }), - }); + ); const { getCachedTools, setCachedTools } = require('~/server/services/Config'); getCachedTools.mockResolvedValue({}); setCachedTools.mockResolvedValue(); @@ -1464,11 +1458,11 @@ describe('MCP Routes', () => { require('~/config').getOAuthReconnectionManager.mockReturnValue({ clearReconnection: jest.fn(), }); - require('~/config').getMCPManager.mockReturnValue({ - getUserConnection: jest.fn().mockResolvedValue({ + require('~/config').getMCPManager.mockReturnValue( + createLeasedMcpManager({ fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: [], complete: true }), }), - }); + ); const { getCachedTools, setCachedTools } = require('~/server/services/Config'); getCachedTools.mockResolvedValue({}); setCachedTools.mockResolvedValue(); @@ -1518,11 +1512,11 @@ describe('MCP Routes', () => { require('~/config').getOAuthReconnectionManager.mockReturnValue({ clearReconnection: jest.fn(), }); - require('~/config').getMCPManager.mockReturnValue({ - getUserConnection: jest.fn().mockResolvedValue({ + require('~/config').getMCPManager.mockReturnValue( + createLeasedMcpManager({ fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: [], complete: true }), }), - }); + ); const { getCachedTools, setCachedTools } = require('~/server/services/Config'); getCachedTools.mockResolvedValue({}); setCachedTools.mockResolvedValue(); @@ -1639,7 +1633,7 @@ describe('MCP Routes', () => { require('~/config').getFlowStateManager.mockReturnValue(mockFlowManager); const mockMcpManager = { - getUserConnection: jest.fn().mockRejectedValue(new Error('Reconnection failed')), + withUserConnectionLease: jest.fn().mockRejectedValue(new Error('Reconnection failed')), }; require('~/config').getMCPManager.mockReturnValue(mockMcpManager); @@ -1692,7 +1686,7 @@ describe('MCP Routes', () => { require('~/config').getFlowStateManager.mockReturnValue(mockFlowManager); const mockMcpManager = { - getUserConnection: jest.fn(), + withUserConnectionLease: jest.fn(), }; require('~/config').getMCPManager.mockReturnValue(mockMcpManager); @@ -1711,7 +1705,7 @@ describe('MCP Routes', () => { expect(response.status).toBe(302); expect(response.headers.location).toBe(`${basePath}/oauth/error?error=callback_failed`); expect(mockFlowManager.completeFlow).not.toHaveBeenCalled(); - expect(mockMcpManager.getUserConnection).not.toHaveBeenCalled(); + expect(mockMcpManager.withUserConnectionLease).not.toHaveBeenCalled(); }); it('should use original flow state credentials when storing tokens', async () => { @@ -1755,9 +1749,7 @@ describe('MCP Routes', () => { const mockUserConnection = { fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: [], complete: true }), }; - const mockMcpManager = { - getUserConnection: jest.fn().mockResolvedValue(mockUserConnection), - }; + const mockMcpManager = createLeasedMcpManager(mockUserConnection); require('~/config').getMCPManager.mockReturnValue(mockMcpManager); require('~/config').getOAuthReconnectionManager = jest.fn().mockReturnValue({ clearReconnection: jest.fn(), @@ -2927,11 +2919,9 @@ describe('MCP Routes', () => { }; require('~/config').getFlowStateManager.mockReturnValue(mockFlowManager); - const mockMcpManager = { - getUserConnection: jest.fn().mockResolvedValue({ - fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: [], complete: true }), - }), - }; + const mockMcpManager = createLeasedMcpManager({ + fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: [], complete: true }), + }); require('~/config').getMCPManager.mockReturnValue(mockMcpManager); const flowId = 'test-user-id:test-server'; @@ -2980,15 +2970,15 @@ describe('MCP Routes', () => { MCPTokenStorage.storeTokens.mockResolvedValue(); mockRegistryInstance.getServerConfig.mockResolvedValue({}); - const mockMcpManager = { - getUserConnection: jest.fn().mockResolvedValue({ + const mockMcpManager = createLeasedMcpManager( + { fetchToolsSnapshot: jest.fn().mockResolvedValue({ tools: [{ name: 'test-tool', description: 'Test tool' }], complete: true, }), - }), - getToolPublicationGeneration: jest.fn().mockReturnValue('oauth-connection-generation'), - }; + }, + { getToolPublicationGeneration: jest.fn().mockReturnValue('oauth-connection-generation') }, + ); require('~/config').getMCPManager.mockReturnValue(mockMcpManager); const flowId = 'test-user-id:test-server'; @@ -3035,15 +3025,19 @@ describe('MCP Routes', () => { mockOAuthCompletion(mockTokens); MCPTokenStorage.storeTokens.mockResolvedValue(); mockRegistryInstance.getServerConfig.mockResolvedValue({}); - require('~/config').getMCPManager.mockReturnValue({ - getUserConnection: jest.fn().mockResolvedValue({ - fetchToolsSnapshot: jest.fn().mockResolvedValue({ - tools: [{ name: 'partial-tool', description: 'Only the first page' }], - complete: false, - }), - }), - getToolPublicationGeneration: jest.fn().mockReturnValue('oauth-connection-generation'), - }); + require('~/config').getMCPManager.mockReturnValue( + createLeasedMcpManager( + { + fetchToolsSnapshot: jest.fn().mockResolvedValue({ + tools: [{ name: 'partial-tool', description: 'Only the first page' }], + complete: false, + }), + }, + { + getToolPublicationGeneration: jest.fn().mockReturnValue('oauth-connection-generation'), + }, + ), + ); const flowId = 'test-user-id:test-server'; const csrfToken = generateTestCsrfToken(flowId); diff --git a/api/server/routes/mcp.js b/api/server/routes/mcp.js index bb4be24618..99174ce59d 100644 --- a/api/server/routes/mcp.js +++ b/api/server/routes/mcp.js @@ -522,33 +522,39 @@ router.get('/:serverName/oauth/callback', async (req, res) => { } const customUserVars = getServerCustomUserVars(userMCPAuthMap, serverName); - const userConnection = await mcpManager.getUserConnection({ - user, - serverName, - flowManager, - serverConfig, - customUserVars, - tokenMethods: { - findToken: db.findToken, - updateToken: db.updateToken, - createToken: db.createToken, - deleteTokens: db.deleteTokens, + const { snapshot, publicationGeneration } = await mcpManager.withUserConnectionLease( + { + user, + serverName, + flowManager, + serverConfig, + customUserVars, + tokenMethods: { + findToken: db.findToken, + updateToken: db.updateToken, + createToken: db.createToken, + deleteTokens: db.deleteTokens, + }, }, - }); + async (userConnection) => { + logger.info( + `[MCP OAuth] Successfully reconnected ${serverName} for user ${flowState.userId}`, + ); - logger.info( - `[MCP OAuth] Successfully reconnected ${serverName} for user ${flowState.userId}`, + const oauthReconnectionManager = getOAuthReconnectionManager(); + oauthReconnectionManager.clearReconnection(flowState.userId, serverName); + + const snapshot = + typeof userConnection.fetchOrderedToolsSnapshot === 'function' + ? await userConnection.fetchOrderedToolsSnapshot() + : await userConnection.fetchToolsSnapshot(); + return { + snapshot, + publicationGeneration: mcpManager.getToolPublicationGeneration?.(userConnection), + }; + }, ); - - const oauthReconnectionManager = getOAuthReconnectionManager(); - oauthReconnectionManager.clearReconnection(flowState.userId, serverName); - - const snapshot = - typeof userConnection.fetchOrderedToolsSnapshot === 'function' - ? await userConnection.fetchOrderedToolsSnapshot() - : await userConnection.fetchToolsSnapshot(); if (snapshot.complete) { - const publicationGeneration = mcpManager.getToolPublicationGeneration?.(userConnection); await updateMCPServerTools({ userId: flowState.userId, serverName, diff --git a/api/server/services/MCP.js b/api/server/services/MCP.js index 6eb0c723a3..d00b4053d3 100644 --- a/api/server/services/MCP.js +++ b/api/server/services/MCP.js @@ -608,24 +608,6 @@ function createOAuthEnd({ res, stepId, toolCall, streamId = null, jobCreatedAt } }; } -/** - * @param {object} params - * @param {string} params.userId - The ID of the user. - * @param {string} params.serverName - The name of the server. - * @param {string} params.toolName - The name of the tool. - * @param {string} [params.tenantId] - The tenant ID for the current request. - * @param {FlowStateManager} params.flowManager - The flow manager instance. - */ -function createAbortHandler({ userId, serverName, toolName, tenantId, flowManager }) { - return function () { - logger.info(`[MCP][User: ${userId}][${serverName}][${toolName}] Tool call aborted`); - const flowId = getOAuthFlowId(userId, serverName, tenantId); - // Clean up both mcp_oauth and mcp_get_tokens flows - flowManager.failFlow(flowId, 'mcp_oauth', new Error('Tool call aborted')); - flowManager.failFlow(flowId, 'mcp_get_tokens', new Error('Tool call aborted')); - }; -} - /** * @param {Object} params * @param {() => Promise} params.runStepEmitter @@ -696,66 +678,43 @@ async function reconnectServer({ serverName, }); - // Set up abort handler to clean up OAuth flows if request is aborted - const tenantId = user?.tenantId ?? getTenantId(); - const oauthFlowId = getOAuthFlowId(user.id, serverName, tenantId); - const abortHandler = () => { - logger.info( - `[MCP][User: ${user.id}][${serverName}] Tool loading aborted, cleaning up OAuth flows`, - ); - // Clean up both mcp_oauth and mcp_get_tokens flows - flowManager.failFlow(oauthFlowId, 'mcp_oauth', new Error('Tool loading aborted')); - flowManager.failFlow(oauthFlowId, 'mcp_get_tokens', new Error('Tool loading aborted')); - }; - - if (signal) { - signal.addEventListener('abort', abortHandler, { once: true }); - } - - try { - const runStepEmitter = createRunStepEmitter({ - res, - index, - runId, - stepId, - toolCall, - streamId, - jobCreatedAt, - }); - const runStepDeltaEmitter = createRunStepDeltaEmitter({ - res, - stepId, - toolCall, - streamId, - jobCreatedAt, - }); - const callback = createOAuthCallback({ runStepEmitter, runStepDeltaEmitter }); - const oauthStart = createOAuthStart({ - res, - flowId, - callback, - flowManager, - }); - return await reinitMCPServer({ - user, - signal, - serverName, - configServers, - oauthStart, - flowManager, - userMCPAuthMap, - requestBody, - requestScopedConnections, - forceNew: true, - returnOnOAuth: false, - connectionTimeout: Time.THIRTY_SECONDS, - }); - } finally { - // Clean up abort handler to prevent memory leaks - if (signal) { - signal.removeEventListener('abort', abortHandler); - } - } + const runStepEmitter = createRunStepEmitter({ + res, + index, + runId, + stepId, + toolCall, + streamId, + jobCreatedAt, + }); + const runStepDeltaEmitter = createRunStepDeltaEmitter({ + res, + stepId, + toolCall, + streamId, + jobCreatedAt, + }); + const callback = createOAuthCallback({ runStepEmitter, runStepDeltaEmitter }); + const oauthStart = createOAuthStart({ + res, + flowId, + callback, + flowManager, + }); + return await reinitMCPServer({ + user, + signal, + serverName, + configServers, + oauthStart, + flowManager, + userMCPAuthMap, + requestBody, + requestScopedConnections, + forceNew: true, + returnOnOAuth: false, + connectionTimeout: Time.THIRTY_SECONDS, + }); } /** @@ -1090,11 +1049,6 @@ function createToolInstance({ const effectiveUser = config?.configurable?.user ?? capturedUser; const permissionUser = effectiveUser; const userId = effectiveUser?.id || config?.configurable?.user_id || capturedUser?.id; - /** @type {ReturnType} */ - let abortHandler = null; - /** @type {AbortSignal} */ - let derivedSignal = null; - try { const provider = (config?.metadata?.provider || capturedProvider)?.toLowerCase(); const canUseMCP = mcpPermissionContext @@ -1105,7 +1059,7 @@ function createToolInstance({ } const flowsCache = getLogStores(CacheKeys.FLOWS); const flowManager = getFlowStateManager(flowsCache); - derivedSignal = config?.signal ? AbortSignal.any([config.signal]) : undefined; + const derivedSignal = config?.signal ? AbortSignal.any([config.signal]) : undefined; const mcpManager = getMCPManager(userId); const { args: _args, stepId, ...toolCall } = config.toolCall ?? {}; @@ -1130,12 +1084,6 @@ function createToolInstance({ jobCreatedAt, }); - if (derivedSignal) { - const tenantId = config?.configurable?.user?.tenantId ?? getTenantId(); - abortHandler = createAbortHandler({ userId, serverName, toolName, tenantId, flowManager }); - derivedSignal.addEventListener('abort', abortHandler, { once: true }); - } - const customUserVars = config?.configurable?.userMCPAuthMap?.[`${Constants.mcp_prefix}${serverName}`]; @@ -1205,11 +1153,6 @@ function createToolInstance({ throw new Error( `[MCP][${serverName}][${toolName}] tool call failed${error?.message ? `: ${error?.message}` : '.'}`, ); - } finally { - // Clean up abort handler to prevent memory leaks - if (abortHandler && derivedSignal) { - derivedSignal.removeEventListener('abort', abortHandler); - } } }; diff --git a/api/server/services/MCP.spec.js b/api/server/services/MCP.spec.js index 23e7067e76..66bfe592ed 100644 --- a/api/server/services/MCP.spec.js +++ b/api/server/services/MCP.spec.js @@ -1411,7 +1411,7 @@ describe('User parameter passing tests', () => { } }); - it('should fail tenant-scoped OAuth flows when tool loading is aborted', async () => { + it('does not fail shared OAuth flows when tool loading is aborted', async () => { const mockUser = { id: 'tenant-user', name: 'Tenant User' }; const mockRes = { write: jest.fn(), flush: jest.fn() }; const abortController = new AbortController(); @@ -1419,9 +1419,7 @@ describe('User parameter passing tests', () => { createFlowWithHandler: jest.fn(), failFlow: jest.fn(), }; - mockGetTenantId.mockReturnValue('tenant/a'); mockGetFlowStateManager.mockReturnValue(mockFlowManager); - MCPOAuthHandler.generateFlowId.mockReturnValue('tenant-flow-id'); let resolveReinit; mockReinitMCPServer.mockImplementation( @@ -1445,21 +1443,7 @@ describe('User parameter passing tests', () => { resolveReinit({ tools: [], availableTools: {} }); await createToolsPromise; - expect(MCPOAuthHandler.generateFlowId).toHaveBeenCalledWith( - mockUser.id, - 'tenant-abort-server', - 'tenant/a', - ); - expect(mockFlowManager.failFlow).toHaveBeenCalledWith( - 'tenant-flow-id', - 'mcp_oauth', - expect.any(Error), - ); - expect(mockFlowManager.failFlow).toHaveBeenCalledWith( - 'tenant-flow-id', - 'mcp_get_tokens', - expect.any(Error), - ); + expect(mockFlowManager.failFlow).not.toHaveBeenCalled(); }); it('should throw error if user is not provided', async () => { @@ -1487,6 +1471,82 @@ describe('User parameter passing tests', () => { }); describe('createMCPTool', () => { + it('keeps shared OAuth recovery alive when one tool caller aborts', async () => { + const mockUser = { id: 'shared-recovery-user', role: 'USER' }; + const mockRes = { write: jest.fn(), flush: jest.fn() }; + const ownerAbort = new AbortController(); + const waiterAbort = new AbortController(); + const flowManager = { + getFlowState: jest.fn().mockResolvedValue(null), + createFlowWithHandler: jest.fn(), + failFlow: jest.fn(), + }; + let completeRecovery; + const sharedRecovery = new Promise((resolve) => { + completeRecovery = resolve; + }); + const callTool = jest.fn(({ options }) => { + const signal = options?.signal; + return new Promise((resolve, reject) => { + const onAbort = () => { + signal?.removeEventListener('abort', onAbort); + reject(new Error('tool caller aborted')); + }; + signal?.addEventListener('abort', onAbort, { once: true }); + sharedRecovery.then(() => { + signal?.removeEventListener('abort', onAbort); + resolve(['ok', null]); + }); + }); + }); + const { getRoleByName } = require('~/models'); + getRoleByName.mockResolvedValue({ + permissions: { + [PermissionTypes.MCP_SERVERS]: { + [Permissions.USE]: true, + }, + }, + }); + mockGetFlowStateManager.mockReturnValue(flowManager); + mockGetMCPManager.mockReturnValue({ callTool }); + + const mcpTool = await createMCPTool({ + res: mockRes, + user: mockUser, + config: { url: 'https://runtime-oauth.example.com/mcp' }, + toolKey: `test-tool${D}test-server`, + provider: 'openai', + userMCPAuthMap: {}, + availableTools: { + [`test-tool${D}test-server`]: { + function: { + description: 'Cached tool', + parameters: { type: 'object', properties: {} }, + }, + }, + }, + }); + const createConfig = (signal) => ({ + signal, + configurable: { user: mockUser }, + metadata: { provider: 'openai', thread_id: 'thread-1', run_id: 'run-1' }, + toolCall: {}, + }); + + const ownerCall = mcpTool.invoke({}, createConfig(ownerAbort.signal)); + const waiterCall = mcpTool.invoke({}, createConfig(waiterAbort.signal)); + await new Promise((resolve) => setImmediate(resolve)); + + ownerAbort.abort(); + + await expect(ownerCall).rejects.toThrow('Aborted'); + expect(flowManager.failFlow).not.toHaveBeenCalled(); + + completeRecovery(); + await expect(waiterCall).resolves.toBe('ok'); + expect(callTool).toHaveBeenCalledTimes(2); + }); + it.each(['OAuth flow initiated - return early', 'Pending OAuth flow reused - return early'])( 'preserves runtime-detected OAuth for the internal signal: %s', async (oauthSignal) => { diff --git a/packages/api/src/mcp/MCPConnectionFactory.ts b/packages/api/src/mcp/MCPConnectionFactory.ts index 977978d9d3..f31cedf0f3 100644 --- a/packages/api/src/mcp/MCPConnectionFactory.ts +++ b/packages/api/src/mcp/MCPConnectionFactory.ts @@ -23,6 +23,7 @@ import { } from '~/mcp/oauth'; import { sanitizeUrlForLogging, isClientRejectionMessage, isOAuthServer } from './utils'; import { PENDING_STALE_MS, normalizeExpiresAt } from '~/flow/manager'; +import { isOAuthAuthenticationError } from './errors'; import { preProcessGraphTokens } from '~/utils/graph'; import { withTimeout } from '~/utils/promise'; import { MCPConnection } from './connection'; @@ -44,6 +45,8 @@ type OAuthRequiredEvent = { skipSilentRefresh?: boolean; }; +type OAuthRecoveryPhase = 'silent-refresh' | 'interactive' | 'terminal'; + /** * Factory for creating MCP connections with optional OAuth authentication. * Handles OAuth flows, token management, and connection retry logic. @@ -988,16 +991,46 @@ export class MCPConnectionFactory { connection: MCPConnection, eventName: 'oauthRequired' | 'oauthReauthenticationRequired' = 'oauthRequired', ): () => void { - const oauthHandler = async (data: OAuthRequiredEvent) => { + const isRequestRecovery = eventName === 'oauthReauthenticationRequired'; + let recoveryPhase: OAuthRecoveryPhase = 'silent-refresh'; + let eventHandling: Promise | null = null; + + const handleOAuthEvent = async (data: OAuthRequiredEvent) => { logger.info(`${this.logPrefix} oauthRequired event received`); - if (!data.skipSilentRefresh && this.shouldAttemptSilentTokenRefresh(data)) { - const refreshedTokens = await this.attemptSilentTokenRefresh(); - if (refreshedTokens) { - connection.setOAuthTokens(refreshedTokens); - connection.emit('oauthHandled'); + if (this.connectionReady) { + const emitted = connection.emit('oauthReauthenticationRequired', { + ...data, + skipSilentRefresh: data.skipSilentRefresh, + }); + if (emitted) { return; } + logger.info(`${this.logPrefix} Cached connection requires a live OAuth request handler`); + connection.emit('oauthFailed', new Error('OAuth reauthentication required')); + return; + } + + if (isRequestRecovery && recoveryPhase === 'terminal') { + logger.warn(`${this.logPrefix} OAuth recovery phase budget exhausted`); + connection.emit('oauthFailed', new Error('OAuth recovery phase budget exhausted')); + return; + } + + if (!isRequestRecovery || recoveryPhase === 'silent-refresh') { + recoveryPhase = 'interactive'; + if (!data.skipSilentRefresh && this.shouldAttemptSilentTokenRefresh(data)) { + const refreshedTokens = await this.attemptSilentTokenRefresh(); + if (refreshedTokens) { + connection.setOAuthTokens(refreshedTokens); + connection.emit('oauthHandled', 'silent-refresh' satisfies t.OAuthHandledSource); + return; + } + } + } + + if (isRequestRecovery) { + recoveryPhase = 'terminal'; } // Silent refresh failed and we're about to fall through to interactive @@ -1007,21 +1040,6 @@ export class MCPConnectionFactory { // window in `handleOAuthRequired`). await this.invalidateCompletedOAuthFlow(); - if (this.connectionReady) { - const emitted = connection.emit('oauthReauthenticationRequired', { - ...data, - skipSilentRefresh: true, - }); - if (emitted) { - return; - } - logger.info( - `${this.logPrefix} Silent refresh did not recover cached connection; requiring fresh OAuth prompt`, - ); - connection.emit('oauthFailed', new Error('OAuth reauthentication required')); - return; - } - if (this.returnOnOAuth) { try { const config = this.serverConfig; @@ -1098,15 +1116,10 @@ export class MCPConnectionFactory { // Start monitoring in background — createFlow will find the existing PENDING state // written by initFlow above, so metadata arg is unused (pass {} to make that explicit) - this.flowManager!.createFlow(newFlowId, 'mcp_oauth', {}, this.signal).catch( - async (error) => { - logger.debug(`${this.logPrefix} OAuth flow monitor ended`, error); - await this.clearStaleClientIfRejected( - flowMetadata.reusedClientCredentialSetId, - error, - ); - }, - ); + this.flowManager!.createFlow(newFlowId, 'mcp_oauth', {}).catch(async (error) => { + logger.debug(`${this.logPrefix} OAuth flow monitor ended`, error); + await this.clearStaleClientIfRejected(flowMetadata.reusedClientCredentialSetId, error); + }); if (this.oauthStart) { logger.info(`${this.logPrefix} OAuth flow started, issuing authorization URL`); @@ -1187,7 +1200,7 @@ export class MCPConnectionFactory { // Only emit oauthHandled if we actually got tokens (OAuth succeeded) if (result?.tokens) { - connection.emit('oauthHandled'); + connection.emit('oauthHandled', 'interactive' satisfies t.OAuthHandledSource); } else { await this.clearStaleClientIfRejected(result?.reusedClientCredentialSetId, result?.error); logger.warn(`${this.logPrefix} OAuth failed, emitting oauthFailed event`); @@ -1195,6 +1208,23 @@ export class MCPConnectionFactory { } }; + const oauthHandler = (data: OAuthRequiredEvent): Promise => { + if (!isRequestRecovery) { + return handleOAuthEvent(data); + } + if (eventHandling) { + return eventHandling; + } + + const handling = handleOAuthEvent(data).finally(() => { + if (eventHandling === handling) { + eventHandling = null; + } + }); + eventHandling = handling; + return handling; + }; + connection.on(eventName, oauthHandler); return () => { @@ -1279,7 +1309,7 @@ export class MCPConnectionFactory { throw error; } - if (this.useOAuth && this.isOAuthError(error)) { + if (this.useOAuth && isOAuthAuthenticationError(error)) { logger.info(`${this.logPrefix} OAuth required, stopping connection attempts`); throw error; } @@ -1332,48 +1362,6 @@ export class MCPConnectionFactory { return false; } - // Determines if an error indicates OAuth authentication is required - private isOAuthError(error: unknown): boolean { - if (!error || typeof error !== 'object') { - return false; - } - - // Check for error code - if ('code' in error) { - const code = (error as { code?: number }).code; - if (code === 401 || code === 403) { - return true; - } - } - - // Check message for various auth error indicators - if ('message' in error && typeof error.message === 'string') { - const message = error.message.toLowerCase(); - // Check for 401 status - if (message.includes('401') || message.includes('non-200 status code (401)')) { - return true; - } - // Check for invalid_token (OAuth servers return this for expired/revoked tokens) - if (message.includes('invalid_token')) { - return true; - } - // Check for invalid_grant (OAuth servers return this for expired/revoked grants) - if (message.includes('invalid_grant')) { - return true; - } - // Check for authentication required - if (message.includes('authentication required') || message.includes('unauthorized')) { - return true; - } - // Check for missing authorization values (e.g., Amazon Ads MCP returns HTTP 400 with this) - if (message.includes('no authorization')) { - return true; - } - } - - return false; - } - /** Manages OAuth flow initiation and completion */ protected async handleOAuthRequired(): Promise<{ tokens: MCPOAuthTokens | null; @@ -1437,7 +1425,7 @@ export class MCPConnectionFactory { reusedStoredClient = flowMeta?.reusedStoredClient === true; reusedClientCredentialSetId = flowMeta?.reusedClientCredentialSetId; - const tokens = await this.flowManager.createFlow(flowId, 'mcp_oauth', {}, this.signal); + const tokens = await this.waitForSharedOAuthFlow(flowId); if (typeof this.oauthEnd === 'function') { await this.oauthEnd(); } @@ -1542,7 +1530,7 @@ export class MCPConnectionFactory { // createFlow will find the existing PENDING state written by initFlow above, // so metadata arg is unused (pass {} to make that explicit) - const tokens = await this.flowManager.createFlow(newFlowId, 'mcp_oauth', {}, this.signal); + const tokens = await this.waitForSharedOAuthFlow(newFlowId); if (typeof this.oauthEnd === 'function') { await this.oauthEnd(); } @@ -1562,4 +1550,40 @@ export class MCPConnectionFactory { return { tokens: null, reusedStoredClient, reusedClientCredentialSetId, error }; } } + + private waitForSharedOAuthFlow(flowId: string): Promise { + const flow = this.flowManager!.createFlow(flowId, 'mcp_oauth', {}); + const signal = this.signal; + if (!signal) { + return flow; + } + + return new Promise((resolve, reject) => { + const cleanup = () => signal.removeEventListener('abort', onAbort); + const onAbort = () => { + cleanup(); + reject( + signal.reason instanceof Error ? signal.reason : new Error('MCP OAuth flow wait aborted'), + ); + }; + + flow.then( + (tokens) => { + cleanup(); + resolve(tokens); + }, + (error: unknown) => { + cleanup(); + reject(error); + }, + ); + + if (signal.aborted) { + onAbort(); + return; + } + + signal.addEventListener('abort', onAbort, { once: true }); + }); + } } diff --git a/packages/api/src/mcp/MCPManager.ts b/packages/api/src/mcp/MCPManager.ts index 5852bcd1ff..162cc32277 100644 --- a/packages/api/src/mcp/MCPManager.ts +++ b/packages/api/src/mcp/MCPManager.ts @@ -28,9 +28,11 @@ import { UserConnectionManager } from './UserConnectionManager'; import { ConnectionsRepository } from './ConnectionsRepository'; import { MCPConnectionFactory } from './MCPConnectionFactory'; import { processMCPEnv, isPluginSourced } from '~/utils/env'; +import { OAuthLifecycleRelay } from './oauth/pending'; import { preProcessGraphTokens } from '~/utils/graph'; import { formatToolContent } from './parsers'; import { MCPConnection } from './connection'; +import { mcpConfig } from './mcpConfig'; function createOboToolCallErrorMessage( logPrefix: string, @@ -48,12 +50,35 @@ function createOboToolCallErrorMessage( return `${logPrefix} ${error.userMessage} Cannot execute tool ${toolName}. ${failureSuffix}`; } +class OAuthRecoveryTakeoverRequired extends Error {} + +type OAuthReconnectResult = + | { connected: true } + | { + connected: false; + error: unknown; + oauthHandled: boolean; + source?: t.OAuthHandledSource; + }; + +const OAUTH_RECOVERY_RECONNECT_ATTEMPTS = 3; +const OAUTH_RECOVERY_RECONNECT_DELAY_MS = 2000; + /** * Centralized manager for MCP server connections and tool execution. * Extends UserConnectionManager to handle both app-level and user-specific connections. */ export class MCPManager extends UserConnectionManager { private static instance: MCPManager | null; + private readonly oauthRecoveries = new WeakMap< + MCPConnection, + { + promise: Promise; + callbacks: OAuthLifecycleRelay; + allowsTakeover: boolean; + takeoverClaimed?: boolean; + } + >(); /** Creates and initializes the singleton MCPManager instance */ public static async createInstance(configs: t.MCPServers): Promise { @@ -75,6 +100,117 @@ export class MCPManager extends UserConnectionManager { this.appConnections = new ConnectionsRepository(undefined); } + public override async getUserConnection( + opts: t.UserMCPConnectionOptions, + ): Promise { + const userId = opts.user?.id; + if (opts.forceNew || !userId) { + return super.getUserConnection(opts); + } + + const connectionKey = `${userId}:${opts.serverName}`; + const requestConnection = opts.requestScopedConnections?.connections.get(connectionKey) as + | MCPConnection + | undefined; + const connection = requestConnection ?? this.userConnections.get(userId)?.get(opts.serverName); + const recovery = connection ? this.oauthRecoveries.get(connection) : undefined; + const providedConfigIsNewer = + connection != null && + opts.serverConfig?.updatedAt != null && + connection.isStale(opts.serverConfig.updatedAt); + if (recovery && !providedConfigIsNewer) { + if (recovery.callbacks) { + await recovery.callbacks.add({ + oauthStart: opts.oauthStart, + oauthEnd: opts.oauthEnd, + flowManager: opts.flowManager, + userId, + serverName: opts.serverName, + }); + } + await this.waitForActiveRecovery(recovery.promise, opts.signal); + } + + return super.getUserConnection(opts); + } + + /** Runs work against a user connection while preventing recovery from replacing its SDK client. */ + public async withUserConnectionLease( + opts: t.UserMCPConnectionOptions, + operation: (connection: MCPConnection) => Promise, + ): Promise { + while (true) { + const connection = await this.getUserConnection(opts); + this.retainConnection(connection); + const recovery = this.oauthRecoveries.get(connection)?.promise; + if (recovery) { + await this.releaseConnection(connection); + await this.waitForActiveRecovery(recovery, opts.signal); + continue; + } + + try { + return await operation(connection); + } finally { + await this.releaseConnection(connection); + } + } + } + + private waitForActiveRecovery(recovery: Promise, signal?: AbortSignal): Promise { + if (!signal) { + return recovery; + } + + return new Promise((resolve, reject) => { + const onRecoveryResolved = () => { + signal.removeEventListener('abort', onAbort); + resolve(); + }; + const onRecoveryRejected = (error: unknown) => { + signal.removeEventListener('abort', onAbort); + reject(error); + }; + const onAbort = () => { + signal.removeEventListener('abort', onAbort); + const reason = signal.reason; + reject(reason instanceof Error ? reason : new Error('OAuth recovery wait aborted')); + }; + + if (signal.aborted) { + onAbort(); + return; + } + + signal.addEventListener('abort', onAbort, { once: true }); + recovery.then(onRecoveryResolved, onRecoveryRejected); + }); + } + + protected override getActiveConnectionRecovery( + connection: MCPConnection, + ): Promise | undefined { + return this.oauthRecoveries.get(connection)?.promise; + } + + protected override waitForConnectionRecovery( + recovery: Promise, + signal?: AbortSignal, + ): Promise { + return this.waitForActiveRecovery(recovery, signal); + } + + private claimRecoveryTakeover(recovery: { + allowsTakeover: boolean; + takeoverClaimed?: boolean; + }): boolean { + if (!recovery.allowsTakeover || recovery.takeoverClaimed) { + return false; + } + recovery.takeoverClaimed = true; + return true; + } + /** Retrieves an app-level or user-specific connection based on provided arguments */ public async getConnection( args: { @@ -302,53 +438,64 @@ export class MCPManager extends UserConnectionManager { }; } - const userConnections = this.getUserConnections(userId); - if (!userConnections || userConnections.size === 0) { - return { tools: null }; - } - if (!userConnections.has(serverName)) { - return { tools: null }; - } + let awaitedRecovery: Promise | undefined; + while (true) { + const userConnections = this.getUserConnections(userId); + const connection = userConnections?.get(serverName); + if (!connection) { + return { tools: null }; + } - const connection = userConnections.get(serverName)!; - if (effectiveConfig == null) { - await this.disconnectUserConnection(userId, serverName); - return { tools: null }; + if (effectiveConfig == null) { + await this.disconnectUserConnection(userId, serverName); + return { tools: null }; + } + const connectionConfigGeneration = this.getToolConfigGeneration(connection); + const effectiveConfigGeneration = getMCPAppToolsPublicationGeneration(effectiveConfig); + if ( + connectionConfigGeneration != null && + effectiveConfigGeneration != null && + connectionConfigGeneration !== effectiveConfigGeneration + ) { + await this.disconnectUserConnection(userId, serverName); + return { tools: null }; + } + const publicationGeneration = this.getToolPublicationGeneration(connection); + const currentGeneration = await getMCPToolsChangedGeneration({ userId, serverName }); + if ( + publicationGeneration != null && + currentGeneration != null && + publicationGeneration !== currentGeneration + ) { + await this.disconnectUserConnection(userId, serverName); + return { tools: null }; + } + + this.retainConnection(connection); + const recovery = this.oauthRecoveries.get(connection)?.promise; + if (recovery && recovery !== awaitedRecovery) { + awaitedRecovery = recovery; + await this.releaseConnection(connection); + await this.waitForConnectionRecovery(recovery); + continue; + } + + try { + const tools = await MCPServerInspector.getToolFunctions(serverName, connection); + const generationAfterFetch = await getMCPToolsChangedGeneration({ userId, serverName }); + if ( + publicationGeneration != null && + generationAfterFetch != null && + publicationGeneration !== generationAfterFetch + ) { + await this.disconnectUserConnection(userId, serverName); + return { tools: null }; + } + return { tools, publicationGeneration }; + } finally { + await this.releaseConnection(connection); + } } - const connectionConfigGeneration = this.getToolConfigGeneration(connection); - const effectiveConfigGeneration = getMCPAppToolsPublicationGeneration(effectiveConfig); - if ( - connectionConfigGeneration != null && - effectiveConfigGeneration != null && - connectionConfigGeneration !== effectiveConfigGeneration - ) { - await this.disconnectUserConnection(userId, serverName); - return { tools: null }; - } - const publicationGeneration = this.getToolPublicationGeneration(connection); - const currentGeneration = await getMCPToolsChangedGeneration({ userId, serverName }); - if ( - publicationGeneration != null && - currentGeneration != null && - publicationGeneration !== currentGeneration - ) { - await this.disconnectUserConnection(userId, serverName); - return { tools: null }; - } - const tools = await MCPServerInspector.getToolFunctions(serverName, connection); - const generationAfterFetch = await getMCPToolsChangedGeneration({ userId, serverName }); - if ( - publicationGeneration != null && - generationAfterFetch != null && - publicationGeneration !== generationAfterFetch - ) { - await this.disconnectUserConnection(userId, serverName); - return { tools: null }; - } - return { - tools, - publicationGeneration, - }; } catch (error) { logger.warn( `[getServerToolFunctions] Error getting tool functions for server ${serverName}`, @@ -422,6 +569,198 @@ ${formattedInstructions} Please follow these instructions when using tools from the respective MCP servers.`; } + private async recoverOAuthConnection( + connection: MCPConnection, + error: unknown, + serverName: string, + userId: string, + attachSharedOAuthHandler: (relay: OAuthLifecycleRelay) => () => void, + oauthStart: t.OAuthStartHandler | undefined, + oauthEnd: (() => Promise) | undefined, + flowManager: FlowStateManager, + signal?: AbortSignal, + allowsTakeover = true, + ): Promise { + const existingRecovery = this.oauthRecoveries.get(connection); + if (existingRecovery) { + if (existingRecovery.callbacks) { + await existingRecovery.callbacks.add({ + oauthStart, + oauthEnd, + flowManager, + userId, + serverName, + }); + } + try { + return await this.waitForActiveRecovery(existingRecovery.promise, signal); + } catch (recoveryError) { + if (signal?.aborted) { + throw recoveryError; + } + if (!allowsTakeover || !this.claimRecoveryTakeover(existingRecovery)) { + throw recoveryError; + } + if (this.oauthRecoveries.get(connection) === existingRecovery) { + this.oauthRecoveries.delete(connection); + } + throw new OAuthRecoveryTakeoverRequired(); + } + } + + const callbacks = new OAuthLifecycleRelay({ + oauthStart, + oauthEnd, + logPrefix: `[MCP][User: ${userId}][${serverName}]`, + }); + const recovery = Promise.resolve().then(async () => { + const cleanupRequestOAuthHandler = attachSharedOAuthHandler(callbacks); + try { + await this.waitForOAuthRecovery(connection, () => + connection.emit('oauthReauthenticationRequired', { + serverName, + error, + serverUrl: connection.url, + userId, + }), + ); + await this.connectAfterOAuthRecovery(connection, async (connectError) => { + await this.waitForOAuthRecovery(connection, () => + connection.emit('oauthReauthenticationRequired', { + serverName, + error: connectError, + serverUrl: connection.url, + userId, + skipSilentRefresh: true, + }), + ); + }); + } finally { + cleanupRequestOAuthHandler(); + } + }); + + const recoveryEntry = { promise: recovery, callbacks, allowsTakeover, takeoverClaimed: false }; + this.oauthRecoveries.set(connection, recoveryEntry); + this.holdDeferredConnectionDisposal(connection); + const clearRecovery = () => { + if (this.oauthRecoveries.get(connection) === recoveryEntry) { + this.oauthRecoveries.delete(connection); + } + }; + const releaseRecoveryDisposal = () => this.releaseDeferredConnectionDisposal(connection); + void recovery.then(clearRecovery, clearRecovery); + void recovery.then(releaseRecoveryDisposal, releaseRecoveryDisposal); + await this.waitForActiveRecovery(recovery, signal); + } + + private async connectAfterOAuthRecovery( + connection: MCPConnection, + requestInteractiveRecovery: (error: unknown) => Promise, + ): Promise { + await this.waitForConnectionBorrowersToDrain(connection); + const firstAttempt = await this.connectWithTransientRetries(connection); + if (firstAttempt.connected) { + return; + } + if (!firstAttempt.oauthHandled) { + throw firstAttempt.error; + } + if (firstAttempt.source === 'silent-refresh') { + await requestInteractiveRecovery(firstAttempt.error); + } + + const secondAttempt = await this.connectWithTransientRetries(connection); + if (!secondAttempt.connected) { + throw secondAttempt.error; + } + } + + private async connectWithTransientRetries( + connection: MCPConnection, + ): Promise { + let result: OAuthReconnectResult | undefined; + for (let attempt = 1; attempt <= OAUTH_RECOVERY_RECONNECT_ATTEMPTS; attempt++) { + result = await this.connectOnceAfterOAuth(connection); + if ( + result.connected || + result.oauthHandled || + connection.isOAuthAuthenticationError(result.error) || + attempt === OAUTH_RECOVERY_RECONNECT_ATTEMPTS + ) { + return result; + } + await this.waitForOAuthReconnectRetry(attempt); + } + return result!; + } + + private waitForOAuthReconnectRetry(attempt: number): Promise { + return new Promise((resolve) => { + setTimeout(resolve, OAUTH_RECOVERY_RECONNECT_DELAY_MS * attempt); + }); + } + + private async connectOnceAfterOAuth(connection: MCPConnection): Promise { + let oauthHandled = false; + let source: t.OAuthHandledSource | undefined; + const handleOAuth = (handledSource?: t.OAuthHandledSource) => { + oauthHandled = true; + source = handledSource; + }; + connection.on('oauthHandled', handleOAuth); + try { + await connection.connect(); + return { connected: true }; + } catch (error) { + return { connected: false, error, oauthHandled, source }; + } finally { + connection.off('oauthHandled', handleOAuth); + } + } + + private waitForOAuthRecovery( + connection: MCPConnection, + requestRecovery: () => boolean, + ): Promise { + return new Promise((resolve, reject) => { + const cleanup = () => { + clearTimeout(timeout); + connection.off('oauthHandled', handleSuccess); + connection.off('oauthFailed', handleFailure); + }; + const handleSuccess = () => { + cleanup(); + resolve(); + }; + const handleFailure = (oauthError: Error) => { + cleanup(); + reject(oauthError); + }; + const timeout = setTimeout(() => { + cleanup(); + reject(new Error(`OAuth recovery timeout after ${mcpConfig.OAUTH_HANDLING_TIMEOUT}ms`)); + }, mcpConfig.OAUTH_HANDLING_TIMEOUT); + + connection.once('oauthHandled', handleSuccess); + connection.once('oauthFailed', handleFailure); + + let recoveryRequested: boolean; + try { + recoveryRequested = requestRecovery(); + } catch (error) { + cleanup(); + reject(error); + return; + } + if (recoveryRequested) { + return; + } + cleanup(); + reject(new Error('OAuth recovery requested without an active request handler')); + }); + } + /** * Calls a tool on an MCP server, using either a user-specific connection * (if userId is provided) or an app-level connection. Updates the last activity timestamp @@ -469,184 +808,351 @@ Please follow these instructions when using tools from the respective MCP server oboTokenResolver?: OboTokenResolver; oboTrustChecker?: OboTrustChecker; }): Promise { - /** User-specific connection */ - let connection: MCPConnection | undefined; - let cleanupRequestOAuthHandler: (() => void) | undefined; - let disconnectAfterCall = false; const userId = user?.id; const logPrefix = userId ? `[MCP][User: ${userId}][${serverName}]` : `[MCP][${serverName}]`; - - try { - connection = await this.getConnection({ - serverName, - user, - flowManager, - tokenMethods, - oauthStart, - oauthEnd, - oboTokenResolver, - oboTrustChecker, - graphTokenResolver, - signal: options?.signal, - customUserVars, - requestBody, - requestScopedConnections, - serverConfig: providedConfig, - }); - - if (!(await connection.isConnected())) { - /** May happen if getUserConnection failed silently or app connection dropped */ - throw new McpError( - ErrorCode.InternalError, // Use InternalError for connection issues - `${logPrefix} Connection is not active. Cannot execute tool ${toolName}.`, - ); - } - - const registry = MCPServersRegistry.getInstance(); - const rawConfig = providedConfig ?? (await registry.getServerConfig(serverName, userId)); - if (!rawConfig) { - throw new McpError( - ErrorCode.InvalidRequest, - `${logPrefix} Configuration for server "${serverName}" not found.`, - ); - } - const isDbSourced = isUserSourced(rawConfig); - const ephemeralConnection = !!userId && requiresEphemeralUserConnection(rawConfig); - disconnectAfterCall = ephemeralConnection && !requestScopedConnections; - - /** - * Pre-process Graph token placeholders (async) before the synchronous processMCPEnv pass. - * Plugin-sourced configs are excluded for the same reason processMCPEnv excludes them: - * a placeholder a plugin authored must never resolve against the user's Graph token. - */ - const graphProcessedConfig = - isDbSourced || isPluginSourced(rawConfig) - ? (rawConfig as t.MCPOptions) - : await preProcessGraphTokens(rawConfig as t.MCPOptions, { - user, - graphTokenResolver, - scopes: process.env.GRAPH_API_SCOPES, - }); - const currentOptions = processMCPEnv({ - user, - body: requestBody, - dbSourced: isDbSourced, - options: graphProcessedConfig, - customUserVars, - }); - - const resolvedHeaders: Record = - 'headers' in currentOptions ? { ...(currentOptions.headers || {}) } : {}; - - /** Refresh OBO token on each tool call to ensure it's current */ - const oboConfig = rawConfig.obo; - if (oboConfig && oboTokenResolver && user) { - const oboTrusted = oboTrustChecker - ? await oboTrustChecker({ - source: rawConfig.source, - author: rawConfig.author, - dbId: rawConfig.dbId, - }) - : true; - if (!oboTrusted) { - logger.warn( - `${logPrefix} OBO config not trusted (author lacks ${PermissionTypes.MCP_SERVERS}.${Permissions.CONFIGURE_OBO}); refusing to mint a downstream token`, - ); - throw new McpError( - ErrorCode.InternalError, - `${logPrefix} OBO is not permitted for server "${serverName}". The user who configured it no longer has permission to use OBO.`, - ); + this.bindRequestScopedConnectionStore(requestScopedConnections); + let recoveryTakeoverConsumed = false; + while (true) { + /** User-specific connection */ + let connection: MCPConnection | undefined; + let connectionRetained = false; + let deferredDisposalHeld = false; + let attachSharedOAuthHandler: ((relay: OAuthLifecycleRelay) => () => void) | undefined; + let disposeAfterCall = false; + const retainConnectionLease = () => { + if (!connection || connectionRetained) { + return; } - let oboTokens: MCPOAuthTokens; + this.retainConnection(connection); + connectionRetained = true; + }; + const releaseConnectionLease = async (preserveDisposalHold = false) => { + if (!connection || !connectionRetained) { + return; + } + if (deferredDisposalHeld && !preserveDisposalHold) { + await this.releaseDeferredConnectionDisposal(connection); + deferredDisposalHeld = false; + } + connectionRetained = false; + await this.releaseConnection(connection); + }; + const waitForRecoveryWithoutLease = async (startRecovery: () => Promise) => { + const recovery = startRecovery(); + // Keep an eviction marker across the temporary lease gap and transfer that + // responsibility back to this caller after recovery. An unrelated final + // borrower may disconnect the old client, but cannot consume the marker. + if (!deferredDisposalHeld) { + this.holdDeferredConnectionDisposal(connection!); + deferredDisposalHeld = true; + } + await releaseConnectionLease(true); try { - oboTokens = await resolveOboToken(user, oboConfig, oboTokenResolver); - } catch (error) { - if (error instanceof OboTokenResolutionError) { - throw new McpError( - ErrorCode.InternalError, - createOboToolCallErrorMessage(logPrefix, toolName, error), - ); - } - throw error; + await recovery; + } finally { + retainConnectionLease(); } + }; - if (!oboTokens.access_token) { - throw new McpError( - ErrorCode.InternalError, - `${logPrefix} OBO token refresh failed. Cannot execute tool ${toolName}. Re-authenticate the user and retry.`, - ); - } - resolvedHeaders['Authorization'] = `Bearer ${oboTokens.access_token}`; - } - if (userId && user && oauthStart && flowManager && isOAuthServer(currentOptions)) { - const { allowedDomains, allowedAddresses, useSSRFProtection } = - await registry.resolveAllowlists({ userId, role: user?.role }); - cleanupRequestOAuthHandler = MCPConnectionFactory.attachRequestOAuthHandler( - { + try { + let awaitedCheckoutRecovery: Promise | undefined; + while (true) { + connection = await this.getConnection({ serverName, - serverConfig: currentOptions, - dbSourced: isDbSourced, - skipEnvProcessing: true, - useSSRFProtection, - allowedDomains, - allowedAddresses, - }, - { - useOAuth: true, user, flowManager, tokenMethods, - signal: options?.signal, oauthStart, oauthEnd, + oboTokenResolver, + oboTrustChecker, + graphTokenResolver, + signal: options?.signal, customUserVars, requestBody, - }, - connection, - ); - } - - connection.setRequestHeaders(resolvedHeaders); - - const result = await connection.client.request( - { - method: 'tools/call', - params: { - name: toolName, - arguments: toolArguments, - }, - }, - CallToolResultSchema, - { - timeout: connection.timeout, - resetTimeoutOnProgress: true, - ...options, - }, - ); - const hasPersistentUserConnections = - !!userId && (this.userConnections.get(userId)?.size ?? 0) > 0; - if (!ephemeralConnection && hasPersistentUserConnections) { - await this.updateUserLastActivity(userId); - } - this.checkIdleConnections(); - return formatToolContent(result as t.MCPToolCallResponse, provider); - } catch (error) { - // Log with context and re-throw or handle as needed - logger.error(`${logPrefix}[${toolName}] Tool call failed`, error); - // Rethrowing allows the caller (createMCPTool) to handle the final user message - throw error; - } finally { - cleanupRequestOAuthHandler?.(); - // Ephemeral connections are never stored in userConnections, so disconnecting - // is the only cleanup needed; removing the map entry here could orphan a - // still-connected cached connection from before a config change. - if (disconnectAfterCall && connection) { - try { - await connection.disconnect(); - } catch (disconnectError) { - logger.warn(`${logPrefix}[${toolName}] Failed to disconnect ephemeral connection`, { - error: disconnectError, + requestScopedConnections, + serverConfig: providedConfig, }); + retainConnectionLease(); + const checkoutRecovery = this.oauthRecoveries.get(connection); + if (!checkoutRecovery || checkoutRecovery.promise === awaitedCheckoutRecovery) { + break; + } + if (checkoutRecovery.callbacks) { + await checkoutRecovery.callbacks.add({ + oauthStart, + oauthEnd, + flowManager, + userId: userId!, + serverName, + }); + } + awaitedCheckoutRecovery = checkoutRecovery.promise; + await releaseConnectionLease(); + try { + await this.waitForConnectionRecovery(checkoutRecovery.promise, options?.signal); + } catch (recoveryError) { + if ( + options?.signal?.aborted || + recoveryTakeoverConsumed || + !this.claimRecoveryTakeover(checkoutRecovery) + ) { + throw recoveryError; + } + recoveryTakeoverConsumed = true; + if (this.oauthRecoveries.get(connection) === checkoutRecovery) { + this.oauthRecoveries.delete(connection); + } + continue; + } + } + + const connectionIsActive = await connection.isConnected(); + const connectionCheckError = connectionIsActive + ? undefined + : connection.getLastConnectionCheckError(); + + if ( + !connectionIsActive && + (!userId || !connection.isOAuthAuthenticationError(connectionCheckError)) + ) { + /** May happen if getUserConnection failed silently or app connection dropped */ + throw new McpError( + ErrorCode.InternalError, + `${logPrefix} Connection is not active. Cannot execute tool ${toolName}.`, + ); + } + + const registry = MCPServersRegistry.getInstance(); + const rawConfig = providedConfig ?? (await registry.getServerConfig(serverName, userId)); + if (!rawConfig) { + throw new McpError( + ErrorCode.InvalidRequest, + `${logPrefix} Configuration for server "${serverName}" not found.`, + ); + } + const isDbSourced = isUserSourced(rawConfig); + const ephemeralConnection = !!userId && requiresEphemeralUserConnection(rawConfig); + disposeAfterCall = ephemeralConnection && !requestScopedConnections; + + /** Plugin-authored placeholders must not resolve against the user's Graph token. */ + const graphProcessedConfig = + isDbSourced || isPluginSourced(rawConfig) + ? (rawConfig as t.MCPOptions) + : await preProcessGraphTokens(rawConfig as t.MCPOptions, { + user, + graphTokenResolver, + scopes: process.env.GRAPH_API_SCOPES, + }); + const currentOptions = processMCPEnv({ + user, + body: requestBody, + dbSourced: isDbSourced, + options: graphProcessedConfig, + customUserVars, + }); + + const resolvedHeaders: Record = + 'headers' in currentOptions ? { ...(currentOptions.headers || {}) } : {}; + + /** Refresh OBO token on each tool call to ensure it's current */ + const oboConfig = rawConfig.obo; + if (oboConfig && oboTokenResolver && user) { + const oboTrusted = oboTrustChecker + ? await oboTrustChecker({ + source: rawConfig.source, + author: rawConfig.author, + dbId: rawConfig.dbId, + }) + : true; + if (!oboTrusted) { + logger.warn( + `${logPrefix} OBO config not trusted (author lacks ${PermissionTypes.MCP_SERVERS}.${Permissions.CONFIGURE_OBO}); refusing to mint a downstream token`, + ); + throw new McpError( + ErrorCode.InternalError, + `${logPrefix} OBO is not permitted for server "${serverName}". The user who configured it no longer has permission to use OBO.`, + ); + } + let oboTokens: MCPOAuthTokens; + try { + oboTokens = await resolveOboToken(user, oboConfig, oboTokenResolver); + } catch (error) { + if (error instanceof OboTokenResolutionError) { + throw new McpError( + ErrorCode.InternalError, + createOboToolCallErrorMessage(logPrefix, toolName, error), + ); + } + throw error; + } + + if (!oboTokens.access_token) { + throw new McpError( + ErrorCode.InternalError, + `${logPrefix} OBO token refresh failed. Cannot execute tool ${toolName}. Re-authenticate the user and retry.`, + ); + } + resolvedHeaders['Authorization'] = `Bearer ${oboTokens.access_token}`; + } + if ( + userId && + user && + oauthStart && + flowManager && + (isOAuthServer(currentOptions) || connection.usesOAuth()) + ) { + const { allowedDomains, allowedAddresses, useSSRFProtection } = + await registry.resolveAllowlists({ userId, role: user?.role }); + attachSharedOAuthHandler = (relay) => + MCPConnectionFactory.attachRequestOAuthHandler( + { + serverName, + serverConfig: currentOptions, + dbSourced: isDbSourced, + skipEnvProcessing: true, + useSSRFProtection, + allowedDomains, + allowedAddresses, + }, + { + useOAuth: true, + user, + flowManager, + tokenMethods, + oauthStart: relay.start, + oauthEnd: relay.end, + customUserVars, + requestBody, + }, + connection!, + ); + } + + connection.setRequestHeaders(resolvedHeaders); + + if (!connectionIsActive) { + const requestOAuthHandler = attachSharedOAuthHandler; + if (!requestOAuthHandler || !userId) { + throw new McpError( + ErrorCode.InternalError, + `${logPrefix} Connection is not active. Cannot execute tool ${toolName}.`, + ); + } + + try { + await waitForRecoveryWithoutLease(() => + this.recoverOAuthConnection( + connection!, + connectionCheckError, + serverName, + userId, + requestOAuthHandler, + oauthStart, + oauthEnd, + flowManager, + options?.signal, + !recoveryTakeoverConsumed, + ), + ); + } catch (recoveryError) { + if (recoveryError instanceof OAuthRecoveryTakeoverRequired) { + throw recoveryError; + } + if (options?.signal?.aborted) { + throw recoveryError; + } + logger.warn( + `${logPrefix}[${toolName}] Connection-check OAuth recovery failed`, + recoveryError, + ); + throw connectionCheckError; + } + } + + const requestTool = () => + connection!.client.request( + { + method: 'tools/call', + params: { + name: toolName, + arguments: toolArguments, + }, + }, + CallToolResultSchema, + { + timeout: connection!.timeout, + resetTimeoutOnProgress: true, + ...options, + }, + ); + + let result: Awaited>; + try { + result = await requestTool(); + } catch (error) { + const requestOAuthHandler = attachSharedOAuthHandler; + if (!requestOAuthHandler || !userId) { + throw error; + } + + if (!connection.isOAuthAuthenticationError(error)) { + throw error; + } + + try { + await waitForRecoveryWithoutLease(() => + this.recoverOAuthConnection( + connection!, + error, + serverName, + userId, + requestOAuthHandler, + oauthStart, + oauthEnd, + flowManager, + options?.signal, + !recoveryTakeoverConsumed, + ), + ); + } catch (recoveryError) { + if (recoveryError instanceof OAuthRecoveryTakeoverRequired) { + throw recoveryError; + } + if (options?.signal?.aborted) { + throw recoveryError; + } + logger.warn(`${logPrefix}[${toolName}] Runtime OAuth recovery failed`, recoveryError); + throw error; + } + result = await requestTool(); + } + const hasPersistentUserConnections = + !!userId && (this.userConnections.get(userId)?.size ?? 0) > 0; + if (!ephemeralConnection && hasPersistentUserConnections) { + await this.updateUserLastActivity(userId); + } + this.checkIdleConnections(); + return formatToolContent(result as t.MCPToolCallResponse, provider); + } catch (error) { + if (error instanceof OAuthRecoveryTakeoverRequired) { + recoveryTakeoverConsumed = true; + continue; + } + // Log with context and re-throw or handle as needed + logger.error(`${logPrefix}[${toolName}] Tool call failed`, error); + // Rethrowing allows the caller (createMCPTool) to handle the final user message + throw error; + } finally { + await releaseConnectionLease(); + // Ephemeral connections are never stored in userConnections, so disposing + // is the only cleanup needed; removing the map entry here could orphan a + // still-connected cached connection from before a config change. + if (disposeAfterCall && connection) { + await this.disposeEvictedConnection( + connection, + `${logPrefix}[${toolName}] Ephemeral connection`, + ); } } } diff --git a/packages/api/src/mcp/UserConnectionManager.ts b/packages/api/src/mcp/UserConnectionManager.ts index afbc2c4b09..c08fd3ebe7 100644 --- a/packages/api/src/mcp/UserConnectionManager.ts +++ b/packages/api/src/mcp/UserConnectionManager.ts @@ -1,7 +1,5 @@ -import { logger, getTenantId } from '@librechat/data-schemas'; +import { logger } from '@librechat/data-schemas'; import { ErrorCode, McpError } from '@modelcontextprotocol/sdk/types.js'; -import type { MCPOAuthFlowMetadata } from '~/mcp/oauth'; -import type { FlowState } from '~/flow/types'; import type * as t from './types'; import { cancelMCPToolsChanged, @@ -18,31 +16,19 @@ import { requiresOAuthMachinery, } from './utils'; import { MCPServersRegistry } from '~/mcp/registry/MCPServersRegistry'; -import { detectOAuthRequirement, MCPOAuthHandler } from '~/mcp/oauth'; import { ConnectionsRepository } from '~/mcp/ConnectionsRepository'; import { MCPConnectionFactory } from '~/mcp/MCPConnectionFactory'; import { processMCPEnv, isPluginSourced } from '~/utils/env'; +import { OAuthLifecycleRelay } from '~/mcp/oauth/pending'; import { preProcessGraphTokens } from '~/utils/graph'; +import { detectOAuthRequirement } from '~/mcp/oauth'; import { isMCPDomainAllowed } from '~/auth/domain'; -import { PENDING_STALE_MS } from '~/flow/manager'; import { MCPConnection } from './connection'; import { mcpConfig } from './mcpConfig'; -type PendingOAuthStart = { - authURL: string; - options?: t.OAuthStartOptions; -}; - -type PendingOAuthState = { - oauthStarts: Set; - emittedAuthUrls: WeakMap; - primaryOAuthStart?: t.OAuthStartHandler; - lastOAuthStart?: PendingOAuthStart; -}; - type PendingConnection = { promise: Promise; - oauth: PendingOAuthState; + oauth: OAuthLifecycleRelay; }; type ConnectionCreationGuard = { cancelled: boolean }; @@ -63,6 +49,10 @@ export abstract class UserConnectionManager { protected userLastActivity: Map = new Map(); /** In-flight connection promises keyed by `userId:serverName` — coalesces concurrent attempts */ protected pendingConnections: Map = new Map(); + private readonly connectionBorrowers = new WeakMap(); + private readonly connectionBorrowerDrainWaiters = new WeakMap void>>(); + private readonly deferredConnectionDisposalHolds = new WeakMap(); + private readonly deferredConnectionDisposals = new WeakMap(); /** All durable creations, including forced replacements, visible to mutation teardown. */ private readonly activeConnectionCreations: Map> = new Map(); /** Serializes explicit durable replacements without coalescing their callers. */ @@ -233,24 +223,35 @@ export abstract class UserConnectionManager { ? opts.requestScopedConnections : undefined; if (requestScopedConnections) { + this.bindRequestScopedConnectionStore(requestScopedConnections); const requestConnectionKey = `${userId}:${serverName}`; const existing = requestScopedConnections.connections.get(requestConnectionKey) as | MCPConnection | undefined; if (existing) { if (!config || (config.updatedAt && existing.isStale(config.updatedAt))) { - await existing.disconnect().catch((error) => { - logger.warn( - `[MCP][User: ${userId}][${serverName}] Failed to disconnect stale request-scoped connection`, - error, - ); - }); requestScopedConnections.connections.delete(requestConnectionKey); - } else if (await existing.isConnected()) { - logger.debug(`[MCP][User: ${userId}][${serverName}] Reusing request-scoped connection`); - return existing; + await this.disposeEvictedConnection(existing, `[MCP][User: ${userId}][${serverName}]`); } else { + const activeRecovery = this.getActiveConnectionRecovery(existing); + let awaitedRecovery = activeRecovery; + if (activeRecovery) { + await this.waitForConnectionRecovery(activeRecovery, opts.signal); + } + let connected = await existing.isConnected(); + let recovery = this.getActiveConnectionRecovery(existing); + while (recovery && recovery !== awaitedRecovery) { + awaitedRecovery = recovery; + await this.waitForConnectionRecovery(recovery, opts.signal); + connected = await existing.isConnected(); + recovery = this.getActiveConnectionRecovery(existing); + } + if (connected) { + logger.debug(`[MCP][User: ${userId}][${serverName}] Reusing request-scoped connection`); + return existing; + } requestScopedConnections.connections.delete(requestConnectionKey); + await this.disposeEvictedConnection(existing, `[MCP][User: ${userId}][${serverName}]`); } } @@ -264,14 +265,19 @@ export abstract class UserConnectionManager { return pending; } - const pendingOAuth = this.createPendingOAuthState(opts.oauthStart); + const pendingOAuth = new OAuthLifecycleRelay({ + oauthStart: opts.oauthStart, + oauthEnd: opts.oauthEnd, + logPrefix: `[MCP][User: ${userId}][${serverName}]`, + }); const connectionPromise = this.createUserConnectionInternal( { ...opts, forceNew: true, ephemeralConnection: true, serverConfig: config, - oauthStart: this.createPendingOAuthStart(serverName, userId, pendingOAuth), + oauthStart: pendingOAuth.start, + oauthEnd: pendingOAuth.end, }, userId, forceNew === true, @@ -303,12 +309,22 @@ export abstract class UserConnectionManager { const pending = this.pendingConnections.get(lockKey); if (pending) { logger.debug(`[MCP][User: ${userId}][${serverName}] Joining in-flight connection attempt`); - await this.addPendingOAuthStart(pending.oauth, opts, userId); + await pending.oauth.add({ + oauthStart: opts.oauthStart, + oauthEnd: opts.oauthEnd, + flowManager: opts.flowManager, + userId, + serverName, + }); return pending.promise; } } - const pendingOAuth = this.createPendingOAuthState(opts.oauthStart); + const pendingOAuth = new OAuthLifecycleRelay({ + oauthStart: opts.oauthStart, + oauthEnd: opts.oauthEnd, + logPrefix: `[MCP][User: ${userId}][${serverName}]`, + }); const creationGuard: ConnectionCreationGuard = { cancelled: false }; this.registerConnectionCreation(lockKey, creationGuard); const createConnection = () => @@ -318,7 +334,8 @@ export abstract class UserConnectionManager { forceNew: forceNewConnection, ephemeralConnection, serverConfig: config, - oauthStart: this.createPendingOAuthStart(serverName, userId, pendingOAuth), + oauthStart: pendingOAuth.start, + oauthEnd: pendingOAuth.end, }, userId, clearCooldown, @@ -348,170 +365,6 @@ export abstract class UserConnectionManager { } } - private createPendingOAuthState(oauthStart?: t.OAuthStartHandler): PendingOAuthState { - return { - oauthStarts: oauthStart ? new Set([oauthStart]) : new Set(), - emittedAuthUrls: new WeakMap(), - primaryOAuthStart: oauthStart, - }; - } - - private createPendingOAuthStart( - serverName: string, - userId: string, - pendingOAuth: PendingOAuthState, - ): t.OAuthStartHandler { - return async (authURL, options) => { - pendingOAuth.lastOAuthStart = { authURL, options }; - - let primaryError: unknown; - const oauthStarts = Array.from(pendingOAuth.oauthStarts); - for (const oauthStart of oauthStarts) { - try { - await this.emitPendingOAuthStart(pendingOAuth, oauthStart, authURL, options); - } catch (error) { - if (oauthStart === pendingOAuth.primaryOAuthStart) { - primaryError = error; - } else { - logger.warn( - `[MCP][User: ${userId}][${serverName}] Failed to notify joined OAuth listener`, - error, - ); - } - } - } - - if (primaryError) { - throw primaryError; - } - }; - } - - private async addPendingOAuthStart( - pendingOAuth: PendingOAuthState, - opts: t.UserMCPConnectionOptions, - userId: string, - ): Promise { - const { oauthStart, serverName } = opts; - if (typeof oauthStart !== 'function') { - return; - } - - pendingOAuth.oauthStarts.add(oauthStart); - const lastOAuthStart = pendingOAuth.lastOAuthStart; - if (lastOAuthStart) { - try { - const pendingOAuthStart = - lastOAuthStart.options?.expiresAt == null - ? await this.getFlowPendingOAuthStart(opts, userId) - : undefined; - const replayOAuthStart = - pendingOAuthStart?.authURL === lastOAuthStart.authURL - ? pendingOAuthStart - : lastOAuthStart; - await this.emitPendingOAuthStart( - pendingOAuth, - oauthStart, - replayOAuthStart.authURL, - replayOAuthStart.options, - ); - } catch (error) { - logger.warn( - `[MCP][User: ${userId}][${serverName}] Failed to re-issue pending OAuth URL`, - error, - ); - } - return; - } - - await this.reissuePendingOAuthStart(opts, userId, pendingOAuth); - } - - private async emitPendingOAuthStart( - pendingOAuth: PendingOAuthState, - oauthStart: t.OAuthStartHandler, - authURL: string, - options?: t.OAuthStartOptions, - ): Promise { - if (pendingOAuth.emittedAuthUrls.get(oauthStart) === authURL) { - return; - } - pendingOAuth.emittedAuthUrls.set(oauthStart, authURL); - await oauthStart(authURL, options); - } - - private getPendingOAuthStart(flow: FlowState | null | undefined): PendingOAuthStart | undefined { - if (flow?.status !== 'PENDING') { - return undefined; - } - - const expiresAt = flow.createdAt + PENDING_STALE_MS; - if (expiresAt <= Date.now()) { - return undefined; - } - - const metadata = flow.metadata as MCPOAuthFlowMetadata | undefined; - const authorizationUrl = metadata?.authorizationUrl; - if (!authorizationUrl) { - return undefined; - } - - return { authURL: authorizationUrl, options: { expiresAt } }; - } - - private async getFlowPendingOAuthStart( - { flowManager, serverName }: Pick, - userId: string, - ): Promise { - if (!flowManager) { - return undefined; - } - - const flowId = MCPOAuthHandler.generateFlowId(userId, serverName, getTenantId()); - const existingFlow = await flowManager.getFlowState(flowId, 'mcp_oauth'); - return this.getPendingOAuthStart(existingFlow); - } - - private async reissuePendingOAuthStart( - { flowManager, oauthStart, serverName }: t.UserMCPConnectionOptions, - userId: string, - pendingOAuth?: PendingOAuthState, - ): Promise { - if (!flowManager || typeof oauthStart !== 'function') { - return; - } - - try { - const pendingOAuthStart = await this.getFlowPendingOAuthStart( - { flowManager, serverName }, - userId, - ); - if (!pendingOAuthStart) { - return; - } - - logger.info( - `[MCP][User: ${userId}][${serverName}] Re-issuing stored authorization URL while joining in-flight connection`, - ); - if (pendingOAuth) { - pendingOAuth.lastOAuthStart = pendingOAuthStart; - await this.emitPendingOAuthStart( - pendingOAuth, - oauthStart, - pendingOAuthStart.authURL, - pendingOAuthStart.options, - ); - } else { - await oauthStart(pendingOAuthStart.authURL, pendingOAuthStart.options); - } - } catch (error) { - logger.warn( - `[MCP][User: ${userId}][${serverName}] Failed to re-issue pending OAuth URL`, - error, - ); - } - } - private async createUserConnectionInternal( { serverName, @@ -621,22 +474,41 @@ export abstract class UserConnectionManager { if (!config || (config.updatedAt && connection.isStale(config.updatedAt))) { if (config) { logger.info( - `[MCP][User: ${userId}][${serverName}] Config was updated, disconnecting stale connection`, + `[MCP][User: ${userId}][${serverName}] Config was updated, evicting stale connection`, ); } await this.disconnectUserConnection(userId, serverName, creationGuard); connection = undefined; - } else if (await connection.isConnected()) { - logger.debug(`[MCP][User: ${userId}][${serverName}] Reusing active connection`); - await this.updateUserLastActivity(userId); - await this.assertToolPublicationLeaseCurrent(connection, userId, serverName, creationGuard); - if (creationGuard?.cancelled) { - throw new Error( - `[MCP][User: ${userId}][${serverName}] Connection creation was cancelled during teardown`, - ); - } - return connection; } else { + const activeRecovery = this.getActiveConnectionRecovery(connection); + let awaitedRecovery = activeRecovery; + if (activeRecovery) { + await this.waitForConnectionRecovery(activeRecovery, signal); + } + let connected = await connection.isConnected(); + let recovery = this.getActiveConnectionRecovery(connection); + while (recovery && recovery !== awaitedRecovery) { + awaitedRecovery = recovery; + await this.waitForConnectionRecovery(recovery, signal); + connected = await connection.isConnected(); + recovery = this.getActiveConnectionRecovery(connection); + } + if (connected) { + logger.debug(`[MCP][User: ${userId}][${serverName}] Reusing active connection`); + await this.updateUserLastActivity(userId); + await this.assertToolPublicationLeaseCurrent( + connection, + userId, + serverName, + creationGuard, + ); + if (creationGuard?.cancelled) { + throw new Error( + `[MCP][User: ${userId}][${serverName}] Connection creation was cancelled during teardown`, + ); + } + return connection; + } logger.warn( `[MCP][User: ${userId}][${serverName}] Found existing but disconnected connection object. Cleaning up.`, ); @@ -964,6 +836,115 @@ export abstract class UserConnectionManager { logger.debug(`[MCP][User: ${userId}][${serverName}] Removed connection entry.`); } + protected retainConnection(connection: MCPConnection): void { + const borrowers = this.connectionBorrowers.get(connection) ?? 0; + this.connectionBorrowers.set(connection, borrowers + 1); + } + + protected getActiveConnectionRecovery(_connection: MCPConnection): Promise | undefined { + return undefined; + } + + protected waitForConnectionRecovery( + recovery: Promise, + _signal?: AbortSignal, + ): Promise { + return recovery; + } + + protected holdDeferredConnectionDisposal(connection: MCPConnection): void { + const holds = this.deferredConnectionDisposalHolds.get(connection) ?? 0; + this.deferredConnectionDisposalHolds.set(connection, holds + 1); + } + + protected async releaseDeferredConnectionDisposal(connection: MCPConnection): Promise { + const holds = this.deferredConnectionDisposalHolds.get(connection) ?? 0; + if (holds > 1) { + this.deferredConnectionDisposalHolds.set(connection, holds - 1); + return; + } + this.deferredConnectionDisposalHolds.delete(connection); + await this.finalizeDeferredConnectionDisposal(connection); + } + + protected async releaseConnection(connection: MCPConnection): Promise { + const borrowers = this.connectionBorrowers.get(connection) ?? 0; + if (borrowers > 1) { + this.connectionBorrowers.set(connection, borrowers - 1); + return; + } + + this.connectionBorrowers.delete(connection); + await this.finalizeDeferredConnectionDisposal(connection); + + const drainWaiters = this.connectionBorrowerDrainWaiters.get(connection); + if (drainWaiters) { + this.connectionBorrowerDrainWaiters.delete(connection); + for (const resolve of drainWaiters) { + resolve(); + } + } + } + + protected waitForConnectionBorrowersToDrain(connection: MCPConnection): Promise { + if ((this.connectionBorrowers.get(connection) ?? 0) === 0) { + return Promise.resolve(); + } + + return new Promise((resolve) => { + const drainWaiters = this.connectionBorrowerDrainWaiters.get(connection) ?? new Set(); + drainWaiters.add(resolve); + this.connectionBorrowerDrainWaiters.set(connection, drainWaiters); + }); + } + + protected bindRequestScopedConnectionStore( + requestScopedConnections?: t.RequestScopedMCPConnectionStore, + ): void { + if (!requestScopedConnections || requestScopedConnections.disposeConnection) { + return; + } + + requestScopedConnections.disposeConnection = async (connectionKey, connection) => { + await this.disposeEvictedConnection( + connection as MCPConnection, + `[MCP][Request-scoped: ${connectionKey}]`, + ); + }; + } + + protected async disposeEvictedConnection( + connection: MCPConnection, + logPrefix: string, + ): Promise { + this.deferredConnectionDisposals.set(connection, logPrefix); + await this.finalizeDeferredConnectionDisposal(connection); + } + + private async finalizeDeferredConnectionDisposal(connection: MCPConnection): Promise { + if ( + (this.connectionBorrowers.get(connection) ?? 0) > 0 || + (this.deferredConnectionDisposalHolds.get(connection) ?? 0) > 0 + ) { + return; + } + + const logPrefix = this.deferredConnectionDisposals.get(connection); + if (!logPrefix) { + return; + } + this.deferredConnectionDisposals.delete(connection); + await this.disposeConnection(connection, logPrefix); + } + + private async disposeConnection(connection: MCPConnection, logPrefix: string): Promise { + try { + await connection.dispose(); + } catch (error) { + logger.warn(`${logPrefix} Failed to dispose evicted connection`, error); + } + } + /** Disconnects and removes a specific user connection */ public async disconnectUserConnection( userId: string, @@ -980,10 +961,11 @@ export abstract class UserConnectionManager { const connection = userMap?.get(serverName); try { if (connection) { - logger.info(`[MCP][User: ${userId}][${serverName}] Disconnecting...`); + const logPrefix = `[MCP][User: ${userId}][${serverName}]`; + logger.info(`${logPrefix} Disconnecting...`); connection.removeAllListeners?.('toolsChanged'); this.removeUserConnection(userId, serverName); - await connection.dispose(); + await this.disposeEvictedConnection(connection, logPrefix); } } finally { await cancelMCPToolsChanged({ userId, serverName }); diff --git a/packages/api/src/mcp/__tests__/MCPConnection.test.ts b/packages/api/src/mcp/__tests__/MCPConnection.test.ts index 6ba9d92171..2ca6a966dc 100644 --- a/packages/api/src/mcp/__tests__/MCPConnection.test.ts +++ b/packages/api/src/mcp/__tests__/MCPConnection.test.ts @@ -1,16 +1,17 @@ /** * Tests for MCPConnection error detection methods. * - * These tests use standalone implementations that mirror the private methods in MCPConnection. - * This approach was chosen because MCPConnection requires complex dependencies (Client, transport) - * that are difficult to mock properly. The standalone implementations are kept in sync with - * the actual implementation in connection.ts. + * Rate-limit and SSE tests use standalone implementations that mirror private methods in + * MCPConnection. OAuth classification exercises the production helper shared by the connection + * and factory. * * Alternative approaches considered: * 1. Reflection/type casting - fragile and breaks with refactoring * 2. Protected methods with test subclass - changes public API for testing * 3. Integration tests - tested separately in the full MCP test suite */ +import { isOAuthAuthenticationError } from '~/mcp/errors'; + describe('MCPConnection Error Detection', () => { /** * Standalone implementation of isRateLimitError for testing. @@ -45,52 +46,6 @@ describe('MCPConnection Error Detection', () => { return false; } - /** - * Standalone implementation of isOAuthError for testing. - * This mirrors the private method in MCPConnection (connection.ts). - * Keep in sync with the actual implementation. - */ - function isOAuthError(error: unknown): boolean { - if (!error || typeof error !== 'object') { - return false; - } - - // Check for error code - if ('code' in error) { - const code = (error as { code?: number }).code; - if (code === 401 || code === 403) { - return true; - } - } - - // Check message for various auth error indicators - if ('message' in error && typeof error.message === 'string') { - const message = error.message.toLowerCase(); - // Check for 401 status - if (message.includes('401') || message.includes('non-200 status code (401)')) { - return true; - } - // Check for invalid_grant (OAuth servers return this for expired/revoked grants) - if (message.includes('invalid_grant')) { - return true; - } - // Check for invalid_token (OAuth servers return this for expired/revoked tokens) - if (message.includes('invalid_token')) { - return true; - } - // Check for authentication required - if (message.includes('authentication required') || message.includes('unauthorized')) { - return true; - } - // Check for missing authorization values (e.g., Amazon Ads MCP returns HTTP 400 with this) - if (message.includes('no authorization')) { - return true; - } - } - - return false; - } - describe('isRateLimitError', () => { it('should detect rate limit error by code 429', () => { const error = { code: 429, message: 'Too many requests' }; @@ -142,30 +97,36 @@ describe('MCPConnection Error Detection', () => { }); }); - describe('isOAuthError', () => { - it('should detect OAuth error by code 401', () => { - const error = { code: 401, message: 'Unauthorized' }; - expect(isOAuthError(error)).toBe(true); + describe('isOAuthAuthenticationError', () => { + it.each([ + { code: 401, message: 'Unauthorized' }, + { status: 403, message: 'Forbidden' }, + { statusCode: 401, message: 'Authentication required' }, + { message: 'Error POSTing to endpoint (HTTP 401): Unauthorized' }, + { message: 'Error POSTing to endpoint (HTTP 403): Forbidden' }, + { message: 'Non-200 status code (403)' }, + { message: '403 Forbidden' }, + { message: 'Unauthorized (401)' }, + { message: 'Forbidden (403)' }, + { message: 'The server rejected the token with insufficient_scope' }, + ])('should detect OAuth authentication error %#', (error) => { + expect(isOAuthAuthenticationError(error)).toBe(true); }); - it('should detect OAuth error by code 403', () => { - const error = { code: 403, message: 'Forbidden' }; - expect(isOAuthError(error)).toBe(true); - }); - - it('should detect OAuth error by message containing 401', () => { - const error = { message: 'Error POSTing to endpoint (HTTP 401): Unauthorized' }; - expect(isOAuthError(error)).toBe(true); - }); - - it('should not detect OAuth error for 429 rate limit', () => { - const error = { code: 429, message: 'Too many requests' }; - expect(isOAuthError(error)).toBe(false); + it.each([ + { code: 429, message: 'Too many requests' }, + { message: 'Customer 401 not found' }, + { message: 'Order 403 is unavailable' }, + { message: 'User is unauthorized to delete this record' }, + { message: 'No authorization to delete this record' }, + { code: 400, message: 'Bad request: missing required field' }, + ])('should ignore non-authentication error %#', (error) => { + expect(isOAuthAuthenticationError(error)).toBe(false); }); it('should detect OAuth error for invalid_token', () => { const error = { message: 'The access token is invalid_token or expired' }; - expect(isOAuthError(error)).toBe(true); + expect(isOAuthAuthenticationError(error)).toBe(true); }); it('should detect OAuth error for invalid_grant', () => { @@ -173,7 +134,7 @@ describe('MCPConnection Error Detection', () => { message: 'Streamable HTTP error: Error POSTing to endpoint: {"error":"invalid_grant","error_description":"The provided authorization grant is invalid, expired, or revoked"}', }; - expect(isOAuthError(error)).toBe(true); + expect(isOAuthAuthenticationError(error)).toBe(true); }); it('should detect OAuth error for "no authorization" in message (HTTP 400)', () => { @@ -181,17 +142,12 @@ describe('MCPConnection Error Detection', () => { message: 'Either no authorization values are specified or it could not be derived from the request', }; - expect(isOAuthError(error)).toBe(true); + expect(isOAuthAuthenticationError(error)).toBe(true); }); it('should detect OAuth error for "No authorization" with different casing', () => { const error = { message: 'No Authorization header provided' }; - expect(isOAuthError(error)).toBe(true); - }); - - it('should not detect OAuth error for unrelated 400 errors', () => { - const error = { code: 400, message: 'Bad request: missing required field' }; - expect(isOAuthError(error)).toBe(false); + expect(isOAuthAuthenticationError(error)).toBe(true); }); }); @@ -202,10 +158,10 @@ describe('MCPConnection Error Detection', () => { // Rate limit error should be detected as rate limit, not OAuth expect(isRateLimitError(rateLimitError)).toBe(true); - expect(isOAuthError(rateLimitError)).toBe(false); + expect(isOAuthAuthenticationError(rateLimitError)).toBe(false); // OAuth error should be detected as OAuth, not rate limit - expect(isOAuthError(oauthError)).toBe(true); + expect(isOAuthAuthenticationError(oauthError)).toBe(true); expect(isRateLimitError(oauthError)).toBe(false); }); }); diff --git a/packages/api/src/mcp/__tests__/MCPConnectionAgentLifecycle.test.ts b/packages/api/src/mcp/__tests__/MCPConnectionAgentLifecycle.test.ts index 19c972f0c6..e1715b421f 100644 --- a/packages/api/src/mcp/__tests__/MCPConnectionAgentLifecycle.test.ts +++ b/packages/api/src/mcp/__tests__/MCPConnectionAgentLifecycle.test.ts @@ -549,12 +549,13 @@ describe('MCPConnection SSE 404 handling – session-aware', () => { conn: MCPConnection, transport: ReturnType, code = 404, - ) { + ): Error { ( conn as unknown as { setupTransportErrorHandlers: (t: unknown) => void } ).setupTransportErrorHandlers(transport); const sseError = Object.assign(new Error('Failed to open SSE stream'), { code }); transport.onerror?.(sseError); + return sseError; } beforeEach(() => { @@ -615,6 +616,19 @@ describe('MCPConnection SSE 404 handling – session-aware', () => { expect(mockLogger.warn).toHaveBeenCalledWith(expect.stringContaining('session lost')); expect(emitSpy).toHaveBeenCalledWith('connectionChange', 'error'); }); + it('marks an OAuth-challenged connection unusable without blindly reconnecting', async () => { + const conn = makeConn(); + const transport = makeTransportStub(); + conn.emit('connectionChange', 'connected'); + const emitSpy = jest.spyOn(conn, 'emit'); + + const oauthError = fireSSEError(conn, transport, 401); + + expect(emitSpy).toHaveBeenCalledWith('oauthError', expect.any(Error)); + expect(emitSpy).not.toHaveBeenCalledWith('connectionChange', 'error'); + expect(await conn.isConnected()).toBe(false); + expect(conn.getLastConnectionCheckError()).toBe(oauthError); + }); }); describe('MCPConnection SSE stream disconnect handling', () => { diff --git a/packages/api/src/mcp/__tests__/MCPConnectionFactory.oauthSdk.integration.test.ts b/packages/api/src/mcp/__tests__/MCPConnectionFactory.oauthSdk.integration.test.ts index 4ea28e1363..95e17af213 100644 --- a/packages/api/src/mcp/__tests__/MCPConnectionFactory.oauthSdk.integration.test.ts +++ b/packages/api/src/mcp/__tests__/MCPConnectionFactory.oauthSdk.integration.test.ts @@ -11,10 +11,12 @@ import { createOAuthMCPServer, type OAuthTestServer, } from './helpers/oauthTestServer'; +import { MCPServersRegistry } from '~/mcp/registry/MCPServersRegistry'; import { MCPConnectionFactory } from '~/mcp/MCPConnectionFactory'; +import { MCPTokenStorage, MCPOAuthHandler } from '~/mcp/oauth'; import { FlowStateManager } from '~/flow/manager'; import { MCPConnection } from '~/mcp/connection'; -import { MCPTokenStorage } from '~/mcp/oauth'; +import { MCPManager } from '~/mcp/MCPManager'; jest.mock('@librechat/data-schemas', () => ({ logger: { @@ -64,6 +66,16 @@ async function safeDisconnect(conn: MCPConnection | null): Promise { await conn.disconnect().catch(() => undefined); } +async function waitFor(condition: () => boolean, timeoutMs = 5000): Promise { + const deadline = Date.now() + timeoutMs; + while (!condition()) { + if (Date.now() > deadline) { + throw new Error('Timed out waiting for condition'); + } + await new Promise((resolve) => setTimeout(resolve, 10)); + } +} + function createFlowManager(): FlowStateManager { return new FlowStateManager(new MockKeyv() as unknown as Keyv, { ttl: 30000, @@ -122,7 +134,7 @@ async function storeTokens( server: OAuthTestServer, tokens: MCPOAuthTokens, scope = 'read', -): Promise { +): Promise { const clientInfo: OAuthClientInformation = { client_id: CLIENT_ID, redirect_uris: ['http://localhost'], @@ -144,7 +156,7 @@ async function storeTokens( resource: server.resourceUrl, }; - await MCPTokenStorage.storeTokens({ + return MCPTokenStorage.storeTokens({ userId: USER_ID, serverName: SERVER_NAME, tokens, @@ -231,6 +243,258 @@ describe('MCPConnectionFactory OAuth against real SDK Streamable HTTP server', ( expect(storedAccessToken?.token).not.toBe(`enc:${initialTokens.access_token}`); }); + it('recovers a tool call rejected after connection and retries with the refreshed token', async () => { + server = await createOAuthMCPServer({ + issueRefreshTokens: true, + requireResourceParameter: true, + tokenScopes: ['read'], + scopesSupported: ['read'], + }); + const initialTokens = await issueTokens(server); + await storeTokens(tokenStore, server, initialTokens); + const flowManager = createFlowManager(); + const tokenMethods = { + findToken: tokenStore.findToken, + createToken: tokenStore.createToken, + updateToken: tokenStore.updateToken, + deleteTokens: tokenStore.deleteTokens, + }; + const serverConfig = { + type: 'streamable-http' as const, + url: server.url, + initTimeout: 15000, + requiresOAuth: true, + }; + + connection = await MCPConnectionFactory.create( + { serverName: SERVER_NAME, serverConfig }, + { + useOAuth: true, + user: { id: USER_ID } as IUser, + flowManager, + tokenMethods, + }, + ); + server.issuedTokens.delete(initialTokens.access_token); + const isConnectedSpy = jest.spyOn(connection, 'isConnected').mockResolvedValueOnce(true); + + const manager = new MCPManager(); + jest.spyOn(manager, 'getConnection').mockResolvedValue(connection); + const registrySpy = jest.spyOn(MCPServersRegistry, 'getInstance').mockReturnValue({ + resolveAllowlists: jest.fn().mockResolvedValue({ + allowedDomains: null, + allowedAddresses: null, + useSSRFProtection: false, + }), + } as unknown as MCPServersRegistry); + const oauthStart = jest.fn(async (_authorizationUrl: string): Promise => undefined); + + try { + await expect( + manager.callTool({ + user: { id: USER_ID } as IUser, + serverName: SERVER_NAME, + serverConfig, + toolName: 'echo', + toolArguments: { message: 'runtime refresh' }, + provider: 'openai', + flowManager, + tokenMethods, + oauthStart, + }), + ).resolves.toBeDefined(); + + expect( + server.tokenRequests.filter((request) => request.grantType === 'refresh_token'), + ).toHaveLength(1); + expect(oauthStart).not.toHaveBeenCalled(); + } finally { + isConnectedSpy.mockRestore(); + registrySpy.mockRestore(); + } + }); + + it('lets an in-flight request finish before a concurrent OAuth reconnect', async () => { + let markSlowRequestStarted: (() => void) | undefined; + let releaseFirstSlowRequest: (() => void) | undefined; + const slowRequestStarted = new Promise((resolve) => { + markSlowRequestStarted = resolve; + }); + const firstSlowRequestBlocked = new Promise((resolve) => { + releaseFirstSlowRequest = resolve; + }); + let slowRequestCount = 0; + server = await createOAuthMCPServer({ + issueRefreshTokens: true, + requireResourceParameter: true, + tokenScopes: ['read'], + scopesSupported: ['read'], + echoHandler: async (message) => { + if (message === 'slow borrower' && slowRequestCount++ === 0) { + markSlowRequestStarted?.(); + await firstSlowRequestBlocked; + } + return `echo: ${message}`; + }, + }); + const initialTokens = await issueTokens(server); + await storeTokens(tokenStore, server, initialTokens); + const flowManager = createFlowManager(); + const tokenMethods = { + findToken: tokenStore.findToken, + createToken: tokenStore.createToken, + updateToken: tokenStore.updateToken, + deleteTokens: tokenStore.deleteTokens, + }; + const serverConfig = { + type: 'streamable-http' as const, + url: server.url, + initTimeout: 15000, + requiresOAuth: true, + }; + + connection = await MCPConnectionFactory.create( + { serverName: SERVER_NAME, serverConfig }, + { + useOAuth: true, + user: { id: USER_ID } as IUser, + flowManager, + tokenMethods, + }, + ); + const connectSpy = jest.spyOn(connection, 'connect'); + const isConnectedSpy = jest.spyOn(connection, 'isConnected').mockResolvedValue(true); + + const manager = new MCPManager(); + jest.spyOn(manager, 'getConnection').mockResolvedValue(connection); + const registrySpy = jest.spyOn(MCPServersRegistry, 'getInstance').mockReturnValue({ + resolveAllowlists: jest.fn().mockResolvedValue({ + allowedDomains: null, + allowedAddresses: null, + useSSRFProtection: false, + }), + } as unknown as MCPServersRegistry); + const oauthStart = jest.fn(async (_authorizationUrl: string): Promise => undefined); + const callTool = (message: string) => + manager.callTool({ + user: { id: USER_ID } as IUser, + serverName: SERVER_NAME, + serverConfig, + toolName: 'echo', + toolArguments: { message }, + provider: 'openai', + flowManager, + tokenMethods, + oauthStart, + }); + + try { + const slowCall = callTool('slow borrower'); + await slowRequestStarted; + server.issuedTokens.delete(initialTokens.access_token); + + const recoveringCall = callTool('recovery owner'); + await waitFor( + () => + server.tokenRequests.filter((request) => request.grantType === 'refresh_token').length === + 1, + ); + + expect(connectSpy).not.toHaveBeenCalled(); + expect(slowRequestCount).toBe(1); + releaseFirstSlowRequest?.(); + + await expect(Promise.all([slowCall, recoveringCall])).resolves.toHaveLength(2); + expect( + server.tokenRequests.filter((request) => request.grantType === 'refresh_token'), + ).toHaveLength(1); + expect(slowRequestCount).toBe(1); + expect(oauthStart).not.toHaveBeenCalled(); + } finally { + releaseFirstSlowRequest?.(); + connectSpy.mockRestore(); + isConnectedSpy.mockRestore(); + registrySpy.mockRestore(); + } + }); + + it('escalates to interactive OAuth when the resource rejects refreshed tokens during reconnect', async () => { + server = await createOAuthMCPServer({ + issueRefreshTokens: true, + requireResourceParameter: true, + tokenScopes: ['read'], + scopesSupported: ['read'], + rejectRefreshTokens: 10, + }); + const initialTokens = await issueTokens(server); + await storeTokens(tokenStore, server, initialTokens); + const flowManager = createFlowManager(); + const tokenMethods = { + findToken: tokenStore.findToken, + createToken: tokenStore.createToken, + updateToken: tokenStore.updateToken, + deleteTokens: tokenStore.deleteTokens, + }; + const serverConfig = { + type: 'streamable-http' as const, + url: server.url, + initTimeout: 15000, + requiresOAuth: true, + }; + + connection = await MCPConnectionFactory.create( + { serverName: SERVER_NAME, serverConfig }, + { + useOAuth: true, + user: { id: USER_ID } as IUser, + flowManager, + tokenMethods, + }, + ); + server.issuedTokens.delete(initialTokens.access_token); + const isConnectedSpy = jest.spyOn(connection, 'isConnected').mockResolvedValueOnce(true); + + const manager = new MCPManager(); + jest.spyOn(manager, 'getConnection').mockResolvedValue(connection); + const registrySpy = jest.spyOn(MCPServersRegistry, 'getInstance').mockReturnValue({ + resolveAllowlists: jest.fn().mockResolvedValue({ + allowedDomains: null, + allowedAddresses: null, + useSSRFProtection: false, + }), + } as unknown as MCPServersRegistry); + const oauthStart = jest.fn(async (): Promise => { + const authorizedTokens = await issueTokens(server); + const storedTokens = await storeTokens(tokenStore, server, authorizedTokens); + const flowId = MCPOAuthHandler.generateFlowId(USER_ID, SERVER_NAME); + await flowManager.completeFlow(flowId, 'mcp_oauth', storedTokens); + }); + + try { + await expect( + manager.callTool({ + user: { id: USER_ID } as IUser, + serverName: SERVER_NAME, + serverConfig, + toolName: 'echo', + toolArguments: { message: 'interactive fallback' }, + provider: 'openai', + flowManager, + tokenMethods, + oauthStart, + }), + ).resolves.toBeDefined(); + + expect( + server.tokenRequests.filter((request) => request.grantType === 'refresh_token'), + ).toHaveLength(1); + expect(oauthStart).toHaveBeenCalledTimes(1); + } finally { + isConnectedSpy.mockRestore(); + registrySpy.mockRestore(); + } + }); + it('does not silently refresh an SDK insufficient_scope challenge before starting OAuth', async () => { server = await createOAuthMCPServer({ issueRefreshTokens: true, diff --git a/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts b/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts index 9b28a0f8e9..97febbfd8a 100644 --- a/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts +++ b/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts @@ -95,6 +95,9 @@ describe('MCPConnectionFactory', () => { beforeEach(() => { jest.clearAllMocks(); + // Cached runtime handlers now delegate before attempting their own refresh, + // so queued one-shot refresh results must not leak into the next test. + mockMCPTokenStorage.forceRefreshTokens.mockReset(); // Clear process-local silent-refresh in-flight map so a leftover entry // from a prior test (e.g. one that errored before its `finally` ran) // cannot cause a later test to join a stale promise. @@ -207,6 +210,95 @@ describe('MCPConnectionFactory', () => { } }); + it('aborts only the local waiter for a shared OAuth flow', async () => { + const abortController = new AbortController(); + const abortReason = new Error('owner request aborted'); + const sseConfig = { + url: 'https://api.example.com/mcp', + type: 'sse' as const, + requiresOAuth: true, + } as t.SSEOptions; + const pendingTokens = new Promise(() => undefined); + const oauthStart = jest.fn().mockResolvedValue(undefined); + + mockProcessMCPEnv.mockReturnValue(sseConfig); + mockFlowManager.getFlowState.mockResolvedValue({ + status: 'PENDING', + type: 'mcp_oauth', + metadata: { + authorizationUrl: 'https://auth.example.com/pending', + serverUrl: sseConfig.url, + }, + createdAt: Date.now(), + }); + mockFlowManager.createFlow.mockReturnValue(pendingTokens); + + const factory = new InspectableMCPConnectionFactory( + { serverName: 'test-server', serverConfig: sseConfig }, + { + useOAuth: true, + user: mockUser, + flowManager: mockFlowManager, + oauthStart, + signal: abortController.signal, + }, + ); + + const resultPromise = factory.handleOAuthRequiredForTest(); + await new Promise((resolve) => setImmediate(resolve)); + + expect(mockFlowManager.createFlow).toHaveBeenCalledWith('user123:test-server', 'mcp_oauth', {}); + + abortController.abort(abortReason); + + await expect(resultPromise).resolves.toEqual( + expect.objectContaining({ tokens: null, error: abortReason }), + ); + expect(mockFlowManager.deleteFlow).not.toHaveBeenCalled(); + }); + + it('observes the shared flow after an already-aborted local wait', async () => { + const abortController = new AbortController(); + const abortReason = new Error('request already aborted'); + const sseConfig = { + url: 'https://api.example.com/mcp', + type: 'sse' as const, + requiresOAuth: true, + } as t.SSEOptions; + let rejectSharedFlow: ((error: Error) => void) | undefined; + const sharedFlow = new Promise((_resolve, reject) => { + rejectSharedFlow = reject; + }); + + abortController.abort(abortReason); + mockProcessMCPEnv.mockReturnValue(sseConfig); + mockFlowManager.getFlowState.mockResolvedValue({ + status: 'PENDING', + type: 'mcp_oauth', + metadata: { authorizationUrl: 'https://auth.example.com/pending' }, + createdAt: Date.now(), + }); + mockFlowManager.createFlow.mockReturnValue(sharedFlow); + + const factory = new InspectableMCPConnectionFactory( + { serverName: 'test-server', serverConfig: sseConfig }, + { + useOAuth: true, + user: mockUser, + flowManager: mockFlowManager, + signal: abortController.signal, + }, + ); + + await expect(factory.handleOAuthRequiredForTest()).resolves.toEqual( + expect.objectContaining({ tokens: null, error: abortReason }), + ); + + rejectSharedFlow?.(new Error('shared flow later failed')); + await new Promise((resolve) => setImmediate(resolve)); + expect(mockFlowManager.deleteFlow).not.toHaveBeenCalled(); + }); + describe('static create method', () => { it('should create a basic connection without OAuth', async () => { const basicOptions = { @@ -1250,12 +1342,7 @@ describe('MCPConnectionFactory', () => { expect(initCallOrder).toBeLessThan(createCallOrder); // createFlow should receive {} since initFlow already persisted metadata - expect(mockFlowManager.createFlow).toHaveBeenCalledWith( - 'flow123', - 'mcp_oauth', - {}, - undefined, - ); + expect(mockFlowManager.createFlow).toHaveBeenCalledWith('flow123', 'mcp_oauth', {}); }); it('should delete stale flow and create new OAuth flow when existing flow is COMPLETED', async () => { @@ -1342,7 +1429,6 @@ describe('MCPConnectionFactory', () => { 'user123:test-server', 'mcp_oauth', {}, - undefined, ); }); @@ -1416,7 +1502,7 @@ describe('MCPConnectionFactory', () => { }), ); expect(mockConnectionInstance.setOAuthTokens).toHaveBeenCalledWith(refreshedTokens); - expect(mockConnectionInstance.emit).toHaveBeenCalledWith('oauthHandled'); + expect(mockConnectionInstance.emit).toHaveBeenCalledWith('oauthHandled', 'silent-refresh'); // Silent refresh succeeded — interactive flow must NOT be initiated. expect(mockMCPOAuthHandler.initiateOAuthFlow).not.toHaveBeenCalled(); }); @@ -1561,7 +1647,7 @@ describe('MCPConnectionFactory', () => { await oauthRequiredHandler!({ serverUrl: 'https://api.example.com' }); expect(mockConnectionInstance.setOAuthTokens).toHaveBeenCalledWith(refreshedTokens); - expect(mockConnectionInstance.emit).toHaveBeenCalledWith('oauthHandled'); + expect(mockConnectionInstance.emit).toHaveBeenCalledWith('oauthHandled', 'silent-refresh'); // returnOnOAuth interactive path must NOT trigger when silent refresh succeeds. expect(mockMCPOAuthHandler.initiateOAuthFlow).not.toHaveBeenCalled(); expect(oauthOptions.oauthStart).not.toHaveBeenCalled(); @@ -1743,7 +1829,7 @@ describe('MCPConnectionFactory', () => { // cached ones — that's the whole point of the fix. expect(mockConnectionInstance.setOAuthTokens).toHaveBeenCalledWith(freshlyRefreshedTokens); expect(mockConnectionInstance.setOAuthTokens).not.toHaveBeenCalledWith(staleCachedTokens); - expect(mockConnectionInstance.emit).toHaveBeenCalledWith('oauthHandled'); + expect(mockConnectionInstance.emit).toHaveBeenCalledWith('oauthHandled', 'silent-refresh'); // The cached `mcp_get_tokens` flow state is dropped so the next // `getOAuthTokens` call reads the freshly persisted tokens from storage. expect(mockFlowManager.deleteFlow).toHaveBeenCalledWith('flow123', 'mcp_get_tokens'); @@ -2345,6 +2431,88 @@ describe('MCPConnectionFactory', () => { cleanup(); }); + it('bounds request recovery to one silent refresh and one interactive flow', async () => { + const sseConfig = { + ...mockServerConfig, + url: 'https://api.example.com', + type: 'sse' as const, + } as t.SSEOptions; + const basicOptions = { + serverName: 'test-server', + serverConfig: sseConfig, + }; + const refreshedTokens: MCPOAuthTokens = { + access_token: 'refreshed-access', + refresh_token: 'refresh-token', + token_type: 'Bearer', + obtained_at: Date.now(), + }; + const interactiveTokens: MCPOAuthTokens = { + access_token: 'interactive-access', + token_type: 'Bearer', + obtained_at: Date.now(), + credential_set_id: 'interactive-generation', + }; + mockProcessMCPEnv.mockReturnValue(sseConfig); + mockMCPOAuthHandler.generateFlowId.mockReturnValue('flow123'); + mockMCPTokenStorage.forceRefreshTokens.mockResolvedValueOnce(refreshedTokens); + mockFlowManager.getFlowState.mockResolvedValue(null); + mockMCPOAuthHandler.initiateOAuthFlow.mockResolvedValueOnce({ + authorizationUrl: 'https://auth.example.com', + flowId: 'flow123', + flowMetadata: { + serverName: 'test-server', + userId: 'user123', + serverUrl: 'https://api.example.com', + state: 'fresh-state', + }, + }); + mockFlowManager.createFlow.mockResolvedValueOnce(interactiveTokens); + + let requestOAuthHandler: ((data: Record) => Promise) | undefined; + mockConnectionInstance.on.mockImplementation((event, handler) => { + if (event === 'oauthReauthenticationRequired') { + requestOAuthHandler = handler as (data: Record) => Promise; + } + return mockConnectionInstance; + }); + const cleanup = MCPConnectionFactory.attachRequestOAuthHandler( + basicOptions, + { + useOAuth: true, + user: mockUser, + flowManager: mockFlowManager, + oauthStart: jest.fn(), + tokenMethods: { + findToken: jest.fn(), + createToken: jest.fn(), + updateToken: jest.fn(), + deleteTokens: jest.fn(), + }, + }, + mockConnectionInstance, + ); + const challenge = { + serverUrl: 'https://api.example.com', + error: new Error('Non-200 status code (401)'), + }; + + await Promise.all([requestOAuthHandler!(challenge), requestOAuthHandler!(challenge)]); + await requestOAuthHandler!(challenge); + await requestOAuthHandler!(challenge); + + expect(mockMCPTokenStorage.forceRefreshTokens).toHaveBeenCalledTimes(1); + expect(mockMCPOAuthHandler.initiateOAuthFlow).toHaveBeenCalledTimes(1); + expect(mockConnectionInstance.setOAuthTokens).toHaveBeenNthCalledWith(1, refreshedTokens); + expect(mockConnectionInstance.setOAuthTokens).toHaveBeenNthCalledWith(2, interactiveTokens); + expect(mockConnectionInstance.emit).toHaveBeenCalledWith( + 'oauthFailed', + expect.objectContaining({ message: 'OAuth recovery phase budget exhausted' }), + ); + + cleanup(); + }); + it('should not reuse request-scoped OAuth callbacks after connection is cached', async () => { const sseConfig = { ...mockServerConfig, @@ -3273,7 +3441,7 @@ describe('MCPConnectionFactory', () => { expect(mockMCPTokenStorage.storeTokens).not.toHaveBeenCalled(); expect(mockConnectionInstance.setOAuthTokens).toHaveBeenCalledWith(callbackTokens); - expect(mockConnectionInstance.emit).toHaveBeenCalledWith('oauthHandled'); + expect(mockConnectionInstance.emit).toHaveBeenCalledWith('oauthHandled', 'interactive'); }); it('rejects callback tokens that do not identify a persisted credential generation', async () => { diff --git a/packages/api/src/mcp/__tests__/MCPManager.test.ts b/packages/api/src/mcp/__tests__/MCPManager.test.ts index 89e082ef5c..38506d7299 100644 --- a/packages/api/src/mcp/__tests__/MCPManager.test.ts +++ b/packages/api/src/mcp/__tests__/MCPManager.test.ts @@ -1,8 +1,10 @@ +import { EventEmitter } from 'events'; import { logger } from '@librechat/data-schemas'; import type { IUser } from '@librechat/data-schemas'; import type { GraphTokenResolver } from '~/utils/graph'; import type * as t from '~/mcp/types'; import { OboTokenResolutionError, detectOAuthRequirement, resolveOboToken } from '~/mcp/oauth'; +import { createMCPRequestContext, cleanupMCPRequestContext } from '~/mcp/request'; import { MCPServersInitializer } from '~/mcp/registry/MCPServersInitializer'; import { MCPServerInspector } from '~/mcp/registry/MCPServerInspector'; import { ConnectionsRepository } from '~/mcp/ConnectionsRepository'; @@ -283,6 +285,82 @@ describe('MCPManager', () => { }); }); + describe('withUserConnectionLease', () => { + it('retains a public checkout until its metadata read completes', async () => { + const connection = {} as MCPConnection; + let finishRead: (() => void) | undefined; + const read = new Promise((resolve) => { + finishRead = resolve; + }); + mockAppConnections({}); + const manager = await MCPManager.createInstance(newMCPServersConfig()); + jest.spyOn(manager, 'getUserConnection').mockResolvedValue(connection); + const lifecycle = manager as unknown as { + waitForConnectionBorrowersToDrain: (connection: MCPConnection) => Promise; + }; + + const checkout = manager.withUserConnectionLease( + { serverName, user: { id: userId } as IUser }, + async () => { + await read; + return 'snapshot'; + }, + ); + await new Promise((resolve) => setImmediate(resolve)); + + let drained = false; + const drain = lifecycle.waitForConnectionBorrowersToDrain(connection).then(() => { + drained = true; + }); + await new Promise((resolve) => setImmediate(resolve)); + expect(drained).toBe(false); + + finishRead?.(); + await expect(checkout).resolves.toBe('snapshot'); + await drain; + expect(drained).toBe(true); + }); + + it('rejoins recovery registered between public checkout and lease acquisition', async () => { + const staleConnection = {} as MCPConnection; + const recoveredConnection = {} as MCPConnection; + let finishCheckout: ((connection: MCPConnection) => void) | undefined; + const checkout = new Promise((resolve) => { + finishCheckout = resolve; + }); + let finishRecovery: (() => void) | undefined; + const recovery = new Promise((resolve) => { + finishRecovery = resolve; + }); + mockAppConnections({}); + const manager = await MCPManager.createInstance(newMCPServersConfig()); + jest + .spyOn(manager, 'getUserConnection') + .mockReturnValueOnce(checkout) + .mockResolvedValue(recoveredConnection); + const internals = manager as unknown as { + oauthRecoveries: WeakMap }>; + }; + const operation = jest.fn().mockResolvedValue('snapshot'); + + const result = manager.withUserConnectionLease( + { serverName, user: { id: userId } as IUser }, + operation, + ); + finishCheckout?.(staleConnection); + internals.oauthRecoveries.set(staleConnection, { promise: recovery }); + await new Promise((resolve) => setImmediate(resolve)); + + expect(operation).not.toHaveBeenCalled(); + + internals.oauthRecoveries.delete(staleConnection); + finishRecovery?.(); + + await expect(result).resolves.toBe('snapshot'); + expect(operation).toHaveBeenCalledWith(recoveredConnection); + }); + }); + describe('connectAppServers', () => { it('opens only operator app connections and refreshes their current catalogs', async () => { const connection = new MCPConnection({ @@ -521,6 +599,82 @@ describe('MCPManager', () => { expect(mockLogger.warn).not.toHaveBeenCalled(); }); + it('leases a cached user connection while reading its tool metadata', async () => { + let finishInspection: ((tools: t.LCAvailableTools) => void) | undefined; + const inspection = new Promise((resolve) => { + finishInspection = resolve; + }); + const connection = {} as MCPConnection; + (MCPServerInspector.getToolFunctions as jest.Mock) = jest.fn().mockReturnValue(inspection); + mockAppConnections({ + get: jest.fn().mockResolvedValue(null), + }); + + const manager = await MCPManager.createInstance(newMCPServersConfig()); + const lifecycle = manager as unknown as { + userConnections: Map>; + waitForConnectionBorrowersToDrain: (connection: MCPConnection) => Promise; + }; + lifecycle.userConnections.set(userId, new Map([[serverName, connection]])); + + const toolsPromise = manager.getServerToolFunctions(userId, serverName); + await new Promise((resolve) => setImmediate(resolve)); + let drained = false; + const drainPromise = lifecycle.waitForConnectionBorrowersToDrain(connection).then(() => { + drained = true; + }); + await new Promise((resolve) => setImmediate(resolve)); + + expect(drained).toBe(false); + + finishInspection?.({}); + await expect(toolsPromise).resolves.toEqual({}); + await drainPromise; + expect(drained).toBe(true); + }); + + it('rejoins recovery after leasing cached tool metadata', async () => { + const staleConnection = {} as MCPConnection; + const recoveredConnection = {} as MCPConnection; + let resolveRecovery: (() => void) | undefined; + const recovery = new Promise((resolve) => { + resolveRecovery = resolve; + }); + (MCPServerInspector.getToolFunctions as jest.Mock) = jest.fn().mockResolvedValue({}); + mockAppConnections({ + get: jest.fn().mockResolvedValue(null), + }); + + const manager = await MCPManager.createInstance(newMCPServersConfig()); + const internals = manager as unknown as { + oauthRecoveries: WeakMap< + MCPConnection, + { promise: Promise; allowsTakeover: boolean } + >; + userConnections: Map>; + }; + internals.userConnections.set(userId, new Map([[serverName, staleConnection]])); + internals.oauthRecoveries.set(staleConnection, { + promise: recovery, + allowsTakeover: true, + }); + + const toolsPromise = manager.getServerToolFunctions(userId, serverName); + await new Promise((resolve) => setImmediate(resolve)); + + expect(MCPServerInspector.getToolFunctions).not.toHaveBeenCalled(); + + internals.userConnections.set(userId, new Map([[serverName, recoveredConnection]])); + internals.oauthRecoveries.delete(staleConnection); + resolveRecovery?.(); + + await expect(toolsPromise).resolves.toEqual({}); + expect(MCPServerInspector.getToolFunctions).toHaveBeenCalledWith( + serverName, + recoveredConnection, + ); + }); + it('should include specific server name in error messages', async () => { const specificServerName = 'github_mcp_server'; @@ -794,7 +948,7 @@ describe('MCPManager', () => { ); }); - it('should attach request OAuth handler without reprocessing resolved config', async () => { + it('should attach a recovery handler without reprocessing resolved config', async () => { const rawServerConfig = { type: 'sse', url: 'https://api.example.com/{{LIBRECHAT_USER_ID}}', @@ -813,18 +967,40 @@ describe('MCPManager', () => { Authorization: 'Bearer ${SHOULD_NOT_EXPAND}', }, }; - const cleanupOAuthHandler = jest.fn(); - + const authError = new Error('Non-200 status code (401)'); + const connection = Object.assign(new EventEmitter(), { + client: { + request: jest + .fn() + .mockRejectedValueOnce(authError) + .mockResolvedValueOnce({ + content: [{ type: 'text', text: 'Recovered result' }], + isError: false, + }), + }, + connect: jest.fn().mockResolvedValue(undefined), + getLastConnectionCheckError: jest.fn().mockReturnValue(undefined), + isConnected: jest.fn().mockResolvedValue(true), + isOAuthAuthenticationError: jest.fn((error: unknown) => error === authError), + setRequestHeaders: jest.fn(), + timeout: 30000, + url: processedServerConfig.url, + }) as unknown as MCPConnection; mockProcessMCPEnv.mockReturnValue(processedServerConfig); - (MCPConnectionFactory.attachRequestOAuthHandler as jest.Mock).mockReturnValue( - cleanupOAuthHandler, + (MCPConnectionFactory.attachRequestOAuthHandler as jest.Mock).mockImplementation( + (_basic, _oauth, currentConnection: MCPConnection) => { + const listener = () => currentConnection.emit('oauthHandled'); + currentConnection.on('oauthReauthenticationRequired', listener); + return () => currentConnection.off('oauthReauthenticationRequired', listener); + }, ); mockAppConnections({ - get: jest.fn().mockResolvedValue(mockConnection), + get: jest.fn().mockResolvedValue(connection), }); (mockRegistryInstance.getServerConfig as jest.Mock).mockResolvedValue(rawServerConfig); const manager = await MCPManager.createInstance(newMCPServersConfig()); + jest.spyOn(manager, 'getConnection').mockResolvedValue(connection); const oauthStart = jest.fn(); await manager.callTool({ @@ -838,21 +1014,20 @@ describe('MCPManager', () => { >[0]['flowManager'], }); - /** Runtime resolution, config-identity binding, and callTool each process once — none from - * attaching the request handler. */ - expect(mockProcessMCPEnv).toHaveBeenCalledTimes(3); + /** callTool resolves the config once; handler attachment does not process it again. */ + expect(mockProcessMCPEnv).toHaveBeenCalledTimes(1); expect(MCPConnectionFactory.attachRequestOAuthHandler).toHaveBeenCalledWith( expect.objectContaining({ serverConfig: processedServerConfig, skipEnvProcessing: true, }), expect.objectContaining({ - oauthStart, + oauthStart: expect.any(Function), user: mockUser, }), - mockConnection, + connection, ); - expect(cleanupOAuthHandler).toHaveBeenCalled(); + expect(connection.listenerCount('oauthReauthenticationRequired')).toBe(0); }); it('should leave graph token placeholders sandboxed for user-sourced configs', async () => { @@ -1152,6 +1327,1227 @@ describe('MCPManager', () => { }); }); + describe('callTool - runtime OAuth recovery', () => { + const mockUser = { id: 'oauth-user' } as IUser; + const mockFlowManager = {} as Parameters[0]['flowManager']; + const serverConfig: t.StreamableHTTPOptions = { + type: 'streamable-http', + url: 'https://mcp.example.com', + requiresOAuth: true, + oauth: { + authorization_url: 'https://auth.example.com/authorize', + }, + }; + const toolResult = { + content: [{ type: 'text', text: 'Recovered result' }], + isError: false, + }; + + function createConnection(request: jest.Mock) { + const emitter = new EventEmitter(); + return Object.assign(emitter, { + client: { request }, + connect: jest.fn().mockResolvedValue(undefined), + disconnect: jest.fn().mockResolvedValue(undefined), + dispose: jest.fn().mockResolvedValue(undefined), + getLastConnectionCheckError: jest.fn().mockReturnValue(undefined), + isConnected: jest.fn().mockResolvedValue(true), + isOAuthAuthenticationError: jest.fn((error: unknown) => { + return error instanceof Error && error.message.includes('401'); + }), + usesOAuth: jest.fn().mockReturnValue(true), + setRequestHeaders: jest.fn(), + timeout: 30000, + url: serverConfig.url, + }) as unknown as MCPConnection; + } + + function attachOAuthHandler( + handler: (connection: MCPConnection) => void = (connection) => { + connection.emit('oauthHandled'); + }, + ) { + (MCPConnectionFactory.attachRequestOAuthHandler as jest.Mock).mockImplementation( + (_basic, _oauth, connection: MCPConnection) => { + const listener = () => handler(connection); + connection.on('oauthReauthenticationRequired', listener); + return () => connection.off('oauthReauthenticationRequired', listener); + }, + ); + } + + async function callTool( + manager: MCPManager, + signal?: AbortSignal, + oauthStart: t.OAuthStartHandler = jest.fn(), + requestScopedConnections?: t.RequestScopedMCPConnectionStore, + oauthEnd?: () => Promise, + ) { + return manager.callTool({ + user: mockUser, + serverName, + toolName: 'oauth_tool', + provider: 'openai', + oauthStart, + oauthEnd, + flowManager: mockFlowManager, + requestScopedConnections, + options: signal ? { signal } : undefined, + }); + } + + async function createManager(connection: MCPConnection) { + mockAppConnections({}); + const manager = await MCPManager.createInstance(newMCPServersConfig()); + jest.spyOn(manager, 'getConnection').mockResolvedValue(connection); + return manager; + } + + beforeEach(() => { + (mockRegistryInstance.getServerConfig as jest.Mock).mockResolvedValue(serverConfig); + mockProcessMCPEnv.mockReturnValue(serverConfig); + }); + + it('refreshes, rebuilds, and retries a rejected tool request once', async () => { + const authError = new Error('Non-200 status code (401)'); + const request = jest.fn().mockRejectedValueOnce(authError).mockResolvedValueOnce(toolResult); + const connection = createConnection(request); + attachOAuthHandler(); + const manager = await createManager(connection); + + await expect(callTool(manager)).resolves.toBeDefined(); + + expect(connection.connect).toHaveBeenCalledTimes(1); + expect(request).toHaveBeenCalledTimes(2); + expect(request.mock.calls[1]).toEqual(request.mock.calls[0]); + }); + + it('joins recovery registered between cached checkout and lease acquisition', async () => { + const request = jest.fn().mockResolvedValue(toolResult); + const connection = createConnection(request); + const manager = await createManager(connection); + let resolveConnection: ((connection: MCPConnection) => void) | undefined; + const firstCheckout = new Promise((resolve) => { + resolveConnection = resolve; + }); + (manager.getConnection as jest.Mock) + .mockReset() + .mockReturnValueOnce(firstCheckout) + .mockResolvedValue(connection); + let resolveRecovery: (() => void) | undefined; + const recovery = new Promise((resolve) => { + resolveRecovery = resolve; + }); + const internals = manager as unknown as { + oauthRecoveries: WeakMap< + MCPConnection, + { promise: Promise; allowsTakeover: boolean } + >; + }; + + const toolPromise = callTool(manager); + resolveConnection?.(connection); + internals.oauthRecoveries.set(connection, { promise: recovery, allowsTakeover: true }); + await new Promise((resolve) => setImmediate(resolve)); + + expect(request).not.toHaveBeenCalled(); + expect(manager.getConnection).toHaveBeenCalledTimes(1); + + internals.oauthRecoveries.delete(connection); + resolveRecovery?.(); + + await expect(toolPromise).resolves.toBeDefined(); + expect(manager.getConnection).toHaveBeenCalledTimes(2); + expect(request).toHaveBeenCalledTimes(1); + }); + + it('reacquires the cached connection after claiming a failed checkout recovery', async () => { + const staleRequest = jest.fn(); + const recoveredRequest = jest.fn().mockResolvedValue(toolResult); + const staleConnection = createConnection(staleRequest); + const recoveredConnection = createConnection(recoveredRequest); + const manager = await createManager(staleConnection); + (manager.getConnection as jest.Mock) + .mockReset() + .mockResolvedValueOnce(staleConnection) + .mockResolvedValue(recoveredConnection); + let rejectRecovery: ((error: Error) => void) | undefined; + const recovery = new Promise((_resolve, reject) => { + rejectRecovery = reject; + }); + const internals = manager as unknown as { + oauthRecoveries: WeakMap< + MCPConnection, + { promise: Promise; allowsTakeover: boolean } + >; + }; + internals.oauthRecoveries.set(staleConnection, { + promise: recovery, + allowsTakeover: true, + }); + + const toolPromise = callTool(manager); + await new Promise((resolve) => setImmediate(resolve)); + rejectRecovery?.(new Error('Shared recovery failed')); + + await expect(toolPromise).resolves.toBeDefined(); + expect(manager.getConnection).toHaveBeenCalledTimes(2); + expect(staleRequest).not.toHaveBeenCalled(); + expect(recoveredRequest).toHaveBeenCalledTimes(1); + }); + + it('disposes a recovered connection that was evicted before recovery', async () => { + const authError = new Error('Non-200 status code (401)'); + const request = jest.fn(); + const connection = createConnection(request); + attachOAuthHandler(); + const manager = await createManager(connection); + const lifecycle = manager as unknown as { + disposeEvictedConnection: (connection: MCPConnection, logPrefix: string) => Promise; + }; + request + .mockImplementationOnce(async () => { + await lifecycle.disposeEvictedConnection(connection, '[MCP][evicted]'); + throw authError; + }) + .mockResolvedValueOnce(toolResult); + + await expect(callTool(manager)).resolves.toBeDefined(); + + expect(connection.connect).toHaveBeenCalledTimes(1); + expect(connection.dispose).toHaveBeenCalledTimes(1); + expect(request).toHaveBeenCalledTimes(2); + }); + + it('preserves recovery eviction when another borrower releases last', async () => { + const authError = new Error('Non-200 status code (401)'); + const request = jest.fn(); + const connection = createConnection(request); + attachOAuthHandler(); + const manager = await createManager(connection); + const lifecycle = manager as unknown as { + disposeEvictedConnection: (connection: MCPConnection, logPrefix: string) => Promise; + retainConnection: (connection: MCPConnection) => void; + releaseConnection: (connection: MCPConnection) => Promise; + }; + request + .mockImplementationOnce(async () => { + lifecycle.retainConnection(connection); + await lifecycle.disposeEvictedConnection(connection, '[MCP][evicted]'); + throw authError; + }) + .mockResolvedValueOnce(toolResult); + + const toolPromise = callTool(manager); + await new Promise((resolve) => setImmediate(resolve)); + + expect(connection.connect).not.toHaveBeenCalled(); + expect(connection.dispose).not.toHaveBeenCalled(); + + await lifecycle.releaseConnection(connection); + await expect(toolPromise).resolves.toBeDefined(); + + expect(connection.connect).toHaveBeenCalledTimes(1); + expect(connection.dispose).toHaveBeenCalledTimes(1); + expect(request).toHaveBeenCalledTimes(2); + }); + + it('preserves eviction requested during a recovery lease gap', async () => { + const connection = createConnection(jest.fn()); + const manager = await createManager(connection); + const lifecycle = manager as unknown as { + disposeEvictedConnection: (connection: MCPConnection, logPrefix: string) => Promise; + holdDeferredConnectionDisposal: (connection: MCPConnection) => void; + releaseDeferredConnectionDisposal: (connection: MCPConnection) => Promise; + retainConnection: (connection: MCPConnection) => void; + releaseConnection: (connection: MCPConnection) => Promise; + }; + + lifecycle.holdDeferredConnectionDisposal(connection); + await lifecycle.disposeEvictedConnection(connection, '[MCP][evicted]'); + + expect(connection.dispose).not.toHaveBeenCalled(); + + lifecycle.retainConnection(connection); + await lifecycle.releaseDeferredConnectionDisposal(connection); + await lifecycle.releaseConnection(connection); + + expect(connection.dispose).toHaveBeenCalledTimes(1); + }); + + it('defers final borrower disposal until the recovery hold is released', async () => { + const connection = createConnection(jest.fn()); + const manager = await createManager(connection); + const lifecycle = manager as unknown as { + disposeEvictedConnection: (connection: MCPConnection, logPrefix: string) => Promise; + holdDeferredConnectionDisposal: (connection: MCPConnection) => void; + releaseDeferredConnectionDisposal: (connection: MCPConnection) => Promise; + retainConnection: (connection: MCPConnection) => void; + releaseConnection: (connection: MCPConnection) => Promise; + }; + + lifecycle.holdDeferredConnectionDisposal(connection); + lifecycle.retainConnection(connection); + await lifecycle.disposeEvictedConnection(connection, '[MCP][evicted]'); + await lifecycle.releaseConnection(connection); + + expect(connection.dispose).not.toHaveBeenCalled(); + + await lifecycle.releaseDeferredConnectionDisposal(connection); + + expect(connection.dispose).toHaveBeenCalledTimes(1); + }); + + it('recovers an OAuth failure found by the connection preflight check', async () => { + const authError = new Error('Non-200 status code (401)'); + const request = jest.fn().mockResolvedValueOnce(toolResult); + const connection = createConnection(request); + (connection.isConnected as jest.Mock).mockResolvedValue(false); + (connection.getLastConnectionCheckError as jest.Mock).mockReturnValue(authError); + attachOAuthHandler(); + const manager = await createManager(connection); + + await expect(callTool(manager)).resolves.toBeDefined(); + + expect(connection.connect).toHaveBeenCalledTimes(1); + expect(request).toHaveBeenCalledTimes(1); + }); + + it('recovers when OAuth was detected from the resolved connection config', async () => { + const authError = new Error('Non-200 status code (401)'); + const request = jest.fn().mockRejectedValueOnce(authError).mockResolvedValueOnce(toolResult); + const connection = createConnection(request); + const runtimeConfig: t.ParsedServerConfig = { + type: 'streamable-http', + url: 'https://api.example.com/users/{{LIBRECHAT_USER_ID}}/mcp', + source: 'yaml', + }; + (mockRegistryInstance.getServerConfig as jest.Mock).mockResolvedValue(runtimeConfig); + mockProcessMCPEnv.mockReturnValue({ + ...runtimeConfig, + url: 'https://api.example.com/users/oauth-user/mcp', + }); + attachOAuthHandler(); + const manager = await createManager(connection); + + await expect(callTool(manager)).resolves.toBeDefined(); + + expect(MCPConnectionFactory.attachRequestOAuthHandler).toHaveBeenCalledTimes(1); + expect(connection.connect).toHaveBeenCalledTimes(1); + expect(request).toHaveBeenCalledTimes(2); + }); + + it('keeps the request OAuth handler attached through reconnect-time OAuth', async () => { + const authError = new Error('Non-200 status code (401)'); + const request = jest.fn().mockRejectedValueOnce(authError).mockResolvedValueOnce(toolResult); + const connection = createConnection(request); + (connection.connect as jest.Mock) + .mockImplementationOnce(async () => { + const emitted = connection.emit('oauthReauthenticationRequired', { + serverName, + error: authError, + serverUrl: serverConfig.url, + userId: mockUser.id, + }); + if (!emitted) { + throw new Error('Reconnect OAuth handler missing'); + } + throw new Error('Connection not established'); + }) + .mockResolvedValueOnce(undefined); + const oauthHandler = jest.fn((currentConnection: MCPConnection) => { + currentConnection.emit('oauthHandled'); + }); + attachOAuthHandler(oauthHandler); + const manager = await createManager(connection); + + await expect(callTool(manager)).resolves.toBeDefined(); + + expect(oauthHandler).toHaveBeenCalledTimes(2); + expect(connection.connect).toHaveBeenCalledTimes(2); + expect(request).toHaveBeenCalledTimes(2); + expect(connection.listenerCount('oauthReauthenticationRequired')).toBe(0); + }); + + it('escalates to interactive OAuth when reconnect rejects a silently refreshed token', async () => { + const authError = new Error('Non-200 status code (401)'); + const reconnectError = new Error('Connection not established'); + const request = jest.fn().mockRejectedValueOnce(authError).mockResolvedValueOnce(toolResult); + const connection = createConnection(request); + (connection.connect as jest.Mock) + .mockImplementationOnce(async () => { + connection.emit('oauthHandled', 'silent-refresh'); + throw reconnectError; + }) + .mockResolvedValueOnce(undefined); + const recoveryEvents: Array<{ skipSilentRefresh?: boolean }> = []; + (MCPConnectionFactory.attachRequestOAuthHandler as jest.Mock).mockImplementation( + (_basic, _oauth, currentConnection: MCPConnection) => { + const listener = (data: { skipSilentRefresh?: boolean }) => { + recoveryEvents.push(data); + currentConnection.emit( + 'oauthHandled', + data.skipSilentRefresh ? 'interactive' : 'silent-refresh', + ); + }; + currentConnection.on('oauthReauthenticationRequired', listener); + return () => currentConnection.off('oauthReauthenticationRequired', listener); + }, + ); + const manager = await createManager(connection); + + await expect(callTool(manager)).resolves.toBeDefined(); + + expect(recoveryEvents).toHaveLength(2); + expect(recoveryEvents[0].skipSilentRefresh).toBeUndefined(); + expect(recoveryEvents[1].skipSilentRefresh).toBe(true); + expect(connection.connect).toHaveBeenCalledTimes(2); + expect(request).toHaveBeenCalledTimes(2); + expect(connection.listenerCount('oauthReauthenticationRequired')).toBe(0); + }); + + it('does not start a second recovery when the retried request is rejected', async () => { + const firstError = new Error('Non-200 status code (401)'); + const secondError = new Error('Non-200 status code (401) again'); + const request = jest + .fn() + .mockRejectedValueOnce(firstError) + .mockRejectedValueOnce(secondError); + const connection = createConnection(request); + attachOAuthHandler(); + const manager = await createManager(connection); + + await expect(callTool(manager)).rejects.toBe(secondError); + + expect(connection.connect).toHaveBeenCalledTimes(1); + expect(request).toHaveBeenCalledTimes(2); + }); + + it('returns the original error when OAuth recovery fails', async () => { + const authError = new Error('Non-200 status code (401)'); + const request = jest.fn().mockRejectedValueOnce(authError); + const connection = createConnection(request); + attachOAuthHandler((currentConnection) => { + currentConnection.emit('oauthFailed', new Error('Refresh rejected')); + }); + const manager = await createManager(connection); + + await expect(callTool(manager)).rejects.toBe(authError); + + expect(connection.connect).not.toHaveBeenCalled(); + expect(request).toHaveBeenCalledTimes(1); + }); + + it('does not recover or retry non-authentication failures', async () => { + const toolError = new Error('Tool execution failed'); + const request = jest.fn().mockRejectedValueOnce(toolError); + const connection = createConnection(request); + attachOAuthHandler(); + const manager = await createManager(connection); + + await expect(callTool(manager)).rejects.toBe(toolError); + + expect(MCPConnectionFactory.attachRequestOAuthHandler).not.toHaveBeenCalled(); + expect(connection.connect).not.toHaveBeenCalled(); + expect(request).toHaveBeenCalledTimes(1); + }); + + it('shares one OAuth handler and connection rebuild across concurrent rejected calls', async () => { + const authError = new Error('Non-200 status code (401)'); + const request = jest + .fn() + .mockRejectedValueOnce(authError) + .mockRejectedValueOnce(authError) + .mockResolvedValue(toolResult); + const connection = createConnection(request); + let finishConnect: (() => void) | undefined; + (connection.connect as jest.Mock).mockImplementation( + () => + new Promise((resolve) => { + finishConnect = resolve; + }), + ); + const oauthHandler = jest.fn((currentConnection: MCPConnection) => { + currentConnection.emit('oauthHandled'); + }); + attachOAuthHandler(oauthHandler); + const manager = await createManager(connection); + + const firstCall = callTool(manager); + const secondCall = callTool(manager); + await new Promise((resolve) => setImmediate(resolve)); + finishConnect?.(); + + await expect(Promise.all([firstCall, secondCall])).resolves.toHaveLength(2); + expect(MCPConnectionFactory.attachRequestOAuthHandler).toHaveBeenCalledTimes(1); + expect(oauthHandler).toHaveBeenCalledTimes(1); + expect(connection.connect).toHaveBeenCalledTimes(1); + expect(request).toHaveBeenCalledTimes(4); + expect(connection.listenerCount('oauthReauthenticationRequired')).toBe(0); + }); + + it('lets an aborted runtime waiter leave without cancelling the recovery owner', async () => { + const authError = new Error('Non-200 status code (401)'); + const request = jest.fn().mockRejectedValue(authError); + const connection = createConnection(request); + attachOAuthHandler(() => undefined); + const manager = await createManager(connection); + const controller = new AbortController(); + const abortReason = new Error('waiter aborted'); + + const ownerCall = callTool(manager); + await new Promise((resolve) => setImmediate(resolve)); + const waiterCall = callTool(manager, controller.signal); + await new Promise((resolve) => setImmediate(resolve)); + + expect(MCPConnectionFactory.attachRequestOAuthHandler).toHaveBeenCalledTimes(1); + controller.abort(abortReason); + + await expect(waiterCall).rejects.toBe(abortReason); + expect(connection.listenerCount('oauthReauthenticationRequired')).toBe(1); + + connection.emit('oauthFailed', new Error('Recovery owner aborted')); + + await expect(ownerCall).rejects.toBe(authError); + expect(connection.connect).not.toHaveBeenCalled(); + expect(request).toHaveBeenCalledTimes(1); + expect(connection.listenerCount('oauthReauthenticationRequired')).toBe(0); + }); + + it('lets an aborted recovery owner leave while a live waiter completes recovery', async () => { + const authError = new Error('Non-200 status code (401)'); + const request = jest.fn().mockRejectedValueOnce(authError).mockResolvedValueOnce(toolResult); + const connection = createConnection(request); + attachOAuthHandler(() => undefined); + const manager = await createManager(connection); + const controller = new AbortController(); + const abortReason = new Error('owner aborted'); + + const ownerCall = callTool(manager, controller.signal); + await new Promise((resolve) => setImmediate(resolve)); + const waiterCall = callTool(manager); + await new Promise((resolve) => setImmediate(resolve)); + + controller.abort(abortReason); + + await expect(ownerCall).rejects.toBe(abortReason); + expect(connection.listenerCount('oauthReauthenticationRequired')).toBe(1); + + connection.emit('oauthHandled'); + + await expect(waiterCall).resolves.toBeDefined(); + expect(MCPConnectionFactory.attachRequestOAuthHandler).toHaveBeenCalledTimes(1); + expect( + (MCPConnectionFactory.attachRequestOAuthHandler as jest.Mock).mock.calls[0][1].signal, + ).toBeUndefined(); + expect(connection.connect).toHaveBeenCalledTimes(1); + expect(request).toHaveBeenCalledTimes(2); + expect(connection.listenerCount('oauthReauthenticationRequired')).toBe(0); + }); + + it('relays an active recovery prompt to a joining request', async () => { + const authError = new Error('Non-200 status code (401)'); + const request = jest.fn().mockRejectedValueOnce(authError).mockResolvedValue(toolResult); + const connection = createConnection(request); + const ownerOAuthStart = jest.fn().mockResolvedValue(undefined); + const waiterOAuthStart = jest.fn().mockResolvedValue(undefined); + (MCPConnectionFactory.attachRequestOAuthHandler as jest.Mock).mockImplementation( + (_basic, oauth: { oauthStart?: t.OAuthStartHandler }, currentConnection: MCPConnection) => { + const listener = () => { + void oauth.oauthStart?.('https://auth.example.com/pending'); + }; + currentConnection.on('oauthReauthenticationRequired', listener); + return () => currentConnection.off('oauthReauthenticationRequired', listener); + }, + ); + const manager = await createManager(connection); + + const ownerCall = callTool(manager, undefined, ownerOAuthStart); + await new Promise((resolve) => setImmediate(resolve)); + expect(ownerOAuthStart).toHaveBeenCalledWith('https://auth.example.com/pending', undefined); + + const waiterCall = callTool(manager, undefined, waiterOAuthStart); + await new Promise((resolve) => setImmediate(resolve)); + expect(waiterOAuthStart).toHaveBeenCalledWith('https://auth.example.com/pending', undefined); + + connection.emit('oauthHandled'); + + await expect(Promise.all([ownerCall, waiterCall])).resolves.toHaveLength(2); + expect(MCPConnectionFactory.attachRequestOAuthHandler).toHaveBeenCalledTimes(1); + expect(connection.connect).toHaveBeenCalledTimes(1); + expect(request).toHaveBeenCalledTimes(3); + }); + + it('relays OAuth completion to live waiters without failing recovery on notification errors', async () => { + const authError = new Error('Non-200 status code (401)'); + const request = jest.fn().mockRejectedValueOnce(authError).mockResolvedValue(toolResult); + const connection = createConnection(request); + const ownerOAuthStart = jest.fn().mockResolvedValue(undefined); + const waiterOAuthStart = jest.fn().mockResolvedValue(undefined); + const ownerOAuthEnd = jest.fn().mockRejectedValue(new Error('owner response is stale')); + const waiterOAuthEnd = jest.fn().mockResolvedValue(undefined); + let completeOAuth: (() => void) | undefined; + const completionGate = new Promise((resolve) => { + completeOAuth = resolve; + }); + (MCPConnectionFactory.attachRequestOAuthHandler as jest.Mock).mockImplementation( + ( + _basic, + oauth: { + oauthStart?: t.OAuthStartHandler; + oauthEnd?: () => Promise; + }, + currentConnection: MCPConnection, + ) => { + const listener = async () => { + await oauth.oauthStart?.('https://auth.example.com/pending'); + await completionGate; + try { + await oauth.oauthEnd?.(); + currentConnection.emit('oauthHandled'); + } catch (error) { + currentConnection.emit('oauthFailed', error); + } + }; + currentConnection.on('oauthReauthenticationRequired', listener); + return () => currentConnection.off('oauthReauthenticationRequired', listener); + }, + ); + const manager = await createManager(connection); + + const ownerCall = callTool(manager, undefined, ownerOAuthStart, undefined, ownerOAuthEnd); + await new Promise((resolve) => setImmediate(resolve)); + const waiterCall = callTool(manager, undefined, waiterOAuthStart, undefined, waiterOAuthEnd); + await new Promise((resolve) => setImmediate(resolve)); + + completeOAuth?.(); + + await expect(Promise.all([ownerCall, waiterCall])).resolves.toHaveLength(2); + expect(ownerOAuthEnd).toHaveBeenCalledTimes(1); + expect(waiterOAuthEnd).toHaveBeenCalledTimes(1); + expect(connection.connect).toHaveBeenCalledTimes(1); + expect(request).toHaveBeenCalledTimes(3); + }); + + it('defers request-context cleanup until shared recovery settles', async () => { + const authError = new Error('Non-200 status code (401)'); + const request = jest.fn().mockRejectedValueOnce(authError).mockResolvedValueOnce(toolResult); + const connection = createConnection(request); + attachOAuthHandler(() => undefined); + const manager = await createManager(connection); + const requestContext = createMCPRequestContext(); + requestContext.connections.set(`${mockUser.id}:${serverName}`, connection); + + const toolPromise = callTool(manager, undefined, jest.fn(), requestContext); + await new Promise((resolve) => setImmediate(resolve)); + + await cleanupMCPRequestContext(requestContext); + expect(connection.dispose).not.toHaveBeenCalled(); + expect(requestContext.connections.size).toBe(0); + + connection.emit('oauthHandled'); + + await expect(toolPromise).resolves.toBeDefined(); + expect(connection.connect).toHaveBeenCalledTimes(1); + expect(connection.dispose).toHaveBeenCalledTimes(1); + expect(request).toHaveBeenCalledTimes(2); + }); + + it('retries transient reconnect failures within the recovery budget', async () => { + const authError = new Error('Non-200 status code (401)'); + const transientError = new Error('Connection reset'); + const request = jest.fn().mockRejectedValueOnce(authError).mockResolvedValueOnce(toolResult); + const connection = createConnection(request); + (connection.connect as jest.Mock) + .mockRejectedValueOnce(transientError) + .mockResolvedValueOnce(undefined); + attachOAuthHandler(); + const manager = await createManager(connection); + const retryPolicy = manager as unknown as { + waitForOAuthReconnectRetry: (attempt: number) => Promise; + }; + const waitForRetry = jest + .spyOn(retryPolicy, 'waitForOAuthReconnectRetry') + .mockResolvedValue(undefined); + + await expect(callTool(manager)).resolves.toBeDefined(); + + expect(waitForRetry).toHaveBeenCalledWith(1); + expect(connection.connect).toHaveBeenCalledTimes(2); + expect(request).toHaveBeenCalledTimes(2); + }); + + it('disposes a recovered connection explicitly disconnected during recovery', async () => { + const authError = new Error('Non-200 status code (401)'); + const request = jest.fn().mockRejectedValueOnce(authError).mockResolvedValueOnce(toolResult); + const connection = createConnection(request); + attachOAuthHandler(() => undefined); + const manager = await createManager(connection); + const internals = manager as unknown as { + userConnections: Map>; + }; + internals.userConnections.set(mockUser.id, new Map([[serverName, connection]])); + + const toolPromise = callTool(manager); + await new Promise((resolve) => setImmediate(resolve)); + + await manager.disconnectUserConnection(mockUser.id, serverName); + expect(internals.userConnections.get(mockUser.id)).toBeUndefined(); + expect(connection.dispose).not.toHaveBeenCalled(); + + connection.emit('oauthHandled'); + + await expect(toolPromise).resolves.toBeDefined(); + expect(connection.connect).toHaveBeenCalledTimes(1); + expect(connection.dispose).toHaveBeenCalledTimes(1); + expect(request).toHaveBeenCalledTimes(2); + expect(connection.listenerCount('oauthReauthenticationRequired')).toBe(0); + }); + + it('keeps recovery disposal owned after the request owner aborts', async () => { + const authError = new Error('Non-200 status code (401)'); + const request = jest.fn().mockRejectedValueOnce(authError); + const connection = createConnection(request); + attachOAuthHandler(() => undefined); + const manager = await createManager(connection); + const lifecycle = manager as unknown as { + disposeEvictedConnection: (connection: MCPConnection, logPrefix: string) => Promise; + }; + const controller = new AbortController(); + const abortReason = new Error('request aborted'); + + const ownerCall = callTool(manager, controller.signal); + await new Promise((resolve) => setImmediate(resolve)); + controller.abort(abortReason); + + await expect(ownerCall).rejects.toBe(abortReason); + await lifecycle.disposeEvictedConnection(connection, '[MCP][evicted]'); + expect(connection.dispose).not.toHaveBeenCalled(); + + connection.emit('oauthHandled'); + await new Promise((resolve) => setImmediate(resolve)); + + expect(connection.connect).toHaveBeenCalledTimes(1); + expect(connection.dispose).toHaveBeenCalledTimes(1); + expect(connection.listenerCount('oauthReauthenticationRequired')).toBe(0); + }); + + it('defers ephemeral disposal until an aborted OAuth recovery settles', async () => { + const authError = new Error('Non-200 status code (401)'); + const request = jest.fn().mockRejectedValueOnce(authError); + const connection = createConnection(request); + attachOAuthHandler(() => undefined); + const manager = await createManager(connection); + const ephemeralConfig: t.StreamableHTTPOptions = { + ...serverConfig, + url: 'https://mcp.example.com/{{LIBRECHAT_BODY_CONVERSATIONID}}', + }; + (mockRegistryInstance.getServerConfig as jest.Mock).mockResolvedValue(ephemeralConfig); + mockProcessMCPEnv.mockReturnValue({ + ...ephemeralConfig, + url: 'https://mcp.example.com/conversation-1', + }); + const controller = new AbortController(); + const abortReason = new Error('request aborted'); + + const toolCall = manager.callTool({ + user: mockUser, + serverName, + toolName: 'oauth_tool', + provider: 'openai', + oauthStart: jest.fn(), + flowManager: mockFlowManager, + requestBody: { conversationId: 'conversation-1' }, + options: { signal: controller.signal }, + }); + await new Promise((resolve) => setImmediate(resolve)); + controller.abort(abortReason); + + await expect(toolCall).rejects.toBe(abortReason); + expect(connection.dispose).not.toHaveBeenCalled(); + + connection.emit('oauthHandled'); + await new Promise((resolve) => setImmediate(resolve)); + + expect(connection.connect).toHaveBeenCalledTimes(1); + expect(connection.dispose).toHaveBeenCalledTimes(1); + expect(connection.listenerCount('oauthReauthenticationRequired')).toBe(0); + }); + + it('lets a live waiter recover after the recovery owner fails', async () => { + const authError = new Error('Non-200 status code (401)'); + const request = jest + .fn() + .mockRejectedValueOnce(authError) + .mockRejectedValueOnce(authError) + .mockResolvedValueOnce(toolResult); + const connection = createConnection(request); + let attachmentCount = 0; + (MCPConnectionFactory.attachRequestOAuthHandler as jest.Mock).mockImplementation( + (_basic, _oauth, currentConnection: MCPConnection) => { + const currentAttachment = ++attachmentCount; + const listener = () => { + if (currentAttachment === 2) { + currentConnection.emit('oauthHandled'); + } + }; + currentConnection.on('oauthReauthenticationRequired', listener); + return () => currentConnection.off('oauthReauthenticationRequired', listener); + }, + ); + const manager = await createManager(connection); + + const ownerCall = callTool(manager); + await new Promise((resolve) => setImmediate(resolve)); + const waiterCall = callTool(manager); + await new Promise((resolve) => setImmediate(resolve)); + + expect(MCPConnectionFactory.attachRequestOAuthHandler).toHaveBeenCalledTimes(1); + + connection.emit('oauthFailed', new Error('Recovery owner aborted')); + + await expect(ownerCall).rejects.toBe(authError); + await expect(waiterCall).resolves.toBeDefined(); + expect(MCPConnectionFactory.attachRequestOAuthHandler).toHaveBeenCalledTimes(2); + expect(connection.connect).toHaveBeenCalledTimes(1); + expect(request).toHaveBeenCalledTimes(3); + expect(connection.listenerCount('oauthReauthenticationRequired')).toBe(0); + }); + + it('restarts checkout when a post-request recovery waiter claims takeover', async () => { + const authError = new Error('Non-200 status code (401)'); + const staleRequest = jest.fn().mockRejectedValue(authError); + const recoveredRequest = jest.fn().mockResolvedValue(toolResult); + const staleConnection = createConnection(staleRequest); + const recoveredConnection = createConnection(recoveredRequest); + attachOAuthHandler(() => undefined); + const manager = await createManager(staleConnection); + let currentConnection = staleConnection; + (manager.getConnection as jest.Mock).mockImplementation(async () => currentConnection); + + const ownerCall = callTool(manager); + const waiterCall = callTool(manager); + await new Promise((resolve) => setImmediate(resolve)); + + expect(MCPConnectionFactory.attachRequestOAuthHandler).toHaveBeenCalledTimes(1); + currentConnection = recoveredConnection; + staleConnection.emit('oauthFailed', new Error('Recovery owner failed')); + + await expect(ownerCall).rejects.toBe(authError); + await expect(waiterCall).resolves.toBeDefined(); + expect(staleRequest).toHaveBeenCalledTimes(2); + expect(recoveredRequest).toHaveBeenCalledTimes(1); + expect(staleConnection.connect).not.toHaveBeenCalled(); + expect(recoveredConnection.connect).not.toHaveBeenCalled(); + expect(staleConnection.listenerCount('oauthReauthenticationRequired')).toBe(0); + }); + + it('caps waiter takeover after a shared recovery keeps failing', async () => { + const authError = new Error('Non-200 status code (401)'); + const request = jest.fn().mockRejectedValue(authError); + const connection = createConnection(request); + let attachmentCount = 0; + (MCPConnectionFactory.attachRequestOAuthHandler as jest.Mock).mockImplementation( + (_basic, _oauth, currentConnection: MCPConnection) => { + attachmentCount += 1; + const listener = () => { + currentConnection.emit('oauthFailed', new Error('OAuth denied')); + }; + currentConnection.on('oauthReauthenticationRequired', listener); + return () => currentConnection.off('oauthReauthenticationRequired', listener); + }, + ); + const manager = await createManager(connection); + + const calls = [callTool(manager), callTool(manager), callTool(manager), callTool(manager)]; + const results = await Promise.allSettled(calls); + + expect(results.every((result) => result.status === 'rejected')).toBe(true); + expect(attachmentCount).toBe(2); + expect(connection.connect).not.toHaveBeenCalled(); + expect(connection.listenerCount('oauthReauthenticationRequired')).toBe(0); + }); + + it('allows only one checkout waiter to claim a failed recovery takeover', async () => { + const authError = new Error('Non-200 status code (401)'); + const request = jest.fn().mockRejectedValue(authError); + const connection = createConnection(request); + attachOAuthHandler((currentConnection) => { + currentConnection.emit('oauthFailed', new Error('OAuth denied')); + }); + const manager = await createManager(connection); + let rejectRecovery: ((error: Error) => void) | undefined; + const recovery = new Promise((_resolve, reject) => { + rejectRecovery = reject; + }); + const internals = manager as unknown as { + oauthRecoveries: WeakMap< + MCPConnection, + { promise: Promise; allowsTakeover: boolean; takeoverClaimed?: boolean } + >; + }; + internals.oauthRecoveries.set(connection, { + promise: recovery, + allowsTakeover: true, + }); + + const calls = [callTool(manager), callTool(manager), callTool(manager), callTool(manager)]; + await new Promise((resolve) => setImmediate(resolve)); + internals.oauthRecoveries.delete(connection); + rejectRecovery?.(new Error('Shared recovery failed')); + const results = await Promise.allSettled(calls); + + expect(results.every((result) => result.status === 'rejected')).toBe(true); + expect(request).toHaveBeenCalledTimes(1); + expect(MCPConnectionFactory.attachRequestOAuthHandler).toHaveBeenCalledTimes(1); + expect(connection.listenerCount('oauthReauthenticationRequired')).toBe(0); + }); + }); + + describe('getUserConnection - recovery lifecycle', () => { + const mockUser = { id: 'recovery-user' } as IUser; + const serverConfig: t.StreamableHTTPOptions = { + type: 'streamable-http', + url: 'https://mcp.example.com', + requiresOAuth: false, + }; + + it('waits for an active recovery before reusing a cached connection', async () => { + mockAppConnections({ + has: jest.fn().mockResolvedValue(false), + }); + const manager = await MCPManager.createInstance(newMCPServersConfig()); + const connection = { + isConnected: jest.fn().mockResolvedValue(true), + } as unknown as MCPConnection; + let resolveRecovery: (() => void) | undefined; + const recovery = new Promise((resolve) => { + resolveRecovery = resolve; + }); + const internals = manager as unknown as { + oauthRecoveries: WeakMap< + MCPConnection, + { promise: Promise; allowsTakeover: boolean } + >; + userConnections: Map>; + }; + internals.userConnections.set(mockUser.id, new Map([[serverName, connection]])); + internals.oauthRecoveries.set(connection, { promise: recovery, allowsTakeover: true }); + + const connectionPromise = manager.getUserConnection({ + serverName, + user: mockUser, + serverConfig, + }); + await new Promise((resolve) => setImmediate(resolve)); + + expect(connection.isConnected).not.toHaveBeenCalled(); + + resolveRecovery?.(); + + await expect(connectionPromise).resolves.toBe(connection); + expect(connection.isConnected).toHaveBeenCalledTimes(1); + expect(MCPConnectionFactory.create).not.toHaveBeenCalled(); + }); + + it('rechecks recovery after yielding to load connection configuration', async () => { + mockAppConnections({ + has: jest.fn().mockResolvedValue(false), + }); + let resolveConfig: ((config: t.ParsedServerConfig) => void) | undefined; + (mockRegistryInstance.getServerConfig as jest.Mock).mockReturnValue( + new Promise((resolve) => { + resolveConfig = resolve; + }), + ); + const manager = await MCPManager.createInstance(newMCPServersConfig()); + let recovered = false; + const connection = { + disconnect: jest.fn().mockResolvedValue(undefined), + isConnected: jest.fn().mockImplementation(async () => recovered), + isStale: jest.fn().mockReturnValue(false), + } as unknown as MCPConnection; + let resolveRecovery: (() => void) | undefined; + const recovery = new Promise((resolve) => { + resolveRecovery = resolve; + }); + const internals = manager as unknown as { + oauthRecoveries: WeakMap< + MCPConnection, + { promise: Promise; allowsTakeover: boolean } + >; + userConnections: Map>; + }; + internals.userConnections.set(mockUser.id, new Map([[serverName, connection]])); + + const connectionPromise = manager.getUserConnection({ + serverName, + user: mockUser, + }); + await new Promise((resolve) => setImmediate(resolve)); + + internals.oauthRecoveries.set(connection, { promise: recovery, allowsTakeover: true }); + resolveConfig?.(serverConfig); + await new Promise((resolve) => setImmediate(resolve)); + + expect(connection.disconnect).not.toHaveBeenCalled(); + expect(MCPConnectionFactory.create).not.toHaveBeenCalled(); + + recovered = true; + resolveRecovery?.(); + + await expect(connectionPromise).resolves.toBe(connection); + expect(connection.isConnected).toHaveBeenCalledTimes(1); + expect(connection.disconnect).not.toHaveBeenCalled(); + expect(MCPConnectionFactory.create).not.toHaveBeenCalled(); + }); + + it('replaces stale configuration without waiting for obsolete recovery', async () => { + mockAppConnections({ + has: jest.fn().mockResolvedValue(false), + }); + const manager = await MCPManager.createInstance(newMCPServersConfig()); + const staleConnection = { + disconnect: jest.fn().mockResolvedValue(undefined), + dispose: jest.fn().mockResolvedValue(undefined), + isConnected: jest.fn().mockResolvedValue(false), + isStale: jest.fn().mockReturnValue(true), + } as unknown as MCPConnection; + const replacementConnection = { + dispose: jest.fn().mockResolvedValue(undefined), + isConnected: jest.fn().mockResolvedValue(true), + on: jest.fn(), + refreshToolList: jest.fn().mockResolvedValue(undefined), + } as unknown as MCPConnection; + const recovery = new Promise(() => undefined); + const internals = manager as unknown as { + oauthRecoveries: WeakMap< + MCPConnection, + { promise: Promise; allowsTakeover: boolean } + >; + retainConnection: (connection: MCPConnection) => void; + releaseConnection: (connection: MCPConnection) => Promise; + userConnections: Map>; + }; + internals.userConnections.set(mockUser.id, new Map([[serverName, staleConnection]])); + internals.retainConnection(staleConnection); + internals.oauthRecoveries.set(staleConnection, { + promise: recovery, + allowsTakeover: true, + }); + (MCPConnectionFactory.create as jest.Mock).mockResolvedValue(replacementConnection); + const updatedConfig = { ...serverConfig, updatedAt: Date.now() }; + + await expect( + manager.getUserConnection({ + serverName, + user: mockUser, + serverConfig: updatedConfig, + }), + ).resolves.toBe(replacementConnection); + + expect(staleConnection.dispose).not.toHaveBeenCalled(); + expect(MCPConnectionFactory.create).toHaveBeenCalledTimes(1); + + await internals.releaseConnection(staleConnection); + expect(staleConnection.dispose).toHaveBeenCalledTimes(1); + }); + + it('lets an aborted waiter leave without cancelling the shared recovery', async () => { + mockAppConnections({ + has: jest.fn().mockResolvedValue(false), + }); + const manager = await MCPManager.createInstance(newMCPServersConfig()); + const connection = { + isConnected: jest.fn().mockResolvedValue(true), + } as unknown as MCPConnection; + const recovery = new Promise(() => undefined); + const internals = manager as unknown as { + oauthRecoveries: WeakMap< + MCPConnection, + { promise: Promise; allowsTakeover: boolean } + >; + userConnections: Map>; + }; + internals.userConnections.set(mockUser.id, new Map([[serverName, connection]])); + internals.oauthRecoveries.set(connection, { promise: recovery, allowsTakeover: true }); + const controller = new AbortController(); + const abortReason = new Error('request aborted'); + + const connectionPromise = manager.getUserConnection({ + serverName, + user: mockUser, + serverConfig, + signal: controller.signal, + }); + controller.abort(abortReason); + + await expect(connectionPromise).rejects.toBe(abortReason); + expect(internals.oauthRecoveries.get(connection)?.promise).toBe(recovery); + expect(connection.isConnected).not.toHaveBeenCalled(); + }); + + it('evicts an inactive cached connection without disconnecting current borrowers', async () => { + const unusableConnection = { + disconnect: jest.fn().mockResolvedValue(undefined), + dispose: jest.fn().mockResolvedValue(undefined), + isConnected: jest.fn().mockResolvedValue(false), + } as unknown as MCPConnection; + const replacementConnection = { + dispose: jest.fn().mockResolvedValue(undefined), + isConnected: jest.fn().mockResolvedValue(true), + on: jest.fn(), + refreshToolList: jest.fn().mockResolvedValue(undefined), + } as unknown as MCPConnection; + mockAppConnections({ + has: jest.fn().mockResolvedValue(false), + }); + (MCPConnectionFactory.create as jest.Mock).mockResolvedValue(replacementConnection); + + const manager = await MCPManager.createInstance(newMCPServersConfig()); + const lifecycle = manager as unknown as { + retainConnection: (connection: MCPConnection) => void; + releaseConnection: (connection: MCPConnection) => Promise; + }; + ( + manager as unknown as { userConnections: Map> } + ).userConnections.set(mockUser.id, new Map([[serverName, unusableConnection]])); + lifecycle.retainConnection(unusableConnection); + + await expect( + manager.getUserConnection({ + serverName, + user: mockUser, + serverConfig, + }), + ).resolves.toBe(replacementConnection); + + expect(unusableConnection.dispose).not.toHaveBeenCalled(); + await lifecycle.releaseConnection(unusableConnection); + expect(unusableConnection.dispose).toHaveBeenCalledTimes(1); + }); + + it('disposes an inactive cached connection immediately when it has no borrowers', async () => { + const unusableConnection = { + disconnect: jest.fn().mockResolvedValue(undefined), + dispose: jest.fn().mockResolvedValue(undefined), + isConnected: jest.fn().mockResolvedValue(false), + } as unknown as MCPConnection; + const replacementConnection = { + dispose: jest.fn().mockResolvedValue(undefined), + isConnected: jest.fn().mockResolvedValue(true), + on: jest.fn(), + refreshToolList: jest.fn().mockResolvedValue(undefined), + } as unknown as MCPConnection; + mockAppConnections({ + has: jest.fn().mockResolvedValue(false), + }); + (MCPConnectionFactory.create as jest.Mock).mockResolvedValue(replacementConnection); + + const manager = await MCPManager.createInstance(newMCPServersConfig()); + ( + manager as unknown as { userConnections: Map> } + ).userConnections.set(mockUser.id, new Map([[serverName, unusableConnection]])); + + await expect( + manager.getUserConnection({ + serverName, + user: mockUser, + serverConfig, + }), + ).resolves.toBe(replacementConnection); + + expect(unusableConnection.dispose).toHaveBeenCalledTimes(1); + }); + + it('evicts an inactive request connection without disconnecting current borrowers', async () => { + const requestConnection = { + disconnect: jest.fn().mockResolvedValue(undefined), + dispose: jest.fn().mockResolvedValue(undefined), + isConnected: jest.fn().mockResolvedValue(false), + } as unknown as MCPConnection; + const replacementConnection = { + dispose: jest.fn().mockResolvedValue(undefined), + isConnected: jest.fn().mockResolvedValue(true), + on: jest.fn(), + refreshToolList: jest.fn().mockResolvedValue(undefined), + } as unknown as MCPConnection; + const requestScopedConnections: t.RequestScopedMCPConnectionStore = { + connections: new Map([[`${mockUser.id}:${serverName}`, requestConnection]]), + pending: new Map(), + }; + const bodyScopedConfig: t.ParsedServerConfig = { + type: 'streamable-http', + url: 'https://mcp.example.com/{{LIBRECHAT_BODY_MESSAGEID}}', + source: 'yaml', + requiresOAuth: false, + }; + mockAppConnections({ + has: jest.fn().mockResolvedValue(false), + }); + mockProcessMCPEnv.mockImplementation(({ options, body }) => ({ + ...options, + ...('url' in options && { + url: options.url?.replace('{{LIBRECHAT_BODY_MESSAGEID}}', body?.messageId ?? ''), + }), + })); + (MCPConnectionFactory.create as jest.Mock).mockResolvedValue(replacementConnection); + + const manager = await MCPManager.createInstance(newMCPServersConfig()); + const lifecycle = manager as unknown as { + retainConnection: (connection: MCPConnection) => void; + releaseConnection: (connection: MCPConnection) => Promise; + }; + lifecycle.retainConnection(requestConnection); + + await expect( + manager.getUserConnection({ + serverName, + user: mockUser, + serverConfig: bodyScopedConfig, + requestBody: { messageId: 'message-1' }, + requestScopedConnections, + }), + ).resolves.toBe(replacementConnection); + + expect(requestConnection.dispose).not.toHaveBeenCalled(); + await lifecycle.releaseConnection(requestConnection); + expect(requestConnection.dispose).toHaveBeenCalledTimes(1); + }); + + it('propagates a failed recovery before checking the cached connection', async () => { + mockAppConnections({ + has: jest.fn().mockResolvedValue(false), + }); + const manager = await MCPManager.createInstance(newMCPServersConfig()); + const connection = { + isConnected: jest.fn().mockResolvedValue(true), + } as unknown as MCPConnection; + const recoveryError = new Error('OAuth denied'); + const recovery = Promise.reject(recoveryError); + const internals = manager as unknown as { + oauthRecoveries: WeakMap< + MCPConnection, + { promise: Promise; allowsTakeover: boolean } + >; + userConnections: Map>; + }; + internals.userConnections.set(mockUser.id, new Map([[serverName, connection]])); + internals.oauthRecoveries.set(connection, { promise: recovery, allowsTakeover: true }); + + await expect( + manager.getUserConnection({ + serverName, + user: mockUser, + serverConfig, + }), + ).rejects.toBe(recoveryError); + + expect(connection.isConnected).not.toHaveBeenCalled(); + expect(MCPConnectionFactory.create).not.toHaveBeenCalled(); + }); + }); + describe('callTool - OBO Integration', () => { const mockUser: Partial = { id: 'user-123', diff --git a/packages/api/src/mcp/__tests__/MCPOAuthRaceCondition.test.ts b/packages/api/src/mcp/__tests__/MCPOAuthRaceCondition.test.ts index 99fc8b2e13..e84731f25b 100644 --- a/packages/api/src/mcp/__tests__/MCPOAuthRaceCondition.test.ts +++ b/packages/api/src/mcp/__tests__/MCPOAuthRaceCondition.test.ts @@ -14,6 +14,7 @@ import type { OAuthTestServer } from './helpers/oauthTestServer'; import type { MCPOAuthTokens } from '~/mcp/oauth'; import { MCPTokenStorage, MCPOAuthHandler, ReauthenticationRequiredError } from '~/mcp/oauth'; import { MockKeyv, createOAuthMCPServer } from './helpers/oauthTestServer'; +import { OAuthLifecycleRelay } from '~/mcp/oauth/pending'; import { FlowStateManager } from '~/flow/manager'; jest.mock('@librechat/data-schemas', () => ({ @@ -53,6 +54,44 @@ describe('MCP OAuth Race Condition Fixes', () => { }); describe('Fix 1: Connection mutex coalesces concurrent attempts', () => { + it('does not overwrite a newer prompt while inspecting stored flow state', async () => { + const ownerOAuthStart = jest.fn().mockResolvedValue(undefined); + const waiterOAuthStart = jest.fn().mockResolvedValue(undefined); + let resolveFlow: ((flow: object) => void) | undefined; + const flowManager = { + getFlowState: jest.fn( + () => + new Promise((resolve) => { + resolveFlow = resolve; + }), + ), + }; + const relay = new OAuthLifecycleRelay({ + oauthStart: ownerOAuthStart, + logPrefix: '[MCP][test]', + }); + + await relay.start('https://auth.example.com/old'); + const addWaiter = relay.add({ + oauthStart: waiterOAuthStart, + flowManager: flowManager as never, + userId: 'user-1', + serverName: 'test-server', + }); + + expect(flowManager.getFlowState).toHaveBeenCalledTimes(1); + await relay.start('https://auth.example.com/new'); + resolveFlow?.({ + createdAt: Date.now(), + metadata: { authorizationUrl: 'https://auth.example.com/old' }, + status: 'PENDING', + }); + await addWaiter; + + expect(waiterOAuthStart).toHaveBeenCalledTimes(1); + expect(waiterOAuthStart).toHaveBeenCalledWith('https://auth.example.com/new', undefined); + }); + it('should return the same pending promise for concurrent getUserConnection calls', async () => { const { UserConnectionManager } = await import('~/mcp/UserConnectionManager'); @@ -283,6 +322,9 @@ describe('MCP OAuth Race Condition Fixes', () => { await oauthOptions.oauthStart?.(authorizationUrl); } await connectionReleased; + if (oauthOptions && 'oauthEnd' in oauthOptions) { + await oauthOptions.oauthEnd?.(); + } return mockConnection as never; }); @@ -297,11 +339,13 @@ describe('MCP OAuth Race Condition Fixes', () => { await flowManager.initFlow(`${user.id}:${serverName}`, 'mcp_oauth', { authorizationUrl }); const firstOAuthStart = jest.fn().mockResolvedValue(undefined); + const firstOAuthEnd = jest.fn().mockRejectedValue(new Error('owner response is stale')); const firstConnection = manager.getUserConnection({ serverName, user: user as never, flowManager: flowManager as never, oauthStart: firstOAuthStart, + oauthEnd: firstOAuthEnd, }); for (let i = 0; i < 20 && firstOAuthStart.mock.calls.length === 0; i++) { await new Promise((resolve) => setTimeout(resolve, 5)); @@ -309,21 +353,29 @@ describe('MCP OAuth Race Condition Fixes', () => { expect(firstOAuthStart).toHaveBeenCalledWith(authorizationUrl, undefined); const joinedOAuthStart = jest.fn().mockResolvedValue(undefined); + const joinedOAuthEnd = jest.fn().mockResolvedValue(undefined); const joinedConnection = manager.getUserConnection({ serverName, user: user as never, flowManager: flowManager as never, oauthStart: joinedOAuthStart, + oauthEnd: joinedOAuthEnd, }); + for (let i = 0; i < 20 && joinedOAuthStart.mock.calls.length === 0; i++) { + await new Promise((resolve) => setTimeout(resolve, 5)); + } + expect(joinedOAuthStart).toHaveBeenCalledWith( + authorizationUrl, + expect.objectContaining({ expiresAt: expect.any(Number) }), + ); + releaseConnection(); const [conn1, conn2] = await Promise.all([firstConnection, joinedConnection]); expect(conn1).toBe(conn2); - expect(joinedOAuthStart).toHaveBeenCalledWith( - authorizationUrl, - expect.objectContaining({ expiresAt: expect.any(Number) }), - ); + expect(firstOAuthEnd).toHaveBeenCalledTimes(1); + expect(joinedOAuthEnd).toHaveBeenCalledTimes(1); expect(createSpy).toHaveBeenCalledTimes(1); } finally { releaseConnection(); diff --git a/packages/api/src/mcp/__tests__/helpers/oauthTestServer.ts b/packages/api/src/mcp/__tests__/helpers/oauthTestServer.ts index 2b01c95716..22f10f4016 100644 --- a/packages/api/src/mcp/__tests__/helpers/oauthTestServer.ts +++ b/packages/api/src/mcp/__tests__/helpers/oauthTestServer.ts @@ -70,6 +70,10 @@ export interface OAuthTestServerOptions { scopesSupported?: string[]; /** When true, /authorize and /token reject requests that omit the MCP resource parameter. */ requireResourceParameter?: boolean; + /** Number of refresh-grant access tokens the MCP resource should reject after issuance. */ + rejectRefreshTokens?: number; + /** Optional test hook for controlling echo-tool completion. */ + echoHandler?: (message: string) => string | Promise; } export interface OAuthTokenRequestRecord { @@ -136,6 +140,8 @@ export async function createOAuthMCPServer( requiredScopes = [], scopesSupported = [...new Set([...tokenScopes, ...requiredScopes])], requireResourceParameter = false, + rejectRefreshTokens = 0, + echoHandler, } = options; const sessions = new Map(); @@ -157,6 +163,7 @@ export async function createOAuthMCPServer( } >(); const registeredClients = new Map(); + let rejectedRefreshTokensRemaining = rejectRefreshTokens; let port = 0; const getBaseUrl = () => `http://127.0.0.1:${port}`; @@ -410,7 +417,11 @@ export async function createOAuthMCPServer( const scopes = params.has('scope') ? parseScopes(params.get('scope')) : (refreshTokenScopes.get(refreshToken) ?? tokenScopes); - issuedTokens.add(newAccessToken); + if (rejectedRefreshTokensRemaining > 0) { + rejectedRefreshTokensRemaining -= 1; + } else { + issuedTokens.add(newAccessToken); + } tokenIssueTimes.set(newAccessToken, Date.now()); accessTokenScopes.set(newAccessToken, scopes); @@ -475,9 +486,10 @@ export async function createOAuthMCPServer( sessionIdGenerator: () => randomUUID(), }); const mcp = new McpServer({ name: 'oauth-test-server', version: '0.0.1' }); - mcp.tool('echo', { message: z.string() }, async (args) => ({ - content: [{ type: 'text' as const, text: `echo: ${args.message}` }], - })); + mcp.tool('echo', { message: z.string() }, async (args) => { + const text = echoHandler ? await echoHandler(args.message) : `echo: ${args.message}`; + return { content: [{ type: 'text' as const, text }] }; + }); await mcp.connect(transport); } diff --git a/packages/api/src/mcp/__tests__/request.test.ts b/packages/api/src/mcp/__tests__/request.test.ts index 218c69f011..e58c972134 100644 --- a/packages/api/src/mcp/__tests__/request.test.ts +++ b/packages/api/src/mcp/__tests__/request.test.ts @@ -66,17 +66,42 @@ describe('MCP request context', () => { const res = createResponse(); const context = getMCPRequestContext(req, res); const disconnect = jest.fn().mockResolvedValue(undefined); + const dispose = jest.fn().mockResolvedValue(undefined); const pendingDisconnect = jest.fn().mockResolvedValue(undefined); + const pendingDispose = jest.fn().mockResolvedValue(undefined); - context?.connections.set('server', { disconnect }); - context?.pending.set('pending-server', Promise.resolve({ disconnect: pendingDisconnect })); + context?.connections.set('server', { disconnect, dispose }); + context?.pending.set( + 'pending-server', + Promise.resolve({ disconnect: pendingDisconnect, dispose: pendingDispose }), + ); res.emit('finish'); await nextTick(); - expect(disconnect).toHaveBeenCalledTimes(1); - expect(pendingDisconnect).toHaveBeenCalledTimes(1); + expect(dispose).toHaveBeenCalledTimes(1); + expect(pendingDispose).toHaveBeenCalledTimes(1); + expect(disconnect).not.toHaveBeenCalled(); + expect(pendingDisconnect).not.toHaveBeenCalled(); expect(context?.connections.size).toBe(0); expect(context?.pending.size).toBe(0); }); + + it('uses the lifecycle disposer supplied by the connection manager', async () => { + const req = {}; + const res = createResponse(); + const context = getMCPRequestContext(req, res); + const connection = { disconnect: jest.fn().mockResolvedValue(undefined) }; + const disposeConnection = jest.fn().mockResolvedValue(undefined); + if (context) { + context.disposeConnection = disposeConnection; + context.connections.set('user:server', connection); + } + + res.emit('close'); + await nextTick(); + + expect(disposeConnection).toHaveBeenCalledWith('user:server', connection); + expect(connection.disconnect).not.toHaveBeenCalled(); + }); }); diff --git a/packages/api/src/mcp/connection.ts b/packages/api/src/mcp/connection.ts index 65c3d74fd0..e186e481f6 100644 --- a/packages/api/src/mcp/connection.ts +++ b/packages/api/src/mcp/connection.ts @@ -26,6 +26,7 @@ import type * as t from './types'; import { createSSRFSafeUndiciConnect, isSSRFTarget, resolveHostnameSSRF } from '~/auth'; import { reserveMCPToolsChangedRevision } from './toolsChanged'; import { isOAuthServer, sanitizeUrlForLogging } from './utils'; +import { isOAuthAuthenticationError } from './errors'; import { runOutsideTracing } from '~/utils/tracing'; import { isAddressAllowed } from '~/auth/domain'; import { withTimeout } from '~/utils/promise'; @@ -1151,6 +1152,7 @@ export class MCPConnection extends EventEmitter { private readonly userId?: string; private lastPingTime: number; private lastConnectionCheckAt: number = 0; + private lastConnectionCheckError?: unknown; private oauthTokens?: MCPOAuthTokens | null; private requestHeaders?: Record | null; private oauthRequired = false; @@ -1777,6 +1779,7 @@ export class MCPConnection extends EventEmitter { this.on('connectionChange', (state: t.ConnectionState) => { this.connectionState = state; if (state === 'connected') { + this.lastConnectionCheckError = undefined; const isReconnect = this.hasConnected; this.hasConnected = true; this.toolListRefreshSuspended = false; @@ -2114,7 +2117,7 @@ export class MCPConnection extends EventEmitter { } // Check if it's an OAuth authentication error - if (this.isOAuthError(error)) { + if (isOAuthAuthenticationError(error)) { logger.warn(`${this.getLogPrefix()} OAuth authentication required`); this.oauthRequired = true; const serverUrl = this.url; @@ -2310,9 +2313,12 @@ export class MCPConnection extends EventEmitter { } // Check if it's an OAuth authentication error - if (this.isOAuthError(error)) { + if (isOAuthAuthenticationError(error)) { logger.warn(`${this.getLogPrefix()} OAuth authentication error detected`); + this.lastConnectionCheckError = error; + this.connectionState = 'error'; this.emit('oauthError', error); + return; } /** @@ -2644,6 +2650,7 @@ export class MCPConnection extends EventEmitter { return true; } this.lastConnectionCheckAt = now; + this.lastConnectionCheckError = undefined; try { // Try ping first as it's the lightest check @@ -2660,6 +2667,7 @@ export class MCPConnection extends EventEmitter { (error as Error)?.message.includes('method not found')); if (!pingUnsupported) { + this.lastConnectionCheckError = error; logger.error(`${this.getLogPrefix()} Ping failed:`, error); return false; } @@ -2692,6 +2700,7 @@ export class MCPConnection extends EventEmitter { } } catch (capabilityError) { // If capability check fails, the connection is likely broken + this.lastConnectionCheckError = capabilityError; logger.error(`${this.getLogPrefix()} Connection verification failed:`, capabilityError); return false; } @@ -2707,6 +2716,14 @@ export class MCPConnection extends EventEmitter { return isOAuthServer(this.options); } + public isOAuthAuthenticationError(error: unknown): boolean { + return isOAuthAuthenticationError(error); + } + + public getLastConnectionCheckError(): unknown { + return this.lastConnectionCheckError; + } + /** * Check if this connection is stale compared to config update time. * A connection is stale if it was created before the config was updated. @@ -2718,47 +2735,6 @@ export class MCPConnection extends EventEmitter { return this.createdAt < configUpdatedAt; } - private isOAuthError(error: unknown): boolean { - if (!error || typeof error !== 'object') { - return false; - } - - // Check for error code - if ('code' in error) { - const code = (error as { code?: number }).code; - if (code === 401 || code === 403) { - return true; - } - } - - // Check message for various auth error indicators - if ('message' in error && typeof error.message === 'string') { - const message = error.message.toLowerCase(); - // Check for 401 status - if (message.includes('401') || message.includes('non-200 status code (401)')) { - return true; - } - // Check for invalid_token (OAuth servers return this for expired/revoked tokens) - if (message.includes('invalid_token')) { - return true; - } - // Check for invalid_grant (OAuth servers return this for expired/revoked grants) - if (message.includes('invalid_grant')) { - return true; - } - // Check for authentication required - if (message.includes('authentication required') || message.includes('unauthorized')) { - return true; - } - // Check for missing authorization values (e.g., Amazon Ads MCP returns HTTP 400 with this) - if (message.includes('no authorization')) { - return true; - } - } - - return false; - } - /** * Checks if an error indicates rate limiting (HTTP 429). * Rate limited requests should stop reconnection attempts to avoid making the situation worse. diff --git a/packages/api/src/mcp/errors.ts b/packages/api/src/mcp/errors.ts index 6e3eaa919d..f503e8dbb5 100644 --- a/packages/api/src/mcp/errors.ts +++ b/packages/api/src/mcp/errors.ts @@ -10,6 +10,47 @@ export const MCPErrorCodes = { export type MCPErrorCode = (typeof MCPErrorCodes)[keyof typeof MCPErrorCodes]; +interface OAuthErrorLike { + code?: number; + status?: number; + statusCode?: number; + message?: string; +} + +const OAUTH_HTTP_STATUS_PATTERN = + /(?:\bhttp\s+(?:401|403)\b|\bnon-2\d\d\s+status\s+code\s*\((?:401|403)\)|^(?:error:\s*)?(?:401|403)\b|\bunauthorized\s*\(\s*401\s*\)|\bforbidden\s*\(\s*403\s*\))/i; +const MISSING_AUTHORIZATION_PATTERN = /\bno authorization (?:headers?|values?)\b/i; + +/** Detects HTTP authentication failures and OAuth protocol errors without matching unrelated IDs. */ +export function isOAuthAuthenticationError(error: unknown): boolean { + if (!error || typeof error !== 'object') { + return false; + } + + const candidate = error as OAuthErrorLike; + if ( + [candidate.status, candidate.statusCode, candidate.code].some( + (status) => status === 401 || status === 403, + ) + ) { + return true; + } + + if (typeof candidate.message !== 'string') { + return false; + } + + const message = candidate.message.toLowerCase(); + return ( + OAUTH_HTTP_STATUS_PATTERN.test(message) || + message.includes('invalid_token') || + message.includes('invalid_grant') || + message.includes('insufficient_scope') || + message.includes('authentication required') || + MISSING_AUTHORIZATION_PATTERN.test(message) + ); +} + /** * Custom error for MCP domain restriction violations. * Thrown when a user attempts to connect to an MCP server whose domain is not in the allowlist. diff --git a/packages/api/src/mcp/oauth/pending.ts b/packages/api/src/mcp/oauth/pending.ts index 76935ca7ef..df12bcec44 100644 --- a/packages/api/src/mcp/oauth/pending.ts +++ b/packages/api/src/mcp/oauth/pending.ts @@ -68,3 +68,142 @@ export async function getReplayablePendingMCPOAuthStart({ return undefined; } } + +type OAuthEndHandler = () => Promise; + +export class OAuthLifecycleRelay { + private readonly oauthStarts = new Set(); + private readonly oauthEnds = new Set(); + private readonly emittedAuthUrls = new WeakMap(); + private readonly emittedOAuthEnds = new WeakSet(); + private lastOAuthStart?: PendingOAuthStart; + private oauthEnded = false; + + constructor({ + oauthStart, + oauthEnd, + logPrefix, + }: { + oauthStart?: t.OAuthStartHandler; + oauthEnd?: OAuthEndHandler; + logPrefix: string; + }) { + this.logPrefix = logPrefix; + if (oauthStart) { + this.oauthStarts.add(oauthStart); + } + if (oauthEnd) { + this.oauthEnds.add(oauthEnd); + } + } + + private readonly logPrefix: string; + + public readonly start: t.OAuthStartHandler = async (authURL, options) => { + this.lastOAuthStart = { authURL, options }; + const errors: unknown[] = []; + let delivered = false; + + for (const oauthStart of Array.from(this.oauthStarts)) { + try { + await this.emit(oauthStart, authURL, options); + delivered = true; + } catch (error) { + errors.push(error); + logger.warn(`${this.logPrefix} Failed to notify OAuth prompt listener`, error); + } + } + + if (!delivered && errors.length > 0) { + throw errors[0]; + } + }; + + /** Completion notifications are best effort and cannot invalidate received OAuth tokens. */ + public readonly end: OAuthEndHandler = async () => { + this.oauthEnded = true; + for (const oauthEnd of Array.from(this.oauthEnds)) { + try { + await this.emitEnd(oauthEnd); + } catch (error) { + logger.warn(`${this.logPrefix} Failed to notify OAuth completion listener`, error); + } + } + }; + + public async add({ + oauthStart, + oauthEnd, + flowManager, + userId, + serverName, + }: ReplayablePendingMCPOAuthStartOptions & { + oauthStart?: t.OAuthStartHandler; + oauthEnd?: OAuthEndHandler; + }): Promise { + if (oauthStart) { + this.oauthStarts.add(oauthStart); + } + if (oauthEnd) { + this.oauthEnds.add(oauthEnd); + } + if (this.oauthEnded) { + if (oauthEnd) { + try { + await this.emitEnd(oauthEnd); + } catch (error) { + logger.warn(`${this.logPrefix} Failed to re-issue OAuth completion`, error); + } + } + return; + } + + if (!oauthStart) { + return; + } + const lastOAuthStart = this.lastOAuthStart; + const storedOAuthStart = + !lastOAuthStart || lastOAuthStart.options?.expiresAt == null + ? await getReplayablePendingMCPOAuthStart({ flowManager, userId, serverName }) + : undefined; + const currentOAuthStart = this.lastOAuthStart; + const replayOAuthStart = + storedOAuthStart && + (!currentOAuthStart || storedOAuthStart.authURL === currentOAuthStart.authURL) + ? storedOAuthStart + : currentOAuthStart; + if (!replayOAuthStart) { + return; + } + if (this.oauthEnded) { + return; + } + + this.lastOAuthStart = replayOAuthStart; + try { + await this.emit(oauthStart, replayOAuthStart.authURL, replayOAuthStart.options); + } catch (error) { + logger.warn(`${this.logPrefix} Failed to re-issue pending OAuth URL`, error); + } + } + + private async emit( + oauthStart: t.OAuthStartHandler, + authURL: string, + options?: t.OAuthStartOptions, + ): Promise { + if (this.emittedAuthUrls.get(oauthStart) === authURL) { + return; + } + this.emittedAuthUrls.set(oauthStart, authURL); + await oauthStart(authURL, options); + } + + private async emitEnd(oauthEnd: OAuthEndHandler): Promise { + if (this.emittedOAuthEnds.has(oauthEnd)) { + return; + } + this.emittedOAuthEnds.add(oauthEnd); + await oauthEnd(); + } +} diff --git a/packages/api/src/mcp/request.ts b/packages/api/src/mcp/request.ts index 834283a01b..68b76027f0 100644 --- a/packages/api/src/mcp/request.ts +++ b/packages/api/src/mcp/request.ts @@ -21,6 +21,7 @@ interface MCPResponseLike { interface Disconnectable { disconnect: () => Promise | unknown; + dispose?: () => Promise | unknown; } const contexts = new WeakMap(); @@ -50,29 +51,36 @@ export async function cleanupMCPRequestContext(context?: MCPRequestContext): Pro } context.cleanupStarted = true; - const connections = new Set(); - for (const connection of context.connections.values()) { + const connections = new Map(); + for (const [connectionKey, connection] of context.connections) { if (isDisconnectable(connection)) { - connections.add(connection); + connections.set(connection, connectionKey); } } - const pending = Array.from(context.pending.values()); + const pending = Array.from(context.pending.entries()); if (pending.length > 0) { - const settled = await Promise.allSettled(pending); - for (const result of settled) { + const settled = await Promise.allSettled(pending.map(([, promise]) => promise)); + for (let index = 0; index < settled.length; index++) { + const result = settled[index]; if (result.status === 'fulfilled' && isDisconnectable(result.value)) { - connections.add(result.value); + connections.set(result.value, pending[index][0]); } } } await Promise.allSettled( - Array.from(connections).map(async (connection) => { + Array.from(connections).map(async ([connection, connectionKey]) => { try { - await connection.disconnect(); + if (context.disposeConnection) { + await context.disposeConnection(connectionKey, connection); + } else if (connection.dispose) { + await connection.dispose(); + } else { + await connection.disconnect(); + } } catch (error) { - logger.warn('[MCP Request Context] Failed to disconnect request-scoped connection', error); + logger.warn('[MCP Request Context] Failed to dispose request-scoped connection', error); } }), ); diff --git a/packages/api/src/mcp/types/index.ts b/packages/api/src/mcp/types/index.ts index c81e31d48b..be8d93e755 100644 --- a/packages/api/src/mcp/types/index.ts +++ b/packages/api/src/mcp/types/index.ts @@ -62,6 +62,8 @@ export interface MCPPrompt { export type ConnectionState = 'disconnected' | 'connecting' | 'connected' | 'error'; +export type OAuthHandledSource = 'silent-refresh' | 'interactive'; + export type MCPTool = Tool; export type MCPToolListResponse = ListToolsResult; export type ToolContentPart = TextContent | ImageContent | EmbeddedResource | AudioContent; @@ -211,6 +213,7 @@ export interface UserConnectionContext { export interface RequestScopedMCPConnectionStore { connections: Map; pending: Map>; + disposeConnection?: (connectionKey: string, connection: unknown) => Promise; } export interface OAuthStartOptions {