🏹 feat: Concurrent MCP Initialization Support (danny-avila#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

Author: dustinhealy

api/server/routes/mcp.js

@@ -4,6 +4,7 @@ const { MCPOAuthHandler } = require('@librechat/api');
 const { CacheKeys, Constants } = require('librechat-data-provider');
 const { findToken, updateToken, createToken, deleteTokens } = require('~/models');
 const { setCachedTools, getCachedTools, loadCustomConfig } = require('~/server/services/Config');
+const { getMCPSetupData, getServerConnectionStatus } = require('~/server/services/MCP');
 const { getUserPluginAuthValue } = require('~/server/services/PluginService');
 const { getMCPManager, getFlowStateManager } = require('~/config');
 const { requireJwtAuth } = require('~/server/middleware');
@@ -468,7 +469,7 @@ router.post('/:serverName/reinitialize', requireJwtAuth, async (req, res) => {
 
 /**
  * Get connection status for all MCP servers
- * This endpoint returns the actual connection status from MCPManager without disconnecting idle connections
+ * This endpoint returns all app level and user-scoped connection statuses from MCPManager without disconnecting idle connections
  */
 router.get('/connection/status', requireJwtAuth, async (req, res) => {
   try {
@@ -478,92 +479,83 @@ router.get('/connection/status', requireJwtAuth, async (req, res) => {
       return res.status(401).json({ error: 'User not authenticated' });
     }
 
-    const mcpManager = getMCPManager(user.id);
+    const { mcpConfig, appConnections, userConnections, oauthServers } = await getMCPSetupData(
+      user.id,
+    );
     const connectionStatus = {};
 
-    const printConfig = false;
-    const config = await loadCustomConfig(printConfig);
-    const mcpConfig = config?.mcpServers;
-
-    const appConnections = mcpManager.getAllConnections() || new Map();
-    const userConnections = mcpManager.getUserConnections(user.id) || new Map();
-    const oauthServers = mcpManager.getOAuthServers() || new Set();
-
-    if (!mcpConfig) {
-      return res.status(404).json({ error: 'MCP config not found' });
+    for (const [serverName] of Object.entries(mcpConfig)) {
+      connectionStatus[serverName] = await getServerConnectionStatus(
+        user.id,
+        serverName,
+        appConnections,
+        userConnections,
+        oauthServers,
+      );
     }
 
-    // Get flow manager to check for active/timed-out OAuth flows
-    const flowsCache = getLogStores(CacheKeys.FLOWS);
-    const flowManager = getFlowStateManager(flowsCache);
-
-    for (const [serverName] of Object.entries(mcpConfig)) {
-      const getConnectionState = (serverName) =>
-        appConnections.get(serverName)?.connectionState ??
-        userConnections.get(serverName)?.connectionState ??
-        'disconnected';
+    res.json({
+      success: true,
+      connectionStatus,
+    });
+  } catch (error) {
+    if (error.message === 'MCP config not found') {
+      return res.status(404).json({ error: error.message });
+    }
+    logger.error('[MCP Connection Status] Failed to get connection status', error);
+    res.status(500).json({ error: 'Failed to get connection status' });
+  }
+});
 
-      const baseConnectionState = getConnectionState(serverName);
+/**
+ * Get connection status for a single MCP server
+ * This endpoint returns the connection status for a specific server for a given user
+ */
+router.get('/connection/status/:serverName', requireJwtAuth, async (req, res) => {
+  try {
+    const user = req.user;
+    const { serverName } = req.params;
 
-      let hasActiveOAuthFlow = false;
-      let hasFailedOAuthFlow = false;
+    if (!user?.id) {
+      return res.status(401).json({ error: 'User not authenticated' });
+    }
 
-      if (baseConnectionState === 'disconnected' && oauthServers.has(serverName)) {
-        try {
-          // Check for user-specific OAuth flows
-          const flowId = MCPOAuthHandler.generateFlowId(user.id, serverName);
-          const flowState = await flowManager.getFlowState(flowId, 'mcp_oauth');
-          if (flowState) {
-            // Check if flow failed or timed out
-            const flowAge = Date.now() - flowState.createdAt;
-            const flowTTL = flowState.ttl || 180000; // Default 3 minutes
-
-            if (flowState.status === 'FAILED' || flowAge > flowTTL) {
-              hasFailedOAuthFlow = true;
-              logger.debug(`[MCP Connection Status] Found failed OAuth flow for ${serverName}`, {
-                flowId,
-                status: flowState.status,
-                flowAge,
-                flowTTL,
-                timedOut: flowAge > flowTTL,
-              });
-            } else if (flowState.status === 'PENDING') {
-              hasActiveOAuthFlow = true;
-              logger.debug(`[MCP Connection Status] Found active OAuth flow for ${serverName}`, {
-                flowId,
-                flowAge,
-                flowTTL,
-              });
-            }
-          }
-        } catch (error) {
-          logger.error(
-            `[MCP Connection Status] Error checking OAuth flows for ${serverName}:`,
-            error,
-          );
-        }
-      }
+    if (!serverName) {
+      return res.status(400).json({ error: 'Server name is required' });
+    }
 
-      // Determine the final connection state
-      let finalConnectionState = baseConnectionState;
-      if (hasFailedOAuthFlow) {
-        finalConnectionState = 'error'; // Report as error if OAuth failed
-      } else if (hasActiveOAuthFlow && baseConnectionState === 'disconnected') {
-        finalConnectionState = 'connecting'; // Still waiting for OAuth
-      }
+    const { mcpConfig, appConnections, userConnections, oauthServers } = await getMCPSetupData(
+      user.id,
+    );
 
-      connectionStatus[serverName] = {
-        requiresOAuth: oauthServers.has(serverName),
-        connectionState: finalConnectionState,
-      };
+    if (!mcpConfig[serverName]) {
+      return res
+        .status(404)
+        .json({ error: `MCP server '${serverName}' not found in configuration` });
     }
 
+    const serverStatus = await getServerConnectionStatus(
+      user.id,
+      serverName,
+      appConnections,
+      userConnections,
+      oauthServers,
+    );
+
     res.json({
       success: true,
-      connectionStatus,
+      serverName,
+      connectionStatus: serverStatus.connectionState,
+      requiresOAuth: serverStatus.requiresOAuth,
     });
   } catch (error) {
-    logger.error('[MCP Connection Status] Failed to get connection status', error);
+    if (error.message === 'MCP config not found') {
+      return res.status(404).json({ error: error.message });
+    }
+    logger.error(
+      `[MCP Per-Server Status] Failed to get connection status for ${req.params.serverName}`,
+      error,
+    );
     res.status(500).json({ error: 'Failed to get connection status' });
   }
 });

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,
 };