💫 feat: MCP OAuth Auto-Reconnect (#9646)

* add oauth reconnect tracker

* add connection tracker to mcp manager

* reconnect oauth mcp servers function

* call reconnection in auth controller

* make sure to check connection in panel

* wait for isConnected

* add const for poll interval

* add logging to tryReconnect

* check expiration

* check mcp manager is not null

* check mcp manager is not null

* add test for reconnecting mcp server

* unify logic inside OAuthReconnectionManager

* test reconnection manager, adjust

* chore: reorder import statements in index.js

* chore: imports

* chore: imports

* chore: imports

* chore: imports

* chore: imports

* chore: imports and use types explicitly

---------

Co-authored-by: Danny Avila <danny@librechat.ai>

api/server/controllers/AuthController.js

@@ -11,8 +11,9 @@ const {
   registerUser,
 } = require('~/server/services/AuthService');
 const { findUser, getUserById, deleteAllUserSessions, findSession } = require('~/models');
-const { getOpenIdConfig } = require('~/strategies');
 const { getGraphApiToken } = require('~/server/services/GraphTokenService');
+const { getOAuthReconnectionManager } = require('~/config');
+const { getOpenIdConfig } = require('~/strategies');
 
 const registrationController = async (req, res) => {
   try {
@@ -107,6 +108,14 @@ const refreshController = async (req, res) => {
 
     if (session && session.expiration > new Date()) {
       const token = await setAuthTokens(userId, res, session);
+
+      // trigger OAuth MCP server reconnection asynchronously (best effort)
+      void getOAuthReconnectionManager()
+        .reconnectServers(userId)
+        .catch((err) => {
+          logger.error('Error reconnecting OAuth MCP servers:', err);
+        });
+
       res.status(200).send({ token, user });
     } else if (req?.query?.retry) {
       // Retrying from a refresh token request that failed (401)

api/server/index.js

@@ -12,6 +12,7 @@ const { logger } = require('@librechat/data-schemas');
 const mongoSanitize = require('express-mongo-sanitize');
 const { isEnabled, ErrorController } = require('@librechat/api');
 const { connectDb, indexSync } = require('~/db');
+const initializeOAuthReconnectManager = require('./services/initializeOAuthReconnectManager');
 const createValidateImageRequest = require('./middleware/validateImageRequest');
 const { jwtLogin, ldapLogin, passportLogin } = require('~/strategies');
 const { updateInterfacePermissions } = require('~/models/interface');
@@ -154,7 +155,7 @@ const startServer = async () => {
     res.send(updatedIndexHtml);
   });
 
-  app.listen(port, host, () => {
+  app.listen(port, host, async () => {
     if (host === '0.0.0.0') {
       logger.info(
         `Server listening on all interfaces at port ${port}. Use http://localhost:${port} to access it`,
@@ -163,7 +164,9 @@ const startServer = async () => {
       logger.info(`Server listening at http://${host == '0.0.0.0' ? 'localhost' : host}:${port}`);
     }
 
-    initializeMCPs().then(() => checkMigrations());
+    await initializeMCPs();
+    await initializeOAuthReconnectManager();
+    await checkMigrations();
   });
 };
 

api/server/routes/mcp.js

@@ -1,12 +1,12 @@
 const { Router } = require('express');
 const { logger } = require('@librechat/data-schemas');
+const { CacheKeys, Constants } = require('librechat-data-provider');
 const { MCPOAuthHandler, getUserMCPAuthMap } = require('@librechat/api');
+const { getMCPManager, getFlowStateManager, getOAuthReconnectionManager } = require('~/config');
 const { getMCPSetupData, getServerConnectionStatus } = require('~/server/services/MCP');
 const { findToken, updateToken, createToken, deleteTokens } = require('~/models');
 const { updateMCPUserTools } = require('~/server/services/Config/mcpToolsCache');
 const { getUserPluginAuthValue } = require('~/server/services/PluginService');
-const { CacheKeys, Constants } = require('librechat-data-provider');
-const { getMCPManager, getFlowStateManager } = require('~/config');
 const { reinitMCPServer } = require('~/server/services/Tools/mcp');
 const { requireJwtAuth } = require('~/server/middleware');
 const { findPluginAuthsByKeys } = require('~/models');
@@ -144,6 +144,10 @@ router.get('/:serverName/oauth/callback', async (req, res) => {
           `[MCP OAuth] Successfully reconnected ${serverName} for user ${flowState.userId}`,
         );
 
+        // clear any reconnection attempts
+        const oauthReconnectionManager = getOAuthReconnectionManager();
+        oauthReconnectionManager.clearReconnection(flowState.userId, serverName);
+
         const tools = await userConnection.fetchTools();
         await updateMCPUserTools({
           userId: flowState.userId,

api/server/services/MCP.js

@@ -20,8 +20,8 @@ const {
   ContentTypes,
   isAssistantsEndpoint,
 } = require('librechat-data-provider');
+const { getMCPManager, getFlowStateManager, getOAuthReconnectionManager } = require('~/config');
 const { findToken, createToken, updateToken } = require('~/models');
-const { getMCPManager, getFlowStateManager } = require('~/config');
 const { getCachedTools, getAppConfig } = require('./Config');
 const { reinitMCPServer } = require('./Tools/mcp');
 const { getLogStores } = require('~/cache');
@@ -538,13 +538,20 @@ async function getServerConnectionStatus(
   const baseConnectionState = getConnectionState();
   let finalConnectionState = baseConnectionState;
 
+  // connection state overrides specific to OAuth servers
   if (baseConnectionState === 'disconnected' && oauthServers.has(serverName)) {
-    const { hasActiveFlow, hasFailedFlow } = await checkOAuthFlowStatus(userId, serverName);
-
-    if (hasFailedFlow) {
-      finalConnectionState = 'error';
-    } else if (hasActiveFlow) {
+    // check if server is actively being reconnected
+    const oauthReconnectionManager = getOAuthReconnectionManager();
+    if (oauthReconnectionManager.isReconnecting(userId, serverName)) {
       finalConnectionState = 'connecting';
+    } else {
+      const { hasActiveFlow, hasFailedFlow } = await checkOAuthFlowStatus(userId, serverName);
+
+      if (hasFailedFlow) {
+        finalConnectionState = 'error';
+      } else if (hasActiveFlow) {
+        finalConnectionState = 'connecting';
+      }
     }
   }
 

api/server/services/MCP.spec.js

@@ -31,6 +31,7 @@ jest.mock('./Config', () => ({
 jest.mock('~/config', () => ({
   getMCPManager: jest.fn(),
   getFlowStateManager: jest.fn(),
+  getOAuthReconnectionManager: jest.fn(),
 }));
 
 jest.mock('~/cache', () => ({
@@ -48,6 +49,7 @@ describe('tests for the new helper functions used by the MCP connection status e
   let mockGetMCPManager;
   let mockGetFlowStateManager;
   let mockGetLogStores;
+  let mockGetOAuthReconnectionManager;
 
   beforeEach(() => {
     jest.clearAllMocks();
@@ -56,6 +58,7 @@ describe('tests for the new helper functions used by the MCP connection status e
     mockGetMCPManager = require('~/config').getMCPManager;
     mockGetFlowStateManager = require('~/config').getFlowStateManager;
     mockGetLogStores = require('~/cache').getLogStores;
+    mockGetOAuthReconnectionManager = require('~/config').getOAuthReconnectionManager;
   });
 
   describe('getMCPSetupData', () => {
@@ -354,6 +357,12 @@ describe('tests for the new helper functions used by the MCP connection status e
       const userConnections = new Map();
       const oauthServers = new Set([mockServerName]);
 
+      // Mock OAuthReconnectionManager
+      const mockOAuthReconnectionManager = {
+        isReconnecting: jest.fn(() => false),
+      };
+      mockGetOAuthReconnectionManager.mockReturnValue(mockOAuthReconnectionManager);
+
       const result = await getServerConnectionStatus(
         mockUserId,
         mockServerName,
@@ -370,6 +379,12 @@ describe('tests for the new helper functions used by the MCP connection status e
       const userConnections = new Map();
       const oauthServers = new Set([mockServerName]);
 
+      // Mock OAuthReconnectionManager
+      const mockOAuthReconnectionManager = {
+        isReconnecting: jest.fn(() => false),
+      };
+      mockGetOAuthReconnectionManager.mockReturnValue(mockOAuthReconnectionManager);
+
       // Mock flow state to return failed flow
       const mockFlowManager = {
         getFlowState: jest.fn(() => ({
@@ -401,6 +416,12 @@ describe('tests for the new helper functions used by the MCP connection status e
       const userConnections = new Map();
       const oauthServers = new Set([mockServerName]);
 
+      // Mock OAuthReconnectionManager
+      const mockOAuthReconnectionManager = {
+        isReconnecting: jest.fn(() => false),
+      };
+      mockGetOAuthReconnectionManager.mockReturnValue(mockOAuthReconnectionManager);
+
       // Mock flow state to return active flow
       const mockFlowManager = {
         getFlowState: jest.fn(() => ({
@@ -432,6 +453,12 @@ describe('tests for the new helper functions used by the MCP connection status e
       const userConnections = new Map();
       const oauthServers = new Set([mockServerName]);
 
+      // Mock OAuthReconnectionManager
+      const mockOAuthReconnectionManager = {
+        isReconnecting: jest.fn(() => false),
+      };
+      mockGetOAuthReconnectionManager.mockReturnValue(mockOAuthReconnectionManager);
+
       // Mock flow state to return no flow
       const mockFlowManager = {
         getFlowState: jest.fn(() => null),
@@ -454,6 +481,35 @@ describe('tests for the new helper functions used by the MCP connection status e
       });
     });
 
+    it('should return connecting state when OAuth server is reconnecting', async () => {
+      const appConnections = new Map();
+      const userConnections = new Map();
+      const oauthServers = new Set([mockServerName]);
+
+      // Mock OAuthReconnectionManager to return true for isReconnecting
+      const mockOAuthReconnectionManager = {
+        isReconnecting: jest.fn(() => true),
+      };
+      mockGetOAuthReconnectionManager.mockReturnValue(mockOAuthReconnectionManager);
+
+      const result = await getServerConnectionStatus(
+        mockUserId,
+        mockServerName,
+        appConnections,
+        userConnections,
+        oauthServers,
+      );
+
+      expect(result).toEqual({
+        requiresOAuth: true,
+        connectionState: 'connecting',
+      });
+      expect(mockOAuthReconnectionManager.isReconnecting).toHaveBeenCalledWith(
+        mockUserId,
+        mockServerName,
+      );
+    });
+
     it('should not check OAuth flow status when server is connected', async () => {
       const mockFlowManager = {
         getFlowState: jest.fn(),

api/server/services/initializeOAuthReconnectManager.js

@@ -0,0 +1,26 @@
+const { logger } = require('@librechat/data-schemas');
+const { CacheKeys } = require('librechat-data-provider');
+const { createOAuthReconnectionManager, getFlowStateManager } = require('~/config');
+const { findToken, updateToken, createToken, deleteTokens } = require('~/models');
+const { getLogStores } = require('~/cache');
+
+/**
+ * Initialize OAuth reconnect manager
+ */
+async function initializeOAuthReconnectManager() {
+  try {
+    const flowManager = getFlowStateManager(getLogStores(CacheKeys.FLOWS));
+    const tokenMethods = {
+      findToken,
+      updateToken,
+      createToken,
+      deleteTokens,
+    };
+    await createOAuthReconnectionManager(flowManager, tokenMethods);
+    logger.info(`OAuth reconnect manager initialized successfully.`);
+  } catch (error) {
+    logger.error('Failed to initialize OAuth reconnect manager:', error);
+  }
+}
+
+module.exports = initializeOAuthReconnectManager;