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

packages/core/src/context/utils/adaptiveTokenCalculator.test.ts

@@ -0,0 +1,125 @@
+/**
+ * @license
+ * Copyright 2026 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { describe, it, expect } from 'vitest';
+import { AdaptiveTokenCalculator } from './adaptiveTokenCalculator.js';
+import { NodeBehaviorRegistry } from '../graph/behaviorRegistry.js';
+import { registerBuiltInBehaviors } from '../graph/builtinBehaviors.js';
+import { ContextEventBus } from '../eventBus.js';
+import { createDummyNode } from '../testing/contextTestUtils.js';
+import { NodeType } from '../graph/types.js';
+
+describe('AdaptiveTokenCalculator', () => {
+  const registry = new NodeBehaviorRegistry();
+  registerBuiltInBehaviors(registry);
+  const charsPerToken = 1; // Simplifies math
+
+  it('should initialize with a learned weight of 1.0', () => {
+    const eventBus = new ContextEventBus();
+    const calculator = new AdaptiveTokenCalculator(
+      charsPerToken,
+      registry,
+      eventBus,
+    );
+    expect(calculator.getLearnedWeight()).toBe(1.0);
+  });
+
+  it('should dynamically update learned weight based on token ground truth events', () => {
+    const eventBus = new ContextEventBus();
+    const calculator = new AdaptiveTokenCalculator(
+      charsPerToken,
+      registry,
+      eventBus,
+    );
+
+    // Initial state: weight = 1.0
+
+    // Simulate an event where the API reported fewer tokens than our base units
+    // targetWeight = 50 / 100 = 0.5
+    // newWeight = 1.0 * 0.8 + 0.5 * 0.2 = 0.8 + 0.1 = 0.9
+    eventBus.emitTokenGroundTruth({
+      actualTokens: 50,
+      promptBaseUnits: 100,
+    });
+
+    // JavaScript floating point precision means we should use toBeCloseTo
+    expect(calculator.getLearnedWeight()).toBeCloseTo(0.9, 5);
+
+    // Simulate another event
+    // newWeight = 0.9 * 0.8 + (150 / 100) * 0.2 = 0.72 + 0.3 = 1.02
+    eventBus.emitTokenGroundTruth({
+      actualTokens: 150,
+      promptBaseUnits: 100,
+    });
+
+    expect(calculator.getLearnedWeight()).toBeCloseTo(1.02, 5);
+  });
+
+  it('should clamp the learned weight between 0.5 and 2.0', () => {
+    const eventBus = new ContextEventBus();
+    const calculator = new AdaptiveTokenCalculator(
+      charsPerToken,
+      registry,
+      eventBus,
+    );
+
+    // Push weight up extremely high (API returns 10x tokens)
+    for (let i = 0; i < 20; i++) {
+      eventBus.emitTokenGroundTruth({
+        actualTokens: 1000,
+        promptBaseUnits: 100,
+      });
+    }
+    expect(calculator.getLearnedWeight()).toBe(2.0);
+
+    // Push weight down extremely low (API returns 0 tokens)
+    for (let i = 0; i < 20; i++) {
+      eventBus.emitTokenGroundTruth({ actualTokens: 0, promptBaseUnits: 100 });
+    }
+    expect(calculator.getLearnedWeight()).toBe(0.5);
+  });
+
+  it('should correctly apply the learned weight to node calculations while keeping raw base units stable', () => {
+    const eventBus = new ContextEventBus();
+    const calculator = new AdaptiveTokenCalculator(
+      charsPerToken,
+      registry,
+      eventBus,
+    );
+
+    // Decrease the weight to exactly 0.5
+    for (let i = 0; i < 20; i++) {
+      eventBus.emitTokenGroundTruth({ actualTokens: 0, promptBaseUnits: 100 });
+    }
+
+    const turn1Id = 'turn-1';
+    const node1 = createDummyNode(turn1Id, NodeType.USER_PROMPT);
+
+    // Get raw base units directly
+    const rawTokens = calculator.calculateTokensAndBaseUnits([node1]).baseUnits;
+
+    // Get adjusted tokens
+    const adjustedTokens = calculator.calculateConcreteListTokens([node1]);
+
+    expect(adjustedTokens).toBe(Math.round(rawTokens * 0.5));
+  });
+
+  it('should ignore ground truth events with 0 promptBaseUnits to prevent division by zero', () => {
+    const eventBus = new ContextEventBus();
+    const calculator = new AdaptiveTokenCalculator(
+      charsPerToken,
+      registry,
+      eventBus,
+    );
+
+    eventBus.emitTokenGroundTruth({
+      actualTokens: 100,
+      promptBaseUnits: 0,
+    });
+
+    expect(calculator.getLearnedWeight()).toBe(1.0);
+  });
+});

packages/core/src/context/utils/adaptiveTokenCalculator.ts

@@ -0,0 +1,163 @@
+/**
+ * @license
+ * Copyright 2026 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import type { Content, Part } from '@google/genai';
+import type { ConcreteNode } from '../graph/types.js';
+import {
+  StaticTokenCalculator,
+  type AdvancedTokenCalculator,
+} from './contextTokenCalculator.js';
+import type { NodeBehaviorRegistry } from '../graph/behaviorRegistry.js';
+import type { ContextEventBus, TokenGroundTruthEvent } from '../eventBus.js';
+import { debugLogger } from '../../utils/debugLogger.js';
+
+/**
+ * An Adaptive Token Calculator that dynamically learns the true token cost of the user's
+ * conversation by applying an Exponential Moving Average (EMA) gradient descent to
+ * real usage metadata returned from the Gemini API.
+ *
+ * It wraps the deterministic `StaticTokenCalculator` base heuristic to ensure
+ * immutable node cost caching while still surfacing a self-corrected estimate
+ * to the pipeline processors.
+ */
+export class AdaptiveTokenCalculator implements AdvancedTokenCalculator {
+  private learnedWeight = 1.0;
+  private readonly baseCalculator: StaticTokenCalculator;
+
+  constructor(
+    charsPerToken: number,
+    registry: NodeBehaviorRegistry,
+    eventBus: ContextEventBus,
+  ) {
+    this.baseCalculator = new StaticTokenCalculator(charsPerToken, registry);
+    eventBus.onTokenGroundTruth((event: TokenGroundTruthEvent) => {
+      this.handleGroundTruth(event.actualTokens, event.promptBaseUnits);
+    });
+  }
+
+  private handleGroundTruth(actualTokens: number, promptBaseUnits: number) {
+    if (promptBaseUnits <= 0) return;
+
+    // Determine what ratio we should have used
+    const targetWeight = actualTokens / promptBaseUnits;
+    const oldWeight = this.learnedWeight;
+
+    // Apply Momentum (Learning Rate)
+    const learningRate = 0.2;
+    const newWeight =
+      oldWeight * (1 - learningRate) + targetWeight * learningRate;
+
+    // Clamp to reasonable safety bounds to prevent rogue metadata poisoning the system
+    this.learnedWeight = Math.max(0.5, Math.min(newWeight, 2.0));
+
+    debugLogger.log(
+      `[AdaptiveTokenCalculator] Learned weight updated to ${this.learnedWeight.toFixed(3)} ` +
+        `(API Tokens: ${actualTokens}, Base Units: ${promptBaseUnits}, Target Ratio: ${targetWeight.toFixed(3)})`,
+    );
+  }
+
+  /**
+   * Retrieves the current learned weight multiplier.
+   */
+  getLearnedWeight(): number {
+    return this.learnedWeight;
+  }
+
+  /**
+   * Returns the exact, unweighted Base Heuristic Units for the graph.
+   * This is used exactly once per interaction to capture the baseline sent to the API.
+   */
+  getRawBaseUnits(nodes: readonly ConcreteNode[]): number {
+    return this.baseCalculator.calculateConcreteListTokens(nodes);
+  }
+
+  /**
+   * Returns the exact, unweighted Base Heuristic Units for a raw content chunk.
+   */
+  getRawBaseUnitsForContent(content: Content): number {
+    return this.baseCalculator.calculateContentTokens(content);
+  }
+
+  calculateTokensAndBaseUnits(nodes: readonly ConcreteNode[]): {
+    tokens: number;
+    baseUnits: number;
+  } {
+    const baseUnits = this.baseCalculator.calculateConcreteListTokens(nodes);
+    return {
+      tokens: Math.round(baseUnits * this.learnedWeight),
+      baseUnits,
+    };
+  }
+
+  calculateContentTokensAndBaseUnits(content: Content): {
+    tokens: number;
+    baseUnits: number;
+  } {
+    const baseUnits = this.baseCalculator.calculateContentTokens(content);
+    return {
+      tokens: Math.round(baseUnits * this.learnedWeight),
+      baseUnits,
+    };
+  }
+
+  // --- Delegation and Weighting ---
+
+  garbageCollectCache(liveNodeIds: ReadonlySet<string>): void {
+    this.baseCalculator.garbageCollectCache(liveNodeIds);
+  }
+
+  cacheNodeTokens(node: ConcreteNode): number {
+    return this.baseCalculator.cacheNodeTokens(node);
+  }
+
+  calculateTokenBreakdown(nodes: readonly ConcreteNode[]): {
+    text: number;
+    media: number;
+    tool: number;
+    overhead: number;
+    total: number;
+  } {
+    const raw = this.baseCalculator.calculateTokenBreakdown(nodes);
+    return {
+      text: Math.round(raw.text * this.learnedWeight),
+      media: Math.round(raw.media * this.learnedWeight),
+      tool: Math.round(raw.tool * this.learnedWeight),
+      overhead: Math.round(raw.overhead * this.learnedWeight),
+      total: Math.round(raw.total * this.learnedWeight),
+    };
+  }
+
+  estimateTokensForParts(parts: Part[]): number {
+    const baseUnits = this.baseCalculator.estimateTokensForParts(parts);
+    return Math.round(baseUnits * this.learnedWeight);
+  }
+
+  getTokenCost(node: ConcreteNode): number {
+    const baseUnits = this.baseCalculator.getTokenCost(node);
+    return Math.round(baseUnits * this.learnedWeight);
+  }
+
+  calculateConcreteListTokens(nodes: readonly ConcreteNode[]): number {
+    const baseUnits = this.baseCalculator.calculateConcreteListTokens(nodes);
+    return Math.round(baseUnits * this.learnedWeight);
+  }
+
+  calculateContentTokens(content: Content): number {
+    const baseUnits = this.baseCalculator.calculateContentTokens(content);
+    return Math.round(baseUnits * this.learnedWeight);
+  }
+
+  estimateTokensForString(text: string): number {
+    const baseUnits = this.baseCalculator.estimateTokensForString(text);
+    return Math.round(baseUnits * this.learnedWeight);
+  }
+
+  tokensToChars(tokens: number): number {
+    // If weight is > 1.0 (we are inflating tokens), a single returned token is worth fewer chars.
+    // We reverse the math: convert requested tokens to target base units, then get chars.
+    return this.baseCalculator.tokensToChars(tokens / this.learnedWeight);
+  }
+}

packages/core/src/context/utils/contextTokenCalculator.test.ts

@@ -5,7 +5,7 @@
  */
 
 import { describe, it, expect } from 'vitest';
-import { ContextTokenCalculator } from './contextTokenCalculator.js';
+import { StaticTokenCalculator } from './contextTokenCalculator.js';
 import { NodeBehaviorRegistry } from '../graph/behaviorRegistry.js';
 import { registerBuiltInBehaviors } from '../graph/builtinBehaviors.js';
 import { createDummyNode } from '../testing/contextTestUtils.js';
@@ -16,7 +16,7 @@ describe('ContextTokenCalculator', () => {
   const registry = new NodeBehaviorRegistry();
   registerBuiltInBehaviors(registry);
   const charsPerToken = 1; // Simplifies math for text nodes in tests
-  const calculator = new ContextTokenCalculator(charsPerToken, registry);
+  const calculator = new StaticTokenCalculator(charsPerToken, registry);
 
   it('should include structural overhead for each unique turn', () => {
     const turn1Id = 'turn-1';
@@ -28,16 +28,16 @@ describe('ContextTokenCalculator', () => {
 
     const nodes = [node1, node2, node3];
 
-    // Estimated tokens (using 0.33 per ASCII char heuristic):
-    // node1: floor(17 chars * 0.33) = 5 tokens
-    // node2: floor(17 chars * 0.33) = 5 tokens
-    // node3: floor(19 chars * 0.33) = 6 tokens
+    // Estimated tokens (using charsPerToken = 1):
+    // node1: 17 chars / 1 = 17 tokens
+    // node2: 17 chars / 1 = 17 tokens
+    // node3: 19 chars / 1 = 19 tokens
     // Turn 1 overhead: 5 tokens
     // Turn 2 overhead: 5 tokens
-    // Total: 5 + 5 + 6 + 5 + 5 = 26
+    // Total: 17 + 17 + 19 + 5 + 5 = 63
 
     const total = calculator.calculateConcreteListTokens(nodes);
-    expect(total).toBe(26);
+    expect(total).toBe(63);
   });
 
   it('should handle categorical breakdown with overhead', () => {

packages/core/src/context/utils/contextTokenCalculator.ts

@@ -17,7 +17,41 @@ import type { NodeBehaviorRegistry } from '../graph/behaviorRegistry.js';
  * by the Gemini API. We use this as a baseline heuristic for inlineData/fileData.
  */
 
-export class ContextTokenCalculator {
+export interface ContextTokenCalculator {
+  estimateTokensForString(text: string): number;
+  tokensToChars(tokens: number): number;
+  garbageCollectCache(liveNodeIds: ReadonlySet<string>): void;
+  cacheNodeTokens(node: ConcreteNode): number;
+  getTokenCost(node: ConcreteNode): number;
+  calculateTokenBreakdown(nodes: readonly ConcreteNode[]): {
+    text: number;
+    media: number;
+    tool: number;
+    overhead: number;
+    total: number;
+  };
+  calculateConcreteListTokens(nodes: readonly ConcreteNode[]): number;
+  calculateContentTokens(content: Content): number;
+  estimateTokensForParts(parts: Part[]): number;
+}
+
+export interface AdvancedTokenCalculator extends ContextTokenCalculator {
+  getRawBaseUnits(nodes: readonly ConcreteNode[]): number;
+  getRawBaseUnitsForContent(content: Content): number;
+  calculateTokensAndBaseUnits(nodes: readonly ConcreteNode[]): {
+    tokens: number;
+    baseUnits: number;
+  };
+  calculateContentTokensAndBaseUnits(content: Content): {
+    tokens: number;
+    baseUnits: number;
+  };
+}
+
+/**
+ * A fast, deterministic token heuristic calculator.
+ */
+export class StaticTokenCalculator implements AdvancedTokenCalculator {
   private readonly tokenCache = new Map<string, number>();
 
   constructor(
@@ -143,6 +177,34 @@ export class ContextTokenCalculator {
     return breakdown;
   }
 
+  /**
+   * For the static calculator, Raw Base Units are exactly the same as the final tokens,
+   * because there is no dynamic learned weight (the multiplier is effectively 1.0).
+   */
+  getRawBaseUnits(nodes: readonly ConcreteNode[]): number {
+    return this.calculateConcreteListTokens(nodes);
+  }
+
+  getRawBaseUnitsForContent(content: Content): number {
+    return this.calculateContentTokens(content);
+  }
+
+  calculateTokensAndBaseUnits(nodes: readonly ConcreteNode[]): {
+    tokens: number;
+    baseUnits: number;
+  } {
+    const baseUnits = this.calculateConcreteListTokens(nodes);
+    return { tokens: baseUnits, baseUnits };
+  }
+
+  calculateContentTokensAndBaseUnits(content: Content): {
+    tokens: number;
+    baseUnits: number;
+  } {
+    const baseUnits = this.calculateContentTokens(content);
+    return { tokens: baseUnits, baseUnits };
+  }
+
   /**
    * Fast calculation for a flat array of ConcreteNodes (The Nodes).
    * It relies entirely on the O(1) sidecar token cache.

packages/core/src/context/utils/snapshotGenerator.test.ts

@@ -21,6 +21,9 @@ describe('SnapshotGenerator', () => {
       llmClient: {
         generateJson: mockGenerateJson,
       },
+      advancedTokenCalculator: {
+        getRawBaseUnits: vi.fn().mockReturnValue(100),
+      },
       tokenCalculator: {
         estimateTokensForString: vi.fn().mockReturnValue(100),
       },