🔧 fix: Update Token Calculations/Mapping, MCP env Initialization (#6406)

* fix: Enhance MCP initialization to process environment variables

* fix: only build tokenCountMap with messages that are being used in the payload

* fix: Adjust maxContextTokens calculation to account for maxOutputTokens

* refactor: Make processMCPEnv optional in MCPManager initialization

* chore: Bump version of librechat-data-provider to 0.7.73

api/app/clients/BaseClient.js

@@ -366,17 +366,14 @@ class BaseClient {
    *  context: TMessage[],
    *  remainingContextTokens: number,
    *  messagesToRefine: TMessage[],
-   *  summaryIndex: number,
+   * }>} An object with three properties: `context`, `remainingContextTokens`, and `messagesToRefine`.
-   * }>} An object with four properties: `context`, `summaryIndex`, `remainingContextTokens`, and `messagesToRefine`.
    *    `context` is an array of messages that fit within the token limit.
-   *    `summaryIndex` is the index of the first message in the `messagesToRefine` array.
    *    `remainingContextTokens` is the number of tokens remaining within the limit after adding the messages to the context.
    *    `messagesToRefine` is an array of messages that were not added to the context because they would have exceeded the token limit.
    */
   async getMessagesWithinTokenLimit({ messages: _messages, maxContextTokens, instructions }) {
     // Every reply is primed with <|start|>assistant<|message|>, so we
     // start with 3 tokens for the label after all messages have been counted.
-    let summaryIndex = -1;
     let currentTokenCount = 3;
     const instructionsTokenCount = instructions?.tokenCount ?? 0;
     let remainingContextTokens =
@@ -409,14 +406,12 @@ class BaseClient {
     }
 
     const prunedMemory = messages;
-    summaryIndex = prunedMemory.length - 1;
     remainingContextTokens -= currentTokenCount;
 
     return {
       context: context.reverse(),
       remainingContextTokens,
       messagesToRefine: prunedMemory,
-      summaryIndex,
     };
   }
 
@@ -459,7 +454,7 @@ class BaseClient {
 
     let orderedWithInstructions = this.addInstructions(orderedMessages, instructions);
 
-    let { context, remainingContextTokens, messagesToRefine, summaryIndex } =
+    let { context, remainingContextTokens, messagesToRefine } =
       await this.getMessagesWithinTokenLimit({
         messages: orderedWithInstructions,
         instructions,
@@ -529,7 +524,7 @@ class BaseClient {
     }
 
     // Make sure to only continue summarization logic if the summary message was generated
-    shouldSummarize = summaryMessage && shouldSummarize;
+    shouldSummarize = summaryMessage != null && shouldSummarize === true;
 
     logger.debug('[BaseClient] Context Count (2/2)', {
       remainingContextTokens,
@@ -539,17 +534,18 @@ class BaseClient {
     /** @type {Record<string, number> | undefined} */
     let tokenCountMap;
     if (buildTokenMap) {
-      tokenCountMap = orderedWithInstructions.reduce((map, message, index) => {
+      const currentPayload = shouldSummarize ? orderedWithInstructions : context;
+      tokenCountMap = currentPayload.reduce((map, message, index) => {
         const { messageId } = message;
         if (!messageId) {
           return map;
         }
 
-        if (shouldSummarize && index === summaryIndex && !usePrevSummary) {
+        if (shouldSummarize && index === messagesToRefine.length - 1 && !usePrevSummary) {
           map.summaryMessage = { ...summaryMessage, messageId, tokenCount: summaryTokenCount };
         }
 
-        map[messageId] = orderedWithInstructions[index].tokenCount;
+        map[messageId] = currentPayload[index].tokenCount;
         return map;
       }, {});
     }

api/app/clients/specs/BaseClient.test.js

@@ -164,7 +164,7 @@ describe('BaseClient', () => {
     const result = await TestClient.getMessagesWithinTokenLimit({ messages });
 
     expect(result.context).toEqual(expectedContext);
-    expect(result.summaryIndex).toEqual(expectedIndex);
+    expect(result.messagesToRefine.length - 1).toEqual(expectedIndex);
     expect(result.remainingContextTokens).toBe(expectedRemainingContextTokens);
     expect(result.messagesToRefine).toEqual(expectedMessagesToRefine);
   });
@@ -200,7 +200,7 @@ describe('BaseClient', () => {
     const result = await TestClient.getMessagesWithinTokenLimit({ messages });
 
     expect(result.context).toEqual(expectedContext);
-    expect(result.summaryIndex).toEqual(expectedIndex);
+    expect(result.messagesToRefine.length - 1).toEqual(expectedIndex);
     expect(result.remainingContextTokens).toBe(expectedRemainingContextTokens);
     expect(result.messagesToRefine).toEqual(expectedMessagesToRefine);
   });

api/server/services/AppService.js

@@ -2,6 +2,7 @@ const {
   FileSources,
   EModelEndpoint,
   loadOCRConfig,
+  processMCPEnv,
   getConfigDefaults,
 } = require('librechat-data-provider');
 const { checkVariables, checkHealth, checkConfig, checkAzureVariables } = require('./start/checks');
@@ -54,7 +55,7 @@ const AppService = async (app) => {
 
   if (config.mcpServers != null) {
     const mcpManager = await getMCPManager();
-    await mcpManager.initializeMCP(config.mcpServers);
+    await mcpManager.initializeMCP(config.mcpServers, processMCPEnv);
     await mcpManager.mapAvailableTools(availableTools);
   }
 

api/server/services/Endpoints/agents/initialize.js

@@ -178,6 +178,7 @@ const initializeAgentOptions = async ({
     agent.provider = options.provider;
   }
 
+  /** @type {import('@librechat/agents').ClientOptions} */
   agent.model_parameters = Object.assign(model_parameters, options.llmConfig);
   if (options.configOptions) {
     agent.model_parameters.configuration = options.configOptions;
@@ -196,6 +197,7 @@ const initializeAgentOptions = async ({
 
   const tokensModel =
     agent.provider === EModelEndpoint.azureOpenAI ? agent.model : agent.model_parameters.model;
+  const maxTokens = agent.model_parameters.maxOutputTokens ?? agent.model_parameters.maxTokens ?? 0;
 
   return {
     ...agent,
@@ -204,7 +206,7 @@ const initializeAgentOptions = async ({
     toolContextMap,
     maxContextTokens:
       agent.max_context_tokens ??
-      (getModelMaxTokens(tokensModel, providerEndpointMap[provider]) ?? 4000) * 0.9,
+      ((getModelMaxTokens(tokensModel, providerEndpointMap[provider]) ?? 4000) - maxTokens) * 0.9,
   };
 };
 

package-lock.json

@@ -40808,7 +40808,7 @@
     },
     "packages/data-provider": {
       "name": "librechat-data-provider",
-      "version": "0.7.72",
+      "version": "0.7.73",
       "license": "ISC",
       "dependencies": {
         "axios": "^1.8.2",

packages/data-provider/package.json

@@ -1,6 +1,6 @@
 {
   "name": "librechat-data-provider",
-  "version": "0.7.72",
+  "version": "0.7.73",
   "description": "data services for librechat apps",
   "main": "dist/index.js",
   "module": "dist/index.es.js",

packages/data-provider/src/mcp.ts

@@ -85,3 +85,26 @@ export const MCPOptionsSchema = z.union([
 ]);
 
 export const MCPServersSchema = z.record(z.string(), MCPOptionsSchema);
+
+export type MCPOptions = z.infer<typeof MCPOptionsSchema>;
+
+/**
+ * Recursively processes an object to replace environment variables in string values
+ * @param {MCPOptions} obj - The object to process
+ * @returns {MCPOptions} - The processed object with environment variables replaced
+ */
+export function processMCPEnv(obj: MCPOptions): MCPOptions {
+  if (obj === null || obj === undefined) {
+    return obj;
+  }
+
+  if ('env' in obj && obj.env) {
+    const processedEnv: Record<string, string> = {};
+    for (const [key, value] of Object.entries(obj.env)) {
+      processedEnv[key] = extractEnvVariable(value);
+    }
+    obj.env = processedEnv;
+  }
+
+  return obj;
+}

packages/mcp/src/manager.ts

@@ -1,5 +1,5 @@
 import { CallToolResultSchema } from '@modelcontextprotocol/sdk/types.js';
-import type { JsonSchemaType } from 'librechat-data-provider';
+import type { JsonSchemaType, MCPOptions } from 'librechat-data-provider';
 import type { Logger } from 'winston';
 import type * as t from './types/mcp';
 import { formatToolContent } from './parsers';
@@ -31,13 +31,17 @@ export class MCPManager {
     return MCPManager.instance;
   }
 
-  public async initializeMCP(mcpServers: t.MCPServers): Promise<void> {
+  public async initializeMCP(
+    mcpServers: t.MCPServers,
+    processMCPEnv?: (obj: MCPOptions) => MCPOptions,
+  ): Promise<void> {
     this.logger.info('[MCP] Initializing servers');
 
     const entries = Object.entries(mcpServers);
     const initializedServers = new Set();
     const connectionResults = await Promise.allSettled(
-      entries.map(async ([serverName, config], i) => {
+      entries.map(async ([serverName, _config], i) => {
+        const config = processMCPEnv ? processMCPEnv(_config) : _config;
         const connection = new MCPConnection(serverName, config, this.logger);
 
         connection.on('connectionChange', (state) => {