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

"""Worker main loop: dequeue → submit to thread pool → repeat.

Plan v9: each worker controller process runs ``worker_main_loop``.
The loop polls the DB for pending attempts (capped at
``CONTROLLER_MAX_CONCURRENT_WORKER_THREADS_PER_MACHINE`` concurrent
attempt threads — one worker *process* with an N-thread pool, NOT N
processes) and submits each to a ThreadPoolExecutor that calls
``run_one_attempt``.

This module focuses on the loop + concurrency cap. The
``run_one_attempt`` runner (runner.py) does the actual work.
Tests can exercise the loop with a synthetic agent_runner and a
small SQLite DB.

Out of scope for this skeleton (Phase 1c follow-on):
- LISTEN/NOTIFY wake-up on attempt enqueue (Postgres only).
  Currently polls on a fixed interval.
- Workspace per-PR umbrella (deferred; agent_runner stand-in
  doesn't need it).
- Real OpenCode + MCP subprocess spawn (deferred; tests inject
  a FakeAgentRunner that returns canned JSON).
"""

from __future__ import annotations

import logging
import os
import threading
import time
from concurrent.futures import Future, ThreadPoolExecutor
from dataclasses import dataclass, field

from sqlalchemy.engine import Engine

from ..db.dequeue import dequeue_one
from ..db.session import session_scope
from .identity import build_instance_id
from .runner import AgentRunner, AttemptOutcome, run_one_attempt

logger = logging.getLogger(__name__)


@dataclass
class WorkerConfig:
    """Per-worker-process configuration."""

    roles: list[str] = field(
        default_factory=lambda: ["implementer", "reviewer", "estimator"]
    )
    max_concurrent: int = int(
        os.environ.get("CONTROLLER_MAX_CONCURRENT_WORKER_THREADS_PER_MACHINE", "1")
    )
    poll_interval_s: float = float(
        os.environ.get("CONTROLLER_WORKER_POLL_INTERVAL_S", "5")
    )
    heartbeat_interval_s: int = int(
        os.environ.get("CONTROLLER_HEARTBEAT_INTERVAL_S", "30")
    )


def worker_main_loop(
    engine: Engine,
    agent_runner: AgentRunner,
    *,
    config: WorkerConfig | None = None,
    stop_event: threading.Event | None = None,
    instance_id: str | None = None,
    on_outcome: callable | None = None,  # type: ignore[assignment]
) -> None:
    """Worker main loop.

    Args:
        engine: SQLAlchemy engine (caller manages lifecycle).
        agent_runner: Callable that runs one attempt — production wires
            this to OpenCode + MCP; tests inject a fake.
        config: Per-worker config (defaults read from env vars).
        stop_event: Event the operator (or test) sets to exit cleanly.
            None = run forever (until SIGTERM kills the process).
        instance_id: Override the worker's identity. Default builds a
            fresh one via ``build_instance_id``.
        on_outcome: Optional callback invoked with each ``AttemptOutcome``
            after the runner returns. Tests use this to observe what
            happened without re-querying the DB.

    The loop yields control to the stop_event every poll_interval.
    Outer signal handlers (in production) set the stop_event on
    SIGTERM; the loop drains in-flight workers then exits.
    """
    cfg = config or WorkerConfig()
    stop = stop_event or threading.Event()
    instance = instance_id or build_instance_id()
    pool = ThreadPoolExecutor(
        max_workers=cfg.max_concurrent,
        thread_name_prefix=f"worker-{instance[-8:]}",
    )
    in_flight: set[Future[AttemptOutcome]] = set()
    logger.info(
        "worker started: instance=%s roles=%s max_concurrent=%d",
        instance,
        cfg.roles,
        cfg.max_concurrent,
    )

    try:
        while not stop.is_set():
            # Reap completed futures (no-op if pool just returned).
            _reap_done(in_flight, on_outcome)

            free_slots = cfg.max_concurrent - len(in_flight)
            picked_any = False
            for _ in range(max(0, free_slots)):
                try:
                    attempt = _dequeue_with_session(engine, instance, cfg.roles)
                except Exception:
                    logger.exception("dequeue failed; sleeping before retry")
                    break
                if attempt is None:
                    break
                future = pool.submit(
                    run_one_attempt,
                    engine,
                    attempt_id=attempt["attempt_id"],
                    role=attempt["role"],
                    tier=attempt["tier"],
                    input_payload=attempt["input_payload"],
                    instance_id=instance,
                    agent_runner=agent_runner,
                    heartbeat_interval_s=cfg.heartbeat_interval_s,
                )
                in_flight.add(future)
                picked_any = True

            # If nothing pending, sleep before next poll.
            if not picked_any:
                stop.wait(cfg.poll_interval_s)
        logger.info(
            "worker stopping: stop_event set; draining in-flight=%d", len(in_flight)
        )
    finally:
        # Drain: wait for in-flight work to finish, then close pool.
        for future in list(in_flight):
            try:
                outcome = future.result(timeout=None)
                if on_outcome is not None:
                    on_outcome(outcome)
            except Exception:
                logger.exception("in-flight runner raised during drain")
        pool.shutdown(wait=True)


def _dequeue_with_session(
    engine: Engine, instance_id: str, roles: list[str]
) -> dict | None:
    """Acquire one pending attempt's metadata + input_payload.

    Returns a plain dict (not the ORM object) because the session
    closes after the dequeue; the runner re-opens its own session
    for writes.
    """
    with session_scope(engine) as session:
        from sqlalchemy import text  # local import to keep top imports tight

        result = dequeue_one(session, instance_id=instance_id, roles=roles)
        if not result.acquired:
            return None
        # Re-read the row inside the same session to fetch input_payload.
        row = session.execute(
            text(
                "SELECT attempt_id, role, tier, input_payload "
                "FROM workflow_attempts WHERE attempt_id = :id"
            ),
            {"id": result.attempt_id},
        ).first()
        if row is None:
            return None
        payload = row.input_payload
        # SQLite returns JSON as a string; Postgres returns dict.
        if isinstance(payload, str):
            import json

            payload = json.loads(payload)
        return {
            "attempt_id": row.attempt_id,
            "role": row.role,
            "tier": row.tier,
            "input_payload": payload,
        }


def _reap_done(in_flight: set[Future], on_outcome: callable | None) -> None:  # type: ignore[type-arg]
    done = {f for f in in_flight if f.done()}
    for f in done:
        in_flight.discard(f)
        try:
            outcome = f.result()
            if on_outcome is not None:
                on_outcome(outcome)
        except Exception:
            logger.exception("runner raised")