Merge pull request #771 from heygen-com/feat/engine-discard-warmup-capture

feat(engine): first-frame warmup capture helper for distributed chunks

packages/engine/src/index.ts

@@ -65,11 +65,13 @@ export {
   closeCaptureSession,
   captureFrame,
   captureFrameToBuffer,
+  discardWarmupCapture,
   getCompositionDuration,
   getCapturePerfSummary,
   prepareCaptureSessionForReuse,
   type CaptureSession,
   type BeforeCaptureHook,
+  type DiscardWarmupInnerCapture,
 } from "./services/frameCapture.js";
 
 // ── Screenshot (BeginFrame) ─────────────────────────────────────────────────────

packages/engine/src/services/frameCapture-discardWarmup.test.ts

@@ -0,0 +1,183 @@
+/**
+ * Tests for `discardWarmupCapture` — the helper distributed chunk workers
+ * run before their first real capture to prime `lastFrameCache`.
+ *
+ * The helper is a thin wrapper around the inner `captureFrameCore`
+ * machinery, so its testable contract is post-conditional rather than
+ * pixel-level:
+ *
+ *   1. The wrapper invokes the inner capture exactly once with the supplied
+ *      `(frameIndex, time)`.
+ *   2. After the wrapper returns, the session's perf and BeginFrame damage
+ *      counters look exactly as they did before — even though the inner
+ *      capture mutated them.
+ *   3. The wrapper writes no file to disk (no path is plumbed through;
+ *      asserted indirectly by observing that `outputDir` is never read).
+ *   4. State is restored even if the inner capture throws.
+ *
+ * The inner capture is stubbed (the helper accepts an injectable
+ * `innerCapture` for exactly this reason). We don't need a real Chrome.
+ */
+
+import { describe, expect, it } from "vitest";
+import { existsSync, readdirSync, mkdtempSync, rmSync } from "node:fs";
+import { join } from "node:path";
+import { tmpdir } from "node:os";
+import { discardWarmupCapture, type CaptureSession } from "./frameCapture.js";
+
+function makeFakeSession(): CaptureSession {
+  // The discardWarmupCapture wrapper only reads `capturePerf`,
+  // `beginFrameHasDamageCount`, `beginFrameNoDamageCount`. Everything else
+  // is unused — leave it as bare-minimum stubs cast through `unknown`.
+  return {
+    browser: {} as unknown,
+    page: {} as unknown,
+    options: { width: 1, height: 1, fps: { num: 30, den: 1 } },
+    serverUrl: "",
+    outputDir: mkdtempSync(join(tmpdir(), "discard-warmup-")),
+    onBeforeCapture: null,
+    isInitialized: true,
+    browserConsoleBuffer: [],
+    capturePerf: {
+      frames: 7,
+      seekMs: 100,
+      beforeCaptureMs: 50,
+      screenshotMs: 200,
+      totalMs: 350,
+    },
+    captureMode: "beginframe",
+    beginFrameTimeTicks: 0,
+    beginFrameIntervalMs: 33,
+    beginFrameHasDamageCount: 4,
+    beginFrameNoDamageCount: 3,
+  } as unknown as CaptureSession;
+}
+
+describe("discardWarmupCapture", () => {
+  it("calls the inner capture exactly once with (frameIndex=0, time=0) by default", async () => {
+    const session = makeFakeSession();
+    try {
+      let calls = 0;
+      let receivedFrameIndex = -1;
+      let receivedTime = -1;
+      await discardWarmupCapture(session, undefined, undefined, async (_s, fi, t) => {
+        calls++;
+        receivedFrameIndex = fi;
+        receivedTime = t;
+        return { buffer: Buffer.alloc(0), quantizedTime: t, captureTimeMs: 0 };
+      });
+      expect(calls).toBe(1);
+      expect(receivedFrameIndex).toBe(0);
+      expect(receivedTime).toBe(0);
+    } finally {
+      rmSync(session.outputDir, { recursive: true, force: true });
+    }
+  });
+
+  it("passes through caller-supplied (frameIndex, time)", async () => {
+    const session = makeFakeSession();
+    try {
+      let received: { fi: number; t: number } | null = null;
+      await discardWarmupCapture(session, 240, 8, async (_s, fi, t) => {
+        received = { fi, t };
+        return { buffer: Buffer.alloc(0), quantizedTime: t, captureTimeMs: 0 };
+      });
+      expect(received).toEqual({ fi: 240, t: 8 });
+    } finally {
+      rmSync(session.outputDir, { recursive: true, force: true });
+    }
+  });
+
+  it("restores perf counters after the inner capture mutates them", async () => {
+    const session = makeFakeSession();
+    const before = { ...session.capturePerf };
+    try {
+      await discardWarmupCapture(session, 0, 0, async (s) => {
+        s.capturePerf.frames += 1;
+        s.capturePerf.seekMs += 12;
+        s.capturePerf.beforeCaptureMs += 5;
+        s.capturePerf.screenshotMs += 33;
+        s.capturePerf.totalMs += 50;
+        return { buffer: Buffer.alloc(0), quantizedTime: 0, captureTimeMs: 50 };
+      });
+      expect(session.capturePerf).toEqual(before);
+    } finally {
+      rmSync(session.outputDir, { recursive: true, force: true });
+    }
+  });
+
+  it("restores BeginFrame damage counters after the inner capture mutates them", async () => {
+    const session = makeFakeSession();
+    const hasBefore = session.beginFrameHasDamageCount;
+    const noBefore = session.beginFrameNoDamageCount;
+    try {
+      await discardWarmupCapture(session, 0, 0, async (s) => {
+        s.beginFrameHasDamageCount += 10;
+        s.beginFrameNoDamageCount += 1;
+        return { buffer: Buffer.alloc(0), quantizedTime: 0, captureTimeMs: 0 };
+      });
+      expect(session.beginFrameHasDamageCount).toBe(hasBefore);
+      expect(session.beginFrameNoDamageCount).toBe(noBefore);
+    } finally {
+      rmSync(session.outputDir, { recursive: true, force: true });
+    }
+  });
+
+  it("restores state even when the inner capture throws", async () => {
+    const session = makeFakeSession();
+    const perfBefore = { ...session.capturePerf };
+    const hasBefore = session.beginFrameHasDamageCount;
+    const noBefore = session.beginFrameNoDamageCount;
+    try {
+      let thrown: unknown;
+      try {
+        await discardWarmupCapture(session, 0, 0, async (s) => {
+          s.capturePerf.frames += 5;
+          s.beginFrameNoDamageCount += 2;
+          throw new Error("simulated capture failure");
+        });
+      } catch (err) {
+        thrown = err;
+      }
+      expect((thrown as Error).message).toBe("simulated capture failure");
+      // The whole point of `finally { restore }`: failure must not leak
+      // inflated counters into the real capture summary.
+      expect(session.capturePerf).toEqual(perfBefore);
+      expect(session.beginFrameHasDamageCount).toBe(hasBefore);
+      expect(session.beginFrameNoDamageCount).toBe(noBefore);
+    } finally {
+      rmSync(session.outputDir, { recursive: true, force: true });
+    }
+  });
+
+  it("writes no output file to the session's outputDir", async () => {
+    const session = makeFakeSession();
+    try {
+      expect(existsSync(session.outputDir)).toBe(true);
+      const before = readdirSync(session.outputDir);
+      await discardWarmupCapture(session, 0, 0, async () => ({
+        buffer: Buffer.from([0xff, 0xff, 0xff]),
+        quantizedTime: 0,
+        captureTimeMs: 0,
+      }));
+      const after = readdirSync(session.outputDir);
+      expect(after).toEqual(before);
+    } finally {
+      rmSync(session.outputDir, { recursive: true, force: true });
+    }
+  });
+
+  it("returns undefined (no result type, so the buffer can't escape)", async () => {
+    const session = makeFakeSession();
+    try {
+      const result = await discardWarmupCapture(session, 0, 0, async () => ({
+        buffer: Buffer.from([0x01]),
+        quantizedTime: 0,
+        captureTimeMs: 1,
+      }));
+      expect(result).toBeUndefined();
+    } finally {
+      rmSync(session.outputDir, { recursive: true, force: true });
+    }
+  });
+});

packages/engine/src/services/frameCapture.ts

@@ -809,6 +809,69 @@ export async function captureFrameToBuffer(
   return { buffer, captureTimeMs };
 }
 
+/**
+ * Type of the "inner capture" function consumed by
+ * {@link discardWarmupCapture}. Matches the real `captureFrameCore` signature
+ * with the buffer-bearing result trimmed to what the caller actually uses
+ * (the wrapper never inspects the buffer). Exposed so unit tests can inject
+ * a stub instead of driving Chrome end-to-end.
+ */
+export type DiscardWarmupInnerCapture = (
+  session: CaptureSession,
+  frameIndex: number,
+  time: number,
+) => Promise<{ buffer: Buffer; quantizedTime: number; captureTimeMs: number }>;
+
+/**
+ * Perform one capture, throw away the buffer, and restore any session
+ * side-effects (perf counters, BeginFrame damage tallies) so downstream
+ * captures see state identical to a fresh session.
+ *
+ * Distributed chunk workers need this because Chrome's BeginFrame screenshot
+ * pipeline maintains a per-process `lastFrameCache`: when a captured frame's
+ * `hasDamage` reports `false`, the screenshot path returns the previously
+ * captured buffer. For chunk N (N > 0) the worker has no prior frame in its
+ * cache, so the very first capture's `hasDamage` reporting diverges from
+ * what an in-process render at the same absolute frame index would see (the
+ * in-process renderer always has frame N-1 cached). One discard capture
+ * before the first real capture primes the cache.
+ *
+ * The function intentionally restores perf state so the warmup capture does
+ * NOT bias `getCapturePerfSummary()`'s per-frame averages.
+ *
+ * No file is written; the buffer is discarded.
+ *
+ * @param session — initialized capture session
+ * @param frameIndex — frame index to warm up with (default 0). Chunk
+ *   workers typically pass their chunk's first absolute frame index.
+ * @param time — time in seconds (default 0). Chunk workers typically pass
+ *   the corresponding `frameIndex / fps`.
+ * @param innerCapture — injectable for tests; defaults to the real
+ *   `captureFrameCore`.
+ */
+export async function discardWarmupCapture(
+  session: CaptureSession,
+  frameIndex: number = 0,
+  time: number = 0,
+  innerCapture: DiscardWarmupInnerCapture = captureFrameCore,
+): Promise<void> {
+  // Snapshot the side-effect counters captureFrameCore mutates. We use a
+  // shallow `{...}` for capturePerf because all five fields are primitive
+  // numbers — no nested state to deep-copy.
+  const perfBefore = { ...session.capturePerf };
+  const hasDamageBefore = session.beginFrameHasDamageCount;
+  const noDamageBefore = session.beginFrameNoDamageCount;
+  try {
+    await innerCapture(session, frameIndex, time);
+  } finally {
+    // Always restore — even on error. A failed warmup capture should not
+    // leak inflated perf counters into the real capture summary.
+    session.capturePerf = perfBefore;
+    session.beginFrameHasDamageCount = hasDamageBefore;
+    session.beginFrameNoDamageCount = noDamageBefore;
+  }
+}
+
 export async function closeCaptureSession(session: CaptureSession): Promise<void> {
   // INVARIANT: closeCaptureSession is idempotent. The renderOrchestrator HDR
   // cleanup path tracks a `domSessionClosed` flag and may still re-call this