diff --git a/packages/producer/src/services/render/stages/captureSegmentedStage.test.ts b/packages/producer/src/services/render/stages/captureSegmentedStage.test.ts index 3245f16c7a..359852ae46 100644 --- a/packages/producer/src/services/render/stages/captureSegmentedStage.test.ts +++ b/packages/producer/src/services/render/stages/captureSegmentedStage.test.ts @@ -315,6 +315,37 @@ describe("runCaptureSegmentedStage", () => { }); }); + it("does not capture on a probe that left frames to a screenshot session", async () => { + closeCaptureSession.mockClear(); + const probeSession = { + ...fakeStageInput({ totalFrames: 1 }).probeSession, + captureMode: "drawelement" as const, + workerEncodeEnabled: true, + deRiskFrames: new Set([2]), + }; + const capturedOn = new Set(); + const captureFrame = mock(async (session: unknown) => { + capturedOn.add(session); + return { buffer: Buffer.alloc(1) }; + }); + const result = await runCaptureSegmentedStage({ + ...fakeStageInput({ totalFrames: 6 }), + probeSession: probeSession as never, + segmentFrames: 3, + segmentDir: join(fixtureRoot, "risk-probe"), + deps: { + spawnEncoder: mock(async () => okEncoder()), + captureFrame, + concat: mock(async () => ({ success: true as const })), + closeSession: mock(async () => {}), + removeFile: () => {}, + }, + }); + expect(result.success).toBe(true); + expect(capturedOn.has(probeSession)).toBe(false); + expect(closeCaptureSession.mock.calls.map((c) => c[0])).toContain(probeSession); + }); + it("skips completed segments and still concats all of them in order", async () => { const stableDir = join(fixtureRoot, "stable"); const spawned: string[] = []; diff --git a/packages/producer/src/services/render/stages/captureSegmentedStage.ts b/packages/producer/src/services/render/stages/captureSegmentedStage.ts index b7704916f9..a4aa740b01 100644 --- a/packages/producer/src/services/render/stages/captureSegmentedStage.ts +++ b/packages/producer/src/services/render/stages/captureSegmentedStage.ts @@ -291,6 +291,11 @@ function defaultSessionFactory(input: CaptureSegmentedStageInput): SessionFactor let probeSession = input.probeSession; return { create: async () => { + const probeCannotCaptureEveryFrame = Boolean(probeSession?.deRiskFrames?.size); + if (probeSession && probeCannotCaptureEveryFrame) { + await closeCaptureSession(probeSession).catch(() => {}); + probeSession = null; + } const session = probeSession ?? (await openSegmentedSession(input)); if (probeSession) { openSegmentedSessionReuse(input, probeSession);