feat/adds-delta-snapshot-token-optimization

CHANGELOG.md

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- **Delta snapshot optimization**: Post-action ARIA snapshots within a step now return only the lines that changed since the last snapshot, instead of the full page tree. This cuts token usage by ~92% on complex pages (large tables, dashboards, CRM grids) — from ~1M tokens to ~82k for a 10-action step. The first snapshot of each step is always full. A built-in savings-ratio threshold (20%) ensures delta mode never sends more tokens than a full snapshot would. Enable/disable globally: `configure({ deltaSnapshot: false })`. Default: `true`.
 - **OpenAI Support**: Direct integration with OpenAI models via `@ai-sdk/openai` and `OPENAI_API_KEY`
 - `maxRetries` option to `AssertionOptions` (default: `1`) to control how many times a failed assertion is retried with a fresh page snapshot and screenshot. Setting it to `0` disables retries.
 - `onRetry` callback to `AssertionOptions` that fires before each retry, receiving the retry index and the full `AssertionResult` from the previous attempt for debugging flaky assertions.

src/__tests__/config.test.ts

@@ -1,5 +1,5 @@
 import { describe, it, expect, beforeEach } from "vitest";
-import { configure, getConfig, getModelId, resetConfig, DEFAULT_MODELS } from "../config";
+import { configure, getConfig, getModelId, getDeltaSnapshotEnabled, resetConfig, DEFAULT_MODELS } from "../config";
 
 describe("config", () => {
   beforeEach(() => {
@@ -73,6 +73,27 @@ describe("config", () => {
     expect(getConfig().uploadBasePath).toBe("/tmp/test-uploads");
   });
 
+  it("getDeltaSnapshotEnabled returns true by default", () => {
+    expect(getDeltaSnapshotEnabled()).toBe(true);
+  });
+
+  it("getDeltaSnapshotEnabled returns false when configured off", () => {
+    configure({ deltaSnapshot: false });
+    expect(getDeltaSnapshotEnabled()).toBe(false);
+  });
+
+  it("getDeltaSnapshotEnabled returns true when explicitly configured on", () => {
+    configure({ deltaSnapshot: true });
+    expect(getDeltaSnapshotEnabled()).toBe(true);
+  });
+
+  it("configure merges deltaSnapshot without overwriting other keys", () => {
+    configure({ uploadBasePath: "./uploads" });
+    configure({ deltaSnapshot: false });
+    expect(getConfig().uploadBasePath).toBe("./uploads");
+    expect(getConfig().deltaSnapshot).toBe(false);
+  });
+
   it("resetConfig clears everything", () => {
     configure({
       uploadBasePath: "./uploads",

src/__tests__/integration/run-steps.test.ts

@@ -21,7 +21,7 @@ vi.mock("ai", async (importOriginal) => {
   const actual = await importOriginal<typeof import("ai")>();
   return {
     ...actual,
-    generateText: vi.fn().mockResolvedValue({ text: "done", steps: [], output: {} }),
+    generateText: vi.fn().mockResolvedValue({ text: "done", steps: [], output: {}, usage: { inputTokens: 100, outputTokens: 50, totalTokens: 150 } }),
     streamText: vi.fn(),
   };
 });
@@ -275,7 +275,7 @@ describe("runSteps", () => {
     vi.mocked(generateText).mockImplementation(async (_opts: unknown) => {
       // Extract the step description from the prompt to track order
       callOrder.push(`generateText-call-${callOrder.length + 1}`);
-      return { text: "done", steps: [] } as unknown as Awaited<ReturnType<typeof generateText>>;
+      return { text: "done", steps: [], usage: { inputTokens: 0, outputTokens: 0, totalTokens: 0 } } as unknown as Awaited<ReturnType<typeof generateText>>;
     });
 
     const steps: Step[] = [

src/__tests__/snapshot-diff.test.ts

@@ -0,0 +1,117 @@
+import { describe, it, expect } from "vitest";
+import { computeSnapshotDiff } from "../utils/snapshot-diff";
+
+// 200-line snapshot large enough for the 20% savings threshold to be met on small diffs
+const makeSnapshot = (lines: string[]) => lines.join("\n");
+const BASE_LINES = Array.from({ length: 200 }, (_, i) => `  row: "Item ${i}"`);
+const LARGE_SNAPSHOT = makeSnapshot(BASE_LINES);
+
+describe("computeSnapshotDiff", () => {
+  it("returns unchanged message when snapshots are identical", () => {
+    const { diff, isFull, savedChars } = computeSnapshotDiff(LARGE_SNAPSHOT, LARGE_SNAPSHOT);
+    expect(diff).toBe("[snapshot unchanged — action may not have had a visible DOM effect]");
+    expect(isFull).toBe(false);
+    expect(savedChars).toBe(LARGE_SNAPSHOT.length);
+  });
+
+  it("identical after trimming whitespace is treated as unchanged", () => {
+    const snap = `${LARGE_SNAPSHOT}   `;
+    const { diff } = computeSnapshotDiff(snap, snap);
+    expect(diff).toBe("[snapshot unchanged — action may not have had a visible DOM effect]");
+  });
+
+  it("returns a diff with + and - prefixes for a single-line change", () => {
+    const afterLines = [...BASE_LINES];
+    afterLines[100] = '  checkbox [checked] "Select Item 100"';
+    const after = makeSnapshot(afterLines);
+
+    const { diff, isFull } = computeSnapshotDiff(LARGE_SNAPSHOT, after);
+
+    expect(isFull).toBe(false);
+    expect(diff).toContain('- ' + BASE_LINES[100]);
+    expect(diff).toContain('+ ' + afterLines[100]);
+  });
+
+  it("includes a header with added/removed counts", () => {
+    const afterLines = [...BASE_LINES];
+    afterLines[50] = '  row: "Changed Item 50"';
+    const after = makeSnapshot(afterLines);
+
+    const { diff } = computeSnapshotDiff(LARGE_SNAPSHOT, after);
+
+    expect(diff).toMatch(/\[delta snapshot: \+\d+ lines added, -\d+ lines removed\]/);
+  });
+
+  it("includes context lines around changes", () => {
+    const afterLines = [...BASE_LINES];
+    afterLines[50] = '  row: "Modified"';
+    const after = makeSnapshot(afterLines);
+
+    const { diff } = computeSnapshotDiff(LARGE_SNAPSHOT, after);
+
+    // 2 lines of context on each side of the changed line
+    expect(diff).toContain(`  ${BASE_LINES[48]}`);
+    expect(diff).toContain(`  ${BASE_LINES[49]}`);
+    expect(diff).toContain(`  ${BASE_LINES[51]}`);
+    expect(diff).toContain(`  ${BASE_LINES[52]}`);
+  });
+
+  it("uses '...' separators for skipped lines far from any change", () => {
+    const afterLines = [...BASE_LINES];
+    afterLines[100] = '  row: "Changed"';
+    const after = makeSnapshot(afterLines);
+
+    const { diff } = computeSnapshotDiff(LARGE_SNAPSHOT, after);
+
+    expect(diff).toContain("...");
+  });
+
+  it("falls back to full snapshot when savings ratio is below threshold", () => {
+    // Completely different content — diff would be larger than the 20% savings threshold
+    const after = Array.from({ length: 200 }, (_, i) => `  row: "Different ${i}"`).join("\n");
+    const { diff, isFull, savedChars } = computeSnapshotDiff(LARGE_SNAPSHOT, after);
+    expect(isFull).toBe(true);
+    expect(savedChars).toBe(0);
+    expect(diff).toBe(after);
+  });
+
+  it("returns diff (not full) when change is small relative to snapshot size", () => {
+    const afterLines = [...BASE_LINES];
+    afterLines[100] = '  row: "Modified Item 100"';
+    const after = makeSnapshot(afterLines);
+
+    const { isFull, savedChars } = computeSnapshotDiff(LARGE_SNAPSHOT, after);
+    expect(isFull).toBe(false);
+    expect(savedChars).toBeGreaterThan(0);
+  });
+
+  it("handles lines being added (new row appearing)", () => {
+    const afterLines = [...BASE_LINES];
+    afterLines.splice(100, 0, '  row: "NEW ITEM"');
+    const after = makeSnapshot(afterLines);
+
+    const { diff, isFull } = computeSnapshotDiff(LARGE_SNAPSHOT, after);
+    expect(isFull).toBe(false);
+    expect(diff).toContain('+ ' + '  row: "NEW ITEM"');
+  });
+
+  it("handles lines being removed (row disappearing)", () => {
+    const afterLines = [...BASE_LINES];
+    afterLines.splice(100, 1);
+    const after = makeSnapshot(afterLines);
+
+    const { diff, isFull } = computeSnapshotDiff(LARGE_SNAPSHOT, after);
+    expect(isFull).toBe(false);
+    expect(diff).toContain('- ' + BASE_LINES[100]);
+  });
+
+  it("savedChars equals after.length minus diff.length when not full", () => {
+    const afterLines = [...BASE_LINES];
+    afterLines[10] = '  row: "Changed Item 10"';
+    const after = makeSnapshot(afterLines);
+
+    const { diff, isFull, savedChars } = computeSnapshotDiff(LARGE_SNAPSHOT, after);
+    expect(isFull).toBe(false);
+    expect(savedChars).toBe(after.length - diff.length);
+  });
+});

src/config.ts

@@ -104,6 +104,21 @@ type Config = {
    * after the assertions consume them.
    */
   videoDir?: string;
+  /**
+   * When true (default), post-action snapshots within a single step only
+   * include DOM nodes that changed since the last snapshot, dramatically
+   * reducing token usage on complex pages (large tables, dashboards, grids).
+   *
+   * The first snapshot of each step is always full so the agent has complete
+   * orientation. Subsequent snapshots are diffed. A built-in savings
+   * threshold falls back to a full snapshot if the diff would not save
+   * enough — delta mode is never worse than full mode.
+   *
+   * Set to `false` to disable globally and always return full snapshots.
+   *
+   * @default true
+   */
+  deltaSnapshot?: boolean;
 };
 
 let globalConfig: Config = {};
@@ -161,6 +176,14 @@ export function getMode(): AIMode {
   return getConfig().ai?.mode ?? "snapshot";
 }
 
+/**
+ * Returns whether delta snapshots are enabled. Defaults to `true` — every
+ * snapshot after the first within a step returns only the diff.
+ */
+export function getDeltaSnapshotEnabled(): boolean {
+  return getConfig().deltaSnapshot ?? true;
+}
+
 /**
  * Effective AI config for a single step / call after merging overrides with
  * the global config. `getModelId` looks up a model with the same precedence

src/constants.ts

@@ -25,6 +25,9 @@ export const THINKING_BUDGET_DEFAULT = 1024;
 // Redis
 export const GLOBAL_VALUES_TTL_SECONDS = 86400;
 
+// Delta snapshot
+export const DELTA_SNAPSHOT_MIN_SAVINGS_RATIO = 0.2;
+
 // Video assertions
 export const VIDEO_DEFAULT_DIR = "/tmp/passmark-recordings";
 export const VIDEO_DEFAULT_WIDTH = 1280;

src/index.ts

@@ -537,6 +537,14 @@ export const runSteps = async ({
           }),
       );
 
+      // Log token usage for the step
+      logger.info(
+        `[token usage] step: "${step.description}" | ` +
+        `prompt: ${result.usage.inputTokens} | ` +
+        `completion: ${result.usage.outputTokens} | ` +
+        `total: ${result.usage.totalTokens}`
+      );
+
       // Cache the step action only if it was a single tool call (simple, deterministic action).
       // Multi-step actions are not cached as they may be non-deterministic.
       const allToolCalls = result.steps

src/tools.ts

@@ -3,10 +3,11 @@ import { z } from "zod";
 import { Locator, type Page } from "@playwright/test";
 import { wrapTool } from "axiom/ai";
 import shortid from "shortid";
-import { getConfig } from "./config";
+import { getConfig, getDeltaSnapshotEnabled } from "./config";
 import { isAxiomEnabled } from "./instrumentation";
 import { logger } from "./logger";
 import { LOCATOR_ACTION_TIMEOUT, SNAPSHOT_TIMEOUT, STOP_DELAY } from "./constants";
+import { computeSnapshotDiff } from "./utils/snapshot-diff";
 import {
   PlaywrightTestArgs,
   PlaywrightTestOptions,
@@ -260,6 +261,8 @@ class PlaywrightTools {
   private tabManager?: TabManager;
   private currentStep;
   private abortController?: AbortController;
+  private deltaSnapshotEnabled: boolean;
+  private lastSnapshot: string | null = null;
   public pendingCacheData: Record<string, string> | null = null;
 
   private get page(): Page {
@@ -273,11 +276,30 @@ class PlaywrightTools {
     this.tabManager = tabManager;
     this.currentStep = currentStep;
     this.abortController = abortController;
+    this.deltaSnapshotEnabled = getDeltaSnapshotEnabled();
   }
 
   public async getSnapshot() {
-    const snapshot = await this.page.ariaSnapshot({ mode: "ai", timeout: SNAPSHOT_TIMEOUT });
-    return `url: ${this.page.url()}\n\n${snapshot}`;
+    const raw = await this.page.ariaSnapshot({ mode: "ai", timeout: SNAPSHOT_TIMEOUT });
+    const full = `url: ${this.page.url()}\n\n${raw}`;
+
+    if (!this.deltaSnapshotEnabled || this.lastSnapshot === null) {
+      this.lastSnapshot = full;
+      return full;
+    }
+
+    const { diff, isFull, savedChars } = computeSnapshotDiff(this.lastSnapshot, full);
+    this.lastSnapshot = full;
+
+    if (isFull) {
+      logger.debug("Delta snapshot: change ratio too high, returning full snapshot");
+      return full;
+    }
+
+    logger.debug(
+      `Delta snapshot: -${savedChars.toLocaleString()} chars saved on step "${this.currentStep?.description}"`,
+    );
+    return diff;
   }
 
   public navigateSchema = z.object({
@@ -291,6 +313,7 @@ class PlaywrightTools {
   });
   public async navigate({ url }: z.infer<typeof this.navigateSchema>) {
     await this.page.goto(url, { waitUntil: "load" });
+    this.lastSnapshot = null;
     return { success: true, url };
   }
 
@@ -392,16 +415,19 @@ class PlaywrightTools {
 
   public async goBack() {
     await this.page.goBack();
+    this.lastSnapshot = null;
     return { success: true };
   }
 
   public async goForward() {
     await this.page.goForward();
+    this.lastSnapshot = null;
     return { success: true };
   }
 
   public async reload() {
     await this.page.reload({ waitUntil: "load" });
+    this.lastSnapshot = null;
     return { success: true };
   }