Agent/cyber_agent/application/candidate_service.py

"""Candidate command service with one root-level consistency boundary."""

from __future__ import annotations

import asyncio
from hashlib import sha256
from typing import Any
from weakref import WeakValueDictionary

from cyber_agent.application.candidate import (
    CandidateLedger,
    CandidateLifecycleRecord,
    CandidateOperationRecord,
    CandidatePointer,
    CandidateRef,
    CandidateRepository,
    CandidateVersionRequest,
)
from cyber_agent.application.models import canonical_json
from cyber_agent.core.agent_mode import require_mutable_trace_policy
from cyber_agent.core.artifacts import ArtifactRef, ArtifactResolver, resolve_artifact_refs
from cyber_agent.core.task_protocol import ensure_task_protocol, task_contract_ref


class CandidateStateError(ValueError):
    pass


class CandidateService:
    """Validate, persist, and recover candidate version operations."""

    _locks: "WeakValueDictionary[str, asyncio.Lock]" = WeakValueDictionary()

    def __init__(
        self,
        *,
        store: Any,
        application_binding: Any,
        repository: CandidateRepository,
        artifact_resolver: ArtifactResolver,
    ) -> None:
        self.store = store
        self.binding = application_binding
        self.repository = repository
        self.artifact_resolver = artifact_resolver

    @classmethod
    def _lock_for(cls, root_trace_id: str) -> asyncio.Lock:
        lock = cls._locks.get(root_trace_id)
        if lock is None:
            lock = asyncio.Lock()
            cls._locks[root_trace_id] = lock
        return lock

    async def manage(
        self,
        trace_id: str,
        *,
        operation: str,
        content: Any,
        parent_refs: list[CandidatePointer],
        effective_at_sequence: int,
    ) -> CandidateRef:
        trace, root = await self._require_owner_trace(trace_id)
        root_trace_id = root.trace_id
        operation_id = f"candidate:{trace_id}:{effective_at_sequence}"
        input_payload = {
            "operation": operation,
            "content": content,
            "parent_refs": [item.model_dump(mode="json") for item in parent_refs],
        }
        input_hash = sha256(
            canonical_json(input_payload).encode("utf-8")
        ).hexdigest()
        async with self._lock_for(root_trace_id):
            ledger = await self._load_ledger(root_trace_id)
            existing_operation = next(
                (
                    item
                    for item in ledger.operations
                    if item.operation_id == operation_id
                ),
                None,
            )
            if existing_operation is not None:
                if existing_operation.input_hash != input_hash:
                    raise CandidateStateError(
                        "Candidate operation payload changed during recovery"
                    )
                if existing_operation.status == "completed":
                    if existing_operation.candidate is None:
                        raise CandidateStateError(
                            "Completed candidate operation has no result"
                        )
                    return ledger.candidate(existing_operation.candidate)
                if existing_operation.status == "failed":
                    raise CandidateStateError(
                        existing_operation.error or "Candidate operation failed"
                    )
            else:
                try:
                    pointer = self._allocate_pointer(
                        operation,
                        trace_id,
                        effective_at_sequence,
                        parent_refs,
                        ledger,
                    )
                except ValueError as exc:
                    raise CandidateStateError(str(exc)) from exc
                operation_record = CandidateOperationRecord(
                    operation_id=operation_id,
                    operation=operation,
                    input_hash=input_hash,
                    candidate=pointer,
                )
                ledger = ledger.model_copy(update={
                    "operations": (*ledger.operations, operation_record),
                })
                await self.store.replace_candidate_ledger(
                    root_trace_id,
                    ledger.model_dump(mode="json"),
                )
                existing_operation = operation_record

            assert existing_operation.candidate is not None
            pointer = existing_operation.candidate
            try:
                parents = [ledger.candidate(item) for item in parent_refs]
                self._validate_parents(trace, root, parents)
                state = ensure_task_protocol(trace.context)
                request = CandidateVersionRequest(
                    operation_id=operation_id,
                    application_ref=self.binding.application_ref,
                    root_trace_id=root_trace_id,
                    owner_trace_id=trace_id,
                    uid=trace.uid,
                    candidate_id=pointer.candidate_id,
                    revision=pointer.revision,
                    task_contract_ref=task_contract_ref(
                        state,
                        root_task_anchor_hash=trace.context.get(
                            "root_task_anchor_hash"
                        ),
                    ),
                    parent_refs=tuple(parent_refs),
                    content=content,
                )
                artifact_ref = ArtifactRef.model_validate(
                    await self.repository.merge(request)
                    if operation == "merge"
                    else await self.repository.put_version(request)
                )
                materials, issues = await resolve_artifact_refs(
                    [artifact_ref],
                    resolver=self.artifact_resolver,
                    root_trace_id=root_trace_id,
                    uid=trace.uid,
                )
                if issues or len(materials) != 1:
                    reason = issues[0].reason if issues else "Artifact was not resolved"
                    raise CandidateStateError(
                        f"CandidateRepository returned an invalid ArtifactRef: {reason}"
                    )
                candidate_ref = CandidateRef(
                    **pointer.model_dump(mode="json"),
                    application_ref=self.binding.application_ref,
                    root_trace_id=root_trace_id,
                    owner_trace_id=trace_id,
                    task_contract_ref=request.task_contract_ref,
                    artifact_ref=artifact_ref,
                    parent_refs=tuple(parent_refs),
                    created_at_sequence=effective_at_sequence,
                )
            except Exception as exc:
                failed = self._fail_operation(ledger, operation_id, str(exc))
                await self.store.replace_candidate_ledger(
                    root_trace_id,
                    failed.model_dump(mode="json"),
                )
                raise
            completed = self._complete_operation(
                ledger,
                operation_id,
                candidate_ref,
            )
            # If this atomic publish fails after the Repository committed, the
            # pending operation remains recoverable. Re-execution uses the same
            # operation_id and the Repository's idempotency contract.
            await self.store.replace_candidate_ledger(
                root_trace_id,
                completed.model_dump(mode="json"),
            )
            return candidate_ref

    async def validate_report_refs(
        self,
        trace_id: str,
        refs: list[CandidateRef],
    ) -> None:
        trace, root = await self._require_owner_trace(trace_id)
        async with self._lock_for(root.trace_id):
            ledger = await self._load_ledger(root.trace_id)
            for ref in refs:
                persisted = ledger.candidate(ref)
                if persisted != ref:
                    raise CandidateStateError(
                        "TaskReport CandidateRef does not match the root ledger"
                    )
                self._validate_candidate_owner(persisted, trace, root)
                if ledger.current_state(ref) != "proposed":
                    raise CandidateStateError(
                        "TaskReport may only reference proposed candidates"
                    )

    async def _require_owner_trace(self, trace_id: str):
        trace = await self.store.get_trace(trace_id)
        if trace is None:
            raise CandidateStateError(f"Trace not found: {trace_id}")
        policy = require_mutable_trace_policy(trace.context)
        if not policy.requires_task_progress:
            raise CandidateStateError(
                "Candidates require Recursive revision 3"
            )
        application_ref = trace.context.get("application_ref")
        if application_ref != self.binding.application_ref.model_dump(mode="json"):
            raise CandidateStateError("Candidate Trace application binding mismatch")
        root_trace_id = trace.context.get("root_trace_id") or trace.trace_id
        root = trace if trace.trace_id == root_trace_id else await self.store.get_trace(
            root_trace_id
        )
        if root is None or root.uid != trace.uid:
            raise CandidateStateError("Candidate root or owner mismatch")
        if root.context.get("application_ref") != application_ref:
            raise CandidateStateError("Candidate root application mismatch")
        state = ensure_task_protocol(trace.context)
        if state.get("task_report") is not None:
            raise CandidateStateError("Candidates cannot change after TaskReport submission")
        if state.get("pending_reviews") or state.get("next_actions"):
            raise CandidateStateError("Resolve protocol lifecycle work before candidates")
        return trace, root

    async def _load_ledger(self, root_trace_id: str) -> CandidateLedger:
        raw = await self.store.get_candidate_ledger(root_trace_id)
        return CandidateLedger.model_validate(raw or {})

    def _allocate_pointer(
        self,
        operation: str,
        trace_id: str,
        sequence: int,
        parents: list[CandidatePointer],
        ledger: CandidateLedger,
    ) -> CandidatePointer:
        if operation == "create":
            if parents:
                raise CandidateStateError("create does not accept parent_refs")
            digest = sha256(f"{trace_id}:{sequence}".encode("utf-8")).hexdigest()
            return CandidatePointer(candidate_id=f"cand_{digest[:24]}", revision=1)
        if operation == "fork":
            if len(parents) != 1:
                raise CandidateStateError("fork requires exactly one parent")
            parent = ledger.candidate(parents[0])
            revisions = [
                item.revision
                for item in ledger.candidates
                if item.candidate_id == parent.candidate_id
            ]
            return CandidatePointer(
                candidate_id=parent.candidate_id,
                revision=max(revisions) + 1,
            )
        if operation == "merge":
            if len(parents) < 2:
                raise CandidateStateError("merge requires at least two parents")
            digest = sha256(f"{trace_id}:{sequence}:merge".encode("utf-8")).hexdigest()
            return CandidatePointer(candidate_id=f"cand_{digest[:24]}", revision=1)
        raise CandidateStateError(f"Unsupported candidate operation: {operation}")

    def _validate_parents(self, trace, root, parents: list[CandidateRef]) -> None:
        for parent in parents:
            self._validate_candidate_owner(parent, trace, root)

    def _validate_candidate_owner(self, candidate, trace, root) -> None:
        if (
            candidate.application_ref != self.binding.application_ref
            or candidate.root_trace_id != root.trace_id
            or candidate.owner_trace_id != trace.trace_id
        ):
            raise CandidateStateError(
                "Candidate belongs to another application, root, or Task"
            )

    @staticmethod
    def _complete_operation(
        ledger: CandidateLedger,
        operation_id: str,
        candidate: CandidateRef,
    ) -> CandidateLedger:
        operations = tuple(
            item.model_copy(update={"status": "completed", "error": None})
            if item.operation_id == operation_id
            else item
            for item in ledger.operations
        )
        lifecycle = (*ledger.lifecycle, CandidateLifecycleRecord(
            operation_id=operation_id,
            candidate=CandidatePointer(
                candidate_id=candidate.candidate_id,
                revision=candidate.revision,
            ),
            state="proposed",
            effective_at_sequence=candidate.created_at_sequence,
        ))
        return CandidateLedger(
            candidates=(*ledger.candidates, candidate),
            lifecycle=lifecycle,
            operations=operations,
            validations=ledger.validations,
        )

    @staticmethod
    def _fail_operation(
        ledger: CandidateLedger,
        operation_id: str,
        error: str,
    ) -> CandidateLedger:
        operations = tuple(
            item.model_copy(update={
                "status": "failed",
                "error": error[:2_000],
            })
            if item.operation_id == operation_id
            else item
            for item in ledger.operations
        )
        return ledger.model_copy(update={"operations": operations})