feat: incremental aria snapshot (#37795)

Author: dgozman

packages/injected/src/ariaSnapshot.ts

@@ -495,21 +495,36 @@ function matchesNodeDeep(root: AriaNode, template: AriaTemplateNode, collectAll:
   return results;
 }
 
-export function renderAriaTree(ariaSnapshot: AriaSnapshot, publicOptions: AriaTreeOptions): string {
+function buildByRefMap(root: AriaNode | undefined, map: Map<string, AriaNode> = new Map()): Map<string, AriaNode> {
+  if (root?.ref)
+    map.set(root.ref, root);
+  for (const child of root?.children || []) {
+    if (typeof child !== 'string')
+      buildByRefMap(child, map);
+  }
+  return map;
+}
+
+function 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]);
+}
+
+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 = buildByRefMap(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 +553,84 @@ 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 getSingleInlinedTextChild = (ariaNode: AriaNode | undefined): string | undefined => {
+    return ariaNode?.children.length === 1 && typeof ariaNode.children[0] === 'string' && !Object.keys(ariaNode.props).length ? ariaNode.children[0] : undefined;
+  };
 
+  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);
     const escapedKey = indent + '- ' + yamlEscapeKeyIfNeeded(key);
-    const hasProps = !!Object.keys(ariaNode.props).length;
-    if (!ariaNode.children.length && !hasProps) {
+    const inCursorPointer = renderCursorPointer && !!ariaNode.ref && hasPointerCursor(ariaNode);
+    const singleInlinedTextChild = getSingleInlinedTextChild(ariaNode);
+
+    // Whether ariaNode's subtree is the same as previousNode's, and can be replaced with just a ref.
+    let unchanged = !!previousNode && key === createKey(previousNode, renderCursorPointer) && arePropsEqual(ariaNode, previousNode);
+
+    if (!ariaNode.children.length && !Object.keys(ariaNode.props).length) {
+      // 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;
-      if (text)
-        lines.push(escapedKey + ': ' + yamlEscapeValueIfNeeded(text));
+    } else if (singleInlinedTextChild !== undefined) {
+      // Leaf node with just some text inside.
+      // Unchanged when the previous node also had the same single text child.
+      unchanged = unchanged && getSingleInlinedTextChild(previousNode) === singleInlinedTextChild;
+
+      const shouldInclude = includeText(ariaNode, singleInlinedTextChild);
+      if (shouldInclude)
+        lines.push(escapedKey + ': ' + yamlEscapeValueIfNeeded(renderString(singleInlinedTextChild)));
       else
         lines.push(escapedKey);
     } 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') {
+          unchanged = unchanged && previousNode?.children[childIndex] === child;
+          if (includeText(ariaNode, child))
+            visitText(child, childIndent);
+        } 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);
-  } else {
-    visit(ariaNode, null, '', !!options.renderCursorPointer);
+  // Do not render the root fragment, just its children.
+  const nodesToRender = ariaSnapshot.root.role === 'fragment' ? ariaSnapshot.root.children : [ariaSnapshot.root];
+  for (const nodeToRender of nodesToRender) {
+    if (typeof nodeToRender === 'string')
+      visitText(nodeToRender, '');
+    else
+      visit(nodeToRender, '', !!options.renderCursorPointer, undefined);
   }
   return lines.join('\n');
 }
@@ -636,5 +694,5 @@ function textContributesInfo(node: AriaNode, text: string): boolean {
 }
 
 function hasPointerCursor(ariaNode: AriaNode): boolean {
-  return ariaNode.box.style?.cursor === 'pointer';
+  return ariaNode.box.cursor === 'pointer';
 }

packages/injected/src/domUtils.ts

@@ -112,28 +112,32 @@ export type Box = {
   visible: boolean;
   inline: boolean;
   rect?: DOMRect;
-  style?: CSSStyleDeclaration;
+  // Note: we do not store the CSSStyleDeclaration object, because it is a live object
+  // and changes values over time. This does not work for caching or comparing to the
+  // old values. Instead, store all the properties separately.
+  cursor?: CSSStyleDeclaration['cursor'];
 };
 
 export function computeBox(element: Element): Box {
   // Note: this logic should be similar to waitForDisplayedAtStablePosition() to avoid surprises.
   const style = getElementComputedStyle(element);
   if (!style)
     return { visible: true, inline: false };
+  const cursor = style.cursor;
   if (style.display === 'contents') {
     // display:contents is not rendered itself, but its child nodes are.
     for (let child = element.firstChild; child; child = child.nextSibling) {
       if (child.nodeType === 1 /* Node.ELEMENT_NODE */ && isElementVisible(child as Element))
-        return { visible: true, inline: false, style };
+        return { visible: true, inline: false, cursor };
       if (child.nodeType === 3 /* Node.TEXT_NODE */ && isVisibleTextNode(child as Text))
-        return { visible: true, inline: true, style };
+        return { visible: true, inline: true, cursor };
     }
-    return { visible: false, inline: false, style };
+    return { visible: false, inline: false, cursor };
   }
   if (!isElementStyleVisibilityVisible(element, style))
-    return { style, visible: false, inline: false };
+    return { cursor, visible: false, inline: false };
   const rect = element.getBoundingClientRect();
-  return { rect, style, visible: rect.width > 0 && rect.height > 0, inline: style.display === 'inline' };
+  return { rect, cursor, visible: rect.width > 0 && rect.height > 0, inline: style.display === 'inline' };
 }
 
 export function isElementVisible(element: Element): boolean {

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

@@ -217,10 +217,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');
   },
 });