aios/active-runtime/app/aios_task_lineage_integration.py

from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime
from enum import Enum
from typing import Mapping, MutableMapping, Protocol, Sequence

if not __package__:
    from aios_task_lineage import (
        AcceptanceEvaluator,
        LineageError,
        NextLegalAction,
        RelationType,
        RootReviewRecord,
        RootReviewStore,
        TaskLineage,
        TaskRecord,
        attach_legacy_lineage,
        derive_safe_owner_alias,
        lineage_from_metadata,
        lineage_to_metadata,
        make_child_lineage,
        make_root_lineage,
        normalize_goal_slug,
        review_root_after_child_terminal,
        root_review_from_record,
        root_review_to_record,
    )
else:
    from .aios_task_lineage import (
        AcceptanceEvaluator,
        LineageError,
        NextLegalAction,
        RelationType,
        RootReviewRecord,
        RootReviewStore,
        TaskLineage,
        TaskRecord,
        attach_legacy_lineage,
        derive_safe_owner_alias,
        lineage_from_metadata,
        lineage_to_metadata,
        make_child_lineage,
        make_root_lineage,
        normalize_goal_slug,
        review_root_after_child_terminal,
        root_review_from_record,
        root_review_to_record,
    )


LINEAGE_METADATA_KEY = "aios_task_lineage"
LINEAGE_REQUEST_DIGEST_KEY = "aios_lineage_request_digest"
ROOT_REVIEW_RECORD_KEY = "aios_root_review"
DISPLAY_NAME_METADATA_KEY = "aios_task_display_name"


class IntegrationError(RuntimeError):
    pass


@dataclass(frozen=True)
class TaskAuthoritySnapshot:
    task_id: str
    revision: int
    status: str
    owner_principal_id: str
    mission_id: str
    created_at: datetime
    metadata: Mapping[str, object]
    evidence_refs: tuple[str, ...] = ()
    terminal_evidence_proven: bool = False


@dataclass(frozen=True)
class NewTaskIdentityRequest:
    task_id: str
    owner_principal_id: str
    preferred_owner_alias: str | None
    created_at: datetime
    mission_id: str
    goal_slug: str
    root_acceptance_ref: str
    relation_type: RelationType = RelationType.ROOT
    parent_task_id: str | None = None


@dataclass(frozen=True)
class PreparedTaskIdentity:
    task_id: str
    display_name: str
    lineage: TaskLineage
    metadata_patch: Mapping[str, object]


@dataclass(frozen=True)
class TerminalReviewOutcome:
    review: RootReviewRecord
    dispatch: "RootActionDispatchReceipt"


@dataclass(frozen=True)
class RootActionDispatchReceipt:
    review_id: str
    root_task_id: str
    action: NextLegalAction
    idempotency_key: str
    status: str
    mission_id: str = ""
    root_revision: int = 0
    plan_digest: str | None = None
    execution_key: str | None = None


PERSISTED_DISPATCH_STATUSES = frozenset(
    {
        "PENDING",
        "CLAIMED",
        "EFFECT_UNKNOWN",
        "EXECUTED_PENDING_REREVIEW",
        "WAITING_DEPENDENCY",
        "WAITING_OWNER",
        "CLOSED",
        "SUPERSEDED",
        "BLOCKED",
        # Read compatibility only.  A consumer must reconcile this legacy
        # state unless a complete durable receipt is present.
        "DISPATCHED",
    }
)


class TaskPlaneLineagePort(Protocol):
    """Narrow authority port required by the lineage integration service.

    The live adapter must implement these methods over the existing TaskPlane /
    mission / journal / CAS mechanisms. This protocol deliberately avoids
    guessing concrete DB tables, HTTP endpoints, or handler names.
    """

    def get_task_snapshot(self, task_id: str) -> TaskAuthoritySnapshot: ...

    def list_descendant_snapshots(
        self, root_task_id: str
    ) -> Sequence[TaskAuthoritySnapshot]: ...

    def reserve_sibling_sequence(
        self,
        *,
        parent_task_id: str,
        relation_type: RelationType,
        child_task_id: str,
    ) -> int:
        """Idempotently reserve one sibling sequence for child_task_id."""

    def compare_and_swap_task_metadata(
        self,
        *,
        task_id: str,
        expected_revision: int,
        metadata_patch: Mapping[str, object],
    ) -> bool: ...

    def get_root_review_record(
        self, review_id: str
    ) -> Mapping[str, object] | None: ...

    def compare_and_swap_root_review_record(
        self,
        *,
        root_task_id: str,
        expected_root_revision: int,
        review_record: Mapping[str, object],
    ) -> bool: ...

    def get_root_action_dispatch(
        self, review_id: str
    ) -> RootActionDispatchReceipt | None: ...

    def dispatch_root_action(
        self,
        *,
        review: RootReviewRecord,
        idempotency_key: str,
    ) -> RootActionDispatchReceipt: ...


def _extract_lineage(snapshot: TaskAuthoritySnapshot) -> TaskLineage:
    raw = snapshot.metadata.get(LINEAGE_METADATA_KEY)
    if not isinstance(raw, dict):
        raise IntegrationError(
            f"task {snapshot.task_id} does not have authoritative lineage metadata"
        )
    lineage = lineage_from_metadata(raw)
    if lineage.task_id != snapshot.task_id:
        raise IntegrationError("lineage task_id does not match authority task")
    if lineage.owner_principal_id != snapshot.owner_principal_id:
        raise IntegrationError("lineage owner does not match authority owner")
    if lineage.mission_id != snapshot.mission_id:
        raise IntegrationError("lineage mission does not match authority mission")
    if snapshot.created_at.tzinfo is None or snapshot.created_at.utcoffset() is None:
        raise IntegrationError("authority created_at must be timezone-aware")
    if lineage.created_at != snapshot.created_at:
        raise IntegrationError("lineage created_at does not match authority task")
    display = snapshot.metadata.get(DISPLAY_NAME_METADATA_KEY)
    if display not in (None, "", lineage.display_name):
        raise IntegrationError("lineage display name conflicts with authority metadata")
    return lineage


def _as_task_record(snapshot: TaskAuthoritySnapshot) -> TaskRecord:
    return TaskRecord(
        task_id=snapshot.task_id,
        revision=snapshot.revision,
        status=snapshot.status,
        lineage=_extract_lineage(snapshot),
        evidence_refs=snapshot.evidence_refs,
        terminal_evidence_proven=snapshot.terminal_evidence_proven,
    )


def _validate_record_root_binding(record: TaskRecord, root: TaskRecord) -> None:
    if root.lineage.relation_type != RelationType.ROOT:
        raise IntegrationError("root_task_id does not resolve to a ROOT task")
    if record.lineage.root_task_id != root.task_id:
        raise IntegrationError("task lineage points to a different root")
    if record.lineage.owner_principal_id != root.lineage.owner_principal_id:
        raise IntegrationError("task lineage owner does not match root owner")
    if record.lineage.mission_id != root.lineage.mission_id:
        raise IntegrationError("task lineage mission does not match root mission")
    if record.lineage.root_acceptance_ref != root.lineage.root_acceptance_ref:
        raise IntegrationError("task lineage acceptance reference does not match root")


def _load_authority_record(
    port: TaskPlaneLineagePort, task_id: str
) -> tuple[TaskRecord, TaskRecord]:
    record = _as_task_record(port.get_task_snapshot(task_id))
    if record.lineage.relation_type == RelationType.ROOT:
        return record, record
    root = _as_task_record(port.get_task_snapshot(record.lineage.root_task_id))
    current = record
    seen = {record.task_id}
    while current.lineage.relation_type != RelationType.ROOT:
        _validate_record_root_binding(current, root)
        parent_id = current.lineage.parent_task_id
        if not parent_id or parent_id in seen:
            raise IntegrationError("task lineage contains a parent cycle")
        seen.add(parent_id)
        parent = _as_task_record(port.get_task_snapshot(parent_id))
        _validate_record_root_binding(parent, root)
        if current.lineage.parent_task_id != parent.task_id:
            raise IntegrationError("task lineage parent does not match authority parent")
        if current.lineage.lineage_depth != parent.lineage.lineage_depth + 1:
            raise IntegrationError("task lineage depth does not match authority parent")
        if current.lineage.parent_revision is None or (
            current.lineage.parent_revision > parent.revision
        ):
            raise IntegrationError("task lineage parent revision is not authoritative")
        current = parent
    if current.task_id != root.task_id:
        raise IntegrationError("task lineage parent chain ends at a different root")
    return record, root


def _validate_dispatch_receipt(
    receipt: RootActionDispatchReceipt,
    *,
    review: RootReviewRecord,
    idempotency_key: str,
    existing: bool,
) -> None:
    if (
        receipt.review_id != review.review_id
        or receipt.root_task_id != review.root_task_id
        or receipt.action != review.next_legal_action
        or receipt.idempotency_key != idempotency_key
        or receipt.status not in PERSISTED_DISPATCH_STATUSES
    ):
        noun = "existing root-action dispatch" if existing else "root-action dispatch receipt"
        raise IntegrationError(f"{noun} conflicts with persisted review")


class _RootReviewStoreAdapter(RootReviewStore):
    def __init__(self, port: TaskPlaneLineagePort):
        self._port = port

    def get_task(self, task_id: str) -> TaskRecord:
        record, _ = _load_authority_record(self._port, task_id)
        return record

    def list_descendants(self, root_task_id: str) -> Sequence[TaskRecord]:
        root, _ = _load_authority_record(self._port, root_task_id)
        if root.lineage.relation_type != RelationType.ROOT:
            raise IntegrationError("descendant lookup key is not a ROOT task")
        records: list[TaskRecord] = []
        for snapshot in self._port.list_descendant_snapshots(root_task_id):
            record, resolved_root = _load_authority_record(self._port, snapshot.task_id)
            if resolved_root.task_id != root_task_id:
                raise IntegrationError("descendant authority returned a cross-root task")
            _validate_record_root_binding(record, root)
            records.append(record)
        return tuple(records)

    def get_root_review(self, review_id: str) -> RootReviewRecord | None:
        raw = self._port.get_root_review_record(review_id)
        if raw is None:
            return None
        if not isinstance(raw, Mapping):
            raise IntegrationError("persisted root-review record must be a mapping")
        try:
            return root_review_from_record(
                raw,
                expected_review_id=review_id,
            )
        except LineageError as exc:
            raise IntegrationError(
                f"invalid persisted root-review record: {exc}"
            ) from exc

    def compare_and_swap_root_review(
        self,
        root_task_id: str,
        expected_root_revision: int,
        review: RootReviewRecord,
    ) -> bool:
        return self._port.compare_and_swap_root_review_record(
            root_task_id=root_task_id,
            expected_root_revision=expected_root_revision,
            review_record=root_review_to_record(review),
        )


class TaskIdentityLineageService:
    """Application service that wires the lineage domain to TaskPlane authority."""

    def __init__(self, port: TaskPlaneLineagePort):
        self._port = port

    def prepare_new_task_identity(
        self, request: NewTaskIdentityRequest
    ) -> PreparedTaskIdentity:
        alias = derive_safe_owner_alias(
            request.owner_principal_id, request.preferred_owner_alias
        )

        if request.relation_type == RelationType.ROOT:
            if request.parent_task_id is not None:
                raise IntegrationError("ROOT task cannot have parent_task_id")
            lineage = make_root_lineage(
                task_id=request.task_id,
                owner_principal_id=request.owner_principal_id,
                owner_alias=alias,
                created_at=request.created_at,
                mission_id=request.mission_id,
                goal_slug=request.goal_slug,
                root_acceptance_ref=request.root_acceptance_ref,
            )
        else:
            if not request.parent_task_id:
                raise IntegrationError("non-root task requires parent_task_id")
            parent, root = _load_authority_record(
                self._port, request.parent_task_id
            )

            if request.owner_principal_id != root.lineage.owner_principal_id:
                raise IntegrationError("cross-owner child creation rejected")
            if request.mission_id != root.lineage.mission_id:
                raise IntegrationError("cross-mission child creation rejected")
            if request.root_acceptance_ref != root.lineage.root_acceptance_ref:
                raise IntegrationError("root acceptance reference mismatch")

            sequence = self._port.reserve_sibling_sequence(
                parent_task_id=parent.task_id,
                relation_type=request.relation_type,
                child_task_id=request.task_id,
            )
            lineage = make_child_lineage(
                task_id=request.task_id,
                created_at=request.created_at,
                parent=parent,
                root=root,
                relation_type=request.relation_type,
                sibling_sequence=sequence,
                goal_slug=request.goal_slug,
            )

        return PreparedTaskIdentity(
            task_id=request.task_id,
            display_name=lineage.display_name,
            lineage=lineage,
            metadata_patch={
                LINEAGE_METADATA_KEY: lineage_to_metadata(lineage),
                DISPLAY_NAME_METADATA_KEY: lineage.display_name,
            },
        )

    def attach_legacy_task(
        self,
        *,
        task_id: str,
        parent_task_id: str,
        relation_type: RelationType,
        goal_slug: str,
        expected_owner_principal_id: str | None = None,
        expected_mission_id: str | None = None,
        max_cas_retries: int = 4,
    ) -> TaskLineage:
        if max_cas_retries < 1:
            raise IntegrationError("max_cas_retries must be >= 1")

        for _ in range(max_cas_retries):
            current = self._port.get_task_snapshot(task_id)
            if (
                expected_owner_principal_id is not None
                and current.owner_principal_id != expected_owner_principal_id
            ):
                raise IntegrationError("legacy task owner changed after authorization")
            if (
                expected_mission_id is not None
                and current.mission_id != expected_mission_id
            ):
                raise IntegrationError("legacy task mission changed after authorization")
            raw = current.metadata.get(LINEAGE_METADATA_KEY)
            if isinstance(raw, Mapping):
                record, _ = _load_authority_record(self._port, task_id)
                lineage = record.lineage
                expected_goal_slug = normalize_goal_slug(goal_slug)
                if (
                    lineage.parent_task_id != parent_task_id
                    or lineage.relation_type != relation_type
                    or lineage.goal_slug != expected_goal_slug
                ):
                    raise IntegrationError(
                        "existing legacy lineage conflicts with requested binding"
                    )
                display = current.metadata.get(DISPLAY_NAME_METADATA_KEY)
                if display == lineage.display_name:
                    return lineage
                if display not in (None, ""):
                    raise IntegrationError(
                        "legacy display name conflicts with authoritative lineage"
                    )
                if self._port.compare_and_swap_task_metadata(
                    task_id=current.task_id,
                    expected_revision=current.revision,
                    metadata_patch={
                        DISPLAY_NAME_METADATA_KEY: lineage.display_name,
                    },
                ):
                    return lineage
                continue

            parent, root = _load_authority_record(self._port, parent_task_id)
            if root.status == "COMPLETED":
                raise IntegrationError("cannot attach lineage below a completed root")
            if (
                expected_owner_principal_id is not None
                and root.lineage.owner_principal_id != expected_owner_principal_id
            ):
                raise IntegrationError("legacy root owner does not match authorization")
            if (
                expected_mission_id is not None
                and root.lineage.mission_id != expected_mission_id
            ):
                raise IntegrationError("legacy root mission does not match authorization")
            if current.owner_principal_id != root.lineage.owner_principal_id:
                raise IntegrationError("legacy task owner does not match root owner")
            if current.mission_id != root.lineage.mission_id:
                raise IntegrationError("legacy task mission does not match root mission")

            sequence = self._port.reserve_sibling_sequence(
                parent_task_id=parent.task_id,
                relation_type=relation_type,
                child_task_id=current.task_id,
            )
            lineage = attach_legacy_lineage(
                existing_task_id=current.task_id,
                created_at=current.created_at,
                parent=parent,
                root=root,
                relation_type=relation_type,
                sibling_sequence=sequence,
                goal_slug=goal_slug,
            )
            if self._port.compare_and_swap_task_metadata(
                task_id=current.task_id,
                expected_revision=current.revision,
                metadata_patch={
                LINEAGE_METADATA_KEY: lineage_to_metadata(lineage),
                DISPLAY_NAME_METADATA_KEY: lineage.display_name,
            },
            ):
                return lineage

        raise IntegrationError("legacy lineage CAS retry budget exhausted")

    def on_proven_terminal_child(
        self,
        *,
        child_task_id: str,
        acceptance_evaluator: AcceptanceEvaluator,
        reviewed_at: datetime,
        max_cas_retries: int = 4,
    ) -> TerminalReviewOutcome:
        store = _RootReviewStoreAdapter(self._port)
        review = review_root_after_child_terminal(
            store=store,
            trigger_child_task_id=child_task_id,
            acceptance_evaluator=acceptance_evaluator,
            reviewed_at=reviewed_at,
            max_cas_retries=max_cas_retries,
        )
        existing = self._port.get_root_action_dispatch(review.review_id)
        if existing is not None:
            idempotency_key = (
                f"root-review:{review.review_id}:{review.next_legal_action.value}"
            )
            _validate_dispatch_receipt(
                existing,
                review=review,
                idempotency_key=idempotency_key,
                existing=True,
            )
            return TerminalReviewOutcome(review=review, dispatch=existing)

        idempotency_key = f"root-review:{review.review_id}:{review.next_legal_action.value}"
        receipt = self._port.dispatch_root_action(
            review=review,
            idempotency_key=idempotency_key,
        )
        _validate_dispatch_receipt(
            receipt,
            review=review,
            idempotency_key=idempotency_key,
            existing=False,
        )
        return TerminalReviewOutcome(review=review, dispatch=receipt)