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

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

@@ -1,5 +1,8 @@
-import type { CdpRunner } from "../shared";
-import type { CapturedViewModel } from "./capture";
+import type { CdpFrame, CdpFrameGraph } from "@/browser-driver/frame-graph";
+import { type FrameProjection, resolveFrameProjection } from "../frame-geometry";
+import { projectRectToViewport } from "../geometry";
+import { type CdpRunner, cdpRunnerForTarget, sendToCdpTarget } from "../shared";
+import { type CapturedNode, type CapturedViewModel, captureViewModel } from "./capture";
 import {
   buildFrameDocuments,
   type FrameAxBatch,
@@ -18,39 +21,133 @@ function isAbortError(err: unknown): boolean {
   return err instanceof Error && err.name === "AbortError";
 }
 
-function capturedFrameIds(captured: CapturedViewModel): string[] {
-  const frameIds = new Set<string>();
-  if (captured.rootFrameId) frameIds.add(captured.rootFrameId);
-  for (const frameId of captured.frameNodes?.keys() ?? []) frameIds.add(frameId);
-  if (frameIds.size === 0) frameIds.add("root");
-  return [...frameIds];
+async function discoverFrameGraph(cdp: CdpRunner, tabId: number): Promise<CdpFrameGraph | null> {
+  if (!cdp.getFrameGraph) return null;
+  try {
+    return await cdp.getFrameGraph(tabId);
+  } catch (err) {
+    console.debug("[bsk observation] frame graph capture failed", err);
+    return null;
+  }
+}
+
+function transformFrameNodes(nodes: CapturedNode[], projection: FrameProjection): CapturedNode[] {
+  return nodes.map((node) => ({
+    ...node,
+    rect: (() => {
+      const bounds = projectRectToViewport(node.rect, projection);
+      return bounds ? { x: bounds.x, y: bounds.y, w: bounds.width, h: bounds.height } : null;
+    })(),
+  }));
+}
+
+function frameOwnerNode(captured: CapturedViewModel, frame: CdpFrame): CapturedNode | undefined {
+  if (frame.ownerBackendNodeId === undefined) return undefined;
+  const parentNodes = frame.parentFrameId
+    ? captured.frameNodes?.get(frame.parentFrameId)
+    : captured.nodes;
+  return parentNodes?.find((node) => node.backendNodeId === frame.ownerBackendNodeId);
+}
+
+async function captureMissingFrameDocuments(
+  cdp: CdpRunner,
+  tabId: number,
+  graph: CdpFrameGraph,
+  captured: CapturedViewModel,
+  signal?: AbortSignal,
+): Promise<void> {
+  if (!captured.frameNodes) captured.frameNodes = new Map();
+  if (!captured.frameOwnerBackendNodeIds) captured.frameOwnerBackendNodeIds = new Map();
+  if (!captured.frameParentIds) captured.frameParentIds = new Map();
+
+  for (const frame of graph.frames) {
+    if (!frame.target.sessionId || captured.frameNodes.has(frame.frameId)) continue;
+    const owner = frameOwnerNode(captured, frame);
+    if (!owner?.rect || frame.ownerBackendNodeId === undefined) continue;
+    try {
+      const child = await captureViewModel(cdpRunnerForTarget(cdp, frame.target), tabId, {
+        signal,
+      });
+      const projection = await resolveFrameProjection(cdp, graph, frame.frameId);
+      if (!projection) continue;
+      const rootFrameId = child.rootFrameId ?? frame.frameId;
+      const childFrames = child.frameNodes ?? new Map<string, CapturedNode[]>();
+      if (!childFrames.has(rootFrameId)) childFrames.set(rootFrameId, child.nodes);
+      if (rootFrameId !== frame.frameId && !childFrames.has(frame.frameId)) {
+        childFrames.set(
+          frame.frameId,
+          child.nodes.map((node) => ({ ...node, frameId: frame.frameId })),
+        );
+      }
+      for (const [childFrameId, nodes] of childFrames) {
+        captured.frameNodes.set(childFrameId, transformFrameNodes(nodes, projection));
+      }
+      captured.iframeNodes.set(
+        frame.ownerBackendNodeId,
+        captured.frameNodes.get(frame.frameId) ?? [],
+      );
+      captured.frameOwnerBackendNodeIds.set(frame.frameId, frame.ownerBackendNodeId);
+      if (frame.parentFrameId) captured.frameParentIds.set(frame.frameId, frame.parentFrameId);
+      for (const [childFrameId, ownerBackendNodeId] of child.frameOwnerBackendNodeIds ?? []) {
+        captured.frameOwnerBackendNodeIds.set(childFrameId, ownerBackendNodeId);
+      }
+      for (const [childFrameId, parentFrameId] of child.frameParentIds ?? []) {
+        captured.frameParentIds.set(childFrameId, parentFrameId);
+      }
+    } catch (err) {
+      if (isAbortError(err)) throw err;
+      console.debug("[bsk observation] child frame DOM capture failed", {
+        frameId: frame.frameId,
+        err,
+      });
+    }
+  }
 }
 
 async function captureAxTrees<T extends FrameAxNode>(
   cdp: CdpRunner,
   tabId: number,
+  discoveredFrames: CdpFrame[],
   captured: CapturedViewModel,
   signal?: AbortSignal,
 ): Promise<FrameAxBatch<T>[]> {
-  const frameIds = capturedFrameIds(captured);
+  let frames = [...discoveredFrames];
+  const capturedFrameIds = [...(captured.frameNodes?.keys() ?? [])];
+  if (frames.length === 0) {
+    const rootFrameId = captured.rootFrameId ?? capturedFrameIds[0] ?? "root";
+    frames = [{ frameId: rootFrameId, target: { tabId } }];
+  }
+  for (const frameId of capturedFrameIds) {
+    if (!frames.some((frame) => frame.frameId === frameId)) {
+      frames.push({ frameId, target: { tabId } });
+    }
+  }
+
   let firstFailure: unknown;
-  await cdp.send(tabId, "Accessibility.enable", {});
   const trees = await Promise.all(
-    frameIds.map(async (frameId): Promise<FrameAxBatch<T>> => {
+    frames.map(async (frame): Promise<FrameAxBatch<T>> => {
       try {
         throwIfAborted(signal);
-        const result = await cdp.send<{ nodes?: T[] }>(
-          tabId,
+        await sendToCdpTarget(cdp, frame.target, "Accessibility.enable", {});
+        const result = await sendToCdpTarget<{ nodes?: T[] }>(
+          cdp,
+          frame.target,
           "Accessibility.getFullAXTree",
-          frameId === "root" ? {} : { frameId },
+          frame.frameId === "root" ? {} : { frameId: frame.frameId },
         );
         throwIfAborted(signal);
-        return { frameId, nodes: result.nodes ?? [] };
+        return {
+          frame,
+          nodes: result.nodes ?? [],
+        };
       } catch (err) {
         if (isAbortError(err)) throw err;
         firstFailure ??= err;
-        console.debug("[bsk observation] frame AX capture failed", { frameId, err });
-        return { frameId, nodes: [] };
+        console.debug("[bsk observation] frame AX capture failed", {
+          frameId: frame.frameId,
+          err,
+        });
+        return { frame, nodes: [] };
       }
     }),
   );
@@ -70,6 +167,10 @@ export async function captureFrameData<T extends FrameAxNode>(
   captured: CapturedViewModel,
   signal?: AbortSignal,
 ): Promise<CapturedFrameDocument<T>[]> {
-  const batches = await captureAxTrees<T>(cdp, tabId, captured, signal);
-  return buildFrameDocuments(batches, captured);
+  const graph = await discoverFrameGraph(cdp, tabId);
+  const frames = graph?.frames ?? [];
+  if (graph) await captureMissingFrameDocuments(cdp, tabId, graph, captured, signal);
+  throwIfAborted(signal);
+  const batches = await captureAxTrees<T>(cdp, tabId, frames, captured, signal);
+  return buildFrameDocuments(graph, batches, captured);
 }