Skip to content

Commit a6b4b22

Browse files
pavelfeldmanpavelfeldman
authored
chore: upstream the frame tree snapshot (#35917)
1 parent 4166fd2 commit a6b4b22

File tree

13 files changed

+134
-56
lines changed

13 files changed

+134
-56
lines changed

‎packages/injected/src/ariaSnapshot.ts

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ import type { Box } from './domUtils';
2626
export type AriaNode = AriaProps & {
2727
role: AriaRole | 'fragment' | 'iframe';
2828
name: string;
29-
ref?: number;
29+
ref?: string;
3030
children: (AriaNode | string)[];
3131
element: Element;
3232
box: Box;
@@ -36,23 +36,23 @@ export type AriaNode = AriaProps & {
3636

3737
export type AriaSnapshot = {
3838
root: AriaNode;
39-
elements: Map<number, Element>;
39+
elements: Map<string, Element>;
4040
};
4141

4242
type AriaRef = {
4343
role: string;
4444
name: string;
45-
ref: number;
45+
ref: string;
4646
};
4747

4848
let lastRef = 0;
4949

50-
export function generateAriaTree(rootElement: Element, options?: { forAI?: boolean }): AriaSnapshot {
50+
export function generateAriaTree(rootElement: Element, options?: { forAI?: boolean, refPrefix?: string }): AriaSnapshot {
5151
const visited = new Set<Node>();
5252

5353
const snapshot: AriaSnapshot = {
5454
root: { role: 'fragment', name: '', children: [], element: rootElement, props: {}, box: box(rootElement), receivesPointerEvents: true },
55-
elements: new Map<number, Element>(),
55+
elements: new Map<string, Element>(),
5656
};
5757

5858
const visit = (ariaNode: AriaNode, node: Node) => {
@@ -149,20 +149,20 @@ export function generateAriaTree(rootElement: Element, options?: { forAI?: boole
149149
return snapshot;
150150
}
151151

152-
function ariaRef(element: Element, role: string, name: string, options?: { forAI?: boolean }): number | undefined {
152+
function ariaRef(element: Element, role: string, name: string, options?: { forAI?: boolean, refPrefix?: string }): string | undefined {
153153
if (!options?.forAI)
154154
return undefined;
155155

156156
let ariaRef: AriaRef | undefined;
157157
ariaRef = (element as any)._ariaRef;
158158
if (!ariaRef || ariaRef.role !== role || ariaRef.name !== name) {
159-
ariaRef = { role, name, ref: ++lastRef };
159+
ariaRef = { role, name, ref: (options?.refPrefix ?? '') + 'e' + (++lastRef) };
160160
(element as any)._ariaRef = ariaRef;
161161
}
162162
return ariaRef.ref;
163163
}
164164

165-
function toAriaNode(element: Element, options?: { forAI?: boolean }): AriaNode | null {
165+
function toAriaNode(element: Element, options?: { forAI?: boolean, refPrefix?: string }): AriaNode | null {
166166
if (element.nodeName === 'IFRAME') {
167167
return {
168168
role: 'iframe',
@@ -443,7 +443,7 @@ export function renderAriaTree(ariaSnapshot: AriaSnapshot, options?: { mode?: 'r
443443
const ref = ariaNode.ref;
444444
const cursor = hasPointerCursor(ariaNode) ? ' [cursor=pointer]' : '';
445445
if (ref)
446-
key += ` [ref=e${ref}]${cursor}`;
446+
key += ` [ref=${ref}]${cursor}`;
447447
}
448448

449449
const escapedKey = indent + '- ' + yamlEscapeKeyIfNeeded(key);

‎packages/injected/src/injectedScript.ts

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -297,7 +297,7 @@ export class InjectedScript {
297297
return new Set<Element>(result.map(r => r.element));
298298
}
299299

300-
ariaSnapshot(node: Node, options?: { mode?: 'raw' | 'regex', forAI?: boolean }): string {
300+
ariaSnapshot(node: Node, options?: { mode?: 'raw' | 'regex', forAI?: boolean, refPrefix?: string }): string {
301301
if (node.nodeType !== Node.ELEMENT_NODE)
302302
throw this.createStacklessError('Can only capture aria snapshot of Element nodes.');
303303
this._lastAriaSnapshot = generateAriaTree(node as Element, options);
@@ -675,10 +675,7 @@ export class InjectedScript {
675675

676676
_createAriaRefEngine() {
677677
const queryAll = (root: SelectorRoot, selector: string): Element[] => {
678-
if (!selector.startsWith('e'))
679-
throw this.createStacklessError(`Invalid aria-ref selector "${selector}"`);
680-
const ref = +selector.substring(1);
681-
const result = this._lastAriaSnapshot?.elements?.get(ref);
678+
const result = this._lastAriaSnapshot?.elements?.get(selector);
682679
return result && result.isConnected ? [result] : [];
683680
};
684681
return { queryAll };

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -303,8 +303,8 @@ export class Locator implements api.Locator {
303303
return await this._withElement((h, timeout) => h.screenshot({ ...options, mask, timeout }), options.timeout);
304304
}
305305

306-
async ariaSnapshot(options?: { _forAI?: boolean } & TimeoutOptions): Promise<string> {
307-
const result = await this._frame._channel.ariaSnapshot({ ...options, forAI: options?._forAI, selector: this._selector });
306+
async ariaSnapshot(options?: TimeoutOptions): Promise<string> {
307+
const result = await this._frame._channel.ariaSnapshot({ ...options, selector: this._selector });
308308
return result.snapshot;
309309
}
310310

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -832,6 +832,11 @@ export class Page extends ChannelOwner<channels.PageChannel> implements api.Page
832832
}
833833
return result.pdf;
834834
}
835+
836+
async _snapshotForAI(): Promise<string> {
837+
const result = await this._channel.snapshotForAI();
838+
return result.snapshot;
839+
}
835840
}
836841

837842
export class BindingCall extends ChannelOwner<channels.BindingCallChannel> {

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,7 @@ export const commandsWithTracingSnapshots = new Set([
9090
'Page.mouseWheel',
9191
'Page.touchscreenTap',
9292
'Page.accessibilitySnapshot',
93+
'Page.snapshotForAI',
9394
'Frame.evalOnSelector',
9495
'Frame.evalOnSelectorAll',
9596
'Frame.addScriptTag',

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1386,6 +1386,10 @@ scheme.PagePdfParams = tObject({
13861386
scheme.PagePdfResult = tObject({
13871387
pdf: tBinary,
13881388
});
1389+
scheme.PageSnapshotForAIParams = tOptional(tObject({}));
1390+
scheme.PageSnapshotForAIResult = tObject({
1391+
snapshot: tString,
1392+
});
13891393
scheme.PageStartJSCoverageParams = tObject({
13901394
resetOnNavigation: tOptional(tBoolean),
13911395
reportAnonymousScripts: tOptional(tBoolean),

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -295,6 +295,10 @@ export class PageDispatcher extends Dispatcher<Page, channels.PageChannel, Brows
295295
return { pdf: buffer };
296296
}
297297

298+
async snapshotForAI(params: channels.PageSnapshotForAIParams, metadata: CallMetadata): Promise<channels.PageSnapshotForAIResult> {
299+
return { snapshot: await this._page.snapshotForAI(metadata) };
300+
}
301+
298302
async bringToFront(params: channels.PageBringToFrontParams, metadata: CallMetadata): Promise<void> {
299303
await this._page.bringToFront();
300304
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -811,7 +811,7 @@ export class ElementHandle<T extends Node = Node> extends js.JSHandle<T> {
811811
return this._page.delegate.getBoundingBox(this);
812812
}
813813

814-
async ariaSnapshot(options: { forAI?: boolean }): Promise<string> {
814+
async ariaSnapshot(options?: { forAI?: boolean, refPrefix?: string }): Promise<string> {
815815
return await this.evaluateInUtility(([injected, element, options]) => injected.ariaSnapshot(element, options), options);
816816
}
817817

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

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -111,7 +111,31 @@ export class FrameSelectors {
111111
return Promise.all(result);
112112
}
113113

114+
private _resolveForLastAISnapshot(selector: string, options: types.StrictOptions = {}): SelectorInFrame | null {
115+
const match = selector.match(/^aria-ref=f(\d+)e\d+$/);
116+
if (!match)
117+
return null;
118+
119+
const frameIndex = +match[1];
120+
const page = this.frame._page;
121+
const frameId = page.lastSnapshotFrameIds[frameIndex - 1];
122+
if (!frameId)
123+
return null;
124+
const frameManager = page.frameManager;
125+
const frame = frameManager.frame(frameId);
126+
if (!frame)
127+
return null;
128+
return {
129+
frame,
130+
info: frame.selectors._parseSelector(selector, options),
131+
};
132+
}
133+
114134
async resolveFrameForSelector(selector: string, options: types.StrictOptions = {}, scope?: ElementHandle): Promise<SelectorInFrame | null> {
135+
const resolvedForSnapshot = this._resolveForLastAISnapshot(selector, options);
136+
if (resolvedForSnapshot)
137+
return resolvedForSnapshot;
138+
115139
let frame: Frame = this.frame;
116140
const frameChunks = splitSelectorByFrame(selector);
117141

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

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -175,6 +175,7 @@ export class Page extends SdkObject {
175175
// When throttling for tracing, 200ms between frames, except for 10 frames around the action.
176176
private _frameThrottler = new FrameThrottler(10, 35, 200);
177177
closeReason: string | undefined;
178+
lastSnapshotFrameIds: string[] = [];
178179

179180
constructor(delegate: PageDelegate, browserContext: BrowserContext) {
180181
super(browserContext, 'page');
@@ -807,6 +808,13 @@ export class Page extends SdkObject {
807808
markAsServerSideOnly() {
808809
this._isServerSideOnly = true;
809810
}
811+
812+
async snapshotForAI(metadata: CallMetadata): Promise<string> {
813+
const frameIds: string[] = [];
814+
const snapshot = await snapshotFrameForAI(this.mainFrame(), 0, frameIds);
815+
this.lastSnapshotFrameIds = frameIds;
816+
return snapshot.join('\n');
817+
}
810818
}
811819

812820
export class Worker extends SdkObject {
@@ -988,3 +996,36 @@ class FrameThrottler {
988996
}
989997
}
990998
}
999+
1000+
async function snapshotFrameForAI(frame: frames.Frame, frameOrdinal: number, frameIds: string[]): Promise<string[]> {
1001+
const context = await frame._utilityContext();
1002+
const injectedScript = await context.injectedScript();
1003+
const snapshot = await injectedScript.evaluate((injected, refPrefix) => {
1004+
return injected.ariaSnapshot(injected.document.body, { forAI: true, refPrefix });
1005+
}, frameOrdinal ? 'f' + frameOrdinal : '');
1006+
1007+
const lines = snapshot.split('\n');
1008+
const result = [];
1009+
for (const line of lines) {
1010+
const match = line.match(/^(\s*)- iframe \[ref=(.*)\]/);
1011+
if (!match) {
1012+
result.push(line);
1013+
continue;
1014+
}
1015+
1016+
const leadingSpace = match[1];
1017+
const ref = match[2];
1018+
const frameSelector = `aria-ref=${ref} >> internal:control=enter-frame`;
1019+
const frameBodySelector = `${frameSelector} >> body`;
1020+
const child = await frame.selectors.resolveFrameForSelector(frameBodySelector, { strict: true });
1021+
if (!child) {
1022+
result.push(line);
1023+
continue;
1024+
}
1025+
const frameOrdinal = frameIds.length + 1;
1026+
frameIds.push(child.frame._id);
1027+
const childSnapshot = await snapshotFrameForAI(child.frame, frameOrdinal, frameIds);
1028+
result.push(line + ':', ...childSnapshot.map(l => leadingSpace + ' ' + l));
1029+
}
1030+
return result;
1031+
}

‎packages/protocol/src/channels.d.ts

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2057,6 +2057,7 @@ export interface PageChannel extends PageEventTarget, EventTargetChannel {
20572057
touchscreenTap(params: PageTouchscreenTapParams, metadata?: CallMetadata): Promise<PageTouchscreenTapResult>;
20582058
accessibilitySnapshot(params: PageAccessibilitySnapshotParams, metadata?: CallMetadata): Promise<PageAccessibilitySnapshotResult>;
20592059
pdf(params: PagePdfParams, metadata?: CallMetadata): Promise<PagePdfResult>;
2060+
snapshotForAI(params?: PageSnapshotForAIParams, metadata?: CallMetadata): Promise<PageSnapshotForAIResult>;
20602061
startJSCoverage(params: PageStartJSCoverageParams, metadata?: CallMetadata): Promise<PageStartJSCoverageResult>;
20612062
stopJSCoverage(params?: PageStopJSCoverageParams, metadata?: CallMetadata): Promise<PageStopJSCoverageResult>;
20622063
startCSSCoverage(params: PageStartCSSCoverageParams, metadata?: CallMetadata): Promise<PageStartCSSCoverageResult>;
@@ -2497,6 +2498,11 @@ export type PagePdfOptions = {
24972498
export type PagePdfResult = {
24982499
pdf: Binary,
24992500
};
2501+
export type PageSnapshotForAIParams = {};
2502+
export type PageSnapshotForAIOptions = {};
2503+
export type PageSnapshotForAIResult = {
2504+
snapshot: string,
2505+
};
25002506
export type PageStartJSCoverageParams = {
25012507
resetOnNavigation?: boolean,
25022508
reportAnonymousScripts?: boolean,

‎packages/protocol/src/protocol.yml

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1795,6 +1795,12 @@ Page:
17951795
returns:
17961796
pdf: binary
17971797

1798+
snapshotForAI:
1799+
returns:
1800+
snapshot: string
1801+
flags:
1802+
snapshot: true
1803+
17981804
startJSCoverage:
17991805
parameters:
18001806
resetOnNavigation: boolean?

0 commit comments

Comments
 (0)