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.approaches import (
    APPROACH_UNMATCHED,
    ApproachResult,
    execute_approach,
    select_approach,
)
from rein_aharness.execution_cancel import (
    ExecutionCancel,
    ExecutionCancelled,
    using_cancel,
)
from rein_aharness.glas_execution import (
    GLAS_APPROACH,
    GlasExecutionError,
    execute_profiled_run,
)
from rein_aharness.lease_monitor import LeaseLoss, LeaseMonitor
from rein_aharness.repository_transaction import (
    DirtyRepositoryError,
    GitRepositoryError,
    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")


@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 process_one(
    client: ActivityCoreOpsClient | None = None,
    *,
    report_to_hub: bool = True,
    commit: bool = True,
    dry_run: bool = False,
) -> ProcessResult:
    """Claim at most one ops_run, execute approach, complete or fail."""
    client = client or ActivityCoreOpsClient()
    cfg = client.config

    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,
        )

    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,
) -> ProcessResult:
    """Execute a profiled row without consulting or falling back to legacy routing."""
    cancel = ExecutionCancel()
    hb = _Heartbeat(client, run.id, client.config.lease_seconds, cancel=cancel)
    hb.start()
    execution_error: str | None = None
    execution_reason = ""
    execution_reopen = False
    gateway_result: dict[str, Any] | None = None
    tx_evidence: dict[str, Any] | None = None
    _set_active_run_cancel(cancel)
    try:
        try:
            target = resolve_ops_target(run, client.config)
            with using_cancel(cancel):
                with RepositoryTransaction(target, correlation_id=run.id) as tx:
                    gateway_result = execute_profiled_run(
                        run,
                        client.config,
                        report_to_hub=report_to_hub,
                        cancel=cancel,
                    )
                    tx_evidence = tx.evidence()
        except ExecutionCancelled as exc:
            execution_error = type(exc).__name__
            execution_reason = str(exc)
        except DirtyRepositoryError as exc:
            execution_error = type(exc).__name__
            execution_reason = f"refused: {exc}"
        except RepositoryBusyError as exc:
            execution_error = type(exc).__name__
            execution_reason = f"refused: {exc}"
            execution_reopen = True
        except (GitRepositoryError, RepositoryTransactionError, TaskSpecError) as exc:
            execution_error = type(exc).__name__
            execution_reason = f"refused: {exc}"
        except GlasExecutionError as exc:
            execution_error = type(exc).__name__
            execution_reason = str(exc)
    finally:
        hb.stop()
        _set_active_run_cancel(None)

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

    if execution_error is not None:
        reason = execution_reason or f"profiled execution failed ({execution_error})"
        try:
            out = client.fail(
                run.id,
                error=reason,
                reopen=execution_reopen,
                result={"ok": False, "approach": GLAS_APPROACH, "reason": reason},
            )
        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}",
            )
        return ProcessResult(
            claimed=True,
            run_id=run.id,
            approach=GLAS_APPROACH,
            ok=False,
            reason=reason,
            ops_state=out.state,
        )

    assert gateway_result is not None
    evidence = gateway_result["evidence"]
    ok = gateway_result["ok"]
    reason = str(evidence.get("error") or evidence.get("outcome") or "Glas execution failed")
    detail = {"execution_evidence": evidence}
    if tx_evidence is not None:
        detail["repository_transaction"] = tx_evidence
        gateway_result = {**gateway_result, "repository_transaction": tx_evidence}
    try:
        if ok:
            out = client.complete(run.id, result=gateway_result)
        else:
            out = client.fail(
                run.id,
                error=reason,
                reopen=False,
                result=gateway_result,
            )
    except OpsRunError as exc:
        return ProcessResult(
            claimed=True,
            run_id=run.id,
            approach=GLAS_APPROACH,
            ok=False,
            reason=f"close ops_run failed: {exc}; gateway_ok={ok} {reason}",
            detail=detail,
        )

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


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()
    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