From 56175af0b5f30a9ed460726beaf071facf4f8af0 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Wed, 5 Aug 2026 19:42:26 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=8E=9F=EF=B8=8F=20fix:=20Reconcile=20MCP?= =?UTF-8?q?=20OAuth=20Readiness=20Across=20Pods=20(#14629)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix: stabilize MCP OAuth readiness across pods * fix: harden MCP readiness review findings * fix: resolve CI type check and terminal OAuth polling * fix: address MCP OAuth readiness review * fix: align MCP OAuth readiness state * test: stabilize MCP OAuth readiness assertion * fix: reject stale MCP OAuth callbacks * fix: close distributed MCP OAuth readiness gaps * style: sort Redis MCP test imports * fix: preserve MCP OAuth polling across rolling pods * fix: finalize distributed MCP OAuth readiness * fix: preserve runtime-detected MCP OAuth * fix: report runtime MCP OAuth readiness * fix: preserve live MCP OAuth classification * style: sort MCP connection imports --- api/config/index.js | 2 + api/server/controllers/mcp.js | 22 +- api/server/routes/__tests__/mcp.spec.js | 267 ++++++- api/server/routes/mcp.js | 117 ++- api/server/services/MCP.js | 119 ++- api/server/services/MCP.spec.js | 681 +++++++++++++++++- api/server/services/Tools/mcp.js | 10 +- api/server/services/Tools/mcp.spec.js | 37 + .../src/hooks/MCP/__tests__/polling.spec.ts | 106 +++ client/src/hooks/MCP/polling.ts | 67 ++ client/src/hooks/MCP/useMCPServerManager.ts | 163 +++-- e2e/specs/mock/mcp-oauth-readiness.spec.ts | 417 +++++++++++ packages/api/src/flow/manager.test.ts | 43 ++ packages/api/src/flow/manager.ts | 34 +- packages/api/src/flow/types.ts | 4 + .../__tests__/MCPConnectionFetchTools.test.ts | 16 + .../MCPFlowRedis.cache_integration.spec.ts | 79 ++ .../__tests__/MCPOAuthTokenStorage.test.ts | 172 +++++ packages/api/src/mcp/connection.ts | 7 +- packages/api/src/mcp/oauth/tokens.ts | 75 ++ packages/data-provider/src/api-endpoints.ts | 3 + packages/data-provider/src/data-service.ts | 4 + .../data-provider/src/types/mcpServers.ts | 11 + packages/data-provider/src/types/queries.ts | 7 + 24 files changed, 2352 insertions(+), 111 deletions(-) create mode 100644 client/src/hooks/MCP/__tests__/polling.spec.ts create mode 100644 client/src/hooks/MCP/polling.ts create mode 100644 e2e/specs/mock/mcp-oauth-readiness.spec.ts create mode 100644 packages/api/src/mcp/__tests__/MCPFlowRedis.cache_integration.spec.ts diff --git a/api/config/index.js b/api/config/index.js index 6d9f70ecbb..804c094d29 100644 --- a/api/config/index.js +++ b/api/config/index.js @@ -25,6 +25,8 @@ function getFlowStateManager(flowsCache) { if (!flowManager) { flowManager = new FlowStateManager(flowsCache, { ttl: mcpConfig.OAUTH_FLOW_TTL, + monitorTimeout: mcpConfig.OAUTH_HANDLING_TIMEOUT, + retainedFailureTypes: ['mcp_oauth'], }); } return flowManager; diff --git a/api/server/controllers/mcp.js b/api/server/controllers/mcp.js index 923dbf013a..10eba42768 100644 --- a/api/server/controllers/mcp.js +++ b/api/server/controllers/mcp.js @@ -5,10 +5,12 @@ * @import { MCPServerRegistry } from '@librechat/api' * @import { MCPServerDocument } from 'librechat-data-provider' */ +const { randomUUID } = require('crypto'); const { logger, SystemCapabilities } = require('@librechat/data-schemas'); const { checkAccess, isUserSourced, + MCPConnection, MCPErrorCodes, splitMCPToolKey, normalizeServerName, @@ -403,13 +405,19 @@ const createMCPServerController = async (req, res) => { const reservedServerNames = [ ...new Set([...configNames, ...configNames.map(normalizeServerName)]), ]; - const result = await getMCPServersRegistry().addServer( - 'temp_server_name', - validation.data, - 'DB', - userId, - reservedServerNames, - ); + const inspectionServerName = `temp_server_${randomUUID()}`; + let result; + try { + result = await getMCPServersRegistry().addServer( + inspectionServerName, + validation.data, + 'DB', + userId, + reservedServerNames, + ); + } finally { + MCPConnection.clearCooldown(inspectionServerName); + } res.status(201).json({ serverName: result.serverName, ...redactServerSecrets(result.config, { canEdit: true }), diff --git a/api/server/routes/__tests__/mcp.spec.js b/api/server/routes/__tests__/mcp.spec.js index 634a1878c6..83dfb7874b 100644 --- a/api/server/routes/__tests__/mcp.spec.js +++ b/api/server/routes/__tests__/mcp.spec.js @@ -61,6 +61,9 @@ jest.mock('@librechat/api', () => { getTokens: jest.fn(), deleteUserTokens: jest.fn(), }, + MCPConnection: { + clearCooldown: jest.fn(), + }, getUserMCPAuthMap: jest.fn(), generateCheckAccess: jest.fn(({ permissionType, permissions }) => (req, res, next) => { const { PermissionTypes, Permissions } = require('librechat-data-provider'); @@ -1181,7 +1184,7 @@ describe('MCP Routes', () => { it('should handle OAuth callback successfully', async () => { // mockRegistryInstance is defined at the top of the file const mockFlowManager = { - getFlowState: jest.fn().mockResolvedValue({ status: 'PENDING' }), + getFlowState: jest.fn().mockResolvedValue({ status: 'PENDING', createdAt: Date.now() }), completeFlow: jest.fn().mockResolvedValue(), deleteFlow: jest.fn().mockResolvedValue(true), }; @@ -1275,7 +1278,7 @@ describe('MCP Routes', () => { it('should clear tenant-scoped token flow state after storing callback tokens', async () => { const mockFlowManager = { - getFlowState: jest.fn().mockResolvedValue({ status: 'PENDING' }), + getFlowState: jest.fn().mockResolvedValue({ status: 'PENDING', createdAt: Date.now() }), completeFlow: jest.fn().mockResolvedValue(), deleteFlow: jest.fn().mockResolvedValue(true), }; @@ -1349,7 +1352,7 @@ describe('MCP Routes', () => { status: 'PENDING', }); } - return Promise.resolve({ status: 'PENDING' }); + return Promise.resolve({ status: 'PENDING', createdAt: Date.now() }); }), completeFlow: jest.fn().mockResolvedValue(true), deleteFlow: jest.fn().mockResolvedValue(true), @@ -1415,7 +1418,7 @@ describe('MCP Routes', () => { it('should use oauthHeaders from flow state when present', async () => { const mockFlowManager = { - getFlowState: jest.fn().mockResolvedValue({ status: 'PENDING' }), + getFlowState: jest.fn().mockResolvedValue({ status: 'PENDING', createdAt: Date.now() }), completeFlow: jest.fn().mockResolvedValue(), deleteFlow: jest.fn().mockResolvedValue(true), }; @@ -1467,7 +1470,7 @@ describe('MCP Routes', () => { it('should fall back to registry oauth_headers when flow state lacks them', async () => { const mockFlowManager = { - getFlowState: jest.fn().mockResolvedValue({ status: 'PENDING' }), + getFlowState: jest.fn().mockResolvedValue({ status: 'PENDING', createdAt: Date.now() }), completeFlow: jest.fn().mockResolvedValue(), deleteFlow: jest.fn().mockResolvedValue(true), }; @@ -1544,7 +1547,7 @@ describe('MCP Routes', () => { it('should handle system-level OAuth completion', async () => { // mockRegistryInstance is defined at the top of the file const mockFlowManager = { - getFlowState: jest.fn().mockResolvedValue({ status: 'PENDING' }), + getFlowState: jest.fn().mockResolvedValue({ status: 'PENDING', createdAt: Date.now() }), completeFlow: jest.fn().mockResolvedValue(), deleteFlow: jest.fn().mockResolvedValue(true), }; @@ -1588,7 +1591,7 @@ describe('MCP Routes', () => { it('should handle reconnection failure after OAuth', async () => { // mockRegistryInstance is defined at the top of the file const mockFlowManager = { - getFlowState: jest.fn().mockResolvedValue({ status: 'PENDING' }), + getFlowState: jest.fn().mockResolvedValue({ status: 'PENDING', createdAt: Date.now() }), completeFlow: jest.fn().mockResolvedValue(), deleteFlow: jest.fn().mockResolvedValue(true), }; @@ -1716,7 +1719,7 @@ describe('MCP Routes', () => { // First call checks idempotency (status PENDING = not completed) // Second call retrieves flow state for processing mockFlowManager.getFlowState - .mockResolvedValueOnce({ status: 'PENDING' }) + .mockResolvedValueOnce({ status: 'PENDING', createdAt: Date.now() }) .mockResolvedValueOnce(flowState); MCPOAuthHandler.getFlowState.mockResolvedValue(flowState); @@ -1803,6 +1806,72 @@ describe('MCP Routes', () => { expect(MCPOAuthHandler.completeOAuthFlow).not.toHaveBeenCalled(); expect(MCPTokenStorage.storeTokens).not.toHaveBeenCalled(); }); + + it('should reject a retained failed flow without exchanging or storing tokens', async () => { + const flowId = 'test-user-id:test-server'; + const mockFlowManager = { + getFlowState: jest.fn().mockResolvedValue({ + status: 'FAILED', + error: 'mcp_oauth flow timed out', + }), + }; + + MCPOAuthHandler.getFlowState.mockResolvedValue({ + state: flowId, + serverName: 'test-server', + userId: 'test-user-id', + metadata: {}, + clientInfo: {}, + codeVerifier: 'test-verifier', + }); + getLogStores.mockReturnValue({}); + require('~/config').getFlowStateManager.mockReturnValue(mockFlowManager); + + const csrfToken = generateTestCsrfToken(flowId); + const response = await request(app) + .get('/api/mcp/test-server/oauth/callback') + .set('Cookie', [`oauth_csrf=${csrfToken}`]) + .query({ code: 'late-auth-code', state: flowId }); + const basePath = getBasePath(); + + expect(response.status).toBe(302); + expect(response.headers.location).toBe(`${basePath}/oauth/error?error=invalid_state`); + expect(MCPOAuthHandler.completeOAuthFlow).not.toHaveBeenCalled(); + expect(MCPTokenStorage.storeTokens).not.toHaveBeenCalled(); + }); + + it('should reject an over-age pending flow without relying on its original monitor', async () => { + const flowId = 'test-user-id:test-server'; + const mockFlowManager = { + getFlowState: jest.fn().mockResolvedValue({ + status: 'PENDING', + createdAt: Date.now() - PENDING_STALE_MS - 1000, + }), + }; + + MCPOAuthHandler.getFlowState.mockResolvedValue({ + state: flowId, + serverName: 'test-server', + userId: 'test-user-id', + metadata: {}, + clientInfo: {}, + codeVerifier: 'test-verifier', + }); + getLogStores.mockReturnValue({}); + require('~/config').getFlowStateManager.mockReturnValue(mockFlowManager); + + const csrfToken = generateTestCsrfToken(flowId); + const response = await request(app) + .get('/api/mcp/test-server/oauth/callback') + .set('Cookie', [`oauth_csrf=${csrfToken}`]) + .query({ code: 'late-auth-code', state: flowId }); + const basePath = getBasePath(); + + expect(response.status).toBe(302); + expect(response.headers.location).toBe(`${basePath}/oauth/error?error=invalid_state`); + expect(MCPOAuthHandler.completeOAuthFlow).not.toHaveBeenCalled(); + expect(MCPTokenStorage.storeTokens).not.toHaveBeenCalled(); + }); }); describe('GET /oauth/tokens/:flowId', () => { @@ -1968,6 +2037,28 @@ describe('MCP Routes', () => { }); }); + it('should return retained timeout failures as terminal status', async () => { + const mockFlowManager = { + getFlowState: jest.fn().mockResolvedValue({ + status: 'FAILED', + error: 'mcp_oauth flow timed out', + }), + }; + + getLogStores.mockReturnValue({}); + require('~/config').getFlowStateManager.mockReturnValue(mockFlowManager); + + const response = await request(app).get('/api/mcp/oauth/status/test-user-id:test-server'); + + expect(response.status).toBe(200); + expect(response.body).toEqual({ + status: 'FAILED', + completed: false, + failed: true, + error: 'mcp_oauth flow timed out', + }); + }); + it('should return flow status for a tenant-prefixed flow owned by the user', async () => { const { getTenantId } = require('@librechat/data-schemas'); const mockFlowManager = { @@ -2183,9 +2274,36 @@ describe('MCP Routes', () => { serverName: 'oauth-server', oauthRequired: true, oauthUrl: 'https://oauth.example.com/auth', + flowId: 'test-user-id:oauth-server', + oauthTimeout: expect.any(Number), }); }); + it('should return the remaining lifetime for a reused OAuth flow', async () => { + const now = 1_800_000_000_000; + const dateNowSpy = jest.spyOn(Date, 'now').mockReturnValue(now); + const mockMcpManager = { + disconnectUserConnection: jest.fn().mockResolvedValue(), + }; + + mockRegistryInstance.getServerConfig.mockResolvedValue({ customUserVars: {} }); + require('~/config').getMCPManager.mockReturnValue(mockMcpManager); + require('~/server/services/Tools/mcp').reinitMCPServer.mockResolvedValue({ + success: true, + message: "MCP server 'oauth-server' ready for OAuth authentication", + serverName: 'oauth-server', + oauthRequired: true, + oauthUrl: 'https://oauth.example.com/auth', + oauthExpiresAt: now + 45_000, + }); + + const response = await request(app).post('/api/mcp/oauth-server/reinitialize'); + dateNowSpy.mockRestore(); + + expect(response.status).toBe(200); + expect(response.body.oauthTimeout).toBe(45_000); + }); + it('should return structured reinitialization failure details', async () => { const mockMcpManager = { disconnectUserConnection: jest.fn().mockResolvedValue(), @@ -2386,17 +2504,19 @@ describe('MCP Routes', () => { mcpConfig: mockMcpConfig, appConnections: {}, userConnections: {}, - oauthServers: [], + oauthServers: new Set(), }); getServerConnectionStatus .mockResolvedValueOnce({ connectionState: 'connected', requiresOAuth: false, + authorizationState: 'not_required', }) .mockResolvedValueOnce({ connectionState: 'disconnected', requiresOAuth: true, + authorizationState: 'needs_authorization', }); const response = await request(app).get('/api/mcp/connection/status'); @@ -2409,10 +2529,12 @@ describe('MCP Routes', () => { server1: { connectionState: 'connected', requiresOAuth: false, + authorizationState: 'not_required', }, server2: { connectionState: 'disconnected', requiresOAuth: true, + authorizationState: 'needs_authorization', }, }, }); @@ -2421,6 +2543,113 @@ describe('MCP Routes', () => { expect(getServerConnectionStatus).toHaveBeenCalledTimes(2); }); + it('should batch user variables and pass runtime context into durable status checks', async () => { + currentUser = { id: 'test-user-id', email: 'user@example.com' }; + const mcpConfig = { + server1: { + url: 'https://mcp.example.com/{{LIBRECHAT_USER_ID}}', + customUserVars: { API_KEY: { title: 'API key' } }, + }, + }; + const userMCPAuthMap = { + 'mcp:server1': { API_KEY: 'secret' }, + }; + getMCPSetupData.mockResolvedValue({ + mcpConfig, + appConnections: new Map(), + userConnections: new Map(), + oauthServers: new Set(), + }); + require('@librechat/api').getUserMCPAuthMap.mockResolvedValue(userMCPAuthMap); + getServerConnectionStatus.mockResolvedValue({ + connectionState: 'connected', + requiresOAuth: true, + authorizationState: 'authorized', + }); + + const response = await request(app).get('/api/mcp/connection/status'); + + expect(response.status).toBe(200); + expect(require('@librechat/api').getUserMCPAuthMap).not.toHaveBeenCalled(); + const runtimeContext = getServerConnectionStatus.mock.calls[0][6]; + await expect( + Promise.all([runtimeContext.loadUserMCPAuthMap(), runtimeContext.loadUserMCPAuthMap()]), + ).resolves.toEqual([userMCPAuthMap, userMCPAuthMap]); + await expect( + Promise.all([runtimeContext.loadMCPAllowlists(), runtimeContext.loadMCPAllowlists()]), + ).resolves.toEqual([ + { allowedDomains: null, allowedAddresses: null, useSSRFProtection: true }, + { allowedDomains: null, allowedAddresses: null, useSSRFProtection: true }, + ]); + expect(require('@librechat/api').getUserMCPAuthMap).toHaveBeenCalledWith({ + userId: 'test-user-id', + servers: ['server1'], + findPluginAuthsByKeys: require('~/models').findPluginAuthsByKeys, + }); + expect(require('@librechat/api').getUserMCPAuthMap).toHaveBeenCalledTimes(1); + expect(mockRegistryInstance.resolveAllowlists).toHaveBeenCalledTimes(1); + expect(mockRegistryInstance.resolveAllowlists).toHaveBeenCalledWith({ + userId: 'test-user-id', + role: undefined, + }); + expect(getServerConnectionStatus).toHaveBeenCalledWith( + 'test-user-id', + 'server1', + mcpConfig.server1, + expect.any(Map), + expect.any(Map), + expect.any(Set), + { + user: expect.objectContaining({ id: 'test-user-id', email: 'user@example.com' }), + loadUserMCPAuthMap: expect.any(Function), + loadMCPAllowlists: expect.any(Function), + }, + ); + }); + + it('should resolve independent server statuses concurrently', async () => { + let releaseFirst; + let markSecondStarted; + const firstRelease = new Promise((resolve) => { + releaseFirst = resolve; + }); + const secondStarted = new Promise((resolve) => { + markSecondStarted = resolve; + }); + + getMCPSetupData.mockResolvedValue({ + mcpConfig: { + server1: { endpoint: 'http://server1.com' }, + server2: { endpoint: 'http://server2.com' }, + }, + appConnections: new Map(), + userConnections: new Map(), + oauthServers: new Set(), + }); + getServerConnectionStatus.mockImplementation(async (_userId, serverName) => { + if (serverName === 'server1') { + await firstRelease; + } else { + markSecondStarted(); + } + return { + connectionState: 'connected', + requiresOAuth: false, + authorizationState: 'not_required', + }; + }); + + const responsePromise = request(app) + .get('/api/mcp/connection/status') + .then((res) => res); + await secondStarted; + releaseFirst(); + const response = await responsePromise; + + expect(response.status).toBe(200); + expect(getServerConnectionStatus).toHaveBeenCalledTimes(2); + }); + it('should return 500 when connection status check fails', async () => { getMCPSetupData.mockRejectedValue(new Error('Database error')); @@ -2458,12 +2687,13 @@ describe('MCP Routes', () => { mcpConfig: mockMcpConfig, appConnections: {}, userConnections: {}, - oauthServers: [], + oauthServers: new Set(), }); getServerConnectionStatus.mockResolvedValue({ connectionState: 'requires_auth', requiresOAuth: true, + authorizationState: 'needs_authorization', }); const response = await request(app).get('/api/mcp/connection/status/oauth-server'); @@ -2474,6 +2704,7 @@ describe('MCP Routes', () => { serverName: 'oauth-server', connectionStatus: 'requires_auth', requiresOAuth: true, + authorizationState: 'needs_authorization', }); }); @@ -2484,7 +2715,7 @@ describe('MCP Routes', () => { }, appConnections: {}, userConnections: {}, - oauthServers: [], + oauthServers: new Set(), }); const response = await request(app).get('/api/mcp/connection/status/non-existent-server'); @@ -2657,7 +2888,7 @@ describe('MCP Routes', () => { mockRegistryInstance.getServerConfig.mockResolvedValue({}); const mockFlowManager = { - getFlowState: jest.fn().mockResolvedValue({ status: 'PENDING' }), + getFlowState: jest.fn().mockResolvedValue({ status: 'PENDING', createdAt: Date.now() }), completeFlow: jest.fn(), deleteFlow: jest.fn().mockResolvedValue(true), }; @@ -3009,6 +3240,8 @@ describe('MCP Routes', () => { }); describe('POST /servers', () => { + const { MCPConnection } = require('@librechat/api'); + it('should create MCP server with valid SSE config', async () => { const validConfig = { type: 'sse', @@ -3030,7 +3263,7 @@ describe('MCP Routes', () => { expect(response.body.url).toBe('https://mcp-server.example.com/sse'); expect(response.body.title).toBe('Test SSE Server'); expect(mockRegistryInstance.addServer).toHaveBeenCalledWith( - 'temp_server_name', + expect.stringMatching(/^temp_server_[0-9a-f-]{36}$/), expect.objectContaining({ type: 'sse', url: 'https://mcp-server.example.com/sse', @@ -3039,6 +3272,8 @@ describe('MCP Routes', () => { 'test-user-id', [], ); + const inspectionServerName = mockRegistryInstance.addServer.mock.calls[0][0]; + expect(MCPConnection.clearCooldown).toHaveBeenCalledWith(inspectionServerName); }); it('should reserve config-managed server names when creating MCP server', async () => { @@ -3058,7 +3293,7 @@ describe('MCP Routes', () => { expect(response.status).toBe(201); expect(mockRegistryInstance.addServer).toHaveBeenCalledWith( - 'temp_server_name', + expect.stringMatching(/^temp_server_[0-9a-f-]{36}$/), expect.objectContaining({ type: 'sse', url: 'https://mcp-server.example.com/sse', @@ -3088,7 +3323,7 @@ describe('MCP Routes', () => { expect(response.status).toBe(201); expect(mockRegistryInstance.addServer).toHaveBeenCalledWith( - 'temp_server_name', + expect.stringMatching(/^temp_server_[0-9a-f-]{36}$/), expect.anything(), 'DB', 'test-user-id', @@ -3223,6 +3458,8 @@ describe('MCP Routes', () => { expect(response.status).toBe(500); expect(response.body).toEqual({ message: 'Database connection failed' }); + const inspectionServerName = mockRegistryInstance.addServer.mock.calls[0][0]; + expect(MCPConnection.clearCooldown).toHaveBeenCalledWith(inspectionServerName); }); describe('OBO permission gate', () => { diff --git a/api/server/routes/mcp.js b/api/server/routes/mcp.js index c044a67372..e6a5eecf73 100644 --- a/api/server/routes/mcp.js +++ b/api/server/routes/mcp.js @@ -373,14 +373,19 @@ router.get('/:serverName/oauth/callback', async (req, res) => { }); return res.redirect(`${basePath}/oauth/success?serverName=${encodeURIComponent(serverName)}`); } + const isStalePendingFlow = + currentFlowState?.status === 'PENDING' && + (!currentFlowState.createdAt || Date.now() - currentFlowState.createdAt >= PENDING_STALE_MS); + if (currentFlowState?.status === 'FAILED' || isStalePendingFlow) { + logger.warn('[MCP OAuth] Refusing token exchange for terminal flow', { + flowId, + serverName, + status: currentFlowState.status, + }); + return res.redirect(`${basePath}/oauth/error?error=invalid_state`); + } logger.debug('[MCP OAuth] Completing OAuth flow'); - if (!flowState.oauthHeaders) { - logger.warn( - '[MCP OAuth] oauthHeaders absent from flow state — config-source server oauth_headers will be empty', - { serverName, flowId }, - ); - } /** * Restore tenant context for the callback body. The callback is a cross-origin * redirect from the OAuth provider, so SameSite=Strict cookies (including the @@ -716,6 +721,43 @@ router.post('/oauth/cancel/:serverName', requireJwtAuth, async (req, res) => { } }); +function createMCPStatusRuntimeContext(user, mcpConfig, serverNames) { + const customUserVarServers = serverNames.filter((serverName) => { + const customUserVars = mcpConfig[serverName]?.customUserVars; + return ( + customUserVars && typeof customUserVars === 'object' && Object.keys(customUserVars).length > 0 + ); + }); + let userMCPAuthMapPromise; + let mcpAllowlistsPromise; + const loadUserMCPAuthMap = () => { + if (!customUserVarServers.length) { + return Promise.resolve(undefined); + } + userMCPAuthMapPromise ??= getUserMCPAuthMap({ + userId: user.id, + servers: customUserVarServers, + findPluginAuthsByKeys: db.findPluginAuthsByKeys, + }); + return userMCPAuthMapPromise; + }; + const loadMCPAllowlists = () => { + mcpAllowlistsPromise ??= getMCPServersRegistry().resolveAllowlists({ + userId: user.id, + role: user.role, + }); + return mcpAllowlistsPromise; + }; + return { user: createSafeUser(user), loadUserMCPAuthMap, loadMCPAllowlists }; +} + +function getMCPReinitializeOAuthTimeout(oauthExpiresAt) { + if (typeof oauthExpiresAt !== 'number' || !Number.isFinite(oauthExpiresAt)) { + return mcpSettings.OAUTH_HANDLING_TIMEOUT; + } + return Math.max(0, oauthExpiresAt - Date.now()); +} + /** * Reinitialize MCP server * This endpoint allows reinitializing a specific MCP server @@ -781,13 +823,15 @@ router.post( message, oauthRequired, oauthUrl, + oauthExpiresAt, failureReason, missingUserVars, connectionDeferred, } = result; + let flowId; if (oauthRequired) { - const flowId = getOAuthFlowId(user.id, serverName); + flowId = getOAuthFlowId(user.id, serverName); setOAuthCsrfCookie(res, flowId, OAUTH_CSRF_COOKIE_PATH); } @@ -795,6 +839,8 @@ router.post( success, message, oauthUrl, + flowId, + oauthTimeout: oauthRequired ? getMCPReinitializeOAuthTimeout(oauthExpiresAt) : undefined, serverName, oauthRequired, failureReason, @@ -824,28 +870,37 @@ router.get('/connection/status', requireJwtAuth, async (req, res) => { user.id, { role: user.role, tenantId: getTenantId() }, ); - const connectionStatus = {}; - - for (const [serverName, config] of Object.entries(mcpConfig)) { - try { - connectionStatus[serverName] = await getServerConnectionStatus( - user.id, - serverName, - config, - appConnections, - userConnections, - oauthServers, - ); - } catch (error) { - const message = `Failed to get status for server "${serverName}"`; - logger.error(`[MCP Connection Status] ${message},`, error); - connectionStatus[serverName] = { - connectionState: 'error', - requiresOAuth: oauthServers.has(serverName), - error: message, - }; - } - } + const runtimeContext = createMCPStatusRuntimeContext(user, mcpConfig, Object.keys(mcpConfig)); + const connectionStatus = Object.fromEntries( + await Promise.all( + Object.entries(mcpConfig).map(async ([serverName, config]) => { + try { + const status = await getServerConnectionStatus( + user.id, + serverName, + config, + appConnections, + userConnections, + oauthServers, + runtimeContext, + ); + return [serverName, status]; + } catch (error) { + const message = `Failed to get status for server "${serverName}"`; + logger.error(`[MCP Connection Status] ${message},`, error); + return [ + serverName, + { + connectionState: 'error', + requiresOAuth: oauthServers.has(serverName), + authorizationState: oauthServers.has(serverName) ? 'error' : 'not_required', + error: message, + }, + ]; + } + }), + ), + ); res.json({ success: true, @@ -882,6 +937,8 @@ router.get('/connection/status/:serverName', requireJwtAuth, async (req, res) => .json({ error: `MCP server '${serverName}' not found in configuration` }); } + const runtimeContext = createMCPStatusRuntimeContext(user, mcpConfig, [serverName]); + const serverStatus = await getServerConnectionStatus( user.id, serverName, @@ -889,6 +946,7 @@ router.get('/connection/status/:serverName', requireJwtAuth, async (req, res) => appConnections, userConnections, oauthServers, + runtimeContext, ); res.json({ @@ -896,6 +954,7 @@ router.get('/connection/status/:serverName', requireJwtAuth, async (req, res) => serverName, connectionStatus: serverStatus.connectionState, requiresOAuth: serverStatus.requiresOAuth, + authorizationState: serverStatus.authorizationState, }); } catch (error) { logger.error( diff --git a/api/server/services/MCP.js b/api/server/services/MCP.js index b2781c9bec..acc8a41034 100644 --- a/api/server/services/MCP.js +++ b/api/server/services/MCP.js @@ -5,6 +5,7 @@ const { sendEvent, PENDING_STALE_MS, MCPOAuthHandler, + MCPTokenStorage, isMCPDomainAllowed, splitMCPToolKey, normalizeServerName, @@ -19,14 +20,19 @@ const { buildMCPAuthStepId, buildMCPAuthToolCall, processMCPEnv, + preProcessGraphTokens, buildMCPAuthRunStepEvent, buildMCPAuthRunStepDeltaEvent, buildMCPAuthRunStepEndDeltaEvent, isUserSourced, checkAccessWithRequestCache, + getMissingCustomUserVars, getServerCustomUserVars, requiresEphemeralUserConnection, + requiresOAuthMachinery, + hasRuntimeUrlPlaceholders, containsGraphTokenPlaceholder, + isOAuthServer, } = require('@librechat/api'); const { Time, @@ -1119,8 +1125,20 @@ function createToolInstance({ error.message?.includes('OAuth') || error.message?.includes('authentication') || error.message?.includes('Non-200 status code (401)'); + const isOAuthFlowSignal = + error.message === 'OAuth flow initiated - return early' || + error.message === 'Pending OAuth flow reused - return early'; if (isOAuthError) { + if ( + capturedServerConfig && + !requiresOAuthMachinery(capturedServerConfig) && + !isOAuthFlowSignal + ) { + throw new Error( + `[MCP][${serverName}][${toolName}] upstream authentication failed; MCP OAuth is not configured for this server.`, + ); + } throw new Error( `[MCP][${serverName}][${toolName}] OAuth authentication required. Please check the server logs for the authentication URL.`, ); @@ -1189,7 +1207,7 @@ async function getMCPSetupData(userId, options = {}) { const userConnections = mcpManager.getUserConnections(userId) || new Map(); const oauthServers = new Set( Object.entries(mcpConfig) - .filter(([, config]) => config.requiresOAuth) + .filter(([, config]) => isOAuthServer(config)) .map(([name]) => name), ); @@ -1206,7 +1224,7 @@ async function getMCPSetupData(userId, options = {}) { * @param {string} userId - The user ID * @param {string} serverName - The server name * @param {string} [tenantId] - The tenant ID for the current request. - * @returns {Object} Object containing hasActiveFlow and hasFailedFlow flags + * @returns {Object} Object containing active and failed flow flags */ async function checkOAuthFlowStatus(userId, serverName, tenantId = getTenantId()) { const flowsCache = getLogStores(CacheKeys.FLOWS); @@ -1225,8 +1243,8 @@ async function checkOAuthFlowStatus(userId, serverName, tenantId = getTenantId() // flow the initiate/callback paths already reject, hiding the connect button. const flowTTL = flowState.ttl || PENDING_STALE_MS; - if (flowState.status === 'FAILED' || flowAge > flowTTL) { - const wasCancelled = flowState.error && flowState.error.includes('cancelled'); + if (flowState.status === 'FAILED' || (flowState.status === 'PENDING' && flowAge > flowTTL)) { + const wasCancelled = /abort|cancel/i.test(flowState.error ?? ''); if (wasCancelled) { logger.debug(`[MCP Connection Status] Found cancelled OAuth flow for ${serverName}`, { @@ -1264,6 +1282,69 @@ async function checkOAuthFlowStatus(userId, serverName, tenantId = getTenantId() } } +async function hasDurableMCPAuthorization(userId, serverName, config, runtimeContext = {}) { + const userMCPAuthMap = + runtimeContext.userMCPAuthMap ?? (await runtimeContext.loadUserMCPAuthMap?.()); + const customUserVars = getServerCustomUserVars(userMCPAuthMap, serverName); + if (getMissingCustomUserVars(config, customUserVars).length > 0) { + return false; + } + + const dbSourced = isUserSourced(config); + const bindingConfig = { + ...config, + args: undefined, + env: undefined, + headers: undefined, + oauth_headers: undefined, + }; + const graphProcessedConfig = dbSourced + ? bindingConfig + : await preProcessGraphTokens(bindingConfig, { + user: runtimeContext.user, + graphTokenResolver: getGraphApiToken, + scopes: process.env.GRAPH_API_SCOPES, + }); + const runtimeConfig = processMCPEnv({ + user: runtimeContext.user, + options: graphProcessedConfig, + dbSourced, + customUserVars, + }); + const allowlists = await (runtimeContext.loadMCPAllowlists?.() ?? + getMCPServersRegistry().resolveAllowlists({ + userId, + role: runtimeContext.user?.role, + })); + if ( + runtimeConfig.url && + !(await isMCPDomainAllowed( + runtimeConfig, + allowlists.allowedDomains, + allowlists.allowedAddresses, + )) + ) { + return false; + } + return MCPTokenStorage.hasStoredAuthorization({ + userId, + serverName, + findToken, + validateClientBinding: (clientInfo, storedMetadata) => + MCPOAuthHandler.assertStoredClientBinding( + serverName, + runtimeConfig.url, + clientInfo, + storedMetadata, + runtimeConfig.oauth, + ), + }); +} + +function canDetectMCPRuntimeOAuth(config) { + return config.requiresOAuth == null && config.apiKey == null && hasRuntimeUrlPlaceholders(config); +} + /** * Get connection status for a specific MCP server * @param {string} userId - The user ID @@ -1272,6 +1353,7 @@ async function checkOAuthFlowStatus(userId, serverName, tenantId = getTenantId() * @param {Map} appConnections - App-level connections * @param {Map} userConnections - User-level connections * @param {Set} oauthServers - Set of OAuth servers + * @param {{ user?: Partial, userMCPAuthMap?: Record>, loadUserMCPAuthMap?: () => Promise> | undefined>, loadMCPAllowlists?: () => Promise<{ allowedDomains?: string[] | null, allowedAddresses?: string[] | null }> }} [runtimeContext] * @returns {Object} Object containing requiresOAuth and connectionState */ async function getServerConnectionStatus( @@ -1281,35 +1363,60 @@ async function getServerConnectionStatus( appConnections, userConnections, oauthServers, + runtimeContext = {}, ) { const connection = appConnections.get(serverName) || userConnections.get(serverName); const isStaleOrDoNotExist = connection ? connection?.isStale(config.updatedAt) : true; + const configuredOAuth = oauthServers.has(serverName); + const liveConnectionOAuth = connection?.usesOAuth?.() === true; + const runtimeOAuthCandidate = canDetectMCPRuntimeOAuth(config); + const effectiveOAuth = configuredOAuth || liveConnectionOAuth; const baseConnectionState = isStaleOrDoNotExist ? 'disconnected' : connection?.connectionState || 'disconnected'; let finalConnectionState = baseConnectionState; + let requiresOAuth = effectiveOAuth; + let authorizationState = effectiveOAuth ? 'needs_authorization' : 'not_required'; // connection state overrides specific to OAuth servers - if (baseConnectionState === 'disconnected' && oauthServers.has(serverName)) { + if (effectiveOAuth && baseConnectionState === 'connected') { + authorizationState = 'authorized'; + } else if (effectiveOAuth && baseConnectionState === 'connecting') { + authorizationState = 'authorizing'; + } else if (effectiveOAuth && baseConnectionState === 'error') { + authorizationState = 'error'; + } else if (baseConnectionState === 'disconnected' && (effectiveOAuth || runtimeOAuthCandidate)) { // check if server is actively being reconnected const oauthReconnectionManager = getOAuthReconnectionManager(); if (oauthReconnectionManager.isReconnecting(userId, serverName)) { + requiresOAuth = true; finalConnectionState = 'connecting'; + authorizationState = 'authorizing'; } else { const { hasActiveFlow, hasFailedFlow } = await checkOAuthFlowStatus(userId, serverName); if (hasFailedFlow) { + requiresOAuth = true; finalConnectionState = 'error'; + authorizationState = 'error'; } else if (hasActiveFlow) { + requiresOAuth = true; finalConnectionState = 'connecting'; + authorizationState = 'authorizing'; + } else if (await hasDurableMCPAuthorization(userId, serverName, config, runtimeContext)) { + /** OAuth readiness is durable even when this pod has no live connection. */ + requiresOAuth = true; + finalConnectionState = 'connected'; + authorizationState = 'authorized'; } } } return { - requiresOAuth: oauthServers.has(serverName), + requiresOAuth, connectionState: finalConnectionState, + authorizationState, }; } diff --git a/api/server/services/MCP.spec.js b/api/server/services/MCP.spec.js index 1e529189c6..23e7067e76 100644 --- a/api/server/services/MCP.spec.js +++ b/api/server/services/MCP.spec.js @@ -9,6 +9,7 @@ jest.mock('@librechat/data-schemas', () => ({ warn: jest.fn(), }, getTenantId: mockGetTenantId, + decryptV2: jest.fn(async (value) => value.replace(/^enc:/, '')), })); // Create mock registry instance @@ -17,6 +18,9 @@ const mockRegistryInstance = { getAllServerConfigs: jest.fn(() => Promise.resolve({})), getServerConfig: jest.fn(() => Promise.resolve(null)), ensureConfigServers: jest.fn(() => Promise.resolve({})), + resolveAllowlists: jest.fn(() => + Promise.resolve({ allowedDomains: null, allowedAddresses: null, useSSRFProtection: true }), + ), }; // Create isMCPDomainAllowed mock that can be configured per-test @@ -97,6 +101,12 @@ describe('tests for the new helper functions used by the MCP connection status e jest.clearAllMocks(); jest.spyOn(MCPOAuthHandler, 'generateFlowId'); mockGetTenantId.mockReturnValue(undefined); + mockIsMCPDomainAllowed.mockResolvedValue(true); + mockRegistryInstance.resolveAllowlists.mockResolvedValue({ + allowedDomains: null, + allowedAddresses: null, + useSSRFProtection: true, + }); mockGetMCPManager = require('~/config').getMCPManager; mockGetFlowStateManager = require('~/config').getFlowStateManager; @@ -203,6 +213,7 @@ describe('tests for the new helper functions used by the MCP connection status e const mockConfigWithOAuth = { server1: { type: 'stdio' }, server2: { type: 'http', requiresOAuth: true }, + server3: { type: 'http', oauth: { client_id: 'configured-client' } }, }; mockRegistryInstance.getAllServerConfigs.mockResolvedValue(mockConfigWithOAuth); @@ -229,7 +240,7 @@ describe('tests for the new helper functions used by the MCP connection status e expect(result.mcpConfig).toEqual(mockConfigWithOAuth); expect(result.appConnections).toEqual(mockAppConnections); expect(result.userConnections).toEqual(mockUserConnections); - expect(result.oauthServers).toEqual(new Set(['server2'])); + expect(result.oauthServers).toEqual(new Set(['server2', 'server3'])); }); it('should return empty data when no servers are configured', async () => { @@ -309,6 +320,21 @@ describe('tests for the new helper functions used by the MCP connection status e ); }); + it('should treat aborted flow cleanup as neutral connection status', async () => { + const mockFlowState = { + status: 'FAILED', + createdAt: Date.now() - 60000, + ttl: 180000, + error: 'Tool loading aborted', + }; + const mockFlowManager = { getFlowState: jest.fn(() => mockFlowState) }; + mockGetFlowStateManager.mockReturnValue(mockFlowManager); + + const result = await checkOAuthFlowStatus(mockUserId, mockServerName); + + expect(result).toEqual({ hasActiveFlow: false, hasFailedFlow: false }); + }); + it('should detect failed flow when flow has timed out', async () => { const mockFlowState = { status: 'PENDING', @@ -385,7 +411,7 @@ describe('tests for the new helper functions used by the MCP connection status e expect(result).toEqual({ hasActiveFlow: true, hasFailedFlow: false }); }); - it('should return false flags for other statuses', async () => { + it('should not treat a completed flow as durable authorization', async () => { const mockFlowState = { status: 'COMPLETED', createdAt: Date.now() - 60000, @@ -399,6 +425,20 @@ describe('tests for the new helper functions used by the MCP connection status e expect(result).toEqual({ hasActiveFlow: false, hasFailedFlow: false }); }); + it('should treat an old completed flow as neutral instead of timed out', async () => { + const mockFlowState = { + status: 'COMPLETED', + createdAt: Date.now() - 200000, + ttl: 180000, + }; + const mockFlowManager = { getFlowState: jest.fn(() => mockFlowState) }; + mockGetFlowStateManager.mockReturnValue(mockFlowManager); + + const result = await checkOAuthFlowStatus(mockUserId, mockServerName); + + expect(result).toEqual({ hasActiveFlow: false, hasFailedFlow: false }); + }); + it('should handle errors gracefully', async () => { const mockError = new Error('Flow state error'); const mockFlowManager = { @@ -423,6 +463,10 @@ describe('tests for the new helper functions used by the MCP connection status e const mockServerName = 'test-server'; const mockConfig = { updatedAt: Date.now() }; + beforeEach(() => { + require('~/models').findToken.mockReset(); + }); + it('should return app connection state when available', async () => { const appConnections = new Map([ [ @@ -448,6 +492,7 @@ describe('tests for the new helper functions used by the MCP connection status e expect(result).toEqual({ requiresOAuth: false, connectionState: 'connected', + authorizationState: 'not_required', }); }); @@ -476,6 +521,7 @@ describe('tests for the new helper functions used by the MCP connection status e expect(result).toEqual({ requiresOAuth: false, connectionState: 'connecting', + authorizationState: 'not_required', }); }); @@ -496,6 +542,7 @@ describe('tests for the new helper functions used by the MCP connection status e expect(result).toEqual({ requiresOAuth: false, connectionState: 'disconnected', + authorizationState: 'not_required', }); }); @@ -532,6 +579,7 @@ describe('tests for the new helper functions used by the MCP connection status e expect(result).toEqual({ requiresOAuth: false, connectionState: 'connected', + authorizationState: 'not_required', }); }); @@ -593,6 +641,7 @@ describe('tests for the new helper functions used by the MCP connection status e expect(result).toEqual({ requiresOAuth: true, connectionState: 'error', + authorizationState: 'error', }); }); @@ -631,6 +680,489 @@ describe('tests for the new helper functions used by the MCP connection status e expect(result).toEqual({ requiresOAuth: true, connectionState: 'connecting', + authorizationState: 'authorizing', + }); + }); + + it('should require bound token storage after a completed OAuth flow on another pod', async () => { + const appConnections = new Map(); + const userConnections = new Map(); + const oauthServers = new Set([mockServerName]); + mockGetOAuthReconnectionManager.mockReturnValue({ isReconnecting: jest.fn(() => false) }); + mockGetFlowStateManager.mockReturnValue({ + getFlowState: jest.fn(() => ({ + status: 'COMPLETED', + createdAt: Date.now() - 60000, + result: { access_token: 'encrypted' }, + })), + }); + mockGetLogStores.mockReturnValue({}); + const { findToken } = require('~/models'); + findToken.mockResolvedValue(null); + + const result = await getServerConnectionStatus( + mockUserId, + mockServerName, + mockConfig, + appConnections, + userConnections, + oauthServers, + ); + + expect(result).toEqual({ + requiresOAuth: true, + connectionState: 'disconnected', + authorizationState: 'needs_authorization', + }); + }); + + it('should derive readiness from bound token storage after the flow record expires', async () => { + const appConnections = new Map(); + const userConnections = new Map(); + const oauthServers = new Set([mockServerName]); + const { findToken } = require('~/models'); + const credentialSetId = 'credential-set-a'; + const config = { ...mockConfig, url: 'https://mcp.example.com/' }; + mockGetOAuthReconnectionManager.mockReturnValue({ isReconnecting: jest.fn(() => false) }); + mockGetFlowStateManager.mockReturnValue({ getFlowState: jest.fn(() => null) }); + mockGetLogStores.mockReturnValue({}); + findToken + .mockResolvedValueOnce({ + expiresAt: new Date(Date.now() + 60000), + metadata: { credential_set_id: credentialSetId }, + }) + .mockResolvedValueOnce({ + token: 'enc:{"client_id":"dynamic-client"}', + metadata: { + credential_set_id: credentialSetId, + authorization_endpoint: 'https://auth.example.com/authorize', + token_endpoint: 'https://auth.example.com/token', + server_url: 'https://mcp.example.com/', + client_source: 'dynamic', + }, + }); + + const result = await getServerConnectionStatus( + mockUserId, + mockServerName, + config, + appConnections, + userConnections, + oauthServers, + ); + + expect(result).toEqual({ + requiresOAuth: true, + connectionState: 'connected', + authorizationState: 'authorized', + }); + }); + + it('should derive runtime-detected OAuth readiness from bound token storage', async () => { + const appConnections = new Map(); + const userConnections = new Map(); + const oauthServers = new Set(); + const { findToken } = require('~/models'); + const credentialSetId = 'credential-set-a'; + const config = { + ...mockConfig, + source: 'yaml', + url: 'https://mcp.example.com/users/{{LIBRECHAT_USER_ID}}/mcp', + }; + mockGetOAuthReconnectionManager.mockReturnValue({ isReconnecting: jest.fn(() => false) }); + mockGetFlowStateManager.mockReturnValue({ getFlowState: jest.fn(() => null) }); + mockGetLogStores.mockReturnValue({}); + findToken + .mockResolvedValueOnce({ + expiresAt: new Date(Date.now() + 60000), + metadata: { credential_set_id: credentialSetId }, + }) + .mockResolvedValueOnce({ + token: 'enc:{"client_id":"dynamic-client"}', + metadata: { + credential_set_id: credentialSetId, + authorization_endpoint: 'https://auth.example.com/authorize', + token_endpoint: 'https://auth.example.com/token', + server_url: `https://mcp.example.com/users/${mockUserId}/mcp`, + client_source: 'dynamic', + }, + }); + + const result = await getServerConnectionStatus( + mockUserId, + mockServerName, + config, + appConnections, + userConnections, + oauthServers, + { user: { id: mockUserId } }, + ); + + expect(result).toEqual({ + requiresOAuth: true, + connectionState: 'connected', + authorizationState: 'authorized', + }); + }); + + it('should preserve OAuth status from a live runtime-resolved connection', async () => { + const { findToken } = require('~/models'); + const appConnections = new Map(); + const userConnections = new Map([ + [ + mockServerName, + { + connectionState: 'connected', + isStale: jest.fn(() => false), + usesOAuth: jest.fn(() => true), + }, + ], + ]); + + const result = await getServerConnectionStatus( + mockUserId, + mockServerName, + { + ...mockConfig, + source: 'yaml', + url: 'https://mcp.example.com/{{LIBRECHAT_USER_ID}}/mcp', + }, + appConnections, + userConnections, + new Set(), + ); + + expect(result).toEqual({ + requiresOAuth: true, + connectionState: 'connected', + authorizationState: 'authorized', + }); + expect(findToken).not.toHaveBeenCalled(); + }); + + it('should derive runtime-detected OAuth state from an active shared flow', async () => { + const appConnections = new Map(); + const userConnections = new Map(); + const oauthServers = new Set(); + const config = { + ...mockConfig, + source: 'yaml', + url: 'https://mcp.example.com/{{LIBRECHAT_BODY_TENANT}}/mcp', + }; + mockGetOAuthReconnectionManager.mockReturnValue({ isReconnecting: jest.fn(() => false) }); + mockGetFlowStateManager.mockReturnValue({ + getFlowState: jest.fn(() => ({ + status: 'PENDING', + createdAt: Date.now() - 1000, + ttl: 180000, + })), + }); + mockGetLogStores.mockReturnValue({}); + + const result = await getServerConnectionStatus( + mockUserId, + mockServerName, + config, + appConnections, + userConnections, + oauthServers, + { user: { id: mockUserId } }, + ); + + expect(result).toEqual({ + requiresOAuth: true, + connectionState: 'connecting', + authorizationState: 'authorizing', + }); + }); + + it('should not inspect OAuth state for an explicitly non-OAuth runtime URL', async () => { + const appConnections = new Map(); + const userConnections = new Map(); + const oauthServers = new Set(); + const { findToken } = require('~/models'); + const mockFlowManager = { getFlowState: jest.fn() }; + mockGetFlowStateManager.mockReturnValue(mockFlowManager); + mockGetLogStores.mockReturnValue({}); + + const result = await getServerConnectionStatus( + mockUserId, + mockServerName, + { + ...mockConfig, + source: 'yaml', + url: 'https://mcp.example.com/{{LIBRECHAT_USER_ID}}/mcp', + requiresOAuth: false, + }, + appConnections, + userConnections, + oauthServers, + ); + + expect(result).toEqual({ + requiresOAuth: false, + connectionState: 'disconnected', + authorizationState: 'not_required', + }); + expect(mockFlowManager.getFlowState).not.toHaveBeenCalled(); + expect(findToken).not.toHaveBeenCalled(); + }); + + it('should not report durable readiness when the runtime URL violates current policy', async () => { + const appConnections = new Map(); + const userConnections = new Map(); + const oauthServers = new Set([mockServerName]); + const { findToken } = require('~/models'); + const config = { ...mockConfig, url: 'https://blocked.example.com/' }; + mockGetOAuthReconnectionManager.mockReturnValue({ isReconnecting: jest.fn(() => false) }); + mockGetFlowStateManager.mockReturnValue({ getFlowState: jest.fn(() => null) }); + mockGetLogStores.mockReturnValue({}); + mockRegistryInstance.resolveAllowlists.mockResolvedValue({ + allowedDomains: ['allowed.example.com'], + allowedAddresses: null, + useSSRFProtection: false, + }); + mockIsMCPDomainAllowed.mockResolvedValue(false); + + const result = await getServerConnectionStatus( + mockUserId, + mockServerName, + config, + appConnections, + userConnections, + oauthServers, + { user: { id: mockUserId, role: 'user' } }, + ); + + expect(result).toEqual({ + requiresOAuth: true, + connectionState: 'disconnected', + authorizationState: 'needs_authorization', + }); + expect(mockRegistryInstance.resolveAllowlists).toHaveBeenCalledWith({ + userId: mockUserId, + role: 'user', + }); + expect(mockIsMCPDomainAllowed).toHaveBeenCalledWith( + expect.objectContaining({ url: 'https://blocked.example.com/' }), + ['allowed.example.com'], + null, + ); + expect(findToken).not.toHaveBeenCalled(); + }); + + it('should validate durable readiness against the user-resolved runtime URL', async () => { + const appConnections = new Map(); + const userConnections = new Map(); + const oauthServers = new Set([mockServerName]); + const { findToken } = require('~/models'); + const credentialSetId = 'credential-set-a'; + const config = { + ...mockConfig, + source: 'yaml', + url: 'https://mcp.example.com/users/{{LIBRECHAT_USER_ID}}/mcp', + }; + mockGetOAuthReconnectionManager.mockReturnValue({ isReconnecting: jest.fn(() => false) }); + mockGetFlowStateManager.mockReturnValue({ getFlowState: jest.fn(() => null) }); + mockGetLogStores.mockReturnValue({}); + findToken + .mockResolvedValueOnce({ + expiresAt: new Date(Date.now() + 60000), + metadata: { credential_set_id: credentialSetId }, + }) + .mockResolvedValueOnce({ + token: 'enc:{"client_id":"dynamic-client"}', + metadata: { + credential_set_id: credentialSetId, + authorization_endpoint: 'https://auth.example.com/authorize', + token_endpoint: 'https://auth.example.com/token', + server_url: `https://mcp.example.com/users/${mockUserId}/mcp`, + client_source: 'dynamic', + }, + }); + + const result = await getServerConnectionStatus( + mockUserId, + mockServerName, + config, + appConnections, + userConnections, + oauthServers, + { user: { id: mockUserId } }, + ); + + expect(result).toEqual({ + requiresOAuth: true, + connectionState: 'connected', + authorizationState: 'authorized', + }); + }); + + it('should validate durable readiness against the Graph-resolved runtime URL', async () => { + const appConnections = new Map(); + const userConnections = new Map(); + const oauthServers = new Set([mockServerName]); + const { findToken } = require('~/models'); + const { getGraphApiToken } = require('./GraphTokenService'); + const credentialSetId = 'credential-set-a'; + const config = { + ...mockConfig, + source: 'yaml', + url: 'https://mcp.example.com/{{LIBRECHAT_GRAPH_ACCESS_TOKEN}}/mcp', + }; + const user = { + id: mockUserId, + provider: 'openid', + openidId: 'openid-user', + federatedTokens: { + access_token: 'federated-access-token', + expires_at: Math.floor(Date.now() / 1000) + 3600, + }, + }; + getGraphApiToken.mockResolvedValue({ + access_token: 'resolved-graph-token', + token_type: 'Bearer', + expires_in: 3600, + scope: 'https://graph.microsoft.com/.default', + }); + mockGetOAuthReconnectionManager.mockReturnValue({ isReconnecting: jest.fn(() => false) }); + mockGetFlowStateManager.mockReturnValue({ getFlowState: jest.fn(() => null) }); + mockGetLogStores.mockReturnValue({}); + findToken + .mockResolvedValueOnce({ + expiresAt: new Date(Date.now() + 60000), + metadata: { credential_set_id: credentialSetId }, + }) + .mockResolvedValueOnce({ + token: 'enc:{"client_id":"dynamic-client"}', + metadata: { + credential_set_id: credentialSetId, + authorization_endpoint: 'https://auth.example.com/authorize', + token_endpoint: 'https://auth.example.com/token', + server_url: 'https://mcp.example.com/resolved-graph-token/mcp', + client_source: 'dynamic', + }, + }); + + const result = await getServerConnectionStatus( + mockUserId, + mockServerName, + config, + appConnections, + userConnections, + oauthServers, + { user }, + ); + + expect(result).toEqual({ + requiresOAuth: true, + connectionState: 'connected', + authorizationState: 'authorized', + }); + expect(getGraphApiToken).toHaveBeenCalledWith( + user, + 'federated-access-token', + 'https://graph.microsoft.com/.default', + true, + ); + }); + + it('should not report durable readiness while required custom user variables are missing', async () => { + const appConnections = new Map(); + const userConnections = new Map(); + const oauthServers = new Set([mockServerName]); + const { findToken } = require('~/models'); + const credentialSetId = 'credential-set-a'; + const config = { + ...mockConfig, + url: 'https://mcp.example.com/', + customUserVars: { API_KEY: { title: 'API key' } }, + }; + mockGetOAuthReconnectionManager.mockReturnValue({ isReconnecting: jest.fn(() => false) }); + mockGetFlowStateManager.mockReturnValue({ getFlowState: jest.fn(() => null) }); + mockGetLogStores.mockReturnValue({}); + findToken + .mockResolvedValueOnce({ + expiresAt: new Date(Date.now() + 60000), + metadata: { credential_set_id: credentialSetId }, + }) + .mockResolvedValueOnce({ + token: 'enc:{"client_id":"dynamic-client"}', + metadata: { + credential_set_id: credentialSetId, + authorization_endpoint: 'https://auth.example.com/authorize', + token_endpoint: 'https://auth.example.com/token', + server_url: 'https://mcp.example.com/', + client_source: 'dynamic', + }, + }); + + const result = await getServerConnectionStatus( + mockUserId, + mockServerName, + config, + appConnections, + userConnections, + oauthServers, + { user: { id: mockUserId } }, + ); + + expect(result).toEqual({ + requiresOAuth: true, + connectionState: 'disconnected', + authorizationState: 'needs_authorization', + }); + expect(findToken).not.toHaveBeenCalled(); + }); + + it('should reject stored authorization bound to an older server configuration', async () => { + const appConnections = new Map(); + const userConnections = new Map(); + const oauthServers = new Set([mockServerName]); + const { findToken } = require('~/models'); + const credentialSetId = 'credential-set-a'; + const config = { + updatedAt: Date.now(), + url: 'https://new-mcp.example.com/', + }; + mockGetOAuthReconnectionManager.mockReturnValue({ isReconnecting: jest.fn(() => false) }); + mockGetFlowStateManager.mockReturnValue({ getFlowState: jest.fn(() => null) }); + mockGetLogStores.mockReturnValue({}); + findToken.mockImplementation(({ type }) => { + if (type === 'mcp_oauth') { + return { + expiresAt: new Date(Date.now() + 60000), + metadata: { credential_set_id: credentialSetId }, + }; + } + if (type === 'mcp_oauth_client') { + return { + token: 'enc:{"client_id":"dynamic-client"}', + metadata: { + credential_set_id: credentialSetId, + authorization_endpoint: 'https://auth.example.com/authorize', + token_endpoint: 'https://auth.example.com/token', + server_url: 'https://old-mcp.example.com/', + client_source: 'dynamic', + }, + }; + } + return null; + }); + + const result = await getServerConnectionStatus( + mockUserId, + mockServerName, + config, + appConnections, + userConnections, + oauthServers, + ); + + expect(result).toEqual({ + requiresOAuth: true, + connectionState: 'disconnected', + authorizationState: 'needs_authorization', }); }); @@ -665,6 +1197,7 @@ describe('tests for the new helper functions used by the MCP connection status e expect(result).toEqual({ requiresOAuth: true, connectionState: 'disconnected', + authorizationState: 'needs_authorization', }); }); @@ -691,6 +1224,7 @@ describe('tests for the new helper functions used by the MCP connection status e expect(result).toEqual({ requiresOAuth: true, connectionState: 'connecting', + authorizationState: 'authorizing', }); expect(mockOAuthReconnectionManager.isReconnecting).toHaveBeenCalledWith( mockUserId, @@ -729,6 +1263,7 @@ describe('tests for the new helper functions used by the MCP connection status e expect(result).toEqual({ requiresOAuth: true, connectionState: 'connected', + authorizationState: 'authorized', }); // Should not call flow manager since server is connected @@ -758,6 +1293,7 @@ describe('tests for the new helper functions used by the MCP connection status e expect(result).toEqual({ requiresOAuth: false, connectionState: 'disconnected', + authorizationState: 'not_required', }); // Should not call flow manager since server doesn't require OAuth @@ -951,6 +1487,147 @@ describe('User parameter passing tests', () => { }); describe('createMCPTool', () => { + it.each(['OAuth flow initiated - return early', 'Pending OAuth flow reused - return early'])( + 'preserves runtime-detected OAuth for the internal signal: %s', + async (oauthSignal) => { + const mockUser = { id: 'runtime-oauth-user', role: 'USER' }; + const mockRes = { write: jest.fn(), flush: jest.fn() }; + const { getRoleByName } = require('~/models'); + getRoleByName.mockResolvedValue({ + permissions: { + [PermissionTypes.MCP_SERVERS]: { + [Permissions.USE]: true, + }, + }, + }); + mockGetMCPManager.mockReturnValue({ + callTool: jest.fn().mockRejectedValue(new Error(oauthSignal)), + }); + + 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: {} }, + }, + }, + }, + }); + + await expect( + mcpTool.invoke( + {}, + { + configurable: { user: mockUser }, + metadata: { provider: 'openai', thread_id: 'thread-1', run_id: 'run-1' }, + toolCall: {}, + }, + ), + ).rejects.toThrow( + '[MCP][test-server][test-tool] OAuth authentication required. Please check the server logs for the authentication URL.', + ); + }, + ); + + it('does not label forwarded-token failures as MCP OAuth', async () => { + const mockUser = { id: 'forwarded-token-user', role: 'USER' }; + const mockRes = { write: jest.fn(), flush: jest.fn() }; + const { getRoleByName } = require('~/models'); + getRoleByName.mockResolvedValue({ + permissions: { + [PermissionTypes.MCP_SERVERS]: { + [Permissions.USE]: true, + }, + }, + }); + mockGetMCPManager.mockReturnValue({ + callTool: jest.fn().mockRejectedValue(new Error('Non-200 status code (401)')), + }); + + const mcpTool = await createMCPTool({ + res: mockRes, + user: mockUser, + config: { requiresOAuth: false }, + toolKey: `test-tool${D}test-server`, + provider: 'openai', + userMCPAuthMap: {}, + availableTools: { + [`test-tool${D}test-server`]: { + function: { + description: 'Cached tool', + parameters: { type: 'object', properties: {} }, + }, + }, + }, + }); + + await expect( + mcpTool.invoke( + {}, + { + configurable: { user: mockUser }, + metadata: { provider: 'openai', thread_id: 'thread-1', run_id: 'run-1' }, + toolCall: {}, + }, + ), + ).rejects.toThrow( + '[MCP][test-server][test-tool] upstream authentication failed; MCP OAuth is not configured for this server.', + ); + }); + + it('does not label OBO authentication failures as unconfigured MCP OAuth', async () => { + const mockUser = { id: 'obo-user', role: 'USER' }; + const mockRes = { write: jest.fn(), flush: jest.fn() }; + const { getRoleByName } = require('~/models'); + getRoleByName.mockResolvedValue({ + permissions: { + [PermissionTypes.MCP_SERVERS]: { + [Permissions.USE]: true, + }, + }, + }); + mockGetMCPManager.mockReturnValue({ + callTool: jest.fn().mockRejectedValue(new Error('Non-200 status code (401)')), + }); + + const mcpTool = await createMCPTool({ + res: mockRes, + user: mockUser, + config: { requiresOAuth: false, obo: {} }, + toolKey: `test-tool${D}test-server`, + provider: 'openai', + userMCPAuthMap: {}, + availableTools: { + [`test-tool${D}test-server`]: { + function: { + description: 'Cached tool', + parameters: { type: 'object', properties: {} }, + }, + }, + }, + }); + + await expect( + mcpTool.invoke( + {}, + { + configurable: { user: mockUser }, + metadata: { provider: 'openai', thread_id: 'thread-1', run_id: 'run-1' }, + toolCall: {}, + }, + ), + ).rejects.toThrow( + '[MCP][test-server][test-tool] OAuth authentication required. Please check the server logs for the authentication URL.', + ); + }); + it('should pass user parameter to reinitMCPServer when tool not in cache', async () => { const mockUser = { id: 'test-user-456', email: 'test@example.com' }; const mockRes = { write: jest.fn(), flush: jest.fn() }; diff --git a/api/server/services/Tools/mcp.js b/api/server/services/Tools/mcp.js index 5475a99173..24b8bcd5eb 100644 --- a/api/server/services/Tools/mcp.js +++ b/api/server/services/Tools/mcp.js @@ -63,6 +63,7 @@ async function reinitMCPServer({ let tools = null; let oauthRequired = false; let oauthUrl = null; + let oauthExpiresAt; let ephemeralServer = false; try { @@ -168,9 +169,15 @@ async function reinitMCPServer({ const oauthStart = _oauthStart ?? - (async (authURL) => { + (async (authURL, options) => { logger.info(`[MCP Reinitialize] OAuth URL received for ${serverName}`); + if (authURL !== oauthUrl) { + oauthExpiresAt = undefined; + } oauthUrl = authURL; + if (typeof options?.expiresAt === 'number' && Number.isFinite(options.expiresAt)) { + oauthExpiresAt = options.expiresAt; + } oauthRequired = true; }); @@ -298,6 +305,7 @@ async function reinitMCPServer({ oauthRequired, serverName, oauthUrl, + oauthExpiresAt, tools, }; diff --git a/api/server/services/Tools/mcp.spec.js b/api/server/services/Tools/mcp.spec.js index 7b20cb0d38..d8e788390e 100644 --- a/api/server/services/Tools/mcp.spec.js +++ b/api/server/services/Tools/mcp.spec.js @@ -310,3 +310,40 @@ describe('reinitMCPServer — runtime BODY placeholder pre-check (issue #14074)' expect(result.message).toBe(`Failed to reinitialize MCP server '${serverName}'`); }); }); + +describe('reinitMCPServer — OAuth attempt lifetime', () => { + const user = { id: 'user-123' }; + const serverName = 'Thingy'; + const serverConfig = { + type: 'streamable-http', + url: 'https://thingy.example.com/mcp', + }; + + beforeEach(() => { + jest.clearAllMocks(); + mockUpdateMCPServerTools.mockResolvedValue({}); + }); + + it('returns the expiry supplied when a pending OAuth URL is replayed', async () => { + const expiresAt = Date.now() + 45_000; + mockGetConnection.mockImplementation(async ({ oauthStart }) => { + await oauthStart('https://oauth.example.com/authorize', { expiresAt }); + await oauthStart('https://oauth.example.com/authorize'); + throw new Error('OAuth flow initiated - return early'); + }); + mockDiscoverServerTools.mockResolvedValue({ tools: [], oauthRequired: true, oauthUrl: null }); + + const result = await reinitMCPServer({ + user, + serverName, + serverConfig, + }); + + expect(result).toMatchObject({ + success: true, + oauthRequired: true, + oauthUrl: 'https://oauth.example.com/authorize', + oauthExpiresAt: expiresAt, + }); + }); +}); diff --git a/client/src/hooks/MCP/__tests__/polling.spec.ts b/client/src/hooks/MCP/__tests__/polling.spec.ts new file mode 100644 index 0000000000..8b9ccec481 --- /dev/null +++ b/client/src/hooks/MCP/__tests__/polling.spec.ts @@ -0,0 +1,106 @@ +import { + getMCPOAuthTimeout, + getMCPOAuthPollingOutcome, + isMCPReadyAfterOAuth, + shouldFailMCPOAuthFallback, + isTerminalMCPOAuthPollingError, + shouldUseMCPConnectionStatus, +} from '../polling'; + +describe('getMCPOAuthTimeout', () => { + it('preserves the remaining timeout of a reused flow over the global server window', () => { + expect(getMCPOAuthTimeout(45_000, 600_000)).toBe(45_000); + }); + + it('uses the global server window when the attempt has no explicit timeout', () => { + expect(getMCPOAuthTimeout(undefined, 300_000)).toBe(300_000); + }); +}); + +describe('getMCPOAuthPollingOutcome', () => { + it('treats shared flow completion as terminal success', () => { + expect(getMCPOAuthPollingOutcome({ status: 'COMPLETED', completed: true, failed: false })).toBe( + 'completed', + ); + }); + + it('treats a retained timeout failure as terminal failure', () => { + expect( + getMCPOAuthPollingOutcome({ + status: 'FAILED', + completed: false, + failed: true, + error: 'mcp_oauth flow timed out', + }), + ).toBe('failed'); + }); + + it('keeps polling a pending flow', () => { + expect(getMCPOAuthPollingOutcome({ status: 'PENDING', completed: false, failed: false })).toBe( + 'pending', + ); + }); + + it('treats missing and unauthorized flow records as terminal polling errors', () => { + expect(isTerminalMCPOAuthPollingError({ response: { status: 404 } })).toBe(true); + expect(isTerminalMCPOAuthPollingError({ response: { status: 403 } })).toBe(true); + expect(isTerminalMCPOAuthPollingError({ response: { status: 500 } })).toBe(false); + }); +}); + +describe('shouldUseMCPConnectionStatus', () => { + it('ignores stale cached errors while a live OAuth flow is pending', () => { + expect(shouldUseMCPConnectionStatus('active-flow', false)).toBe(false); + }); + + it('uses durable connection status when no usable flow endpoint remains', () => { + expect(shouldUseMCPConnectionStatus(undefined, false)).toBe(true); + expect(shouldUseMCPConnectionStatus('missing-flow', true)).toBe(true); + }); +}); + +describe('shouldFailMCPOAuthFallback', () => { + it('keeps polling when a fallback pod still reports active authorization', () => { + expect( + shouldFailMCPOAuthFallback(true, { + requiresOAuth: true, + connectionState: 'error', + authorizationState: 'authorizing', + }), + ).toBe(false); + expect( + shouldFailMCPOAuthFallback(true, { + requiresOAuth: true, + connectionState: 'connecting', + authorizationState: 'needs_authorization', + }), + ).toBe(false); + }); + + it('fails a missing flow only after fallback status is no longer authorizing', () => { + expect( + shouldFailMCPOAuthFallback(true, { + requiresOAuth: true, + connectionState: 'disconnected', + authorizationState: 'needs_authorization', + }), + ).toBe(true); + expect(shouldFailMCPOAuthFallback(false, undefined)).toBe(false); + }); +}); + +describe('isMCPReadyAfterOAuth', () => { + const response = { + success: true, + message: 'ready', + serverName: 'test-server', + }; + + it('accepts a successful post-OAuth reinitialization', () => { + expect(isMCPReadyAfterOAuth(response)).toBe(true); + }); + + it('does not treat another OAuth challenge as connection readiness', () => { + expect(isMCPReadyAfterOAuth({ ...response, oauthRequired: true })).toBe(false); + }); +}); diff --git a/client/src/hooks/MCP/polling.ts b/client/src/hooks/MCP/polling.ts new file mode 100644 index 0000000000..325b3bb666 --- /dev/null +++ b/client/src/hooks/MCP/polling.ts @@ -0,0 +1,67 @@ +import type { + MCPServerStatus, + MCPOAuthStatusResponse, + MCPReinitializeResponse, +} from 'librechat-data-provider'; + +export type MCPOAuthPollingOutcome = 'pending' | 'completed' | 'failed'; + +export function getMCPOAuthTimeout( + attemptTimeout: number | undefined, + connectionTimeout: number | undefined, + fallback = 600_000, +): number { + return attemptTimeout ?? connectionTimeout ?? fallback; +} + +/** + * A missing or unauthorized flow is terminal for this browser poll. Retrying it + * forever leaves the OAuth spinner active after the shared flow record is gone. + */ +export function isTerminalMCPOAuthPollingError(error: unknown): boolean { + if (!error || typeof error !== 'object') { + return false; + } + + const responseStatus = (error as { response?: { status?: unknown } }).response?.status; + const status = + typeof responseStatus === 'number' ? responseStatus : (error as { status?: unknown }).status; + return status === 403 || status === 404; +} + +export function getMCPOAuthPollingOutcome(status: MCPOAuthStatusResponse): MCPOAuthPollingOutcome { + if (status.completed || status.status === 'COMPLETED') { + return 'completed'; + } + if (status.failed || status.status === 'FAILED') { + return 'failed'; + } + return 'pending'; +} + +/** An active flow endpoint is newer and more specific than cached connection status. */ +export function shouldUseMCPConnectionStatus( + flowId: string | undefined, + terminalFlowError: boolean, +): boolean { + return !flowId || terminalFlowError; +} + +/** A legacy pod's missing flow route is not terminal while shared fallback state is active. */ +export function shouldFailMCPOAuthFallback( + terminalFlowError: boolean, + serverStatus: MCPServerStatus | undefined, +): boolean { + if (!terminalFlowError) { + return false; + } + return ( + serverStatus?.authorizationState !== 'authorizing' && + serverStatus?.connectionState !== 'connecting' + ); +} + +/** OAuth completion proves credentials were stored; reinitialization proves this request pod can use them. */ +export function isMCPReadyAfterOAuth(response: MCPReinitializeResponse): boolean { + return response.success && response.oauthRequired !== true; +} diff --git a/client/src/hooks/MCP/useMCPServerManager.ts b/client/src/hooks/MCP/useMCPServerManager.ts index ead897c439..6ebeed6632 100644 --- a/client/src/hooks/MCP/useMCPServerManager.ts +++ b/client/src/hooks/MCP/useMCPServerManager.ts @@ -4,6 +4,7 @@ import { useToastContext } from '@librechat/client'; import { useQueryClient } from '@tanstack/react-query'; import { Constants, + dataService, QueryKeys, MCPOptions, Permissions, @@ -21,9 +22,18 @@ import type { TPlugin, MCPServersResponse, MCPConnectionStatusResponse, + MCPOAuthStatusResponse, } from 'librechat-data-provider'; import type { MCPServerInitState } from '~/store/mcp'; import type { ConfigFieldDetail } from '~/common'; +import { + getMCPOAuthTimeout, + getMCPOAuthPollingOutcome, + isMCPReadyAfterOAuth, + shouldFailMCPOAuthFallback, + isTerminalMCPOAuthPollingError, + shouldUseMCPConnectionStatus, +} from './polling'; import { useLocalize, useHasAccess, useMCPSelect, useMCPConnectionStatus } from '~/hooks'; import { useGetStartupConfig, useMCPServersQuery } from '~/data-provider'; import { mcpServerInitStatesAtom, getServerInitState } from '~/store/mcp'; @@ -193,7 +203,7 @@ export function useMCPServerManager({ ); const startServerPolling = useCallback( - (serverName: string) => { + (serverName: string, flowId?: string, initialOAuthTimeout?: number) => { // Prevent duplicate polling for the same server if (pollIntervalsRef.current[serverName]) { console.debug(`[MCP Manager] Polling already active for ${serverName}, skipping duplicate`); @@ -202,6 +212,7 @@ export function useMCPServerManager({ let pollAttempts = 0; let timeoutId: NodeJS.Timeout | null = null; + const pollingStartedAt = Date.now(); /** OAuth can take several minutes if the user steps away from the consent screen. * Poll for the full server-side handling window (MCP_OAUTH_HANDLING_TIMEOUT @@ -214,58 +225,96 @@ export function useMCPServerManager({ return 7500; // Thereafter: every 7.5s }; - /** Honor the server's configured MCP_OAUTH_HANDLING_TIMEOUT (surfaced on the - * connection-status response) so a tuned deadline isn't capped at the default. - * The cache may be empty at start, so this is refreshed from the first status - * refetch below rather than captured once. */ + /** Honor the server's configured MCP_OAUTH_HANDLING_TIMEOUT from the + * reinitialize response. The connection-status cache remains a fallback for + * rolling deployments where the backend does not yet return a flow ID. */ const connectionData = queryClient.getQueryData([ QueryKeys.mcpConnectionStatus, ]); - let oauthTimeoutMs = connectionData?.oauthTimeout ?? 600000; // default 10 minutes + let oauthTimeoutMs = getMCPOAuthTimeout(initialOAuthTimeout, connectionData?.oauthTimeout); // Backstop only; the elapsed-time guard governs. Sized above the worst-case poll count. let maxAttempts = Math.ceil(oauthTimeoutMs / 5000) + 5; + const stopForOAuthTimeout = (elapsedTime: number) => { + console.warn( + `[MCP Manager] OAuth timeout for ${serverName} after ${(elapsedTime / 1000).toFixed(0)}s (attempt ${pollAttempts})`, + ); + showToast({ + message: localize('com_ui_mcp_oauth_timeout', { 0: serverName }), + status: 'error', + }); + if (timeoutId) { + clearTimeout(timeoutId); + } + cleanupServerState(serverName); + }; const pollOnce = async () => { try { pollAttempts++; - const state = getServerInitState(serverInitStates, serverName); + const elapsedTime = Date.now() - pollingStartedAt; - /** Stop polling once the handling window or max attempts is exceeded */ - const elapsedTime = state?.oauthStartTime - ? Date.now() - state.oauthStartTime - : pollAttempts * 5000; // Rough estimate if no start time - - if (pollAttempts > maxAttempts || elapsedTime > oauthTimeoutMs) { - console.warn( - `[MCP Manager] OAuth timeout for ${serverName} after ${(elapsedTime / 1000).toFixed(0)}s (attempt ${pollAttempts})`, - ); - showToast({ - message: localize('com_ui_mcp_oauth_timeout', { 0: serverName }), - status: 'error', - }); - if (timeoutId) { - clearTimeout(timeoutId); + let flowStatus: MCPOAuthStatusResponse | null = null; + let terminalFlowError = false; + if (flowId) { + try { + flowStatus = await dataService.getMCPOAuthStatus(flowId); + } catch (error) { + if (!isTerminalMCPOAuthPollingError(error)) { + throw error; + } + terminalFlowError = true; } - cleanupServerState(serverName); - return; } - - await queryClient.refetchQueries([QueryKeys.mcpConnectionStatus]); + const flowOutcome = flowStatus ? getMCPOAuthPollingOutcome(flowStatus) : null; + const canUseConnectionStatus = shouldUseMCPConnectionStatus(flowId, terminalFlowError); + let isReady = false; + if (flowOutcome === 'completed') { + /** Flow completion is durable credential readiness, not proof that tool discovery + * finished on this pod. Reinitialize once more through the normal API so the + * selected server and its tools are usable before the UI reports success. */ + const readiness = await reinitializeMutation.mutateAsync(serverName); + if (!isMCPReadyAfterOAuth(readiness)) { + showToast({ + message: getMCPReinitializeErrorMessage(readiness, localize), + status: 'error', + }); + if (timeoutId) { + clearTimeout(timeoutId); + } + cleanupServerState(serverName); + return; + } + isReady = true; + } + if (canUseConnectionStatus) { + // The flow record may have expired or been denied on another pod. Re-read + // durable connection state before deciding whether the UI can finish. + await queryClient.refetchQueries([QueryKeys.mcpConnectionStatus]); + } const freshConnectionData = queryClient.getQueryData([ QueryKeys.mcpConnectionStatus, ]); - // Pick up the configured timeout once the status response lands (cache may have - // been empty when polling started), so a tuned deadline is honored mid-flight. + // Pick up the configured timeout when the attempt response had no deadline. + // A reused flow's remaining lifetime must stay authoritative over this global value. if (typeof freshConnectionData?.oauthTimeout === 'number') { - oauthTimeoutMs = freshConnectionData.oauthTimeout; + oauthTimeoutMs = getMCPOAuthTimeout( + initialOAuthTimeout, + freshConnectionData.oauthTimeout, + oauthTimeoutMs, + ); maxAttempts = Math.ceil(oauthTimeoutMs / 5000) + 5; } const freshConnectionStatus = freshConnectionData?.connectionStatus || {}; const serverStatus = freshConnectionStatus[serverName]; + if (canUseConnectionStatus) { + isReady = + serverStatus?.authorizationState === 'authorized' || + serverStatus?.connectionState === 'connected'; + } - if (serverStatus?.connectionState === 'connected') { + if (isReady) { if (timeoutId) { clearTimeout(timeoutId); } @@ -280,7 +329,12 @@ export function useMCPServerManager({ setMCPValues([...currentValues, serverName]); } - await queryClient.invalidateQueries([QueryKeys.mcpTools]); + await Promise.all([ + queryClient.invalidateQueries([QueryKeys.mcpServers]), + queryClient.invalidateQueries([QueryKeys.mcpTools]), + queryClient.invalidateQueries([QueryKeys.mcpAuthValues]), + queryClient.invalidateQueries([QueryKeys.mcpConnectionStatus]), + ]); // This delay is to ensure UI has updated with new connection status before cleanup // Otherwise servers will show as disconnected for a second after OAuth flow completes @@ -290,10 +344,9 @@ export function useMCPServerManager({ return; } - // Check for OAuth timeout (should align with maxAttempts) - if (state?.oauthStartTime && Date.now() - state.oauthStartTime > oauthTimeoutMs) { + if (shouldFailMCPOAuthFallback(terminalFlowError, serverStatus)) { showToast({ - message: localize('com_ui_mcp_oauth_timeout', { 0: serverName }), + message: localize('com_ui_mcp_init_failed'), status: 'error', }); if (timeoutId) { @@ -303,7 +356,30 @@ export function useMCPServerManager({ return; } - if (serverStatus?.connectionState === 'error') { + if (flowOutcome === 'failed') { + showToast({ + message: localize('com_ui_mcp_init_failed'), + status: 'error', + }); + if (timeoutId) { + clearTimeout(timeoutId); + } + cleanupServerState(serverName); + return; + } + + /** Make one final shared-state read before declaring a short reused attempt timed out. */ + if (pollAttempts > maxAttempts || elapsedTime > oauthTimeoutMs) { + stopForOAuthTimeout(elapsedTime); + return; + } + + if ( + !terminalFlowError && + canUseConnectionStatus && + (serverStatus?.authorizationState === 'error' || + serverStatus?.connectionState === 'error') + ) { showToast({ message: localize('com_ui_mcp_init_failed'), status: 'error', @@ -328,12 +404,15 @@ export function useMCPServerManager({ timeoutId = setTimeout(pollOnce, nextInterval); pollIntervalsRef.current[serverName] = timeoutId; } catch (error) { - console.error(`[MCP Manager] Error polling server ${serverName}:`, error); - if (timeoutId) { - clearTimeout(timeoutId); + console.warn(`[MCP Manager] Transient error polling server ${serverName}:`, error); + const elapsedTime = Date.now() - pollingStartedAt; + if (pollAttempts > maxAttempts || elapsedTime > oauthTimeoutMs) { + stopForOAuthTimeout(elapsedTime); + return; } - cleanupServerState(serverName); - return; + const nextInterval = getPollInterval(pollAttempts); + timeoutId = setTimeout(pollOnce, nextInterval); + pollIntervalsRef.current[serverName] = timeoutId; } }; @@ -341,7 +420,7 @@ export function useMCPServerManager({ timeoutId = setTimeout(pollOnce, getPollInterval(0)); pollIntervalsRef.current[serverName] = timeoutId; }, - [queryClient, serverInitStates, showToast, localize, setMCPValues, cleanupServerState], + [queryClient, showToast, localize, setMCPValues, cleanupServerState, reinitializeMutation], ); const initializeServer = useCallback( @@ -378,7 +457,7 @@ export function useMCPServerManager({ window.open(response.oauthUrl, '_blank', 'noopener,noreferrer'); } - startServerPolling(serverName); + startServerPolling(serverName, response.flowId, response.oauthTimeout); } else { await Promise.all([ queryClient.invalidateQueries([QueryKeys.mcpServers]), diff --git a/e2e/specs/mock/mcp-oauth-readiness.spec.ts b/e2e/specs/mock/mcp-oauth-readiness.spec.ts new file mode 100644 index 0000000000..85cdf17b29 --- /dev/null +++ b/e2e/specs/mock/mcp-oauth-readiness.spec.ts @@ -0,0 +1,417 @@ +import { expect, test } from '@playwright/test'; + +const SERVER_NAME = 'e2e-memory'; +const SERVER_TITLE = 'E2E Memory'; +const FLOW_ID = 'e2e-user:e2e-memory'; + +test.describe('MCP OAuth readiness', () => { + test('keeps the server unselected until post-OAuth tool readiness completes', async ({ + page, + }) => { + test.setTimeout(120000); + + let reinitializeCalls = 0; + let flowStatusCalls = 0; + let readinessComplete = false; + let markPendingPolled!: () => void; + let markReadinessStarted!: () => void; + let releaseReadiness!: () => void; + const readinessStarted = new Promise((resolve) => { + markReadinessStarted = resolve; + }); + const pendingPolled = new Promise((resolve) => { + markPendingPolled = resolve; + }); + const readinessGate = new Promise((resolve) => { + releaseReadiness = resolve; + }); + + await page.route('**/api/mcp/connection/status', async (route) => { + await route.fulfill({ + status: 200, + contentType: 'application/json', + body: JSON.stringify({ + success: true, + oauthTimeout: 30000, + connectionStatus: { + [SERVER_NAME]: readinessComplete + ? { + connectionState: 'connected', + requiresOAuth: true, + authorizationState: 'authorized', + } + : { + /** A retry can begin while React Query still holds the previous attempt's + * terminal status. The live PENDING flow must supersede this stale error. */ + connectionState: 'error', + requiresOAuth: true, + authorizationState: 'error', + }, + }, + }), + }); + }); + + await page.route(`**/api/mcp/${SERVER_NAME}/reinitialize`, async (route) => { + reinitializeCalls++; + if (reinitializeCalls === 1) { + await route.fulfill({ + status: 200, + contentType: 'application/json', + body: JSON.stringify({ + success: true, + message: 'OAuth authorization required', + serverName: SERVER_NAME, + oauthRequired: true, + oauthUrl: 'https://oauth.example.test/authorize', + flowId: FLOW_ID, + oauthTimeout: 30000, + }), + }); + return; + } + + markReadinessStarted(); + await readinessGate; + readinessComplete = true; + await route.fulfill({ + status: 200, + contentType: 'application/json', + body: JSON.stringify({ + success: true, + message: 'MCP server reinitialized successfully', + serverName: SERVER_NAME, + oauthRequired: false, + }), + }); + }); + + await page.route('**/api/mcp/oauth/status/**', async (route) => { + flowStatusCalls++; + if (flowStatusCalls === 1) { + markPendingPolled(); + await route.fulfill({ + status: 200, + contentType: 'application/json', + body: JSON.stringify({ status: 'PENDING', completed: false, failed: false }), + }); + return; + } + await route.fulfill({ + status: 200, + contentType: 'application/json', + body: JSON.stringify({ status: 'COMPLETED', completed: true, failed: false }), + }); + }); + + await page.goto('/c/new', { timeout: 10000 }); + await page.getByRole('button', { name: 'MCP Servers', exact: true }).click(); + const serverItem = page.getByRole('menuitemcheckbox', { name: new RegExp(SERVER_TITLE) }); + await expect(serverItem).toHaveAttribute('aria-checked', 'false'); + await serverItem.getByRole('button', { name: `Connect ${SERVER_NAME}` }).click(); + + await page.getByRole('button', { name: 'Authenticate', exact: true }).click(); + await expect(page.getByRole('button', { name: 'Continue with OAuth' })).toBeVisible(); + await pendingPolled; + + await page.keyboard.press('Escape'); + await page.getByRole('button', { name: 'MCP Servers', exact: true }).click(); + await expect(serverItem.getByRole('button', { name: 'Cancel' })).toBeVisible(); + await expect(page.getByText('Failed to initialize MCP server')).toHaveCount(0); + + await readinessStarted; + + await page.keyboard.press('Escape'); + await page.getByRole('button', { name: 'MCP Servers', exact: true }).click(); + await expect(serverItem).toHaveAttribute('aria-checked', 'false'); + await expect(serverItem.getByRole('button', { name: 'Cancel' })).toBeVisible(); + await expect( + page.getByText(`MCP server '${SERVER_NAME}' authenticated successfully`), + ).toHaveCount(0); + + releaseReadiness(); + + await expect( + page.getByText(`MCP server '${SERVER_NAME}' authenticated successfully`).first(), + ).toBeVisible(); + await expect(serverItem).toHaveAttribute('aria-checked', 'true'); + expect(reinitializeCalls).toBe(2); + }); + + test('stops a reused OAuth spinner at the attempt remaining lifetime', async ({ page }) => { + test.setTimeout(30000); + + let flowStatusCalls = 0; + await page.route('**/api/mcp/connection/status', async (route) => { + await route.fulfill({ + status: 200, + contentType: 'application/json', + body: JSON.stringify({ + success: true, + oauthTimeout: 30000, + connectionStatus: { + [SERVER_NAME]: { + connectionState: 'error', + requiresOAuth: true, + authorizationState: 'error', + }, + }, + }), + }); + }); + await page.route(`**/api/mcp/${SERVER_NAME}/reinitialize`, async (route) => { + await route.fulfill({ + status: 200, + contentType: 'application/json', + body: JSON.stringify({ + success: true, + message: 'OAuth authorization required', + serverName: SERVER_NAME, + oauthRequired: true, + oauthUrl: 'https://oauth.example.test/authorize', + flowId: FLOW_ID, + oauthTimeout: 5500, + }), + }); + }); + await page.route('**/api/mcp/oauth/status/**', async (route) => { + flowStatusCalls++; + await route.fulfill({ + status: 200, + contentType: 'application/json', + body: JSON.stringify({ status: 'PENDING', completed: false, failed: false }), + }); + }); + + await page.goto('/c/new', { timeout: 10000 }); + await page.getByRole('button', { name: 'MCP Servers', exact: true }).click(); + const serverItem = page.getByRole('menuitemcheckbox', { name: new RegExp(SERVER_TITLE) }); + await serverItem.getByRole('button', { name: `Connect ${SERVER_NAME}` }).click(); + await page.getByRole('button', { name: 'Authenticate', exact: true }).click(); + + await expect(page.getByText(`OAuth login timed out for ${SERVER_NAME}`).first()).toBeVisible({ + timeout: 15000, + }); + expect(flowStatusCalls).toBe(2); + + await page.keyboard.press('Escape'); + await page.getByRole('button', { name: 'MCP Servers', exact: true }).click(); + await expect(serverItem.getByRole('button', { name: `Connect ${SERVER_NAME}` })).toBeVisible(); + }); + + test('accepts completion found by the final poll after the attempt deadline', async ({ + page, + }) => { + test.setTimeout(30000); + + let reinitializeCalls = 0; + let flowStatusCalls = 0; + await page.route('**/api/mcp/connection/status', async (route) => { + const readinessComplete = reinitializeCalls > 1; + await route.fulfill({ + status: 200, + contentType: 'application/json', + body: JSON.stringify({ + success: true, + oauthTimeout: 30000, + connectionStatus: { + [SERVER_NAME]: { + connectionState: readinessComplete ? 'connected' : 'error', + requiresOAuth: true, + authorizationState: readinessComplete ? 'authorized' : 'error', + }, + }, + }), + }); + }); + await page.route(`**/api/mcp/${SERVER_NAME}/reinitialize`, async (route) => { + reinitializeCalls++; + await route.fulfill({ + status: 200, + contentType: 'application/json', + body: JSON.stringify( + reinitializeCalls === 1 + ? { + success: true, + message: 'OAuth authorization required', + serverName: SERVER_NAME, + oauthRequired: true, + oauthUrl: 'https://oauth.example.test/authorize', + flowId: FLOW_ID, + oauthTimeout: 1000, + } + : { + success: true, + message: 'MCP server reinitialized successfully', + serverName: SERVER_NAME, + oauthRequired: false, + }, + ), + }); + }); + await page.route('**/api/mcp/oauth/status/**', async (route) => { + flowStatusCalls++; + await route.fulfill({ + status: 200, + contentType: 'application/json', + body: JSON.stringify({ status: 'COMPLETED', completed: true, failed: false }), + }); + }); + + await page.goto('/c/new', { timeout: 10000 }); + await page.getByRole('button', { name: 'MCP Servers', exact: true }).click(); + const serverItem = page.getByRole('menuitemcheckbox', { name: new RegExp(SERVER_TITLE) }); + await serverItem.getByRole('button', { name: `Connect ${SERVER_NAME}` }).click(); + await page.getByRole('button', { name: 'Authenticate', exact: true }).click(); + + await expect( + page.getByText(`MCP server '${SERVER_NAME}' authenticated successfully`).first(), + ).toBeVisible({ timeout: 15000 }); + await expect(page.getByText(`OAuth login timed out for ${SERVER_NAME}`)).toHaveCount(0); + await expect( + page.getByRole('menuitemcheckbox', { + name: new RegExp(SERVER_TITLE), + includeHidden: true, + }), + ).toHaveAttribute('aria-checked', 'true'); + expect(flowStatusCalls).toBe(1); + expect(reinitializeCalls).toBe(2); + }); + + test('stops polling at the attempt deadline during repeated transient errors', async ({ + page, + }) => { + test.setTimeout(30000); + + let flowStatusCalls = 0; + await page.route('**/api/mcp/connection/status', async (route) => { + await route.fulfill({ + status: 200, + contentType: 'application/json', + body: JSON.stringify({ + success: true, + oauthTimeout: 30000, + connectionStatus: { + [SERVER_NAME]: { + connectionState: 'error', + requiresOAuth: true, + authorizationState: 'error', + }, + }, + }), + }); + }); + await page.route(`**/api/mcp/${SERVER_NAME}/reinitialize`, async (route) => { + await route.fulfill({ + status: 200, + contentType: 'application/json', + body: JSON.stringify({ + success: true, + message: 'OAuth authorization required', + serverName: SERVER_NAME, + oauthRequired: true, + oauthUrl: 'https://oauth.example.test/authorize', + flowId: FLOW_ID, + oauthTimeout: 1000, + }), + }); + }); + await page.route('**/api/mcp/oauth/status/**', async (route) => { + flowStatusCalls++; + await route.fulfill({ + status: 503, + contentType: 'application/json', + body: JSON.stringify({ error: 'Temporary shared-state failure' }), + }); + }); + + await page.goto('/c/new', { timeout: 10000 }); + await page.getByRole('button', { name: 'MCP Servers', exact: true }).click(); + const serverItem = page.getByRole('menuitemcheckbox', { name: new RegExp(SERVER_TITLE) }); + await serverItem.getByRole('button', { name: `Connect ${SERVER_NAME}` }).click(); + await page.getByRole('button', { name: 'Authenticate', exact: true }).click(); + + await expect(page.getByText(`OAuth login timed out for ${SERVER_NAME}`).first()).toBeVisible({ + timeout: 15000, + }); + expect(flowStatusCalls).toBe(1); + }); + + test('keeps polling when an older fallback pod still reports authorization in progress', async ({ + page, + }) => { + test.setTimeout(30000); + + let flowStatusCalls = 0; + await page.route('**/api/mcp/connection/status', async (route) => { + let serverStatus = { + connectionState: 'connected', + requiresOAuth: true, + authorizationState: 'authorized', + }; + if (flowStatusCalls === 0) { + serverStatus = { + connectionState: 'error', + requiresOAuth: true, + authorizationState: 'error', + }; + } else if (flowStatusCalls === 1) { + serverStatus = { + connectionState: 'error', + requiresOAuth: true, + authorizationState: 'authorizing', + }; + } + await route.fulfill({ + status: 200, + contentType: 'application/json', + body: JSON.stringify({ + success: true, + oauthTimeout: 30000, + connectionStatus: { [SERVER_NAME]: serverStatus }, + }), + }); + }); + await page.route(`**/api/mcp/${SERVER_NAME}/reinitialize`, async (route) => { + await route.fulfill({ + status: 200, + contentType: 'application/json', + body: JSON.stringify({ + success: true, + message: 'OAuth authorization required', + serverName: SERVER_NAME, + oauthRequired: true, + oauthUrl: 'https://oauth.example.test/authorize', + flowId: FLOW_ID, + oauthTimeout: 30000, + }), + }); + }); + await page.route('**/api/mcp/oauth/status/**', async (route) => { + flowStatusCalls++; + await route.fulfill({ + status: 404, + contentType: 'application/json', + body: JSON.stringify({ error: 'Route not found' }), + }); + }); + + await page.goto('/c/new', { timeout: 10000 }); + await page.getByRole('button', { name: 'MCP Servers', exact: true }).click(); + const serverItem = page.getByRole('menuitemcheckbox', { name: new RegExp(SERVER_TITLE) }); + await serverItem.getByRole('button', { name: `Connect ${SERVER_NAME}` }).click(); + await page.getByRole('button', { name: 'Authenticate', exact: true }).click(); + + await page.keyboard.press('Escape'); + await page.getByRole('button', { name: 'MCP Servers', exact: true }).click(); + await expect(serverItem.getByRole('button', { name: 'Cancel' })).toBeVisible({ + timeout: 8000, + }); + await expect(page.getByText('Failed to initialize MCP server')).toHaveCount(0); + await expect( + page.getByText(`MCP server '${SERVER_NAME}' authenticated successfully`).first(), + ).toBeVisible({ timeout: 20000 }); + await expect(serverItem).toHaveAttribute('aria-checked', 'true'); + expect(flowStatusCalls).toBeGreaterThanOrEqual(2); + }); +}); diff --git a/packages/api/src/flow/manager.test.ts b/packages/api/src/flow/manager.test.ts index d955f98f4e..2e89effbe1 100644 --- a/packages/api/src/flow/manager.test.ts +++ b/packages/api/src/flow/manager.test.ts @@ -104,6 +104,29 @@ describe('FlowStateManager', () => { await expect(flowPromise).rejects.toThrow('test-type flow timed out'); }); + it('should retain a terminal timeout for the remaining storage TTL', async () => { + const flowId = 'retained-timeout-flow'; + const type = 'mcp_oauth'; + const shortTimeoutManager = new FlowStateManager(store as unknown as Keyv, { + ttl: 5000, + monitorTimeout: 100, + retainedFailureTypes: ['mcp_oauth'], + ci: true, + }); + + await expect(shortTimeoutManager.createFlow(flowId, type)).rejects.toThrow( + 'mcp_oauth flow timed out', + ); + + await expect(shortTimeoutManager.getFlowState(flowId, type)).resolves.toEqual( + expect.objectContaining({ + status: 'FAILED', + error: 'mcp_oauth flow timed out', + failedAt: expect.any(Number), + }), + ); + }); + it('should maintain flow state consistency under high concurrency', async () => { const flowId = 'concurrent-flow'; const type = 'test-type'; @@ -174,6 +197,26 @@ describe('FlowStateManager', () => { await flowManager.failFlow(flowId, type, new Error('failure')); await expect(flowPromise).rejects.toThrow('failure'); + await expect(flowManager.getFlowState(flowId, type)).resolves.toBeUndefined(); + }, 15000); + + it('should retain configured failed flow types for status polling', async () => { + const flowId = 'retained-failure-flow'; + const type = 'mcp_oauth'; + const retainedManager = new FlowStateManager(store as unknown as Keyv, { + ttl: 5000, + retainedFailureTypes: [type], + ci: true, + }); + const flowPromise = retainedManager.createFlow(flowId, type); + + await new Promise((resolve) => setTimeout(resolve, 500)); + await retainedManager.failFlow(flowId, type, new Error('provider rejected request')); + + await expect(flowPromise).rejects.toThrow('provider rejected request'); + await expect(retainedManager.getFlowState(flowId, type)).resolves.toEqual( + expect.objectContaining({ status: 'FAILED', error: 'provider rejected request' }), + ); }, 15000); it('should not overwrite a completed flow with a late failure', async () => { diff --git a/packages/api/src/flow/manager.ts b/packages/api/src/flow/manager.ts index bbf335fd97..ef52367ee3 100644 --- a/packages/api/src/flow/manager.ts +++ b/packages/api/src/flow/manager.ts @@ -28,19 +28,23 @@ export function normalizeExpiresAt(timestamp: number): number { export class FlowStateManager { private keyv: Keyv; private ttl: number; + private monitorTimeout: number; + private retainedFailureTypes: Set; private intervals: Set; constructor(store: Keyv, options?: FlowManagerOptions) { if (!options) { options = { ttl: 60000 * 3 }; } - const { ci = false, ttl } = options; + const { ci = false, ttl, monitorTimeout = ttl, retainedFailureTypes = [] } = options; if (!ci && !(store instanceof Keyv)) { throw new Error('Invalid store provided to FlowStateManager'); } this.ttl = ttl; + this.monitorTimeout = monitorTimeout; + this.retainedFailureTypes = new Set(retainedFailureTypes); this.keyv = store; this.intervals = new Set(); @@ -145,7 +149,6 @@ export class FlowStateManager { private monitorFlow(flowKey: string, type: string, signal?: AbortSignal): Promise { return new Promise((resolve, reject) => { const checkInterval = 2000; - let elapsedTime = 0; let isCleanedUp = false; let intervalId: NodeJS.Timeout | null = null; let missingStateRetried = false; @@ -230,20 +233,35 @@ export class FlowStateManager { if (flowState.status === 'COMPLETED' && flowState.result !== undefined) { resolve(flowState.result); } else if (flowState.status === 'FAILED') { - await this.keyv.delete(flowKey); + if (!this.retainedFailureTypes.has(type)) { + await this.keyv.delete(flowKey); + } reject(new Error(flowState.error ?? `${type} flow failed`)); } return; } - elapsedTime += checkInterval; - if (elapsedTime >= this.ttl) { + const elapsedTime = Date.now() - flowState.createdAt; + if (elapsedTime >= this.monitorTimeout) { cleanup(); logger.error( - `[${flowKey}] Flow timed out | Elapsed time: ${elapsedTime} | TTL: ${this.ttl}`, + `[${flowKey}] Flow timed out | Elapsed time: ${elapsedTime} | Timeout: ${this.monitorTimeout}`, ); - await this.keyv.delete(flowKey); - reject(new Error(`${type} flow timed out`)); + const message = `${type} flow timed out`; + if (this.retainedFailureTypes.has(type)) { + const remainingTtl = Math.max(1, this.ttl - elapsedTime); + const timedOutState: FlowState = { + ...flowState, + status: 'FAILED', + error: message, + failedAt: Date.now(), + }; + await this.keyv.set(flowKey, timedOutState, remainingTtl); + } else { + await this.keyv.delete(flowKey); + } + reject(new Error(message)); + return; } logger.debug(`[${flowKey}] Flow state elapsed time: ${elapsedTime}, checking again...`); } catch (error) { diff --git a/packages/api/src/flow/types.ts b/packages/api/src/flow/types.ts index 491dcb7896..3a4d295f67 100644 --- a/packages/api/src/flow/types.ts +++ b/packages/api/src/flow/types.ts @@ -18,6 +18,10 @@ export interface FlowState { export interface FlowManagerOptions { ttl: number; + /** Maximum time a flow may remain PENDING. Defaults to the storage TTL. */ + monitorTimeout?: number; + /** Flow types whose FAILED state should remain readable until the storage TTL expires. */ + retainedFailureTypes?: readonly string[]; ci?: boolean; logger?: Logger; } diff --git a/packages/api/src/mcp/__tests__/MCPConnectionFetchTools.test.ts b/packages/api/src/mcp/__tests__/MCPConnectionFetchTools.test.ts index d0aec66fc9..212f4dc498 100644 --- a/packages/api/src/mcp/__tests__/MCPConnectionFetchTools.test.ts +++ b/packages/api/src/mcp/__tests__/MCPConnectionFetchTools.test.ts @@ -296,3 +296,19 @@ describe('MCPConnection.fetchTools pagination', () => { expect(mockLogger.error).toHaveBeenCalledWith(expect.stringContaining('Failed to fetch tools')); }); }); + +describe('MCPConnection.usesOAuth', () => { + it.each([ + [{ type: 'streamable-http', url: 'https://example.com/mcp', requiresOAuth: true }, true], + [{ type: 'streamable-http', url: 'https://example.com/mcp', oauth: {} }, true], + [{ type: 'streamable-http', url: 'https://example.com/mcp', requiresOAuth: false }, false], + ] as const)('reports OAuth from the resolved connection config', (serverConfig, expected) => { + const connection = new MCPConnection({ + serverName: 'oauth-status-test', + serverConfig, + useSSRFProtection: false, + }); + + expect(connection.usesOAuth()).toBe(expected); + }); +}); diff --git a/packages/api/src/mcp/__tests__/MCPFlowRedis.cache_integration.spec.ts b/packages/api/src/mcp/__tests__/MCPFlowRedis.cache_integration.spec.ts new file mode 100644 index 0000000000..5183e731a3 --- /dev/null +++ b/packages/api/src/mcp/__tests__/MCPFlowRedis.cache_integration.spec.ts @@ -0,0 +1,79 @@ +import { randomUUID } from 'crypto'; +import type { Keyv } from 'keyv'; +import { keyvRedisClient, keyvRedisClientReady } from '~/cache/redisClients'; +import { closeRedisClients } from '~/cache/__tests__/redisClients.helper'; +import { standardCache } from '~/cache/cacheFactory'; +import { FlowStateManager } from '~/flow/manager'; + +const FLOW_TYPE = 'mcp_oauth'; +const FLOW_TTL = 30_000; + +describe('MCP OAuth flow state across Redis-backed instances', () => { + let podAStore: Keyv; + let podBStore: Keyv; + let podA: FlowStateManager; + let podB: FlowStateManager; + const flowIds = new Set(); + + beforeAll(async () => { + if (!keyvRedisClient || !keyvRedisClientReady) { + throw new Error('MCP cross-instance flow tests require a real Redis client'); + } + await keyvRedisClientReady; + const namespace = `MCPFlowRedis-${process.pid}-${randomUUID()}`; + podAStore = standardCache(namespace, FLOW_TTL); + podBStore = standardCache(namespace, FLOW_TTL); + podA = new FlowStateManager(podAStore, { + ttl: FLOW_TTL, + retainedFailureTypes: [FLOW_TYPE], + ci: true, + }); + podB = new FlowStateManager(podBStore, { + ttl: FLOW_TTL, + retainedFailureTypes: [FLOW_TYPE], + ci: true, + }); + }); + + afterEach(async () => { + await Promise.all([...flowIds].map((flowId) => podA.deleteFlow(flowId, FLOW_TYPE))); + flowIds.clear(); + }); + + afterAll(async () => { + await closeRedisClients(); + }); + + function createFlowId(): string { + const flowId = randomUUID(); + flowIds.add(flowId); + return flowId; + } + + it('shares pending and completed OAuth state between pod instances', async () => { + const flowId = createFlowId(); + + await podA.initFlow(flowId, FLOW_TYPE, { authorizationUrl: 'https://oauth.example/auth' }); + await expect(podB.getFlowState(flowId, FLOW_TYPE)).resolves.toEqual( + expect.objectContaining({ status: 'PENDING' }), + ); + + await podB.completeFlow(flowId, FLOW_TYPE, 'authorized'); + await expect(podA.getFlowState(flowId, FLOW_TYPE)).resolves.toEqual( + expect.objectContaining({ status: 'COMPLETED', result: 'authorized' }), + ); + }); + + it('retains a terminal OAuth failure for polling on another pod', async () => { + const flowId = createFlowId(); + + await podA.initFlow(flowId, FLOW_TYPE); + const waiter = podA.createFlow(flowId, FLOW_TYPE); + await podB.failFlow(flowId, FLOW_TYPE, new Error('provider rejected request')); + + await expect(waiter).rejects.toThrow('provider rejected request'); + await expect(podB.getFlowState(flowId, FLOW_TYPE)).resolves.toEqual( + expect.objectContaining({ status: 'FAILED', error: 'provider rejected request' }), + ); + }); +}); diff --git a/packages/api/src/mcp/__tests__/MCPOAuthTokenStorage.test.ts b/packages/api/src/mcp/__tests__/MCPOAuthTokenStorage.test.ts index eb75db12de..451d782de7 100644 --- a/packages/api/src/mcp/__tests__/MCPOAuthTokenStorage.test.ts +++ b/packages/api/src/mcp/__tests__/MCPOAuthTokenStorage.test.ts @@ -50,6 +50,178 @@ describe('MCPTokenStorage', () => { jest.clearAllMocks(); }); + describe('hasStoredAuthorization', () => { + const validateClientBinding = () => undefined; + + async function storeClient( + metadata: Record = storedBindingMetadata, + clientInfo = { client_id: 'dynamic-client' }, + ) { + await store.createToken({ + userId: 'u1', + type: 'mcp_oauth_client', + identifier: 'mcp:srv1:client', + token: `enc:${JSON.stringify(clientInfo)}`, + expiresIn: 3600, + metadata, + }); + } + + it('accepts a current access token bound to its stored client generation', async () => { + await createBoundToken(store, { + userId: 'u1', + type: 'mcp_oauth', + identifier: 'mcp:srv1', + token: 'enc:access-token', + expiresIn: 3600, + }); + await storeClient(); + + await expect( + MCPTokenStorage.hasStoredAuthorization({ + userId: 'u1', + serverName: 'srv1', + findToken: store.findToken, + validateClientBinding, + }), + ).resolves.toBe(true); + }); + + it('rejects legacy credentials without binding metadata', async () => { + await store.createToken({ + userId: 'u1', + type: 'mcp_oauth', + identifier: 'mcp:srv1', + token: 'enc:legacy-access-token', + expiresIn: 3600, + }); + await storeClient({}); + + await expect( + MCPTokenStorage.hasStoredAuthorization({ + userId: 'u1', + serverName: 'srv1', + findToken: store.findToken, + validateClientBinding, + }), + ).resolves.toBe(false); + }); + + it('accepts an expired access token only when a bound refresh token remains usable', async () => { + await createBoundToken(store, { + userId: 'u1', + type: 'mcp_oauth', + identifier: 'mcp:srv1', + token: 'enc:expired-access-token', + expiresIn: -60, + }); + await createBoundToken(store, { + userId: 'u1', + type: 'mcp_oauth_refresh', + identifier: 'mcp:srv1:refresh', + token: 'enc:refresh-token', + expiresIn: 3600, + }); + await storeClient(); + + await expect( + MCPTokenStorage.hasStoredAuthorization({ + userId: 'u1', + serverName: 'srv1', + findToken: store.findToken, + validateClientBinding, + }), + ).resolves.toBe(true); + }); + + it('accepts a bound refresh token after the expired access-token record is removed', async () => { + await createBoundToken(store, { + userId: 'u1', + type: 'mcp_oauth_refresh', + identifier: 'mcp:srv1:refresh', + token: 'enc:refresh-token', + expiresIn: 3600, + }); + await storeClient(); + + await expect( + MCPTokenStorage.hasStoredAuthorization({ + userId: 'u1', + serverName: 'srv1', + findToken: store.findToken, + validateClientBinding, + }), + ).resolves.toBe(true); + }); + + it('rejects a refresh-only credential from a different client generation', async () => { + await createBoundToken(store, { + userId: 'u1', + type: 'mcp_oauth_refresh', + identifier: 'mcp:srv1:refresh', + token: 'enc:refresh-token', + expiresIn: 3600, + }); + await storeClient({ ...storedBindingMetadata, credential_set_id: 'different-generation' }); + + await expect( + MCPTokenStorage.hasStoredAuthorization({ + userId: 'u1', + serverName: 'srv1', + findToken: store.findToken, + validateClientBinding, + }), + ).resolves.toBe(false); + }); + + it('rejects an expired access token without a usable refresh token', async () => { + await createBoundToken(store, { + userId: 'u1', + type: 'mcp_oauth', + identifier: 'mcp:srv1', + token: 'enc:expired-access-token', + expiresIn: -60, + }); + await storeClient(); + + await expect( + MCPTokenStorage.hasStoredAuthorization({ + userId: 'u1', + serverName: 'srv1', + findToken: store.findToken, + validateClientBinding, + }), + ).resolves.toBe(false); + }); + + it('rejects usable credentials that are bound to an older server configuration', async () => { + await createBoundToken(store, { + userId: 'u1', + type: 'mcp_oauth', + identifier: 'mcp:srv1', + token: 'enc:access-token', + expiresIn: 3600, + }); + await storeClient({ + ...storedBindingMetadata, + server_url: 'https://old-mcp.example.com/', + }); + + await expect( + MCPTokenStorage.hasStoredAuthorization({ + userId: 'u1', + serverName: 'srv1', + findToken: store.findToken, + validateClientBinding: (_clientInfo, storedMetadata) => { + if (storedMetadata.server_url !== 'https://new-mcp.example.com/') { + throw new Error('stored server binding changed'); + } + }, + }), + ).resolves.toBe(false); + }); + }); + describe('isCurrentAccessToken', () => { it('rejects a flow-cached token after persistent storage has rotated it', async () => { await createBoundToken(store, { diff --git a/packages/api/src/mcp/connection.ts b/packages/api/src/mcp/connection.ts index c479ea4bf0..121ec7480d 100644 --- a/packages/api/src/mcp/connection.ts +++ b/packages/api/src/mcp/connection.ts @@ -21,9 +21,9 @@ import type { Transport } from '@modelcontextprotocol/sdk/shared/transport.js'; import type { MCPOAuthTokens } from './oauth/types'; import type * as t from './types'; import { createSSRFSafeUndiciConnect, isSSRFTarget, resolveHostnameSSRF } from '~/auth'; +import { isOAuthServer, sanitizeUrlForLogging } from './utils'; import { runOutsideTracing } from '~/utils/tracing'; import { isAddressAllowed } from '~/auth/domain'; -import { sanitizeUrlForLogging } from './utils'; import { withTimeout } from '~/utils/promise'; import { mcpConfig } from './mcpConfig'; @@ -2463,6 +2463,11 @@ export class MCPConnection extends EventEmitter { this.oauthTokens = tokens; } + /** Whether this connection's resolved runtime config uses MCP OAuth. */ + public usesOAuth(): boolean { + return isOAuthServer(this.options); + } + /** * Check if this connection is stale compared to config update time. * A connection is stale if it was created before the config was updated. diff --git a/packages/api/src/mcp/oauth/tokens.ts b/packages/api/src/mcp/oauth/tokens.ts index b4dd633cc4..3419a68211 100644 --- a/packages/api/src/mcp/oauth/tokens.ts +++ b/packages/api/src/mcp/oauth/tokens.ts @@ -164,6 +164,81 @@ export class MCPTokenStorage { : `[MCP][User: ${userId}][${serverName}]`; } + /** Returns whether storage contains a currently usable, generation-bound authorization. */ + static async hasStoredAuthorization({ + userId, + serverName, + findToken, + validateClientBinding, + }: { + userId: string; + serverName: string; + findToken: TokenMethods['findToken']; + validateClientBinding: ( + clientInfo: OAuthClientInformation, + storedMetadata: Partial, + ) => void; + }): Promise { + const identifier = `mcp:${serverName}`; + try { + const [accessTokenData, clientInfoData] = await Promise.all([ + findToken({ userId, type: 'mcp_oauth', identifier }), + findToken({ userId, type: 'mcp_oauth_client', identifier: `${identifier}:client` }), + ]); + const clientCredentialSetId = getCredentialSetId(clientInfoData); + const accessCredentialSetId = getCredentialSetId(accessTokenData); + let hasUsableAuthorization = false; + if (accessTokenData) { + if (!accessCredentialSetId || clientCredentialSetId !== accessCredentialSetId) { + return false; + } + if (!accessTokenData.expiresAt || accessTokenData.expiresAt > new Date()) { + hasUsableAuthorization = true; + } + } + + if (!hasUsableAuthorization) { + const refreshTokenData = await findToken({ + userId, + type: 'mcp_oauth_refresh', + identifier: `${identifier}:refresh`, + }); + const refreshCredentialSetId = getCredentialSetId(refreshTokenData); + hasUsableAuthorization = + !!refreshCredentialSetId && + refreshCredentialSetId === clientCredentialSetId && + (!accessCredentialSetId || refreshCredentialSetId === accessCredentialSetId) && + (!refreshTokenData?.expiresAt || refreshTokenData.expiresAt > new Date()); + } + + if (!hasUsableAuthorization || !clientInfoData?.token) { + return false; + } + + const clientInfo = JSON.parse( + await decryptV2(clientInfoData.token), + ) as OAuthClientInformation; + try { + validateClientBinding(clientInfo, getTokenMetadata(clientInfoData)); + } catch (error) { + logger.debug( + `${this.getLogPrefix(userId, serverName)} Stored authorization no longer matches the configured OAuth binding`, + { error }, + ); + return false; + } + return true; + } catch (error) { + logger.warn( + `${this.getLogPrefix(userId, serverName)} Failed to inspect stored authorization`, + { + error, + }, + ); + return false; + } + } + /** * Confirms a flow-cached access token is still the token in persistent storage. This prevents * an old `mcp_get_tokens` result from being paired with newer client-binding metadata. diff --git a/packages/data-provider/src/api-endpoints.ts b/packages/data-provider/src/api-endpoints.ts index e0ba1320ae..aae3b03fcd 100644 --- a/packages/data-provider/src/api-endpoints.ts +++ b/packages/data-provider/src/api-endpoints.ts @@ -230,6 +230,9 @@ export const cancelMCPOAuth = (serverName: string) => { return `${BASE_URL}/api/mcp/oauth/cancel/${serverName}`; }; +export const mcpOAuthStatus = (flowId: string) => + `${BASE_URL}/api/mcp/oauth/status/${encodeURIComponent(flowId)}`; + export const mcpOAuthBind = (serverName: string) => `${BASE_URL}/api/mcp/${serverName}/oauth/bind`; export const actionOAuthBind = (actionId: string) => diff --git a/packages/data-provider/src/data-service.ts b/packages/data-provider/src/data-service.ts index 4db596a867..f046f9c1b7 100644 --- a/packages/data-provider/src/data-service.ts +++ b/packages/data-provider/src/data-service.ts @@ -251,6 +251,10 @@ export function cancelMCPOAuth(serverName: string): Promise { + return request.get(endpoints.mcpOAuthStatus(flowId)); +} + /* Config */ export type StartupConfigOptions = { diff --git a/packages/data-provider/src/types/mcpServers.ts b/packages/data-provider/src/types/mcpServers.ts index 5aa45ab242..c2b1d3d833 100644 --- a/packages/data-provider/src/types/mcpServers.ts +++ b/packages/data-provider/src/types/mcpServers.ts @@ -60,9 +60,20 @@ export interface MCPReinitializeResponse { serverName: string; oauthRequired?: boolean; oauthUrl?: string | null; + /** Shared OAuth attempt identifier used to poll durable flow state. */ + flowId?: string; + /** Remaining OAuth completion window for this attempt, in milliseconds. */ + oauthTimeout?: number; failureReason?: MCPReinitializeFailureReason; missingUserVars?: string[]; /** True when the server uses request-scoped placeholders and the connection * was deferred to the next chat turn (tools are not enumerable up front). */ connectionDeferred?: boolean; } + +export interface MCPOAuthStatusResponse { + status: 'PENDING' | 'COMPLETED' | 'FAILED'; + completed: boolean; + failed: boolean; + error?: string; +} diff --git a/packages/data-provider/src/types/queries.ts b/packages/data-provider/src/types/queries.ts index b7c968c4e4..029b5f51fd 100644 --- a/packages/data-provider/src/types/queries.ts +++ b/packages/data-provider/src/types/queries.ts @@ -200,6 +200,12 @@ export type ListRolesResponse = { export interface MCPServerStatus { requiresOAuth: boolean; connectionState: 'disconnected' | 'connecting' | 'connected' | 'error'; + authorizationState?: + | 'not_required' + | 'authorizing' + | 'authorized' + | 'needs_authorization' + | 'error'; } export interface MCPConnectionStatusResponse { @@ -214,6 +220,7 @@ export interface MCPServerConnectionStatusResponse { serverName: string; requiresOAuth: boolean; connectionStatus: 'disconnected' | 'connecting' | 'connected' | 'error'; + authorizationState?: MCPServerStatus['authorizationState']; } export interface MCPAuthValuesResponse {