feat(context): Introduce adaptive token calculator to more accurately calculate content sizes. (#26888)

packages/core/src/context/graph/render.test.ts

@@ -8,6 +8,7 @@ import { describe, it, expect, vi } from 'vitest';
 import { render } from './render.js';
 import type { ConcreteNode } from './types.js';
 import { NodeType } from './types.js';
+import type { AdvancedTokenCalculator } from '../utils/contextTokenCalculator.js';
 import type { ContextEnvironment } from '../pipeline/environment.js';
 import type { ContextTracer } from '../tracer.js';
 import type { ContextProfile } from '../config/profiles.js';
@@ -37,7 +38,20 @@ describe('render', () => {
 
     const orchestrator = {} as PipelineOrchestrator;
     const sidecar = { config: {} } as ContextProfile; // No budget
+    const mockAdvancedTokenCalculator = {
+      calculateTokensAndBaseUnits: vi.fn().mockReturnValue({
+        tokens: 100,
+        baseUnits: 100,
+      }),
+      getRawBaseUnits: vi.fn().mockReturnValue(100),
+      getRawBaseUnitsForContent: vi.fn().mockReturnValue(0),
+    };
+
     const env = {
+      tokenCalculator: {
+        calculateConcreteListTokens: vi.fn().mockReturnValue(100),
+        calculateTokenBreakdown: vi.fn().mockReturnValue({}),
+      },
       graphMapper: {
         fromGraph: vi.fn((nodes: readonly ConcreteNode[]) =>
           nodes.map((n) => ({ text: n.id })),
@@ -54,12 +68,14 @@ describe('render', () => {
       sidecar,
       tracer,
       env,
+      mockAdvancedTokenCalculator as unknown as AdvancedTokenCalculator,
       new Map(),
-      0,
+      undefined,
       previewNodeIds,
     );
 
     expect(result.history).toEqual([{ text: '1' }, { text: '2' }]);
+    expect(result.baseUnits).toBe(100);
   });
 
   it('simulates the boundary knapsack problem (loose boundary policy)', async () => {
@@ -108,12 +124,24 @@ describe('render', () => {
 
     const currentTokens = 160000;
 
+    const mockAdvancedTokenCalculator = {
+      calculateTokensAndBaseUnits: vi.fn((nodes: readonly ConcreteNode[]) => {
+        const tokens =
+          nodes.length === 1 ? tokenMap[nodes[0].id] : currentTokens;
+        return { tokens, baseUnits: tokens };
+      }),
+      getRawBaseUnits: vi.fn((nodes: readonly ConcreteNode[]) => {
+        if (nodes.length === 1) return tokenMap[nodes[0].id];
+        return currentTokens;
+      }),
+    };
+
     const env = {
       llmClient: {
         countTokens: vi.fn().mockResolvedValue({ totalTokens: 1000 }),
       },
       tokenCalculator: {
-        calculateConcreteListTokens: vi.fn((nodes) => {
+        calculateConcreteListTokens: vi.fn((nodes: readonly ConcreteNode[]) => {
           if (nodes.length === 1) return tokenMap[nodes[0].id];
           return currentTokens;
         }),
@@ -136,8 +164,9 @@ describe('render', () => {
       sidecar,
       tracer,
       env,
+      mockAdvancedTokenCalculator as unknown as AdvancedTokenCalculator,
       new Map(),
-      0,
+      undefined,
       new Set(),
     );
 
@@ -147,6 +176,7 @@ describe('render', () => {
     // Adding C pushes rolling total (70k) above retainedTokens (65k).
     // Under loose policy, C survives. D is strictly older and drops.
     expect(surviving).toEqual(['C', 'B', 'A']); // D is dropped
+    expect(result.baseUnits).toBe(160000);
   });
 
   it('drops nodes that are STRICTLY older than the boundary node', async () => {
@@ -188,12 +218,24 @@ describe('render', () => {
 
     const currentTokens = 160000;
 
+    const mockAdvancedTokenCalculator = {
+      calculateTokensAndBaseUnits: vi.fn((nodes: readonly ConcreteNode[]) => {
+        const tokens =
+          nodes.length === 1 ? tokenMap[nodes[0].id] : currentTokens;
+        return { tokens, baseUnits: tokens };
+      }),
+      getRawBaseUnits: vi.fn((nodes: readonly ConcreteNode[]) => {
+        if (nodes.length === 1) return tokenMap[nodes[0].id];
+        return currentTokens;
+      }),
+    };
+
     const env = {
       llmClient: {
         countTokens: vi.fn().mockResolvedValue({ totalTokens: 1000 }),
       },
       tokenCalculator: {
-        calculateConcreteListTokens: vi.fn((nodes) => {
+        calculateConcreteListTokens: vi.fn((nodes: readonly ConcreteNode[]) => {
           if (nodes.length === 1) return tokenMap[nodes[0].id];
           return currentTokens;
         }),
@@ -216,8 +258,9 @@ describe('render', () => {
       sidecar,
       tracer,
       env,
+      mockAdvancedTokenCalculator as unknown as AdvancedTokenCalculator,
       new Map(),
-      0,
+      undefined,
       new Set(),
     );
 
@@ -225,5 +268,6 @@ describe('render', () => {
     const surviving = result.history.map((c: any) => c.text);
     // C(40k), B(40k). Adding B pushes total to 80k. B is the boundary node and survives. A drops.
     expect(surviving).toEqual(['B', 'C']); // A is dropped
+    expect(result.baseUnits).toBe(160000);
   });
 });