Agent/cyber_agent/application/events.py

"""Canonical root event journal and ordered, failure-isolated projection pump."""

from __future__ import annotations

import asyncio
import logging
from typing import Any
from weakref import WeakValueDictionary

from pydantic import Field

from cyber_agent.application.candidate import CandidateLedger
from cyber_agent.application.models import ApplicationModel, ApplicationRef
from cyber_agent.application.ports import RunEventProjector
from cyber_agent.core.task_protocol import ensure_task_protocol


logger = logging.getLogger(__name__)
CURRENT_RUN_EVENT_SCHEMA_VERSION = 1
SUPPORTED_RUN_EVENT_TYPES = frozenset({
    "task.progress.revised",
    "candidate.version_registered",
    "candidate.lifecycle_changed",
    "validation.completed",
    "task.review.recorded",
    "run.rewound",
})
_ROOT_LOCKS: "WeakValueDictionary[tuple[int, str], asyncio.Lock]" = (
    WeakValueDictionary()
)


class RunEvent(ApplicationModel):
    event_id: int = Field(gt=0)
    schema_version: int = Field(ge=1)
    event_type: str = Field(min_length=1, max_length=200)
    event_key: str = Field(min_length=1, max_length=500)
    application_ref: ApplicationRef
    root_trace_id: str = Field(min_length=1)
    source_trace_id: str = Field(min_length=1)
    effective_at_sequence: int = Field(ge=0)
    payload: dict[str, Any]


class RunEventPump:
    """Project canonical events in journal order without touching Agent state."""

    def __init__(
        self,
        *,
        store: Any,
        application_ref: ApplicationRef,
        projector: RunEventProjector | None,
    ) -> None:
        self.store = store
        self.application_ref = application_ref
        self.projector = projector

    async def pump(self, root_trace_id: str) -> int:
        if self.projector is None:
            return 0
        cursor = await self.projector.load_cursor(
            self.application_ref,
            root_trace_id,
        )
        projected = 0
        for raw in await self.store.get_events(root_trace_id, cursor):
            if raw.get("event") != "run_event":
                continue
            event = RunEvent.model_validate({
                key: raw[key]
                for key in RunEvent.model_fields
                if key in raw
            })
            if (
                event.application_ref != self.application_ref
                or event.root_trace_id != root_trace_id
            ):
                raise ValueError("run_event application/root identity mismatch")
            if (
                event.schema_version != CURRENT_RUN_EVENT_SCHEMA_VERSION
                or event.event_type not in SUPPORTED_RUN_EVENT_TYPES
            ):
                await self.projector.advance_cursor(
                    self.application_ref,
                    root_trace_id,
                    event.event_id,
                    cursor,
                )
            else:
                await self.projector.project(event, cursor)
            next_cursor = await self.projector.load_cursor(
                self.application_ref,
                root_trace_id,
            )
            if next_cursor != event.event_id:
                raise ValueError(
                    "RunEventProjector did not atomically advance its cursor"
                )
            cursor = next_cursor
            projected += 1
        return projected


class RunEventService:
    """Append idempotent canonical events after authoritative state commits."""

    def __init__(
        self,
        *,
        store: Any,
        application_ref: ApplicationRef,
        projector: RunEventProjector | None = None,
    ) -> None:
        self.store = store
        self.application_ref = application_ref
        self.pump = RunEventPump(
            store=store,
            application_ref=application_ref,
            projector=projector,
        )

    @staticmethod
    def _lock(root_trace_id: str) -> asyncio.Lock:
        key = (id(asyncio.get_running_loop()), root_trace_id)
        lock = _ROOT_LOCKS.get(key)
        if lock is None:
            lock = asyncio.Lock()
            _ROOT_LOCKS[key] = lock
        return lock

    async def _emit_record(
        self,
        *,
        source_trace_id: str,
        event_type: str,
        event_key: str,
        effective_at_sequence: int,
        payload: dict[str, Any],
        project: bool = True,
    ) -> tuple[int, bool]:
        source = await self.store.get_trace(source_trace_id)
        if source is None:
            raise ValueError(f"Trace not found: {source_trace_id}")
        root_trace_id = source.context.get("root_trace_id") or source.trace_id
        if source.context.get("application_ref") != self.application_ref.model_dump(
            mode="json"
        ):
            raise ValueError("run_event source application binding mismatch")
        async with self._lock(root_trace_id):
            existing = next(
                (
                    item for item in await self.store.get_events(root_trace_id)
                    if item.get("event") == "run_event"
                    and item.get("event_key") == event_key
                ),
                None,
            )
            if existing is not None:
                expected = {
                    "schema_version": CURRENT_RUN_EVENT_SCHEMA_VERSION,
                    "event_type": event_type,
                    "event_key": event_key,
                    "application_ref": self.application_ref.model_dump(
                        mode="json"
                    ),
                    "root_trace_id": root_trace_id,
                    "source_trace_id": source_trace_id,
                    "effective_at_sequence": effective_at_sequence,
                    "payload": payload,
                }
                actual = {
                    key: existing.get(key)
                    for key in expected
                }
                if actual != expected:
                    raise ValueError(
                        f"run_event key was reused with conflicting payload: {event_key}"
                    )
                event_id = int(existing["event_id"])
                created = False
            else:
                event_id = await self.store.append_event(
                    root_trace_id,
                    "run_event",
                    {
                        "schema_version": CURRENT_RUN_EVENT_SCHEMA_VERSION,
                        "event_type": event_type,
                        "event_key": event_key,
                        "application_ref": self.application_ref.model_dump(
                            mode="json"
                        ),
                        "root_trace_id": root_trace_id,
                        "source_trace_id": source_trace_id,
                        "effective_at_sequence": effective_at_sequence,
                        "payload": payload,
                    },
                )
                created = True
        if project:
            await self.try_pump(root_trace_id)
        return event_id, created

    async def emit(self, **kwargs: Any) -> int:
        event_id, _created = await self._emit_record(**kwargs)
        return event_id

    async def _emit_missing(self, **kwargs: Any) -> int:
        _event_id, created = await self._emit_record(**kwargs)
        return int(created)

    async def emit_after_commit(self, **kwargs: Any) -> None:
        """Never roll back protocol state because journaling/projection failed."""
        try:
            await self.emit(project=False, **kwargs)
        except Exception:
            logger.exception("Failed to emit/project canonical run event")

    async def try_pump(self, root_trace_id: str) -> None:
        try:
            await self.pump.pump(root_trace_id)
        except Exception:
            logger.exception(
                "RunEventProjector stopped at its current cursor for root %s",
                root_trace_id,
            )

    async def reconcile_root(self, root_trace_id: str) -> int:
        """Rebuild missing canonical events from current authoritative records."""
        traces = [
            item for item in await self.store.list_traces(limit=10_000)
            if (item.context.get("root_trace_id") or item.trace_id) == root_trace_id
            and item.context.get("application_ref")
            == self.application_ref.model_dump(mode="json")
        ]
        emitted = 0
        for trace in traces:
            state = ensure_task_protocol(trace.context)
            for revision in state.get("task_progress_revisions", []):
                emitted += await self._emit_missing(
                    source_trace_id=trace.trace_id,
                    event_type="task.progress.revised",
                    event_key=(
                        f"task.progress.revised:{trace.trace_id}:"
                        f"{revision['revision']}"
                    ),
                    effective_at_sequence=revision["effective_at_sequence"],
                    payload={"revision": revision},
                    project=False,
                )
            cache = state.get("task_report_validation")
            if isinstance(cache, dict) and cache.get("aggregate_result"):
                emitted += await self._emit_missing(
                    source_trace_id=trace.trace_id,
                    event_type="validation.completed",
                    event_key=(
                        f"validation.completed:trace:{trace.trace_id}:"
                        f"{cache['plan_hash']}"
                    ),
                    effective_at_sequence=int(
                        cache.get("validated_at_sequence", 0) or 0
                    ),
                    payload={
                        "subject_type": "trace",
                        "validation_result": cache["aggregate_result"],
                    },
                    project=False,
                )
            for review in state.get("reviews", []):
                sequence = int(review.get("reviewed_at_sequence", 0) or 0)
                emitted += await self._emit_missing(
                    source_trace_id=trace.trace_id,
                    event_type="task.review.recorded",
                    event_key=(
                        f"task.review.recorded:{trace.trace_id}:"
                        f"{review.get('child_trace_id')}:{sequence}"
                    ),
                    effective_at_sequence=sequence,
                    payload={"task_review": review},
                    project=False,
                )
            for operation in trace.context.get("run_event_operations", []):
                if operation.get("event_type") != "run.rewound":
                    continue
                emitted += await self._emit_missing(
                    source_trace_id=trace.trace_id,
                    event_type="run.rewound",
                    event_key=f"run.rewound:{operation['operation_id']}",
                    effective_at_sequence=operation["effective_at_sequence"],
                    payload=operation["payload"],
                    project=False,
                )
            for raw in await self.store.get_events(trace.trace_id):
                if raw.get("event") != "rewind":
                    continue
                sequence = int(raw.get("after_sequence", 0) or 0)
                emitted += await self._emit_missing(
                    source_trace_id=trace.trace_id,
                    event_type="run.rewound",
                    event_key=(
                        f"run.rewound:{raw['operation_id']}"
                        if raw.get("operation_id")
                        else f"run.rewound:{trace.trace_id}:"
                        f"{raw['event_id']}:{sequence}"
                    ),
                    effective_at_sequence=sequence,
                    payload={"after_sequence": sequence},
                    project=False,
                )
        ledger = CandidateLedger.model_validate(
            await self.store.get_candidate_ledger(root_trace_id) or {}
        )
        for candidate in ledger.candidates:
            emitted += await self._emit_missing(
                source_trace_id=candidate.owner_trace_id,
                event_type="candidate.version_registered",
                event_key=(
                    f"candidate.version_registered:{candidate.candidate_id}:"
                    f"{candidate.revision}"
                ),
                effective_at_sequence=candidate.created_at_sequence,
                payload={"candidate_ref": candidate.model_dump(mode="json")},
                project=False,
            )
        for lifecycle in ledger.lifecycle:
            candidate = ledger.candidate(lifecycle.candidate)
            emitted += await self._emit_missing(
                source_trace_id=lifecycle.source_trace_id or candidate.owner_trace_id,
                event_type="candidate.lifecycle_changed",
                event_key=(
                    f"candidate.lifecycle_changed:{lifecycle.operation_id}:"
                    f"{lifecycle.state}"
                ),
                effective_at_sequence=lifecycle.effective_at_sequence,
                payload={"lifecycle": lifecycle.model_dump(mode="json")},
                project=False,
            )
        for raw in ledger.validations:
            candidate = raw["candidate_ref"]
            emitted += await self._emit_missing(
                source_trace_id=candidate["owner_trace_id"],
                event_type="validation.completed",
                event_key=(
                    f"validation.completed:candidate:{candidate['candidate_id']}:"
                    f"{candidate['revision']}:{raw['plan_hash']}"
                ),
                effective_at_sequence=raw["validated_at_sequence"],
                payload={
                    "subject_type": "candidate",
                    "candidate_ref": candidate,
                    "validation_result": raw["validation_result"],
                },
                project=False,
            )
        await self.try_pump(root_trace_id)
        return emitted