🚦 feat: Auto-reinitialize MCP Servers on Request (#9226)

api/server/controllers/agents/client.js

@@ -33,18 +33,13 @@ const {
   bedrockInputSchema,
   removeNullishValues,
 } = require('librechat-data-provider');
-const {
-  findPluginAuthsByKeys,
-  getFormattedMemories,
-  deleteMemory,
-  setMemory,
-} = require('~/models');
-const { getMCPAuthMap, checkCapability, hasCustomUserVars } = require('~/server/services/Config');
 const { addCacheControl, createContextHandlers } = require('~/app/clients/prompts');
 const { initializeAgent } = require('~/server/services/Endpoints/agents/agent');
 const { spendTokens, spendStructuredTokens } = require('~/models/spendTokens');
+const { getFormattedMemories, deleteMemory, setMemory } = require('~/models');
 const { encodeAndFormat } = require('~/server/services/Files/images/encode');
 const { getProviderConfig } = require('~/server/services/Endpoints');
+const { checkCapability } = require('~/server/services/Config');
 const BaseClient = require('~/app/clients/BaseClient');
 const { getRoleByName } = require('~/models/Role');
 const { loadAgent } = require('~/models/Agent');
@@ -615,6 +610,7 @@ class AgentClient extends BaseClient {
     await this.chatCompletion({
       payload,
       onProgress: opts.onProgress,
+      userMCPAuthMap: opts.userMCPAuthMap,
       abortController: opts.abortController,
     });
     return this.contentParts;
@@ -747,7 +743,13 @@ class AgentClient extends BaseClient {
     return currentMessageTokens > 0 ? currentMessageTokens : originalEstimate;
   }
 
-  async chatCompletion({ payload, abortController = null }) {
+  /**
+   * @param {object} params
+   * @param {string | ChatCompletionMessageParam[]} params.payload
+   * @param {Record<string, Record<string, string>>} [params.userMCPAuthMap]
+   * @param {AbortController} [params.abortController]
+   */
+  async chatCompletion({ payload, userMCPAuthMap, abortController = null }) {
     /** @type {Partial<GraphRunnableConfig>} */
     let config;
     /** @type {ReturnType<createRun>} */
@@ -903,21 +905,9 @@ class AgentClient extends BaseClient {
           run.Graph.contentData = contentData;
         }
 
-        try {
-          if (await hasCustomUserVars()) {
-            config.configurable.userMCPAuthMap = await getMCPAuthMap({
-              tools: agent.tools,
-              userId: this.options.req.user.id,
-              findPluginAuthsByKeys,
-            });
-          }
-        } catch (err) {
-          logger.error(
-            `[api/server/controllers/agents/client.js #chatCompletion] Error getting custom user vars for agent ${agent.id}`,
-            err,
-          );
+        if (userMCPAuthMap != null) {
+          config.configurable.userMCPAuthMap = userMCPAuthMap;
         }
-
         await run.processStream({ messages }, config, {
           keepContent: i !== 0,
           tokenCounter: createTokenCounter(this.getEncoding()),

api/server/controllers/agents/request.js

@@ -9,6 +9,24 @@ const {
 const { disposeClient, clientRegistry, requestDataMap } = require('~/server/cleanup');
 const { saveMessage } = require('~/models');
 
+function createCloseHandler(abortController) {
+  return function (manual) {
+    if (!manual) {
+      logger.debug('[AgentController] Request closed');
+    }
+    if (!abortController) {
+      return;
+    } else if (abortController.signal.aborted) {
+      return;
+    } else if (abortController.requestCompleted) {
+      return;
+    }
+
+    abortController.abort();
+    logger.debug('[AgentController] Request aborted on close');
+  };
+}
+
 const AgentController = async (req, res, next, initializeClient, addTitle) => {
   let {
     text,
@@ -31,7 +49,6 @@ const AgentController = async (req, res, next, initializeClient, addTitle) => {
   let userMessagePromise;
   let getAbortData;
   let client = null;
-  // Initialize as an array
   let cleanupHandlers = [];
 
   const newConvo = !conversationId;
@@ -62,9 +79,7 @@ const AgentController = async (req, res, next, initializeClient, addTitle) => {
   // Create a function to handle final cleanup
   const performCleanup = () => {
     logger.debug('[AgentController] Performing cleanup');
-    // Make sure cleanupHandlers is an array before iterating
     if (Array.isArray(cleanupHandlers)) {
-      // Execute all cleanup handlers
       for (const handler of cleanupHandlers) {
         try {
           if (typeof handler === 'function') {
@@ -105,8 +120,33 @@ const AgentController = async (req, res, next, initializeClient, addTitle) => {
   };
 
   try {
-    /** @type {{ client: TAgentClient }} */
-    const result = await initializeClient({ req, res, endpointOption });
+    let prelimAbortController = new AbortController();
+    const prelimCloseHandler = createCloseHandler(prelimAbortController);
+    res.on('close', prelimCloseHandler);
+    const removePrelimHandler = (manual) => {
+      try {
+        prelimCloseHandler(manual);
+        res.removeListener('close', prelimCloseHandler);
+      } catch (e) {
+        logger.error('[AgentController] Error removing close listener', e);
+      }
+    };
+    cleanupHandlers.push(removePrelimHandler);
+    /** @type {{ client: TAgentClient; userMCPAuthMap?: Record<string, Record<string, string>> }} */
+    const result = await initializeClient({
+      req,
+      res,
+      endpointOption,
+      signal: prelimAbortController.signal,
+    });
+    if (prelimAbortController.signal?.aborted) {
+      prelimAbortController = null;
+      throw new Error('Request was aborted before initialization could complete');
+    } else {
+      prelimAbortController = null;
+      removePrelimHandler(true);
+      cleanupHandlers.pop();
+    }
     client = result.client;
 
     // Register client with finalization registry if available
@@ -138,22 +178,7 @@ const AgentController = async (req, res, next, initializeClient, addTitle) => {
     };
 
     const { abortController, onStart } = createAbortController(req, res, getAbortData, getReqData);
-
-    // Simple handler to avoid capturing scope
-    const closeHandler = () => {
-      logger.debug('[AgentController] Request closed');
-      if (!abortController) {
-        return;
-      } else if (abortController.signal.aborted) {
-        return;
-      } else if (abortController.requestCompleted) {
-        return;
-      }
-
-      abortController.abort();
-      logger.debug('[AgentController] Request aborted on close');
-    };
-
+    const closeHandler = createCloseHandler(abortController);
     res.on('close', closeHandler);
     cleanupHandlers.push(() => {
       try {
@@ -175,6 +200,7 @@ const AgentController = async (req, res, next, initializeClient, addTitle) => {
       abortController,
       overrideParentMessageId,
       isEdited: !!editedContent,
+      userMCPAuthMap: result.userMCPAuthMap,
       responseMessageId: editedResponseMessageId,
       progressOptions: {
         res,