@librechat/agents 3.1.39 → 3.1.41

package/src/agents/__tests__/AgentContext.test.ts

@@ -7,10 +7,11 @@ describe('AgentContext', () => {
   type ContextOptions = {
     agentConfig?: Partial<t.AgentInputs>;
     tokenCounter?: t.TokenCounter;
+    indexTokenCountMap?: Record<string, number>;
   };
 
   const createBasicContext = (options: ContextOptions = {}): AgentContext => {
-    const { agentConfig = {}, tokenCounter } = options;
+    const { agentConfig = {}, tokenCounter, indexTokenCountMap } = options;
     return AgentContext.fromConfig(
       {
         agentId: 'test-agent',
@@ -18,7 +19,8 @@ describe('AgentContext', () => {
         instructions: 'Test instructions',
         ...agentConfig,
       },
-      tokenCounter
+      tokenCounter,
+      indexTokenCountMap
     );
   };
 
@@ -393,6 +395,23 @@ describe('AgentContext', () => {
       expect(ctx.currentUsage).toBeUndefined();
     });
 
+    it('rebuilds indexTokenCountMap from base map after reset', async () => {
+      const tokenCounter = jest.fn(() => 5);
+      const ctx = createBasicContext({
+        tokenCounter,
+        indexTokenCountMap: { '0': 10, '1': 20 },
+      });
+
+      await ctx.tokenCalculationPromise;
+      ctx.indexTokenCountMap = {};
+
+      ctx.reset();
+      await ctx.tokenCalculationPromise;
+
+      expect(ctx.indexTokenCountMap['1']).toBe(20);
+      expect(ctx.indexTokenCountMap['0'] ?? 0).toBeGreaterThanOrEqual(10);
+    });
+
     it('forces rebuild on next systemRunnable access', () => {
       const ctx = createBasicContext({ agentConfig: { instructions: 'Test' } });
 
@@ -780,8 +799,9 @@ describe('AgentContext', () => {
 
       // Verify state is cleared
       expect(ctx.discoveredToolNames.size).toBe(0);
-      expect(ctx.instructionTokens).toBe(0);
-      expect(ctx.indexTokenCountMap).toEqual({});
+      const resetTokens = ctx.instructionTokens;
+      expect(resetTokens).toBeGreaterThan(0);
+      expect(resetTokens).toBeLessThan(run1Tokens);
 
       // ========== RUN 2 ==========
       // Re-initialize (as fromConfig would do)
@@ -790,6 +810,7 @@ describe('AgentContext', () => {
       // System runnable should NOT include the previously discovered tool
       // (because discoveredToolNames was cleared)
       const run2Tokens = ctx.instructionTokens;
+      expect(run2Tokens).toBe(resetTokens);
 
       // Token count should be lower than run 1 (no discovered tool in system message)
       expect(run2Tokens).toBeLessThan(run1Tokens);

package/src/graphs/MultiAgentGraph.ts

@@ -15,6 +15,7 @@ import {
   getCurrentTaskInput,
   messagesStateReducer,
 } from '@langchain/langgraph';
+import type { LangGraphRunnableConfig } from '@langchain/langgraph';
 import type { BaseMessage, AIMessageChunk } from '@langchain/core/messages';
 import type { ToolRunnableConfig } from '@langchain/core/tools';
 import type * as t from '@/types';
@@ -745,7 +746,8 @@ export class MultiAgentGraph extends StandardGraph {
 
       /** Wrapper function that handles agentMessages channel, handoff reception, and conditional routing */
       const agentWrapper = async (
-        state: t.MultiAgentGraphState
+        state: t.MultiAgentGraphState,
+        config?: LangGraphRunnableConfig
       ): Promise<t.MultiAgentGraphState | Command> => {
         let result: t.MultiAgentGraphState;
 
@@ -817,7 +819,7 @@ export class MultiAgentGraph extends StandardGraph {
             ...state,
             messages: messagesForAgent,
           };
-          result = await agentSubgraph.invoke(transformedState);
+          result = await agentSubgraph.invoke(transformedState, config);
           result = {
             ...result,
             agentMessages: [],
@@ -863,14 +865,14 @@ export class MultiAgentGraph extends StandardGraph {
             ...state,
             messages: state.agentMessages,
           };
-          result = await agentSubgraph.invoke(transformedState);
+          result = await agentSubgraph.invoke(transformedState, config);
           result = {
             ...result,
             /** Clear agentMessages for next agent */
             agentMessages: [],
           };
         } else {
-          result = await agentSubgraph.invoke(state);
+          result = await agentSubgraph.invoke(state, config);
         }
 
         /** If agent has both handoff and direct edges, use Command for exclusive routing */

package/src/messages/format.ts

@@ -288,7 +288,12 @@ function formatAssistantMessage(
   let hasReasoning = false;
 
   if (Array.isArray(message.content)) {
-    for (const part of message.content) {
+    for (const part of message.content as Array<
+      MessageContentComplex | undefined | null
+    >) {
+      if (part == null) {
+        continue;
+      }
       if (part.type === ContentTypes.TEXT && part.tool_call_ids) {
         /*
         If there's pending content, it needs to be aggregated as a single string to prepare for tool calls.
@@ -865,7 +870,6 @@ export const formatAgentMessages = (
     indexMapping[i] = resultIndices;
   }
 
-  // Update the token count map if it was provided
   if (indexTokenCountMap) {
     for (
       let originalIndex = 0;
@@ -875,24 +879,77 @@ export const formatAgentMessages = (
       const resultIndices = indexMapping[originalIndex] || [];
       const tokenCount = indexTokenCountMap[originalIndex];
 
-      if (tokenCount !== undefined) {
-        if (resultIndices.length === 1) {
-          // Simple 1:1 mapping
-          updatedIndexTokenCountMap[resultIndices[0]] = tokenCount;
-        } else if (resultIndices.length > 1) {
-          // If one message was split into multiple, distribute the token count
-          // This is a simplification - in reality, you might want a more sophisticated distribution
-          const countPerMessage = Math.floor(tokenCount / resultIndices.length);
-          resultIndices.forEach((resultIndex, idx) => {
-            if (idx === resultIndices.length - 1) {
-              // Give any remainder to the last message
-              updatedIndexTokenCountMap[resultIndex] =
-                tokenCount - countPerMessage * (resultIndices.length - 1);
-            } else {
-              updatedIndexTokenCountMap[resultIndex] = countPerMessage;
+      if (tokenCount === undefined) {
+        continue;
+      }
+
+      const msgCount = resultIndices.length;
+      if (msgCount === 1) {
+        updatedIndexTokenCountMap[resultIndices[0]] = tokenCount;
+        continue;
+      }
+
+      if (msgCount < 2) {
+        continue;
+      }
+
+      let totalLength = 0;
+      const lastIdx = msgCount - 1;
+      const lengths = new Array<number>(msgCount);
+      for (let k = 0; k < msgCount; k++) {
+        const msg = messages[resultIndices[k]];
+        const { content } = msg;
+        let len = 0;
+        if (typeof content === 'string') {
+          len = content.length;
+        } else if (Array.isArray(content)) {
+          for (const part of content as Array<
+            Record<string, unknown> | string | undefined
+          >) {
+            if (typeof part === 'string') {
+              len += part.length;
+            } else if (part != null && typeof part === 'object') {
+              const val = part.text ?? part.content;
+              if (typeof val === 'string') {
+                len += val.length;
+              }
+            }
+          }
+        }
+        const toolCalls = (msg as AIMessage).tool_calls;
+        if (Array.isArray(toolCalls)) {
+          for (const tc of toolCalls as Array<Record<string, unknown>>) {
+            if (typeof tc.name === 'string') {
+              len += tc.name.length;
+            }
+            const { args } = tc;
+            if (typeof args === 'string') {
+              len += args.length;
+            } else if (args != null) {
+              len += JSON.stringify(args).length;
             }
-          });
+          }
+        }
+        lengths[k] = len;
+        totalLength += len;
+      }
+
+      if (totalLength === 0) {
+        const countPerMessage = Math.floor(tokenCount / msgCount);
+        for (let k = 0; k < lastIdx; k++) {
+          updatedIndexTokenCountMap[resultIndices[k]] = countPerMessage;
+        }
+        updatedIndexTokenCountMap[resultIndices[lastIdx]] =
+          tokenCount - countPerMessage * lastIdx;
+      } else {
+        let distributed = 0;
+        for (let k = 0; k < lastIdx; k++) {
+          const share = Math.floor((lengths[k] / totalLength) * tokenCount);
+          updatedIndexTokenCountMap[resultIndices[k]] = share;
+          distributed += share;
         }
+        updatedIndexTokenCountMap[resultIndices[lastIdx]] =
+          tokenCount - distributed;
       }
     }
   }