chore: upstream the frame tree snapshot (#35917)

Author: pavelfeldman

packages/injected/src/ariaSnapshot.ts

@@ -26,7 +26,7 @@ import type { Box } from './domUtils';
 export type AriaNode = AriaProps & {
   role: AriaRole | 'fragment' | 'iframe';
   name: string;
-  ref?: number;
+  ref?: string;
   children: (AriaNode | string)[];
   element: Element;
   box: Box;
@@ -36,23 +36,23 @@ export type AriaNode = AriaProps & {
 
 export type AriaSnapshot = {
   root: AriaNode;
-  elements: Map<number, Element>;
+  elements: Map<string, Element>;
 };
 
 type AriaRef = {
   role: string;
   name: string;
-  ref: number;
+  ref: string;
 };
 
 let lastRef = 0;
 
-export function generateAriaTree(rootElement: Element, options?: { forAI?: boolean }): AriaSnapshot {
+export function generateAriaTree(rootElement: Element, options?: { forAI?: boolean, refPrefix?: string }): AriaSnapshot {
   const visited = new Set<Node>();
 
   const snapshot: AriaSnapshot = {
     root: { role: 'fragment', name: '', children: [], element: rootElement, props: {}, box: box(rootElement), receivesPointerEvents: true },
-    elements: new Map<number, Element>(),
+    elements: new Map<string, Element>(),
   };
 
   const visit = (ariaNode: AriaNode, node: Node) => {
@@ -149,20 +149,20 @@ export function generateAriaTree(rootElement: Element, options?: { forAI?: boole
   return snapshot;
 }
 
-function ariaRef(element: Element, role: string, name: string, options?: { forAI?: boolean }): number | undefined {
+function ariaRef(element: Element, role: string, name: string, options?: { forAI?: boolean, refPrefix?: string }): string | undefined {
   if (!options?.forAI)
     return undefined;
 
   let ariaRef: AriaRef | undefined;
   ariaRef = (element as any)._ariaRef;
   if (!ariaRef || ariaRef.role !== role || ariaRef.name !== name) {
-    ariaRef = { role, name, ref: ++lastRef };
+    ariaRef = { role, name, ref: (options?.refPrefix ?? '') + 'e' + (++lastRef) };
     (element as any)._ariaRef = ariaRef;
   }
   return ariaRef.ref;
 }
 
-function toAriaNode(element: Element, options?: { forAI?: boolean }): AriaNode | null {
+function toAriaNode(element: Element, options?: { forAI?: boolean, refPrefix?: string }): AriaNode | null {
   if (element.nodeName === 'IFRAME') {
     return {
       role: 'iframe',
@@ -443,7 +443,7 @@ export function renderAriaTree(ariaSnapshot: AriaSnapshot, options?: { mode?: 'r
       const ref = ariaNode.ref;
       const cursor = hasPointerCursor(ariaNode) ? ' [cursor=pointer]' : '';
       if (ref)
-        key += ` [ref=e${ref}]${cursor}`;
+        key += ` [ref=${ref}]${cursor}`;
     }
 
     const escapedKey = indent + '- ' + yamlEscapeKeyIfNeeded(key);

packages/injected/src/injectedScript.ts

@@ -297,7 +297,7 @@ export class InjectedScript {
     return new Set<Element>(result.map(r => r.element));
   }
 
-  ariaSnapshot(node: Node, options?: { mode?: 'raw' | 'regex', forAI?: boolean }): string {
+  ariaSnapshot(node: Node, options?: { mode?: 'raw' | 'regex', forAI?: boolean, refPrefix?: string }): 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);
@@ -675,10 +675,7 @@ export class InjectedScript {
 
   _createAriaRefEngine() {
     const queryAll = (root: SelectorRoot, selector: string): Element[] => {
-      if (!selector.startsWith('e'))
-        throw this.createStacklessError(`Invalid aria-ref selector "${selector}"`);
-      const ref = +selector.substring(1);
-      const result = this._lastAriaSnapshot?.elements?.get(ref);
+      const result = this._lastAriaSnapshot?.elements?.get(selector);
       return result && result.isConnected ? [result] : [];
     };
     return { queryAll };

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

@@ -303,8 +303,8 @@ export class Locator implements api.Locator {
     return await this._withElement((h, timeout) => h.screenshot({ ...options, mask, timeout }), options.timeout);
   }
 
-  async ariaSnapshot(options?: { _forAI?: boolean } & TimeoutOptions): Promise<string> {
-    const result = await this._frame._channel.ariaSnapshot({ ...options, forAI: options?._forAI, selector: this._selector });
+  async ariaSnapshot(options?: TimeoutOptions): Promise<string> {
+    const result = await this._frame._channel.ariaSnapshot({ ...options, selector: this._selector });
     return result.snapshot;
   }
 

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

@@ -832,6 +832,11 @@ export class Page extends ChannelOwner<channels.PageChannel> implements api.Page
     }
     return result.pdf;
   }
+
+  async _snapshotForAI(): Promise<string> {
+    const result = await this._channel.snapshotForAI();
+    return result.snapshot;
+  }
 }
 
 export class BindingCall extends ChannelOwner<channels.BindingCallChannel> {

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

@@ -90,6 +90,7 @@ export const commandsWithTracingSnapshots = new Set([
   'Page.mouseWheel',
   'Page.touchscreenTap',
   'Page.accessibilitySnapshot',
+  'Page.snapshotForAI',
   'Frame.evalOnSelector',
   'Frame.evalOnSelectorAll',
   'Frame.addScriptTag',

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

@@ -1386,6 +1386,10 @@ scheme.PagePdfParams = tObject({
 scheme.PagePdfResult = tObject({
   pdf: tBinary,
 });
+scheme.PageSnapshotForAIParams = tOptional(tObject({}));
+scheme.PageSnapshotForAIResult = tObject({
+  snapshot: tString,
+});
 scheme.PageStartJSCoverageParams = tObject({
   resetOnNavigation: tOptional(tBoolean),
   reportAnonymousScripts: tOptional(tBoolean),

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

@@ -295,6 +295,10 @@ export class PageDispatcher extends Dispatcher<Page, channels.PageChannel, Brows
     return { pdf: buffer };
   }
 
+  async snapshotForAI(params: channels.PageSnapshotForAIParams, metadata: CallMetadata): Promise<channels.PageSnapshotForAIResult> {
+    return { snapshot: await this._page.snapshotForAI(metadata) };
+  }
+
   async bringToFront(params: channels.PageBringToFrontParams, metadata: CallMetadata): Promise<void> {
     await this._page.bringToFront();
   }

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

@@ -811,7 +811,7 @@ export class ElementHandle<T extends Node = Node> extends js.JSHandle<T> {
     return this._page.delegate.getBoundingBox(this);
   }
 
-  async ariaSnapshot(options: { forAI?: boolean }): Promise<string> {
+  async ariaSnapshot(options?: { forAI?: boolean, refPrefix?: string }): Promise<string> {
     return await this.evaluateInUtility(([injected, element, options]) => injected.ariaSnapshot(element, options), options);
   }
 

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

@@ -111,7 +111,31 @@ export class FrameSelectors {
     return Promise.all(result);
   }
 
+  private _resolveForLastAISnapshot(selector: string, options: types.StrictOptions = {}): SelectorInFrame | null {
+    const match = selector.match(/^aria-ref=f(\d+)e\d+$/);
+    if (!match)
+      return null;
+
+    const frameIndex = +match[1];
+    const page = this.frame._page;
+    const frameId = page.lastSnapshotFrameIds[frameIndex - 1];
+    if (!frameId)
+      return null;
+    const frameManager = page.frameManager;
+    const frame = frameManager.frame(frameId);
+    if (!frame)
+      return null;
+    return {
+      frame,
+      info: frame.selectors._parseSelector(selector, options),
+    };
+  }
+
   async resolveFrameForSelector(selector: string, options: types.StrictOptions = {}, scope?: ElementHandle): Promise<SelectorInFrame | null> {
+    const resolvedForSnapshot = this._resolveForLastAISnapshot(selector, options);
+    if (resolvedForSnapshot)
+      return resolvedForSnapshot;
+
     let frame: Frame = this.frame;
     const frameChunks = splitSelectorByFrame(selector);
 

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

@@ -175,6 +175,7 @@ export class Page extends SdkObject {
   // When throttling for tracing, 200ms between frames, except for 10 frames around the action.
   private _frameThrottler = new FrameThrottler(10, 35, 200);
   closeReason: string | undefined;
+  lastSnapshotFrameIds: string[] = [];
 
   constructor(delegate: PageDelegate, browserContext: BrowserContext) {
     super(browserContext, 'page');
@@ -807,6 +808,13 @@ export class Page extends SdkObject {
   markAsServerSideOnly() {
     this._isServerSideOnly = true;
   }
+
+  async snapshotForAI(metadata: CallMetadata): Promise<string> {
+    const frameIds: string[] = [];
+    const snapshot = await snapshotFrameForAI(this.mainFrame(), 0, frameIds);
+    this.lastSnapshotFrameIds = frameIds;
+    return snapshot.join('\n');
+  }
 }
 
 export class Worker extends SdkObject {
@@ -988,3 +996,36 @@ class FrameThrottler {
     }
   }
 }
+
+async function snapshotFrameForAI(frame: frames.Frame, frameOrdinal: number, frameIds: string[]): Promise<string[]> {
+  const context = await frame._utilityContext();
+  const injectedScript = await context.injectedScript();
+  const snapshot = await injectedScript.evaluate((injected, refPrefix) => {
+    return injected.ariaSnapshot(injected.document.body, { forAI: true, refPrefix });
+  }, frameOrdinal ? 'f' + frameOrdinal : '');
+
+  const lines = snapshot.split('\n');
+  const result = [];
+  for (const line of lines) {
+    const match = line.match(/^(\s*)- iframe \[ref=(.*)\]/);
+    if (!match) {
+      result.push(line);
+      continue;
+    }
+
+    const leadingSpace = match[1];
+    const ref = match[2];
+    const frameSelector = `aria-ref=${ref} >> internal:control=enter-frame`;
+    const frameBodySelector = `${frameSelector} >> body`;
+    const child = await frame.selectors.resolveFrameForSelector(frameBodySelector, { strict: true });
+    if (!child) {
+      result.push(line);
+      continue;
+    }
+    const frameOrdinal = frameIds.length + 1;
+    frameIds.push(child.frame._id);
+    const childSnapshot = await snapshotFrameForAI(child.frame, frameOrdinal, frameIds);
+    result.push(line + ':', ...childSnapshot.map(l => leadingSpace + '  ' + l));
+  }
+  return result;
+}