Agent/cyber_agent/application/candidate.py

"""Application-owned candidate content behind framework-owned immutable references."""

from __future__ import annotations

from typing import Any, Literal, Protocol, runtime_checkable

from pydantic import Field, model_validator

from cyber_agent.application.models import ApplicationModel, ApplicationRef, canonical_json
from cyber_agent.core.artifacts import ArtifactRef
from cyber_agent.core.task_contract import TaskContractRef


MAX_CANDIDATES_PER_ROOT = 256
MAX_CANDIDATE_OPERATIONS_PER_ROOT = 2_048
MAX_CANDIDATE_VALIDATIONS_PER_ROOT = 2_048


class CandidatePointer(ApplicationModel):
    candidate_id: str = Field(min_length=1, max_length=200)
    revision: int = Field(ge=1)


class CandidateRef(CandidatePointer):
    schema_version: Literal[1] = 1
    application_ref: ApplicationRef
    root_trace_id: str = Field(min_length=1)
    owner_trace_id: str = Field(min_length=1)
    task_contract_ref: TaskContractRef
    artifact_ref: ArtifactRef
    parent_refs: tuple[CandidatePointer, ...] = ()
    created_at_sequence: int = Field(ge=0)

    @model_validator(mode="after")
    def validate_parent_refs(self) -> "CandidateRef":
        identities = [
            (item.candidate_id, item.revision)
            for item in self.parent_refs
        ]
        if len(identities) != len(set(identities)):
            raise ValueError("CandidateRef parent_refs must be unique")
        if (self.candidate_id, self.revision) in set(identities):
            raise ValueError("CandidateRef cannot reference itself")
        return self


class CandidateVersionRequest(ApplicationModel):
    operation_id: str = Field(min_length=1, max_length=300)
    application_ref: ApplicationRef
    root_trace_id: str = Field(min_length=1)
    owner_trace_id: str = Field(min_length=1)
    uid: str | None = None
    candidate_id: str = Field(min_length=1, max_length=200)
    revision: int = Field(ge=1)
    task_contract_ref: TaskContractRef
    parent_refs: tuple[CandidatePointer, ...] = ()
    content: Any

    @model_validator(mode="after")
    def validate_content(self) -> "CandidateVersionRequest":
        canonical_json(self.content)
        return self


class CandidateAdoptionRequest(ApplicationModel):
    operation_id: str = Field(min_length=1, max_length=300)
    candidate_ref: CandidateRef


class CandidateAdoptionReceipt(ApplicationModel):
    operation_id: str = Field(min_length=1, max_length=300)
    external_id: str = Field(min_length=1, max_length=500)
    committed: bool = True


class CandidateReviewAction(ApplicationModel):
    action: Literal["adopt", "discard", "revise"]
    candidate_ref: CandidateRef
    reason: str = Field(min_length=1, max_length=2_000)


@runtime_checkable
class CandidateRepository(Protocol):
    async def put_version(
        self,
        request: CandidateVersionRequest,
    ) -> ArtifactRef: ...

    async def load_version(self, candidate_ref: CandidateRef) -> Any: ...

    async def merge(
        self,
        request: CandidateVersionRequest,
    ) -> ArtifactRef: ...

    async def commit_adoption(
        self,
        request: CandidateAdoptionRequest,
    ) -> CandidateAdoptionReceipt: ...


class CandidateLifecycleRecord(ApplicationModel):
    operation_id: str = Field(min_length=1, max_length=300)
    candidate: CandidatePointer
    state: Literal[
        "proposed",
        "adoption_pending",
        "adopted",
        "discarded",
        "adoption_failed",
    ]
    effective_at_sequence: int = Field(ge=0)
    source_trace_id: str | None = None
    reason: str | None = Field(default=None, max_length=2_000)
    receipt: CandidateAdoptionReceipt | None = None


class CandidateOperationRecord(ApplicationModel):
    operation_id: str = Field(min_length=1, max_length=300)
    # Stable framework command identity.  Candidate tool commands persist the
    # assistant tool_call_id here so retries of one call are idempotent while
    # two distinct calls with identical arguments remain distinct commands.
    command_id: str | None = Field(default=None, min_length=1, max_length=300)
    operation: Literal[
        "create", "fork", "merge", "adopt", "discard", "revise",
        "review_batch",
    ]
    input_hash: str = Field(pattern=r"^[0-9a-f]{64}$")
    status: Literal["pending", "completed", "failed"] = "pending"
    candidate: CandidatePointer | None = None
    error: str | None = Field(default=None, max_length=2_000)


class CandidateLedger(ApplicationModel):
    schema_version: Literal[1] = 1
    candidates: tuple[CandidateRef, ...] = ()
    lifecycle: tuple[CandidateLifecycleRecord, ...] = ()
    operations: tuple[CandidateOperationRecord, ...] = ()
    validations: tuple[dict[str, Any], ...] = ()
    validation_checkpoints: tuple[dict[str, Any], ...] = ()

    @model_validator(mode="after")
    def validate_ledger(self) -> "CandidateLedger":
        if len(self.candidates) > MAX_CANDIDATES_PER_ROOT:
            raise ValueError("Candidate ledger candidate limit exceeded")
        if len(self.operations) > MAX_CANDIDATE_OPERATIONS_PER_ROOT:
            raise ValueError("Candidate ledger operation limit exceeded")
        if len(self.validations) > MAX_CANDIDATE_VALIDATIONS_PER_ROOT:
            raise ValueError("Candidate ledger validation limit exceeded")
        if len(self.validation_checkpoints) > MAX_CANDIDATE_VALIDATIONS_PER_ROOT:
            raise ValueError("Candidate ledger validation checkpoint limit exceeded")
        identities = [
            (item.candidate_id, item.revision)
            for item in self.candidates
        ]
        if len(identities) != len(set(identities)):
            raise ValueError("Candidate ledger contains duplicate revisions")
        operation_ids = [item.operation_id for item in self.operations]
        if len(operation_ids) != len(set(operation_ids)):
            raise ValueError("Candidate ledger contains duplicate operation IDs")
        from cyber_agent.application.quality import (
            CandidateValidationCheckpoint,
            CandidateValidationRecord,
        )

        candidate_by_key = {
            (item.candidate_id, item.revision): item
            for item in self.candidates
        }
        validation_keys = []
        for raw in self.validations:
            validation = CandidateValidationRecord.model_validate(raw)
            key = (
                validation.candidate_ref.candidate_id,
                validation.candidate_ref.revision,
            )
            if candidate_by_key.get(key) != validation.candidate_ref:
                raise ValueError(
                    "Candidate ledger validation references an altered candidate"
                )
            validation_keys.append((key, validation.plan_hash))
        if len(validation_keys) != len(set(validation_keys)):
            raise ValueError("Candidate ledger contains duplicate validations")

        checkpoint_keys = []
        for raw in self.validation_checkpoints:
            checkpoint = CandidateValidationCheckpoint.model_validate(raw)
            candidate_key = (
                checkpoint.candidate_ref.candidate_id,
                checkpoint.candidate_ref.revision,
            )
            if candidate_by_key.get(candidate_key) != checkpoint.candidate_ref:
                raise ValueError(
                    "Candidate ledger checkpoint references an altered candidate"
                )
            checkpoint_keys.append((
                checkpoint.candidate_ref.candidate_id,
                checkpoint.candidate_ref.revision,
                checkpoint.plan_hash,
            ))
        if len(checkpoint_keys) != len(set(checkpoint_keys)):
            raise ValueError("Candidate ledger contains duplicate validation checkpoints")
        known = set(identities)
        parents_by_candidate: dict[
            tuple[str, int], tuple[tuple[str, int], ...]
        ] = {}
        for candidate in self.candidates:
            candidate_key = (candidate.candidate_id, candidate.revision)
            parent_keys: list[tuple[str, int]] = []
            for parent in candidate.parent_refs:
                parent_key = (parent.candidate_id, parent.revision)
                if parent_key not in known:
                    raise ValueError("Candidate ledger contains an unknown parent")
                if (
                    parent.candidate_id == candidate.candidate_id
                    and parent.revision >= candidate.revision
                ):
                    raise ValueError("Candidate lineage references a future revision")
                parent_keys.append(parent_key)
            parents_by_candidate[candidate_key] = tuple(parent_keys)

        visited: set[tuple[str, int]] = set()
        visiting: set[tuple[str, int]] = set()

        def visit(candidate_key: tuple[str, int]) -> None:
            if candidate_key in visiting:
                raise ValueError("Candidate ledger contains a lineage cycle")
            if candidate_key in visited:
                return
            visiting.add(candidate_key)
            for parent_key in parents_by_candidate[candidate_key]:
                visit(parent_key)
            visiting.remove(candidate_key)
            visited.add(candidate_key)

        for candidate_key in parents_by_candidate:
            visit(candidate_key)
        return self

    def candidate(self, pointer: CandidatePointer) -> CandidateRef:
        for item in self.candidates:
            if (
                item.candidate_id == pointer.candidate_id
                and item.revision == pointer.revision
            ):
                return item
        raise ValueError(
            f"Candidate revision not found: {pointer.candidate_id}@{pointer.revision}"
        )

    def current_state(self, pointer: CandidatePointer) -> str | None:
        state = None
        for item in self.lifecycle:
            if (
                item.candidate.candidate_id == pointer.candidate_id
                and item.candidate.revision == pointer.revision
            ):
                state = item.state
        return state