Merge pull request #325 from 0xDevNinja/fix/subtitle-ts-overflow-and-checkpoint-keyerror

fix: subtitle timestamp ms overflow; checkpoint KeyError on manifest-only stages

lib/checkpoint.py

@@ -106,8 +106,17 @@ def _validate_artifacts_for_stage(
     status: str,
     artifacts: dict[str, Any],
 ) -> None:
-    required_artifact = CANONICAL_STAGE_ARTIFACTS[stage]
-    if status in {"completed", "awaiting_human"} and required_artifact not in artifacts:
+    # Valid stages come from the pipeline manifest (get_pipeline_stages), which
+    # can declare stages beyond the 9 canonical ones (e.g. character-animation's
+    # `character_design`/`rig_plan`, screen-demo's `real_capture`). Those have no
+    # canonical artifact, so look it up defensively — a missing entry means the
+    # stage simply has no required artifact, not a crash.
+    required_artifact = CANONICAL_STAGE_ARTIFACTS.get(stage)
+    if (
+        required_artifact is not None
+        and status in {"completed", "awaiting_human"}
+        and required_artifact not in artifacts
+    ):
         raise CheckpointValidationError(
             f"Stage {stage!r} with status {status!r} must include "
             f"canonical artifact {required_artifact!r}"

tests/lib/test_checkpoint_noncanonical_stage.py

@@ -0,0 +1,62 @@
+"""Regression test: checkpoint validation must not crash on manifest-only stages.
+
+`_validate_artifacts_for_stage` looked up `CANONICAL_STAGE_ARTIFACTS[stage]`
+unconditionally. Valid stages, however, come from the pipeline manifest
+(`get_pipeline_stages`), which declares stages beyond the 9 canonical ones —
+e.g. `character-animation` adds `character_design` / `rig_plan`. Those pass the
+`stage in valid_stages` guard, then raised an unhandled `KeyError` on the
+canonical lookup, so those stages could never be checkpointed. The lookup is now
+defensive (`.get`), treating a missing entry as "no required artifact".
+"""
+
+import sys
+from pathlib import Path
+
+import pytest
+
+PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent
+sys.path.insert(0, str(PROJECT_ROOT))
+
+from lib.checkpoint import (  # noqa: E402
+    CheckpointValidationError,
+    get_pipeline_stages,
+    validate_checkpoint,
+)
+
+
+def _checkpoint(stage, status, artifacts, pipeline_type):
+    return {
+        "version": "1.0",
+        "project_id": "proj",
+        "pipeline_type": pipeline_type,
+        "stage": stage,
+        "status": status,
+        "timestamp": "2026-01-01T00:00:00Z",
+        "artifacts": artifacts,
+    }
+
+
+def test_manifest_declares_noncanonical_stage():
+    # Guard the premise: the manifest really does add stages the canonical map
+    # doesn't know about.
+    stages = get_pipeline_stages("character-animation")
+    assert "character_design" in stages
+
+
+def test_noncanonical_stage_does_not_raise_keyerror():
+    # character_design has no canonical artifact; completing it with no
+    # artifacts must validate cleanly rather than crash.
+    validate_checkpoint(
+        _checkpoint("character_design", "completed", {}, "character-animation")
+    )
+    validate_checkpoint(
+        _checkpoint("rig_plan", "in_progress", {}, "character-animation")
+    )
+
+
+def test_canonical_stage_still_requires_its_artifact():
+    # The fix must not weaken enforcement for canonical stages.
+    with pytest.raises(CheckpointValidationError):
+        validate_checkpoint(
+            _checkpoint("compose", "completed", {}, "character-animation")
+        )

tests/tools/test_subtitle_timestamps.py

@@ -0,0 +1,49 @@
+"""Regression tests: SRT/VTT timestamp formatting must not overflow milliseconds.
+
+`_ts_srt` / `_ts_vtt` computed the seconds and millisecond fields independently:
+`ms = int(round((seconds % 1) * 1000))`. When the fractional part is >= 0.9995,
+that rounds to 1000, emitting a malformed 4-digit `…,1000` with no carry into
+the seconds (and, at boundaries, minutes/hours) field — e.g. 0.9999s ->
+`00:00:00,1000` instead of `00:00:01,000`. Such timestamps are rejected by
+strict SRT/VTT parsers (ffmpeg subtitles filter, VLC, browser WebVTT).
+"""
+
+import re
+import sys
+from pathlib import Path
+
+PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent
+sys.path.insert(0, str(PROJECT_ROOT))
+
+from tools.subtitle.subtitle_gen import SubtitleGen  # noqa: E402
+
+_SRT_RE = re.compile(r"^\d{2}:\d{2}:\d{2},\d{3}$")
+_VTT_RE = re.compile(r"^\d{2}:\d{2}:\d{2}\.\d{3}$")
+
+
+def test_millisecond_carry_does_not_overflow():
+    # Every fractional part >= 0.9995 must carry into the next second, never
+    # emit a 4-digit millisecond field.
+    assert SubtitleGen._ts_srt(0.9999) == "00:00:01,000"
+    assert SubtitleGen._ts_vtt(0.9999) == "00:00:01.000"
+    assert SubtitleGen._ts_srt(1.9996) == "00:00:02,000"
+
+
+def test_carry_propagates_across_minute_and_hour_boundaries():
+    assert SubtitleGen._ts_srt(59.9999) == "00:01:00,000"
+    assert SubtitleGen._ts_srt(3599.9999) == "01:00:00,000"
+    assert SubtitleGen._ts_srt(7261.9999) == "02:01:02,000"
+
+
+def test_normal_values_unchanged():
+    assert SubtitleGen._ts_srt(0.0) == "00:00:00,000"
+    assert SubtitleGen._ts_srt(1.5) == "00:00:01,500"
+    assert SubtitleGen._ts_srt(0.4994) == "00:00:00,499"
+    assert SubtitleGen._ts_vtt(83.25) == "00:01:23.250"
+
+
+def test_all_outputs_are_well_formed():
+    # Sweep values that land on and around the dangerous boundary.
+    for t in (0.0, 0.4995, 0.9995, 0.9999, 1.0, 59.9995, 3599.9999, 12345.9999):
+        assert _SRT_RE.match(SubtitleGen._ts_srt(t)), SubtitleGen._ts_srt(t)
+        assert _VTT_RE.match(SubtitleGen._ts_vtt(t)), SubtitleGen._ts_vtt(t)

tools/subtitle/subtitle_gen.py

@@ -309,19 +309,27 @@ class SubtitleGen(BaseTool):
         return "\n".join(lines)
 
     @staticmethod
-    def _ts_srt(seconds: float) -> str:
+    def _hmsms(seconds: float) -> tuple[int, int, int, int]:
+        """Decompose seconds into (h, m, s, ms), rounding to whole ms first.
+
+        Rounding to total milliseconds before splitting the fields lets the
+        carry propagate: 0.9995s+ must become the next second (…,000), not a
+        malformed 4-digit …,1000 with the seconds field left unincremented.
+        """
+        total_ms = int(round(max(0.0, seconds) * 1000))
+        h, rem = divmod(total_ms, 3_600_000)
+        m, rem = divmod(rem, 60_000)
+        s, ms = divmod(rem, 1_000)
+        return h, m, s, ms
+
+    @classmethod
+    def _ts_srt(cls, seconds: float) -> str:
         """Format seconds as SRT timestamp: HH:MM:SS,mmm"""
-        h = int(seconds // 3600)
-        m = int((seconds % 3600) // 60)
-        s = int(seconds % 60)
-        ms = int(round((seconds % 1) * 1000))
+        h, m, s, ms = cls._hmsms(seconds)
         return f"{h:02d}:{m:02d}:{s:02d},{ms:03d}"
 
-    @staticmethod
-    def _ts_vtt(seconds: float) -> str:
+    @classmethod
+    def _ts_vtt(cls, seconds: float) -> str:
         """Format seconds as VTT timestamp: HH:MM:SS.mmm"""
-        h = int(seconds // 3600)
-        m = int((seconds % 3600) // 60)
-        s = int(seconds % 60)
-        ms = int(round((seconds % 1) * 1000))
+        h, m, s, ms = cls._hmsms(seconds)
         return f"{h:02d}:{m:02d}:{s:02d}.{ms:03d}"