Agent/tests/test_candidate_service.py

import asyncio
import json
import tempfile
import unittest
from unittest.mock import patch

from pydantic import ValidationError

from cyber_agent.application import (
    AgentApplication,
    AgentRole,
    ApplicationRegistry,
    ApplicationServices,
    ProviderRef,
)
from cyber_agent.application.candidate import (
    CandidateAdoptionReceipt,
    CandidateLedger,
    CandidateOperationRecord,
    CandidatePointer,
    CandidateReviewAction,
    MAX_CANDIDATES_PER_ROOT,
    MAX_CANDIDATE_OPERATIONS_PER_ROOT,
)
from cyber_agent.application.quality import CandidateValidationRecord
from cyber_agent.application.candidate_service import (
    CandidateService,
    CandidateStateError,
)
from cyber_agent.core.runner import AgentRunner, RunConfig
from cyber_agent.core.artifacts import (
    ArtifactRef,
    ValidationMaterial,
    material_content_hash,
)
from cyber_agent.core.task_protocol import TaskBrief, new_task_protocol
from cyber_agent.core.validation import (
    ScopeValidationResult,
    ValidationCheck,
    ValidationResult,
)
from cyber_agent.tools import get_tool_registry
from cyber_agent.tools.errors import RecoverableToolExecutionError
from cyber_agent.tools.registry import ToolRegistry
from cyber_agent.trace.models import Message as TraceMessage, Trace
from cyber_agent.trace.store import (
    FileSystemTraceStore,
    TraceStoreCorruptionError,
)


class MemoryCandidateRepository:
    def __init__(self):
        self.contents = {}
        self.put_calls = 0
        self.merge_calls = 0
        self.bad_artifact = False
        self.adoptions = {}
        self.adoption_calls = 0

    async def put_version(self, request):
        self.put_calls += 1
        return self._save(request)

    async def merge(self, request):
        self.merge_calls += 1
        return self._save(request)

    def _save(self, request):
        key = request.operation_id
        self.contents.setdefault(key, request)
        persisted = self.contents[key]
        digest = material_content_hash(persisted.content)
        if self.bad_artifact:
            digest = "0" * 64
        return ArtifactRef(
            artifact_id=(
                f"candidate:{persisted.candidate_id}:{persisted.revision}"
            ),
            version=str(persisted.revision),
            content_hash=digest,
            kind="candidate.output",
            mime_type="application/json",
        )

    async def resolve(self, ref, root_trace_id, uid):
        request = next(
            item
            for item in self.contents.values()
            if (
                ref.artifact_id
                == f"candidate:{item.candidate_id}:{item.revision}"
            )
        )
        return ValidationMaterial(
            **ref.model_dump(),
            root_trace_id=root_trace_id,
            uid=uid,
            content=request.content,
        )

    async def load_version(self, candidate_ref):
        return next(
            item.content
            for item in self.contents.values()
            if item.candidate_id == candidate_ref.candidate_id
            and item.revision == candidate_ref.revision
        )

    async def commit_adoption(self, request):
        self.adoption_calls += 1
        self.adoptions.setdefault(request.operation_id, request.candidate_ref)
        return CandidateAdoptionReceipt(
            operation_id=request.operation_id,
            external_id=f"published:{request.candidate_ref.candidate_id}",
        )


def application():
    return AgentApplication(
        application_id="candidate.test",
        application_version="1",
        root_role="writer",
        roles=(AgentRole(
            role_id="writer",
            model="fake",
            system_prompt="write",
        ),),
        artifact_resolver_ref=ProviderRef(
            provider_id="artifacts",
            provider_version="1",
        ),
        candidate_repository_ref=ProviderRef(
            provider_id="candidates",
            provider_version="1",
        ),
    )


class CandidateServiceTest(unittest.IsolatedAsyncioTestCase):
    async def asyncSetUp(self):
        self.temp = tempfile.TemporaryDirectory()
        self.store = FileSystemTraceStore(self.temp.name)
        self.repository = MemoryCandidateRepository()
        self.binding = ApplicationRegistry().register(
            application(),
            ApplicationServices(
                tool_registry=ToolRegistry(),
                artifact_resolver=self.repository,
                candidate_repository=self.repository,
            ),
        )
        self.service = CandidateService(
            store=self.store,
            application_binding=self.binding,
            repository=self.repository,
            artifact_resolver=self.repository,
        )
        await self._create_root("root-a")

    async def asyncTearDown(self):
        self.temp.cleanup()

    def _context(self, root_trace_id, *, task_brief=None):
        return {
            "agent_mode": "recursive",
            "agent_mode_revision": 3,
            "root_trace_id": root_trace_id,
            "root_task_anchor_hash": "a" * 64,
            "application_ref": self.binding.application_ref.model_dump(mode="json"),
            "application_role_id": "writer",
            "application_role_hash": self.binding.role("writer").role_hash,
            "task_protocol": new_task_protocol(task_brief),
        }

    async def _create_root(self, trace_id, *, uid="user-1"):
        await self.store.create_trace(Trace(
            trace_id=trace_id,
            mode="agent",
            agent_type="writer",
            uid=uid,
            context=self._context(trace_id),
        ))

    async def _create_child(self, trace_id, root_trace_id="root-a"):
        await self.store.create_trace(Trace(
            trace_id=trace_id,
            mode="agent",
            agent_type="writer",
            uid="user-1",
            parent_trace_id=root_trace_id,
            context=self._context(
                root_trace_id,
                task_brief=TaskBrief(
                    objective="write candidate",
                    reason="the root needs an option",
                    completion_criteria=["candidate is complete"],
                    expected_outputs=["one candidate"],
                ),
            ),
        ))

    async def _record_validation(self, trace_id, candidate, outcome="passed"):
        plan_hash = material_content_hash({
            "candidate": candidate.model_dump(mode="json"),
            "outcome": outcome,
        })
        status = "passed" if outcome == "passed" else "failed"
        scope = ScopeValidationResult(
            validator_trace_id=f"{trace_id}@validator",
            scope="output",
            outcome=outcome,
            checks=[ValidationCheck(
                check_id="quality.rule",
                status=status,
                issue=None if status == "passed" else "needs revision",
            )],
            reason="checked",
            retry_from=None if outcome == "passed" else "output",
            plan_hash=plan_hash,
        )
        result = ValidationResult(
            evaluated_trace_id=trace_id,
            outcome=outcome,
            scope_results=[scope],
            issues=[] if outcome == "passed" else ["needs revision"],
            retry_from=None if outcome == "passed" else "output",
            plan_hash=plan_hash,
        )
        await self.service.record_validation(
            trace_id,
            CandidateValidationRecord(
                candidate_ref=candidate,
                plan_hash=plan_hash,
                validation_result=result.model_dump(mode="json"),
                validated_at_sequence=candidate.created_at_sequence,
            ),
        )

    async def test_create_fork_merge_and_idempotent_recovery(self):
        first, duplicate = await asyncio.gather(
            self.service.manage(
                "root-a",
                operation="create",
                content={"text": "A"},
                parent_refs=[],
                effective_at_sequence=1,
            ),
            self.service.manage(
                "root-a",
                operation="create",
                content={"text": "A"},
                parent_refs=[],
                effective_at_sequence=1,
            ),
        )
        self.assertEqual(first, duplicate)
        self.assertEqual(1, self.repository.put_calls)

        second = await self.service.manage(
            "root-a",
            operation="create",
            content={"text": "B"},
            parent_refs=[],
            effective_at_sequence=2,
        )
        fork = await self.service.manage(
            "root-a",
            operation="fork",
            content={"text": "A2"},
            parent_refs=[CandidatePointer(
                candidate_id=first.candidate_id,
                revision=first.revision,
            )],
            effective_at_sequence=3,
        )
        self.assertEqual(first.candidate_id, fork.candidate_id)
        self.assertEqual(2, fork.revision)
        merged = await self.service.manage(
            "root-a",
            operation="merge",
            content={"text": "AB"},
            parent_refs=[
                CandidatePointer(
                    candidate_id=fork.candidate_id,
                    revision=fork.revision,
                ),
                CandidatePointer(
                    candidate_id=second.candidate_id,
                    revision=second.revision,
                ),
            ],
            effective_at_sequence=4,
        )
        self.assertNotIn(
            merged.candidate_id,
            {first.candidate_id, second.candidate_id},
        )
        self.assertEqual(2, len(merged.parent_refs))
        ledger = CandidateLedger.model_validate(
            await self.store.get_candidate_ledger("root-a")
        )
        self.assertEqual(4, len(ledger.candidates))
        self.assertTrue(all(
            ledger.current_state(item) == "proposed"
            for item in ledger.candidates
        ))

    async def test_finalize_crash_retries_same_semantic_command_identity(self):
        original_replace = self.store.replace_candidate_ledger
        writes = 0

        async def fail_completed_publish(root_trace_id, ledger):
            nonlocal writes
            writes += 1
            if writes == 2:
                raise OSError("simulated ledger finalize crash")
            await original_replace(root_trace_id, ledger)

        with patch.object(
            self.store,
            "replace_candidate_ledger",
            side_effect=fail_completed_publish,
        ):
            with self.assertRaisesRegex(
                RecoverableToolExecutionError,
                "original command ID",
            ):
                await self.service.manage(
                    "root-a",
                    operation="create",
                    content={"text": "survives retry"},
                    parent_refs=[],
                    effective_at_sequence=2,
                    command_id="candidate-crash-call",
                )

        pending = CandidateLedger.model_validate(
            await self.store.get_candidate_ledger("root-a")
        )
        self.assertEqual("pending", pending.operations[0].status)
        operation_id = pending.operations[0].operation_id
        recovered = await self.service.manage(
            "root-a",
            operation="create",
            content={"text": "survives retry"},
            parent_refs=[],
            effective_at_sequence=99,
            command_id="candidate-crash-call",
        )
        ledger = CandidateLedger.model_validate(
            await self.store.get_candidate_ledger("root-a")
        )
        self.assertEqual(operation_id, ledger.operations[0].operation_id)
        self.assertEqual("completed", ledger.operations[0].status)
        self.assertEqual(1, len(ledger.candidates))
        self.assertEqual(1, len(self.repository.contents))
        self.assertEqual(
            recovered,
            ledger.candidate(ledger.operations[0].candidate),
        )

    async def test_distinct_tool_calls_with_identical_payload_create_versions(self):
        first = await self.service.manage(
            "root-a",
            operation="create",
            content={"text": "same payload"},
            parent_refs=[],
            effective_at_sequence=10,
            command_id="candidate-call-one",
        )
        await self.store.add_message(TraceMessage.create(
            trace_id="root-a",
            role="tool",
            sequence=10,
            tool_call_id="candidate-call-one",
            content={"tool_name": "manage_candidate", "result": "registered"},
        ))
        await self.store.update_trace("root-a", head_sequence=10)
        second = await self.service.manage(
            "root-a",
            operation="create",
            content={"text": "same payload"},
            parent_refs=[],
            effective_at_sequence=20,
            command_id="candidate-call-two",
        )
        await self.store.add_message(TraceMessage.create(
            trace_id="root-a",
            role="tool",
            sequence=20,
            parent_sequence=10,
            tool_call_id="candidate-call-two",
            content={"tool_name": "manage_candidate", "result": "registered"},
        ))
        await self.store.update_trace("root-a", head_sequence=20)

        self.assertNotEqual(first, second)
        self.assertEqual(2, self.repository.put_calls)
        ledger = CandidateLedger.model_validate(
            await self.store.get_candidate_ledger("root-a")
        )
        self.assertEqual(2, len(ledger.candidates))
        self.assertEqual(2, len(ledger.operations))

    async def test_rewind_makes_future_branch_candidate_unreportable(self):
        await self.store.add_message(TraceMessage.create(
            trace_id="root-a",
            role="user",
            sequence=1,
            content="before candidate",
        ))
        await self.store.update_trace("root-a", head_sequence=1)
        candidate = await self.service.manage(
            "root-a",
            operation="create",
            content={"text": "future branch"},
            parent_refs=[],
            effective_at_sequence=10,
            command_id="future-candidate-call",
        )
        await self.store.add_message(TraceMessage.create(
            trace_id="root-a",
            role="tool",
            sequence=10,
            parent_sequence=1,
            tool_call_id="future-candidate-call",
            content={"tool_name": "manage_candidate", "result": "registered"},
        ))
        await self.store.update_trace("root-a", head_sequence=10)
        await self.service.validate_report_refs("root-a", [candidate])

        await self.store.update_trace("root-a", head_sequence=1)
        with self.assertRaisesRegex(
            CandidateStateError,
            "current owner Trace branch",
        ):
            await self.service.validate_report_refs("root-a", [candidate])

    async def test_checkpoint_update_cannot_mutate_identity_or_sequence(self):
        candidate = await self.service.manage(
            "root-a",
            operation="create",
            content={"text": "checkpoint"},
            parent_refs=[],
            effective_at_sequence=3,
        )
        plan_hash = "a" * 64
        await self.service.begin_validation_checkpoint(
            "root-a",
            candidate,
            plan={"checks": []},
            plan_hash=plan_hash,
            validated_at_sequence=3,
            material_usage_recorded=False,
        )

        with self.assertRaisesRegex(
            CandidateStateError,
            "fields are immutable",
        ):
            await self.service.update_validation_checkpoint(
                "root-a",
                candidate,
                plan_hash,
                validated_at_sequence=-1,
            )
        ledger = CandidateLedger.model_validate(
            await self.store.get_candidate_ledger("root-a")
        )
        self.assertEqual(
            3,
            ledger.validation_checkpoints[0]["validated_at_sequence"],
        )

    async def test_cross_root_and_cross_task_parent_access_is_rejected(self):
        first = await self.service.manage(
            "root-a",
            operation="create",
            content={"text": "A"},
            parent_refs=[],
            effective_at_sequence=1,
        )
        await self._create_root("root-b")
        with self.assertRaisesRegex(CandidateStateError, "not found"):
            await self.service.manage(
                "root-b",
                operation="fork",
                content={"text": "stolen"},
                parent_refs=[CandidatePointer(
                    candidate_id=first.candidate_id,
                    revision=first.revision,
                )],
                effective_at_sequence=2,
            )

        await self._create_child("child-a")
        with self.assertRaisesRegex(CandidateStateError, "another application, root, or Task"):
            await self.service.manage(
                "child-a",
                operation="fork",
                content={"text": "stolen"},
                parent_refs=[CandidatePointer(
                    candidate_id=first.candidate_id,
                    revision=first.revision,
                )],
                effective_at_sequence=3,
            )

    async def test_bad_artifact_fails_without_registering_candidate(self):
        self.repository.bad_artifact = True
        with self.assertRaisesRegex(CandidateStateError, "invalid ArtifactRef"):
            await self.service.manage(
                "root-a",
                operation="create",
                content={"text": "bad"},
                parent_refs=[],
                effective_at_sequence=1,
            )
        ledger = CandidateLedger.model_validate(
            await self.store.get_candidate_ledger("root-a")
        )
        self.assertEqual(0, len(ledger.candidates))
        self.assertEqual("failed", ledger.operations[0].status)

    async def test_report_validation_rejects_tamper_and_wrong_owner(self):
        candidate = await self.service.manage(
            "root-a",
            operation="create",
            content={"text": "A"},
            parent_refs=[],
            effective_at_sequence=1,
        )
        await self.service.validate_report_refs("root-a", [candidate])
        tampered = candidate.model_copy(update={
            "artifact_ref": candidate.artifact_ref.model_copy(update={
                "version": "tampered",
            }),
        })
        with self.assertRaisesRegex(CandidateStateError, "does not match"):
            await self.service.validate_report_refs("root-a", [tampered])
        await self._create_child("child-a")
        with self.assertRaisesRegex(CandidateStateError, "another application, root, or Task"):
            await self.service.validate_report_refs("child-a", [candidate])

    async def test_ledger_rejects_future_lineage_and_store_reports_corruption(self):
        candidate = await self.service.manage(
            "root-a",
            operation="create",
            content={"text": "A"},
            parent_refs=[],
            effective_at_sequence=1,
        )
        raw = (await self.store.get_candidate_ledger("root-a"))
        future = candidate.model_copy(update={
            "revision": 2,
            "parent_refs": (CandidatePointer(
                candidate_id=candidate.candidate_id,
                revision=3,
            ),),
        })
        raw["candidates"].append(future.model_dump(mode="json"))
        with self.assertRaisesRegex(ValidationError, "unknown parent"):
            CandidateLedger.model_validate(raw)

        ledger_path = self.store._get_candidate_ledger_file("root-a")
        ledger_path.write_text("{broken", encoding="utf-8")
        with self.assertRaisesRegex(TraceStoreCorruptionError, "candidate ledger"):
            await self.store.get_candidate_ledger("root-a")

    async def test_cyclic_lineage_is_rejected_before_repository_side_effect(self):
        first = await self.service.manage(
            "root-a",
            operation="create",
            content={"text": "A"},
            parent_refs=[],
            effective_at_sequence=1,
        )
        raw = await self.store.get_candidate_ledger("root-a")
        second = first.model_copy(update={
            "candidate_id": "candidate-b",
            "parent_refs": (CandidatePointer(
                candidate_id=first.candidate_id,
                revision=first.revision,
            ),),
        })
        raw["candidates"][0]["parent_refs"] = [{
            "candidate_id": second.candidate_id,
            "revision": second.revision,
        }]
        raw["candidates"].append(second.model_dump(mode="json"))
        ledger_path = self.store._get_candidate_ledger_file("root-a")
        ledger_path.write_text(json.dumps(raw), encoding="utf-8")

        put_calls = self.repository.put_calls
        with self.assertRaisesRegex(ValidationError, "lineage cycle"):
            await self.service.manage(
                "root-a",
                operation="create",
                content={"text": "must not be persisted"},
                parent_refs=[],
                effective_at_sequence=2,
            )
        self.assertEqual(put_calls, self.repository.put_calls)

    async def test_review_adoption_is_idempotent_and_creates_rewind_barrier(self):
        await self._create_child("child-a")
        candidate = await self.service.manage(
            "child-a",
            operation="create",
            content={"text": "A"},
            parent_refs=[],
            effective_at_sequence=4,
        )
        await self._record_validation("child-a", candidate)
        action = CandidateReviewAction(
            action="adopt",
            candidate_ref=candidate,
            reason="selected exact revision",
        )
        first = await self.service.apply_review_actions(
            "root-a",
            "child-a",
            report_refs=[candidate],
            actions=[action],
            effective_at_sequence=9,
        )
        second = await self.service.apply_review_actions(
            "root-a",
            "child-a",
            report_refs=[candidate],
            actions=[action],
            effective_at_sequence=9,
        )
        self.assertEqual(first, second)
        self.assertEqual(1, self.repository.adoption_calls)
        self.assertEqual(1, len(self.repository.adoptions))
        with self.assertRaisesRegex(CandidateStateError, "committed"):
            await self.service.assert_rewind_allowed("root-a", 8)
        await self.service.assert_rewind_allowed("root-a", 9)
        with self.assertRaisesRegex(CandidateStateError, "terminal"):
            await self.service.apply_review_actions(
                "root-a",
                "child-a",
                report_refs=[candidate],
                actions=[CandidateReviewAction(
                    action="discard",
                    candidate_ref=candidate,
                    reason="too late",
                )],
                effective_at_sequence=10,
            )

    async def test_adoption_retry_uses_same_command_after_parent_sequence_changes(self):
        await self._create_child("child-a")
        candidate = await self.service.manage(
            "child-a",
            operation="create",
            content={"text": "publish once"},
            parent_refs=[],
            effective_at_sequence=4,
        )
        await self._record_validation("child-a", candidate)
        peer = await self.service.manage(
            "child-a",
            operation="create",
            content={"text": "peer"},
            parent_refs=[],
            effective_at_sequence=5,
        )
        await self._record_validation("child-a", peer)
        action = CandidateReviewAction(
            action="adopt",
            candidate_ref=candidate,
            reason="stable parent decision",
        )
        first = await self.service.apply_review_actions(
            "root-a",
            "child-a",
            report_refs=[candidate, peer],
            actions=[action],
            effective_at_sequence=10,
            command_id="adoption-review-call",
        )
        recovered = await self.service.apply_review_actions(
            "root-a",
            "child-a",
            report_refs=[candidate, peer],
            actions=[action],
            effective_at_sequence=999,
            command_id="adoption-review-call",
        )
        self.assertEqual(first, recovered)
        self.assertEqual(1, self.repository.adoption_calls)
        self.assertEqual(1, len(self.repository.adoptions))
        with self.assertRaisesRegex(
            CandidateStateError,
            "review batch payload changed",
        ):
            await self.service.apply_review_actions(
                "root-a",
                "child-a",
                report_refs=[candidate, peer],
                actions=[action.model_copy(update={
                    "reason": "tampered decision payload",
                })],
                effective_at_sequence=1_000,
                command_id="adoption-review-call",
            )
        self.assertEqual(1, self.repository.adoption_calls)
        with self.assertRaisesRegex(
            CandidateStateError,
            "review batch payload changed",
        ):
            await self.service.apply_review_actions(
                "root-a",
                "child-a",
                report_refs=[candidate, peer],
                actions=[
                    action,
                    CandidateReviewAction(
                        action="discard",
                        candidate_ref=peer,
                        reason="injected tail action",
                    ),
                ],
                effective_at_sequence=1_001,
                command_id="adoption-review-call",
            )
        self.assertEqual("proposed", (
            CandidateLedger.model_validate(
                await self.store.get_candidate_ledger("root-a")
            ).current_state(peer)
        ))
        with self.assertRaisesRegex(
            CandidateStateError,
            "review batch payload changed",
        ):
            await self.service.apply_review_actions(
                "root-a",
                "child-a",
                report_refs=[candidate, peer],
                actions=[],
                effective_at_sequence=1_002,
                command_id="adoption-review-call",
            )

    async def test_review_batch_replay_rejects_deleted_and_reordered_actions(self):
        await self._create_child("child-a")
        candidates = []
        for sequence, text in ((4, "A"), (5, "B")):
            candidate = await self.service.manage(
                "child-a",
                operation="create",
                content={"text": text},
                parent_refs=[],
                effective_at_sequence=sequence,
            )
            await self._record_validation("child-a", candidate)
            candidates.append(candidate)
        actions = [
            CandidateReviewAction(
                action="discard",
                candidate_ref=candidate,
                reason=f"discard {candidate.candidate_id}",
            )
            for candidate in candidates
        ]
        await self.service.apply_review_actions(
            "root-a",
            "child-a",
            report_refs=candidates,
            actions=actions,
            effective_at_sequence=10,
            command_id="two-action-review",
        )

        for altered in (actions[:1], list(reversed(actions))):
            with self.assertRaisesRegex(
                CandidateStateError,
                "review batch payload changed",
            ):
                await self.service.apply_review_actions(
                    "root-a",
                    "child-a",
                    report_refs=candidates,
                    actions=altered,
                    effective_at_sequence=999,
                    command_id="two-action-review",
                )

    async def test_empty_review_batch_cannot_gain_an_action_during_replay(self):
        await self._create_child("child-a")
        candidate = await self.service.manage(
            "child-a",
            operation="create",
            content={"text": "A"},
            parent_refs=[],
            effective_at_sequence=4,
        )
        await self._record_validation("child-a", candidate)
        self.assertEqual([], await self.service.apply_review_actions(
            "root-a",
            "child-a",
            report_refs=[candidate],
            actions=[],
            effective_at_sequence=10,
            command_id="empty-review-batch",
        ))

        with self.assertRaisesRegex(
            CandidateStateError,
            "review batch payload changed",
        ):
            await self.service.apply_review_actions(
                "root-a",
                "child-a",
                report_refs=[candidate],
                actions=[CandidateReviewAction(
                    action="discard",
                    candidate_ref=candidate,
                    reason="injected after empty batch",
                )],
                effective_at_sequence=999,
                command_id="empty-review-batch",
            )
        self.assertEqual(
            "proposed",
            CandidateLedger.model_validate(
                await self.store.get_candidate_ledger("root-a")
            ).current_state(candidate),
        )

    async def test_review_capacity_rejects_before_envelope_or_repository_write(self):
        await self._create_child("child-a")
        candidate = await self.service.manage(
            "child-a",
            operation="create",
            content={"text": "capacity"},
            parent_refs=[],
            effective_at_sequence=4,
        )
        await self._record_validation("child-a", candidate)
        ledger = CandidateLedger.model_validate(
            await self.store.get_candidate_ledger("root-a")
        )
        fillers = [
            CandidateOperationRecord(
                operation_id=f"capacity-filler-{index}",
                operation="review_batch",
                input_hash="0" * 64,
                status="completed",
            )
            for index in range(
                MAX_CANDIDATE_OPERATIONS_PER_ROOT - len(ledger.operations)
            )
        ]
        saturated = CandidateLedger.model_validate({
            **ledger.model_dump(mode="json"),
            "operations": [
                *[item.model_dump(mode="json") for item in ledger.operations],
                *[item.model_dump(mode="json") for item in fillers],
            ],
        })
        await self.store.replace_candidate_ledger(
            "root-a",
            saturated.model_dump(mode="json"),
        )
        before = await self.store.get_candidate_ledger("root-a")

        with self.assertRaisesRegex(
            CandidateStateError,
            "operation limit exceeded before execution",
        ):
            await self.service.apply_review_actions(
                "root-a",
                "child-a",
                report_refs=[candidate],
                actions=[CandidateReviewAction(
                    action="adopt",
                    candidate_ref=candidate,
                    reason="must not reach repository",
                )],
                effective_at_sequence=10,
                command_id="capacity-overflow-review",
            )
        self.assertEqual(0, self.repository.adoption_calls)
        self.assertEqual(before, await self.store.get_candidate_ledger("root-a"))

    async def test_review_capacity_allows_exact_envelope_and_action_boundary(self):
        await self._create_child("child-a")
        candidate = await self.service.manage(
            "child-a",
            operation="create",
            content={"text": "exact capacity"},
            parent_refs=[],
            effective_at_sequence=4,
        )
        await self._record_validation("child-a", candidate)
        ledger = CandidateLedger.model_validate(
            await self.store.get_candidate_ledger("root-a")
        )
        target = MAX_CANDIDATE_OPERATIONS_PER_ROOT - 2
        fillers = [
            CandidateOperationRecord(
                operation_id=f"exact-filler-{index}",
                operation="review_batch",
                input_hash="1" * 64,
                status="completed",
            )
            for index in range(target - len(ledger.operations))
        ]
        near_limit = CandidateLedger.model_validate({
            **ledger.model_dump(mode="json"),
            "operations": [
                *[item.model_dump(mode="json") for item in ledger.operations],
                *[item.model_dump(mode="json") for item in fillers],
            ],
        })
        await self.store.replace_candidate_ledger(
            "root-a",
            near_limit.model_dump(mode="json"),
        )

        await self.service.apply_review_actions(
            "root-a",
            "child-a",
            report_refs=[candidate],
            actions=[CandidateReviewAction(
                action="discard",
                candidate_ref=candidate,
                reason="fits exact operation boundary",
            )],
            effective_at_sequence=10,
            command_id="exact-capacity-review",
        )
        completed = CandidateLedger.model_validate(
            await self.store.get_candidate_ledger("root-a")
        )
        self.assertEqual(MAX_CANDIDATE_OPERATIONS_PER_ROOT, len(completed.operations))
        self.assertEqual("discarded", completed.current_state(candidate))

    async def test_review_batch_reserves_action_slot_across_crash_recovery(self):
        await self._create_child("child-a")
        candidate = await self.service.manage(
            "child-a",
            operation="create",
            content={"text": "reserved review"},
            parent_refs=[],
            effective_at_sequence=4,
        )
        await self._record_validation("child-a", candidate)
        ledger = CandidateLedger.model_validate(
            await self.store.get_candidate_ledger("root-a")
        )
        target = MAX_CANDIDATE_OPERATIONS_PER_ROOT - 2
        fillers = [
            CandidateOperationRecord(
                operation_id=f"review-reservation-filler-{index}",
                operation="review_batch",
                input_hash="2" * 64,
                status="completed",
            )
            for index in range(target - len(ledger.operations))
        ]
        near_limit = CandidateLedger.model_validate({
            **ledger.model_dump(mode="json"),
            "operations": [
                *[item.model_dump(mode="json") for item in ledger.operations],
                *[item.model_dump(mode="json") for item in fillers],
            ],
        })
        await self.store.replace_candidate_ledger(
            "root-a",
            near_limit.model_dump(mode="json"),
        )
        action = CandidateReviewAction(
            action="adopt",
            candidate_ref=candidate,
            reason="reserve the complete review batch",
        )
        original_replace = self.store.replace_candidate_ledger
        writes = 0

        async def crash_after_batch_reservation(root_trace_id, raw_ledger):
            nonlocal writes
            writes += 1
            if writes == 2:
                raise OSError("crash after review reservation")
            await original_replace(root_trace_id, raw_ledger)

        with patch.object(
            self.store,
            "replace_candidate_ledger",
            side_effect=crash_after_batch_reservation,
        ):
            with self.assertRaisesRegex(OSError, "review reservation"):
                await self.service.apply_review_actions(
                    "root-a",
                    "child-a",
                    report_refs=[candidate],
                    actions=[action],
                    effective_at_sequence=10,
                    command_id="capacity-reserved-review",
                )
        reserved = CandidateLedger.model_validate(
            await self.store.get_candidate_ledger("root-a")
        )
        self.assertEqual(MAX_CANDIDATE_OPERATIONS_PER_ROOT, len(reserved.operations))
        self.assertEqual(2, sum(
            item.status == "pending" for item in reserved.operations
        ))
        repository_version_calls = self.repository.put_calls
        with self.assertRaisesRegex(
            CandidateStateError,
            "operation limit exceeded before execution",
        ):
            await self.service.manage(
                "child-a",
                operation="create",
                content={"text": "must not steal review capacity"},
                parent_refs=[],
                effective_at_sequence=11,
                command_id="sibling-after-review-reservation",
            )
        self.assertEqual(repository_version_calls, self.repository.put_calls)

        records = await self.service.apply_review_actions(
            "root-a",
            "child-a",
            report_refs=[candidate],
            actions=[action],
            effective_at_sequence=999,
            command_id="capacity-reserved-review",
        )
        self.assertEqual("adopted", records[0].state)
        self.assertEqual(1, self.repository.adoption_calls)
        completed = CandidateLedger.model_validate(
            await self.store.get_candidate_ledger("root-a")
        )
        self.assertEqual(MAX_CANDIDATE_OPERATIONS_PER_ROOT, len(completed.operations))
        self.assertTrue(all(
            item.status == "completed"
            for item in completed.operations
            if item.command_id == "capacity-reserved-review"
        ))

    async def test_candidate_capacity_rejects_before_repository_or_intent_write(self):
        await self._create_child("child-a")
        candidate = await self.service.manage(
            "child-a",
            operation="create",
            content={"text": "seed"},
            parent_refs=[],
            effective_at_sequence=4,
        )
        ledger = CandidateLedger.model_validate(
            await self.store.get_candidate_ledger("root-a")
        )
        saturated_candidates = [candidate]
        saturated_candidates.extend(
            candidate.model_copy(update={
                "candidate_id": f"capacity-candidate-{index}",
            })
            for index in range(1, MAX_CANDIDATES_PER_ROOT)
        )
        saturated = CandidateLedger.model_validate({
            **ledger.model_dump(mode="json"),
            "candidates": [
                item.model_dump(mode="json") for item in saturated_candidates
            ],
        })
        await self.store.replace_candidate_ledger(
            "root-a",
            saturated.model_dump(mode="json"),
        )
        before = await self.store.get_candidate_ledger("root-a")
        repository_calls = self.repository.put_calls

        with self.assertRaisesRegex(
            CandidateStateError,
            "candidate limit exceeded before execution",
        ):
            await self.service.manage(
                "child-a",
                operation="create",
                content={"text": "must remain external-orphan free"},
                parent_refs=[],
                effective_at_sequence=5,
                command_id="candidate-capacity-overflow",
            )

        self.assertEqual(repository_calls, self.repository.put_calls)
        self.assertEqual(before, await self.store.get_candidate_ledger("root-a"))

    async def test_pending_candidate_reserves_capacity_until_recovery(self):
        await self._create_child("child-a")
        seed = await self.service.manage(
            "child-a",
            operation="create",
            content={"text": "seed"},
            parent_refs=[],
            effective_at_sequence=4,
        )
        ledger = CandidateLedger.model_validate(
            await self.store.get_candidate_ledger("root-a")
        )
        near_limit_candidates = [seed]
        near_limit_candidates.extend(
            seed.model_copy(update={
                "candidate_id": f"reserved-candidate-{index}",
            })
            for index in range(1, MAX_CANDIDATES_PER_ROOT - 1)
        )
        near_limit = CandidateLedger.model_validate({
            **ledger.model_dump(mode="json"),
            "candidates": [
                item.model_dump(mode="json") for item in near_limit_candidates
            ],
        })
        await self.store.replace_candidate_ledger(
            "root-a",
            near_limit.model_dump(mode="json"),
        )
        original_replace = self.store.replace_candidate_ledger
        writes = 0

        async def crash_after_repository(root_trace_id, raw_ledger):
            nonlocal writes
            writes += 1
            if writes == 2:
                raise OSError("crash before candidate publish")
            await original_replace(root_trace_id, raw_ledger)

        with patch.object(
            self.store,
            "replace_candidate_ledger",
            side_effect=crash_after_repository,
        ):
            with self.assertRaisesRegex(
                RecoverableToolExecutionError,
                "original command ID",
            ):
                await self.service.manage(
                    "child-a",
                    operation="create",
                    content={"text": "reserved A"},
                    parent_refs=[],
                    effective_at_sequence=5,
                    command_id="reserved-candidate-A",
                )
        pending = CandidateLedger.model_validate(
            await self.store.get_candidate_ledger("root-a")
        )
        self.assertEqual(MAX_CANDIDATES_PER_ROOT - 1, len(pending.candidates))
        self.assertEqual(
            1,
            sum(
                item.status == "pending" and item.operation == "create"
                for item in pending.operations
            ),
        )
        calls_after_a = self.repository.put_calls

        with self.assertRaisesRegex(
            CandidateStateError,
            "candidate limit exceeded before execution",
        ):
            await self.service.manage(
                "child-a",
                operation="create",
                content={"text": "must not steal A's slot"},
                parent_refs=[],
                effective_at_sequence=6,
                command_id="reserved-candidate-B",
            )
        self.assertEqual(calls_after_a, self.repository.put_calls)

        recovered = await self.service.manage(
            "child-a",
            operation="create",
            content={"text": "reserved A"},
            parent_refs=[],
            effective_at_sequence=99,
            command_id="reserved-candidate-A",
        )
        completed = CandidateLedger.model_validate(
            await self.store.get_candidate_ledger("root-a")
        )
        self.assertEqual(MAX_CANDIDATES_PER_ROOT, len(completed.candidates))
        self.assertEqual(recovered, completed.candidate(recovered))

    async def test_runner_replays_pending_candidate_with_original_tool_call_id(self):
        await self._create_child("child-a")
        registry = get_tool_registry().clone(["manage_candidate"])
        runner = AgentRunner(
            trace_store=self.store,
            tool_registry=registry,
            llm_call=lambda **_kwargs: None,
            application_binding=self.binding,
            candidate_service=self.service,
        )
        config = RunConfig(
            tools=["manage_candidate"],
            tool_groups=[],
            auto_execute_tools=True,
        )
        request = {
            "request": {
                "operation": "create",
                "content": {"text": "recover with the original command"},
                "parent_refs": [],
            },
        }
        assistant = TraceMessage.create(
            trace_id="child-a",
            role="assistant",
            sequence=4,
            content={
                "text": "",
                "tool_calls": [{
                    "id": "recoverable-candidate-call",
                    "type": "function",
                    "function": {
                        "name": "manage_candidate",
                        "arguments": json.dumps(request),
                    },
                }],
            },
        )
        await self.store.add_message(assistant)
        await self.store.update_trace(
            "child-a",
            head_sequence=4,
            last_sequence=4,
        )
        trace = await self.store.get_trace("child-a")
        original_replace = self.store.replace_candidate_ledger
        writes = 0

        async def fail_first_publish(root_trace_id, raw_ledger):
            nonlocal writes
            writes += 1
            if writes == 2:
                raise OSError("candidate publish interrupted")
            await original_replace(root_trace_id, raw_ledger)

        with patch.object(
            self.store,
            "replace_candidate_ledger",
            side_effect=fail_first_publish,
        ):
            with self.assertRaises(RecoverableToolExecutionError):
                await registry.execute(
                    "manage_candidate",
                    request,
                    context=runner._build_tool_context(
                        config=config,
                        trace=trace,
                        trace_id=trace.trace_id,
                        goal_id=None,
                        goal_tree=None,
                        sequence=5,
                        head_sequence=4,
                        tool_call_id="recoverable-candidate-call",
                        side_branch_ctx=None,
                        trigger_event=None,
                    ),
                    allowed_tool_names={"manage_candidate"},
                    tool_call_id="recoverable-candidate-call",
                )
        pending = CandidateLedger.model_validate(
            await self.store.get_candidate_ledger("root-a")
        )
        self.assertEqual("pending", pending.operations[-1].status)
        self.assertEqual(1, len(self.repository.contents))
        self.assertFalse(any(
            item.role == "tool"
            for item in await self.store.get_trace_messages("child-a")
        ))

        healed, next_sequence = await runner._heal_orphaned_tool_calls(
            [assistant],
            await self.store.get_trace("child-a"),
            None,
            config,
            5,
        )
        self.assertEqual(6, next_sequence)
        self.assertEqual("tool", healed[-1].role)
        self.assertEqual(
            "recoverable-candidate-call",
            healed[-1].tool_call_id,
        )
        self.assertEqual(1, len(self.repository.contents))
        completed = CandidateLedger.model_validate(
            await self.store.get_candidate_ledger("root-a")
        )
        matching = [
            item for item in completed.operations
            if item.command_id == "recoverable-candidate-call"
        ]
        self.assertEqual(1, len(matching))
        self.assertEqual("completed", matching[0].status)
        self.assertEqual(1, len(completed.candidates))

    async def test_descendant_adoption_blocks_ancestor_rewind_without_mapping(self):
        await self._create_child("child-a")
        await self.store.create_trace(Trace(
            trace_id="grandchild-a",
            mode="agent",
            agent_type="writer",
            uid="user-1",
            parent_trace_id="child-a",
            context=self._context(
                "root-a",
                task_brief=TaskBrief(
                    objective="write descendant candidate",
                    reason="the child needs an option",
                    completion_criteria=["candidate is complete"],
                    expected_outputs=["one candidate"],
                ),
            ),
        ))
        candidate = await self.service.manage(
            "grandchild-a",
            operation="create",
            content={"text": "nested candidate"},
            parent_refs=[],
            effective_at_sequence=4,
        )
        await self._record_validation("grandchild-a", candidate)
        await self.service.apply_review_actions(
            "child-a",
            "grandchild-a",
            report_refs=[candidate],
            actions=[CandidateReviewAction(
                action="adopt",
                candidate_ref=candidate,
                reason="commit descendant output",
            )],
            effective_at_sequence=9,
        )

        with self.assertRaisesRegex(CandidateStateError, "committed"):
            await self.service.assert_rewind_allowed("root-a", 100)

    async def test_failed_candidate_can_only_be_revised_or_discarded(self):
        await self._create_child("child-a")
        candidate = await self.service.manage(
            "child-a",
            operation="create",
            content={"text": "placeholder"},
            parent_refs=[],
            effective_at_sequence=4,
        )
        await self._record_validation("child-a", candidate, outcome="failed")
        with self.assertRaisesRegex(CandidateStateError, "passed"):
            await self.service.apply_review_actions(
                "root-a",
                "child-a",
                report_refs=[candidate],
                actions=[CandidateReviewAction(
                    action="adopt",
                    candidate_ref=candidate,
                    reason="invalid",
                )],
                effective_at_sequence=9,
            )
        records = await self.service.apply_review_actions(
            "root-a",
            "child-a",
            report_refs=[candidate],
            actions=[CandidateReviewAction(
                action="revise",
                candidate_ref=candidate,
                reason="remove placeholder",
            )],
            effective_at_sequence=10,
        )
        self.assertEqual("discarded", records[0].state)
        self.assertIn("retry_from=output", records[0].reason)

    async def test_adoption_replays_same_operation_after_finalize_crash(self):
        await self._create_child("child-a")
        candidate = await self.service.manage(
            "child-a",
            operation="create",
            content={"text": "A"},
            parent_refs=[],
            effective_at_sequence=4,
        )
        await self._record_validation("child-a", candidate)
        action = CandidateReviewAction(
            action="adopt",
            candidate_ref=candidate,
            reason="publish once",
        )
        original_replace = self.store.replace_candidate_ledger
        calls = 0

        async def fail_finalize(root_trace_id, ledger):
            nonlocal calls
            calls += 1
            if calls == 3:
                raise OSError("crash before framework finalize")
            await original_replace(root_trace_id, ledger)

        with patch.object(
            self.store,
            "replace_candidate_ledger",
            side_effect=fail_finalize,
        ):
            with self.assertRaisesRegex(OSError, "finalize"):
                await self.service.apply_review_actions(
                    "root-a",
                    "child-a",
                    report_refs=[candidate],
                    actions=[action],
                    effective_at_sequence=9,
                )
        pending = CandidateLedger.model_validate(
            await self.store.get_candidate_ledger("root-a")
        )
        self.assertEqual("adoption_pending", pending.current_state(candidate))
        await self.service.apply_review_actions(
            "root-a",
            "child-a",
            report_refs=[candidate],
            actions=[action],
            effective_at_sequence=9,
        )
        self.assertEqual(2, self.repository.adoption_calls)
        self.assertEqual(1, len(self.repository.adoptions))
        completed = CandidateLedger.model_validate(
            await self.store.get_candidate_ledger("root-a")
        )
        self.assertEqual("adopted", completed.current_state(candidate))


if __name__ == "__main__":
    unittest.main()