feat(extension): support OOPIF refs with unified frame geometry

apps/extension/src/tools/vom/frame-document.ts

@@ -1,3 +1,9 @@
+import {
+  type CdpFrame,
+  type CdpFrameGraph,
+  type CdpTarget,
+  cdpTargetKey,
+} from "@/browser-driver/frame-graph";
 import type { CapturedNode, CapturedViewModel } from "./capture";
 
 export interface FrameOwnedAxNode {
@@ -9,14 +15,11 @@ export interface FrameOwnedAxNode {
 }
 
 export interface FrameAxBatch<T extends FrameOwnedAxNode> {
-  frameId: string;
+  frame: CdpFrame;
   nodes: T[];
 }
 
-export interface FrameDocument<T extends FrameOwnedAxNode> {
-  frameId: string;
-  parentFrameId?: string;
-  ownerBackendNodeId?: number;
+export interface FrameDocument<T extends FrameOwnedAxNode> extends CdpFrame {
   contextScopeId: string;
   axNodes: T[];
   domNodes: CapturedNode[];
@@ -29,30 +32,55 @@ interface Ownership {
 
 interface OwnedCandidate<T> {
   node: T;
+  target: CdpTarget;
   ownership: Ownership;
 }
 
-function frameIds<T extends FrameOwnedAxNode>(
+function targetNodeKey(target: CdpTarget, nodeId: string): string {
+  return `${cdpTargetKey(target)}:${nodeId}`;
+}
+
+function targetBackendKey(target: CdpTarget, backendNodeId: number): string {
+  return `${cdpTargetKey(target)}:${backendNodeId}`;
+}
+
+function frameList<T extends FrameOwnedAxNode>(
+  graph: CdpFrameGraph | null,
   batches: FrameAxBatch<T>[],
   captured: CapturedViewModel,
-): string[] {
-  const ids = new Set<string>();
-  if (captured.rootFrameId) ids.add(captured.rootFrameId);
-  for (const frameId of captured.frameNodes?.keys() ?? []) ids.add(frameId);
-  for (const batch of batches) ids.add(batch.frameId);
-  return [...ids];
+): CdpFrame[] {
+  const frames = new Map<string, CdpFrame>();
+  for (const frame of graph?.frames ?? []) frames.set(frame.frameId, frame);
+  for (const batch of batches) {
+    if (!frames.has(batch.frame.frameId)) frames.set(batch.frame.frameId, batch.frame);
+  }
+  for (const frameId of captured.frameNodes?.keys() ?? []) {
+    if (!frames.has(frameId))
+      frames.set(frameId, { frameId, target: { tabId: batches[0]?.frame.target.tabId ?? 0 } });
+  }
+  return [...frames.values()].map((frame) => {
+    const ownerBackendNodeId =
+      frame.ownerBackendNodeId ?? captured.frameOwnerBackendNodeIds?.get(frame.frameId);
+    const parentFrameId = frame.parentFrameId ?? captured.frameParentIds?.get(frame.frameId);
+    return {
+      ...frame,
+      ...(ownerBackendNodeId !== undefined ? { ownerBackendNodeId } : {}),
+      ...(parentFrameId ? { parentFrameId } : {}),
+    };
+  });
 }
 
 export function buildFrameDocuments<T extends FrameOwnedAxNode>(
+  graph: CdpFrameGraph | null,
   batches: FrameAxBatch<T>[],
   captured: CapturedViewModel,
 ): FrameDocument<T>[] {
-  const frames = frameIds(batches, captured);
-  const knownFrames = new Set(frames);
-  const backendOwner = new Map<number, string>();
-  for (const frameId of frames) {
-    for (const node of captured.frameNodes?.get(frameId) ?? []) {
-      backendOwner.set(node.backendNodeId, frameId);
+  const frames = frameList(graph, batches, captured);
+  const frameById = new Map(frames.map((frame) => [frame.frameId, frame]));
+  const backendOwner = new Map<string, string>();
+  for (const frame of frames) {
+    for (const node of captured.frameNodes?.get(frame.frameId) ?? []) {
+      backendOwner.set(targetBackendKey(frame.target, node.backendNodeId), frame.frameId);
     }
   }
 
@@ -64,15 +92,17 @@ export function buildFrameDocuments<T extends FrameOwnedAxNode>(
     const resolveOwnership = (node: T): Ownership => {
       const cached = ownershipByNodeId.get(node.nodeId);
       if (cached) return cached;
-      if (node.frameId && knownFrames.has(node.frameId)) {
+      if (node.frameId && frameById.has(node.frameId)) {
         const ownership = { frameId: node.frameId, strength: 4 };
         ownershipByNodeId.set(node.nodeId, ownership);
         return ownership;
       }
       if (typeof node.backendDOMNodeId === "number") {
-        const owner = backendOwner.get(node.backendDOMNodeId);
-        if (owner) {
-          const ownership = { frameId: owner, strength: 3 };
+        const frameId = backendOwner.get(
+          targetBackendKey(batch.frame.target, node.backendDOMNodeId),
+        );
+        if (frameId) {
+          const ownership = { frameId, strength: 3 };
           ownershipByNodeId.set(node.nodeId, ownership);
           return ownership;
         }
@@ -91,28 +121,32 @@ export function buildFrameDocuments<T extends FrameOwnedAxNode>(
           return ownership;
         }
       }
-      const ownership = { frameId: batch.frameId, strength: 1 };
+      const ownership = { frameId: batch.frame.frameId, strength: 1 };
       ownershipByNodeId.set(node.nodeId, ownership);
       return ownership;
     };
 
     for (const node of batch.nodes) {
       const ownership = resolveOwnership(node);
-      const existing = candidates.get(node.nodeId);
+      const key = targetNodeKey(batch.frame.target, node.nodeId);
+      const existing = candidates.get(key);
       if (!existing || ownership.strength > existing.ownership.strength) {
-        candidates.set(node.nodeId, { node, ownership });
+        candidates.set(key, { node, target: batch.frame.target, ownership });
       }
     }
   }
 
-  const ownershipByNodeId = new Map(
-    [...candidates].map(([nodeId, candidate]) => [nodeId, candidate.ownership.frameId]),
+  const ownershipByTargetNode = new Map(
+    [...candidates].map(([key, candidate]) => [key, candidate.ownership.frameId]),
   );
   const axNodesByFrame = new Map<string, T[]>();
-  for (const { node, ownership } of candidates.values()) {
-    const parentFrameId = node.parentId ? ownershipByNodeId.get(node.parentId) : undefined;
+  for (const candidate of candidates.values()) {
+    const { node, target, ownership } = candidate;
+    const parentFrameId = node.parentId
+      ? ownershipByTargetNode.get(targetNodeKey(target, node.parentId))
+      : undefined;
     const childIds = node.childIds?.filter(
-      (childId) => ownershipByNodeId.get(childId) === ownership.frameId,
+      (childId) => ownershipByTargetNode.get(targetNodeKey(target, childId)) === ownership.frameId,
     );
     const ownedNode = {
       ...node,
@@ -127,16 +161,10 @@ export function buildFrameDocuments<T extends FrameOwnedAxNode>(
     axNodesByFrame.set(ownership.frameId, nodes);
   }
 
-  return frames.map((frameId) => ({
-    frameId,
-    ...(captured.frameParentIds?.get(frameId)
-      ? { parentFrameId: captured.frameParentIds.get(frameId) }
-      : {}),
-    ...(captured.frameOwnerBackendNodeIds?.get(frameId) !== undefined
-      ? { ownerBackendNodeId: captured.frameOwnerBackendNodeIds.get(frameId) }
-      : {}),
-    contextScopeId: frameId,
-    axNodes: axNodesByFrame.get(frameId) ?? [],
-    domNodes: captured.frameNodes?.get(frameId) ?? [],
+  return frames.map((frame) => ({
+    ...frame,
+    contextScopeId: frame.frameId,
+    axNodes: axNodesByFrame.get(frame.frameId) ?? [],
+    domNodes: captured.frameNodes?.get(frame.frameId) ?? [],
   }));
 }