Parcourir la source

fix(host): make phase-three continuation replay-safe

SamLee il y a 1 jour
Parent
commit
ea1f99803a

+ 131 - 25
script_build_host/src/script_build_host/application/phase_three.py

@@ -41,6 +41,7 @@ from script_build_host.infrastructure.ownership import OwnerLease
 PHASE_TWO_CANDIDATE_PORTFOLIO_READY = "PHASE_TWO_CANDIDATE_PORTFOLIO_READY"
 PHASE_THREE_ROOT_ACCEPTED = "PHASE_THREE_ROOT_ACCEPTED"
 PHASE_THREE_POLICY = "script-build-phase-three/v1"
+PHASE_THREE_EXECUTION_STARTED = "planner_policy_execution_started"
 
 
 @dataclass(frozen=True, slots=True)
@@ -112,28 +113,42 @@ class PhaseThreeContinuationService:
             completion = await self.coordinator.root_completion(binding.root_trace_id)
             status = await self.legacy_state.get_status(script_build_id)
             if completion.get("status") == TaskStatus.COMPLETED.value:
+                if status is BuildStatus.RUNNING:
+                    lease = self.owner_lease_factory(script_build_id)
+                    token = await lease.acquire(script_build_id)
+                    try:
+                        await self._set_root_accepted_checkpoint(token)
+                    finally:
+                        await lease.release()
+                    status = BuildStatus.PARTIAL
                 return PhaseThreeAdvanceResult(
                     script_build_id,
                     status,
                     binding.root_trace_id,
                     str(binding.input_snapshot_id),
                 )
-            if status is not BuildStatus.PARTIAL:
+            reentry = False
+            if status in {BuildStatus.PARTIAL, BuildStatus.RUNNING} and completion.get(
+                "status"
+            ) in {TaskStatus.NEEDS_REPLAN.value, TaskStatus.PENDING.value}:
+                reentry = await self._phase_three_reentry_ready(binding.root_trace_id)
+            if status is not BuildStatus.PARTIAL and not reentry:
                 raise PhaseThreeBoundaryNotReady("phase-three advance requires a partial build")
-            if (
+            if not reentry and (
                 completion.get("status") != TaskStatus.BLOCKED.value
                 or completion.get("blocked_reason") != PHASE_TWO_CANDIDATE_PORTFOLIO_READY
             ):
                 raise PhaseThreeBoundaryNotReady(
                     "phase-three advance requires the phase-two portfolio checkpoint"
                 )
-            if self.phase_two_boundary_verifier is None:
+            if not reentry and self.phase_two_boundary_verifier is None:
                 raise PhaseThreeBoundaryNotReady("phase-two closure verifier is unavailable")
-            await self.phase_two_boundary_verifier.verify_checkpoint(
-                script_build_id=script_build_id,
-                root_trace_id=binding.root_trace_id,
-                input_snapshot_id=str(binding.input_snapshot_id),
-            )
+            if not reentry:
+                await self.phase_two_boundary_verifier.verify_checkpoint(
+                    script_build_id=script_build_id,
+                    root_trace_id=binding.root_trace_id,
+                    input_snapshot_id=str(binding.input_snapshot_id),
+                )
             ledger = await self.coordinator.task_store.load(binding.root_trace_id)
             if any(
                 item.status
@@ -178,15 +193,23 @@ class PhaseThreeContinuationService:
         if binding.root_trace_id != owner_token.root_trace_id:
             raise PhaseThreeBoundaryNotReady("owner token is bound to another Root")
         status = await self.legacy_state.get_status(script_build_id)
-        if status is not BuildStatus.PARTIAL:
-            raise PhaseThreeBoundaryNotReady(f"cannot continue a {status.value} build")
         snapshot = await self.input_snapshots.get(
             str(binding.input_snapshot_id), script_build_id=script_build_id
         )
         ledger = await self.coordinator.task_store.load(binding.root_trace_id)
         root = ledger.tasks[ledger.root_task_id]
         if root.status is TaskStatus.COMPLETED:
-            return
+            if status is BuildStatus.RUNNING:
+                await self._set_root_accepted_checkpoint(owner_token)
+                return
+            if status is BuildStatus.PARTIAL:
+                return
+            raise PhaseThreeBoundaryNotReady(f"cannot restore a completed Root from {status.value}")
+        reentry = status is BuildStatus.RUNNING and await self._phase_three_reentry_ready(
+            binding.root_trace_id
+        )
+        if status is not BuildStatus.PARTIAL and not reentry:
+            raise PhaseThreeBoundaryNotReady(f"cannot continue a {status.value} build")
         if root.status is TaskStatus.BLOCKED:
             if root.blocked_reason != PHASE_TWO_CANDIDATE_PORTFOLIO_READY:
                 raise PhaseThreeBoundaryNotReady("Root is not at the phase-two checkpoint")
@@ -207,7 +230,7 @@ class PhaseThreeContinuationService:
             raise PhaseThreeBoundaryNotReady("phase two still has active operations")
         closure = await self._accepted_root_closure(script_build_id, ledger)
         await self._verify_owner(owner_token)
-        snapshot, binding = await self._ensure_snapshot(snapshot, binding)
+        snapshot, binding = await self._ensure_snapshot(snapshot, binding, owner_token)
         root_contract = await self._execute_fenced(
             owner_token,
             lambda: self._freeze_root_contract(
@@ -273,8 +296,12 @@ class PhaseThreeContinuationService:
                     f"phase-three-contract:{script_build_id}",
                 ),
             )
-        await self.legacy_state.begin_phase(script_build_id)
+        await self._begin_phase(owner_token)
         await self._verify_owner(owner_token)
+        await self._execute_fenced(
+            owner_token,
+            lambda: self._claim_policy_execution(binding.root_trace_id, str(snapshot.snapshot_id)),
+        )
         if iterator is None:
             await self.runner.run_result(
                 messages=[], config=self.factory.build_phase_three_run_config(binding, snapshot)
@@ -286,12 +313,7 @@ class PhaseThreeContinuationService:
         await self._verify_owner(owner_token)
         if completion.get("status") != TaskStatus.COMPLETED.value:
             raise ProtocolViolation("phase-three Root was not explicitly ACCEPTed")
-        await self.legacy_state.set_checkpoint(
-            script_build_id,
-            checkpoint_code=PHASE_THREE_ROOT_ACCEPTED,
-            summary="Root accepted one validated immutable delivery manifest.",
-            root_trace_id=binding.root_trace_id,
-        )
+        await self._set_root_accepted_checkpoint(owner_token)
 
     async def _verify_owner(self, token: MissionOwnerToken) -> None:
         if self.fenced_command_gate is not None:
@@ -303,7 +325,9 @@ class PhaseThreeContinuationService:
         await self._verify_owner(token)
         return await mutation()
 
-    async def _ensure_snapshot(self, snapshot: Any, binding: Any) -> tuple[Any, Any]:
+    async def _ensure_snapshot(
+        self, snapshot: Any, binding: Any, owner_token: MissionOwnerToken
+    ) -> tuple[Any, Any]:
         presets = {str(item.get("preset")) for item in snapshot.prompt_manifest}
         required = {item.preset for item in self.prompt_requests}
         if required <= presets:
@@ -314,13 +338,46 @@ class PhaseThreeContinuationService:
             model_manifest=self.model_manifest or snapshot.model_manifest,
             version=3,
         )
-        binding = await self.bindings.compare_and_set_input_snapshot(
-            script_build_id=binding.script_build_id,
-            expected_snapshot_id=int(snapshot.snapshot_id),
-            new_snapshot_id=int(extended.snapshot_id),
-        )
+        if self.fenced_command_gate is not None:
+            await self.fenced_command_gate.compare_and_set_input_snapshot(
+                owner_token,
+                expected_snapshot_id=int(snapshot.snapshot_id),
+                new_snapshot_id=int(extended.snapshot_id),
+            )
+            binding = await self.bindings.get_by_build(binding.script_build_id)
+        else:
+            await self._verify_owner(owner_token)
+            binding = await self.bindings.compare_and_set_input_snapshot(
+                script_build_id=binding.script_build_id,
+                expected_snapshot_id=int(snapshot.snapshot_id),
+                new_snapshot_id=int(extended.snapshot_id),
+            )
         return extended, binding
 
+    async def _begin_phase(self, owner_token: MissionOwnerToken) -> None:
+        if self.fenced_command_gate is not None:
+            await self.fenced_command_gate.begin_phase(owner_token)
+            return
+        await self._verify_owner(owner_token)
+        await self.legacy_state.begin_phase(owner_token.script_build_id)
+
+    async def _set_root_accepted_checkpoint(self, owner_token: MissionOwnerToken) -> None:
+        summary = "Root accepted one validated immutable delivery manifest."
+        if self.fenced_command_gate is not None:
+            await self.fenced_command_gate.set_checkpoint(
+                owner_token,
+                checkpoint_code=PHASE_THREE_ROOT_ACCEPTED,
+                summary=summary,
+            )
+            return
+        await self._verify_owner(owner_token)
+        await self.legacy_state.set_checkpoint(
+            owner_token.script_build_id,
+            checkpoint_code=PHASE_THREE_ROOT_ACCEPTED,
+            summary=summary,
+            root_trace_id=owner_token.root_trace_id,
+        )
+
     async def _accepted_root_closure(
         self, script_build_id: int, ledger: Any
     ) -> _AcceptedRootClosure:
@@ -522,6 +579,54 @@ class PhaseThreeContinuationService:
         )
         return iterator
 
+    async def _claim_policy_execution(self, root_trace_id: str, snapshot_id: str) -> None:
+        """Persist the no-replay boundary immediately before invoking the model."""
+
+        events = await self.runner.trace_store.get_events(root_trace_id, 0)
+        starts = []
+        for event in events:
+            if (
+                event.get("event") != PHASE_THREE_EXECUTION_STARTED
+                and event.get("event_type") != PHASE_THREE_EXECUTION_STARTED
+            ):
+                continue
+            payload = event.get("data") or event.get("payload") or event
+            if payload.get("policy_version") == PHASE_THREE_POLICY:
+                starts.append(payload)
+        if starts:
+            # Internal-test mode deliberately fails closed here.  Replaying an
+            # uncertain model call is less safe than requiring explicit recovery.
+            raise MissionRecoveryRequired()
+        event_id = await self.runner.trace_store.append_event(
+            root_trace_id,
+            PHASE_THREE_EXECUTION_STARTED,
+            {
+                "policy_version": PHASE_THREE_POLICY,
+                "input_snapshot_id": snapshot_id,
+            },
+        )
+        if not event_id:
+            raise PhaseThreeBoundaryNotReady("phase-three execution marker was not durable")
+
+    async def _phase_three_reentry_ready(self, root_trace_id: str) -> bool:
+        events = await self.runner.trace_store.get_events(root_trace_id, 0)
+        migrations = []
+        starts = []
+        for event in events:
+            payload = event.get("data") or event.get("payload") or event
+            if payload.get("policy_version") != PHASE_THREE_POLICY:
+                continue
+            event_type = event.get("event") or event.get("event_type")
+            if event_type == "planner_policy_migrated":
+                migrations.append(payload)
+            elif event_type == PHASE_THREE_EXECUTION_STARTED:
+                starts.append(payload)
+        if starts:
+            raise MissionRecoveryRequired()
+        if len(migrations) > 1:
+            raise PhaseThreeBoundaryNotReady("multiple phase-three policy migrations exist")
+        return len(migrations) == 1
+
     async def _find_policy_messages(
         self, root_trace_id: str, policy: str, message: str
     ) -> tuple[int, int] | None:
@@ -576,6 +681,7 @@ class PhaseThreeContinuationService:
 
 
 __all__ = [
+    "PHASE_THREE_EXECUTION_STARTED",
     "PHASE_THREE_POLICY",
     "PHASE_THREE_ROOT_ACCEPTED",
     "PhaseThreeAdvanceResult",

+ 96 - 0
script_build_host/tests/test_phase_three_contracts.py

@@ -1,11 +1,21 @@
 from __future__ import annotations
 
 from dataclasses import replace
+from types import SimpleNamespace
+from unittest.mock import AsyncMock
 
 import pytest
+from agent.orchestration import TaskStatus
 
+from script_build_host.application.phase_three import (
+    PHASE_THREE_EXECUTION_STARTED,
+    PHASE_THREE_POLICY,
+    PhaseThreeContinuationService,
+)
 from script_build_host.domain.errors import (
     LegacyCanonicalIdentityConflict,
+    MissionRecoveryRequired,
+    PhaseThreeBoundaryNotReady,
     ProtocolViolation,
 )
 from script_build_host.domain.phase_three_artifacts import (
@@ -18,6 +28,7 @@ from script_build_host.domain.phase_two_artifacts import (
     ScriptParagraphV1,
     StructuredScriptArtifactV1,
 )
+from script_build_host.domain.records import BuildStatus
 
 _DIGEST = "sha256:" + "a" * 64
 
@@ -145,3 +156,88 @@ def test_root_delivery_manifest_is_strict_bounded_and_publishable() -> None:
         ).require_publishable()
     with pytest.raises(ProtocolViolation, match="2,000"):
         replace(manifest, build_summary="x" * 2_001, canonical_sha256="")
+
+
+@pytest.mark.asyncio
+async def test_phase_three_model_execution_marker_is_single_use() -> None:
+    class TraceStore:
+        def __init__(self) -> None:
+            self.events: list[dict[str, str]] = []
+
+        async def get_events(self, _trace_id: str, _since: int):
+            return list(self.events)
+
+        async def append_event(self, _trace_id: str, event: str, payload: dict[str, str]):
+            self.events.append({"event": event, **payload})
+            return len(self.events)
+
+    service = object.__new__(PhaseThreeContinuationService)
+    trace_store = TraceStore()
+    service.runner = SimpleNamespace(trace_store=trace_store)
+
+    await service._claim_policy_execution("root", "snapshot-3")
+    assert trace_store.events == [
+        {
+            "event": PHASE_THREE_EXECUTION_STARTED,
+            "policy_version": PHASE_THREE_POLICY,
+            "input_snapshot_id": "snapshot-3",
+        }
+    ]
+    with pytest.raises(MissionRecoveryRequired):
+        await service._claim_policy_execution("root", "snapshot-3")
+
+
+@pytest.mark.asyncio
+async def test_phase_three_reentry_is_allowed_only_before_model_execution() -> None:
+    class TraceStore:
+        def __init__(self) -> None:
+            self.events = [
+                {
+                    "event": "planner_policy_migrated",
+                    "policy_version": PHASE_THREE_POLICY,
+                    "input_snapshot_id": "snapshot-3",
+                }
+            ]
+
+        async def get_events(self, _trace_id: str, _since: int):
+            return list(self.events)
+
+        async def append_event(self, _trace_id: str, event: str, payload: dict[str, str]):
+            self.events.append({"event": event, **payload})
+            return len(self.events)
+
+    service = object.__new__(PhaseThreeContinuationService)
+    trace_store = TraceStore()
+    service.runner = SimpleNamespace(trace_store=trace_store)
+
+    assert await service._phase_three_reentry_ready("root") is True
+    await service._claim_policy_execution("root", "snapshot-3")
+    with pytest.raises(MissionRecoveryRequired):
+        await service._phase_three_reentry_ready("root")
+
+
+@pytest.mark.asyncio
+async def test_completed_root_does_not_restore_a_failed_build() -> None:
+    service = object.__new__(PhaseThreeContinuationService)
+    service._verify_owner = AsyncMock()
+    service._set_root_accepted_checkpoint = AsyncMock()
+    service.bindings = SimpleNamespace(
+        get_by_build=AsyncMock(
+            return_value=SimpleNamespace(root_trace_id="root", input_snapshot_id=3)
+        )
+    )
+    service.legacy_state = SimpleNamespace(get_status=AsyncMock(return_value=BuildStatus.FAILED))
+    service.input_snapshots = SimpleNamespace(get=AsyncMock(return_value=SimpleNamespace()))
+    root = SimpleNamespace(status=TaskStatus.COMPLETED)
+    service.coordinator = SimpleNamespace(
+        task_store=SimpleNamespace(
+            load=AsyncMock(
+                return_value=SimpleNamespace(root_task_id="root-task", tasks={"root-task": root})
+            )
+        )
+    )
+    owner_token = SimpleNamespace(root_trace_id="root")
+
+    with pytest.raises(PhaseThreeBoundaryNotReady, match="failed"):
+        await service.run(7, owner_token=owner_token)
+    service._set_root_accepted_checkpoint.assert_not_awaited()