feat: incremental aria snapshot

Author: dgozman

packages/injected/src/ariaSnapshot.ts

@@ -495,21 +495,31 @@ function matchesNodeDeep(root: AriaNode, template: AriaTemplateNode, collectAll:
   return results;
 }
 
-export function renderAriaTree(ariaSnapshot: AriaSnapshot, publicOptions: AriaTreeOptions): string {
+export function renderAriaTree(ariaSnapshot: AriaSnapshot, publicOptions: AriaTreeOptions, previous?: AriaSnapshot): string {
   const options = toInternalOptions(publicOptions);
   const lines: string[] = [];
   const includeText = options.renderStringsAsRegex ? textContributesInfo : () => true;
   const renderString = options.renderStringsAsRegex ? convertToBestGuessRegex : (str: string) => str;
-  const visit = (ariaNode: AriaNode | string, parentAriaNode: AriaNode | null, indent: string, renderCursorPointer: boolean) => {
-    if (typeof ariaNode === 'string') {
-      if (parentAriaNode && !includeText(parentAriaNode, ariaNode))
-        return;
-      const text = yamlEscapeValueIfNeeded(renderString(ariaNode));
-      if (text)
-        lines.push(indent + '- text: ' + text);
-      return;
+
+  const previousByRef = new Map<string, AriaNode>();
+  const visitPrevious = (ariaNode: AriaNode) => {
+    if (ariaNode.ref)
+      previousByRef.set(ariaNode.ref, ariaNode);
+    for (const child of ariaNode.children) {
+      if (typeof child !== 'string')
+        visitPrevious(child);
     }
+  };
+  if (previous)
+    visitPrevious(previous.root);
+
+  const visitText = (text: string, indent: string) => {
+    const escaped = yamlEscapeValueIfNeeded(renderString(text));
+    if (escaped)
+      lines.push(indent + '- text: ' + escaped);
+  };
 
+  const createKey = (ariaNode: AriaNode, renderCursorPointer: boolean): string => {
     let key = ariaNode.role;
     // Yaml has a limit of 1024 characters per key, and we leave some space for role and attributes.
     if (ariaNode.name && ariaNode.name.length <= 900) {
@@ -538,41 +548,93 @@ export function renderAriaTree(ariaSnapshot: AriaSnapshot, publicOptions: AriaTr
     if (ariaNode.selected === true)
       key += ` [selected]`;
 
-    let inCursorPointer = false;
     if (ariaNode.ref) {
       key += ` [ref=${ariaNode.ref}]`;
-      if (renderCursorPointer && hasPointerCursor(ariaNode)) {
-        inCursorPointer = true;
+      if (renderCursorPointer && hasPointerCursor(ariaNode))
         key += ' [cursor=pointer]';
-      }
     }
+    return key;
+  };
+
+  const arePropsEqual = (a: AriaNode, b: AriaNode): boolean => {
+    const aKeys = Object.keys(a.props);
+    const bKeys = Object.keys(b.props);
+    return aKeys.length === bKeys.length && aKeys.every(k => a.props[k] === b.props[k]);
+  };
+
+  const visit = (ariaNode: AriaNode, indent: string, renderCursorPointer: boolean, previousNode: AriaNode | undefined): { unchanged: boolean } => {
+    if (ariaNode.ref)
+      previousNode = previousByRef.get(ariaNode.ref);
+    const linesBefore = lines.length;
+
+    const key = createKey(ariaNode, renderCursorPointer);
+    let unchanged = !!previousNode && key === createKey(previousNode, renderCursorPointer) && arePropsEqual(ariaNode, previousNode);
 
     const escapedKey = indent + '- ' + yamlEscapeKeyIfNeeded(key);
     const hasProps = !!Object.keys(ariaNode.props).length;
+    const inCursorPointer = renderCursorPointer && !!ariaNode.ref && hasPointerCursor(ariaNode);
+
     if (!ariaNode.children.length && !hasProps) {
+      // Leaf node without children.
       lines.push(escapedKey);
     } else if (ariaNode.children.length === 1 && typeof ariaNode.children[0] === 'string' && !hasProps) {
-      const text = includeText(ariaNode, ariaNode.children[0]) ? renderString(ariaNode.children[0] as string) : null;
+      // Leaf node with only some text inside.
+      const shouldInclude = includeText(ariaNode, ariaNode.children[0]);
+      const text = shouldInclude ? renderString(ariaNode.children[0]) : null;
       if (text)
         lines.push(escapedKey + ': ' + yamlEscapeValueIfNeeded(text));
       else
         lines.push(escapedKey);
+
+      // Node is unchanged only when previous node also had the same single text child.
+      unchanged = unchanged && !!previousNode &&
+          previousNode.children.length === 1 && typeof previousNode.children[0] === 'string' &&
+          !Object.keys(previousNode.props).length && ariaNode.children[0] === previousNode.children[0];
     } else {
+      // Node with (optional) props and some children.
       lines.push(escapedKey + ':');
       for (const [name, value] of Object.entries(ariaNode.props))
         lines.push(indent + '  - /' + name + ': ' + yamlEscapeValueIfNeeded(value));
-      for (const child of ariaNode.children || [])
-        visit(child, ariaNode, indent + '  ', renderCursorPointer && !inCursorPointer);
+
+      // All children must be the same.
+      unchanged = unchanged && previousNode?.children.length === ariaNode.children.length;
+
+      const childIndent = indent + '  ';
+      for (let childIndex = 0 ; childIndex < ariaNode.children.length; childIndex++) {
+        const child = ariaNode.children[childIndex];
+        if (typeof child === 'string') {
+          const shouldInclude = includeText(ariaNode, child);
+          if (shouldInclude)
+            visitText(child, childIndent);
+          unchanged = unchanged && previousNode?.children[childIndex] === child && shouldInclude === includeText(previousNode, child);
+        } else {
+          const previousChild = previousNode?.children[childIndex];
+          const childResult = visit(child, childIndent, renderCursorPointer && !inCursorPointer, typeof previousChild !== 'string' ? previousChild : undefined);
+          unchanged = unchanged && childResult.unchanged;
+        }
+      }
     }
+
+    if (unchanged && ariaNode.ref) {
+      // Replace the whole subtree with a single reference.
+      lines.splice(linesBefore);
+      lines.push(indent + `- ref=${ariaNode.ref} [unchanged]`);
+    }
+
+    return { unchanged };
   };
 
   const ariaNode = ariaSnapshot.root;
   if (ariaNode.role === 'fragment') {
     // Render fragment.
-    for (const child of ariaNode.children || [])
-      visit(child, ariaNode, '', !!options.renderCursorPointer);
+    for (const child of ariaNode.children || []) {
+      if (typeof child === 'string')
+        visitText(child, '');
+      else
+        visit(child, '', !!options.renderCursorPointer, undefined);
+    }
   } else {
-    visit(ariaNode, null, '', !!options.renderCursorPointer);
+    visit(ariaNode, '', !!options.renderCursorPointer, undefined);
   }
   return lines.join('\n');
 }

packages/injected/src/injectedScript.ts

@@ -92,7 +92,8 @@ export class InjectedScript {
   readonly window: Window & typeof globalThis;
   readonly document: Document;
   readonly consoleApi: ConsoleAPI;
-  private _lastAriaSnapshot: AriaSnapshot | undefined;
+  private _lastAriaSnapshotForTrack = new Map<string, AriaSnapshot>();
+  private _lastAriaSnapshotForQuery: AriaSnapshot | undefined;
 
   // Recorder must use any external dependencies through InjectedScript.
   // Otherwise it will end up with a copy of all modules it uses, and any
@@ -299,11 +300,18 @@ export class InjectedScript {
     return new Set<Element>(result.map(r => r.element));
   }
 
-  ariaSnapshot(node: Node, options: AriaTreeOptions): string {
+  ariaSnapshot(node: Node, options: AriaTreeOptions & { track?: string, incremental?: boolean }): string {
     if (node.nodeType !== Node.ELEMENT_NODE)
       throw this.createStacklessError('Can only capture aria snapshot of Element nodes.');
-    this._lastAriaSnapshot = generateAriaTree(node as Element, options);
-    return renderAriaTree(this._lastAriaSnapshot, options);
+    const ariaSnapshot = generateAriaTree(node as Element, options);
+    let previous: AriaSnapshot | undefined;
+    if (options.incremental)
+      previous = options.track ? this._lastAriaSnapshotForTrack.get(options.track) : this._lastAriaSnapshotForQuery;
+    const result = renderAriaTree(ariaSnapshot, options, previous);
+    if (options.track)
+      this._lastAriaSnapshotForTrack.set(options.track, ariaSnapshot);
+    this._lastAriaSnapshotForQuery = ariaSnapshot;
+    return result;
   }
 
   ariaSnapshotForRecorder(): { ariaSnapshot: string, refs: Map<Element, string> } {
@@ -692,7 +700,7 @@ export class InjectedScript {
 
   _createAriaRefEngine() {
     const queryAll = (root: SelectorRoot, selector: string): Element[] => {
-      const result = this._lastAriaSnapshot?.elements?.get(selector);
+      const result = this._lastAriaSnapshotForQuery?.elements?.get(selector);
       return result && result.isConnected ? [result] : [];
     };
     return { queryAll };

packages/playwright-core/src/client/page.ts

@@ -847,8 +847,8 @@ export class Page extends ChannelOwner<channels.PageChannel> implements api.Page
     return result.pdf;
   }
 
-  async _snapshotForAI(options: TimeoutOptions = {}): Promise<string> {
-    const result = await this._channel.snapshotForAI({ timeout: this._timeoutSettings.timeout(options) });
+  async _snapshotForAI(options: TimeoutOptions & { track?: string, mode?: 'full' | 'incremental' } = {}): Promise<string> {
+    const result = await this._channel.snapshotForAI({ timeout: this._timeoutSettings.timeout(options), track: options.track, mode: options.mode });
     return result.snapshot;
   }
 }

packages/playwright-core/src/protocol/validator.ts

@@ -1465,6 +1465,8 @@ scheme.PageRequestsResult = tObject({
   requests: tArray(tChannel(['Request'])),
 });
 scheme.PageSnapshotForAIParams = tObject({
+  track: tOptional(tString),
+  mode: tOptional(tEnum(['full', 'incremental'])),
   timeout: tFloat,
 });
 scheme.PageSnapshotForAIResult = tObject({

packages/playwright-core/src/server/dispatchers/pageDispatcher.ts

@@ -352,7 +352,7 @@ export class PageDispatcher extends Dispatcher<Page, channels.PageChannel, Brows
   }
 
   async snapshotForAI(params: channels.PageSnapshotForAIParams, progress: Progress): Promise<channels.PageSnapshotForAIResult> {
-    return { snapshot: await this._page.snapshotForAI(progress) };
+    return { snapshot: await this._page.snapshotForAI(progress, params) };
   }
 
   async bringToFront(params: channels.PageBringToFrontParams, progress: Progress): Promise<void> {

packages/playwright-core/src/server/page.ts

@@ -859,9 +859,9 @@ export class Page extends SdkObject {
     await Promise.all(this.frames().map(frame => frame.hideHighlight().catch(() => {})));
   }
 
-  async snapshotForAI(progress: Progress): Promise<string> {
+  async snapshotForAI(progress: Progress, options: { track?: string, mode?: 'full' | 'incremental' }): Promise<string> {
     this.lastSnapshotFrameIds = [];
-    const snapshot = await snapshotFrameForAI(progress, this.mainFrame(), 0, this.lastSnapshotFrameIds);
+    const snapshot = await snapshotFrameForAI(progress, this.mainFrame(), 0, this.lastSnapshotFrameIds, options);
     return snapshot.join('\n');
   }
 }
@@ -1037,18 +1037,18 @@ class FrameThrottler {
   }
 }
 
-async function snapshotFrameForAI(progress: Progress, frame: frames.Frame, frameOrdinal: number, frameIds: string[]): Promise<string[]> {
+async function snapshotFrameForAI(progress: Progress, frame: frames.Frame, frameOrdinal: number, frameIds: string[], options: { track?: string, mode?: 'full' | 'incremental' }): Promise<string[]> {
   // Only await the topmost navigations, inner frames will be empty when racing.
   const snapshot = await frame.retryWithProgressAndTimeouts(progress, [1000, 2000, 4000, 8000], async continuePolling => {
     try {
       const context = await progress.race(frame._utilityContext());
       const injectedScript = await progress.race(context.injectedScript());
-      const snapshotOrRetry = await progress.race(injectedScript.evaluate((injected, refPrefix) => {
+      const snapshotOrRetry = await progress.race(injectedScript.evaluate((injected, options) => {
         const node = injected.document.body;
         if (!node)
           return true;
-        return injected.ariaSnapshot(node, { mode: 'ai', refPrefix });
-      }, frameOrdinal ? 'f' + frameOrdinal : ''));
+        return injected.ariaSnapshot(node, { mode: 'ai', ...options });
+      }, { refPrefix: frameOrdinal ? 'f' + frameOrdinal : '', incremental: options.mode === 'incremental', track: options.track }));
       if (snapshotOrRetry === true)
         return continuePolling;
       return snapshotOrRetry;
@@ -1080,7 +1080,7 @@ async function snapshotFrameForAI(progress: Progress, frame: frames.Frame, frame
     const frameOrdinal = frameIds.length + 1;
     frameIds.push(child.frame._id);
     try {
-      const childSnapshot = await snapshotFrameForAI(progress, child.frame, frameOrdinal, frameIds);
+      const childSnapshot = await snapshotFrameForAI(progress, child.frame, frameOrdinal, frameIds, options);
       result.push(line + ':', ...childSnapshot.map(l => leadingSpace + '  ' + l));
     } catch {
       result.push(line);

packages/playwright/src/mcp/browser/response.ts

@@ -28,7 +28,7 @@ export class Response {
   private _code: string[] = [];
   private _images: { contentType: string, data: Buffer }[] = [];
   private _context: Context;
-  private _includeSnapshot = false;
+  private _includeSnapshot: 'none' | 'full' | 'incremental' = 'none';
   private _includeTabs = false;
   private _tabSnapshot: TabSnapshot | undefined;
 
@@ -75,8 +75,8 @@ export class Response {
     return this._images;
   }
 
-  setIncludeSnapshot() {
-    this._includeSnapshot = true;
+  setIncludeSnapshot(full?: 'full') {
+    this._includeSnapshot = full ?? 'incremental';
   }
 
   setIncludeTabs() {
@@ -86,8 +86,8 @@ export class Response {
   async finish() {
     // All the async snapshotting post-action is happening here.
     // Everything below should race against modal states.
-    if (this._includeSnapshot && this._context.currentTab())
-      this._tabSnapshot = await this._context.currentTabOrDie().captureSnapshot();
+    if (this._includeSnapshot !== 'none' && this._context.currentTab())
+      this._tabSnapshot = await this._context.currentTabOrDie().captureSnapshot(this._includeSnapshot);
     for (const tab of this._context.tabs())
       await tab.updateTitle();
   }
@@ -126,7 +126,7 @@ ${this._code.join('\n')}
     }
 
     // List browser tabs.
-    if (this._includeSnapshot || this._includeTabs)
+    if (this._includeSnapshot !== 'none' || this._includeTabs)
       response.push(...renderTabsMarkdown(this._context.tabs(), this._includeTabs));
 
     // Add snapshot if provided.

packages/playwright/src/mcp/browser/tab.ts

@@ -62,6 +62,7 @@ export class Tab extends EventEmitter<TabEventsInterface> {
     this.context = context;
     this.page = page as Page;
     this._onPageClose = onPageClose;
+    page.on('console', console.error);
     page.on('console', event => this._handleConsoleMessage(messageToConsoleMessage(event)));
     page.on('pageerror', error => this._handleConsoleMessage(pageErrorToConsoleMessage(error)));
     page.on('request', request => this._requests.add(request));
@@ -217,10 +218,10 @@ export class Tab extends EventEmitter<TabEventsInterface> {
     return this._requests;
   }
 
-  async captureSnapshot(): Promise<TabSnapshot> {
+  async captureSnapshot(mode: 'full' | 'incremental'): Promise<TabSnapshot> {
     let tabSnapshot: TabSnapshot | undefined;
     const modalStates = await this._raceAgainstModalStates(async () => {
-      const snapshot = await this.page._snapshotForAI();
+      const snapshot = await this.page._snapshotForAI({ mode, track: 'response' });
       tabSnapshot = {
         url: this.page.url(),
         title: await this.page.title(),

packages/playwright/src/mcp/browser/tools/snapshot.ts

@@ -30,7 +30,7 @@ const snapshot = defineTool({
 
   handle: async (context, params, response) => {
     await context.ensureTab();
-    response.setIncludeSnapshot();
+    response.setIncludeSnapshot('full');
   },
 });
 

packages/protocol/src/channels.d.ts

@@ -2548,10 +2548,13 @@ export type PageRequestsResult = {
   requests: RequestChannel[],
 };
 export type PageSnapshotForAIParams = {
+  track?: string,
+  mode?: 'full' | 'incremental',
   timeout: number,
 };
 export type PageSnapshotForAIOptions = {
-
+  track?: string,
+  mode?: 'full' | 'incremental',
 };
 export type PageSnapshotForAIResult = {
   snapshot: string,