From fc3189b7182d642033f8999f521e43587fac3e17 Mon Sep 17 00:00:00 2001 From: Dustin Healy <54083382+dustinhealy@users.noreply.github.com> Date: Wed, 22 Apr 2026 14:05:51 -0700 Subject: [PATCH] =?UTF-8?q?=F0=9F=94=90=20fix:=20Restore=20Tenant=20Contex?= =?UTF-8?q?t=20in=20MCP=20OAuth=20Callback=20(#12782)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix: restore tenant context in MCP OAuth callback for multi-tenant deployments The MCP OAuth callback is a cross-origin redirect from the OAuth provider. SameSite=Strict cookies (including the JWT) are not sent, leaving the callback with no tenant context. With TENANT_ISOLATION_STRICT=true, all DB writes fail. Stores tenantId in flow metadata at OAuth initiation time (when the user is authenticated), then restores it via tenantStorage.run in the callback, wrapping the entire post-validation body. * test: address review findings for tenant context tests - Assert tenantId flows through to initFlow in MCPConnectionFactory test - Add beforeEach to tenant context tests to reset mocks independently --- api/server/routes/__tests__/mcp.spec.js | 85 ++++++++ api/server/routes/mcp.js | 182 ++++++++++-------- packages/api/src/mcp/MCPConnectionFactory.ts | 6 +- .../__tests__/MCPConnectionFactory.test.ts | 6 +- packages/api/src/mcp/oauth/types.ts | 2 + 5 files changed, 196 insertions(+), 85 deletions(-) diff --git a/api/server/routes/__tests__/mcp.spec.js b/api/server/routes/__tests__/mcp.spec.js index 9e3ed7a351..7fabcc910a 100644 --- a/api/server/routes/__tests__/mcp.spec.js +++ b/api/server/routes/__tests__/mcp.spec.js @@ -60,6 +60,9 @@ jest.mock('@librechat/api', () => { jest.mock('@librechat/data-schemas', () => ({ getTenantId: jest.fn(), + tenantStorage: { + run: jest.fn((store, fn) => fn()), + }, logger: { debug: jest.fn(), info: jest.fn(), @@ -1862,6 +1865,88 @@ describe('MCP Routes', () => { }); }); + describe('GET /:serverName/oauth/callback - Tenant Context', () => { + beforeEach(() => { + const { getTenantId, tenantStorage } = require('@librechat/data-schemas'); + const { MCPOAuthHandler, MCPTokenStorage } = require('@librechat/api'); + getTenantId.mockReset(); + tenantStorage.run.mockReset(); + tenantStorage.run.mockImplementation((store, fn) => fn()); + MCPOAuthHandler.resolveStateToFlowId.mockReset(); + MCPOAuthHandler.getFlowState.mockReset(); + MCPOAuthHandler.completeOAuthFlow.mockReset(); + MCPTokenStorage.storeTokens.mockReset(); + }); + + it('should wrap callback body in tenantStorage.run when flowState has tenantId and no current context', async () => { + const { getTenantId, tenantStorage } = require('@librechat/data-schemas'); + const { MCPOAuthHandler, MCPTokenStorage } = require('@librechat/api'); + const flowId = 'user123:test-server'; + const csrfToken = generateTestCsrfToken(flowId); + + getTenantId.mockReturnValue(undefined); + + MCPOAuthHandler.resolveStateToFlowId.mockResolvedValue(flowId); + MCPOAuthHandler.getFlowState.mockResolvedValue({ + serverName: 'test-server', + userId: 'user123', + tenantId: 'tenant-abc', + metadata: {}, + clientInfo: {}, + codeVerifier: 'test-verifier', + }); + MCPOAuthHandler.completeOAuthFlow.mockResolvedValue({ + access_token: 'token', + token_type: 'bearer', + }); + MCPTokenStorage.storeTokens.mockResolvedValue(); + + const response = await request(app) + .get(`/api/mcp/test-server/oauth/callback?code=test-code&state=${flowId}`) + .set('Cookie', [`oauth_csrf=${csrfToken}`]) + .expect(302); + + expect(tenantStorage.run).toHaveBeenCalledWith( + { tenantId: 'tenant-abc' }, + expect.any(Function), + ); + expect(MCPTokenStorage.storeTokens).toHaveBeenCalled(); + + const basePath = getBasePath(); + expect(response.headers.location).toContain(`${basePath}/oauth/success`); + }); + + it('should not call tenantStorage.run when flowState has no tenantId', async () => { + const { getTenantId, tenantStorage } = require('@librechat/data-schemas'); + const { MCPOAuthHandler, MCPTokenStorage } = require('@librechat/api'); + const flowId = 'user123:test-server'; + const csrfToken = generateTestCsrfToken(flowId); + + getTenantId.mockReturnValue(undefined); + + MCPOAuthHandler.resolveStateToFlowId.mockResolvedValue(flowId); + MCPOAuthHandler.getFlowState.mockResolvedValue({ + serverName: 'test-server', + userId: 'user123', + metadata: {}, + clientInfo: {}, + codeVerifier: 'test-verifier', + }); + MCPOAuthHandler.completeOAuthFlow.mockResolvedValue({ + access_token: 'token', + token_type: 'bearer', + }); + MCPTokenStorage.storeTokens.mockResolvedValue(); + + await request(app) + .get(`/api/mcp/test-server/oauth/callback?code=test-code&state=${flowId}`) + .set('Cookie', [`oauth_csrf=${csrfToken}`]) + .expect(302); + + expect(tenantStorage.run).not.toHaveBeenCalled(); + }); + }); + describe('GET /servers', () => { // mockRegistryInstance is defined at the top of the file diff --git a/api/server/routes/mcp.js b/api/server/routes/mcp.js index b747e6f5ed..e2c1df738c 100644 --- a/api/server/routes/mcp.js +++ b/api/server/routes/mcp.js @@ -1,5 +1,5 @@ const { Router } = require('express'); -const { logger, getTenantId } = require('@librechat/data-schemas'); +const { logger, getTenantId, tenantStorage } = require('@librechat/data-schemas'); const { CacheKeys, Constants, @@ -267,101 +267,121 @@ router.get('/:serverName/oauth/callback', async (req, res) => { { serverName, flowId }, ); } - const oauthHeaders = - flowState.oauthHeaders ?? (await getOAuthHeaders(serverName, flowState.userId)); - const tokens = await MCPOAuthHandler.completeOAuthFlow(flowId, code, flowManager, oauthHeaders); - logger.info('[MCP OAuth] OAuth flow completed, tokens received in callback route'); - - /** Persist tokens immediately so reconnection uses fresh credentials */ - if (flowState?.userId && tokens) { - try { - await MCPTokenStorage.storeTokens({ - userId: flowState.userId, - serverName, - tokens, - createToken: db.createToken, - updateToken: db.updateToken, - findToken: db.findToken, - clientInfo: flowState.clientInfo, - metadata: flowState.metadata, - }); - logger.debug('[MCP OAuth] Stored OAuth tokens prior to reconnection', { - serverName, - userId: flowState.userId, - }); - } catch (error) { - logger.error('[MCP OAuth] Failed to store OAuth tokens after callback', error); - throw error; + /** + * Restore tenant context for the callback body. The callback is a cross-origin + * redirect from the OAuth provider, so SameSite=Strict cookies (including the + * JWT) are not sent. The tenantId was stored in the flow metadata at initiation + * time when the user was authenticated. + */ + const runWithTenant = async (fn) => { + const flowTenantId = flowState.tenantId; + if (flowTenantId && !getTenantId()) { + return tenantStorage.run({ tenantId: flowTenantId }, fn); } + return fn(); + }; - /** - * Clear any cached `mcp_get_tokens` flow result so subsequent lookups - * re-fetch the freshly stored credentials instead of returning stale nulls. - */ - if (typeof flowManager?.deleteFlow === 'function') { + await runWithTenant(async () => { + const oauthHeaders = + flowState.oauthHeaders ?? (await getOAuthHeaders(serverName, flowState.userId)); + const tokens = await MCPOAuthHandler.completeOAuthFlow( + flowId, + code, + flowManager, + oauthHeaders, + ); + logger.info('[MCP OAuth] OAuth flow completed, tokens received in callback route'); + + /** Persist tokens immediately so reconnection uses fresh credentials */ + if (flowState?.userId && tokens) { try { - await flowManager.deleteFlow(flowId, 'mcp_get_tokens'); + await MCPTokenStorage.storeTokens({ + userId: flowState.userId, + serverName, + tokens, + createToken: db.createToken, + updateToken: db.updateToken, + findToken: db.findToken, + clientInfo: flowState.clientInfo, + metadata: flowState.metadata, + }); + logger.debug('[MCP OAuth] Stored OAuth tokens prior to reconnection', { + serverName, + userId: flowState.userId, + }); } catch (error) { - logger.warn('[MCP OAuth] Failed to clear cached token flow state', error); + logger.error('[MCP OAuth] Failed to store OAuth tokens after callback', error); + throw error; + } + + /** + * Clear any cached `mcp_get_tokens` flow result so subsequent lookups + * re-fetch the freshly stored credentials instead of returning stale nulls. + */ + if (typeof flowManager?.deleteFlow === 'function') { + try { + await flowManager.deleteFlow(flowId, 'mcp_get_tokens'); + } catch (error) { + logger.warn('[MCP OAuth] Failed to clear cached token flow state', error); + } } } - } - try { - const mcpManager = getMCPManager(flowState.userId); - logger.debug(`[MCP OAuth] Attempting to reconnect ${serverName} with new OAuth tokens`); + try { + const mcpManager = getMCPManager(flowState.userId); + logger.debug(`[MCP OAuth] Attempting to reconnect ${serverName} with new OAuth tokens`); - if (flowState.userId !== 'system') { - const user = { id: flowState.userId }; + if (flowState.userId !== 'system') { + const user = { id: flowState.userId }; - const userConnection = await mcpManager.getUserConnection({ - user, - serverName, - flowManager, - tokenMethods: { - findToken: db.findToken, - updateToken: db.updateToken, - createToken: db.createToken, - deleteTokens: db.deleteTokens, - }, - }); + const userConnection = await mcpManager.getUserConnection({ + user, + serverName, + flowManager, + tokenMethods: { + findToken: db.findToken, + updateToken: db.updateToken, + createToken: db.createToken, + deleteTokens: db.deleteTokens, + }, + }); - logger.info( - `[MCP OAuth] Successfully reconnected ${serverName} for user ${flowState.userId}`, - ); + logger.info( + `[MCP OAuth] Successfully reconnected ${serverName} for user ${flowState.userId}`, + ); - // clear any reconnection attempts - const oauthReconnectionManager = getOAuthReconnectionManager(); - oauthReconnectionManager.clearReconnection(flowState.userId, serverName); + const oauthReconnectionManager = getOAuthReconnectionManager(); + oauthReconnectionManager.clearReconnection(flowState.userId, serverName); - const tools = await userConnection.fetchTools(); - await updateMCPServerTools({ - userId: flowState.userId, - serverName, - tools, - }); - } else { - logger.debug(`[MCP OAuth] System-level OAuth completed for ${serverName}`); - } - } catch (error) { - logger.warn( - `[MCP OAuth] Failed to reconnect ${serverName} after OAuth, but tokens are saved:`, - error, - ); - } - - /** ID of the flow that the tool/connection is waiting for */ - const toolFlowId = flowState.metadata?.toolFlowId; - if (toolFlowId) { - logger.debug('[MCP OAuth] Completing tool flow', { toolFlowId }); - const completed = await flowManager.completeFlow(toolFlowId, 'mcp_oauth', tokens); - if (!completed) { + const tools = await userConnection.fetchTools(); + await updateMCPServerTools({ + userId: flowState.userId, + serverName, + tools, + }); + } else { + logger.debug(`[MCP OAuth] System-level OAuth completed for ${serverName}`); + } + } catch (error) { logger.warn( - '[MCP OAuth] Tool flow state not found during completion — waiter will time out', - { toolFlowId }, + `[MCP OAuth] Failed to reconnect ${serverName} after OAuth, but tokens are saved:`, + error, ); } - } + + /** ID of the flow that the tool/connection is waiting for */ + const toolFlowId = flowState.metadata?.toolFlowId; + if (toolFlowId) { + logger.debug('[MCP OAuth] Completing tool flow', { toolFlowId }); + const completed = await flowManager.completeFlow(toolFlowId, 'mcp_oauth', tokens); + if (!completed) { + logger.warn( + '[MCP OAuth] Tool flow state not found during completion — waiter will time out', + { toolFlowId }, + ); + } + } + }); /* end runWithTenant */ /** Redirect to success page with flowId and serverName */ const redirectUrl = `${basePath}/oauth/success?serverName=${encodeURIComponent(serverName)}`; diff --git a/packages/api/src/mcp/MCPConnectionFactory.ts b/packages/api/src/mcp/MCPConnectionFactory.ts index c308fee72e..03ae50385f 100644 --- a/packages/api/src/mcp/MCPConnectionFactory.ts +++ b/packages/api/src/mcp/MCPConnectionFactory.ts @@ -1,4 +1,4 @@ -import { logger } from '@librechat/data-schemas'; +import { logger, getTenantId } from '@librechat/data-schemas'; import type { OAuthClientInformation } from '@modelcontextprotocol/sdk/shared/auth.js'; import type { Tool } from '@modelcontextprotocol/sdk/types.js'; import type { TokenMethods } from '@librechat/data-schemas'; @@ -366,7 +366,7 @@ export class MCPConnectionFactory { } // Store flow state BEFORE redirecting so the callback can find it - const metadataWithUrl = { ...flowMetadata, authorizationUrl }; + const metadataWithUrl = { ...flowMetadata, authorizationUrl, tenantId: getTenantId() }; await this.flowManager!.initFlow(newFlowId, 'mcp_oauth', metadataWithUrl); await MCPOAuthHandler.storeStateMapping(flowMetadata.state, newFlowId, this.flowManager!); @@ -670,7 +670,7 @@ export class MCPConnectionFactory { reusedStoredClient = flowMetadata.reusedStoredClient === true; // Store flow state BEFORE redirecting so the callback can find it - const metadataWithUrl = { ...flowMetadata, authorizationUrl }; + const metadataWithUrl = { ...flowMetadata, authorizationUrl, tenantId: getTenantId() }; await this.flowManager.initFlow(newFlowId, 'mcp_oauth', metadataWithUrl); await MCPOAuthHandler.storeStateMapping(flowMetadata.state, newFlowId, this.flowManager); diff --git a/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts b/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts index d90ca5b345..33decd5c0c 100644 --- a/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts +++ b/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts @@ -18,6 +18,7 @@ jest.mock('@librechat/data-schemas', () => ({ error: jest.fn(), debug: jest.fn(), }, + getTenantId: jest.fn(), })); const mockLogger = logger as jest.Mocked; @@ -241,6 +242,9 @@ describe('MCPConnectionFactory', () => { }, }; + const { getTenantId } = require('@librechat/data-schemas'); + (getTenantId as jest.Mock).mockReturnValue('test-tenant'); + mockMCPOAuthHandler.initiateOAuthFlow.mockResolvedValue(mockFlowData); // createFlow runs as a background monitor — simulate it staying pending mockFlowManager.createFlow.mockReturnValue(new Promise(() => {})); @@ -278,7 +282,7 @@ describe('MCPConnectionFactory', () => { expect(mockFlowManager.initFlow).toHaveBeenCalledWith( 'flow123', 'mcp_oauth', - expect.objectContaining(mockFlowData.flowMetadata), + expect.objectContaining({ ...mockFlowData.flowMetadata, tenantId: 'test-tenant' }), ); const initCallOrder = mockFlowManager.initFlow.mock.invocationCallOrder[0]; const oauthStartCallOrder = (oauthOptions.oauthStart as jest.Mock).mock diff --git a/packages/api/src/mcp/oauth/types.ts b/packages/api/src/mcp/oauth/types.ts index ee8ce2d76d..20db2bc2a7 100644 --- a/packages/api/src/mcp/oauth/types.ts +++ b/packages/api/src/mcp/oauth/types.ts @@ -93,6 +93,8 @@ export interface MCPOAuthFlowMetadata extends FlowMetadata { oauthHeaders?: Record; /** True when the flow reused a stored client registration from a prior successful OAuth flow */ reusedStoredClient?: boolean; + /** Tenant context captured at flow initiation for callback replay (SameSite cookies unavailable on cross-origin redirects) */ + tenantId?: string; } export interface MCPOAuthTokens extends OAuthTokens {