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.glas_execution import (
    GLAS_APPROACH,
    GlasExecutionError,
    execute_profiled_run,
)
from rein_aharness.ops_run_client import (
    ActivityCoreOpsClient,
    OpsRun,
    OpsRunConfig,
    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))


class _Heartbeat:
    def __init__(
        self,
        client: ActivityCoreOpsClient,
        run_id: str,
        lease_seconds: int,
    ):
        self._client = client
        self._run_id = run_id
        self._lease = lease_seconds
        self._stop = threading.Event()
        self._thread: threading.Thread | None = None

    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:
                    logger.warning("heartbeat failed run_id=%s: %s", self._run_id, exc)

        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)


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

    hb = _Heartbeat(client, run.id, cfg.lease_seconds)
    hb.start()
    try:
        ar: ApproachResult = execute_approach(
            run,
            approach=approach,
            config=cfg,
            report_to_hub=report_to_hub,
            commit=commit,
        )
    finally:
        hb.stop()

    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."""
    hb = _Heartbeat(client, run.id, client.config.lease_seconds)
    hb.start()
    try:
        try:
            gateway_result = execute_profiled_run(
                run,
                client.config,
                report_to_hub=report_to_hub,
            )
        except GlasExecutionError as exc:
            reason = str(exc)
            try:
                out = client.fail(
                    run.id,
                    error=reason,
                    reopen=False,
                    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,
            )
    finally:
        hb.stop()

    evidence = gateway_result["evidence"]
    ok = gateway_result["ok"]
    reason = str(evidence.get("error") or evidence.get("outcome") or "Glas execution failed")
    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={"execution_evidence": evidence},
        )

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


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

    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