diff --git a/api/server/routes/__tests__/mcp.spec.js b/api/server/routes/__tests__/mcp.spec.js index 6109346446..856b73f75c 100644 --- a/api/server/routes/__tests__/mcp.spec.js +++ b/api/server/routes/__tests__/mcp.spec.js @@ -830,6 +830,131 @@ describe('MCP Routes', () => { }); }); + it('should resolve and forward customUserVars so header templates are substituted on first post-callback connection', async () => { + const flowId = 'test-user-id:test-server'; + const mockFlowManager = { + getFlowState: jest.fn().mockResolvedValue({ + status: 'PENDING', + createdAt: Date.now(), + }), + completeFlow: jest.fn().mockResolvedValue(true), + deleteFlow: jest.fn().mockResolvedValue(true), + }; + const mockFlowState = { + serverName: 'test-server', + userId: 'test-user-id', + metadata: {}, + clientInfo: {}, + codeVerifier: 'test-verifier', + }; + const mergedServerConfig = { + type: 'streamable-http', + url: 'https://override.example.com/mcp', + source: 'config', + customUserVars: { + LITELLM_KEY: { title: 'LiteLLM Key' }, + }, + }; + const fetchedTools = [{ name: 'search', inputSchema: { type: 'object' } }]; + + getLogStores.mockReturnValue({}); + require('~/config').getFlowStateManager.mockReturnValue(mockFlowManager); + MCPOAuthHandler.getFlowState.mockResolvedValue(mockFlowState); + MCPOAuthHandler.completeOAuthFlow.mockResolvedValue({ + access_token: 'test-token', + }); + MCPTokenStorage.storeTokens.mockResolvedValue(); + mockRegistryInstance.getServerConfig.mockResolvedValue({}); + mockResolveAllMcpConfigs.mockResolvedValueOnce({ 'test-server': mergedServerConfig }); + require('@librechat/api').getUserMCPAuthMap.mockResolvedValueOnce({ + [`mcp_test-server`]: { LITELLM_KEY: 'sk-real-user-key' }, + }); + + const mockMcpManager = { + getUserConnection: jest.fn().mockResolvedValue({ + fetchTools: jest.fn().mockResolvedValue(fetchedTools), + }), + }; + require('~/config').getMCPManager.mockReturnValue(mockMcpManager); + require('~/config').getOAuthReconnectionManager.mockReturnValue({ + clearReconnection: jest.fn(), + }); + const { updateMCPServerTools } = require('~/server/services/Config/mcp'); + updateMCPServerTools.mockResolvedValue(); + + const response = await request(app) + .get('/api/mcp/test-server/oauth/callback') + .query({ code: 'test-code', state: flowId }); + + expect(response.status).toBe(302); + expect(require('@librechat/api').getUserMCPAuthMap).toHaveBeenCalledWith({ + userId: 'test-user-id', + servers: ['test-server'], + findPluginAuthsByKeys: require('~/models').findPluginAuthsByKeys, + }); + expect(mockMcpManager.getUserConnection).toHaveBeenCalledWith( + expect.objectContaining({ customUserVars: { LITELLM_KEY: 'sk-real-user-key' } }), + ); + }); + + it('should not call getUserMCPAuthMap when the server has no customUserVars', async () => { + const flowId = 'test-user-id:test-server'; + const mockFlowManager = { + getFlowState: jest.fn().mockResolvedValue({ + status: 'PENDING', + createdAt: Date.now(), + }), + completeFlow: jest.fn().mockResolvedValue(true), + deleteFlow: jest.fn().mockResolvedValue(true), + }; + const mockFlowState = { + serverName: 'test-server', + userId: 'test-user-id', + metadata: {}, + clientInfo: {}, + codeVerifier: 'test-verifier', + }; + const mergedServerConfig = { + type: 'streamable-http', + url: 'https://override.example.com/mcp', + source: 'config', + }; + const fetchedTools = [{ name: 'search', inputSchema: { type: 'object' } }]; + + getLogStores.mockReturnValue({}); + require('~/config').getFlowStateManager.mockReturnValue(mockFlowManager); + MCPOAuthHandler.getFlowState.mockResolvedValue(mockFlowState); + MCPOAuthHandler.completeOAuthFlow.mockResolvedValue({ + access_token: 'test-token', + }); + MCPTokenStorage.storeTokens.mockResolvedValue(); + mockRegistryInstance.getServerConfig.mockResolvedValue({}); + mockResolveAllMcpConfigs.mockResolvedValueOnce({ 'test-server': mergedServerConfig }); + require('@librechat/api').getUserMCPAuthMap.mockClear(); + + const mockMcpManager = { + getUserConnection: jest.fn().mockResolvedValue({ + fetchTools: jest.fn().mockResolvedValue(fetchedTools), + }), + }; + require('~/config').getMCPManager.mockReturnValue(mockMcpManager); + require('~/config').getOAuthReconnectionManager.mockReturnValue({ + clearReconnection: jest.fn(), + }); + const { updateMCPServerTools } = require('~/server/services/Config/mcp'); + updateMCPServerTools.mockResolvedValue(); + + const response = await request(app) + .get('/api/mcp/test-server/oauth/callback') + .query({ code: 'test-code', state: flowId }); + + expect(response.status).toBe(302); + expect(require('@librechat/api').getUserMCPAuthMap).not.toHaveBeenCalled(); + expect(mockMcpManager.getUserConnection).toHaveBeenCalledWith( + expect.objectContaining({ customUserVars: undefined }), + ); + }); + it('should reject when no PENDING flow exists and no cookies are present', async () => { const flowId = 'test-user-id:test-server'; const mockFlowManager = { diff --git a/api/server/routes/mcp.js b/api/server/routes/mcp.js index bcd3392c7d..83b5710756 100644 --- a/api/server/routes/mcp.js +++ b/api/server/routes/mcp.js @@ -3,9 +3,9 @@ const { logger, getTenantId, tenantStorage } = require('@librechat/data-schemas' const { CacheKeys, Constants, + Permissions, PermissionBits, PermissionTypes, - Permissions, } = require('librechat-data-provider'); const { getBasePath, @@ -14,7 +14,6 @@ const { MCPTokenStorage, setOAuthSession, PENDING_STALE_MS, - mcpConfig: mcpSettings, getUserMCPAuthMap, validateOAuthCsrf, OAUTH_CSRF_COOKIE, @@ -22,6 +21,8 @@ const { generateCheckAccess, validateOAuthSession, OAUTH_SESSION_COOKIE, + mcpConfig: mcpSettings, + getServerCustomUserVars, } = require('@librechat/api'); const { createMCPServerController, @@ -459,11 +460,38 @@ router.get('/:serverName/oauth/callback', async (req, res) => { ); } + /** + * Without this, getUserConnection resolves `headers`/`oauth_headers` + * customUserVars templates (e.g. `{{MY_VAR}}`) with no substitution + * data, so the literal placeholder is sent on this first post-callback + * connection attempt even though the user's value is already saved - + * surfaces upstream as a generic auth rejection from the MCP server. + * The other reconnect path (oauth/reinitialize route below) already + * resolves this the same way; this one was missing it. + */ + let userMCPAuthMap; + if (serverConfig?.customUserVars && typeof serverConfig.customUserVars === 'object') { + try { + userMCPAuthMap = await getUserMCPAuthMap({ + userId: flowState.userId, + servers: [serverName], + findPluginAuthsByKeys: db.findPluginAuthsByKeys, + }); + } catch (error) { + logger.warn( + `[MCP OAuth] Could not resolve customUserVars for ${serverName} before reconnecting:`, + error, + ); + } + } + const customUserVars = getServerCustomUserVars(userMCPAuthMap, serverName); + const userConnection = await mcpManager.getUserConnection({ user, serverName, flowManager, serverConfig, + customUserVars, tokenMethods: { findToken: db.findToken, updateToken: db.updateToken, diff --git a/api/server/services/MCP.js b/api/server/services/MCP.js index 18355d99a9..9fea5617d8 100644 --- a/api/server/services/MCP.js +++ b/api/server/services/MCP.js @@ -21,6 +21,7 @@ const { buildMCPAuthRunStepEndDeltaEvent, isUserSourced, checkAccessWithRequestCache, + getServerCustomUserVars, requiresEphemeralUserConnection, containsGraphTokenPlaceholder, } = require('@librechat/api'); @@ -215,10 +216,6 @@ async function resolveAllMcpConfigs(userId, user) { return await registry.getAllServerConfigs(userId, configServers); } -function getServerCustomUserVars(userMCPAuthMap, serverName) { - return userMCPAuthMap?.[`${Constants.mcp_prefix}${serverName}`]; -} - /** * Best-effort early gate; the authoritative check is * `assertResolvedRuntimeConfigAllowed` in `@librechat/api`, whose resolution diff --git a/packages/api/src/mcp/__tests__/auth.test.ts b/packages/api/src/mcp/__tests__/auth.test.ts index 2a06d5dd6a..c5b59a5a41 100644 --- a/packages/api/src/mcp/__tests__/auth.test.ts +++ b/packages/api/src/mcp/__tests__/auth.test.ts @@ -1,7 +1,7 @@ import type { PluginAuthMethods } from '@librechat/data-schemas'; import type { GenericTool } from '@librechat/agents'; +import { getUserMCPAuthMap, getServerCustomUserVars } from '../auth'; import { getPluginAuthMap } from '~/agents/auth'; -import { getUserMCPAuthMap } from '../auth'; jest.mock('~/agents/auth', () => ({ getPluginAuthMap: jest.fn(), @@ -301,3 +301,20 @@ describe('getUserMCPAuthMap', () => { }); }); }); + +describe('getServerCustomUserVars', () => { + it('reads a server entry by its mcp_-prefixed key', () => { + const authMap = { + 'mcp_my-server': { API_KEY: 'sk-123' }, + 'mcp_other-server': { TOKEN: 'abc' }, + }; + expect(getServerCustomUserVars(authMap, 'my-server')).toEqual({ API_KEY: 'sk-123' }); + }); + + it('returns undefined for a missing server or map', () => { + expect( + getServerCustomUserVars({ 'mcp_my-server': { API_KEY: 'sk-123' } }, 'absent'), + ).toBeUndefined(); + expect(getServerCustomUserVars(undefined, 'my-server')).toBeUndefined(); + }); +}); diff --git a/packages/api/src/mcp/auth.ts b/packages/api/src/mcp/auth.ts index 912e0fad41..6c7b3b72a7 100644 --- a/packages/api/src/mcp/auth.ts +++ b/packages/api/src/mcp/auth.ts @@ -5,6 +5,14 @@ import type { GenericTool } from '@librechat/agents'; import { getPluginAuthMap } from '~/agents/auth'; import { splitMCPToolKey } from './utils'; +/** Reads one server's customUserVars from a `getUserMCPAuthMap` result */ +export function getServerCustomUserVars( + userMCPAuthMap: Record> | undefined, + serverName: string, +): Record | undefined { + return userMCPAuthMap?.[`${Constants.mcp_prefix}${serverName}`]; +} + export async function getUserMCPAuthMap({ userId, tools,