Fix browser screenshot capture for parked webviews

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

@@ -1,8 +1,8 @@
-import type { WebContents } from "electron";
+import type { Rectangle } from "electron";
 import { ipcMain } from "electron";
 import { BrowserAutomationExecuteRequestSchema } from "@getpaseo/protocol/browser-automation/rpc-schemas";
 import type { BrowserAutomationConsoleLogEntry } from "@getpaseo/protocol/browser-automation/rpc-schemas";
-import type { TabContents, BrowserRegistry } from "./service.js";
+import type { TabContents, BrowserRegistry, TabImage } from "./service.js";
 import { executeAutomationCommand } from "./service.js";
 import {
   listRegisteredPaseoBrowserIds,
@@ -13,14 +13,86 @@ import {
 } from "../browser-webviews/index.js";
 
 const MAX_CONSOLE_MESSAGES_PER_TAB = 200;
+const PIXEL_CAPTURE_BRIDGE_TIMEOUT_MS = 5_000;
 const consoleMessagesByContentsId = new Map<number, BrowserAutomationConsoleLogEntry[]>();
 const observedContentsIds = new Set<number>();
+let nextPixelCaptureBridgeRequest = 0;
 
 interface IpcHandlerRegistry {
   handle(channel: string, listener: (event: unknown, ...args: unknown[]) => unknown): void;
 }
 
-function adaptWebContents(contents: WebContents): TabContents {
+type IpcListener = (event: unknown, payload: unknown) => void;
+
+interface IpcCaptureBridge {
+  on(channel: string, listener: IpcListener): void;
+  removeListener(channel: string, listener: IpcListener): void;
+}
+
+interface HostWebContents {
+  readonly id: number;
+  isDestroyed(): boolean;
+  send(channel: string, payload: unknown): void;
+}
+
+interface WebContentsDebugger {
+  isAttached(): boolean;
+  attach(protocolVersion?: string): void;
+  sendCommand(command: string, params?: Record<string, unknown>): Promise<unknown>;
+}
+
+interface ConsoleMessageEmitter {
+  on(
+    event: "console-message",
+    listener: (
+      event: unknown,
+      level: unknown,
+      message: unknown,
+      line: unknown,
+      sourceId: unknown,
+    ) => void,
+  ): void;
+  once(event: "destroyed", listener: () => void): void;
+}
+
+interface BrowserAutomationWebContents extends ConsoleMessageEmitter {
+  readonly id: number;
+  readonly hostWebContents: HostWebContents | null;
+  readonly debugger: WebContentsDebugger;
+  getURL(): string;
+  getTitle(): string;
+  canGoBack(): boolean;
+  canGoForward(): boolean;
+  isLoading(): boolean;
+  isDestroyed(): boolean;
+  executeJavaScript(code: string): Promise<unknown>;
+  loadURL(url: string): Promise<void>;
+  goBack(): void;
+  goForward(): void;
+  reload(): void;
+  capturePage(rect?: Rectangle, options?: { stayHidden?: boolean }): Promise<TabImage>;
+  invalidate(): void;
+  getBackgroundThrottling(): boolean;
+  setBackgroundThrottling(allowed: boolean): void;
+}
+
+type PixelCaptureBridgeKind = "prepare" | "restore";
+
+interface PixelCaptureBridgeSuccess {
+  token?: string;
+}
+
+interface PixelCaptureBridgeOptions {
+  ipc?: IpcCaptureBridge;
+  createRequestId?: () => string;
+  timeoutMs?: number;
+}
+
+export function adaptWebContents(
+  contents: BrowserAutomationWebContents,
+  browserId: string,
+  options?: PixelCaptureBridgeOptions,
+): TabContents {
   observeConsoleMessages(contents);
   return {
     id: contents.id,
@@ -35,7 +107,19 @@ function adaptWebContents(contents: WebContents): TabContents {
     goBack: () => contents.goBack(),
     goForward: () => contents.goForward(),
     reload: () => contents.reload(),
-    capturePage: (options) => contents.capturePage(undefined, options),
+    capturePage: (captureOptions) => contents.capturePage(undefined, captureOptions),
+    prepareForPixelCapture: async () => {
+      const result = await requestPixelCaptureBridge(contents, browserId, "prepare", options);
+      if (!result.token) {
+        throw new Error("Browser pixel capture preparation did not return a token.");
+      }
+      return { token: result.token };
+    },
+    restorePixelCapture: async (preparation) => {
+      await requestPixelCaptureBridge(contents, browserId, "restore", options, {
+        token: preparation.token,
+      });
+    },
     invalidate: () => contents.invalidate(),
     isBackgroundThrottlingAllowed: () => contents.getBackgroundThrottling(),
     setBackgroundThrottling: (allowed) => contents.setBackgroundThrottling(allowed),
@@ -49,7 +133,101 @@ function adaptWebContents(contents: WebContents): TabContents {
   };
 }
 
-function observeConsoleMessages(contents: WebContents): void {
+function requestPixelCaptureBridge(
+  contents: BrowserAutomationWebContents,
+  browserId: string,
+  kind: PixelCaptureBridgeKind,
+  options: PixelCaptureBridgeOptions | undefined,
+  extraPayload?: { token: string },
+): Promise<PixelCaptureBridgeSuccess> {
+  const host = contents.hostWebContents;
+  if (!host || host.isDestroyed()) {
+    return Promise.reject(new Error("Browser host renderer is not available."));
+  }
+
+  const ipc = options?.ipc ?? ipcMain;
+  const requestId =
+    options?.createRequestId?.() ?? `browser-pixel-capture-${++nextPixelCaptureBridgeRequest}`;
+  const timeoutMs = options?.timeoutMs ?? PIXEL_CAPTURE_BRIDGE_TIMEOUT_MS;
+  const requestChannel =
+    kind === "prepare" ? "paseo:browser:capture-prepare" : "paseo:browser:capture-restore";
+  const responseChannel =
+    kind === "prepare" ? "paseo:browser:capture-prepared" : "paseo:browser:capture-restored";
+
+  return new Promise((resolve, reject) => {
+    let timeoutId: ReturnType<typeof setTimeout> | undefined;
+    const cleanup = () => {
+      if (timeoutId) {
+        clearTimeout(timeoutId);
+      }
+      ipc.removeListener(responseChannel, listener);
+    };
+    const listener: IpcListener = (_event, payload) => {
+      const response = readPixelCaptureBridgeResponse(payload, requestId);
+      if (!response) {
+        return;
+      }
+      cleanup();
+      if (response.ok) {
+        resolve(response);
+      } else {
+        reject(new Error(response.message));
+      }
+    };
+
+    ipc.on(responseChannel, listener);
+    timeoutId = setTimeout(() => {
+      cleanup();
+      reject(new Error(`Browser pixel capture ${kind} timed out.`));
+    }, timeoutMs);
+
+    try {
+      host.send(requestChannel, {
+        requestId,
+        browserId,
+        ...(extraPayload ? { token: extraPayload.token } : {}),
+      });
+    } catch (error) {
+      cleanup();
+      reject(error);
+    }
+  });
+}
+
+function readPixelCaptureBridgeResponse(
+  payload: unknown,
+  requestId: string,
+): ({ ok: true } & PixelCaptureBridgeSuccess) | { ok: false; message: string } | null {
+  if (!isRecord(payload)) {
+    return null;
+  }
+  const record = payload;
+  if (record.requestId !== requestId) {
+    return null;
+  }
+  if (record.ok === true) {
+    return {
+      ok: true,
+      ...(typeof record.token === "string" && record.token.length > 0
+        ? { token: record.token }
+        : {}),
+    };
+  }
+  if (record.ok === false) {
+    const message =
+      typeof record.message === "string" && record.message.length > 0
+        ? record.message
+        : "Browser pixel capture bridge failed.";
+    return { ok: false, message };
+  }
+  return null;
+}
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+  return typeof value === "object" && value !== null && !Array.isArray(value);
+}
+
+function observeConsoleMessages(contents: BrowserAutomationWebContents): void {
   if (observedContentsIds.has(contents.id)) {
     return;
   }
@@ -89,7 +267,7 @@ function createRegistry(): BrowserRegistry {
     listRegisteredBrowserIdsForWorkspace: listRegisteredPaseoBrowserIdsForWorkspace,
     getTabContents(browserId: string): TabContents | null {
       const contents = getPaseoBrowserWebContents(browserId);
-      return contents ? adaptWebContents(contents) : null;
+      return contents ? adaptWebContents(contents, browserId) : null;
     },
     getBrowserWorkspaceId: getPaseoBrowserWorkspaceId,
     getWorkspaceActiveBrowserId: getWorkspaceActivePaseoBrowserId,