diff --git a/CHANGELOG.md b/CHANGELOG.md index ca22584..6f4714c 100644 --- a/CHANGELOG.md +++ b/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. diff --git a/src/__tests__/integration/run-steps.test.ts b/src/__tests__/integration/run-steps.test.ts index 5dd590b..428fe42 100644 --- a/src/__tests__/integration/run-steps.test.ts +++ b/src/__tests__/integration/run-steps.test.ts @@ -21,7 +21,7 @@ vi.mock("ai", async (importOriginal) => { const actual = await importOriginal(); 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>; + return { text: "done", steps: [], usage: { inputTokens: 0, outputTokens: 0, totalTokens: 0 } } as unknown as Awaited>; }); const steps: Step[] = [ diff --git a/src/__tests__/snapshot-diff.test.ts b/src/__tests__/snapshot-diff.test.ts new file mode 100644 index 0000000..fc19b44 --- /dev/null +++ b/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); + }); +}); diff --git a/src/constants.ts b/src/constants.ts index 8c14452..c5c096d 100644 --- a/src/constants.ts +++ b/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; diff --git a/src/index.ts b/src/index.ts index 75ee9d4..cefdaa6 100644 --- a/src/index.ts +++ b/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, diff --git a/src/tools.ts b/src/tools.ts index e02f8b3..3de2dc7 100644 --- a/src/tools.ts +++ b/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 | 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) { 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 }; } diff --git a/src/utils/snapshot-diff.ts b/src/utils/snapshot-diff.ts new file mode 100644 index 0000000..ffd5fe3 --- /dev/null +++ b/src/utils/snapshot-diff.ts @@ -0,0 +1,133 @@ +import { DELTA_SNAPSHOT_MIN_SAVINGS_RATIO } from "../constants"; + +/** + * Computes a human-readable structural diff between two ariaSnapshot strings. + * + * Returns a compact diff showing added/removed lines with 2 lines of context. + * Falls back to the full snapshot if the diff would not save at least + * DELTA_SNAPSHOT_MIN_SAVINGS_RATIO of the original size — delta mode is + * never worse than full mode. + */ +export function computeSnapshotDiff( + before: string, + after: string, +): { diff: string; isFull: boolean; savedChars: number } { + if (before.trim() === after.trim()) { + return { + diff: "[snapshot unchanged — action may not have had a visible DOM effect]", + isFull: false, + savedChars: after.length, + }; + } + + const beforeLines = before.split("\n"); + const afterLines = after.split("\n"); + + const diff = computeLineDiff(beforeLines, afterLines); + const contextLines = 2; + + const changedIndices = new Set(); + diff.forEach((entry, i) => { + if (entry.type !== "equal") { + for ( + let c = Math.max(0, i - contextLines); + c <= Math.min(diff.length - 1, i + contextLines); + c++ + ) { + changedIndices.add(c); + } + } + }); + + const output: string[] = []; + let prevWasGap = false; + let addedLines = 0; + let removedLines = 0; + + diff.forEach((entry, i) => { + if (!changedIndices.has(i)) { + if (!prevWasGap) { + output.push("..."); + prevWasGap = true; + } + return; + } + prevWasGap = false; + + if (entry.type === "add") { + output.push(`+ ${entry.line}`); + addedLines++; + } else if (entry.type === "remove") { + output.push(`- ${entry.line}`); + removedLines++; + } else { + output.push(` ${entry.line}`); + } + }); + + const header = `[delta snapshot: +${addedLines} lines added, -${removedLines} lines removed]\n\n`; + const diffStr = header + output.join("\n"); + + const savedChars = after.length - diffStr.length; + const savingsRatio = savedChars / after.length; + + if (savingsRatio < DELTA_SNAPSHOT_MIN_SAVINGS_RATIO) { + return { diff: after, isFull: true, savedChars: 0 }; + } + + return { diff: diffStr, isFull: false, savedChars }; +} + +type DiffEntry = { type: "equal" | "add" | "remove"; line: string }; + +function computeLineDiff(before: string[], after: string[]): DiffEntry[] { + const result: DiffEntry[] = []; + let i = 0, + j = 0; + + const afterIndex = new Map(); + after.forEach((line, idx) => { + if (!afterIndex.has(line)) afterIndex.set(line, []); + afterIndex.get(line)!.push(idx); + }); + + while (i < before.length && j < after.length) { + if (before[i] === after[j]) { + result.push({ type: "equal", line: before[i] }); + i++; + j++; + } else { + const nextInAfter = (afterIndex.get(before[i]) ?? []).find((idx) => idx >= j); + const nextInBefore = before.findIndex((l, idx) => idx > i && l === after[j]); + + const distToAfterMatch = nextInAfter !== undefined ? nextInAfter - j : Infinity; + const distToBeforeMatch = nextInBefore !== -1 ? nextInBefore - i : Infinity; + + if (distToAfterMatch <= distToBeforeMatch && nextInAfter !== undefined) { + while (j < nextInAfter) { + result.push({ type: "add", line: after[j] }); + j++; + } + } else if (nextInBefore !== -1) { + while (i < nextInBefore) { + result.push({ type: "remove", line: before[i] }); + i++; + } + } else { + result.push({ type: "remove", line: before[i] }); + result.push({ type: "add", line: after[j] }); + i++; + j++; + } + } + } + + while (i < before.length) { + result.push({ type: "remove", line: before[i++] }); + } + while (j < after.length) { + result.push({ type: "add", line: after[j++] }); + } + + return result; +}