Harden browser capture prep cleanup

2026-07-03 03:53:53 +02:00
parent 6781521c8e
commit 33ffeb4d42
7 changed files with 397 additions and 43 deletions
@@ -1,5 +1,5 @@
 import type { Rectangle } from "electron";
-import { describe, expect, test } from "vitest";
+import { describe, expect, test, vi } from "vitest";
 import type { TabImage } from "./service.js";
 import { adaptWebContents } from "./ipc.js";

@@ -40,7 +40,13 @@ class FakeHostWebContents {
  }
 }

-type IpcListener = (event: unknown, payload: unknown) => void;
+interface FakeIpcEvent {
+  sender: {
+    id: number;
+  };
+}
+
+type IpcListener = (event: FakeIpcEvent, payload: unknown) => void;

 class FakeIpcBridge {
  private readonly listeners = new Map<string, IpcListener[]>();
@@ -59,9 +65,10 @@ class FakeIpcBridge {
    );
  }

-  public emit(channel: string, payload: unknown): void {
+  public emit(channel: string, payload: unknown, input: { senderId?: number } = {}): void {
+    const event = { sender: { id: input.senderId ?? 10 } };
    for (const listener of this.listeners.get(channel) ?? []) {
-      listener({}, payload);
+      listener(event, payload);
    }
  }

@@ -190,14 +197,33 @@ describe("browser automation IPC adapter", () => {
    const host = new FakeHostWebContents(10);
    const contents = new FakeWebContents(20, host);
    const ipc = new FakeIpcBridge();
+    const requestIds = ["prepare-1", "restore-1"];
    const tab = adaptWebContents(contents, "browser-a", {
      ipc,
-      createRequestId: () => "restore-1",
+      createRequestId: () => {
+        const requestId = requestIds.shift();
+        if (!requestId) {
+          throw new Error("Missing request id");
+        }
+        return requestId;
+      },
    });

+    const preparation = tab.prepareForPixelCapture();
+    ipc.emit("paseo:browser:capture-prepared", {
+      requestId: "prepare-1",
+      ok: true,
+      token: "token-a",
+    });
+    await expect(preparation).resolves.toEqual({ token: "token-a" });
+
    const restored = tab.restorePixelCapture({ token: "token-a" });

    expect(host.sentMessages).toEqual([
+      {
+        channel: "paseo:browser:capture-prepare",
+        payload: { requestId: "prepare-1", browserId: "browser-a" },
+      },
      {
        channel: "paseo:browser:capture-restore",
        payload: { requestId: "restore-1", browserId: "browser-a", token: "token-a" },
@@ -208,6 +234,51 @@ describe("browser automation IPC adapter", () => {
    await expect(restored).resolves.toBeUndefined();
  });

+  test("restorePixelCapture uses the host captured during preparation when the guest detaches", async () => {
+    const host = new FakeHostWebContents(10);
+    const contents = new FakeWebContents(20, host);
+    const ipc = new FakeIpcBridge();
+    const requestIds = ["prepare-1", "restore-1"];
+    const tab = adaptWebContents(contents, "browser-a", {
+      ipc,
+      createRequestId: () => {
+        const requestId = requestIds.shift();
+        if (!requestId) {
+          throw new Error("Missing request id");
+        }
+        return requestId;
+      },
+    });
+
+    const preparationPromise = tab.prepareForPixelCapture();
+    ipc.emit("paseo:browser:capture-prepared", {
+      requestId: "prepare-1",
+      ok: true,
+      token: "token-a",
+    });
+    const preparation = await preparationPromise;
+    contents.hostWebContents = null;
+
+    const restored = tab.restorePixelCapture(preparation);
+
+    expect(host.sentMessages).toEqual([
+      {
+        channel: "paseo:browser:capture-prepare",
+        payload: { requestId: "prepare-1", browserId: "browser-a" },
+      },
+      {
+        channel: "paseo:browser:capture-restore",
+        payload: { requestId: "restore-1", browserId: "browser-a", token: "token-a" },
+      },
+    ]);
+    ipc.emit("paseo:browser:capture-restored", { requestId: "restore-1", ok: true });
+
+    await expect(restored).resolves.toBeUndefined();
+    await expect(tab.restorePixelCapture(preparation)).rejects.toThrow(
+      "Browser pixel capture preparation is no longer active.",
+    );
+  });
+
  test("prepareForPixelCapture rejects when the renderer reports preparation failure", async () => {
    const host = new FakeHostWebContents(10);
    const contents = new FakeWebContents(20, host);
@@ -228,6 +299,71 @@ describe("browser automation IPC adapter", () => {
    expect(ipc.listenerCount("paseo:browser:capture-prepared")).toBe(0);
  });

+  test("prepareForPixelCapture ignores matching responses from the wrong sender", async () => {
+    const host = new FakeHostWebContents(10);
+    const contents = new FakeWebContents(20, host);
+    const ipc = new FakeIpcBridge();
+    const tab = adaptWebContents(contents, "browser-a", {
+      ipc,
+      createRequestId: () => "prepare-1",
+    });
+
+    const preparation = tab.prepareForPixelCapture();
+    ipc.emit(
+      "paseo:browser:capture-prepared",
+      {
+        requestId: "prepare-1",
+        ok: true,
+        token: "spoofed-token",
+      },
+      { senderId: 99 },
+    );
+    expect(ipc.listenerCount("paseo:browser:capture-prepared")).toBe(1);
+
+    ipc.emit("paseo:browser:capture-prepared", {
+      requestId: "prepare-1",
+      ok: true,
+      token: "token-a",
+    });
+
+    await expect(preparation).resolves.toEqual({ token: "token-a" });
+  });
+
+  test("prepareForPixelCapture cancels the renderer request when the ack times out", async () => {
+    vi.useFakeTimers();
+    try {
+      const host = new FakeHostWebContents(10);
+      const contents = new FakeWebContents(20, host);
+      const ipc = new FakeIpcBridge();
+      const tab = adaptWebContents(contents, "browser-a", {
+        ipc,
+        createRequestId: () => "prepare-1",
+        timeoutMs: 25,
+      });
+
+      const preparation = tab.prepareForPixelCapture();
+      const rejection = expect(preparation).rejects.toThrow(
+        "Browser pixel capture prepare timed out.",
+      );
+      await vi.advanceTimersByTimeAsync(25);
+
+      await rejection;
+      expect(host.sentMessages).toEqual([
+        {
+          channel: "paseo:browser:capture-prepare",
+          payload: { requestId: "prepare-1", browserId: "browser-a" },
+        },
+        {
+          channel: "paseo:browser:capture-cancel",
+          payload: { requestId: "prepare-1", browserId: "browser-a" },
+        },
+      ]);
+      expect(ipc.listenerCount("paseo:browser:capture-prepared")).toBe(0);
+    } finally {
+      vi.useRealTimers();
+    }
+  });
+
  test("prepareForPixelCapture rejects when the guest has no embedder renderer", async () => {
    const contents = new FakeWebContents(20, null);
    const tab = adaptWebContents(contents, "browser-a");
@@ -22,7 +22,13 @@ interface IpcHandlerRegistry {
  handle(channel: string, listener: (event: unknown, ...args: unknown[]) => unknown): void;
 }

-type IpcListener = (event: unknown, payload: unknown) => void;
+interface IpcBridgeEvent {
+  sender?: {
+    id?: number;
+  };
+}
+
+type IpcListener = (event: IpcBridgeEvent, payload: unknown) => void;

 interface IpcCaptureBridge {
  on(channel: string, listener: IpcListener): void;
@@ -88,12 +94,18 @@ interface PixelCaptureBridgeOptions {
  timeoutMs?: number;
 }

+interface PreparedPixelCapture {
+  browserId: string;
+  host: HostWebContents;
+}
+
 export function adaptWebContents(
  contents: BrowserAutomationWebContents,
  browserId: string,
  options?: PixelCaptureBridgeOptions,
 ): TabContents {
  observeConsoleMessages(contents);
+  const preparedPixelCapturesByToken = new Map<string, PreparedPixelCapture>();
  return {
    id: contents.id,
    getURL: () => contents.getURL(),
@@ -109,16 +121,30 @@ export function adaptWebContents(
    reload: () => contents.reload(),
    capturePage: (captureOptions) => contents.capturePage(undefined, captureOptions),
    prepareForPixelCapture: async () => {
-      const result = await requestPixelCaptureBridge(contents, browserId, "prepare", options);
+      const host = getPixelCaptureHost(contents);
+      const result = await requestPixelCaptureBridge({ host, browserId, kind: "prepare", options });
      if (!result.token) {
        throw new Error("Browser pixel capture preparation did not return a token.");
      }
+      preparedPixelCapturesByToken.set(result.token, { browserId, host });
      return { token: result.token };
    },
    restorePixelCapture: async (preparation) => {
-      await requestPixelCaptureBridge(contents, browserId, "restore", options, {
-        token: preparation.token,
-      });
+      const prepared = preparedPixelCapturesByToken.get(preparation.token);
+      if (!prepared) {
+        throw new Error("Browser pixel capture preparation is no longer active.");
+      }
+      try {
+        await requestPixelCaptureBridge({
+          host: prepared.host,
+          browserId: prepared.browserId,
+          kind: "restore",
+          options,
+          extraPayload: { token: preparation.token },
+        });
+      } finally {
+        preparedPixelCapturesByToken.delete(preparation.token);
+      }
    },
    invalidate: () => contents.invalidate(),
    isBackgroundThrottlingAllowed: () => contents.getBackgroundThrottling(),
@@ -133,15 +159,23 @@ export function adaptWebContents(
  };
 }

-function requestPixelCaptureBridge(
-  contents: BrowserAutomationWebContents,
-  browserId: string,
-  kind: PixelCaptureBridgeKind,
-  options: PixelCaptureBridgeOptions | undefined,
-  extraPayload?: { token: string },
-): Promise<PixelCaptureBridgeSuccess> {
+function getPixelCaptureHost(contents: BrowserAutomationWebContents): HostWebContents {
  const host = contents.hostWebContents;
  if (!host || host.isDestroyed()) {
+    throw new Error("Browser host renderer is not available.");
+  }
+  return host;
+}
+
+function requestPixelCaptureBridge(input: {
+  host: HostWebContents;
+  browserId: string;
+  kind: PixelCaptureBridgeKind;
+  options: PixelCaptureBridgeOptions | undefined;
+  extraPayload?: { token: string };
+}): Promise<PixelCaptureBridgeSuccess> {
+  const { host, browserId, kind, options, extraPayload } = input;
+  if (host.isDestroyed()) {
    return Promise.reject(new Error("Browser host renderer is not available."));
  }

@@ -162,7 +196,10 @@ function requestPixelCaptureBridge(
      }
      ipc.removeListener(responseChannel, listener);
    };
-    const listener: IpcListener = (_event, payload) => {
+    const listener: IpcListener = (event, payload) => {
+      if (readSenderId(event) !== host.id) {
+        return;
+      }
      const response = readPixelCaptureBridgeResponse(payload, requestId);
      if (!response) {
        return;
@@ -178,6 +215,11 @@ function requestPixelCaptureBridge(
    ipc.on(responseChannel, listener);
    timeoutId = setTimeout(() => {
      cleanup();
+      sendPixelCaptureCancel(host, {
+        requestId,
+        browserId,
+        ...(extraPayload ? { token: extraPayload.token } : {}),
+      });
      reject(new Error(`Browser pixel capture ${kind} timed out.`));
    }, timeoutMs);

@@ -194,6 +236,25 @@ function requestPixelCaptureBridge(
  });
 }

+function sendPixelCaptureCancel(
+  host: HostWebContents,
+  payload: { browserId: string; requestId?: string; token?: string },
+): void {
+  if (host.isDestroyed()) {
+    return;
+  }
+  try {
+    host.send("paseo:browser:capture-cancel", payload);
+  } catch {
+    // The original prepare/restore request owns the user-visible error.
+  }
+}
+
+function readSenderId(event: IpcBridgeEvent): number | null {
+  const senderId = event.sender?.id;
+  return typeof senderId === "number" ? senderId : null;
+}
+
 function readPixelCaptureBridgeResponse(
  payload: unknown,
  requestId: string,
@@ -159,7 +159,7 @@ async function runPreparedPixelCapture<T>(
    try {
      // Offscreen-parked webview guests have no compositor surface for
      // capturePage/CDP to copy. The renderer must briefly make the host
-      // paintable before capture; see ~/.paseo/plans/browser-capture-harness-results.md.
+      // paintable before capture; see docs/browser-capture-harness.md.
      preparation = await prepareForPixelCapture(contents);
      contents.setBackgroundThrottling(false);
      contents.invalidate();