fix(cli): snapshot shoots a frame again when Chrome drops the screenshot mid-paint (#6941)

* fix(cli): snapshot shoots a frame again when Chrome drops the screenshot mid-paint

* fix(cli): keep the first screenshot error when the retry fails differently

packages/cli/src/commands/snapshot.screenshotRetry.test.ts

@@ -0,0 +1,37 @@
+import { describe, expect, it, vi } from "vitest";
+import { screenshotOnceMore } from "./snapshot.js";
+
+const dropped = () =>
+  new Error("Protocol error (Page.captureScreenshot): Unable to capture screenshot");
+
+describe("screenshotOnceMore", () => {
+  it("shoots again when Chrome drops the screenshot mid-paint", async () => {
+    const shoot = vi.fn().mockRejectedValueOnce(dropped()).mockResolvedValueOnce("png");
+    await expect(screenshotOnceMore(shoot)).resolves.toBe("png");
+    expect(shoot).toHaveBeenCalledTimes(2);
+  });
+
+  it("gives up after the second drop", async () => {
+    const shoot = vi.fn().mockRejectedValue(dropped());
+    await expect(screenshotOnceMore(shoot)).rejects.toThrow(/Unable to capture screenshot/);
+    expect(shoot).toHaveBeenCalledTimes(2);
+  });
+
+  it("reports the first drop when the second try fails differently", async () => {
+    const first = dropped();
+    const shoot = vi
+      .fn()
+      .mockRejectedValueOnce(first)
+      .mockRejectedValueOnce(new Error("Target closed"));
+    await expect(screenshotOnceMore(shoot)).rejects.toBe(first);
+  });
+
+  it("does not retry a screenshot that fails the same way every time", async () => {
+    const zero = new Error(
+      "Protocol error (Page.captureScreenshot): Cannot take screenshot with 0 width.",
+    );
+    const shoot = vi.fn().mockRejectedValue(zero);
+    await expect(screenshotOnceMore(shoot)).rejects.toBe(zero);
+    expect(shoot).toHaveBeenCalledTimes(1);
+  });
+});

packages/cli/src/commands/snapshot.test.ts

@@ -142,6 +142,13 @@ describe("transparent snapshot capture", () => {
     expect(source).toContain('page.screenshot({ type: "png", omitBackground: true })');
   });
 
+  it("shoots each frame through the one-retry wrapper", () => {
+    const source = readFileSync(new URL("./snapshot.ts", import.meta.url), "utf8");
+    expect(source).toContain(
+      'screenshotOnceMore(() => page.screenshot({ type: "png", omitBackground: true }))',
+    );
+  });
+
   it("exposes --proxy and forwards it to the snapshot page server", () => {
     const source = readFileSync(new URL("./snapshot.ts", import.meta.url), "utf8");
     expect(source).toContain("proxy: {");

packages/cli/src/commands/snapshot.ts

@@ -1,6 +1,7 @@
 import {
   DEFAULT_CONFIG,
   buildZeroDurationDiagnostic,
+  isTransientBrowserError,
   downloadToTemp,
   evaluateHfDiagnostic,
   extractMediaMetadata,
@@ -342,6 +343,17 @@ interface SnapshotCompositePage {
   }): Promise<Uint8Array>;
 }
 
+export async function screenshotOnceMore<T>(shoot: () => Promise<T>): Promise<T> {
+  try {
+    return await shoot();
+  } catch (err) {
+    if (!isTransientBrowserError(err)) throw err;
+    return shoot().catch(() => {
+      throw err;
+    });
+  }
+}
+
 /** Recapture canvas composites once the decoded video overlays have been injected. */
 export async function recaptureSnapshotComposite(page: SnapshotCompositePage): Promise<void> {
   const hasPageComposite = await page.evaluate(async () => {
@@ -1371,7 +1383,7 @@ export async function captureSnapshots(
           writeCaptureFileSync(
             framePath,
             await atStep(`capturing the frame at ${time}s`, () =>
-              page.screenshot({ type: "png", omitBackground: true }),
+              screenshotOnceMore(() => page.screenshot({ type: "png", omitBackground: true })),
             ),
           );
         }