From 5dcd026ef77fc402e36aa630dd5f1357811c5039 Mon Sep 17 00:00:00 2001
From: 0xDevNinja <manmit0x@gmail.com>
Date: Tue, 7 Jul 2026 15:29:03 +0530
Subject: [PATCH] fix(checkpoint): don't KeyError on manifest stages without a
 canonical artifact
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

_validate_artifacts_for_stage looked up CANONICAL_STAGE_ARTIFACTS[stage]
unconditionally, but the valid stage list comes from the pipeline manifest via
get_pipeline_stages(), which declares stages beyond the 9 canonical ones — e.g.
character-animation adds `character_design`/`rig_plan`. Such a stage passes the
`stage in valid_stages` guard, then raised an unhandled KeyError on the
canonical lookup, so those stages could never be checkpointed (the crash hits
write/read_checkpoint and friends, even for in_progress checkpoints).

Look the canonical artifact up defensively with `.get()` and skip the
required-artifact check when there is none. Canonical stages still require their
artifact when completed.
---
 lib/checkpoint.py                             | 13 +++-
 .../lib/test_checkpoint_noncanonical_stage.py | 62 +++++++++++++++++++
 2 files changed, 73 insertions(+), 2 deletions(-)
 create mode 100644 tests/lib/test_checkpoint_noncanonical_stage.py

diff --git a/lib/checkpoint.py b/lib/checkpoint.py
index b2b597a0..e35e6164 100644
--- a/lib/checkpoint.py
+++ b/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}"
diff --git a/tests/lib/test_checkpoint_noncanonical_stage.py b/tests/lib/test_checkpoint_noncanonical_stage.py
new file mode 100644
index 00000000..15225bc9
--- /dev/null
+++ b/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")
+        )