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);
   });
 });

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

@@ -11,6 +11,7 @@ import type { ContextProfile } from '../config/profiles.js';
 import type { PipelineOrchestrator } from '../pipeline/orchestrator.js';
 import type { ContextEnvironment } from '../pipeline/environment.js';
 import { performCalibration } from '../utils/tokenCalibration.js';
+import type { AdvancedTokenCalculator } from '../utils/contextTokenCalculator.js';
 
 /**
  * Maps the Episodic Context Graph back into a raw Gemini Content[] array for transmission.
@@ -22,21 +23,43 @@ export async function render(
   sidecar: ContextProfile,
   tracer: ContextTracer,
   env: ContextEnvironment,
+  advancedTokenCalculator: AdvancedTokenCalculator,
   protectionReasons: Map<string, string> = new Map(),
-  headerTokens: number = 0,
+  header?: Content,
   previewNodeIds: ReadonlySet<string> = new Set(),
-): Promise<{ history: Content[]; didApplyManagement: boolean }> {
+): Promise<{
+  history: Content[];
+  didApplyManagement: boolean;
+  baseUnits: number;
+}> {
+  let headerTokens = 0;
+  let headerBaseUnits = 0;
+  if (header) {
+    const costs =
+      advancedTokenCalculator.calculateContentTokensAndBaseUnits(header);
+    headerTokens = costs.tokens;
+    headerBaseUnits = costs.baseUnits;
+  }
+
   if (!sidecar.config.budget) {
     const visibleNodes = nodes.filter((n) => !previewNodeIds.has(n.id));
     const contents = env.graphMapper.fromGraph(visibleNodes);
     tracer.logEvent('Render', 'Render Context to LLM (No Budget)', {
       renderedContext: contents,
     });
-    return { history: contents, didApplyManagement: false };
+
+    // In all cases, retrieve raw base units from the token calculator interface
+    const baseUnits =
+      advancedTokenCalculator.getRawBaseUnits(nodes) + headerBaseUnits;
+
+    return { history: contents, didApplyManagement: false, baseUnits };
   }
 
   const maxTokens = sidecar.config.budget.maxTokens;
-  const graphTokens = env.tokenCalculator.calculateConcreteListTokens(nodes);
+
+  const { tokens: graphTokens, baseUnits: graphBaseUnits } =
+    advancedTokenCalculator.calculateTokensAndBaseUnits(nodes);
+
   const currentTokens = graphTokens + headerTokens;
 
   const protectedIds = new Set(protectionReasons.keys());
@@ -70,7 +93,11 @@ export async function render(
       renderedContext: contents,
     });
     performCalibration(env, visibleNodes, contents);
-    return { history: contents, didApplyManagement: false };
+    return {
+      history: contents,
+      didApplyManagement: false,
+      baseUnits: graphBaseUnits + headerBaseUnits,
+    };
   }
   const targetDelta = currentTokens - sidecar.config.budget.retainedTokens;
   tracer.logEvent(
@@ -119,5 +146,10 @@ export async function render(
     renderedContextSanitized: contents,
   });
   performCalibration(env, visibleNodes, contents);
-  return { history: contents, didApplyManagement: true };
+  return {
+    history: contents,
+    didApplyManagement: true,
+    baseUnits:
+      advancedTokenCalculator.getRawBaseUnits(visibleNodes) + headerBaseUnits,
+  };
 }