fix(context): Fix snapshot recovery across sessions. (#26939)

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

@@ -122,4 +122,69 @@ describe('AdaptiveTokenCalculator', () => {
 
     expect(calculator.getLearnedWeight()).toBe(1.0);
   });
+
+  it('should subtract overhead tokens from actual tokens when determining target weight', () => {
+    const eventBus = new ContextEventBus();
+    const getOverheadTokens = () => 40;
+    const calculator = new AdaptiveTokenCalculator(
+      charsPerToken,
+      registry,
+      eventBus,
+      getOverheadTokens,
+    );
+
+    // Initial state: weight = 1.0
+
+    // Simulate an event where the API reported 100 tokens, and our base units were 100
+    // But overhead is 40.
+    // actualGraphTokens = 100 - 40 = 60
+    // rawTargetWeight = 60 / 100 = 0.6
+    // targetWeight = Math.max(0.5, 0.6) = 0.6
+    // newWeight = 1.0 * 0.8 + 0.6 * 0.2 = 0.8 + 0.12 = 0.92
+    eventBus.emitTokenGroundTruth({
+      actualTokens: 100,
+      promptBaseUnits: 100,
+    });
+
+    expect(calculator.getLearnedWeight()).toBeCloseTo(0.92, 5);
+  });
+
+  it('should enforce the maxStep limit to prevent violent oscillation from massive outliers', () => {
+    const eventBus = new ContextEventBus();
+    const maxStep = 0.05; // Tight limit
+    const calculator = new AdaptiveTokenCalculator(
+      charsPerToken,
+      registry,
+      eventBus,
+      undefined,
+      { maxStep },
+    );
+
+    // Initial state: weight = 1.0
+
+    // Simulate a massive outlier where the API reports 10,000 tokens for 100 base units.
+    // rawTargetWeight = 100
+    // targetWeight = Math.min(100, 1.0 * 2.0) = 2.0
+    // emaWeight = 1.0 * 0.8 + 2.0 * 0.2 = 1.2
+    // BUT maxStep is 0.05, so the actual step is clamped.
+    // finalWeight = 1.0 + 0.05 = 1.05
+    eventBus.emitTokenGroundTruth({
+      actualTokens: 10000,
+      promptBaseUnits: 100,
+    });
+
+    expect(calculator.getLearnedWeight()).toBeCloseTo(1.05, 5);
+
+    // Simulate a massive under-estimation
+    // rawTargetWeight = 0
+    // targetWeight = Math.max(0, 1.05 * 0.5) = 0.525
+    // emaWeight = 1.05 * 0.8 + 0.525 * 0.2 = 0.84 + 0.105 = 0.945
+    // BUT maxStep is 0.05, so step is clamped: 1.05 - 0.05 = 1.0
+    eventBus.emitTokenGroundTruth({
+      actualTokens: 0,
+      promptBaseUnits: 100,
+    });
+
+    expect(calculator.getLearnedWeight()).toBeCloseTo(1.0, 5);
+  });
 });

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

@@ -14,6 +14,13 @@ import type { NodeBehaviorRegistry } from '../graph/behaviorRegistry.js';
 import type { ContextEventBus, TokenGroundTruthEvent } from '../eventBus.js';
 import { debugLogger } from '../../utils/debugLogger.js';
 
+export interface AdaptiveLearningConfig {
+  /** The momentum factor for the Exponential Moving Average (EMA). Defaults to 0.2. */
+  learningRate?: number;
+  /** The absolute maximum change allowed to the weight in a single turn. Defaults to 0.15. */
+  maxStep?: number;
+}
+
 /**
  * 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
@@ -26,12 +33,18 @@ import { debugLogger } from '../../utils/debugLogger.js';
 export class AdaptiveTokenCalculator implements AdvancedTokenCalculator {
   private learnedWeight = 1.0;
   private readonly baseCalculator: StaticTokenCalculator;
+  private readonly learningRate: number;
+  private readonly maxStep: number;
 
   constructor(
     charsPerToken: number,
     registry: NodeBehaviorRegistry,
     eventBus: ContextEventBus,
+    private readonly getOverheadTokens?: () => number,
+    config?: AdaptiveLearningConfig,
   ) {
+    this.learningRate = config?.learningRate ?? 0.2;
+    this.maxStep = config?.maxStep ?? 0.15;
     this.baseCalculator = new StaticTokenCalculator(charsPerToken, registry);
     eventBus.onTokenGroundTruth((event: TokenGroundTruthEvent) => {
       this.handleGroundTruth(event.actualTokens, event.promptBaseUnits);
@@ -41,21 +54,44 @@ export class AdaptiveTokenCalculator implements AdvancedTokenCalculator {
   private handleGroundTruth(actualTokens: number, promptBaseUnits: number) {
     if (promptBaseUnits <= 0) return;
 
+    const overheadTokens = this.getOverheadTokens
+      ? this.getOverheadTokens()
+      : 0;
+
+    // The Gemini API token count includes the static overhead (system instruction + tools)
+    // and the dynamic chat history (which we measure as promptBaseUnits).
+    // We subtract the overhead so the adaptive calculator is comparing "apples to apples"
+    // when learning the weight multiplier for the graph nodes.
+    const actualGraphTokens = Math.max(0, actualTokens - overheadTokens);
+
     // Determine what ratio we should have used
-    const targetWeight = actualTokens / promptBaseUnits;
+    const rawTargetWeight = actualGraphTokens / promptBaseUnits;
     const oldWeight = this.learnedWeight;
 
-    // Apply Momentum (Learning Rate)
-    const learningRate = 0.2;
-    const newWeight =
-      oldWeight * (1 - learningRate) + targetWeight * learningRate;
+    // Dampen extreme outliers *before* applying the EMA by capping the target weight
+    // to a reasonable multiple of the current weight. This prevents a single massive
+    // anomaly from destroying the running average.
+    const targetWeight = Math.max(
+      oldWeight * 0.5,
+      Math.min(rawTargetWeight, oldWeight * 2.0),
+    );
 
-    // Clamp to reasonable safety bounds to prevent rogue metadata poisoning the system
+    // Apply Momentum (Learning Rate)
+    let newWeight =
+      oldWeight * (1 - this.learningRate) + targetWeight * this.learningRate;
+
+    // Hard limit the maximum step size per turn to prevent violent oscillation
+    if (newWeight > oldWeight + this.maxStep)
+      newWeight = oldWeight + this.maxStep;
+    if (newWeight < oldWeight - this.maxStep)
+      newWeight = oldWeight - this.maxStep;
+
+    // Clamp to reasonable absolute safety bounds
     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)})`,
+        `(API Tokens: ${actualTokens}, Overhead: ${overheadTokens}, Graph Tokens: ${actualGraphTokens}, Base Units: ${promptBaseUnits}, Target Ratio: ${targetWeight.toFixed(3)})`,
     );
   }
 

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

@@ -5,12 +5,132 @@
  */
 
 import { describe, it, expect, vi, beforeEach } from 'vitest';
-import { SnapshotGenerator, type SnapshotState } from './snapshotGenerator.js';
+import {
+  SnapshotGenerator,
+  type SnapshotState,
+  SnapshotStateHelper,
+} from './snapshotGenerator.js';
 import type { ContextEnvironment } from '../pipeline/environment.js';
 import { NodeType, type ConcreteNode } from '../graph/types.js';
 
 import type { Mock } from 'vitest';
 
+describe('SnapshotStateHelper', () => {
+  describe('exportState', () => {
+    it('should flatten nested abstractsIds to pristine IDs', () => {
+      // Setup a graph with nested snapshots
+      // S3 abstracts [S2, N5]
+      // S2 abstracts [S1, N3, N4]
+      // S1 abstracts [N1, N2]
+
+      const nodes: ConcreteNode[] = [
+        {
+          id: 'N1',
+          type: NodeType.USER_PROMPT,
+          timestamp: 10,
+          role: 'user',
+          payload: { text: '1' },
+          turnId: 'T1',
+        },
+        {
+          id: 'N2',
+          type: NodeType.AGENT_THOUGHT,
+          timestamp: 20,
+          role: 'model',
+          payload: { text: '2' },
+          turnId: 'T1',
+        },
+        {
+          id: 'S1',
+          type: NodeType.SNAPSHOT,
+          timestamp: 30,
+          role: 'user',
+          payload: { text: 'State 1' },
+          turnId: 'S1',
+          abstractsIds: ['N1', 'N2'],
+        },
+        {
+          id: 'N3',
+          type: NodeType.USER_PROMPT,
+          timestamp: 40,
+          role: 'user',
+          payload: { text: '3' },
+          turnId: 'T2',
+        },
+        {
+          id: 'N4',
+          type: NodeType.AGENT_THOUGHT,
+          timestamp: 50,
+          role: 'model',
+          payload: { text: '4' },
+          turnId: 'T2',
+        },
+        {
+          id: 'S2',
+          type: NodeType.SNAPSHOT,
+          timestamp: 60,
+          role: 'user',
+          payload: { text: 'State 2' },
+          turnId: 'S2',
+          abstractsIds: ['S1', 'N3', 'N4'],
+        },
+        {
+          id: 'N5',
+          type: NodeType.USER_PROMPT,
+          timestamp: 70,
+          role: 'user',
+          payload: { text: '5' },
+          turnId: 'T3',
+        },
+        {
+          id: 'S3',
+          type: NodeType.SNAPSHOT,
+          timestamp: 80,
+          role: 'user',
+          payload: { text: 'State 3' },
+          turnId: 'S3',
+          abstractsIds: ['S2', 'N5'],
+        },
+      ];
+
+      const state = SnapshotStateHelper.exportState(nodes);
+
+      expect(state.snapshot).toBeDefined();
+      expect(state.snapshot?.text).toBe('State 3');
+
+      // Should be flattened to only the "pristine" (non-snapshot) IDs
+      const consumedIds = state.snapshot?.consumedIds;
+      expect(consumedIds).toContain('N1');
+      expect(consumedIds).toContain('N2');
+      expect(consumedIds).toContain('N3');
+      expect(consumedIds).toContain('N4');
+      expect(consumedIds).toContain('N5');
+
+      // Should NOT contain the intermediate snapshot IDs
+      expect(consumedIds).not.toContain('S1');
+      expect(consumedIds).not.toContain('S2');
+
+      expect(consumedIds?.length).toBe(5);
+    });
+
+    it('should return empty state if no snapshot baseline is found', () => {
+      const nodes: ConcreteNode[] = [
+        {
+          id: 'N1',
+          type: NodeType.USER_PROMPT,
+          timestamp: 10,
+          role: 'user',
+          payload: { text: '1' },
+          turnId: 'T1',
+        },
+      ];
+
+      const state = SnapshotStateHelper.exportState(nodes);
+      expect(state).toEqual({});
+    });
+  });
+});
+
 describe('SnapshotGenerator', () => {
   let mockEnv: ContextEnvironment;
   let mockGenerateJson: Mock;

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

@@ -48,10 +48,39 @@ export interface SnapshotState {
   recent_arc: string[];
 }
 
+import { debugLogger } from '../../utils/debugLogger.js';
+
+export function isSnapshotState(text: string): boolean {
+  const trimmed = text.trim();
+  if (!trimmed.startsWith('{') || !trimmed.endsWith('}')) {
+    return false;
+  }
+  try {
+    const parsed: unknown = JSON.parse(trimmed);
+    if (!isRecord(parsed)) return false;
+    const isSnap =
+      Array.isArray(parsed['active_tasks']) &&
+      Array.isArray(parsed['discovered_facts']) &&
+      Array.isArray(parsed['constraints_and_preferences']) &&
+      Array.isArray(parsed['recent_arc']);
+    if (!isSnap) {
+      debugLogger.log(
+        '[isSnapshotState] FAILED FOR JSON:',
+        JSON.stringify(parsed),
+      );
+    }
+    return isSnap;
+  } catch {
+    debugLogger.log('[isSnapshotState] PARSE FAILED FOR:', trimmed);
+    return false;
+  }
+}
+
 export interface BaselineSnapshotInfo {
   text: string;
   abstractsIds: string[];
   id: string;
+  timestamp: number;
 }
 
 /**
@@ -61,6 +90,20 @@ export interface BaselineSnapshotInfo {
 export function findLatestSnapshotBaseline(
   targets: readonly ConcreteNode[],
 ): BaselineSnapshotInfo | undefined {
+  debugLogger.log(
+    '[findLatestSnapshotBaseline] Targets:',
+    targets.map((t) => ({
+      id: t.id,
+      type: t.type,
+      text:
+        t.payload &&
+        typeof t.payload === 'object' &&
+        'text' in t.payload &&
+        typeof t.payload.text === 'string'
+          ? t.payload.text.substring(0, 20)
+          : '',
+    })),
+  );
   const lastSnapshotNode = [...targets]
     .reverse()
     .find((n) => n.type === NodeType.SNAPSHOT && n.payload.text);
@@ -72,8 +115,10 @@ export function findLatestSnapshotBaseline(
         ? [...lastSnapshotNode.abstractsIds]
         : [],
       id: lastSnapshotNode.id,
+      timestamp: lastSnapshotNode.timestamp,
     };
   }
+
   return undefined;
 }
 
@@ -326,3 +371,61 @@ ${formatNodesForLlm(nodes)}`;
     return JSON.stringify(newState);
   }
 }
+
+/**
+ * Shared logic for working with Snapshot node state.
+ */
+export class SnapshotStateHelper {
+  /**
+   * Flatten nested abstract IDs to only the "pristine" (non-snapshot) IDs.
+   */
+  static flattenAbstracts(
+    nodes: ConcreteNode[],
+    abstractsIds: readonly string[],
+  ): string[] {
+    const pristineIds: string[] = [];
+    const nodeMap = new Map(nodes.map((n) => [n.id, n]));
+
+    const walk = (ids: readonly string[]) => {
+      for (const id of ids) {
+        const node = nodeMap.get(id);
+        if (!node) {
+          // Fallback: if node not in map, treat as pristine ID
+          pristineIds.push(id);
+          continue;
+        }
+
+        if (node.type === NodeType.SNAPSHOT && node.abstractsIds) {
+          walk(node.abstractsIds);
+        } else {
+          pristineIds.push(id);
+        }
+      }
+    };
+
+    walk(abstractsIds);
+    return Array.from(new Set(pristineIds)); // Dedupe
+  }
+
+  /**
+   * Helper to extract state from the most recent snapshot in a list of nodes.
+   */
+  static exportState(nodes: ConcreteNode[]): {
+    snapshot?: { text: string; consumedIds: string[] };
+  } {
+    const baseline = findLatestSnapshotBaseline(nodes);
+    if (!baseline) return {};
+
+    const node = nodes.find((n) => n.id === baseline.id);
+    if (!node || node.type !== NodeType.SNAPSHOT) return {};
+
+    const consumedIds = this.flattenAbstracts(nodes, node.abstractsIds || []);
+
+    return {
+      snapshot: {
+        text: baseline.text,
+        consumedIds,
+      },
+    };
+  }
+}