Add native desktop fallback for screenshot capture

apps/desktop/src/backendStartup.test.ts

@@ -0,0 +1,79 @@
+import * as Net from "node:net";
+
+import { afterEach, describe, expect, it } from "vitest";
+
+import { waitForTcpServer } from "./backendStartup";
+
+const servers = new Set<Net.Server>();
+
+async function listenOnLoopback(delayMs = 0): Promise<{ server: Net.Server; port: number }> {
+  const server = Net.createServer();
+  servers.add(server);
+
+  await new Promise<void>((resolve, reject) => {
+    server.once("error", reject);
+    setTimeout(() => {
+      server.listen(0, "127.0.0.1", () => {
+        resolve();
+      });
+    }, delayMs);
+  });
+
+  const address = server.address();
+  const port = typeof address === "object" && address !== null ? address.port : 0;
+  return { server, port };
+}
+
+afterEach(async () => {
+  await Promise.all(
+    [...servers].map(
+      (server) =>
+        new Promise<void>((resolve) => {
+          server.close(() => resolve());
+        }),
+    ),
+  );
+  servers.clear();
+});
+
+describe("waitForTcpServer", () => {
+  it("resolves when the server is already listening", async () => {
+    const { port } = await listenOnLoopback();
+
+    await expect(waitForTcpServer({ host: "127.0.0.1", port, timeoutMs: 500 })).resolves.toBe(
+      undefined,
+    );
+  });
+
+  it("waits for a delayed server startup", async () => {
+    let attempts = 0;
+
+    await expect(
+      waitForTcpServer({
+        host: "127.0.0.1",
+        port: 3773,
+        timeoutMs: 1_000,
+        retryDelayMs: 20,
+        tryConnect: async () => {
+          attempts += 1;
+          return attempts >= 4;
+        },
+      }),
+    ).resolves.toBe(undefined);
+    expect(attempts).toBe(4);
+  });
+
+  it("fails when the server never starts", async () => {
+    const probe = Net.createServer();
+    await new Promise<void>((resolve) => {
+      probe.listen(0, "127.0.0.1", () => resolve());
+    });
+    const address = probe.address();
+    const port = typeof address === "object" && address !== null ? address.port : 0;
+    await new Promise<void>((resolve) => probe.close(() => resolve()));
+
+    await expect(
+      waitForTcpServer({ host: "127.0.0.1", port, timeoutMs: 120, retryDelayMs: 20 }),
+    ).rejects.toThrow(`Timed out waiting for backend on 127.0.0.1:${port}`);
+  });
+});

apps/desktop/src/backendStartup.ts

@@ -0,0 +1,60 @@
+import * as Net from "node:net";
+
+export interface WaitForTcpServerInput {
+  readonly host: string;
+  readonly port: number;
+  readonly timeoutMs?: number;
+  readonly retryDelayMs?: number;
+  readonly tryConnect?: (host: string, port: number) => Promise<boolean>;
+}
+
+export async function waitForTcpServer(input: WaitForTcpServerInput): Promise<void> {
+  const timeoutMs = input.timeoutMs ?? 15_000;
+  const retryDelayMs = input.retryDelayMs ?? 100;
+  const deadline = Date.now() + timeoutMs;
+
+  while (Date.now() < deadline) {
+    const connected = await (input.tryConnect ?? tryConnect)(input.host, input.port);
+    if (connected) {
+      return;
+    }
+    await delay(retryDelayMs);
+  }
+
+  throw new Error(`Timed out waiting for backend on ${input.host}:${input.port}`);
+}
+
+function tryConnect(host: string, port: number): Promise<boolean> {
+  return new Promise((resolve) => {
+    const socket = Net.createConnection({ host, port });
+    let settled = false;
+
+    const settle = (value: boolean) => {
+      if (settled) {
+        return;
+      }
+      settled = true;
+      socket.destroy();
+      resolve(value);
+    };
+
+    socket.setTimeout(500, () => {
+      settle(false);
+    });
+    socket.once("connect", () => {
+      settle(true);
+    });
+    socket.once("error", () => {
+      settle(false);
+    });
+    socket.once("close", () => {
+      settle(false);
+    });
+  });
+}
+
+function delay(ms: number): Promise<void> {
+  return new Promise((resolve) => {
+    setTimeout(resolve, ms);
+  });
+}

apps/desktop/src/main.ts

@@ -32,6 +32,7 @@ import type { ContextMenuItem } from "@okcode/contracts";
 import { NetService } from "@okcode/shared/Net";
 import { RotatingFileSink } from "@okcode/shared/logging";
 import { showDesktopConfirmDialog } from "./confirmDialog";
+import { waitForTcpServer } from "./backendStartup";
 import { createEmptyTabsState } from "./preview";
 import { DesktopPreviewController } from "./previewController";
 import { resolveDesktopRendererUrl } from "./rendererUrl";
@@ -54,6 +55,7 @@ import { isArm64HostRunningIntelBuild, resolveDesktopRuntimeInfo } from "./runti
 syncShellEnvironment();
 
 const PICK_FOLDER_CHANNEL = "desktop:pick-folder";
+const CAPTURE_WINDOW_CHANNEL = "desktop:capture-window";
 const CONFIRM_CHANNEL = "desktop:confirm";
 const SET_THEME_CHANNEL = "desktop:set-theme";
 const SET_SIDEBAR_OPACITY_CHANNEL = "desktop:set-sidebar-opacity";
@@ -1153,6 +1155,19 @@ function registerIpcHandlers(): void {
     return result.filePaths[0] ?? null;
   });
 
+  ipcMain.removeHandler(CAPTURE_WINDOW_CHANNEL);
+  ipcMain.handle(CAPTURE_WINDOW_CHANNEL, async (event) => {
+    try {
+      const image = await event.sender.capturePage();
+      if (image.isEmpty()) {
+        return null;
+      }
+      return image.toDataURL();
+    } catch {
+      return null;
+    }
+  });
+
   ipcMain.removeHandler(CONFIRM_CHANNEL);
   ipcMain.handle(CONFIRM_CHANNEL, async (_event, message: unknown) => {
     if (typeof message !== "string") {
@@ -1660,6 +1675,8 @@ async function bootstrap(): Promise<void> {
   writeDesktopLogHeader("bootstrap ipc handlers registered");
   startBackend();
   writeDesktopLogHeader("bootstrap backend start requested");
+  await waitForTcpServer({ host: "127.0.0.1", port: backendPort });
+  writeDesktopLogHeader(`bootstrap backend ready port=${backendPort}`);
   mainWindow = createWindow();
   writeDesktopLogHeader("bootstrap main window created");
 }

apps/desktop/src/preload.ts

@@ -2,6 +2,7 @@ import { contextBridge, ipcRenderer } from "electron";
 import type { DesktopBridge } from "@okcode/contracts";
 
 const PICK_FOLDER_CHANNEL = "desktop:pick-folder";
+const CAPTURE_WINDOW_CHANNEL = "desktop:capture-window";
 const CONFIRM_CHANNEL = "desktop:confirm";
 const SET_THEME_CHANNEL = "desktop:set-theme";
 const SET_SIDEBAR_OPACITY_CHANNEL = "desktop:set-sidebar-opacity";
@@ -34,6 +35,7 @@ const wsUrl = process.env.OKCODE_DESKTOP_WS_URL ?? null;
 
 contextBridge.exposeInMainWorld("desktopBridge", {
   getWsUrl: () => wsUrl,
+  captureWindow: () => ipcRenderer.invoke(CAPTURE_WINDOW_CHANNEL),
   pickFolder: () => ipcRenderer.invoke(PICK_FOLDER_CHANNEL),
   confirm: (message) => ipcRenderer.invoke(CONFIRM_CHANNEL, message),
   setTheme: (theme) => ipcRenderer.invoke(SET_THEME_CHANNEL, theme),

apps/web/src/components/ScreenshotTool.tsx

@@ -7,6 +7,8 @@ import { readDesktopPreviewBridge } from "~/desktopPreview";
 import { toastManager } from "~/components/ui/toast";
 import { Button } from "~/components/ui/button";
 import { Tooltip, TooltipTrigger, TooltipPopup } from "~/components/ui/tooltip";
+import { readDesktopBridge } from "~/lib/runtimeBridge";
+import { buildDomCaptureOptions, captureBaseScreenshotDataUrl } from "~/lib/screenshotCapture";
 import { cn, isMacPlatform } from "~/lib/utils";
 
 // ── Types ───────────────────────────────────────────────────────────
@@ -60,17 +62,17 @@ async function captureRegion(rect: {
 
   // Capture the full page at device resolution (DOM only — native BrowserView is excluded)
   const rootElement = document.documentElement;
-  const dataUrl = await toPng(rootElement, {
-    width: rootElement.scrollWidth,
-    height: rootElement.scrollHeight,
-    pixelRatio: dpr,
-    // Exclude our own overlay from the capture
-    filter: (node) => {
-      if (node instanceof HTMLElement && node.dataset.screenshotOverlay === "true") {
-        return false;
-      }
-      return true;
-    },
+  const desktopBridge = readDesktopBridge();
+  const dataUrl = await captureBaseScreenshotDataUrl({
+    captureWindow: () => desktopBridge?.captureWindow() ?? Promise.resolve(null),
+    captureDom: () =>
+      toPng(
+        rootElement,
+        buildDomCaptureOptions({
+          rootElement,
+          pixelRatio: dpr,
+        }),
+      ),
   });
 
   // Load into an Image to crop

apps/web/src/lib/screenshotCapture.test.ts

@@ -0,0 +1,49 @@
+import { describe, expect, it, vi } from "vitest";
+
+import { buildDomCaptureOptions, captureBaseScreenshotDataUrl } from "./screenshotCapture";
+
+describe("screenshotCapture", () => {
+  it("prefers native desktop window capture when available", async () => {
+    const captureWindow = vi.fn(async () => "data:image/png;base64,desktop");
+    const captureDom = vi.fn(async () => "data:image/png;base64,dom");
+
+    const result = await captureBaseScreenshotDataUrl({ captureWindow, captureDom });
+
+    expect(result).toBe("data:image/png;base64,desktop");
+    expect(captureWindow).toHaveBeenCalledOnce();
+    expect(captureDom).not.toHaveBeenCalled();
+  });
+
+  it("falls back to DOM capture when native window capture is unavailable", async () => {
+    const captureWindow = vi.fn(async () => null);
+    const captureDom = vi.fn(async () => "data:image/png;base64,dom");
+
+    const result = await captureBaseScreenshotDataUrl({ captureWindow, captureDom });
+
+    expect(result).toBe("data:image/png;base64,dom");
+    expect(captureWindow).toHaveBeenCalledOnce();
+    expect(captureDom).toHaveBeenCalledOnce();
+  });
+
+  it("skips font embedding and excludes the screenshot overlay from DOM capture", () => {
+    const rootElement = {
+      scrollWidth: 1440,
+      scrollHeight: 900,
+    } as HTMLElement;
+
+    const options = buildDomCaptureOptions({
+      rootElement,
+      pixelRatio: 2,
+    });
+
+    const overlay = { dataset: { screenshotOverlay: "true" } } as unknown as HTMLElement;
+    const content = { dataset: {} } as unknown as HTMLElement;
+
+    expect(options.width).toBe(1440);
+    expect(options.height).toBe(900);
+    expect(options.pixelRatio).toBe(2);
+    expect(options.skipFonts).toBe(true);
+    expect(options.filter?.(overlay as HTMLElement)).toBe(false);
+    expect(options.filter?.(content as HTMLElement)).toBe(true);
+  });
+});

apps/web/src/lib/screenshotCapture.ts

@@ -0,0 +1,49 @@
+const TRANSPARENT_IMAGE_PLACEHOLDER = "data:image/gif;base64,R0lGODlhAQABAAAAACwAAAAAAQABAAA=";
+
+export interface ScreenshotCaptureRect {
+  readonly x: number;
+  readonly y: number;
+  readonly width: number;
+  readonly height: number;
+}
+
+export interface DomCaptureOptions {
+  readonly width: number;
+  readonly height: number;
+  readonly pixelRatio: number;
+  readonly skipFonts: boolean;
+  readonly imagePlaceholder: string;
+  readonly onImageErrorHandler: () => undefined;
+  readonly filter: (node: HTMLElement) => boolean;
+}
+
+export function buildDomCaptureOptions(input: {
+  readonly rootElement: HTMLElement;
+  readonly pixelRatio: number;
+}): DomCaptureOptions {
+  return {
+    width: input.rootElement.scrollWidth,
+    height: input.rootElement.scrollHeight,
+    pixelRatio: input.pixelRatio,
+    skipFonts: true,
+    imagePlaceholder: TRANSPARENT_IMAGE_PLACEHOLDER,
+    onImageErrorHandler: () => undefined,
+    filter: (node: HTMLElement) => {
+      if ("dataset" in node && node.dataset?.screenshotOverlay === "true") {
+        return false;
+      }
+      return true;
+    },
+  };
+}
+
+export async function captureBaseScreenshotDataUrl(input: {
+  readonly captureWindow?: (() => Promise<string | null>) | null;
+  readonly captureDom: () => Promise<string>;
+}): Promise<string> {
+  const nativeCapture = await input.captureWindow?.();
+  if (typeof nativeCapture === "string" && nativeCapture.length > 0) {
+    return nativeCapture;
+  }
+  return input.captureDom();
+}

packages/contracts/src/ipc.ts

@@ -275,6 +275,7 @@ export interface PreviewNavigateResult {
 
 export interface DesktopBridge {
   getWsUrl: () => string | null;
+  captureWindow: () => Promise<string | null>;
   pickFolder: () => Promise<string | null>;
   confirm: (message: string) => Promise<boolean>;
   setTheme: (theme: DesktopTheme) => Promise<void>;