[Bugfix] Stop caching synthetic videos without requested audio

tests/helpers/media.py

@@ -20,6 +20,7 @@
 import tempfile
 import threading
 import time
+import uuid
 from concurrent.futures.process import BrokenProcessPool
 from contextlib import contextmanager
 from pathlib import Path
@@ -308,16 +309,12 @@ def _mux_mp4_bytes_with_synthetic_audio(
     duration_sec = num_frames / fps if fps > 0 else 0.0
     duration_int = max(1, int(math.ceil(duration_sec)))
 
-    try:
-        audio_result = generate_synthetic_audio(
-            duration=duration_int,
-            num_channels=1,
-            sample_rate=sample_rate,
-        )
-        audio_pcm = audio_result["np_array"]
-    except Exception as e:
-        logger.warning("Synthetic video: generate_synthetic_audio failed (%s); using video-only MP4.", e)
-        return video_mp4_bytes
+    audio_result = generate_synthetic_audio(
+        duration=duration_int,
+        num_channels=1,
+        sample_rate=sample_rate,
+    )
+    audio_pcm = audio_result["np_array"]
 
     try:
         import imageio_ffmpeg
@@ -326,47 +323,38 @@ def _mux_mp4_bytes_with_synthetic_audio(
     except Exception:
         ffmpeg_exe = "ffmpeg"
 
-    try:
-        with tempfile.TemporaryDirectory(prefix="syn_vid_mux_") as tmp:
-            vid_path = os.path.join(tmp, "video.mp4")
-            wav_path = os.path.join(tmp, "audio.wav")
-            out_path = os.path.join(tmp, "out.mp4")
-            with open(vid_path, "wb") as f:
-                f.write(video_mp4_bytes)
-            sf.write(wav_path, audio_pcm, sample_rate, format="WAV", subtype="PCM_16")
-            cmd = [
-                ffmpeg_exe,
-                "-y",
-                "-nostdin",
-                "-hide_banner",
-                "-loglevel",
-                "error",
-                "-i",
-                vid_path,
-                "-i",
-                wav_path,
-                "-c:v",
-                "copy",
-                "-c:a",
-                "aac",
-                "-b:a",
-                "128k",
-                "-shortest",
-                "-movflags",
-                "+faststart",
-                out_path,
-            ]
-            subprocess.run(cmd, check=True, stdin=subprocess.DEVNULL, timeout=300)
-            with open(out_path, "rb") as f:
-                return f.read()
-    except (
-        FileNotFoundError,
-        subprocess.CalledProcessError,
-        subprocess.TimeoutExpired,
-        OSError,
-    ) as e:
-        logger.warning("Synthetic video: audio mux failed (%s); using video-only MP4.", e)
-        return video_mp4_bytes
+    with tempfile.TemporaryDirectory(prefix="syn_vid_mux_") as tmp:
+        vid_path = os.path.join(tmp, "video.mp4")
+        wav_path = os.path.join(tmp, "audio.wav")
+        out_path = os.path.join(tmp, "out.mp4")
+        with open(vid_path, "wb") as f:
+            f.write(video_mp4_bytes)
+        sf.write(wav_path, audio_pcm, sample_rate, format="WAV", subtype="PCM_16")
+        cmd = [
+            ffmpeg_exe,
+            "-y",
+            "-nostdin",
+            "-hide_banner",
+            "-loglevel",
+            "error",
+            "-i",
+            vid_path,
+            "-i",
+            wav_path,
+            "-c:v",
+            "copy",
+            "-c:a",
+            "aac",
+            "-b:a",
+            "128k",
+            "-shortest",
+            "-movflags",
+            "+faststart",
+            out_path,
+        ]
+        subprocess.run(cmd, check=True, stdin=subprocess.DEVNULL, timeout=300)
+        with open(out_path, "rb") as f:
+            return f.read()
 
 
 def generate_synthetic_video(
@@ -383,18 +371,49 @@ def generate_synthetic_video(
     ``width`` / ``height`` / ``num_frames`` / ``embed_audio`` unless
     ``force_regenerate`` is true. Cache root: ``cache_dir`` if given, else the
     default temp subdirectory.
+
+    ``embed_audio=True`` requires successful TTS generation and FFmpeg muxing.
+    Cached MP4s with corrupt/incomplete media or without the requested audio
+    stream are regenerated. Cache files are published atomically.
     """
     root = _resolve_synthetic_media_cache_dir(cache_dir)
     root.mkdir(parents=True, exist_ok=True)
     cache_path = root / f"synth_video_w{width}_h{height}_nf{num_frames}_ea{int(embed_audio)}.mp4"
 
     if not force_regenerate and cache_path.is_file():
         video_bytes = cache_path.read_bytes()
-        return {
-            "np_array": _np_array_from_mp4_bytes(video_bytes),
-            "base64": base64.b64encode(video_bytes).decode("utf-8"),
-            "file_path": str(cache_path.resolve()),
-        }
+        has_required_audio = True
+        if embed_audio:
+            import av
+
+            try:
+                with av.open(io.BytesIO(video_bytes)) as container:
+                    has_required_audio = False
+                    if container.streams.audio:
+                        audio_stream = container.streams.audio[0]
+                        samples = sum(frame.samples for frame in container.decode(audio=0))
+                        expected_samples = 0
+                        if audio_stream.duration is not None and audio_stream.time_base is not None:
+                            expected_samples = math.ceil(
+                                audio_stream.duration * audio_stream.time_base * audio_stream.codec_context.sample_rate
+                            )
+                        has_required_audio = samples > 0 and samples >= expected_samples
+            except av.error.FFmpegError:
+                has_required_audio = False
+        if has_required_audio:
+            try:
+                cached_frames = _np_array_from_mp4_bytes(video_bytes)
+            except RuntimeError:
+                cached_frames = None
+            if cached_frames is not None and len(cached_frames) == num_frames:
+                return {
+                    "np_array": cached_frames,
+                    "base64": base64.b64encode(video_bytes).decode("utf-8"),
+                    "file_path": str(cache_path.resolve()),
+                }
+        logger.warning(
+            "Cached synthetic video %s is incomplete or lacks the requested audio; regenerating.", cache_path
+        )
 
     import cv2
     import imageio
@@ -475,7 +494,15 @@ class _BouncingBall(TypedDict):
         else video_only_bytes
     )
 
-    cache_path.write_bytes(video_bytes)
+    temp_cache_path = root / f".{cache_path.name}.{uuid.uuid4().hex}.tmp"
+    # Exclusive creation preserves the normal 0666 & ~umask cache permissions.
+    temp_cache = temp_cache_path.open("xb")
+    try:
+        with temp_cache:
+            temp_cache.write(video_bytes)
+        os.replace(temp_cache_path, cache_path)
+    finally:
+        temp_cache_path.unlink(missing_ok=True)
 
     return {
         "np_array": np.array(video_frames),