aios/active-runtime/app/aios_task_lineage_hooks.py

from __future__ import annotations

from dataclasses import dataclass, replace
from datetime import datetime
from typing import Mapping

if not __package__:
    from aios_task_lineage import AcceptanceEvaluator, RelationType, TaskLineage, lineage_to_metadata, lineage_from_metadata, root_review_to_record
    from aios_task_lineage_integration import (
        DISPLAY_NAME_METADATA_KEY,
        IntegrationError,
        LINEAGE_METADATA_KEY,
        NewTaskIdentityRequest,
        PreparedTaskIdentity,
        TaskIdentityLineageService,
        TerminalReviewOutcome,
    )
    from aios_task_lineage_taskplane_adapter import (
        AuthenticatedTaskContext,
        InProcessTaskPlaneLineageAdapter,
        require_context_matches_task,
    )
else:
    from .aios_task_lineage import AcceptanceEvaluator, RelationType, TaskLineage, lineage_to_metadata, lineage_from_metadata, root_review_to_record
    from .aios_task_lineage_integration import (
        DISPLAY_NAME_METADATA_KEY,
        IntegrationError,
        LINEAGE_METADATA_KEY,
        NewTaskIdentityRequest,
        PreparedTaskIdentity,
        TaskIdentityLineageService,
        TerminalReviewOutcome,
    )
    from .aios_task_lineage_taskplane_adapter import (
        AuthenticatedTaskContext,
        InProcessTaskPlaneLineageAdapter,
        require_context_matches_task,
    )


class LifecycleHookError(IntegrationError):
    pass


@dataclass(frozen=True)
class IdentityPersistenceOutcome:
    prepared: PreparedTaskIdentity
    persisted_revision_at_least: int
    independently_verified: bool


@dataclass(frozen=True)
class LegacyAttachOutcome:
    lineage: TaskLineage
    independently_verified: bool


@dataclass(frozen=True)
class TerminalHookOutcome:
    outcome: TerminalReviewOutcome
    root_review_persisted: bool
    dispatch_persisted: bool
    independently_verified: bool


class TaskLineageLifecycleHooks:
    """Concrete lifecycle hook coordinator.

    It does not own TaskPlane state. It orchestrates existing TaskPlane CAS,
    journal/continuation and optional independent authority readback.
    """

    def __init__(self, adapter: InProcessTaskPlaneLineageAdapter):
        self._adapter = adapter
        self._service = TaskIdentityLineageService(adapter)

    def persist_new_task_identity(
        self,
        *,
        request: NewTaskIdentityRequest,
        context: AuthenticatedTaskContext,
        max_cas_retries: int = 4,
        require_independent_readback: bool = True,
    ) -> IdentityPersistenceOutcome:
        if max_cas_retries < 1:
            raise LifecycleHookError("max_cas_retries must be >= 1")
        if require_independent_readback and not self._adapter.has_independent_readback:
            raise LifecycleHookError("independent TaskPlane readback is not configured")
        if request.owner_principal_id != context.principal_id:
            raise LifecycleHookError("request principal does not match authenticated context")
        if request.mission_id != context.mission_id:
            raise LifecycleHookError("request mission does not match authenticated context")

        for _ in range(max_cas_retries):
            current = self._adapter.get_task_snapshot(request.task_id)
            require_context_matches_task(context, current)
            if current.created_at != request.created_at:
                raise LifecycleHookError("request created_at does not match authority")
            if (request.relation_type == RelationType.ROOT
                    and request.preferred_owner_alias not in (None, context.preferred_alias)):
                raise LifecycleHookError("request alias does not match authenticated context")
            request = replace(request, preferred_owner_alias=context.preferred_alias)
            prepared = self._service.prepare_new_task_identity(request)

            # Idempotent replay and legacy display-name repair are fail-closed:
            # never overwrite an existing conflicting lineage record.
            raw = current.metadata.get(LINEAGE_METADATA_KEY)
            display = current.metadata.get(DISPLAY_NAME_METADATA_KEY)
            if raw is not None and not isinstance(raw, Mapping):
                raise LifecycleHookError("existing lineage metadata is malformed")
            if isinstance(raw, Mapping):
                # parent_revision records the creation observation, not the
                # parent's current progress revision on an idempotent replay.
                expected = dict(prepared.metadata_patch[LINEAGE_METADATA_KEY])
                if request.parent_task_id is not None:
                    stored = lineage_from_metadata(dict(raw))
                    if (stored.parent_revision is None
                            or stored.parent_revision > prepared.lineage.parent_revision):
                        raise LifecycleHookError("existing parent revision conflicts with authority")
                    expected["parent_revision"] = stored.parent_revision
                if raw != expected:
                    raise LifecycleHookError(
                        "existing lineage metadata conflicts with requested identity"
                    )
                prepared = replace(prepared, lineage=lineage_from_metadata(dict(raw)),
                                   metadata_patch={LINEAGE_METADATA_KEY: dict(raw),
                                                   DISPLAY_NAME_METADATA_KEY: prepared.display_name})
                if display == prepared.display_name:
                    verified = False
                    if require_independent_readback:
                        self._adapter.verify_task_readback(
                            task_id=request.task_id,
                            expected_revision_at_least=current.revision,
                            expected_metadata=prepared.metadata_patch,
                        )
                        verified = True
                    return IdentityPersistenceOutcome(
                        prepared=prepared,
                        persisted_revision_at_least=current.revision,
                        independently_verified=verified,
                    )
                if display not in (None, ""):
                    raise LifecycleHookError(
                        "existing display name conflicts with requested identity"
                    )
                metadata_patch = {
                    DISPLAY_NAME_METADATA_KEY: prepared.display_name,
                }
            else:
                if display not in (None, ""):
                    raise LifecycleHookError(
                        "display name exists without authoritative lineage"
                    )
                metadata_patch = prepared.metadata_patch

            ok = self._adapter.compare_and_swap_task_metadata(
                task_id=request.task_id,
                expected_revision=current.revision,
                metadata_patch=metadata_patch,
            )
            if not ok:
                continue

            verified = False
            if require_independent_readback:
                self._adapter.verify_task_readback(
                    task_id=request.task_id,
                    expected_revision_at_least=current.revision + 1,
                    expected_metadata=prepared.metadata_patch,
                )
                verified = True
            return IdentityPersistenceOutcome(
                prepared=prepared,
                persisted_revision_at_least=current.revision + 1,
                independently_verified=verified,
            )

        raise LifecycleHookError("identity persistence CAS retry budget exhausted")

    def attach_legacy_identity(
        self,
        *,
        task_id: str,
        parent_task_id: str,
        relation_type: RelationType,
        goal_slug: str,
        context: AuthenticatedTaskContext,
        max_cas_retries: int = 4,
        require_independent_readback: bool = True,
    ) -> LegacyAttachOutcome:
        if require_independent_readback and not self._adapter.has_independent_readback:
            raise LifecycleHookError("independent TaskPlane readback is not configured")
        before = self._adapter.get_task_snapshot(task_id)
        require_context_matches_task(context, before)
        lineage = self._service.attach_legacy_task(
            task_id=task_id,
            parent_task_id=parent_task_id,
            relation_type=relation_type,
            goal_slug=goal_slug,
            max_cas_retries=max_cas_retries,
            expected_owner_principal_id=context.principal_id,
            expected_mission_id=context.mission_id,
        )
        verified = False
        if require_independent_readback:
            self._adapter.verify_task_readback(
                task_id=task_id,
                expected_revision_at_least=before.revision,
                expected_metadata={
                    LINEAGE_METADATA_KEY: lineage_to_metadata(lineage),
                    DISPLAY_NAME_METADATA_KEY: lineage.display_name,
                },
            )
            # Domain serialization compatibility check below avoids requiring
            # TaskLineage to expose a method in older versions.
            verified = True
        return LegacyAttachOutcome(lineage=lineage, independently_verified=verified)

    def on_proven_terminal_child(
        self,
        *,
        child_task_id: str,
        context: AuthenticatedTaskContext,
        acceptance_evaluator: AcceptanceEvaluator,
        reviewed_at: datetime,
        max_cas_retries: int = 4,
        require_independent_readback: bool = True,
    ) -> TerminalHookOutcome:
        if require_independent_readback and not self._adapter.has_independent_readback:
            raise LifecycleHookError("independent TaskPlane readback is not configured")
        child = self._adapter.get_task_snapshot(child_task_id)
        require_context_matches_task(context, child)
        if not child.terminal_evidence_proven:
            raise LifecycleHookError(
                "terminal child hook requires proven terminal evidence"
            )
        raw = child.metadata.get(LINEAGE_METADATA_KEY)
        if not isinstance(raw, Mapping):
            raise LifecycleHookError("terminal child lineage missing")
        lineage = lineage_from_metadata(dict(raw))
        root = self._adapter.get_task_snapshot(lineage.root_task_id)
        require_context_matches_task(context, root)

        def authorized_acceptance(root_record, descendants):
            if (root_record.lineage.owner_principal_id != context.principal_id
                    or root_record.lineage.mission_id != context.mission_id):
                raise LifecycleHookError("root authority changed during terminal review")
            return acceptance_evaluator(root_record, descendants)

        outcome = self._service.on_proven_terminal_child(
            child_task_id=child_task_id,
            acceptance_evaluator=authorized_acceptance,
            reviewed_at=reviewed_at,
            max_cas_retries=max_cas_retries,
        )

        persisted_review = self._adapter.get_root_review_record(
            outcome.review.review_id
        )
        if persisted_review is None:
            raise LifecycleHookError("root review was not persisted")
        if persisted_review != root_review_to_record(outcome.review):
            raise LifecycleHookError("persisted review does not match returned review")

        dispatch = self._adapter.get_root_action_dispatch(outcome.review.review_id)
        if dispatch is None:
            raise LifecycleHookError("root action dispatch was not persisted")
        if dispatch != outcome.dispatch:
            raise LifecycleHookError("persisted dispatch does not match returned receipt")

        verified = False
        if require_independent_readback:
            self._adapter.verify_root_readback(
                root_task_id=outcome.review.root_task_id,
                expected_review_id=outcome.review.review_id,
            )
            verified = True

        return TerminalHookOutcome(
            outcome=outcome,
            root_review_persisted=True,
            dispatch_persisted=True,
            independently_verified=verified,
        )