fix(vom): bug fix

apps/extension/src/tools/__tests__/observation.test.ts

@@ -36,6 +36,14 @@ function fakeAgentWindow(ids: number[]) {
 const TINY_PNG =
   "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGNgYGBgAAAABQABh6FO1AAAAABJRU5ErkJggg==";
 
+function pngWithIhdrDimensions(width: number, height: number): string {
+  const bytes = Uint8Array.from(atob(TINY_PNG), (character) => character.charCodeAt(0));
+  const view = new DataView(bytes.buffer);
+  view.setUint32(16, width);
+  view.setUint32(20, height);
+  return btoa(String.fromCharCode(...bytes));
+}
+
 function makeScreenshotDeps(
   opts: {
     cdp?: CdpRunner;
@@ -86,6 +94,7 @@ describe("stripDataUrlPrefix", () => {
 describe("parsePngDimensions", () => {
   it("parses width/height from the IHDR chunk", () => {
     expect(parsePngDimensions(TINY_PNG)).toEqual({ width: 1, height: 1 });
+    expect(parsePngDimensions(pngWithIhdrDimensions(100, 80))).toEqual({ width: 100, height: 80 });
   });
   it("returns null on non-PNG input", () => {
     expect(parsePngDimensions("not-a-png-payload-just-random-base64-text-zzzzzzzzz")).toBeNull();
@@ -120,6 +129,20 @@ describe("handleScreenshot", () => {
     expect(get).not.toHaveBeenCalled();
   });
 
+  it("fails closed when a full-tab capture is not a valid PNG", async () => {
+    const sm = new SessionManager({ agentWindow: fakeAgentWindow([100]) });
+    await sm.start("aa11");
+    const res = await handleScreenshot(
+      sm,
+      { session_id: "aa11" },
+      makeScreenshotDeps({
+        captureVisibleTab: vi.fn(async () => "data:image/png;base64,bm90LWEtcG5n"),
+      }),
+    );
+
+    expect(res).toMatchObject({ code: "cdp_failed", message: /invalid PNG data/ });
+  });
+
   it("returns not_found when Agent Window has no active tab", async () => {
     const sm = new SessionManager({ agentWindow: fakeAgentWindow([100]) });
     await sm.start("aa11");
@@ -349,7 +372,7 @@ describe("handleScreenshot", () => {
         cssLayoutViewport: { clientWidth: 800, clientHeight: 600 },
       }),
       "DOM.getContentQuads": () => ({ quads: [[0, 0, 200, 0, 200, 100, 0, 100]] }),
-      "Page.captureScreenshot": () => ({ data: TINY_PNG }),
+      "Page.captureScreenshot": () => ({ data: pngWithIhdrDimensions(100, 80) }),
     });
 
     const res = await handleScreenshot(
@@ -368,6 +391,95 @@ describe("handleScreenshot", () => {
       clip?: { x: number; y: number; width: number; height: number };
     };
     expect(clip.clip).toMatchObject({ x: 25, y: 30, width: 50, height: 40 });
+    expect(res).toMatchObject({ width: 100, height: 80 });
+  });
+
+  it("routes OOPIF surface geometry through its child session and captures top coordinates", async () => {
+    const sm = new SessionManager({ agentWindow: fakeAgentWindow([100]) });
+    const ctx = await sm.start("aa11");
+    ctx.refStore.set("e5", 999, {
+      tabId: 7,
+      frameId: "child-frame",
+      cdpSessionId: "child-session",
+      kind: "surface",
+      visibleRect: { x: 224, y: 346, w: 200, h: 80 },
+    });
+    const { cdp, sent } = makeFakeCdp({
+      "DOM.getBoxModel": () => ({
+        model: { content: [204, 306, 604, 306, 604, 506, 204, 506] },
+      }),
+      "Page.getLayoutMetrics": () => ({
+        cssLayoutViewport: { clientWidth: 1000, clientHeight: 800 },
+      }),
+      "Page.captureScreenshot": () => ({ data: TINY_PNG }),
+    });
+    cdp.getFrameGraph = vi.fn(async () => ({
+      rootFrameId: "main",
+      frames: [
+        { frameId: "main", target: { tabId: 7 } },
+        {
+          frameId: "child-frame",
+          parentFrameId: "main",
+          ownerBackendNodeId: 99,
+          target: { tabId: 7, sessionId: "child-session" },
+        },
+      ],
+    }));
+    const targetCalls: Array<{ sessionId?: string; method: string }> = [];
+    cdp.sendToTarget = vi.fn(async (target, method) => {
+      targetCalls.push({ sessionId: target.sessionId, method });
+      if (method === "DOM.getContentQuads") {
+        return { quads: [[10, 20, 110, 20, 110, 60, 10, 60]] };
+      }
+      if (method === "Page.getLayoutMetrics") {
+        return { cssLayoutViewport: { clientWidth: 200, clientHeight: 100 } };
+      }
+      throw new Error(`unexpected child CDP call ${method}`);
+    }) as CdpRunner["sendToTarget"];
+
+    const res = await handleScreenshot(
+      sm,
+      { session_id: "aa11", ref: "@e5", tab_id: 7 },
+      makeScreenshotDeps({
+        cdp,
+        get: vi.fn(async () => ({ id: 7, windowId: 100, active: false }) as chrome.tabs.Tab),
+        query: vi.fn(),
+        captureVisibleTab: vi.fn(),
+      }),
+    );
+
+    if ("code" in res) throw new Error(`unexpected error: ${JSON.stringify(res)}`);
+    expect(targetCalls).toEqual([
+      { sessionId: "child-session", method: "DOM.getContentQuads" },
+      { sessionId: "child-session", method: "Page.getLayoutMetrics" },
+    ]);
+    const clip = sent.find((call) => call.method === "Page.captureScreenshot")?.params as {
+      clip?: { x: number; y: number; width: number; height: number };
+    };
+    expect(clip.clip).toMatchObject({ x: 224, y: 346, width: 200, height: 80 });
+  });
+
+  it("fails closed when a ref capture is not a valid PNG", async () => {
+    const sm = new SessionManager({ agentWindow: fakeAgentWindow([100]) });
+    const ctx = await sm.start("aa11");
+    ctx.refStore.set("e5", 999, { tabId: 7, kind: "surface" });
+    const { cdp } = makeFakeCdp({
+      "DOM.getContentQuads": () => ({ quads: [[0, 0, 50, 0, 50, 50, 0, 50]] }),
+      "Page.captureScreenshot": () => ({ data: "bm90LWEtcG5n" }),
+    });
+
+    const res = await handleScreenshot(
+      sm,
+      { session_id: "aa11", ref: "@e5", tab_id: 7 },
+      makeScreenshotDeps({
+        cdp,
+        get: vi.fn(async () => ({ id: 7, windowId: 100, active: false }) as chrome.tabs.Tab),
+        query: vi.fn(),
+        captureVisibleTab: vi.fn(),
+      }),
+    );
+
+    expect(res).toMatchObject({ code: "cdp_failed", message: /invalid PNG data/ });
   });
 
   it("returns not_found for unknown ref", async () => {

apps/extension/src/tools/observation.ts

@@ -116,8 +116,8 @@ export function stripDataUrlPrefix(dataUrl: string): string {
 
 /**
  * Parse a PNG's IHDR chunk and return `(width, height)`. Returns
- * `null` on any malformed input so callers fall back to `0/0` instead
- * of throwing.
+ * `null` on malformed input so callers can reject dimensions that are
+ * not authoritative instead of guessing from a CSS-space rectangle.
  *
  * PNG layout: 8-byte signature, then a 4-byte length, 4-byte type
  * ("IHDR"), then the chunk data — width is bytes 16-19 BE, height is
@@ -129,9 +129,9 @@ export function parsePngDimensions(base64: string): { width: number; height: num
     const head = base64.length > 64 ? base64.slice(0, 64) : base64;
     const bin = atob(head);
     if (bin.length < 24) return null;
-    if (bin.charCodeAt(0) !== 0x89 || bin.charCodeAt(1) !== 0x50 || bin.charCodeAt(2) !== 0x4e) {
-      return null;
-    }
+    const expectedSignature = [0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a];
+    if (expectedSignature.some((byte, index) => bin.charCodeAt(index) !== byte)) return null;
+    if (bin.slice(12, 16) !== "IHDR") return null;
     const u32 = (off: number) =>
       (bin.charCodeAt(off) << 24) |
       (bin.charCodeAt(off + 1) << 16) |
@@ -250,10 +250,10 @@ async function captureElementScreenshot(
     if (!image_base64) {
       return { code: "cdp_failed", message: "Page.captureScreenshot returned no data" };
     }
-    const dims = parsePngDimensions(image_base64) ?? {
-      width: Math.round(rect.width),
-      height: Math.round(rect.height),
-    };
+    const dims = parsePngDimensions(image_base64);
+    if (!dims) {
+      return { code: "cdp_failed", message: "Page.captureScreenshot returned invalid PNG data" };
+    }
     return { image_base64, width: dims.width, height: dims.height };
   } catch (err) {
     return {
@@ -402,7 +402,10 @@ export async function handleScreenshot(
   if (isRpcError(captured)) return captured;
   if (signal?.aborted) return cancelled("screenshot");
   const image_base64 = captured;
-  const dims = parsePngDimensions(image_base64) ?? { width: 0, height: 0 };
+  const dims = parsePngDimensions(image_base64);
+  if (!dims) {
+    return { code: "cdp_failed", message: "screenshot capture returned invalid PNG data" };
+  }
   return withShotDialogs({
     image_base64,
     width: dims.width,

apps/extension/src/tools/vom/__tests__/capture.test.ts

@@ -1,6 +1,7 @@
 import { describe, expect, it, vi } from "vitest";
 import { OVERLAY_HOST_MARKER_ATTR, OVERLAY_HOST_NAME } from "../../../lib/overlay-bridge";
 import { captureViewModel, collectOverlayExcludedBackendIds, probeHoverSurfaces } from "../capture";
+import { discoverRenderedSurfaces } from "../rendered-surfaces";
 
 // Minimal but format-accurate captureSnapshot reply: a body with one
 // fixed full-screen overlay div carrying a password input.
@@ -1188,6 +1189,97 @@ describe("captureViewModel", () => {
     expect(input?.rect).toEqual({ x: 16, y: 28, w: 40, h: 20 });
   });
 
+  it("keeps same-origin iframe canvas discovery in top-level CSS coordinates at dpr 1 and 2", async () => {
+    const captureAtDpr = async (dpr: number) => {
+      const S = ["html", "body", "iframe", "canvas", "static", "auto"];
+      const i = (s: string) => S.indexOf(s);
+      const scale = (values: number[]) => values.map((value) => value * dpr);
+      const snapshot = {
+        strings: S,
+        documents: [
+          {
+            frameId: "main",
+            nodes: {
+              parentIndex: [-1, 0, 1],
+              nodeName: [i("html"), i("body"), i("iframe")],
+              backendNodeId: [10, 11, 13],
+              attributes: [[], [], []],
+              contentDocumentIndex: { index: [2], value: [1] },
+            },
+            layout: {
+              nodeIndex: [1, 2],
+              styles: [
+                [i("static"), i("auto")],
+                [i("static"), i("auto")],
+              ],
+              bounds: [scale([0, 0, 1000, 800]), scale([100, 150, 300, 200])],
+              paintOrders: [0, 1],
+            },
+          },
+          {
+            frameId: "child",
+            nodes: {
+              parentIndex: [-1, 0],
+              nodeName: [i("body"), i("canvas")],
+              backendNodeId: [20, 21],
+              attributes: [[], []],
+            },
+            layout: {
+              nodeIndex: [0, 1],
+              styles: [
+                [i("static"), i("auto")],
+                [i("static"), i("auto")],
+              ],
+              bounds: [scale([0, 0, 300, 200]), scale([20, 30, 120, 60])],
+              paintOrders: [0, 1],
+            },
+          },
+        ],
+      };
+      const cdp = {
+        send: vi.fn(async (_tabId: number, method: string) => {
+          if (method === "DOMSnapshot.enable") return {};
+          if (method === "DOMSnapshot.captureSnapshot") return snapshot;
+          if (method === "Page.getLayoutMetrics") {
+            return {
+              cssLayoutViewport: { clientWidth: 1000, clientHeight: 800, pageX: 0, pageY: 0 },
+              layoutViewport: { clientWidth: 1000 * dpr, clientHeight: 800 * dpr },
+            };
+          }
+          throw new Error(method);
+        }) as unknown as <T>(tabId: number, method: string, params?: object) => Promise<T>,
+      };
+      const captured = await captureViewModel(cdp, 4);
+      const childNodes = captured.frameNodes?.get("child") ?? [];
+      return discoverRenderedSurfaces(
+        [
+          {
+            frameId: "child",
+            contextScopeId: "child",
+            parentFrameId: "main",
+            ownerBackendNodeId: 13,
+            target: { tabId: 4 },
+            domNodes: childNodes,
+            axNodes: [],
+          },
+        ],
+        captured.viewport,
+      );
+    };
+
+    const atDpr1 = await captureAtDpr(1);
+    const atDpr2 = await captureAtDpr(2);
+
+    expect(atDpr1).toEqual([
+      expect.objectContaining({
+        backendNodeId: 21,
+        frameId: "child",
+        visibleRect: { x: 120, y: 180, w: 120, h: 60 },
+      }),
+    ]);
+    expect(atDpr2).toEqual(atDpr1);
+  });
+
   it("normalizes bounds then subtracts CSS scroll at dpr>1", async () => {
     const S = ["html", "body", "div", "position", "fixed", "static", "pointer-events", "auto"];
     const i = (s: string) => S.indexOf(s);