fix(screenshot): hide overlays during capture so shots only contain page content (#64)

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

@@ -1,5 +1,6 @@
 import { renderVom } from "@browser-skill/vom";
 import { describe, expect, it, vi } from "vitest";
+import { CAPTURE_SUPPRESS, type CaptureSuppressMessage } from "@/lib/capture-suppress-bridge";
 import { OVERLAY_HOST_MARKER_ATTR, OVERLAY_HOST_NAME } from "@/lib/overlay-bridge";
 import { SessionManager } from "@/session-manager/manager";
 import type { CdpRunner } from "@/tools/shared";
@@ -345,6 +346,98 @@ describe("handleScreenshot", () => {
   });
 });
 
+describe("handleScreenshot overlay suppression", () => {
+  /**
+   * Fake background→content bridge that mirrors the content script's
+   * contract: `begin` hides the overlay, `end` restores it. Records the
+   * phase order so tests can assert begin → capture → end.
+   */
+  function fakeSuppressBridge(events: string[]) {
+    let overlayHidden = false;
+    const sendToTab = vi.fn(async (_tabId: number, message: CaptureSuppressMessage) => {
+      events.push(message.phase);
+      overlayHidden = message.phase === "begin";
+      return { type: CAPTURE_SUPPRESS, ok: true };
+    });
+    return { sendToTab, isHidden: () => overlayHidden };
+  }
+
+  it("hides the overlay around a visible-tab capture", async () => {
+    const sm = new SessionManager({ agentWindow: fakeAgentWindow([100]) });
+    await sm.start("aa11");
+    const events: string[] = [];
+    const bridge = fakeSuppressBridge(events);
+    const capture = vi.fn(async (_w: number) => {
+      events.push("capture");
+      // The overlay must already be hidden when the capture runs.
+      expect(bridge.isHidden()).toBe(true);
+      return `data:image/png;base64,${TINY_PNG}`;
+    });
+    const res = await handleScreenshot(
+      sm,
+      { session_id: "aa11" },
+      { ...makeScreenshotDeps({ captureVisibleTab: capture }), sendToTab: bridge.sendToTab },
+    );
+    if ("code" in res) throw new Error(`unexpected error: ${JSON.stringify(res)}`);
+    expect(res.image_base64).toBe(TINY_PNG);
+    expect(events).toEqual(["begin", "capture", "end"]);
+    expect(bridge.isHidden()).toBe(false);
+  });
+
+  it("hides the overlay around a CDP element capture", async () => {
+    const sm = new SessionManager({ agentWindow: fakeAgentWindow([100]) });
+    const ctx = await sm.start("aa11");
+    ctx.refStore.set("e5", 999, { tabId: 7 });
+    const events: string[] = [];
+    const bridge = fakeSuppressBridge(events);
+    const { cdp } = makeFakeCdp({
+      "DOM.scrollIntoViewIfNeeded": () => ({}),
+      "DOM.getContentQuads": () => ({ quads: [[10, 20, 110, 20, 110, 60, 10, 60]] }),
+      "Page.captureScreenshot": () => {
+        events.push("capture");
+        expect(bridge.isHidden()).toBe(true);
+        return { data: TINY_PNG };
+      },
+    });
+    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(),
+        }),
+        sendToTab: bridge.sendToTab,
+      },
+    );
+    if ("code" in res) throw new Error(`unexpected error: ${JSON.stringify(res)}`);
+    expect(res.image_base64).toBe(TINY_PNG);
+    expect(events).toEqual(["begin", "capture", "end"]);
+    expect(bridge.isHidden()).toBe(false);
+  });
+
+  it("still captures when the tab has no content script to ack", async () => {
+    const sm = new SessionManager({ agentWindow: fakeAgentWindow([100]) });
+    await sm.start("aa11");
+    const sendToTab = vi.fn(async () => {
+      throw new Error("Could not establish connection. Receiving end does not exist.");
+    });
+    const capture = vi.fn(async (_w: number) => `data:image/png;base64,${TINY_PNG}`);
+    const res = await handleScreenshot(
+      sm,
+      { session_id: "aa11" },
+      { ...makeScreenshotDeps({ captureVisibleTab: capture }), sendToTab },
+    );
+    if ("code" in res) throw new Error(`unexpected error: ${JSON.stringify(res)}`);
+    expect(res.image_base64).toBe(TINY_PNG);
+    expect(capture).toHaveBeenCalledTimes(1);
+    // The failed begin must not be followed by an end.
+    expect(sendToTab).toHaveBeenCalledTimes(1);
+  });
+});
+
 // ---------------------------------------------------------------------------
 // buildVomScene
 // ---------------------------------------------------------------------------

apps/extension/src/tools/observation.ts

@@ -11,6 +11,10 @@ import {
   type VomScene,
 } from "@browser-skill/vom";
 import { ChromiumCdp } from "@/browser-driver/chromium-cdp";
+import {
+  type CaptureSuppressSendToTab,
+  withOverlaysHiddenForCapture,
+} from "@/lib/capture-suppress-bridge";
 import type { SessionManager } from "@/session-manager/manager";
 import type {
   GetHtmlParams,
@@ -114,6 +118,12 @@ export interface ScreenshotDeps {
   cdp?: SharedCdpRunner;
   tabsApi: ChromeTabsApi;
   captureApi: ChromeTabsCaptureApi;
+  /**
+   * Bridge used to hide the in-page overlay while a screenshot is taken,
+   * so captured frames only contain page content. Defaults to
+   * `chrome.tabs.sendMessage`; tests inject a fake.
+   */
+  sendToTab?: CaptureSuppressSendToTab;
 }
 
 function defaultScreenshotDeps(): ScreenshotDeps {
@@ -192,7 +202,12 @@ export async function handleScreenshot(
     if (isRpcError(node)) return node;
     deps.cdp.trackSessionTab?.(ctx.sessionId, target.tabId);
     await deps.cdp.ensureAttachedToUrl?.(target.tabId, target.url);
-    const captured = await captureElementScreenshot(deps.cdp, target.tabId, node.backendNodeId);
+    const cdp = deps.cdp;
+    const captured = await withOverlaysHiddenForCapture(
+      target.tabId,
+      () => captureElementScreenshot(cdp, target.tabId, node.backendNodeId),
+      deps.sendToTab,
+    );
     if (isRpcError(captured)) return captured;
     return withShotDialogs({
       image_base64: captured.image_base64,
@@ -212,7 +227,11 @@ export async function handleScreenshot(
   }
 
   try {
-    const dataUrl = await deps.captureApi.captureVisibleTab(target.windowId, { format: "png" });
+    const dataUrl = await withOverlaysHiddenForCapture(
+      target.tabId,
+      () => deps.captureApi.captureVisibleTab(target.windowId, { format: "png" }),
+      deps.sendToTab,
+    );
     const image_base64 = stripDataUrlPrefix(dataUrl);
     const dims = parsePngDimensions(image_base64) ?? { width: 0, height: 0 };
     return withShotDialogs({