Align five-stage pipeline with manuscript

5 files changed, 315 insertions, 10 deletions

diff --git a/tests/test_e2e.py b/tests/test_e2e.py
index 9672527..ebf23a4 100644
--- a/tests/test_e2e.py
+++ b/tests/test_e2e.py
@@ -29,6 +29,14 @@ def test_offline_end_to_end_smoke(tmp_path: Path) -> None:
         "kernel_variant",
     }
     kernel_meta = record["variants"]["kernel_variant"]["_meta"]
-    assert kernel_meta["proof_dag"]["terminal_node_id"] == "n3"
-    assert kernel_meta["replacement_plan"]["changes"][0]["slot_id"] == "s1"
-    assert kernel_meta["diffused_proof"]["terminal_answer"] == "4"
+    assert kernel_meta["core_steps"] == [
+        "use nonnegativity of a square",
+        "expand the square",
+        "divide by a positive quantity",
+    ]
+    assert kernel_meta["mutable_slots"] == {
+        "slot1": {
+            "description": "positive square root at equality",
+            "original": "1",
+        }
+    }
diff --git a/tests/test_kernel_models.py b/tests/test_kernel_models.py
index 864d6d5..0dc653d 100644
--- a/tests/test_kernel_models.py
+++ b/tests/test_kernel_models.py
@@ -40,7 +40,7 @@ def test_replacement_must_reference_a_real_dag_node() -> None:
         {
             "changes": [
                 {
-                    "slot_id": "s1",
+                    "slot_id": "slot1",
                     "source_node_id": "n9",
                     "description": "constant",
                     "original_value": "1",
@@ -56,6 +56,70 @@ def test_replacement_must_reference_a_real_dag_node() -> None:
         replacements.validate_against(branched_dag())
 
 
+def test_replacement_must_target_a_source_leaf() -> None:
+    replacements = ReplacementPlan.model_validate(
+        {
+            "changes": [
+                {
+                    "slot_id": "slot1",
+                    "source_node_id": "n2",
+                    "description": "derived value",
+                    "original_value": "1",
+                    "replacement_value": "2",
+                    "guard_condition": "positive",
+                    "guard_justification": "2 is positive",
+                }
+            ],
+            "closure_statement": "No undeclared changes.",
+        }
+    )
+    with pytest.raises(ValueError, match="must target source leaf nodes.*n2"):
+        replacements.validate_against(branched_dag())
+
+
+def test_replacement_slot_id_matches_release_schema() -> None:
+    with pytest.raises(ValueError, match="must be slot1, slot2"):
+        ReplacementPlan.model_validate(
+            {
+                "changes": [
+                    {
+                        "slot_id": "s1",
+                        "source_node_id": "n1",
+                        "description": "constant",
+                        "original_value": "1",
+                        "replacement_value": "2",
+                        "guard_condition": "positive",
+                        "guard_justification": "2 is positive",
+                    }
+                ],
+                "closure_statement": "No undeclared changes.",
+            }
+        )
+
+
+def test_replacement_repair_preserves_slot_and_leaf_identity() -> None:
+    previous = ReplacementPlan.model_validate(
+        {
+            "changes": [
+                {
+                    "slot_id": "slot1",
+                    "source_node_id": "n1",
+                    "description": "constant",
+                    "original_value": "1",
+                    "replacement_value": "2",
+                    "guard_condition": "positive",
+                    "guard_justification": "2 is positive",
+                }
+            ],
+            "closure_statement": "No undeclared changes.",
+        }
+    )
+    changed_target = previous.model_copy(deep=True)
+    changed_target.changes[0].source_node_id = "n2"
+    with pytest.raises(ValueError, match="preserve replacement slot IDs"):
+        changed_target.validate_repair_of(previous)
+
+
 def test_diffusion_preserves_dependencies_and_methods() -> None:
     dag = branched_dag()
     methods = MethodPlan.model_validate(
@@ -130,7 +194,6 @@ def test_accept_verdict_normalizes_explicit_none_sentinel() -> None:
     verdict = JudgeVerdict(
         verdict="accept",
         step_by_step_check="n1 valid",
-        replacement_check="s1 valid",
         blocking_issues="None detected.",
         patch_suggestion="N/A",
     )
diff --git a/tests/test_paper_pipeline.py b/tests/test_paper_pipeline.py
new file mode 100644
index 0000000..7f1de46
--- /dev/null
+++ b/tests/test_paper_pipeline.py
@@ -0,0 +1,192 @@
+from __future__ import annotations
+
+import asyncio
+import json
+
+from gap_pipeline.clients import ScriptedClient
+from gap_pipeline.paper_pipeline import PaperKernelPipeline, PaperPipelineConfig
+from gap_pipeline.prompts import JUDGE_SYSTEM_PROMPT
+from gap_pipeline.store import RunStore
+
+
+def _dag() -> dict:
+    return {
+        "nodes": [
+            {
+                "node_id": "n1",
+                "claim": "(a-1)^2 >= 0",
+                "dependencies": [],
+            },
+            {
+                "node_id": "n2",
+                "claim": "a+1/a >= 2",
+                "dependencies": ["n1"],
+            },
+        ],
+        "terminal_node_id": "n2",
+    }
+
+
+def _methods() -> dict:
+    return {
+        "nodes": [
+            {
+                "node_id": "n1",
+                "method_label": "use nonnegativity of a square",
+            },
+            {
+                "node_id": "n2",
+                "method_label": "expand and divide by a positive quantity",
+            },
+        ]
+    }
+
+
+def _replacement() -> dict:
+    return {
+        "changes": [
+            {
+                "slot_id": "slot1",
+                "source_node_id": "n1",
+                "description": "positive equality value",
+                "original_value": "1",
+                "replacement_value": "2",
+                "guard_condition": "replacement is positive",
+                "guard_justification": "2 is positive",
+            }
+        ],
+        "closure_statement": "No undeclared changes.",
+    }
+
+
+def _diffused() -> dict:
+    return {
+        "nodes": [
+            {
+                "node_id": "n1",
+                "dependencies": [],
+                "method_label": "use nonnegativity of a square",
+                "instantiated_claim": "(x-2)^2 >= 0",
+                "justification": "squares are nonnegative",
+            },
+            {
+                "node_id": "n2",
+                "dependencies": ["n1"],
+                "method_label": "expand and divide by a positive quantity",
+                "instantiated_claim": "x+4/x >= 4",
+                "justification": "expand and divide by x>0",
+            },
+        ],
+        "terminal_node_id": "n2",
+        "terminal_answer": "4",
+    }
+
+
+def _variant(suffix: str = "") -> dict:
+    return {
+        "question": f"Let x>0. Prove that x+4/x >= 4.{suffix}",
+        "solution": (
+            "[n1] Since (x-2)^2 >= 0. "
+            "[n2] Expand and divide by x>0 to obtain x+4/x >= 4."
+        ),
+        "node_order": ["n1", "n2"],
+        "terminal_answer": "4",
+    }
+
+
+def _accept() -> dict:
+    return {
+        "verdict": "accept",
+        "step_by_step_check": "n1 is valid; n2 is valid",
+        "blocking_issues": "",
+        "patch_suggestion": "",
+    }
+
+
+def _reject() -> dict:
+    return {
+        "verdict": "reject",
+        "step_by_step_check": "n1 is valid; n2 needs a wording correction",
+        "blocking_issues": "the terminal wording is ambiguous",
+        "patch_suggestion": "clarify the terminal wording",
+    }
+
+
+def test_five_stage_repair_uses_prior_bundle_and_appendix_judge(
+    tmp_path,
+    item,
+) -> None:
+    proposer = ScriptedClient(
+        {
+            f"{item.item_id}.stage1.dag": _dag(),
+            f"{item.item_id}.stage2.methods": _methods(),
+            f"{item.item_id}.stage3.replacement": [
+                _replacement(),
+                _replacement(),
+            ],
+            f"{item.item_id}.stage4.diffusion": [
+                _diffused(),
+                _diffused(),
+            ],
+            f"{item.item_id}.stage5.render": [
+                _variant(),
+                _variant(" The requested bound is explicit."),
+            ],
+        }
+    )
+    judges = []
+    for judge_id in range(1, 6):
+        first = _reject() if judge_id == 1 else _accept()
+        judges.append(
+            ScriptedClient(
+                {f"{item.item_id}.verify": [first, _accept(), _accept()]}
+            )
+        )
+
+    run_root = tmp_path / "run"
+    result = asyncio.run(
+        PaperKernelPipeline(
+            proposer=proposer,
+            judges=judges,
+            store=RunStore(run_root, item.item_id),
+            config=PaperPipelineConfig(
+                proposer_model="scripted",
+                judge_model="scripted",
+            ),
+        ).run(item)
+    )
+
+    assert result.status == "accepted"
+    assert [row.pass_streak_after for row in result.iterations] == [0, 1, 2]
+    assert result.iterations[0].bundle_sha256 != result.iterations[1].bundle_sha256
+    assert result.iterations[1].bundle_sha256 == result.iterations[2].bundle_sha256
+
+    calls_dir = run_root / "items" / item.item_id / "calls"
+    repair_stage_calls = [
+        json.loads(
+            (
+                calls_dir
+                / f"{item.item_id}.stage{stage}.{name}.v02.json"
+            ).read_text()
+        )
+        for stage, name in [
+            (3, "replacement"),
+            (4, "diffusion"),
+            (5, "render"),
+        ]
+    ]
+    assert "PREVIOUS REPLACEMENT PLAN" in repair_stage_calls[0]["user_prompt"]
+    assert "PREVIOUS DIFFUSED PROOF" in repair_stage_calls[1]["user_prompt"]
+    assert "PREVIOUS RENDERED VARIANT" in repair_stage_calls[2]["user_prompt"]
+    assert all(
+        "the terminal wording is ambiguous" in call["user_prompt"]
+        for call in repair_stage_calls
+    )
+
+    judge_calls = sorted(calls_dir.glob(f"{item.item_id}.verify.*.json"))
+    assert len(judge_calls) == 15
+    assert not any(".a1." in path.name for path in judge_calls)
+    first_judge_call = json.loads(judge_calls[0].read_text())
+    assert first_judge_call["system_prompt"] == JUDGE_SYSTEM_PROMPT
+    assert "METHOD-LABEL SEQUENCE (abstract plan):" in first_judge_call["user_prompt"]
+    assert "SOURCE PROOF DAG:" not in first_judge_call["user_prompt"]
diff --git a/tests/test_prompts.py b/tests/test_prompts.py
index 4f6568b..9125d26 100644
--- a/tests/test_prompts.py
+++ b/tests/test_prompts.py
@@ -3,7 +3,7 @@ from __future__ import annotations
 import hashlib
 import inspect
 
-from gap_pipeline import prompts
+from gap_pipeline import kernel_prompts, prompts
 from gap_pipeline.clients import OpenAIJsonClient
 from gap_pipeline.kernel_models import (
     DiffusedProof,
@@ -12,6 +12,7 @@ from gap_pipeline.kernel_models import (
     ReplacementPlan,
 )
 from gap_pipeline.kernel_prompts import (
+    JUDGE_SYSTEM,
     dag_user,
     diffusion_user,
     judge_user,
@@ -40,6 +41,19 @@ EXPECTED = {
     "SURFACE_USER_TEMPLATE": "5d9043d7d1c1033db1aea0696812a2c4cbb6f3e28ca9d7b6436c036b40831fe7",
 }
 
+EXPECTED_FIVE_STAGE = {
+    "DAG_SYSTEM": "86e596a0c3aea07100ba642bb07d22e25385c2a51007b8f7578b662502ae4650",
+    "DAG_USER": "28757954499cee438cb3676fbb8aff39b2bbb4644a1f9dac2604b71d24b03547",
+    "METHOD_SYSTEM": "1091d4db002673abf67f4407418d0d7e34f648878f907d21e1ee910318d447ce",
+    "METHOD_USER": "79901153a7e2ef16ca4ade9763e4fd05a6d9d248a76085998723389f9934dca6",
+    "REPLACEMENT_SYSTEM": "2bf314a8cdcf986b344d45af081f7a1f675da6be5687ae4cd594db2ed7c0b25b",
+    "REPLACEMENT_USER": "6e88657245d951e47893eb9bcd6dbc5cd7a826fc67b8b2ce849da5b559f19fd9",
+    "DIFFUSION_SYSTEM": "f6ce12280b592830f5bacca3c8bc8131da26c9c0f94ffb9bfb7e9f53347f63ac",
+    "DIFFUSION_USER": "e92b39f15d1af1acc12daf34fabc0dfaf63bcebb4c947821fadb121aec44c9c1",
+    "RENDER_SYSTEM": "ead495a3f7df7d9655dc24ac2a437d4358ec080cd5e4477b3e97641828c3233e",
+    "RENDER_USER": "9e3eabd0fecf6946552db3217efa2eea7938099bcb093c9ca02fabf5b7d8f70c",
+}
+
 
 def test_prompt_values_are_byte_locked() -> None:
     actual = {
@@ -49,6 +63,16 @@ def test_prompt_values_are_byte_locked() -> None:
     assert actual == EXPECTED
 
 
+def test_five_stage_prompt_values_are_byte_locked() -> None:
+    actual = {
+        name: hashlib.sha256(
+            getattr(kernel_prompts, name).encode("utf-8")
+        ).hexdigest()
+        for name in EXPECTED_FIVE_STAGE
+    }
+    assert actual == EXPECTED_FIVE_STAGE
+
+
 def test_o3_adapter_does_not_send_temperature() -> None:
     source = inspect.getsource(OpenAIJsonClient.generate_json)
     assert '"temperature"' not in source
@@ -68,7 +92,7 @@ def test_literal_five_stage_prompts_render(item) -> None:
         {
             "changes": [
                 {
-                    "slot_id": "s1",
+                    "slot_id": "slot1",
                     "source_node_id": "n1",
                     "description": "constant",
                     "original_value": "1",
@@ -108,6 +132,15 @@ def test_literal_five_stage_prompts_render(item) -> None:
         replacement_user(item, dag, methods),
         diffusion_user(item, dag, methods, replacements),
         render_user(replacements, diffused),
-        judge_user(item, dag, methods, replacements, diffused, variant),
+        judge_user(item, methods, replacements, variant),
     ]
     assert all("{" in value and "}" in value for value in rendered)
+    assert JUDGE_SYSTEM == prompts.JUDGE_SYSTEM_PROMPT
+    assert rendered[-1] == prompts.JUDGE_USER_TEMPLATE.format(
+        original_problem=item.problem,
+        original_solution=item.solution,
+        method_labels='["n1: method"]',
+        slot_replacement=replacements.model_dump_json(indent=2),
+        candidate_problem="question",
+        candidate_proof="[n1] solution",
+    )
diff --git a/tests/test_release.py b/tests/test_release.py
index 8927b14..e122340 100644
--- a/tests/test_release.py
+++ b/tests/test_release.py
@@ -65,7 +65,7 @@ def test_offline_release_export_verifies_and_assembles(
             "accepted_replacement_plan": {
                 "changes": [
                     {
-                        "slot_id": "s1",
+                        "slot_id": "slot1",
                         "source_node_id": "n1",
                         "description": "value",
                         "original_value": "1",
@@ -109,7 +109,16 @@ def test_offline_release_export_verifies_and_assembles(
     )
     assert set(output["variants"]) == {*SURFACE_FAMILIES, "kernel_variant"}
     assert output["variants"]["kernel_variant"]["question"] == candidate["question"]
-    assert output["variants"]["kernel_variant"]["_meta"]["replacement_plan"]
+    kernel_meta = output["variants"]["kernel_variant"]["_meta"]
+    assert kernel_meta == {
+        "core_steps": ["method"],
+        "mutable_slots": {
+            "slot1": {
+                "description": "value",
+                "original": "1",
+            }
+        },
+    }
     with pytest.raises(FileExistsError):
         export_release(
             source_dataset=source_dir,