🧮 perf: Enable Agent Context Count Reuse (#15130)

* perf: enable agent context count reuse

* fix: declare token counter return type

* fix: initialize cached token counters

* style: sort token counter imports

* fix: Keep cached token counts exact

api/package.json

@@ -46,7 +46,7 @@
     "@azure/storage-blob": "^12.30.0",
     "@google/genai": "^2.8.0",
     "@keyv/redis": "5.1.6",
-    "@librechat/agents": "^3.6.12",
+    "@librechat/agents": "^3.6.13",
     "@librechat/api": "*",
     "@librechat/data-schemas": "*",
     "@microsoft/microsoft-graph-client": "^3.0.7",

api/server/controllers/agents/client.js

@@ -17,7 +17,7 @@ const {
   omitTitleOptions,
   getProviderConfig,
   memoryInstructions,
-  createTokenCounter,
+  createCachedTokenCounter,
   applyContextToAgent,
   isMemoryAgentEnabled,
   recordCollectedUsage,
@@ -3249,7 +3249,7 @@ class AgentClient extends BaseClient {
       };
 
       const toolSet = buildToolSet(this.options.agent);
-      const tokenCounter = createTokenCounter(this.getEncoding());
+      const tokenCounter = await createCachedTokenCounter(this.getEncoding());
 
       /** Pre-resolve invoked skill bodies + re-prime files before formatting messages */
       const skillPrimeResult = this.options.primeInvokedSkills
@@ -3889,7 +3889,7 @@ class AgentClient extends BaseClient {
         this.contentParts.push(...seedContent);
       }
 
-      const tokenCounter = createTokenCounter(this.getEncoding());
+      const tokenCounter = await createCachedTokenCounter(this.getEncoding());
       const agents = collectReachableAgents([
         this.options.agent,
         ...(this.agentConfigs?.size > 0 ? this.agentConfigs.values() : []),

api/server/controllers/agents/client.test.js

@@ -47,7 +47,7 @@ jest.mock('@librechat/api', () => ({
   createRun: (...args) => mockCreateRun(...args),
   countFormattedMessageTokens: jest.fn(() => 42),
   countTokens: jest.fn((text) => Math.ceil(String(text ?? '').length / 4)),
-  createTokenCounter: jest.fn(() => jest.fn(() => 0)),
+  createCachedTokenCounter: jest.fn(async () => jest.fn(() => 0)),
   createDetachedSubagentUsageRecorder: (...args) =>
     mockCreateDetachedSubagentUsageRecorder(...args),
   captureAgentCheckpointGeneration: (...args) => mockCaptureAgentCheckpointGeneration(...args),

package-lock.json

@@ -63,7 +63,7 @@
         "@azure/storage-blob": "^12.30.0",
         "@google/genai": "^2.8.0",
         "@keyv/redis": "5.1.6",
-        "@librechat/agents": "^3.6.12",
+        "@librechat/agents": "^3.6.13",
         "@librechat/api": "*",
         "@librechat/data-schemas": "*",
         "@microsoft/microsoft-graph-client": "^3.0.7",
@@ -10630,9 +10630,9 @@
       }
     },
     "node_modules/@librechat/agents": {
-      "version": "3.6.12",
-      "resolved": "https://registry.npmjs.org/@librechat/agents/-/agents-3.6.12.tgz",
-      "integrity": "sha512-S73bRfCJ4bklASObib7gpSn9Vm0IkJ9ZmVgtzcKzUYju8lkb18vU9FcK1ZFcyZl7/FZAQbHIHRYBDKHMqEn7xg==",
+      "version": "3.6.13",
+      "resolved": "https://registry.npmjs.org/@librechat/agents/-/agents-3.6.13.tgz",
+      "integrity": "sha512-nrcxEKwJDmIDsGf7txLqTgqqSKkJKl90xS5vqTtAoYSi4/EVPZ6ZFpJHLT3WMn2nWL5jjxdOmstet3JbSF1B+A==",
       "license": "MIT",
       "dependencies": {
         "@anthropic-ai/sdk": "^0.115.0",
@@ -42822,7 +42822,7 @@
         "@azure/storage-blob": "^12.30.0",
         "@google/genai": "^2.8.0",
         "@keyv/redis": "5.1.6",
-        "@librechat/agents": "^3.6.12",
+        "@librechat/agents": "^3.6.13",
         "@librechat/data-schemas": "*",
         "@modelcontextprotocol/sdk": "^1.30.0",
         "@opentelemetry/api": "^1.9.0",

packages/api/package.json

@@ -113,7 +113,7 @@
     "@azure/storage-blob": "^12.30.0",
     "@google/genai": "^2.8.0",
     "@keyv/redis": "5.1.6",
-    "@librechat/agents": "^3.6.12",
+    "@librechat/agents": "^3.6.13",
     "@librechat/data-schemas": "*",
     "@modelcontextprotocol/sdk": "^1.30.0",
     "@opentelemetry/api": "^1.9.0",

packages/api/src/agents/client.spec.ts

@@ -1,6 +1,10 @@
 import { ContentTypes } from 'librechat-data-provider';
+import { Tokenizer as AiTokenizer } from 'ai-tokenizer';
+import { Providers, StandardGraph } from '@librechat/agents';
+import { HumanMessage } from '@librechat/agents/langchain/messages';
 import type { TMessage } from 'librechat-data-provider';
 import {
+  createCachedTokenCounter,
   prependQuotes,
   prependFileContext,
   applyAttachmentOnlyText,
@@ -8,6 +12,39 @@ import {
 } from './client';
 import { ATTACHMENT_ONLY_TEXT } from '~/files/context';
 
+describe('createCachedTokenCounter', () => {
+  it('enables stable-message reuse in the agents runtime', async () => {
+    const getTokenCount = jest.spyOn(AiTokenizer.prototype, 'count');
+    try {
+      const tokenCounter = await createCachedTokenCounter('o200k_base');
+      const graph = new StandardGraph({
+        runId: 'token-cache-integration',
+        agents: [
+          {
+            agentId: 'primary',
+            provider: Providers.OPENAI,
+            instructions: 'Test instructions',
+          },
+        ],
+        tokenCounter,
+      });
+      const agentContext = graph.agentContexts.get('primary');
+      await agentContext?.tokenCalculationPromise;
+      getTokenCount.mockClear();
+      const message = new HumanMessage('Stable retained context');
+
+      agentContext?.contextPressureTokenCounts?.count(message);
+      const callsAfterFirstCount = getTokenCount.mock.calls.length;
+      agentContext?.contextPressureTokenCounts?.count(message);
+
+      expect(callsAfterFirstCount).toBeGreaterThan(0);
+      expect(getTokenCount).toHaveBeenCalledTimes(callsAfterFirstCount);
+    } finally {
+      getTokenCount.mockRestore();
+    }
+  });
+});
+
 describe('prependFileContext', () => {
   it('prepends file context to string content', () => {
     const message: FormattedMessageWithContent = { content: 'Answer this question.' };

packages/api/src/agents/client.ts

@@ -6,9 +6,10 @@ import {
   getTokenCountForMessage,
   estimateOpenAIImageTokens,
   estimateAnthropicImageTokens,
+  markTokenCounterCacheCompatible,
 } from '@librechat/agents';
+import type { MessageContentComplex, TokenCounter } from '@librechat/agents';
 import type { BaseMessage } from '@librechat/agents/langchain/messages';
-import type { MessageContentComplex } from '@librechat/agents';
 import type { Agent, TMessage } from 'librechat-data-provider';
 import type { ServerRequest } from '~/types';
 import { getSafeErrorMetadata, mergeQuotedText, formatQuotesAsMarkdown } from '~/utils';
@@ -378,9 +379,19 @@ export function countFormattedMessageTokens(
   return isClaude ? Math.ceil(numTokens * CLAUDE_TOKEN_CORRECTION) : numTokens;
 }
 
-export function createTokenCounter(encoding: Parameters<typeof Tokenizer.getTokenCount>[1]) {
+export function createTokenCounter(
+  encoding: Parameters<typeof Tokenizer.getTokenCount>[1],
+): TokenCounter {
+  return createMessageTokenCounter(encoding, (text: string) =>
+    Tokenizer.getTokenCount(text, encoding),
+  );
+}
+
+function createMessageTokenCounter(
+  encoding: Parameters<typeof Tokenizer.getTokenCount>[1],
+  countTokens: (text: string) => number,
+): TokenCounter {
   const isClaude = encoding === 'claude';
-  const countTokens = (text: string) => Tokenizer.getTokenCount(text, encoding);
   return function (message: BaseMessage): number {
     const count = getTokenCountForMessage(
       message,
@@ -391,6 +402,13 @@ export function createTokenCounter(encoding: Parameters<typeof Tokenizer.getToke
   };
 }
 
+export async function createCachedTokenCounter(
+  encoding: Parameters<typeof Tokenizer.getTokenCount>[1],
+): Promise<TokenCounter> {
+  const countTokens = await Tokenizer.createExactTokenCounter(encoding ?? 'o200k_base');
+  return markTokenCounterCacheCompatible(createMessageTokenCounter(encoding, countTokens));
+}
+
 export function logToolError(_graph: unknown, error: unknown, toolId: string): void {
   /**
    * A GraphInterrupt unwinding out of a tool body is the HITL pause working as

packages/api/src/utils/tokenizer.spec.ts

@@ -1,3 +1,4 @@
+import { Tokenizer as AiTokenizer } from 'ai-tokenizer';
 import Tokenizer from './tokenizer';
 
 jest.mock('@librechat/data-schemas', () => ({
@@ -53,4 +54,19 @@ describe('Tokenizer', () => {
       expect(count).toBeGreaterThan(0);
     });
   });
+
+  describe('createExactTokenCounter', () => {
+    it('throws instead of returning a fallback estimate after a tokenizer failure', async () => {
+      const counter = await Tokenizer.createExactTokenCounter('o200k_base');
+      const count = jest.spyOn(AiTokenizer.prototype, 'count').mockImplementationOnce(() => {
+        throw new Error('tokenizer failed');
+      });
+      try {
+        expect(() => counter('Stable cached context')).toThrow('tokenizer failed');
+      } finally {
+        count.mockRestore();
+        await Tokenizer.initEncoding('o200k_base');
+      }
+    });
+  });
 });

packages/api/src/utils/tokenizer.ts

@@ -27,6 +27,23 @@ class Tokenizer {
     return this.loadingPromises[encoding];
   }
 
+  /** Returns a counter that never substitutes an estimate for tokenizer errors. */
+  async createExactTokenCounter(encoding: EncodingName): Promise<(text: string) => number> {
+    await this.initEncoding(encoding);
+    const tokenizer = this.tokenizersCache[encoding];
+    if (!tokenizer) {
+      throw new Error(`Tokenizer encoding failed to initialize: ${encoding}`);
+    }
+    return (text: string): number => {
+      try {
+        return tokenizer.count(text);
+      } catch (error) {
+        this.handleCountError(encoding, tokenizer, error);
+        throw error;
+      }
+    };
+  }
+
   getTokenCount(text: string, encoding: EncodingName = 'o200k_base'): number {
     const tokenizer = this.tokenizersCache[encoding];
     if (!tokenizer) {
@@ -36,13 +53,22 @@ class Tokenizer {
     try {
       return tokenizer.count(text);
     } catch (error) {
-      logger.error('[Tokenizer] Error getting token count:', error);
-      delete this.tokenizersCache[encoding];
-      delete this.loadingPromises[encoding];
-      this.initEncoding(encoding);
+      this.handleCountError(encoding, tokenizer, error);
       return Math.ceil(text.length / 4);
     }
   }
+
+  private handleCountError(encoding: EncodingName, tokenizer: AiTokenizer, error: unknown): void {
+    logger.error('[Tokenizer] Error getting token count:', error);
+    if (this.tokenizersCache[encoding] !== tokenizer) {
+      return;
+    }
+    delete this.tokenizersCache[encoding];
+    delete this.loadingPromises[encoding];
+    void this.initEncoding(encoding).catch((initError: unknown) => {
+      logger.error(`[Tokenizer] Error reloading ${encoding} encoding:`, initError);
+    });
+  }
 }
 
 const TokenizerSingleton: Tokenizer = new Tokenizer();