@librechat/agents 2.4.322 → 3.0.0-rc1

package/src/specs/openai.simple.test.ts

@@ -4,12 +4,19 @@
 import { config } from 'dotenv';
 config();
 import { Calculator } from '@langchain/community/tools/calculator';
-import { HumanMessage, BaseMessage, UsageMetadata } from '@langchain/core/messages';
-import type { StandardGraph } from '@/graphs';
+import {
+  HumanMessage,
+  BaseMessage,
+  UsageMetadata,
+} from '@langchain/core/messages';
 import type * as t from '@/types';
-import { ToolEndHandler, ModelEndHandler, createMetadataAggregator } from '@/events';
+import {
+  ToolEndHandler,
+  ModelEndHandler,
+  createMetadataAggregator,
+} from '@/events';
+import { ContentTypes, GraphEvents, Providers, TitleMethod } from '@/common';
 import { ChatModelStreamHandler, createContentAggregator } from '@/stream';
-import { ContentTypes, GraphEvents, Providers } from '@/common';
 import { capitalizeFirstLetter } from './spec.utils';
 import { getLLMConfig } from '@/utils/llmConfig';
 import { getArgs } from '@/scripts/args';
@@ -36,7 +43,8 @@ describe(`${capitalizeFirstLetter(provider)} Streaming Tests`, () => {
   beforeEach(async () => {
     conversationHistory = [];
     collectedUsage = [];
-    const { contentParts: cp, aggregateContent: ac } = createContentAggregator();
+    const { contentParts: cp, aggregateContent: ac } =
+      createContentAggregator();
     contentParts = cp as t.MessageContentComplex[];
     aggregateContent = ac;
   });
@@ -49,36 +57,62 @@ describe(`${capitalizeFirstLetter(provider)} Streaming Tests`, () => {
     onRunStepSpy.mockReset();
   });
 
-  const setupCustomHandlers = (): Record<string | GraphEvents, t.EventHandler> => ({
+  const setupCustomHandlers = (): Record<
+    string | GraphEvents,
+    t.EventHandler
+  > => ({
     [GraphEvents.TOOL_END]: new ToolEndHandler(),
     [GraphEvents.CHAT_MODEL_END]: new ModelEndHandler(collectedUsage),
     [GraphEvents.CHAT_MODEL_STREAM]: new ChatModelStreamHandler(),
     [GraphEvents.ON_RUN_STEP_COMPLETED]: {
-      handle: (event: GraphEvents.ON_RUN_STEP_COMPLETED, data: t.StreamEventData): void => {
-        aggregateContent({ event, data: data as unknown as { result: t.ToolEndEvent; } });
-      }
+      handle: (
+        event: GraphEvents.ON_RUN_STEP_COMPLETED,
+        data: t.StreamEventData
+      ): void => {
+        aggregateContent({
+          event,
+          data: data as unknown as { result: t.ToolEndEvent },
+        });
+      },
     },
     [GraphEvents.ON_RUN_STEP]: {
-      handle: (event: GraphEvents.ON_RUN_STEP, data: t.StreamEventData, metadata, graph): void => {
+      handle: (
+        event: GraphEvents.ON_RUN_STEP,
+        data: t.StreamEventData,
+        metadata,
+        graph
+      ): void => {
         onRunStepSpy(event, data, metadata, graph);
         aggregateContent({ event, data: data as t.RunStep });
-      }
+      },
     },
     [GraphEvents.ON_RUN_STEP_DELTA]: {
-      handle: (event: GraphEvents.ON_RUN_STEP_DELTA, data: t.StreamEventData): void => {
+      handle: (
+        event: GraphEvents.ON_RUN_STEP_DELTA,
+        data: t.StreamEventData
+      ): void => {
         aggregateContent({ event, data: data as t.RunStepDeltaEvent });
-      }
+      },
     },
     [GraphEvents.ON_MESSAGE_DELTA]: {
-      handle: (event: GraphEvents.ON_MESSAGE_DELTA, data: t.StreamEventData, metadata, graph): void => {
+      handle: (
+        event: GraphEvents.ON_MESSAGE_DELTA,
+        data: t.StreamEventData,
+        metadata,
+        graph
+      ): void => {
         onMessageDeltaSpy(event, data, metadata, graph);
         aggregateContent({ event, data: data as t.MessageDeltaEvent });
-      }
+      },
     },
     [GraphEvents.TOOL_START]: {
-      handle: (_event: string, _data: t.StreamEventData, _metadata?: Record<string, unknown>): void => {
+      handle: (
+        _event: string,
+        _data: t.StreamEventData,
+        _metadata?: Record<string, unknown>
+      ): void => {
         // Handle tool start
-      }
+      },
     },
   });
 
@@ -93,7 +127,8 @@ describe(`${capitalizeFirstLetter(provider)} Streaming Tests`, () => {
         type: 'standard',
         llmConfig,
         tools: [new Calculator()],
-        instructions: 'You are a friendly AI assistant. Always address the user by their name.',
+        instructions:
+          'You are a friendly AI assistant. Always address the user by their name.',
         additional_instructions: `The user's name is ${userName} and they are located in ${location}.`,
       },
       returnContent: true,
@@ -109,7 +144,9 @@ describe(`${capitalizeFirstLetter(provider)} Streaming Tests`, () => {
 
     const finalContentParts = await run.processStream(inputs, config);
     expect(finalContentParts).toBeDefined();
-    const allTextParts = finalContentParts?.every((part) => part.type === ContentTypes.TEXT);
+    const allTextParts = finalContentParts?.every(
+      (part) => part.type === ContentTypes.TEXT
+    );
     expect(allTextParts).toBe(true);
     expect(collectedUsage.length).toBeGreaterThan(0);
     expect(collectedUsage[0].input_tokens).toBeGreaterThan(0);
@@ -117,26 +154,30 @@ describe(`${capitalizeFirstLetter(provider)} Streaming Tests`, () => {
 
     const finalMessages = run.getRunMessages();
     expect(finalMessages).toBeDefined();
-    conversationHistory.push(...finalMessages ?? []);
+    conversationHistory.push(...(finalMessages ?? []));
     expect(conversationHistory.length).toBeGreaterThan(1);
     runningHistory = conversationHistory.slice();
 
     expect(onMessageDeltaSpy).toHaveBeenCalled();
     expect(onMessageDeltaSpy.mock.calls.length).toBeGreaterThan(1);
-    expect((onMessageDeltaSpy.mock.calls[0][3] as StandardGraph).provider).toBeDefined();
+    expect(onMessageDeltaSpy.mock.calls[0][3]).toBeDefined(); // Graph exists
 
     expect(onRunStepSpy).toHaveBeenCalled();
     expect(onRunStepSpy.mock.calls.length).toBeGreaterThan(0);
-    expect((onRunStepSpy.mock.calls[0][3] as StandardGraph).provider).toBeDefined();
+    expect(onRunStepSpy.mock.calls[0][3]).toBeDefined(); // Graph exists
 
     const { handleLLMEnd, collected } = createMetadataAggregator();
     const titleResult = await run.generateTitle({
+      provider,
       inputText: userMessage,
+      titleMethod: TitleMethod.STRUCTURED,
       contentParts,
       chainOptions: {
-        callbacks: [{
-          handleLLMEnd,
-        }],
+        callbacks: [
+          {
+            handleLLMEnd,
+          },
+        ],
       },
     });
 
@@ -146,9 +187,66 @@ describe(`${capitalizeFirstLetter(provider)} Streaming Tests`, () => {
     expect(collected).toBeDefined();
   });
 
+  test(`${capitalizeFirstLetter(provider)}: should generate title using completion method`, async () => {
+    const { userName, location } = await getArgs();
+    const llmConfig = getLLMConfig(provider);
+    const customHandlers = setupCustomHandlers();
+
+    run = await Run.create<t.IState>({
+      runId: 'test-run-id-completion',
+      graphConfig: {
+        type: 'standard',
+        llmConfig,
+        tools: [new Calculator()],
+        instructions:
+          'You are a friendly AI assistant. Always address the user by their name.',
+        additional_instructions: `The user's name is ${userName} and they are located in ${location}.`,
+      },
+      returnContent: true,
+      customHandlers,
+    });
+
+    const userMessage = 'What is the weather like today?';
+    conversationHistory = [];
+    conversationHistory.push(new HumanMessage(userMessage));
+
+    const inputs = {
+      messages: conversationHistory,
+    };
+
+    const finalContentParts = await run.processStream(inputs, config);
+    expect(finalContentParts).toBeDefined();
+
+    const { handleLLMEnd, collected } = createMetadataAggregator();
+    const titleResult = await run.generateTitle({
+      provider,
+      inputText: userMessage,
+      titleMethod: TitleMethod.COMPLETION, // Using completion method
+      contentParts,
+      chainOptions: {
+        callbacks: [
+          {
+            handleLLMEnd,
+          },
+        ],
+      },
+    });
+
+    expect(titleResult).toBeDefined();
+    expect(titleResult.title).toBeDefined();
+    expect(titleResult.title).not.toBe('');
+    // Completion method doesn't return language
+    expect(titleResult.language).toBeUndefined();
+    expect(collected).toBeDefined();
+    console.log(`Completion method generated title: "${titleResult.title}"`);
+  });
+
   test(`${capitalizeFirstLetter(provider)}: should follow-up`, async () => {
     console.log('Previous conversation length:', runningHistory.length);
-    console.log('Last message:', runningHistory[runningHistory.length - 1].content);
+    console.log(
+      'Last message:',
+      runningHistory[runningHistory.length - 1].content
+    );
     const { userName, location } = await getArgs();
     const llmConfig = getLLMConfig(provider);
     const customHandlers = setupCustomHandlers();
@@ -159,7 +257,8 @@ describe(`${capitalizeFirstLetter(provider)} Streaming Tests`, () => {
         type: 'standard',
         llmConfig,
         tools: [new Calculator()],
-        instructions: 'You are a friendly AI assistant. Always address the user by their name.',
+        instructions:
+          'You are a friendly AI assistant. Always address the user by their name.',
         additional_instructions: `The user's name is ${userName} and they are located in ${location}.`,
       },
       returnContent: true,
@@ -175,7 +274,9 @@ describe(`${capitalizeFirstLetter(provider)} Streaming Tests`, () => {
 
     const finalContentParts = await run.processStream(inputs, config);
     expect(finalContentParts).toBeDefined();
-    const allTextParts = finalContentParts?.every((part) => part.type === ContentTypes.TEXT);
+    const allTextParts = finalContentParts?.every(
+      (part) => part.type === ContentTypes.TEXT
+    );
     expect(allTextParts).toBe(true);
     expect(collectedUsage.length).toBeGreaterThan(0);
     expect(collectedUsage[0].input_tokens).toBeGreaterThan(0);
@@ -184,7 +285,10 @@ describe(`${capitalizeFirstLetter(provider)} Streaming Tests`, () => {
     const finalMessages = run.getRunMessages();
     expect(finalMessages).toBeDefined();
     expect(finalMessages?.length).toBeGreaterThan(0);
-    console.log(`${capitalizeFirstLetter(provider)} follow-up message:`, finalMessages?.[finalMessages.length - 1]?.content);
+    console.log(
+      `${capitalizeFirstLetter(provider)} follow-up message:`,
+      finalMessages?.[finalMessages.length - 1]?.content
+    );
 
     expect(onMessageDeltaSpy).toHaveBeenCalled();
     expect(onMessageDeltaSpy.mock.calls.length).toBeGreaterThan(1);
@@ -196,9 +300,12 @@ describe(`${capitalizeFirstLetter(provider)} Streaming Tests`, () => {
   test('should handle errors appropriately', async () => {
     // Test error scenarios
     await expect(async () => {
-      await run.processStream({
-        messages: [],
-      }, {} as any);
+      await run.processStream(
+        {
+          messages: [],
+        },
+        {} as any
+      );
     }).rejects.toThrow();
   });
-});
+});

package/src/specs/openrouter.simple.test.ts

@@ -0,0 +1,107 @@
+import { config } from 'dotenv';
+config();
+import { Calculator } from '@langchain/community/tools/calculator';
+import {
+  HumanMessage,
+  BaseMessage,
+  UsageMetadata,
+} from '@langchain/core/messages';
+import type * as t from '@/types';
+import { ToolEndHandler, ModelEndHandler } from '@/events';
+import { ContentTypes, GraphEvents, Providers, TitleMethod } from '@/common';
+import { ChatModelStreamHandler, createContentAggregator } from '@/stream';
+import { capitalizeFirstLetter } from './spec.utils';
+import { getLLMConfig } from '@/utils/llmConfig';
+import { getArgs } from '@/scripts/args';
+import { Run } from '@/run';
+
+// Auto-skip if OpenRouter env is missing
+const hasOpenRouter = (process.env.OPENROUTER_API_KEY ?? '').trim() !== '';
+const describeIf = hasOpenRouter ? describe : describe.skip;
+
+const provider = Providers.OPENROUTER;
+describeIf(`${capitalizeFirstLetter(provider)} Streaming Tests`, () => {
+  jest.setTimeout(60000);
+  let run: Run<t.IState>;
+  let collectedUsage: UsageMetadata[];
+  let conversationHistory: BaseMessage[];
+  let contentParts: t.MessageContentComplex[];
+
+  const configV2 = {
+    configurable: { thread_id: 'or-convo-1' },
+    streamMode: 'values',
+    version: 'v2' as const,
+  };
+
+  beforeEach(async () => {
+    conversationHistory = [];
+    collectedUsage = [];
+    const { contentParts: cp } = createContentAggregator();
+    contentParts = cp as t.MessageContentComplex[];
+  });
+
+  const onMessageDeltaSpy = jest.fn();
+  const onRunStepSpy = jest.fn();
+
+  afterAll(() => {
+    onMessageDeltaSpy.mockReset();
+    onRunStepSpy.mockReset();
+  });
+
+  const setupCustomHandlers = (): Record<
+    string | GraphEvents,
+    t.EventHandler
+  > => ({
+    [GraphEvents.TOOL_END]: new ToolEndHandler(),
+    [GraphEvents.CHAT_MODEL_END]: new ModelEndHandler(collectedUsage),
+    [GraphEvents.CHAT_MODEL_STREAM]: new ChatModelStreamHandler(),
+  });
+
+  test(`${capitalizeFirstLetter(provider)}: simple stream + title`, async () => {
+    const { userName, location } = await getArgs();
+    const llmConfig = getLLMConfig(provider);
+    const customHandlers = setupCustomHandlers();
+
+    run = await Run.create<t.IState>({
+      runId: 'or-run-1',
+      graphConfig: {
+        type: 'standard',
+        llmConfig,
+        tools: [new Calculator()],
+        instructions: 'You are a friendly AI assistant.',
+        additional_instructions: `The user's name is ${userName} and they are located in ${location}.`,
+      },
+      returnContent: true,
+      customHandlers,
+    });
+
+    const userMessage = 'hi';
+    conversationHistory.push(new HumanMessage(userMessage));
+
+    const finalContentParts = await run.processStream(
+      { messages: conversationHistory },
+      configV2
+    );
+    expect(finalContentParts).toBeDefined();
+    const allTextParts = finalContentParts?.every(
+      (part) => part.type === ContentTypes.TEXT
+    );
+    expect(allTextParts).toBe(true);
+    expect(
+      (collectedUsage[0]?.input_tokens ?? 0) +
+        (collectedUsage[0]?.output_tokens ?? 0)
+    ).toBeGreaterThan(0);
+
+    const finalMessages = run.getRunMessages();
+    expect(finalMessages).toBeDefined();
+    conversationHistory.push(...(finalMessages ?? []));
+
+    const titleRes = await run.generateTitle({
+      provider,
+      inputText: userMessage,
+      titleMethod: TitleMethod.COMPLETION,
+      contentParts,
+    });
+    expect(titleRes.title).toBeDefined();
+  });
+});

package/src/specs/prune.test.ts

@@ -725,8 +725,11 @@ describe('Prune Messages Tests', () => {
           type: 'standard',
           llmConfig,
           instructions: 'You are a helpful assistant.',
+          maxContextTokens: 1000,
         },
         returnContent: true,
+        tokenCounter,
+        indexTokenCountMap: {},
       });
 
       // Override the model to use a fake LLM
@@ -734,10 +737,6 @@ describe('Prune Messages Tests', () => {
 
       const messages = [new HumanMessage('Hello, how are you?')];
 
-      const indexTokenCountMap = {
-        0: tokenCounter(messages[0]),
-      };
-
       const config: Partial<RunnableConfig> & {
         version: 'v1' | 'v2';
         streamMode: string;
@@ -749,11 +748,7 @@ describe('Prune Messages Tests', () => {
         version: 'v2' as const,
       };
 
-      await run.processStream({ messages }, config, {
-        maxContextTokens: 1000,
-        indexTokenCountMap,
-        tokenCounter,
-      });
+      await run.processStream({ messages }, config);
 
       const finalMessages = run.getRunMessages();
       expect(finalMessages).toBeDefined();

package/src/specs/reasoning.test.ts

@@ -3,9 +3,12 @@
 // src/scripts/cli.test.ts
 import { config } from 'dotenv';
 config();
-import { HumanMessage, BaseMessage, MessageContentText } from '@langchain/core/messages';
+import {
+  HumanMessage,
+  BaseMessage,
+  MessageContentText,
+} from '@langchain/core/messages';
 import type { RunnableConfig } from '@langchain/core/runnables';
-import type { StandardGraph } from '@/graphs';
 import type * as t from '@/types';
 import { ChatModelStreamHandler, createContentAggregator } from '@/stream';
 import { capitalizeFirstLetter } from './spec.utils';
@@ -28,40 +31,49 @@ describe(`${capitalizeFirstLetter(provider)} Streaming Tests`, () => {
   let aggregateContent: t.ContentAggregator;
   let runSteps: Set<string>;
 
-  const config: Partial<RunnableConfig> & { version: 'v1' | 'v2'; run_id?: string; streamMode: string } = {
+  const config: Partial<RunnableConfig> & {
+    version: 'v1' | 'v2';
+    run_id?: string;
+    streamMode: string;
+  } = {
     configurable: {
       thread_id: 'conversation-num-1',
     },
     streamMode: 'values',
     version: 'v2' as const,
-    callbacks: [{
-      async handleCustomEvent(event, data, metadata): Promise<void> {
-        if (event !== GraphEvents.ON_MESSAGE_DELTA) {
-          return;
-        }
-        const messageDeltaData = data as t.MessageDeltaEvent;
-
-        // Wait until we see the run step (with timeout for safety)
-        const maxAttempts = 50; // 5 seconds total
-        let attempts = 0;
-        while (!runSteps.has(messageDeltaData.id) && attempts < maxAttempts) {
-          await new Promise(resolve => setTimeout(resolve, 100));
-          attempts++;
-        }
-
-        if (!runSteps.has(messageDeltaData.id)) {
-          console.warn(`Timeout waiting for run step: ${messageDeltaData.id}`);
-        }
-
-        onMessageDeltaSpy(event, data, metadata, run.Graph);
-        aggregateContent({ event, data: messageDeltaData });
+    callbacks: [
+      {
+        async handleCustomEvent(event, data, metadata): Promise<void> {
+          if (event !== GraphEvents.ON_MESSAGE_DELTA) {
+            return;
+          }
+          const messageDeltaData = data as t.MessageDeltaEvent;
+
+          // Wait until we see the run step (with timeout for safety)
+          const maxAttempts = 50; // 5 seconds total
+          let attempts = 0;
+          while (!runSteps.has(messageDeltaData.id) && attempts < maxAttempts) {
+            await new Promise((resolve) => setTimeout(resolve, 100));
+            attempts++;
+          }
+
+          if (!runSteps.has(messageDeltaData.id)) {
+            console.warn(
+              `Timeout waiting for run step: ${messageDeltaData.id}`
+            );
+          }
+
+          onMessageDeltaSpy(event, data, metadata, run.Graph);
+          aggregateContent({ event, data: messageDeltaData });
+        },
       },
-    }],
+    ],
   };
 
   beforeEach(async () => {
     conversationHistory = [];
-    const { contentParts: parts, aggregateContent: ac } = createContentAggregator();
+    const { contentParts: parts, aggregateContent: ac } =
+      createContentAggregator();
     aggregateContent = ac;
     runSteps = new Set();
     contentParts = parts as t.MessageContentComplex[];
@@ -81,32 +93,54 @@ describe(`${capitalizeFirstLetter(provider)} Streaming Tests`, () => {
     onRunStepSpy.mockReset();
   });
 
-  const setupCustomHandlers = (): Record<string | GraphEvents, t.EventHandler> => ({
+  const setupCustomHandlers = (): Record<
+    string | GraphEvents,
+    t.EventHandler
+  > => ({
     [GraphEvents.CHAT_MODEL_STREAM]: new ChatModelStreamHandler(),
     [GraphEvents.ON_RUN_STEP_COMPLETED]: {
-      handle: (event: GraphEvents.ON_RUN_STEP_COMPLETED, data: t.StreamEventData): void => {
-        aggregateContent({ event, data: data as unknown as { result: t.ToolEndEvent; } });
-      }
+      handle: (
+        event: GraphEvents.ON_RUN_STEP_COMPLETED,
+        data: t.StreamEventData
+      ): void => {
+        aggregateContent({
+          event,
+          data: data as unknown as { result: t.ToolEndEvent },
+        });
+      },
     },
     [GraphEvents.ON_RUN_STEP]: {
-      handle: (event: GraphEvents.ON_RUN_STEP, data: t.StreamEventData, metadata, graph): void => {
+      handle: (
+        event: GraphEvents.ON_RUN_STEP,
+        data: t.StreamEventData,
+        metadata,
+        graph
+      ): void => {
         const runStepData = data as t.RunStep;
         runSteps.add(runStepData.id);
 
         onRunStepSpy(event, runStepData, metadata, graph);
         aggregateContent({ event, data: runStepData });
-      }
+      },
     },
     [GraphEvents.ON_RUN_STEP_DELTA]: {
-      handle: (event: GraphEvents.ON_RUN_STEP_DELTA, data: t.StreamEventData): void => {
+      handle: (
+        event: GraphEvents.ON_RUN_STEP_DELTA,
+        data: t.StreamEventData
+      ): void => {
         aggregateContent({ event, data: data as t.RunStepDeltaEvent });
-      }
+      },
     },
     [GraphEvents.ON_REASONING_DELTA]: {
-      handle: (event: GraphEvents.ON_REASONING_DELTA, data: t.StreamEventData, metadata, graph): void => {
+      handle: (
+        event: GraphEvents.ON_REASONING_DELTA,
+        data: t.StreamEventData,
+        metadata,
+        graph
+      ): void => {
         onReasoningDeltaSpy(event, data, metadata, graph);
         aggregateContent({ event, data: data as t.ReasoningDeltaEvent });
-      }
+      },
     },
   });
 
@@ -120,7 +154,8 @@ describe(`${capitalizeFirstLetter(provider)} Streaming Tests`, () => {
       graphConfig: {
         type: 'standard',
         llmConfig,
-        instructions: 'You are a friendly AI assistant. Always address the user by their name.',
+        instructions:
+          'You are a friendly AI assistant. Always address the user by their name.',
         additional_instructions: `The user's name is ${userName} and they are located in ${location}.`,
       },
       returnContent: true,
@@ -141,25 +176,26 @@ describe(`${capitalizeFirstLetter(provider)} Streaming Tests`, () => {
     expect(contentParts.length).toBe(2);
     const reasoningContent = reasoningText.match(/<think>(.*)<\/think>/s)?.[0];
     const content = reasoningText.split(/<\/think>/)[1];
-    expect((contentParts[0] as t.ReasoningContentText).think).toBe(reasoningContent);
+    expect((contentParts[0] as t.ReasoningContentText).think).toBe(
+      reasoningContent
+    );
     expect((contentParts[1] as MessageContentText).text).toBe(content);
 
     const finalMessages = run.getRunMessages();
     expect(finalMessages).toBeDefined();
-    conversationHistory.push(...finalMessages ?? []);
+    conversationHistory.push(...(finalMessages ?? []));
     expect(conversationHistory.length).toBeGreaterThan(1);
 
     expect(onMessageDeltaSpy).toHaveBeenCalled();
     expect(onMessageDeltaSpy.mock.calls.length).toBeGreaterThan(1);
-    expect((onMessageDeltaSpy.mock.calls[0][3] as StandardGraph).provider).toBeDefined();
+    expect(onMessageDeltaSpy.mock.calls[0][3]).toBeDefined(); // Graph exists
 
     expect(onReasoningDeltaSpy).toHaveBeenCalled();
     expect(onReasoningDeltaSpy.mock.calls.length).toBeGreaterThan(1);
-    expect((onReasoningDeltaSpy.mock.calls[0][3] as StandardGraph).provider).toBeDefined();
+    expect(onReasoningDeltaSpy.mock.calls[0][3]).toBeDefined(); // Graph exists
 
     expect(onRunStepSpy).toHaveBeenCalled();
     expect(onRunStepSpy.mock.calls.length).toBeGreaterThan(0);
-    expect((onRunStepSpy.mock.calls[0][3] as StandardGraph).provider).toBeDefined();
-
+    expect(onRunStepSpy.mock.calls[0][3]).toBeDefined(); // Graph exists
   });
-});
+});

package/src/specs/token-memoization.test.ts

@@ -0,0 +1,39 @@
+import { HumanMessage } from '@langchain/core/messages';
+import { createTokenCounter } from '@/utils/tokens';
+
+describe('Token encoder memoization', () => {
+  jest.setTimeout(45000);
+
+  test('fetches BPE once and reuses encoder across counters', async () => {
+    const originalFetch = global.fetch;
+    let fetchCalls = 0;
+    global.fetch = (async (...args: Parameters<typeof fetch>) => {
+      fetchCalls += 1;
+      // Delegate to real fetch
+      return await originalFetch(...args);
+    }) as typeof fetch;
+
+    try {
+      const counter1 = await createTokenCounter();
+      const counter2 = await createTokenCounter();
+
+      const m1 = new HumanMessage('hello world');
+      const m2 = new HumanMessage('another short text');
+
+      const c11 = counter1(m1);
+      const c12 = counter1(m2);
+      const c21 = counter2(m1);
+      const c22 = counter2(m2);
+
+      expect(c11).toBeGreaterThan(0);
+      expect(c12).toBeGreaterThan(0);
+      expect(c21).toBe(c11);
+      expect(c22).toBe(c12);
+
+      // Only one fetch for the shared encoder
+      expect(fetchCalls).toBe(1);
+    } finally {
+      global.fetch = originalFetch;
+    }
+  });
+});