aios/active-runtime/app/aios_task_lineage_taskplane_adapter.py

from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime, timezone
import json
from pathlib import Path
import sqlite3
from typing import Callable, Mapping, Protocol, Sequence, runtime_checkable

if not __package__:
    from aios_task_lineage import RelationType, RootReviewRecord, lineage_from_metadata
    from aios_task_lineage_integration import (
        IntegrationError,
        LINEAGE_METADATA_KEY,
        RootActionDispatchReceipt,
        TaskAuthoritySnapshot,
        TaskPlaneLineagePort,
    )
else:
    from .aios_task_lineage import RelationType, RootReviewRecord, lineage_from_metadata
    from .aios_task_lineage_integration import (
        IntegrationError,
        LINEAGE_METADATA_KEY,
        RootActionDispatchReceipt,
        TaskAuthoritySnapshot,
        TaskPlaneLineagePort,
    )


class AuthorityAdapterError(IntegrationError):
    pass


def _parse_datetime(value: object) -> datetime:
    if isinstance(value, datetime):
        if value.tzinfo is None or value.utcoffset() is None:
            raise AuthorityAdapterError("authority datetime must be timezone-aware")
        return value.astimezone(timezone.utc)
    if not isinstance(value, str) or not value:
        raise AuthorityAdapterError("authority datetime missing")
    try:
        dt = datetime.fromisoformat(value.replace("Z", "+00:00"))
    except ValueError as exc:
        raise AuthorityAdapterError("authority datetime invalid") from exc
    if dt.tzinfo is None or dt.utcoffset() is None:
        raise AuthorityAdapterError("authority datetime must include timezone")
    return dt.astimezone(timezone.utc)


def _load_json_object(value: object, *, field: str) -> dict[str, object]:
    if value in (None, ""):
        return {}
    if isinstance(value, Mapping):
        return dict(value)
    if not isinstance(value, str):
        raise AuthorityAdapterError(f"{field} must be a JSON object")
    try:
        obj = json.loads(value)
    except json.JSONDecodeError as exc:
        raise AuthorityAdapterError(f"{field} is invalid JSON") from exc
    if not isinstance(obj, dict):
        raise AuthorityAdapterError(f"{field} must decode to an object")
    return obj


def _metadata_identity(metadata: Mapping[str, object]) -> tuple[str, str]:
    lineage = metadata.get(LINEAGE_METADATA_KEY)
    if isinstance(lineage, Mapping):
        owner = lineage.get("owner_principal_id")
        mission = lineage.get("mission_id")
        if isinstance(owner, str) and owner and isinstance(mission, str) and mission:
            return owner, mission

    # Pre-lineage/legacy compatibility. These keys are accepted only as
    # authority-provided metadata; missing values fail closed rather than being
    # invented by this adapter.
    owner = metadata.get("owner_principal_id")
    mission = metadata.get("mission_id")
    if not isinstance(owner, str) or not owner:
        raise AuthorityAdapterError("task owner_principal_id is not discoverable")
    if not isinstance(mission, str) or not mission:
        raise AuthorityAdapterError("task mission_id is not discoverable")
    return owner, mission


def _metadata_evidence(metadata: Mapping[str, object]) -> tuple[str, ...]:
    raw = metadata.get("evidence_refs")
    if raw is None:
        return ()
    if not isinstance(raw, list) or not all(isinstance(x, str) for x in raw):
        raise AuthorityAdapterError("evidence_refs must be a string list")
    return tuple(dict.fromkeys(raw))


@runtime_checkable
class TaskPlaneReadback(Protocol):
    """Independent authority readback surface.

    This may be backed by the TaskPlane query service, a read replica, or a
    read-only SQLite authority. It must not mutate task state.
    """

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

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

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

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


@dataclass(frozen=True)
class CallableTaskPlaneReadback:
    """Readback adapter for an existing official TaskPlane query/service."""

    get_task: Callable[[str], TaskAuthoritySnapshot]
    list_descendants: Callable[[str], Sequence[TaskAuthoritySnapshot]]
    get_review: Callable[[str], Mapping[str, object] | None] | None = None
    get_dispatch: Callable[[str], RootActionDispatchReceipt | None] | None = None

    def get_task_snapshot(self, task_id: str) -> TaskAuthoritySnapshot:
        snapshot = self.get_task(task_id)
        if not isinstance(snapshot, TaskAuthoritySnapshot):
            raise AuthorityAdapterError("readback returned invalid task snapshot")
        if snapshot.task_id != task_id:
            raise AuthorityAdapterError("readback returned wrong task identity")
        return snapshot

    def list_descendant_snapshots(
        self, root_task_id: str
    ) -> Sequence[TaskAuthoritySnapshot]:
        values = tuple(self.list_descendants(root_task_id))
        if not all(isinstance(v, TaskAuthoritySnapshot) for v in values):
            raise AuthorityAdapterError("readback returned invalid descendant snapshot")
        return values

    def get_root_review_record(self, review_id: str) -> Mapping[str, object] | None:
        if self.get_review is None:
            raise AuthorityAdapterError("independent review readback is not configured")
        return self.get_review(review_id)

    def get_root_action_dispatch(self, review_id: str) -> RootActionDispatchReceipt | None:
        if self.get_dispatch is None:
            raise AuthorityAdapterError("independent dispatch readback is not configured")
        return self.get_dispatch(review_id)


@dataclass(frozen=True)
class AuthoritySchemaReport:
    task_table: str
    columns: tuple[str, ...]
    required_columns_present: bool
    missing_columns: tuple[str, ...]
    task_count: int


class ReadOnlySQLiteTaskPlane:
    """Read-only discovery/readback helper for the current TaskPlane SQLite authority.

    This class never writes. It opens SQLite read-only, including committed WAL
    data, and exists only for discovery, diagnostics and independent readback.

    Execution writes must go through the live TaskPlane repository/service and
    are supplied separately through ``TaskPlaneServiceBindings``.
    """

    REQUIRED_TASK_COLUMNS = (
        "task_id",
        "status",
        "revision",
        "metadata_json",
        "created_at",
    )

    def __init__(self, db_path: Path):
        self.db_path = Path(db_path)

    def _connect(self) -> sqlite3.Connection:
        if not self.db_path.is_file():
            raise AuthorityAdapterError(f"TaskPlane DB missing: {self.db_path}")
        # A live WAL database is not immutable: ignoring its WAL can return
        # stale state. as_uri also quotes URI metacharacters in file names.
        uri = self.db_path.resolve().as_uri() + "?mode=ro"
        con = sqlite3.connect(uri, uri=True)
        con.row_factory = sqlite3.Row
        return con

    def inspect_schema(self) -> AuthoritySchemaReport:
        con = self._connect()
        try:
            rows = list(con.execute("PRAGMA table_info(tasks)"))
            columns = tuple(str(r["name"]) for r in rows)
            missing = tuple(c for c in self.REQUIRED_TASK_COLUMNS if c not in columns)
            count = int(con.execute("SELECT COUNT(*) FROM tasks").fetchone()[0]) if rows else 0
            return AuthoritySchemaReport(
                task_table="tasks",
                columns=columns,
                required_columns_present=not missing,
                missing_columns=missing,
                task_count=count,
            )
        finally:
            con.close()

    def _row_to_snapshot(self, row: sqlite3.Row) -> TaskAuthoritySnapshot:
        keys = set(row.keys())
        metadata = _load_json_object(
            row["metadata_json"] if "metadata_json" in keys else None,
            field="metadata_json",
        )
        owner, mission = _metadata_identity(metadata)
        created_value = row["created_at"] if "created_at" in keys else None
        return TaskAuthoritySnapshot(
            task_id=str(row["task_id"]),
            revision=int(row["revision"]),
            status=str(row["status"]),
            owner_principal_id=owner,
            mission_id=mission,
            created_at=_parse_datetime(created_value),
            metadata=metadata,
            evidence_refs=_metadata_evidence(metadata),
            # A metadata flag cannot stand in for execution-integrity proof.
            # The official service adapter must resolve verified receipts.
            terminal_evidence_proven=False,
        )

    def get_task_snapshot(self, task_id: str) -> TaskAuthoritySnapshot:
        con = self._connect()
        try:
            row = con.execute("SELECT * FROM tasks WHERE task_id=? LIMIT 1", (task_id,)).fetchone()
            if row is None:
                raise AuthorityAdapterError(f"task not found: {task_id}")
            return self._row_to_snapshot(row)
        finally:
            con.close()

    def list_descendant_snapshots(self, root_task_id: str) -> tuple[TaskAuthoritySnapshot, ...]:
        con = self._connect()
        try:
            rows = list(con.execute("SELECT * FROM tasks ORDER BY task_id"))
        finally:
            con.close()
        out: list[TaskAuthoritySnapshot] = []
        for row in rows:
            try:
                snap = self._row_to_snapshot(row)
            except AuthorityAdapterError:
                # Legacy tasks without discoverable identity cannot safely be
                # attached to a root by this read-only adapter.
                continue
            raw = snap.metadata.get(LINEAGE_METADATA_KEY)
            if (
                isinstance(raw, Mapping)
                and raw.get("root_task_id") == root_task_id
                and snap.task_id != root_task_id
            ):
                out.append(snap)
        return tuple(out)


@dataclass(frozen=True)
class TaskPlaneServiceBindings:
    """In-process bindings to the existing TaskPlane authority.

    These callables must be backed by the real TaskPlane repository/service,
    including its existing transactions, revision/CAS, journal and continuation
    mechanisms. They must NOT be implemented as out-of-band SQLite mutations.
    """

    get_task_snapshot: Callable[[str], TaskAuthoritySnapshot]
    list_descendant_snapshots: Callable[[str], Sequence[TaskAuthoritySnapshot]]
    reserve_sibling_sequence: Callable[[str, RelationType, str], int]
    compare_and_swap_task_metadata: Callable[[str, int, Mapping[str, object]], bool]
    get_root_review_record: Callable[[str], Mapping[str, object] | None]
    compare_and_swap_root_review_record: Callable[
        [str, int, Mapping[str, object]], bool
    ]
    get_root_action_dispatch: Callable[[str], RootActionDispatchReceipt | None]
    dispatch_root_action: Callable[
        [RootReviewRecord, str], RootActionDispatchReceipt
    ]


class InProcessTaskPlaneLineageAdapter(TaskPlaneLineagePort):
    """Concrete TaskPlaneLineagePort backed by existing TaskPlane service calls."""

    def __init__(
        self,
        bindings: TaskPlaneServiceBindings,
        *,
        readback: TaskPlaneReadback | None = None,
    ):
        self._b = bindings
        self._readback = readback

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

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

    def reserve_sibling_sequence(
        self,
        *,
        parent_task_id: str,
        relation_type: RelationType,
        child_task_id: str,
    ) -> int:
        seq = self._b.reserve_sibling_sequence(
            parent_task_id, relation_type, child_task_id
        )
        if not isinstance(seq, int) or seq < 1:
            raise AuthorityAdapterError("TaskPlane returned invalid sibling sequence")
        return seq

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

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

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

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

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

    def verify_task_readback(
        self,
        *,
        task_id: str,
        expected_revision_at_least: int,
        expected_metadata: Mapping[str, object],
    ) -> TaskAuthoritySnapshot:
        if self._readback is None:
            raise AuthorityAdapterError("independent TaskPlane readback is not configured")
        snap = self._readback.get_task_snapshot(task_id)
        if snap.task_id != task_id:
            raise AuthorityAdapterError("TaskPlane readback task identity mismatch")
        if snap.revision < expected_revision_at_least:
            raise AuthorityAdapterError("TaskPlane readback revision did not advance")
        for key, expected in expected_metadata.items():
            if snap.metadata.get(key) != expected:
                raise AuthorityAdapterError(
                    f"TaskPlane readback metadata mismatch for {key}"
                )
        return snap

    def verify_root_readback(
        self,
        *,
        root_task_id: str,
        expected_review_id: str,
    ) -> tuple[TaskAuthoritySnapshot, ...]:
        """Verify that independent authority descendants remain bound to root.

        The persisted ROOT_REVIEW record itself may live in a journal rather
        than the task row. This method verifies the task/lineage projection
        independently while the service binding verifies the journal receipt.
        """
        if self._readback is None:
            raise AuthorityAdapterError("independent TaskPlane readback is not configured")
        review = self.get_root_review_record(expected_review_id)
        dispatch = self.get_root_action_dispatch(expected_review_id)
        if (review is None or review.get("review_id") != expected_review_id
                or review.get("root_task_id") != root_task_id or dispatch is None):
            raise AuthorityAdapterError("root review or dispatch identity mismatch")
        read_review = getattr(self._readback, "get_root_review_record", None)
        read_dispatch = getattr(self._readback, "get_root_action_dispatch", None)
        if not callable(read_review) or not callable(read_dispatch):
            raise AuthorityAdapterError("independent review/dispatch readback is not configured")
        if read_review(expected_review_id) != review:
            raise AuthorityAdapterError("independent root review readback mismatch")
        if read_dispatch(expected_review_id) != dispatch:
            raise AuthorityAdapterError("independent dispatch readback mismatch")
        root = self._readback.get_task_snapshot(root_task_id)
        raw_root = root.metadata.get(LINEAGE_METADATA_KEY)
        if root.task_id != root_task_id or not isinstance(raw_root, Mapping):
            raise AuthorityAdapterError("root readback identity or lineage missing")
        root_lineage = lineage_from_metadata(dict(raw_root))
        if (root_lineage.relation_type != RelationType.ROOT
                or root_lineage.task_id != root_task_id
                or root_lineage.owner_principal_id != root.owner_principal_id
                or root_lineage.mission_id != root.mission_id):
            raise AuthorityAdapterError("root readback authority mismatch")
        revision = review.get("root_revision_after")
        if type(revision) is not int or root.revision < revision:
            raise AuthorityAdapterError("root readback revision is stale")
        descendants = tuple(self._readback.list_descendant_snapshots(root_task_id))
        seen = set()
        for snap in descendants:
            raw = snap.metadata.get(LINEAGE_METADATA_KEY)
            if not isinstance(raw, Mapping):
                raise AuthorityAdapterError("descendant lineage missing")
            lineage = lineage_from_metadata(dict(raw))
            if (snap.task_id in seen or snap.task_id == root_task_id
                    or lineage.task_id != snap.task_id
                    or lineage.root_task_id != root_task_id
                    or lineage.owner_principal_id != root.owner_principal_id
                    or lineage.mission_id != root.mission_id
                    or snap.owner_principal_id != root.owner_principal_id
                    or snap.mission_id != root.mission_id
                    or lineage.root_acceptance_ref != root_lineage.root_acceptance_ref):
                raise AuthorityAdapterError(
                    f"descendant {snap.task_id} is projected under wrong root"
                )
            seen.add(snap.task_id)
        if review.get("trigger_child_task_id") not in seen:
            raise AuthorityAdapterError("trigger child missing from independent readback")
        return descendants

    @property
    def has_independent_readback(self) -> bool:
        return self._readback is not None


@dataclass(frozen=True)
class AuthenticatedTaskContext:
    principal_id: str
    preferred_alias: str | None
    mission_id: str


def require_context_matches_task(
    context: AuthenticatedTaskContext,
    task: TaskAuthoritySnapshot,
) -> None:
    if context.principal_id != task.owner_principal_id:
        raise AuthorityAdapterError("authenticated principal does not own task")
    if context.mission_id != task.mission_id:
        raise AuthorityAdapterError("authenticated mission does not match task")