rein-aharness/rein_aharness/claim_loop.py

"""Continuous ops_run claim worker (REIN-A-0002-T03).

  rein-aharness claim-loop
  rein-aharness claim-loop --once
  rein-aharness poll --source=ops-run

Concurrency default 1. Heartbeats while an approach runs.
"""

from __future__ import annotations

import logging
import os
import signal
import threading
import time
from dataclasses import dataclass, field
from typing import Any

from rein_aharness import metrics
from rein_aharness.approaches import (
    APPROACH_UNMATCHED,
    ApproachResult,
    execute_approach,
    select_approach,
)
from rein_aharness.close_outbox import (
    CloseOutbox,
    CloseRequest,
    PermanentCloseDeliveryError,
    ReplayReport,
)
from rein_aharness.execution_cancel import (
    ExecutionCancel,
    ExecutionCancelled,
    using_cancel,
)
from rein_aharness.glas_execution import (
    GLAS_APPROACH,
    GlasExecutionError,
    execute_profiled_run,
    normalise_execution_evidence_for_close,
)
from rein_aharness.lease_monitor import LeaseLoss, LeaseMonitor
from rein_aharness.repository_transaction import (
    DirtyRepositoryError,
    GitRepositoryError,
    RepositoryAcceptanceError,
    RepositoryBusyError,
    RepositoryTransaction,
    RepositoryTransactionError,
)
from rein_aharness.taskspec import TaskSpecError
from rein_aharness.ops_run_client import (
    ActivityCoreOpsClient,
    OpsRun,
    OpsRunConfig,
    resolve_ops_target,
    OpsRunError,
)

logger = logging.getLogger("rein_aharness.claim_loop")

_PERMANENT_CLOSE_CODES = frozenset(
    {
        "not_found",
        "wrong_owner",
        "expired_lease",
        "state_conflict",
        "terminal_conflict",
        "evidence_conflict",
    }
)


@dataclass
class ProcessResult:
    claimed: bool
    empty: bool = False
    retry_full_interval: bool = False
    run_id: str | None = None
    approach: str | None = None
    ok: bool | None = None
    reason: str = ""
    ops_state: str | None = None
    detail: dict[str, Any] = field(default_factory=dict)


def _heartbeat_interval(lease_seconds: int) -> float:
    # Heartbeat at 1/3 lease, min 30s, max 300s
    return max(30.0, min(300.0, lease_seconds / 3.0))


_active_run_cancel: ExecutionCancel | None = None
_active_run_cancel_lock = threading.Lock()


def _set_active_run_cancel(cancel: ExecutionCancel | None) -> None:
    global _active_run_cancel
    with _active_run_cancel_lock:
        _active_run_cancel = cancel


def _cancel_active_run(reason: str) -> None:
    with _active_run_cancel_lock:
        cancel = _active_run_cancel
    if cancel is not None:
        cancel.cancel(reason)


class _Heartbeat:
    def __init__(
        self,
        client: ActivityCoreOpsClient,
        run_id: str,
        lease_seconds: int,
        cancel: ExecutionCancel | None = None,
    ):
        self._client = client
        self._run_id = run_id
        self._lease = lease_seconds
        self._stop = threading.Event()
        self._thread: threading.Thread | None = None
        self.cancel = cancel if cancel is not None else ExecutionCancel()
        self.monitor = LeaseMonitor(self._on_lost)

    def _on_lost(self, _loss: LeaseLoss) -> None:
        self.cancel.cancel("lease-loss")

    def start(self) -> None:
        interval = _heartbeat_interval(self._lease)

        def _loop() -> None:
            while not self._stop.wait(interval):
                try:
                    self._client.heartbeat(self._run_id, lease_seconds=self._lease)
                    logger.info("heartbeat ok run_id=%s", self._run_id)
                except OpsRunError as exc:
                    if exc.lease_rejected:
                        evidence = self.monitor.mark_lost(
                            f"{type(exc).__name__}:{exc.status_code}"
                        )
                        logger.warning(
                            "heartbeat lost run_id=%s error_type=%s observed_at=%s",
                            self._run_id,
                            evidence.error_type,
                            evidence.observed_at,
                        )
                    else:
                        logger.warning(
                            "heartbeat transient failure run_id=%s error_type=%s",
                            self._run_id,
                            type(exc).__name__,
                        )

        self._thread = threading.Thread(
            target=_loop, name=f"ops-hb-{self._run_id[:8]}", daemon=True
        )
        self._thread.start()

    def stop(self) -> None:
        self._stop.set()
        if self._thread is not None:
            self._thread.join(timeout=5.0)
        self.monitor.stop()


def _cancelled_result(
    run_id: str,
    approach: str | None,
    heartbeat: _Heartbeat,
) -> ProcessResult | None:
    """Skip Activity Core close after lease loss or execution cancellation."""
    loss = heartbeat.monitor.evidence
    if heartbeat.monitor.lost:
        assert loss is not None
        return ProcessResult(
            claimed=True,
            run_id=run_id,
            approach=approach,
            ok=False,
            reason=f"lease lost ({loss.error_type})",
            detail={
                "lease_loss": {
                    "error_type": loss.error_type,
                    "observed_at": loss.observed_at,
                }
            },
        )
    if heartbeat.cancel.cancelled:
        reason = heartbeat.cancel.reason or "cancelled"
        return ProcessResult(
            claimed=True,
            run_id=run_id,
            approach=approach,
            ok=False,
            reason=f"execution cancelled ({reason})",
            detail={"cancellation": heartbeat.cancel.evidence()},
        )
    return None


def _deliver_close(client: ActivityCoreOpsClient, request: CloseRequest) -> OpsRun:
    if request.worker_id != client.config.worker_id:
        raise PermanentCloseDeliveryError("outbox worker identity no longer matches")
    try:
        if request.action == "complete":
            return client.complete(request.run_id, result=request.result)
        return client.fail(
            request.run_id,
            error=request.error,
            reopen=request.reopen,
            result=request.result,
        )
    except OpsRunError as exc:
        if exc.code in _PERMANENT_CLOSE_CODES:
            raise PermanentCloseDeliveryError(
                f"Activity Core refused close with {exc.code}"
            ) from exc
        raise


def replay_close_outbox(
    client: ActivityCoreOpsClient,
    outbox: CloseOutbox | None = None,
    *,
    limit: int = 100,
) -> ReplayReport:
    """Replay close-only work; this function never dispatches a workload."""
    store = outbox or CloseOutbox()
    return store.replay(lambda request: _deliver_close(client, request), limit=limit)


def _durable_close(
    client: ActivityCoreOpsClient,
    outbox: CloseOutbox,
    request: CloseRequest,
) -> tuple[bool, str | None, str]:
    """Enqueue before delivery and report the durable outcome for this entry."""
    receipt = outbox.enqueue(request)
    delivered: dict[str, OpsRun] = {}

    def deliver(item: CloseRequest) -> None:
        delivered[item.entry_id] = _deliver_close(client, item)

    report = outbox.replay(deliver)
    state = outbox.entry_state(receipt.entry_id)
    if state != "delivered":
        reason = (
            "close evidence quarantined"
            if state == "quarantined"
            else "close evidence remains pending"
        )
        return False, None, reason
    row = delivered.get(receipt.entry_id)
    if row is not None:
        return True, row.state, ""
    expected_state = "succeeded" if request.action == "complete" else "failed"
    if report.delivered:
        logger.info("replayed %s close evidence entries", report.delivered)
    return True, expected_state, ""


def _bounded_reason(value: str, *, default: str) -> str:
    cleaned = " ".join(str(value).split())
    return cleaned[:500] or default


def process_one(
    client: ActivityCoreOpsClient | None = None,
    *,
    report_to_hub: bool = True,
    commit: bool = True,
    dry_run: bool = False,
    outbox: CloseOutbox | None = None,
) -> ProcessResult:
    """Claim at most one ops_run, execute approach, complete or fail."""
    client = client or ActivityCoreOpsClient()
    cfg = client.config
    outbox = outbox or CloseOutbox()

    replay = replay_close_outbox(client, outbox)
    if replay.remaining:
        return ProcessResult(
            claimed=False,
            retry_full_interval=True,
            reason=(
                "required Activity Core close evidence remains pending "
                f"(attempted={replay.attempted} remaining={replay.remaining})"
            ),
        )
    if replay.quarantined:
        logger.error("quarantined %s close evidence entries", replay.quarantined)

    try:
        claimed = client.claim(limit=1)
    except OpsRunError as exc:
        return ProcessResult(
            claimed=False,
            retry_full_interval=True,
            reason=f"claim error: {exc}",
        )

    if not claimed:
        return ProcessResult(claimed=False, empty=True, reason="queue empty")

    run = claimed[0]
    approach = GLAS_APPROACH if run.harness_profile_ref else select_approach(run)
    logger.info(
        "claimed run_id=%s approach=%s title=%r labels=%s",
        run.id,
        approach,
        run.title,
        run.labels,
    )

    if dry_run:
        try:
            client.fail(
                run.id,
                error="dry-run: not executing",
                reopen=True,
                result={"approach": approach, "dry_run": True},
            )
        except OpsRunError as exc:
            return ProcessResult(
                claimed=True,
                run_id=run.id,
                approach=approach,
                ok=False,
                reason=f"dry-run reopen failed: {exc}",
            )
        return ProcessResult(
            claimed=True,
            run_id=run.id,
            approach=approach,
            ok=True,
            reason="dry-run reopened",
            ops_state="open",
            detail={"dry_run": True},
        )

    if run.harness_profile_ref:
        return _process_profiled_run(
            client,
            run,
            report_to_hub=report_to_hub,
            outbox=outbox,
        )

    if run.repository_grant is not None:
        reason = (
            "refused: repository_grant requires an authoritative "
            "harness_profile_ref"
        )
        try:
            out = client.fail(
                run.id,
                error=reason,
                reopen=False,
                result={"ok": False, "approach": approach, "reason": reason},
            )
        except OpsRunError as exc:
            return ProcessResult(
                claimed=True,
                run_id=run.id,
                approach=approach,
                ok=False,
                reason=f"grant-route refusal close failed: {exc}",
            )
        return ProcessResult(
            claimed=True,
            run_id=run.id,
            approach=approach,
            ok=False,
            reason=reason,
            ops_state=out.state,
        )

    cancel = ExecutionCancel()
    hb = _Heartbeat(client, run.id, cfg.lease_seconds, cancel=cancel)
    hb.start()
    execution_error: str | None = None
    _set_active_run_cancel(cancel)
    try:
        try:
            with using_cancel(cancel):
                ar: ApproachResult = execute_approach(
                    run,
                    approach=approach,
                    config=cfg,
                    report_to_hub=report_to_hub,
                    commit=commit,
                )
        except Exception as exc:
            execution_error = type(exc).__name__
    finally:
        hb.stop()
        _set_active_run_cancel(None)

    refused = _cancelled_result(run.id, approach, hb)
    if refused is not None:
        return refused

    if execution_error is not None:
        return ProcessResult(
            claimed=True,
            run_id=run.id,
            approach=approach,
            ok=False,
            reason=f"approach failed ({execution_error})",
            detail={"execution_error_type": execution_error},
        )

    payload = {
        "approach": ar.approach,
        "ok": ar.ok,
        "reason": ar.reason,
        "target_repo": run.target_repo,
        **(ar.result or {}),
    }
    # Prefer approach-provided target_repo if set
    if ar.result and ar.result.get("target_repo"):
        payload["target_repo"] = ar.result["target_repo"]

    try:
        if ar.ok:
            # skipped_existing still succeeds the ops_run (idempotent day)
            out = client.complete(run.id, result=payload)
            state = out.state
            logger.info("completed run_id=%s approach=%s", run.id, ar.approach)
        else:
            reopen = ar.reopen and ar.approach != APPROACH_UNMATCHED
            out = client.fail(
                run.id,
                error=ar.reason or "approach failed",
                reopen=reopen,
                result=payload,
            )
            state = out.state
            logger.warning(
                "failed run_id=%s approach=%s reopen=%s reason=%s",
                run.id,
                ar.approach,
                reopen,
                ar.reason,
            )
    except OpsRunError as exc:
        return ProcessResult(
            claimed=True,
            run_id=run.id,
            approach=ar.approach,
            ok=False,
            reason=f"close ops_run failed: {exc}; approach_ok={ar.ok} {ar.reason}",
            detail=payload,
        )

    return ProcessResult(
        claimed=True,
        run_id=run.id,
        approach=ar.approach,
        ok=ar.ok,
        reason=ar.reason,
        ops_state=state,
        detail=payload,
    )


def _process_profiled_run(
    client: ActivityCoreOpsClient,
    run: OpsRun,
    *,
    report_to_hub: bool,
    outbox: CloseOutbox,
) -> ProcessResult:
    """Keep the claim heartbeat alive through evidence persistence and close."""
    cancel = ExecutionCancel()
    hb = _Heartbeat(client, run.id, client.config.lease_seconds, cancel=cancel)
    hb.start()
    _set_active_run_cancel(cancel)
    try:
        return _process_profiled_run_active(
            client,
            run,
            report_to_hub=report_to_hub,
            outbox=outbox,
            cancel=cancel,
            hb=hb,
        )
    finally:
        hb.stop()
        _set_active_run_cancel(None)


def _process_profiled_run_active(
    client: ActivityCoreOpsClient,
    run: OpsRun,
    *,
    report_to_hub: bool,
    outbox: CloseOutbox,
    cancel: ExecutionCancel,
    hb: _Heartbeat,
) -> ProcessResult:
    """Execute a profiled row without consulting or falling back to legacy routing."""
    execution_error: str | None = None
    execution_reason = ""
    execution_reopen = False
    gateway_result: dict[str, Any] | None = None
    safe_evidence: dict[str, Any] = {}
    execution_ok = False
    tx: RepositoryTransaction | None = None
    tx_evidence: dict[str, Any] | None = None
    try:
        target = resolve_ops_target(run, client.config)
        with using_cancel(cancel):
            with RepositoryTransaction(target, correlation_id=run.id) as active_tx:
                tx = active_tx
                gateway_result = execute_profiled_run(
                    run,
                    client.config,
                    report_to_hub=report_to_hub,
                    cancel=cancel,
                )
                safe_evidence = normalise_execution_evidence_for_close(
                    gateway_result["evidence"]
                )
                execution_ok = gateway_result["ok"]
                execution_reason = _bounded_reason(
                    str(
                        safe_evidence.get("error")
                        or safe_evidence.get("outcome")
                        or ""
                    ),
                    default="Glas execution failed",
                )
                if run.repository_grant is not None:
                    try:
                        active_tx.validate_acceptance(
                            run.repository_grant.acceptance_policy()
                        )
                    except RepositoryAcceptanceError as exc:
                        if execution_ok or exc.code != "head-unchanged":
                            execution_ok = False
                            execution_error = type(exc).__name__
                            execution_reason = _bounded_reason(
                                str(exc),
                                default="repository acceptance failed",
                            )
    except ExecutionCancelled as exc:
        execution_error = type(exc).__name__
        execution_reason = _bounded_reason(str(exc), default="execution cancelled")
    except DirtyRepositoryError as exc:
        execution_error = type(exc).__name__
        execution_reason = _bounded_reason(
            f"refused: {exc}", default="dirty repository refused"
        )
    except RepositoryBusyError as exc:
        execution_error = type(exc).__name__
        execution_reason = _bounded_reason(
            f"refused: {exc}", default="busy repository refused"
        )
        execution_reopen = True
    except (GitRepositoryError, RepositoryTransactionError, TaskSpecError) as exc:
        execution_error = type(exc).__name__
        execution_reason = _bounded_reason(
            f"refused: {exc}", default="repository transaction refused"
        )
    except GlasExecutionError:
        execution_error = "GlasExecutionError"
        execution_reason = "profiled execution failed (GlasExecutionError)"

    if tx is not None and tx.baseline is not None:
        tx_evidence = tx.evidence()
        if run.repository_grant is not None:
            tx_evidence["repository_grant"] = run.repository_grant.evidence()

    if run.repository_grant is not None and tx is not None:
        try:
            accepted = tx.acceptance
            commit_sha = safe_evidence.get("commit_sha")
            metrics.record_external_execution(
                tx.repo,
                "rein-aharness",
                success=(
                    execution_error is None
                    and execution_ok
                    and not cancel.cancelled
                    and not hb.monitor.lost
                ),
                committed=bool(accepted or commit_sha),
                head_after=(
                    accepted.head
                    if accepted is not None
                    else commit_sha if isinstance(commit_sha, str) else None
                ),
                reason=execution_reason or None,
                metadata={
                    "ops_run_id": run.id,
                    "harness_profile_ref": run.harness_profile_ref,
                    "repository_grant_id": run.repository_grant.grant_id,
                },
                session_id=tx.transaction_id,
            )
            assert tx_evidence is not None
            tx_evidence["metrics"] = {
                "storage": "external",
                "session_id": tx.transaction_id,
                "projection_ready": True,
            }
        except OSError:
            execution_ok = False
            execution_error = "ExternalMetricsError"
            execution_reason = "required external metrics persistence failed"

    # Re-check cancellation after metrics persistence. The heartbeat remains
    # active through this point and through close delivery, so a late signal or
    # lease loss cannot turn a cancelled execution into a successful close.
    refused = _cancelled_result(run.id, GLAS_APPROACH, hb)
    if refused is not None:
        if tx_evidence is not None:
            refused.detail["repository_transaction"] = tx_evidence
        if execution_error == "ExternalMetricsError":
            refused.reason = (
                f"{refused.reason}; required external metrics persistence failed"
            )
        if hb.monitor.lost or tx is None or tx_evidence is None:
            return refused
        payload = {
            "ok": False,
            "approach": GLAS_APPROACH,
            "reason": refused.reason,
            "repository_transaction": tx_evidence,
        }
        try:
            delivered, state, close_reason = _durable_close(
                client,
                outbox,
                CloseRequest(
                    run_id=run.id,
                    transaction_id=tx.transaction_id,
                    worker_id=client.config.worker_id,
                    action="fail",
                    result=payload,
                    error=refused.reason,
                    reopen=False,
                ),
            )
        except (OSError, RuntimeError) as exc:
            refused.reason = f"close evidence failure ({type(exc).__name__})"
            return refused
        if delivered:
            refused.ops_state = state
        else:
            refused.reason = close_reason
        return refused

    ok = execution_error is None and execution_ok
    reason = execution_reason or (
        "profiled execution failed"
        if not ok
        else ""
    )
    payload: dict[str, Any] = {
        "ok": ok,
        "approach": GLAS_APPROACH,
    }
    if safe_evidence:
        payload["execution_evidence"] = safe_evidence
    if reason:
        payload["reason"] = reason
    if tx_evidence is not None:
        payload["repository_transaction"] = tx_evidence

    if tx is None or tx_evidence is None:
        try:
            out = client.fail(
                run.id,
                error=reason,
                reopen=execution_reopen,
                result=payload,
            )
        except OpsRunError as close_exc:
            return ProcessResult(
                claimed=True,
                run_id=run.id,
                approach=GLAS_APPROACH,
                ok=False,
                reason=f"close ops_run failed: {close_exc}; {reason}",
                detail=payload,
            )
        return ProcessResult(
            claimed=True,
            run_id=run.id,
            approach=GLAS_APPROACH,
            ok=False,
            reason=reason,
            ops_state=out.state,
            detail=payload,
        )

    action = "complete" if ok else "fail"
    try:
        close_request = CloseRequest(
            run_id=run.id,
            transaction_id=tx.transaction_id,
            worker_id=client.config.worker_id,
            action=action,
            result=payload,
            error=reason if action == "fail" else "",
            reopen=False,
        )
        delivered, state, close_reason = _durable_close(
            client,
            outbox,
            close_request,
        )
    except (OSError, RuntimeError) as exc:
        return ProcessResult(
            claimed=True,
            run_id=run.id,
            approach=GLAS_APPROACH,
            ok=False,
            reason=f"close evidence failure ({type(exc).__name__})",
            detail=payload,
        )

    if not delivered:
        return ProcessResult(
            claimed=True,
            run_id=run.id,
            approach=GLAS_APPROACH,
            ok=False,
            reason=close_reason,
            detail=payload,
        )

    return ProcessResult(
        claimed=True,
        run_id=run.id,
        approach=GLAS_APPROACH,
        ok=ok,
        reason="" if ok else reason,
        ops_state=state,
        detail=payload,
    )


def poll_peek(client: ActivityCoreOpsClient | None = None) -> list[dict[str, Any]]:
    """List open ops_runs with selected approach (no claim)."""
    client = client or ActivityCoreOpsClient()
    rows = client.list_open()
    out = []
    for run in rows:
        out.append(
            {
                "id": run.id,
                "title": run.title,
                "state": run.state,
                "labels": run.labels,
                "target_repo": run.target_repo,
                "approach": (
                    GLAS_APPROACH if run.harness_profile_ref else select_approach(run)
                ),
                "harness_profile_ref": run.harness_profile_ref,
                "created_at": run.raw.get("created_at"),
            }
        )
    return out


def run_claim_loop(
    *,
    once: bool = False,
    interval_seconds: float | None = None,
    report_to_hub: bool = True,
    commit: bool = True,
    dry_run: bool = False,
    max_iterations: int | None = None,
) -> int:
    """Poll forever (or once). Returns process exit code."""
    if interval_seconds is None:
        try:
            interval_seconds = float(
                os.environ.get("AGENT_HARNESS_CLAIM_INTERVAL", "30")
            )
        except ValueError:
            interval_seconds = 30.0
    interval_seconds = max(1.0, interval_seconds)

    stop = threading.Event()

    def _handle_sig(*_args: Any) -> None:
        logger.info("shutdown signal received")
        stop.set()
        _cancel_active_run("signal")

    signal.signal(signal.SIGINT, _handle_sig)
    signal.signal(signal.SIGTERM, _handle_sig)

    client = ActivityCoreOpsClient()
    from rein_aharness.readiness import run_readiness_checks

    readiness = run_readiness_checks(client)
    if not readiness.ok:
        failed = ", ".join(
            check.name for check in readiness.checks if not check.ok
        )
        logger.error("claim-loop readiness failed: %s", failed)
        return 2
    logger.info(
        "claim-loop start worker_id=%s url=%s labels=%s interval=%ss once=%s",
        client.config.worker_id,
        client.config.base_url,
        client.config.claim_labels,
        interval_seconds,
        once,
    )

    iterations = 0
    exit_code = 0
    while not stop.is_set():
        iterations += 1
        t0 = time.monotonic()
        try:
            result = process_one(
                client,
                report_to_hub=report_to_hub,
                commit=commit,
                dry_run=dry_run,
            )
        except Exception as exc:  # noqa: BLE001
            logger.exception("process_one crashed: %s", exc)
            result = ProcessResult(claimed=False, reason=str(exc))
            exit_code = 1

        elapsed = time.monotonic() - t0
        if result.empty:
            logger.debug("queue empty (%.2fs)", elapsed)
        else:
            logger.info(
                "cycle claimed=%s run_id=%s ok=%s approach=%s state=%s reason=%s (%.2fs)",
                result.claimed,
                result.run_id,
                result.ok,
                result.approach,
                result.ops_state,
                result.reason,
                elapsed,
            )
            if result.claimed and result.ok is False:
                exit_code = 1

        if once:
            break
        if max_iterations is not None and iterations >= max_iterations:
            break
        # Empty queues and upstream errors both use the configured backoff.
        # The short pause is only for a cycle that actually claimed work.
        sleep_for = (
            interval_seconds
            if result.empty or result.retry_full_interval
            else min(2.0, interval_seconds)
        )
        stop.wait(sleep_for)

    logger.info("claim-loop stop iterations=%s exit=%s", iterations, exit_code)
    return 0 if once and exit_code == 0 else exit_code if once else 0