fix(aws-lambda): make webm and mov audio in the separate audio task like mp4 (#6098)

* fix(aws-lambda): make webm and mov audio in the separate audio task like mp4

A WebM or MOV plan mixed and encoded its own audio, so on a stack whose plan runs on an arm64 function the audio came out of a different ffmpeg build than an x86 render's. The plan now leaves that audio to the separate audio task, as mp4 already does, and the assembler joins the chunks with that file. Only mp4 chunks stitch; cfr and png-sequence keep their audio in the plan.

* test(aws-lambda): join mov chunks with the audio task's track too

* refactor(producer): keep the separate-audio gate's one-line doc comment

packages/aws-lambda/src/handler.test.ts

@@ -1099,6 +1099,7 @@ describe("handler dispatch", () => {
     chunkCount: number,
     separateAudio = false,
     workerConcurrency?: number,
+    format: "mp4" | "webm" | "mov" = "mp4",
   ) {
     const tmpRoot = makeTmpRoot();
     const s3 = new FakeS3Client();
@@ -1114,8 +1115,9 @@ describe("handler dispatch", () => {
         const v1Dir = join(tmpRoot, "v1");
         makeMinimalV1PlanDir(v1Dir, !separateAudio, {
           chunkCount,
-          chunksCarryAudio: !separateAudio,
+          chunksCarryAudio: !separateAudio && format === "mp4",
           separateAudio,
+          format,
         });
         return publishPlanV2FromExecutionPlan(v1Dir, publisher);
       },
@@ -1136,7 +1138,7 @@ describe("handler dispatch", () => {
         Action: "plan",
         ProjectS3Uri: "s3://bucket/project.tar.gz",
         PlanOutputS3Prefix: "s3://bucket/renders/f/",
-        Config: { fps: 30, width: 640, height: 360, format: "mp4", workerConcurrency },
+        Config: { fps: 30, width: 640, height: 360, format, workerConcurrency },
       },
       deps,
     );
@@ -1151,8 +1153,8 @@ describe("handler dispatch", () => {
           PlanHash: planned.PlanHash,
           ChunkS3Uris: chunkS3Uris,
           AudioS3Uri: audioS3Uri,
-          OutputS3Uri: "s3://bucket/renders/f/output.mp4",
-          Format: "mp4",
+          OutputS3Uri: `s3://bucket/renders/f/output.${format}`,
+          Format: format,
         },
         deps,
       );
@@ -1256,6 +1258,30 @@ describe("handler dispatch", () => {
     },
   );
 
+  it.each(["webm", "mov"] as const)(
+    "joins %s chunks with the audio task's track when the plan left its audio out",
+    async (format) => {
+      const { s3, assembleMock, assemble, planned } = await planStitched(2, true, 4, format);
+      expect(planned.SeparateAudio).toBe(true);
+      const audioUri = "s3://bucket/renders/f/audio/audio.m4a";
+      s3.objects.set(audioUri, Buffer.from("AAC-FROM-AUDIO-TASK"));
+      let audio: string | null = null;
+      assembleMock.mockImplementationOnce(async (_p, _c, audioPath, outputPath) => {
+        audio = audioPath && readFileSync(audioPath, "utf-8");
+        writeFileSync(outputPath, "LOCAL-JOIN");
+        return { outputPath, durationMs: 1, framesEncoded: 60, fileSize: 10 };
+      });
+      const chunks = [0, 1].map((index) => {
+        const uri = `s3://bucket/renders/f/chunks/000${index}.${format}`;
+        s3.objects.set(uri, Buffer.from("CHUNK"));
+        return uri;
+      });
+      await assemble(chunks, audioUri);
+      expect(assembleMock).toHaveBeenCalledTimes(1);
+      expect(audio).toBe("AAC-FROM-AUDIO-TASK");
+    },
+  );
+
   it("refuses an audio URI for a v2 plan that ships its own audio", async () => {
     const { assemble } = await planStitched(1);
     await expect(
@@ -1366,12 +1392,24 @@ describe("handler — audio", () => {
     expect(buildPlanAudio).toHaveBeenCalledTimes(1);
   });
 
-  it("leaves the audio to the plan when the stitch cannot take it (cfr)", async () => {
-    const { result, s3, buildPlanAudio } = await runAudio({ cfr: true });
-    expect(result).toMatchObject({ Action: "audio", HasAudio: false, AudioS3Uri: null });
-    expect(buildPlanAudio).not.toHaveBeenCalled();
-    expect(s3.ops).toHaveLength(0);
-  });
+  it.each(["webm", "mov"])(
+    "mixes %s audio too, so one function makes every audio byte",
+    async (format) => {
+      const { result, buildPlanAudio } = await runAudio({ format });
+      expect(result).toMatchObject({ Action: "audio", HasAudio: true });
+      expect(buildPlanAudio).toHaveBeenCalledTimes(1);
+    },
+  );
+
+  it.each([{ cfr: true }, { format: "png-sequence" }])(
+    "leaves the audio to the plan for %o",
+    async (config) => {
+      const { result, s3, buildPlanAudio } = await runAudio(config);
+      expect(result).toMatchObject({ Action: "audio", HasAudio: false, AudioS3Uri: null });
+      expect(buildPlanAudio).not.toHaveBeenCalled();
+      expect(s3.ops).toHaveLength(0);
+    },
+  );
 });
 
 /** A real client (so it can pre-sign offline) whose requests go to the fake. */
@@ -1619,6 +1657,7 @@ function makeMinimalV1PlanDir(
     chunksCarryAudio?: boolean;
     separateAudio?: boolean;
     withDeferredVideo?: boolean;
+    format?: string;
   } = {},
 ): void {
   const chunkCount = options.chunkCount ?? 1;
@@ -1629,7 +1668,13 @@ function makeMinimalV1PlanDir(
     planHash: "a".repeat(64),
     chunkCount,
     totalFrames: 30 * chunkCount,
-    dimensions: { fpsNum: 30, fpsDen: 1, width: 640, height: 360, format: "mp4" },
+    dimensions: {
+      fpsNum: 30,
+      fpsDen: 1,
+      width: 640,
+      height: 360,
+      format: options.format ?? "mp4",
+    },
     ffmpegVersion: "6.0",
     producerVersion: "test",
     fontSnapshotSha: "font-snapshot-test",

packages/aws-lambda/src/handler.ts

@@ -566,7 +566,7 @@ async function handleAudio(event: AudioEvent, deps?: HandlerDeps): Promise<Audio
     AudioS3Uri: null,
     DurationMs: Date.now() - started,
   });
-  // The plan mixed it (non-mp4, cfr): the assembler reads it from the manifest.
+  // The plan mixed it (png-sequence, cfr): the assembler reads it from the manifest.
   if (!plansSeparateAudio(event.Config, { deferAudio: true })) return none();
   const s3 = deps?.s3 ?? getS3Client();
   await ensureChromePath(deps);
@@ -835,7 +835,11 @@ async function handleAssembleV2(
     }
     // `downloadAndMaterializePlanV2` materializes atomically. Audio is
     // assembler-only and lives at the familiar v1-compatible location.
-    const audioPath = resolvePlanAudioPath(planDir);
+    let audioPath = resolvePlanAudioPath(planDir);
+    if (event.AudioS3Uri !== null) {
+      audioPath = join(planDir, PLAN_AUDIO_RELATIVE_PATH);
+      await downloadS3ObjectToFile(s3, event.AudioS3Uri, audioPath);
+    }
     const chunksStarted = Date.now();
     const chunkPaths = await downloadChunkObjects(s3, event.ChunkS3Uris, work, event.Format);
     const chunksMs = Date.now() - chunksStarted;

packages/producer/src/services/distributed/assemble.ts

@@ -400,10 +400,11 @@ async function runAssembleFfmpeg(
   }
 }
 
-/** Chunks carrying audio, or with it as a separate file, stitch under one index like S3's. */
+/** Mp4 chunks carrying audio, or with it as a separate file, stitch under one index like S3's. */
 export function chunksAreStitchedMp4(planDir: string, cfr: boolean): boolean {
   const { chunksCarryAudio, separateAudio } = lockedAudio(planDir);
   if (!chunksCarryAudio && !separateAudio) return false;
+  if (planFormat(planDir) !== "mp4") return false;
   if (cfr) {
     throw new Error("[assemble] cfr=true cannot re-encode a plan whose mp4 chunks are stitched");
   }
@@ -414,6 +415,11 @@ export function stitchesSeparateAudio(planDir: string): boolean {
   return lockedAudio(planDir).separateAudio === true;
 }
 
+function planFormat(planDir: string): DistributedFormat {
+  const plan = JSON.parse(readFileSync(join(planDir, "plan.json"), "utf-8")) as PlanJsonForAssemble;
+  return plan.dimensions.format;
+}
+
 function lockedAudio(
   planDir: string,
 ): Pick<LockedRenderConfig, "chunksCarryAudio" | "separateAudio"> {

packages/producer/src/services/distributed/plan.test.ts

@@ -1112,6 +1112,25 @@ describe.skipIf(spawnSync(getFfmpegBinary(), ["-version"]).status !== 0)(
       expect(existsSync(join(withCfr, "audio.m4a"))).toBe(true);
     }, 60_000);
 
+    it("leaves webm and mov audio to buildPlanAudio too; png-sequence keeps its own", async () => {
+      const project = audioProject("project-separate-audio-formats");
+      const config = { fps: 30, width: 320, height: 240 } as const;
+      for (const [format, separate] of [
+        ["webm", true],
+        ["mov", true],
+        ["png-sequence", false],
+      ] as const) {
+        const dir = join(runRoot, `plan-separate-audio-${format}`);
+        await buildLocalExecutionPlan(project, { ...config, format }, dir, { deferAudio: true });
+        const encoder = JSON.parse(readFileSync(join(dir, "meta", "encoder.json"), "utf-8"));
+        expect({ format, separate: encoder.separateAudio === true }).toEqual({ format, separate });
+        expect({ format, mixed: existsSync(join(dir, "audio.m4a")) }).toEqual({
+          format,
+          mixed: !separate,
+        });
+      }
+    }, 60_000);
+
     it("gives buildPlanAudio the sound of a video whose src the scene script sets", async () => {
       const project = join(runRoot, "project-script-src");
       mkdirSync(project, { recursive: true });

packages/producer/src/services/distributed/plan.ts

@@ -1089,7 +1089,7 @@ export function plansSeparateAudio(
   config: Pick<DistributedRenderConfig, "format" | "cfr">,
   options: Pick<BuildLocalExecutionPlanOptions, "deferAudio">,
 ): boolean {
-  return options.deferAudio === true && config.format === "mp4" && config.cfr !== true;
+  return options.deferAudio === true && config.format !== "png-sequence" && config.cfr !== true;
 }
 
 // The plan's copy, compile and probe, shared by the plan and by buildPlanAudio so both see the