From b443254151a2f96d791fc47441c5ce70d7966341 Mon Sep 17 00:00:00 2001 From: Sean McGrath Date: Tue, 11 Nov 2025 13:51:20 +1300 Subject: [PATCH] =?UTF-8?q?=F0=9F=94=90=20fix:=20persist=20new=20MCP=20oau?= =?UTF-8?q?th=20tokens=20properly=20(#10439)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix: re-fetch OAuth flow state after completeOAuthFlow * test: add tests for MCP OAuth flow state bugs --- api/server/routes/__tests__/mcp.spec.js | 87 ++++++++++++++++++++ api/server/routes/mcp.js | 7 +- packages/api/src/mcp/MCPConnectionFactory.ts | 35 ++++++-- packages/api/src/mcp/oauth/tokens.ts | 6 +- 4 files changed, 127 insertions(+), 8 deletions(-) diff --git a/api/server/routes/__tests__/mcp.spec.js b/api/server/routes/__tests__/mcp.spec.js index 43e086f7b3..bb5c7c7027 100644 --- a/api/server/routes/__tests__/mcp.spec.js +++ b/api/server/routes/__tests__/mcp.spec.js @@ -498,6 +498,93 @@ describe('MCP Routes', () => { expect(response.headers.location).toBe('/oauth/error?error=callback_failed'); expect(mockMcpManager.getUserConnection).not.toHaveBeenCalled(); }); + + it('should re-fetch flow state after completeOAuthFlow to capture DCR updates', async () => { + const { mcpServersRegistry } = require('@librechat/api'); + const mockFlowManager = { + completeFlow: jest.fn().mockResolvedValue(), + deleteFlow: jest.fn().mockResolvedValue(true), + }; + const initialClientInfo = { + client_id: 'initial123', + client_secret: 'initial_secret', + }; + const updatedClientInfo = { + client_id: 'updated456', + client_secret: 'updated_secret', + }; + const initialFlowState = { + serverName: 'test-server', + userId: 'test-user-id', + metadata: { toolFlowId: 'tool-flow-123', serverUrl: 'http://example.com' }, + clientInfo: initialClientInfo, + codeVerifier: 'test-verifier', + }; + const updatedFlowState = { + serverName: 'test-server', + userId: 'test-user-id', + metadata: { toolFlowId: 'tool-flow-123', serverUrl: 'http://example.com' }, + clientInfo: updatedClientInfo, // DCR re-registration changed credentials + codeVerifier: 'test-verifier', + }; + const mockTokens = { + access_token: 'test-access-token', + refresh_token: 'test-refresh-token', + }; + + // First call returns initial state, second call returns updated state + MCPOAuthHandler.getFlowState + .mockResolvedValueOnce(initialFlowState) + .mockResolvedValueOnce(updatedFlowState); + MCPOAuthHandler.completeOAuthFlow.mockResolvedValue(mockTokens); + MCPTokenStorage.storeTokens.mockResolvedValue(); + mcpServersRegistry.getServerConfig.mockResolvedValue({}); + getLogStores.mockReturnValue({}); + require('~/config').getFlowStateManager.mockReturnValue(mockFlowManager); + + const mockUserConnection = { + fetchTools: jest.fn().mockResolvedValue([]), + }; + const mockMcpManager = { + getUserConnection: jest.fn().mockResolvedValue(mockUserConnection), + }; + require('~/config').getMCPManager.mockReturnValue(mockMcpManager); + require('~/config').getOAuthReconnectionManager = jest.fn().mockReturnValue({ + clearReconnection: jest.fn(), + }); + + const response = await request(app).get('/api/mcp/test-server/oauth/callback').query({ + code: 'test-auth-code', + state: 'test-flow-id', + }); + + expect(response.status).toBe(302); + expect(response.headers.location).toBe('/oauth/success?serverName=test-server'); + + // Verify MCPOAuthHandler.getFlowState was called TWICE (before and after completion) + expect(MCPOAuthHandler.getFlowState).toHaveBeenCalledTimes(2); + expect(MCPOAuthHandler.getFlowState).toHaveBeenNthCalledWith( + 1, + 'test-flow-id', + mockFlowManager, + ); + expect(MCPOAuthHandler.getFlowState).toHaveBeenNthCalledWith( + 2, + 'test-flow-id', + mockFlowManager, + ); + + // Verify storeTokens was called with UPDATED credentials, not initial ones + expect(MCPTokenStorage.storeTokens).toHaveBeenCalledWith( + expect.objectContaining({ + userId: 'test-user-id', + serverName: 'test-server', + tokens: mockTokens, + clientInfo: updatedClientInfo, // Should use updated, not initial + metadata: updatedFlowState.metadata, // Should use updated metadata + }), + ); + }); }); describe('GET /oauth/tokens/:flowId', () => { diff --git a/api/server/routes/mcp.js b/api/server/routes/mcp.js index 8d6d91e8d9..f232004158 100644 --- a/api/server/routes/mcp.js +++ b/api/server/routes/mcp.js @@ -139,6 +139,9 @@ router.get('/:serverName/oauth/callback', async (req, res) => { const tokens = await MCPOAuthHandler.completeOAuthFlow(flowId, code, flowManager, oauthHeaders); logger.info('[MCP OAuth] OAuth flow completed, tokens received in callback route'); + // Re-fetch flow state after completeOAuthFlow to capture any DCR updates + const updatedFlowState = await MCPOAuthHandler.getFlowState(flowId, flowManager); + /** Persist tokens immediately so reconnection uses fresh credentials */ if (flowState?.userId && tokens) { try { @@ -149,8 +152,8 @@ router.get('/:serverName/oauth/callback', async (req, res) => { createToken, updateToken, findToken, - clientInfo: flowState.clientInfo, - metadata: flowState.metadata, + clientInfo: updatedFlowState?.clientInfo || flowState.clientInfo, + metadata: updatedFlowState?.metadata || flowState.metadata, }); logger.debug('[MCP OAuth] Stored OAuth tokens prior to reconnection', { serverName, diff --git a/packages/api/src/mcp/MCPConnectionFactory.ts b/packages/api/src/mcp/MCPConnectionFactory.ts index 4425788cc9..e931a9eb5b 100644 --- a/packages/api/src/mcp/MCPConnectionFactory.ts +++ b/packages/api/src/mcp/MCPConnectionFactory.ts @@ -329,6 +329,7 @@ export class MCPConnectionFactory { /** Check if there's already an ongoing OAuth flow for this flowId */ const existingFlow = await this.flowManager.getFlowState(flowId, 'mcp_oauth'); + if (existingFlow && existingFlow.status === 'PENDING') { logger.debug( `${this.logPrefix} OAuth flow already exists for ${flowId}, waiting for completion`, @@ -342,13 +343,31 @@ export class MCPConnectionFactory { `${this.logPrefix} OAuth flow completed, tokens received for ${this.serverName}`, ); - /** Client information from the existing flow metadata */ + // Re-fetch flow state after completion to get updated credentials + const updatedFlowState = await MCPOAuthHandler.getFlowState( + flowId, + this.flowManager as FlowStateManager, + ); + + /** Client information from the updated flow metadata */ const existingMetadata = existingFlow.metadata as unknown as MCPOAuthFlowMetadata; - const clientInfo = existingMetadata?.clientInfo; + const clientInfo = updatedFlowState?.clientInfo || existingMetadata?.clientInfo; return { tokens, clientInfo }; } + // Clean up old completed flows: createFlow() may return cached results otherwise + if (existingFlow && existingFlow.status !== 'PENDING') { + try { + await this.flowManager.deleteFlow(flowId, 'mcp_oauth'); + logger.debug( + `${this.logPrefix} Cleared stale ${existingFlow.status} OAuth flow for ${flowId}`, + ); + } catch (error) { + logger.warn(`${this.logPrefix} Failed to clear stale OAuth flow`, error); + } + } + logger.debug(`${this.logPrefix} Initiating new OAuth flow for ${this.serverName}...`); const { authorizationUrl, @@ -383,9 +402,15 @@ export class MCPConnectionFactory { } logger.info(`${this.logPrefix} OAuth flow completed, tokens received for ${this.serverName}`); - /** Client information from the flow metadata */ - const clientInfo = flowMetadata?.clientInfo; - const metadata = flowMetadata?.metadata; + // Re-fetch flow state after completion to get updated credentials + const updatedFlowState = await MCPOAuthHandler.getFlowState( + newFlowId, + this.flowManager as FlowStateManager, + ); + + /** Client information from the updated flow state */ + const clientInfo = updatedFlowState?.clientInfo || flowMetadata?.clientInfo; + const metadata = updatedFlowState?.metadata || flowMetadata?.metadata; return { tokens, diff --git a/packages/api/src/mcp/oauth/tokens.ts b/packages/api/src/mcp/oauth/tokens.ts index a212a7f333..1f615ee48b 100644 --- a/packages/api/src/mcp/oauth/tokens.ts +++ b/packages/api/src/mcp/oauth/tokens.ts @@ -217,7 +217,11 @@ export class MCPTokenStorage { } } - logger.debug(`${logPrefix} Stored OAuth tokens`); + logger.debug(`${logPrefix} Stored OAuth tokens`, { + client_id: clientInfo?.client_id, + has_refresh_token: !!tokens.refresh_token, + expires_at: 'expires_at' in tokens ? tokens.expires_at : 'N/A', + }); } catch (error) { const logPrefix = this.getLogPrefix(userId, serverName); logger.error(`${logPrefix} Failed to store tokens`, error);