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`.
 - **OpenCode Zen gateway support**: set `gateway: "opencodezen"` in `configure()` and provide `OPENCODEZEN_API_KEY` to route all model requests through [OpenCode Zen](https://opencode.ai/docs/ko/zen/) (`https://opencode.ai/zen/v1`), an OpenAI-compatible gateway with 30+ curated models including Claude, Gemini, GPT, Qwen, and more.
 - **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.

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/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

@@ -526,6 +526,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
@@ -753,7 +761,7 @@ export const runUserFlow = async ({
       ? resolveModel(effectiveAi.getModelId("userFlowLow"), effectiveAi.gateway)
       : resolveModel(effectiveAi.getModelId("userFlowHigh"), effectiveAi.gateway);
 
-  const { tools } = getAItools(page, {
+  const { tools, resetLastSnapshot } = getAItools(page, {
     abortController,
   });
 
@@ -783,6 +791,7 @@ export const runUserFlow = async ({
           prepareStep: async ({ messages }) => {
             // Remove older messages to keep the context window small
             if (messages.length > 11) {
+              resetLastSnapshot();
               const modifiedMessages = [messages[0], ...messages.slice(-10)];
               return {
                 messages: modifiedMessages,

src/tools.ts

@@ -7,6 +7,7 @@ import { getConfig } 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,
@@ -58,7 +59,7 @@ export function getAItools(page: Page, settings?: ToolSettings) {
       const snapshot = await playwrightTools.getSnapshot();
       return { ...result, snapshot };
     } catch (_error) {
-      return `Error executing this action. Retry the action or try a different one.\n\nLatest Snapshot:\n\n${await playwrightTools.getSnapshot()}`;
+      return `Error executing this action. Retry the action or try a different one.\n\nLatest Snapshot:\n\n${await playwrightTools.getSnapshot({ forceFull: true })}`;
     }
   };
 
@@ -151,7 +152,7 @@ export function getAItools(page: Page, settings?: ToolSettings) {
           reasoning: z.string().describe("A quick one-line reasoning behind this action"),
         }),
         execute: async (_args) => {
-          return await playwrightTools.getSnapshot();
+          return await playwrightTools.getSnapshot({ forceFull: true });
         },
       }),
     ),
@@ -258,6 +259,7 @@ export function getAItools(page: Page, settings?: ToolSettings) {
     clearPendingCacheData: () => {
       playwrightTools.pendingCacheData = null;
     },
+    resetLastSnapshot: () => playwrightTools.resetLastSnapshot(),
   };
 }
 
@@ -266,6 +268,7 @@ class PlaywrightTools {
   private tabManager?: TabManager;
   private currentStep;
   private abortController?: AbortController;
+  private lastSnapshot: string | null = null;
   public pendingCacheData: Record<string, string> | null = null;
 
   private get page(): Page {
@@ -281,9 +284,31 @@ class PlaywrightTools {
     this.abortController = abortController;
   }
 
-  public async getSnapshot() {
-    const snapshot = await this.page.ariaSnapshot({ mode: "ai", timeout: SNAPSHOT_TIMEOUT });
-    return `url: ${this.page.url()}\n\n${snapshot}`;
+  public async getSnapshot({ forceFull = false }: { forceFull?: boolean } = {}) {
+    const raw = await this.page.ariaSnapshot({ mode: "ai", timeout: SNAPSHOT_TIMEOUT });
+    const full = `url: ${this.page.url()}\n\n${raw}`;
+
+    if (this.lastSnapshot === null || forceFull) {
+      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 resetLastSnapshot() {
+    this.lastSnapshot = null;
   }
 
   public navigateSchema = z.object({
@@ -297,6 +322,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 };
   }
 
@@ -398,16 +424,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 };
   }