Harden browser capture prep cleanup

packages/desktop/src/features/browser-automation/ipc.test.ts

@@ -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");

packages/desktop/src/features/browser-automation/ipc.ts

@@ -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,

packages/desktop/src/features/browser-automation/service.ts

@@ -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();

packages/desktop/src/preload.ts

@@ -2,12 +2,19 @@ import { contextBridge, ipcRenderer, webUtils } from "electron";
 
 type EventHandler = (payload: unknown) => void;
 type BrowserPixelCapturePrepareHandler = (input: {
+  requestId: string;
   browserId: string;
 }) => Promise<{ token: string }>;
 type BrowserPixelCaptureRestoreHandler = (input: { token: string }) => Promise<void>;
+type BrowserPixelCaptureCancelHandler = (input: {
+  requestId?: string;
+  token?: string;
+}) => Promise<void>;
 
 let prepareForPixelCaptureHandler: BrowserPixelCapturePrepareHandler | null = null;
 let restorePixelCaptureHandler: BrowserPixelCaptureRestoreHandler | null = null;
+let cancelPixelCaptureHandler: BrowserPixelCaptureCancelHandler | null = null;
+const canceledPixelCaptureRequestIds = new Set<string>();
 
 function readStringField(payload: unknown, key: string): string | null {
   if (!isRecord(payload)) {
@@ -38,13 +45,23 @@ ipcRenderer.on("paseo:browser:capture-prepare", async (_event, payload: unknown)
   }
 
   try {
-    const preparation = await prepareForPixelCaptureHandler({ browserId });
+    const preparation = await prepareForPixelCaptureHandler({ requestId, browserId });
+    if (canceledPixelCaptureRequestIds.delete(requestId)) {
+      await restorePixelCaptureHandler?.({ token: preparation.token });
+      ipcRenderer.send("paseo:browser:capture-prepared", {
+        requestId,
+        ok: false,
+        message: "Browser pixel capture preparation was canceled.",
+      });
+      return;
+    }
     ipcRenderer.send("paseo:browser:capture-prepared", {
       requestId,
       ok: true,
       token: preparation.token,
     });
   } catch (error) {
+    canceledPixelCaptureRequestIds.delete(requestId);
     ipcRenderer.send("paseo:browser:capture-prepared", {
       requestId,
       ok: false,
@@ -53,6 +70,25 @@ ipcRenderer.on("paseo:browser:capture-prepare", async (_event, payload: unknown)
   }
 });
 
+ipcRenderer.on("paseo:browser:capture-cancel", async (_event, payload: unknown) => {
+  const requestId = readStringField(payload, "requestId");
+  const token = readStringField(payload, "token");
+  if (requestId) {
+    canceledPixelCaptureRequestIds.add(requestId);
+  }
+  if (!requestId && !token) {
+    return;
+  }
+  try {
+    await cancelPixelCaptureHandler?.({
+      ...(requestId ? { requestId } : {}),
+      ...(token ? { token } : {}),
+    });
+  } catch {
+    // The original prepare/restore request owns the user-visible error.
+  }
+});
+
 ipcRenderer.on("paseo:browser:capture-restore", async (_event, payload: unknown) => {
   const requestId = readStringField(payload, "requestId");
   const token = readStringField(payload, "token");
@@ -181,5 +217,13 @@ contextBridge.exposeInMainWorld("paseoDesktop", {
         }
       };
     },
+    onCancelPixelCapture: (handler: BrowserPixelCaptureCancelHandler): (() => void) => {
+      cancelPixelCaptureHandler = handler;
+      return () => {
+        if (cancelPixelCaptureHandler === handler) {
+          cancelPixelCaptureHandler = null;
+        }
+      };
+    },
   },
 });