🏹 feat: Concurrent MCP Initialization Support (#8677)

* ✨ feat: Enhance MCP Connection Status Management - Introduced new functions to retrieve and manage connection status for multiple MCP servers, including OAuth flow checks and server-specific status retrieval. - Refactored the MCP connection status endpoints to support both all servers and individual server queries. - Replaced the old server initialization hook with a new `useMCPServerManager` hook for improved state management and handling of multiple OAuth flows. - Updated the MCPPanel component to utilize the new context provider for better state handling and UI updates. - Fixed a number of UI bugs when initializing servers * 🗣️ i18n: Remove unused strings from translation.json * refactor: move helper functions out of the route module into mcp service file * ci: add tests for newly added functions in mcp service file * fix: memoize setMCPValues to avoid render loop
2026-01-08 19:48:51 +01:00 · 2025-07-28 09:25:34 -07:00 · 2025-07-28 09:25:34 -07:00 · 0ef3fefaec
commit 0ef3fefaec
parent 37aba18a96
15 changed files with 1092 additions and 542 deletions
--- a/api/server/services/MCP.js
+++ b/api/server/services/MCP.js
@ -12,7 +12,7 @@ const {
 } = require('@librechat/api');
 const { findToken, createToken, updateToken } = require('~/models');
 const { getMCPManager, getFlowStateManager } = require('~/config');
-const { getCachedTools } = require('./Config');
+const { getCachedTools, loadCustomConfig } = require('./Config');
 const { getLogStores } = require('~/cache');

 /**
@ -239,6 +239,123 @@ async function createMCPTool({ req, res, toolKey, provider: _provider }) {
  return toolInstance;
 }

+/**
+ * Get MCP setup data including config, connections, and OAuth servers
+ * @param {string} userId - The user ID
+ * @returns {Object} Object containing mcpConfig, appConnections, userConnections, and oauthServers
+ */
+async function getMCPSetupData(userId) {
+  const printConfig = false;
+  const config = await loadCustomConfig(printConfig);
+  const mcpConfig = config?.mcpServers;
+
+  if (!mcpConfig) {
+    throw new Error('MCP config not found');
+  }
+
+  const mcpManager = getMCPManager(userId);
+  const appConnections = mcpManager.getAllConnections() || new Map();
+  const userConnections = mcpManager.getUserConnections(userId) || new Map();
+  const oauthServers = mcpManager.getOAuthServers() || new Set();
+
+  return {
+    mcpConfig,
+    appConnections,
+    userConnections,
+    oauthServers,
+  };
+}
+
+/**
+ * Check OAuth flow status for a user and server
+ * @param {string} userId - The user ID
+ * @param {string} serverName - The server name
+ * @returns {Object} Object containing hasActiveFlow and hasFailedFlow flags
+ */
+async function checkOAuthFlowStatus(userId, serverName) {
+  const flowsCache = getLogStores(CacheKeys.FLOWS);
+  const flowManager = getFlowStateManager(flowsCache);
+  const flowId = MCPOAuthHandler.generateFlowId(userId, serverName);
+
+  try {
+    const flowState = await flowManager.getFlowState(flowId, 'mcp_oauth');
+    if (!flowState) {
+      return { hasActiveFlow: false, hasFailedFlow: false };
+    }
+
+    const flowAge = Date.now() - flowState.createdAt;
+    const flowTTL = flowState.ttl || 180000; // Default 3 minutes
+
+    if (flowState.status === 'FAILED' || flowAge > flowTTL) {
+      logger.debug(`[MCP Connection Status] Found failed OAuth flow for ${serverName}`, {
+        flowId,
+        status: flowState.status,
+        flowAge,
+        flowTTL,
+        timedOut: flowAge > flowTTL,
+      });
+      return { hasActiveFlow: false, hasFailedFlow: true };
+    }
+
+    if (flowState.status === 'PENDING') {
+      logger.debug(`[MCP Connection Status] Found active OAuth flow for ${serverName}`, {
+        flowId,
+        flowAge,
+        flowTTL,
+      });
+      return { hasActiveFlow: true, hasFailedFlow: false };
+    }
+
+    return { hasActiveFlow: false, hasFailedFlow: false };
+  } catch (error) {
+    logger.error(`[MCP Connection Status] Error checking OAuth flows for ${serverName}:`, error);
+    return { hasActiveFlow: false, hasFailedFlow: false };
+  }
+}
+
+/**
+ * Get connection status for a specific MCP server
+ * @param {string} userId - The user ID
+ * @param {string} serverName - The server name
+ * @param {Map} appConnections - App-level connections
+ * @param {Map} userConnections - User-level connections
+ * @param {Set} oauthServers - Set of OAuth servers
+ * @returns {Object} Object containing requiresOAuth and connectionState
+ */
+async function getServerConnectionStatus(
+  userId,
+  serverName,
+  appConnections,
+  userConnections,
+  oauthServers,
+) {
+  const getConnectionState = () =>
+    appConnections.get(serverName)?.connectionState ??
+    userConnections.get(serverName)?.connectionState ??
+    'disconnected';
+
+  const baseConnectionState = getConnectionState();
+  let finalConnectionState = baseConnectionState;
+
+  if (baseConnectionState === 'disconnected' && oauthServers.has(serverName)) {
+    const { hasActiveFlow, hasFailedFlow } = await checkOAuthFlowStatus(userId, serverName);
+
+    if (hasFailedFlow) {
+      finalConnectionState = 'error';
+    } else if (hasActiveFlow) {
+      finalConnectionState = 'connecting';
+    }
+  }
+
+  return {
+    requiresOAuth: oauthServers.has(serverName),
+    connectionState: finalConnectionState,
+  };
+}
+
 module.exports = {
  createMCPTool,
+  getMCPSetupData,
+  checkOAuthFlowStatus,
+  getServerConnectionStatus,
 };
--- a/api/server/services/MCP.spec.js
+++ b/api/server/services/MCP.spec.js
@ -0,0 +1,510 @@
+const { logger } = require('@librechat/data-schemas');
+const { MCPOAuthHandler } = require('@librechat/api');
+const { CacheKeys } = require('librechat-data-provider');
+const { getMCPSetupData, checkOAuthFlowStatus, getServerConnectionStatus } = require('./MCP');
+
+// Mock all dependencies
+jest.mock('@librechat/data-schemas', () => ({
+  logger: {
+    debug: jest.fn(),
+    error: jest.fn(),
+  },
+}));
+
+jest.mock('@librechat/api', () => ({
+  MCPOAuthHandler: {
+    generateFlowId: jest.fn(),
+  },
+}));
+
+jest.mock('librechat-data-provider', () => ({
+  CacheKeys: {
+    FLOWS: 'flows',
+  },
+}));
+
+jest.mock('./Config', () => ({
+  loadCustomConfig: jest.fn(),
+}));
+
+jest.mock('~/config', () => ({
+  getMCPManager: jest.fn(),
+  getFlowStateManager: jest.fn(),
+}));
+
+jest.mock('~/cache', () => ({
+  getLogStores: jest.fn(),
+}));
+
+jest.mock('~/models', () => ({
+  findToken: jest.fn(),
+  createToken: jest.fn(),
+  updateToken: jest.fn(),
+}));
+
+describe('tests for the new helper functions used by the MCP connection status endpoints', () => {
+  let mockLoadCustomConfig;
+  let mockGetMCPManager;
+  let mockGetFlowStateManager;
+  let mockGetLogStores;
+
+  beforeEach(() => {
+    jest.clearAllMocks();
+
+    mockLoadCustomConfig = require('./Config').loadCustomConfig;
+    mockGetMCPManager = require('~/config').getMCPManager;
+    mockGetFlowStateManager = require('~/config').getFlowStateManager;
+    mockGetLogStores = require('~/cache').getLogStores;
+  });
+
+  describe('getMCPSetupData', () => {
+    const mockUserId = 'user-123';
+    const mockConfig = {
+      mcpServers: {
+        server1: { type: 'stdio' },
+        server2: { type: 'http' },
+      },
+    };
+
+    beforeEach(() => {
+      mockGetMCPManager.mockReturnValue({
+        getAllConnections: jest.fn(() => new Map()),
+        getUserConnections: jest.fn(() => new Map()),
+        getOAuthServers: jest.fn(() => new Set()),
+      });
+    });
+
+    it('should successfully return MCP setup data', async () => {
+      mockLoadCustomConfig.mockResolvedValue(mockConfig);
+
+      const mockAppConnections = new Map([['server1', { status: 'connected' }]]);
+      const mockUserConnections = new Map([['server2', { status: 'disconnected' }]]);
+      const mockOAuthServers = new Set(['server2']);
+
+      const mockMCPManager = {
+        getAllConnections: jest.fn(() => mockAppConnections),
+        getUserConnections: jest.fn(() => mockUserConnections),
+        getOAuthServers: jest.fn(() => mockOAuthServers),
+      };
+      mockGetMCPManager.mockReturnValue(mockMCPManager);
+
+      const result = await getMCPSetupData(mockUserId);
+
+      expect(mockLoadCustomConfig).toHaveBeenCalledWith(false);
+      expect(mockGetMCPManager).toHaveBeenCalledWith(mockUserId);
+      expect(mockMCPManager.getAllConnections).toHaveBeenCalled();
+      expect(mockMCPManager.getUserConnections).toHaveBeenCalledWith(mockUserId);
+      expect(mockMCPManager.getOAuthServers).toHaveBeenCalled();
+
+      expect(result).toEqual({
+        mcpConfig: mockConfig.mcpServers,
+        appConnections: mockAppConnections,
+        userConnections: mockUserConnections,
+        oauthServers: mockOAuthServers,
+      });
+    });
+
+    it('should throw error when MCP config not found', async () => {
+      mockLoadCustomConfig.mockResolvedValue({});
+      await expect(getMCPSetupData(mockUserId)).rejects.toThrow('MCP config not found');
+    });
+
+    it('should handle null values from MCP manager gracefully', async () => {
+      mockLoadCustomConfig.mockResolvedValue(mockConfig);
+
+      const mockMCPManager = {
+        getAllConnections: jest.fn(() => null),
+        getUserConnections: jest.fn(() => null),
+        getOAuthServers: jest.fn(() => null),
+      };
+      mockGetMCPManager.mockReturnValue(mockMCPManager);
+
+      const result = await getMCPSetupData(mockUserId);
+
+      expect(result).toEqual({
+        mcpConfig: mockConfig.mcpServers,
+        appConnections: new Map(),
+        userConnections: new Map(),
+        oauthServers: new Set(),
+      });
+    });
+  });
+
+  describe('checkOAuthFlowStatus', () => {
+    const mockUserId = 'user-123';
+    const mockServerName = 'test-server';
+    const mockFlowId = 'flow-123';
+
+    beforeEach(() => {
+      const mockFlowsCache = {};
+      const mockFlowManager = {
+        getFlowState: jest.fn(),
+      };
+
+      mockGetLogStores.mockReturnValue(mockFlowsCache);
+      mockGetFlowStateManager.mockReturnValue(mockFlowManager);
+      MCPOAuthHandler.generateFlowId.mockReturnValue(mockFlowId);
+    });
+
+    it('should return false flags when no flow state exists', async () => {
+      const mockFlowManager = { getFlowState: jest.fn(() => null) };
+      mockGetFlowStateManager.mockReturnValue(mockFlowManager);
+
+      const result = await checkOAuthFlowStatus(mockUserId, mockServerName);
+
+      expect(mockGetLogStores).toHaveBeenCalledWith(CacheKeys.FLOWS);
+      expect(MCPOAuthHandler.generateFlowId).toHaveBeenCalledWith(mockUserId, mockServerName);
+      expect(mockFlowManager.getFlowState).toHaveBeenCalledWith(mockFlowId, 'mcp_oauth');
+      expect(result).toEqual({ hasActiveFlow: false, hasFailedFlow: false });
+    });
+
+    it('should detect failed flow when status is FAILED', async () => {
+      const mockFlowState = {
+        status: 'FAILED',
+        createdAt: Date.now() - 60000, // 1 minute ago
+        ttl: 180000,
+      };
+      const mockFlowManager = { getFlowState: jest.fn(() => mockFlowState) };
+      mockGetFlowStateManager.mockReturnValue(mockFlowManager);
+
+      const result = await checkOAuthFlowStatus(mockUserId, mockServerName);
+
+      expect(result).toEqual({ hasActiveFlow: false, hasFailedFlow: true });
+      expect(logger.debug).toHaveBeenCalledWith(
+        expect.stringContaining('Found failed OAuth flow'),
+        expect.objectContaining({
+          flowId: mockFlowId,
+          status: 'FAILED',
+        }),
+      );
+    });
+
+    it('should detect failed flow when flow has timed out', async () => {
+      const mockFlowState = {
+        status: 'PENDING',
+        createdAt: Date.now() - 200000, // 200 seconds ago (> 180s TTL)
+        ttl: 180000,
+      };
+      const mockFlowManager = { getFlowState: jest.fn(() => mockFlowState) };
+      mockGetFlowStateManager.mockReturnValue(mockFlowManager);
+
+      const result = await checkOAuthFlowStatus(mockUserId, mockServerName);
+
+      expect(result).toEqual({ hasActiveFlow: false, hasFailedFlow: true });
+      expect(logger.debug).toHaveBeenCalledWith(
+        expect.stringContaining('Found failed OAuth flow'),
+        expect.objectContaining({
+          timedOut: true,
+        }),
+      );
+    });
+
+    it('should detect failed flow when TTL not specified and flow exceeds default TTL', async () => {
+      const mockFlowState = {
+        status: 'PENDING',
+        createdAt: Date.now() - 200000, // 200 seconds ago (> 180s default TTL)
+        // ttl not specified, should use 180000 default
+      };
+      const mockFlowManager = { getFlowState: jest.fn(() => mockFlowState) };
+      mockGetFlowStateManager.mockReturnValue(mockFlowManager);
+
+      const result = await checkOAuthFlowStatus(mockUserId, mockServerName);
+
+      expect(result).toEqual({ hasActiveFlow: false, hasFailedFlow: true });
+    });
+
+    it('should detect active flow when status is PENDING and within TTL', async () => {
+      const mockFlowState = {
+        status: 'PENDING',
+        createdAt: Date.now() - 60000, // 1 minute ago (< 180s TTL)
+        ttl: 180000,
+      };
+      const mockFlowManager = { getFlowState: jest.fn(() => mockFlowState) };
+      mockGetFlowStateManager.mockReturnValue(mockFlowManager);
+
+      const result = await checkOAuthFlowStatus(mockUserId, mockServerName);
+
+      expect(result).toEqual({ hasActiveFlow: true, hasFailedFlow: false });
+      expect(logger.debug).toHaveBeenCalledWith(
+        expect.stringContaining('Found active OAuth flow'),
+        expect.objectContaining({
+          flowId: mockFlowId,
+        }),
+      );
+    });
+
+    it('should return false flags for other statuses', async () => {
+      const mockFlowState = {
+        status: 'COMPLETED',
+        createdAt: Date.now() - 60000,
+        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 = {
+        getFlowState: jest.fn(() => {
+          throw mockError;
+        }),
+      };
+      mockGetFlowStateManager.mockReturnValue(mockFlowManager);
+
+      const result = await checkOAuthFlowStatus(mockUserId, mockServerName);
+
+      expect(result).toEqual({ hasActiveFlow: false, hasFailedFlow: false });
+      expect(logger.error).toHaveBeenCalledWith(
+        expect.stringContaining('Error checking OAuth flows'),
+        mockError,
+      );
+    });
+  });
+
+  describe('getServerConnectionStatus', () => {
+    const mockUserId = 'user-123';
+    const mockServerName = 'test-server';
+
+    it('should return app connection state when available', async () => {
+      const appConnections = new Map([[mockServerName, { connectionState: 'connected' }]]);
+      const userConnections = new Map();
+      const oauthServers = new Set();
+
+      const result = await getServerConnectionStatus(
+        mockUserId,
+        mockServerName,
+        appConnections,
+        userConnections,
+        oauthServers,
+      );
+
+      expect(result).toEqual({
+        requiresOAuth: false,
+        connectionState: 'connected',
+      });
+    });
+
+    it('should fallback to user connection state when app connection not available', async () => {
+      const appConnections = new Map();
+      const userConnections = new Map([[mockServerName, { connectionState: 'connecting' }]]);
+      const oauthServers = new Set();
+
+      const result = await getServerConnectionStatus(
+        mockUserId,
+        mockServerName,
+        appConnections,
+        userConnections,
+        oauthServers,
+      );
+
+      expect(result).toEqual({
+        requiresOAuth: false,
+        connectionState: 'connecting',
+      });
+    });
+
+    it('should default to disconnected when no connections exist', async () => {
+      const appConnections = new Map();
+      const userConnections = new Map();
+      const oauthServers = new Set();
+
+      const result = await getServerConnectionStatus(
+        mockUserId,
+        mockServerName,
+        appConnections,
+        userConnections,
+        oauthServers,
+      );
+
+      expect(result).toEqual({
+        requiresOAuth: false,
+        connectionState: 'disconnected',
+      });
+    });
+
+    it('should prioritize app connection over user connection', async () => {
+      const appConnections = new Map([[mockServerName, { connectionState: 'connected' }]]);
+      const userConnections = new Map([[mockServerName, { connectionState: 'disconnected' }]]);
+      const oauthServers = new Set();
+
+      const result = await getServerConnectionStatus(
+        mockUserId,
+        mockServerName,
+        appConnections,
+        userConnections,
+        oauthServers,
+      );
+
+      expect(result).toEqual({
+        requiresOAuth: false,
+        connectionState: 'connected',
+      });
+    });
+
+    it('should indicate OAuth requirement when server is in OAuth servers set', async () => {
+      const appConnections = new Map();
+      const userConnections = new Map();
+      const oauthServers = new Set([mockServerName]);
+
+      const result = await getServerConnectionStatus(
+        mockUserId,
+        mockServerName,
+        appConnections,
+        userConnections,
+        oauthServers,
+      );
+
+      expect(result.requiresOAuth).toBe(true);
+    });
+
+    it('should handle OAuth flow status when disconnected and requires OAuth with failed flow', async () => {
+      const appConnections = new Map();
+      const userConnections = new Map();
+      const oauthServers = new Set([mockServerName]);
+
+      // Mock flow state to return failed flow
+      const mockFlowManager = {
+        getFlowState: jest.fn(() => ({
+          status: 'FAILED',
+          createdAt: Date.now() - 60000,
+          ttl: 180000,
+        })),
+      };
+      mockGetFlowStateManager.mockReturnValue(mockFlowManager);
+      mockGetLogStores.mockReturnValue({});
+      MCPOAuthHandler.generateFlowId.mockReturnValue('test-flow-id');
+
+      const result = await getServerConnectionStatus(
+        mockUserId,
+        mockServerName,
+        appConnections,
+        userConnections,
+        oauthServers,
+      );
+
+      expect(result).toEqual({
+        requiresOAuth: true,
+        connectionState: 'error',
+      });
+    });
+
+    it('should handle OAuth flow status when disconnected and requires OAuth with active flow', async () => {
+      const appConnections = new Map();
+      const userConnections = new Map();
+      const oauthServers = new Set([mockServerName]);
+
+      // Mock flow state to return active flow
+      const mockFlowManager = {
+        getFlowState: jest.fn(() => ({
+          status: 'PENDING',
+          createdAt: Date.now() - 60000, // 1 minute ago
+          ttl: 180000, // 3 minutes TTL
+        })),
+      };
+      mockGetFlowStateManager.mockReturnValue(mockFlowManager);
+      mockGetLogStores.mockReturnValue({});
+      MCPOAuthHandler.generateFlowId.mockReturnValue('test-flow-id');
+
+      const result = await getServerConnectionStatus(
+        mockUserId,
+        mockServerName,
+        appConnections,
+        userConnections,
+        oauthServers,
+      );
+
+      expect(result).toEqual({
+        requiresOAuth: true,
+        connectionState: 'connecting',
+      });
+    });
+
+    it('should handle OAuth flow status when disconnected and requires OAuth with no flow', async () => {
+      const appConnections = new Map();
+      const userConnections = new Map();
+      const oauthServers = new Set([mockServerName]);
+
+      // Mock flow state to return no flow
+      const mockFlowManager = {
+        getFlowState: jest.fn(() => null),
+      };
+      mockGetFlowStateManager.mockReturnValue(mockFlowManager);
+      mockGetLogStores.mockReturnValue({});
+      MCPOAuthHandler.generateFlowId.mockReturnValue('test-flow-id');
+
+      const result = await getServerConnectionStatus(
+        mockUserId,
+        mockServerName,
+        appConnections,
+        userConnections,
+        oauthServers,
+      );
+
+      expect(result).toEqual({
+        requiresOAuth: true,
+        connectionState: 'disconnected',
+      });
+    });
+
+    it('should not check OAuth flow status when server is connected', async () => {
+      const mockFlowManager = {
+        getFlowState: jest.fn(),
+      };
+      mockGetFlowStateManager.mockReturnValue(mockFlowManager);
+      mockGetLogStores.mockReturnValue({});
+
+      const appConnections = new Map([[mockServerName, { connectionState: 'connected' }]]);
+      const userConnections = new Map();
+      const oauthServers = new Set([mockServerName]);
+
+      const result = await getServerConnectionStatus(
+        mockUserId,
+        mockServerName,
+        appConnections,
+        userConnections,
+        oauthServers,
+      );
+
+      expect(result).toEqual({
+        requiresOAuth: true,
+        connectionState: 'connected',
+      });
+
+      // Should not call flow manager since server is connected
+      expect(mockFlowManager.getFlowState).not.toHaveBeenCalled();
+    });
+
+    it('should not check OAuth flow status when server does not require OAuth', async () => {
+      const mockFlowManager = {
+        getFlowState: jest.fn(),
+      };
+      mockGetFlowStateManager.mockReturnValue(mockFlowManager);
+      mockGetLogStores.mockReturnValue({});
+
+      const appConnections = new Map();
+      const userConnections = new Map();
+      const oauthServers = new Set(); // Server not in OAuth servers
+
+      const result = await getServerConnectionStatus(
+        mockUserId,
+        mockServerName,
+        appConnections,
+        userConnections,
+        oauthServers,
+      );
+
+      expect(result).toEqual({
+        requiresOAuth: false,
+        connectionState: 'disconnected',
+      });
+
+      // Should not call flow manager since server doesn't require OAuth
+      expect(mockFlowManager.getFlowState).not.toHaveBeenCalled();
+    });
+  });
+});