chore: experiment with stable aria refs (#35900)

packages/injected/src/ariaSnapshot.ts

@@ -26,6 +26,7 @@ import type { Box } from './domUtils';
 export type AriaNode = AriaProps & {
   role: AriaRole | 'fragment' | 'iframe';
   name: string;
+  ref?: number;
   children: (AriaNode | string)[];
   element: Element;
   box: Box;
@@ -36,28 +37,24 @@ export type AriaNode = AriaProps & {
 export type AriaSnapshot = {
   root: AriaNode;
   elements: Map<number, Element>;
-  generation: number;
-  ids: Map<Element, number>;
 };
 
-export function generateAriaTree(rootElement: Element, generation: number, options?: { forAI?: boolean }): AriaSnapshot {
+type AriaRef = {
+  role: string;
+  name: string;
+  ref: number;
+};
+
+let lastRef = 0;
+
+export function generateAriaTree(rootElement: Element, options?: { forAI?: boolean }): 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>(),
-    generation,
-    ids: new Map<Element, number>(),
   };
 
-  const addElement = (element: Element) => {
-    const id = snapshot.elements.size + 1;
-    snapshot.elements.set(id, element);
-    snapshot.ids.set(element, id);
-  };
-
-  addElement(rootElement);
-
   const visit = (ariaNode: AriaNode, node: Node) => {
     if (visited.has(node))
       return;
@@ -91,10 +88,12 @@ export function generateAriaTree(rootElement: Element, generation: number, optio
       }
     }
 
-    addElement(element);
     const childAriaNode = toAriaNode(element, options);
-    if (childAriaNode)
+    if (childAriaNode) {
+      if (childAriaNode.ref)
+        snapshot.elements.set(childAriaNode.ref, element);
       ariaNode.children.push(childAriaNode);
+    }
     processElement(childAriaNode || ariaNode, element, ariaChildren);
   };
 
@@ -150,9 +149,32 @@ export function generateAriaTree(rootElement: Element, generation: number, optio
   return snapshot;
 }
 
+function ariaRef(element: Element, role: string, name: string, options?: { forAI?: boolean }): number | 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 };
+    (element as any)._ariaRef = ariaRef;
+  }
+  return ariaRef.ref;
+}
+
 function toAriaNode(element: Element, options?: { forAI?: boolean }): AriaNode | null {
-  if (element.nodeName === 'IFRAME')
-    return { role: 'iframe', name: '', children: [], props: {}, element, box: box(element), receivesPointerEvents: true };
+  if (element.nodeName === 'IFRAME') {
+    return {
+      role: 'iframe',
+      name: '',
+      ref: ariaRef(element, 'iframe', '', options),
+      children: [],
+      props: {},
+      element,
+      box: box(element),
+      receivesPointerEvents: true
+    };
+  }
 
   const defaultRole = options?.forAI ? 'generic' : null;
   const role = roleUtils.getAriaRole(element) ?? defaultRole;
@@ -161,7 +183,17 @@ function toAriaNode(element: Element, options?: { forAI?: boolean }): AriaNode |
 
   const name = normalizeWhiteSpace(roleUtils.getElementAccessibleName(element, false) || '');
   const receivesPointerEvents = roleUtils.receivesPointerEvents(element);
-  const result: AriaNode = { role, name, children: [], props: {}, element, box: box(element), receivesPointerEvents };
+
+  const result: AriaNode = {
+    role,
+    name,
+    ref: ariaRef(element, role, name, options),
+    children: [],
+    props: {},
+    element,
+    box: box(element),
+    receivesPointerEvents
+  };
 
   if (roleUtils.kAriaCheckedRoles.includes(role))
     result.checked = roleUtils.getAriaChecked(element);
@@ -266,7 +298,7 @@ export type MatcherReceived = {
 };
 
 export function matchesAriaTree(rootElement: Element, template: AriaTemplateNode): { matches: AriaNode[], received: MatcherReceived } {
-  const snapshot = generateAriaTree(rootElement, 0);
+  const snapshot = generateAriaTree(rootElement);
   const matches = matchesNodeDeep(snapshot.root, template, false, false);
   return {
     matches,
@@ -278,7 +310,7 @@ export function matchesAriaTree(rootElement: Element, template: AriaTemplateNode
 }
 
 export function getAllByAria(rootElement: Element, template: AriaTemplateNode): Element[] {
-  const root = generateAriaTree(rootElement, 0).root;
+  const root = generateAriaTree(rootElement).root;
   const matches = matchesNodeDeep(root, template, true, false);
   return matches.map(n => n.element);
 }
@@ -408,10 +440,10 @@ export function renderAriaTree(ariaSnapshot: AriaSnapshot, options?: { mode?: 'r
     if (ariaNode.selected === true)
       key += ` [selected]`;
     if (options?.forAI && receivesPointerEvents(ariaNode)) {
-      const id = ariaSnapshot.ids.get(ariaNode.element);
+      const ref = ariaNode.ref;
       const cursor = hasPointerCursor(ariaNode) ? ' [cursor=pointer]' : '';
-      if (id)
-        key += ` [ref=s${ariaSnapshot.generation}e${id}]${cursor}`;
+      if (ref)
+        key += ` [ref=e${ref}]${cursor}`;
     }
 
     const escapedKey = indent + '- ' + yamlEscapeKeyIfNeeded(key);

packages/injected/src/injectedScript.ts

@@ -227,7 +227,7 @@ export class InjectedScript {
     this._engines.set('internal:attr', this._createNamedAttributeEngine());
     this._engines.set('internal:testid', this._createNamedAttributeEngine());
     this._engines.set('internal:role', createRoleEngine(true));
-    this._engines.set('aria-ref', this._createAriaIdEngine());
+    this._engines.set('aria-ref', this._createAriaRefEngine());
 
     for (const { name, source } of options.customEngines)
       this._engines.set(name, this.eval(source));
@@ -300,15 +300,10 @@ export class InjectedScript {
   ariaSnapshot(node: Node, options?: { mode?: 'raw' | 'regex', forAI?: boolean }): string {
     if (node.nodeType !== Node.ELEMENT_NODE)
       throw this.createStacklessError('Can only capture aria snapshot of Element nodes.');
-    const generation = (this._lastAriaSnapshot?.generation || 0) + 1;
-    this._lastAriaSnapshot = generateAriaTree(node as Element, generation, options);
+    this._lastAriaSnapshot = generateAriaTree(node as Element, options);
     return renderAriaTree(this._lastAriaSnapshot, options);
   }
 
-  ariaSnapshotElement(snapshot: AriaSnapshot, elementId: number): Element | null {
-    return snapshot.elements.get(elementId) || null;
-  }
-
   getAllByAria(document: Document, template: AriaTemplateNode): Element[] {
     return getAllByAria(document.documentElement, template);
   }
@@ -678,15 +673,12 @@ export class InjectedScript {
     return result;
   }
 
-  _createAriaIdEngine() {
+  _createAriaRefEngine() {
     const queryAll = (root: SelectorRoot, selector: string): Element[] => {
-      const match = selector.match(/^s(\d+)e(\d+)$/);
-      if (!match)
-        throw this.createStacklessError('Invalid aria-ref selector, should be of form s<number>e<number>');
-      const [, generation, elementId] = match;
-      if (this._lastAriaSnapshot?.generation !== +generation)
-        throw this.createStacklessError(`Stale aria-ref, expected s${this._lastAriaSnapshot?.generation}e{number}, got ${selector}`);
-      const result = this._lastAriaSnapshot?.elements?.get(+elementId);
+      if (!selector.startsWith('e'))
+        throw this.createStacklessError(`Invalid aria-ref selector "${selector}"`);
+      const ref = +selector.substring(1);
+      const result = this._lastAriaSnapshot?.elements?.get(ref);
       return result && result.isConnected ? [result] : [];
     };
     return { queryAll };

tests/page/page-aria-snapshot-ai.spec.ts

@@ -27,20 +27,27 @@ it('should generate refs', async ({ page }) => {
   `);
 
   const snapshot1 = await page.locator('body').ariaSnapshot(forAI);
-  expect(snapshot1).toContain('- button "One" [ref=s1e3]');
-  expect(snapshot1).toContain('- button "Two" [ref=s1e4]');
-  expect(snapshot1).toContain('- button "Three" [ref=s1e5]');
+  expect(snapshot1).toContainYaml(`
+    - generic [ref=e1]:
+      - button "One" [ref=e2]
+      - button "Two" [ref=e3]
+      - button "Three" [ref=e4]
+  `);
+  await expect(page.locator('aria-ref=e2')).toHaveText('One');
+  await expect(page.locator('aria-ref=e3')).toHaveText('Two');
+  await expect(page.locator('aria-ref=e4')).toHaveText('Three');
 
-  await expect(page.locator('aria-ref=s1e3')).toHaveText('One');
-  await expect(page.locator('aria-ref=s1e4')).toHaveText('Two');
-  await expect(page.locator('aria-ref=s1e5')).toHaveText('Three');
+  await page.locator('aria-ref=e3').evaluate((e: HTMLElement) => {
+    e.textContent = 'Not Two';
+  });
 
   const snapshot2 = await page.locator('body').ariaSnapshot(forAI);
-  expect(snapshot2).toContain('- button "One" [ref=s2e3]');
-  await expect(page.locator('aria-ref=s2e3')).toHaveText('One');
-
-  const e = await expect(page.locator('aria-ref=s1e3')).toHaveText('One').catch(e => e);
-  expect(e.message).toContain('Error: Stale aria-ref, expected s2e{number}, got s1e3');
+  expect(snapshot2).toContainYaml(`
+    - generic [ref=e1]:
+      - button "One" [ref=e2]
+      - button "Not Two" [ref=e5]
+      - button "Three" [ref=e4]
+  `);
 });
 
 it('should list iframes', async ({ page }) => {
@@ -80,15 +87,15 @@ it('ref mode can be used to stitch all frame snapshots', async ({ page, server }
   }
 
   expect(await allFrameSnapshot(page)).toContainYaml(`
-    - generic [ref=s1e2]:
-      - iframe [ref=s1e3]:
-        - generic [ref=s1e2]:
-          - iframe [ref=s1e3]:
-            - generic [ref=s1e3]: Hi, I'm frame
-          - iframe [ref=s1e4]:
-            - generic [ref=s1e3]: Hi, I'm frame
-      - iframe [ref=s1e4]:
-        - generic [ref=s1e3]: Hi, I'm frame
+    - generic [ref=e1]:
+      - iframe [ref=e2]:
+        - generic [ref=e1]:
+          - iframe [ref=e2]:
+            - generic [ref=e2]: Hi, I'm frame
+          - iframe [ref=e3]:
+            - generic [ref=e2]: Hi, I'm frame
+      - iframe [ref=e3]:
+        - generic [ref=e2]: Hi, I'm frame
   `);
 });
 
@@ -101,10 +108,10 @@ it('should not generate refs for hidden elements', async ({ page }) => {
 
   const snapshot = await page.locator('body').ariaSnapshot(forAI);
   expect(snapshot).toContainYaml(`
-    - generic [ref=s1e2]:
-      - button "One" [ref=s1e3]
+    - generic [ref=e1]:
+      - button "One" [ref=e2]
       - button "Two"
-      - button "Three" [ref=s1e5]
+      - button "Three" [ref=e4]
   `);
 });
 
@@ -133,12 +140,12 @@ it('should not generate refs for elements with pointer-events:none', async ({ pa
 
   const snapshot = await page.locator('body').ariaSnapshot(forAI);
   expect(snapshot).toContainYaml(`
-    - generic [ref=s1e2]:
+    - generic [ref=e1]:
       - button "no-ref"
-      - button "with-ref" [ref=s1e5]
-      - button "with-ref" [ref=s1e8]
-      - button "with-ref" [ref=s1e11]
-      - generic [ref=s1e12]:
+      - button "with-ref" [ref=e4]
+      - button "with-ref" [ref=e7]
+      - button "with-ref" [ref=e10]
+      - generic [ref=e11]:
         - generic:
           - button "no-ref"
   `);
@@ -178,19 +185,19 @@ it('emit generic roles for nodes w/o roles', async ({ page }) => {
   const snapshot = await page.locator('body').ariaSnapshot(forAI);
 
   expect(snapshot).toContainYaml(`
-    - generic [ref=s1e3]:
-      - generic [ref=s1e4]:
-        - generic [ref=s1e5]:
+    - generic [ref=e2]:
+      - generic [ref=e3]:
+        - generic [ref=e4]:
           - radio "Apple" [checked]
-        - generic [ref=s1e7]: Apple
-      - generic [ref=s1e8]:
-        - generic [ref=s1e9]:
+        - generic [ref=e6]: Apple
+      - generic [ref=e7]:
+        - generic [ref=e8]:
           - radio "Pear"
-        - generic [ref=s1e11]: Pear
-      - generic [ref=s1e12]:
-        - generic [ref=s1e13]:
+        - generic [ref=e10]: Pear
+      - generic [ref=e11]:
+        - generic [ref=e12]:
           - radio "Orange"
-        - generic [ref=s1e15]: Orange
+        - generic [ref=e14]: Orange
   `);
 });
 
@@ -207,7 +214,7 @@ it('should collapse generic nodes', async ({ page }) => {
 
   const snapshot = await page.locator('body').ariaSnapshot(forAI);
   expect(snapshot).toContainYaml(`
-    - button \"Button\" [ref=s1e6]
+    - button \"Button\" [ref=e5]
   `);
 });
 
@@ -218,6 +225,6 @@ it('should include cursor pointer hint', async ({ page }) => {
 
   const snapshot = await page.locator('body').ariaSnapshot(forAI);
   expect(snapshot).toContainYaml(`
-    - button \"Button\" [ref=s1e3] [cursor=pointer]
+    - button \"Button\" [ref=e2] [cursor=pointer]
   `);
 });