cleveragents-core/tools/controller/worker/runner.py

"""Per-attempt runner: dequeue → run → write output → release lock.

Plan v9 worker mainline:
1. The dequeue (in loop.py) gives us an attempt_id + role + input_payload
2. Start a heartbeat thread keeping the lock alive
3. Invoke the agent (Phase 1c-2: dispatch via OpenCode using the role's
   builder MCP; this skeleton emits a stand-in for the test harness)
4. On success: write status='complete' + output_payload; release lock
5. On lost lock (heartbeat detected reap): abort silently — reaper has
   already re-pended; another worker will retry
6. On worker-internal error: write status='failed' + outcome classification

This module's ``run_one_attempt`` is the pure function the test
harness exercises. The loop module wraps it with the dequeue +
ThreadPoolExecutor + retry semantics.

For Phase 1c skeleton: the actual OpenCode invocation is parameterised
as ``agent_runner``. Tests inject a FakeAgentRunner; the production
adapter (Phase 1c-2) will spawn the MCP subprocess + the OpenCode
session.
"""
from __future__ import annotations

import logging
from collections.abc import Callable
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Any

from sqlalchemy import text
from sqlalchemy.engine import Engine

from ..db.session import session_scope
from .heartbeat import Heartbeat, HeartbeatStopped

logger = logging.getLogger(__name__)


# ─── result types ─────────────────────────────────────────────────────


@dataclass
class AttemptOutcome:
    """What ``run_one_attempt`` returns.

    Mirrors the status='complete' vs 'failed' write the runner
    performs. Tests can introspect this without re-querying the DB.
    """

    attempt_id: int
    status: str             # 'complete' | 'failed' | 'aborted'
    outcome: str | None     # role-output outcome OR worker-side classification
    output_payload: dict[str, Any] | None
    wallclock_seconds: float


class WorkerError(Exception):
    """Internal worker-side error (transport, MCP crash, etc.).

    The runner catches this and writes ``status='failed'`` with
    ``outcome='worker-internal-error'``; the master re-enqueues
    with pickup_count++.
    """

    def __init__(self, message: str, outcome: str = "worker-internal-error"):
        self.outcome = outcome
        super().__init__(message)


class WorkerLostLock(Exception):
    """Raised when the heartbeat detected the lock was reaped.

    The runner catches this and silently aborts — does NOT write
    output. The reaper has already re-pended the attempt; another
    worker will pick it up.
    """


# ─── agent_runner protocol ────────────────────────────────────────────


# Type alias for the callable that drives one attempt's role-MCP +
# OpenCode session. The production adapter wraps OpenCode + spawns
# the matching builder MCP subprocess. Tests inject a fake.
#
# Inputs:
#   attempt_id, role, tier, input_payload, instance_id, lost_lock_event_check
#
# The agent_runner reads ``input_payload`` (already validated against
# {Role}InputV1 contract by the master), invokes the role-builder MCP,
# and returns the canonical output dict (already round-tripped through
# Pydantic by the controller's strict-parse layer).
#
# ``lost_lock_event_check()`` returns True if the heartbeat thread has
# signaled lost lock; the agent_runner SHOULD check this periodically
# and raise ``WorkerLostLock`` to bail out early. The runner also
# catches the raised event for defense-in-depth.
AgentRunner = Callable[..., dict[str, Any]]


# ─── runner ────────────────────────────────────────────────────────────


def run_one_attempt(
    engine: Engine,
    *,
    attempt_id: int,
    role: str,
    tier: int | None,
    input_payload: dict[str, Any],
    instance_id: str,
    agent_runner: AgentRunner,
    heartbeat_interval_s: int = 30,
) -> AttemptOutcome:
    """Run one attempt end-to-end.

    Caller has already dequeued the attempt (status='in_progress',
    lock columns set). This function:
    - Starts the heartbeat
    - Invokes ``agent_runner`` to do the actual work
    - On success: writes status='complete' + output payload
    - On lost lock: aborts silently (no DB write — reaper handled it)
    - On WorkerError: writes status='failed' with the outcome label
    - On other Exception: writes status='failed' outcome='worker-internal-error'
    - Stops the heartbeat in finally
    """
    started_at = datetime.now(timezone.utc)
    heartbeat = Heartbeat(
        engine, attempt_id, instance_id, interval_s=heartbeat_interval_s
    )
    heartbeat.start()
    output_payload: dict[str, Any] | None = None
    outcome: str | None = None
    final_status = "failed"
    try:
        # Wrap the agent_runner so it can check lost_lock without
        # depending on the heartbeat module directly.
        def lost_lock_check() -> bool:
            return heartbeat.lost_lock_event.is_set()

        try:
            output_payload = agent_runner(
                attempt_id=attempt_id,
                role=role,
                tier=tier,
                input_payload=input_payload,
                instance_id=instance_id,
                lost_lock_check=lost_lock_check,
            )
        except WorkerLostLock:
            # Agent observed lock loss + bailed. Honor the contract:
            # silent abort.
            logger.info(
                "attempt_id=%s aborted: lost lock during agent run",
                attempt_id,
            )
            return AttemptOutcome(
                attempt_id=attempt_id,
                status="aborted",
                outcome="lost-lock",
                output_payload=None,
                wallclock_seconds=_elapsed(started_at),
            )

        # Defense in depth: even if the agent didn't raise, check the
        # event before writing.
        if heartbeat.lost_lock_event.is_set():
            logger.info(
                "attempt_id=%s lock lost during agent run; not writing output",
                attempt_id,
            )
            return AttemptOutcome(
                attempt_id=attempt_id,
                status="aborted",
                outcome="lost-lock",
                output_payload=None,
                wallclock_seconds=_elapsed(started_at),
            )

        # Success path: write status='complete'.
        outcome = (output_payload or {}).get("outcome") if isinstance(output_payload, dict) else None
        final_status = "complete"
    except WorkerError as exc:
        logger.warning(
            "attempt_id=%s worker error: %s (outcome=%s)",
            attempt_id, exc, exc.outcome,
        )
        outcome = exc.outcome
        output_payload = {"error": str(exc), "worker_outcome": exc.outcome}
        final_status = "failed"
    except Exception as exc:  # noqa: BLE001 — last-resort catch
        logger.exception("attempt_id=%s unexpected exception", attempt_id)
        outcome = "worker-internal-error"
        output_payload = {"error": str(exc), "exception_type": type(exc).__name__}
        final_status = "failed"
    finally:
        heartbeat.stop()

    # Write the result. If the lock was lost between agent return and
    # this write, the UPDATE WHERE locked_by_instance=us returns
    # rowcount=0 and we treat as aborted.
    wrote = _write_outcome(
        engine,
        attempt_id=attempt_id,
        instance_id=instance_id,
        status=final_status,
        outcome=outcome,
        output_payload=output_payload,
        wallclock_seconds=_elapsed(started_at),
    )
    if not wrote:
        logger.info(
            "attempt_id=%s output write found no matching locked row "
            "(lost lock); treating as aborted",
            attempt_id,
        )
        return AttemptOutcome(
            attempt_id=attempt_id,
            status="aborted",
            outcome="lost-lock-at-write",
            output_payload=None,
            wallclock_seconds=_elapsed(started_at),
        )
    return AttemptOutcome(
        attempt_id=attempt_id,
        status=final_status,
        outcome=outcome,
        output_payload=output_payload,
        wallclock_seconds=_elapsed(started_at),
    )


def _elapsed(started_at: datetime) -> float:
    return (datetime.now(timezone.utc) - started_at).total_seconds()


def _write_outcome(
    engine: Engine,
    *,
    attempt_id: int,
    instance_id: str,
    status: str,
    outcome: str | None,
    output_payload: dict[str, Any] | None,
    wallclock_seconds: float,
) -> bool:
    """UPDATE the attempt row with terminal state + release lock.

    Returns True if the row was updated (we still held the lock).
    Returns False if rowcount=0 (lost lock; another worker / reaper
    took over).
    """
    now = datetime.now(timezone.utc)
    import json
    with session_scope(engine) as session:
        result = session.execute(
            text(
                "UPDATE workflow_attempts SET "
                "  status = :status, "
                "  outcome = :outcome, "
                "  output_payload = :output_payload, "
                "  output_version = :output_version, "
                "  finished_at = :now, "
                "  wallclock_seconds = :wallclock, "
                "  locked_by_instance = NULL, "
                "  locked_at = NULL, "
                "  lock_heartbeat_at = NULL "
                "WHERE attempt_id = :attempt_id "
                "  AND locked_by_instance = :instance"
            ),
            {
                "status": status,
                "outcome": outcome,
                "output_payload": (
                    json.dumps(output_payload) if output_payload is not None else None
                ),
                "output_version": (
                    output_payload.get("output_version")
                    if isinstance(output_payload, dict) else None
                ),
                "now": now,
                "wallclock": wallclock_seconds,
                "attempt_id": attempt_id,
                "instance": instance_id,
            },
        )
        return result.rowcount > 0