Fix full-page screenshots for pages using viewport units

desktop-app/src/main/preload-webview.ts

@@ -96,6 +96,10 @@
   );
 
   window.addEventListener('scroll', () => {
+    if ((window as Window & {__responsivelyScreenshotInProgress?: boolean})
+      .__responsivelyScreenshotInProgress) {
+      return;
+    }
     scrollPending = true;
     requestFlush();
   });

desktop-app/src/main/screenshot/bitmap-stitcher.ts

@@ -0,0 +1,46 @@
+import {nativeImage} from 'electron';
+
+export class BitmapStitcher {
+  private result: Buffer | undefined;
+
+  private width = 0;
+
+  private height = 0;
+
+  private pixelsPerCssPixel = 1;
+
+  constructor(
+    private readonly documentHeight: number,
+    private readonly viewportHeight: number
+  ) {}
+
+  append(tile: Electron.NativeImage, cssY: number) {
+    const {width, height} = tile.getSize();
+    if (this.result === undefined) {
+      this.pixelsPerCssPixel = height / this.viewportHeight;
+      this.width = width;
+      this.height = Math.ceil(this.documentHeight * this.pixelsPerCssPixel);
+      this.result = Buffer.alloc(this.width * this.height * 4);
+    }
+
+    const tileBitmap = tile.toBitmap();
+    const destinationY = Math.round(cssY * this.pixelsPerCssPixel);
+    const rowsToCopy = Math.min(height, this.height - destinationY);
+    const rowBytes = this.width * 4;
+    for (let row = 0; row < rowsToCopy; row += 1) {
+      const sourceStart = row * rowBytes;
+      tileBitmap.copy(
+        this.result,
+        (destinationY + row) * rowBytes,
+        sourceStart,
+        sourceStart + rowBytes
+      );
+    }
+  }
+
+  toImage() {
+    return this.result === undefined
+      ? undefined
+      : nativeImage.createFromBitmap(this.result, {width: this.width, height: this.height});
+  }
+}

desktop-app/src/main/screenshot/capture-full-page.ts

@@ -0,0 +1,91 @@
+import {BitmapStitcher} from './bitmap-stitcher';
+import {captureViewport, delay} from './capture-viewport';
+
+const SCROLL_PAINT_DELAY_MS = 100;
+
+interface PageMetrics {
+  height: number;
+  viewportHeight: number;
+  scrollX: number;
+  scrollY: number;
+}
+
+const getPageMetrics = async (contents: Electron.WebContents) =>
+  contents.executeJavaScript(`({
+    height: Math.max(
+      document.body?.scrollHeight ?? 0,
+      document.documentElement.scrollHeight,
+      document.documentElement.offsetHeight
+    ),
+    viewportHeight: window.innerHeight,
+    scrollX: window.scrollX,
+    scrollY: window.scrollY
+  })`) as Promise<PageMetrics>;
+
+const getCapturePositions = ({height, viewportHeight}: PageMetrics) => {
+  const maxScrollY = Math.max(0, height - viewportHeight);
+  const positions: number[] = [];
+  for (let y = 0; y < maxScrollY; y += viewportHeight) {
+    positions.push(y);
+  }
+  positions.push(maxScrollY);
+  return positions;
+};
+
+const beginCapture = (contents: Electron.WebContents) =>
+  contents.executeJavaScript(`
+    window.__responsivelyScreenshotInProgress = true;
+    window.__responsivelyScreenshotStyleRestore = {
+      html: document.documentElement.getAttribute('style'),
+      body: document.body?.getAttribute('style') ?? null
+    };
+    document.documentElement.style.setProperty('scroll-behavior', 'auto', 'important');
+    document.documentElement.style.setProperty('scroll-snap-type', 'none', 'important');
+    document.body?.style.setProperty('scroll-behavior', 'auto', 'important');
+    document.body?.style.setProperty('scroll-snap-type', 'none', 'important');
+  `);
+
+const endCapture = (contents: Electron.WebContents, metrics: PageMetrics) =>
+  contents.executeJavaScript(`
+    const restore = window.__responsivelyScreenshotStyleRestore;
+    if (restore?.html === null) document.documentElement.removeAttribute('style');
+    else if (restore) document.documentElement.setAttribute('style', restore.html);
+    if (document.body) {
+      if (restore?.body === null) document.body.removeAttribute('style');
+      else if (restore) document.body.setAttribute('style', restore.body);
+    }
+    delete window.__responsivelyScreenshotStyleRestore;
+    window.scrollTo({left: ${metrics.scrollX}, top: ${metrics.scrollY}, behavior: 'instant'});
+    window.__responsivelyScreenshotInProgress = false;
+  `);
+
+const scrollTo = (contents: Electron.WebContents, y: number) =>
+  contents.executeJavaScript(`new Promise((resolve) => {
+    window.scrollTo({left: 0, top: ${y}, behavior: 'instant'});
+    requestAnimationFrame(() => requestAnimationFrame(() => resolve(window.scrollY)));
+  })`) as Promise<number>;
+
+export const captureFullPage = async (contents: Electron.WebContents) => {
+  const metrics = await getPageMetrics(contents);
+  const stitcher = new BitmapStitcher(metrics.height, metrics.viewportHeight);
+
+  await beginCapture(contents);
+  try {
+    for (const requestedY of getCapturePositions(metrics)) {
+      const actualY = await scrollTo(contents, requestedY);
+      // The DOM observes scrolling before the guest compositor presents it.
+      // Force a repaint so capturePage does not return the previous tile.
+      contents.invalidate();
+      await delay(SCROLL_PAINT_DELAY_MS);
+      const tile = await captureViewport(contents);
+      if (tile === undefined) {
+        return undefined;
+      }
+      stitcher.append(tile, actualY);
+    }
+  } finally {
+    await endCapture(contents, metrics);
+  }
+
+  return stitcher.toImage();
+};

desktop-app/src/main/screenshot/capture-viewport.ts

@@ -0,0 +1,30 @@
+const CAPTURE_ATTEMPTS = 3;
+const CAPTURE_RETRY_DELAY_MS = 250;
+
+export const delay = (ms: number) =>
+  new Promise<void>((resolve) => {
+    setTimeout(resolve, ms);
+  });
+
+export const captureViewport = async (contents: Electron.WebContents) => {
+  let lastError: unknown;
+  let image: Electron.NativeImage | undefined;
+  for (let attempt = 0; attempt < CAPTURE_ATTEMPTS; attempt += 1) {
+    if (attempt > 0) {
+      await delay(CAPTURE_RETRY_DELAY_MS);
+    }
+    try {
+      image = await contents.capturePage();
+      lastError = undefined;
+      if (!image.isEmpty()) {
+        return image;
+      }
+    } catch (error) {
+      lastError = error;
+    }
+  }
+  if (lastError !== undefined) {
+    throw lastError;
+  }
+  return image;
+};

desktop-app/src/main/screenshot/index.ts

@@ -5,50 +5,42 @@ import path from 'path';
 import {IPC_MAIN_CHANNELS} from '../../common/constants';
 import store from '../../store';
 import {isRegisteredWebview} from '../webview-registry';
+import {captureFullPage} from './capture-full-page';
+import {captureViewport} from './capture-viewport';
 
 export interface ScreenshotArgs {
   webContentsId: number;
   fullPage?: boolean;
   device: Device;
 }
 
-export interface ScreenshotAllArgs {
-  webContentsId: number;
-  device: Device;
-  previousHeight: string;
-  previousTransform: string;
-  pageHeight: number;
-}
+export type ScreenshotAllArgs = ScreenshotArgs;
 
 export interface ScreenshotResult {
   done: boolean;
 }
 
-const CAPTURE_ATTEMPTS = 3;
-const CAPTURE_RETRY_DELAY_MS = 250;
-
-const delay = (ms: number) =>
-  new Promise<void>((resolve) => {
-    setTimeout(resolve, ms);
-  });
-
 export const captureImage = async (
-  webContentsId: number
+  webContentsId: number,
+  fullPage = false
 ): Promise<Electron.NativeImage | undefined> => {
   // Single choke point for both the IPC and MCP screenshot paths.
   if (!isRegisteredWebview(webContentsId)) {
     return undefined;
   }
   const WebContents = webContents.fromId(webContentsId);
+  if (WebContents === undefined) {
+    return undefined;
+  }
 
-  const isExecuted = await WebContents?.executeJavaScript(`
+  const isExecuted = await WebContents.executeJavaScript(`
     if (window.isExecuted) {
       true;
     }
   `);
 
   if (!isExecuted) {
-    await WebContents?.executeJavaScript(`
+    await WebContents.executeJavaScript(`
       const bgColor = window.getComputedStyle(document.body).backgroundColor;
       if (bgColor === 'rgba(0, 0, 0, 0)') {
         document.body.style.backgroundColor = 'white';
@@ -57,38 +49,23 @@ export const captureImage = async (
     `);
   }
 
+  if (fullPage) {
+    return captureFullPage(WebContents);
+  }
+
   // capturePage throws (e.g. UnknownVizError) or hands back an empty frame
   // while the guest's compositor surface is still settling — typically right
   // after a navigation or a freshly attached preview. That state is
   // transient, so retry briefly before reporting the capture as failed.
-  let lastError: unknown;
-  let image: Electron.NativeImage | undefined;
-  for (let attempt = 0; attempt < CAPTURE_ATTEMPTS; attempt += 1) {
-    if (attempt > 0) {
-      await delay(CAPTURE_RETRY_DELAY_MS);
-    }
-    try {
-      image = await WebContents?.capturePage();
-      lastError = undefined;
-      if (image !== undefined && !image.isEmpty()) {
-        return image;
-      }
-    } catch (error) {
-      lastError = error;
-    }
-  }
-  if (lastError !== undefined) {
-    throw lastError;
-  }
-  return image;
+  return captureViewport(WebContents);
 };
 
 const quickScreenshot = async (arg: ScreenshotArgs): Promise<ScreenshotResult> => {
   const {
     webContentsId,
     device: {name},
   } = arg;
-  const image = await captureImage(webContentsId);
+  const image = await captureImage(webContentsId, arg.fullPage);
   if (image === undefined) {
     return {done: false};
   }
@@ -110,13 +87,11 @@ const quickScreenshot = async (arg: ScreenshotArgs): Promise<ScreenshotResult> =
 };
 
 const captureAllDecies = async (args: Array<ScreenshotAllArgs>): Promise<ScreenshotResult> => {
-  const screenShots = args.map((arg) => {
-    const {device, webContentsId} = arg;
-    const screenShotArg: ScreenshotArgs = {device, webContentsId};
-    return quickScreenshot(screenShotArg);
-  });
-
-  await Promise.all(screenShots);
+  // Full-page capture scrolls each guest while stitching its viewport. Capture
+  // sequentially so mirrored scrolling cannot make previews race each other.
+  for (const arg of args) {
+    await quickScreenshot(arg);
+  }
   return {done: true};
 };
 

desktop-app/src/renderer/components/ToolBar/index.tsx

@@ -12,11 +12,7 @@ import {Icon} from '@iconify/react';
 import {ScreenshotAllArgs} from 'main/screenshot';
 import {selectActiveSuite} from 'renderer/store/features/device-manager';
 import {getDevicesMap} from 'common/deviceList';
-import {
-  FULL_PAGE_SETTLE_MS,
-  prepareFullPageCapture,
-  useShutterSound,
-} from 'renderer/hooks/useScreenshot';
+import {useShutterSound} from 'renderer/hooks/useScreenshot';
 import {APP_VIEWS, setAppView} from 'renderer/store/features/ui';
 import NavigationControls from './NavigationControls';
 import Menu from './Menu';
@@ -51,37 +47,22 @@ const ToolBar = () => {
     dispatch(setIsCapturingScreenshot(true));
     const webViews: NodeListOf<Electron.WebviewTag> = document.querySelectorAll('webView');
     const screens: Array<ScreenshotAllArgs> = [];
-    const restores: Array<() => void> = [];
     const devices = activeSuite.devices.map((d) => getDevicesMap()[d]);
-    // Sequential await: every webview must be measured and resized before the
-    // capture fires (the old forEach(async) raced the capture call).
     for (const webview of Array.from(webViews)) {
       const device = devices.find((d) => d.name === webview.id);
       if (device != null) {
-        const prep = await prepareFullPageCapture(webview);
         screens.push({
           webContentsId: webview.getWebContentsId(),
           device,
-          previousHeight: prep.previousHeight,
-          previousTransform: prep.previousTransform,
-          pageHeight: prep.pageHeight,
+          fullPage: true,
         });
-        restores.push(prep.restore);
       }
     }
-    await new Promise((resolve) => {
-      setTimeout(resolve, FULL_PAGE_SETTLE_MS);
-    });
     await window.electron.ipcRenderer.invoke<Array<ScreenshotAllArgs>, unknown>(
       IPC_MAIN_CHANNELS.SCREENSHOT_ALL,
       screens
     );
 
-    // reset webviews to original size
-    restores.forEach((restore) => {
-      restore();
-    });
-
     dispatch(setIsCapturingScreenshot(false));
     playShutter();
   };