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

    R-round4 P2: defensive normalization of input_payload — if the
    DB returned NULL (legacy/seeded data, schema bug, etc.) we
    substitute an empty dict so the agent_runner doesn't crash on
    ``dict(None)`` and infinite-loop pickup_count → STUCK.
    """
    if not isinstance(input_payload, dict):
        logger.warning(
            "attempt_id=%s received non-dict input_payload (%r); "
            "substituting empty dict so the agent_runner doesn't crash",
            attempt_id, type(input_payload).__name__,
        )
        input_payload = {}
    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()

    # Phase 1k++++ (run-fix): head_sha bookkeeping. tick.py reads
    # head_sha_before/after to compute head_sha_advanced; without
    # writing them, implementer 'resolved' outcomes would map to a
    # no-op event and workflows would stall after the agent ran.
    #
    # head_sha_before comes from the input_payload (what the
    # prefetch saw + the agent was told to start from).
    # head_sha_after comes from the agent's output_payload — the
    # last commit_sha if any commits were produced; falls back to
    # head_sha_before when the agent didn't commit (no advance).
    hs_before = input_payload.get("head_sha") if isinstance(input_payload, dict) else None
    hs_after = hs_before
    if isinstance(output_payload, dict):
        commits = output_payload.get("commit_shas") or []
        if isinstance(commits, list) and commits:
            last = commits[-1]
            if isinstance(last, str) and last:
                hs_after = last

    # 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),
        head_sha_before=hs_before,
        head_sha_after=hs_after,
    )
    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,
    head_sha_before: str | None = None,
    head_sha_after: str | None = None,
) -> 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).

    ``head_sha_before`` / ``head_sha_after`` (Phase 1k++++ run-fix):
    tick.py reads these to compute ``head_sha_advanced``, which the
    outcome mapper requires to distinguish ``implementer_pushed``
    (true push happened) from ``implementer_blocked`` (worker said
    resolved but git didn't move). Without writing them, the tick
    sees head_sha_advanced=False and the workflow doesn't progress.
    """
    now = datetime.now(timezone.utc)
    from .._json_safe import safe_json_dumps as _safe_json_dumps
    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, "
                "  head_sha_before = :head_sha_before, "
                "  head_sha_after = :head_sha_after, "
                "  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": (
                    # Restricted encoder (datetime / Decimal / UUID /
                    # Path / set) — anything else raises so worker
                    # output regressions surface loudly instead of
                    # silently stringifying to "<MyObj at 0x...>".
                    _safe_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,
                "head_sha_before": head_sha_before,
                "head_sha_after": head_sha_after,
                "attempt_id": attempt_id,
                "instance": instance_id,
            },
        )
        return result.rowcount > 0