atlas-iac/services/hermes/scripts/execution_pool_maintenance.py

#!/usr/bin/env python3
"""Recoverable maintenance passes for the Hermes execution-pool coordinator.

Every pass here runs against a shared Kanban SQLite database that other writers
(the agent UI, the local CLI lane, crons) hold concurrently, so *any* board call
can fail transiently. The rules this module enforces:

* No pass ever leaves a claimed run without a durable row that some later pass
  will revisit. A Kanban write that fails leaves the row in its retryable state.
* No single poisoned row, task, or board aborts the work queued behind it.
* A durable row only reaches a terminal state on authoritative evidence about
  its Kanban run, because terminal rows are the only ones ever collected.
* A coordinator-side fault is never converted into a Kanban mutation. Only a
  genuine capability failure of the run itself parks that exact run.
"""

from __future__ import annotations

import sqlite3
import sys
from typing import Any

import cli_lane_dispatch
from cli_lane_config import canonical_run_id
from execution_pool_project import ProjectPolicyError, distributed_workspace_eligible
from execution_pool_protocol import ProtocolError


RECOVERABLE_BOARD_ERRORS = (OSError, sqlite3.Error)
DEFERRALS: dict[str, str] = {}
MAX_DEFERRALS = 64


def _task_value(task: Any, name: str, default: Any = None) -> Any:
    return getattr(task, name, default)


def _defer(context: str, error: BaseException) -> None:
    """Report one recoverable pool-side failure once per distinct detail.

    The maintenance loop re-runs every few seconds, so an unchanged failure is
    reported once rather than on every pass. The table is bounded because its
    keys include run identities.
    """
    detail = f"{type(error).__name__}: {error}"
    if DEFERRALS.get(context) == detail:
        return
    if len(DEFERRALS) >= MAX_DEFERRALS:
        DEFERRALS.clear()
    DEFERRALS[context] = detail
    print(f"pool work deferred for {context}: {detail}", file=sys.stderr, flush=True)


def _settled(context: str) -> None:
    DEFERRALS.pop(context, None)


def _context(record: dict[str, Any]) -> str:
    return f"{record['board']}/{record['task_id']}#{record['run_id']}"


def _block_exact(
    kanban_db: Any,
    board: str,
    task_id: str,
    run_id: int | None,
    reason: str,
    *,
    kind: str = "capability",
) -> None:
    """Park exactly one run, never letting the write itself escape."""
    try:
        with kanban_db.scoped_current_board(board):
            connection = kanban_db.connect(board=board)
            try:
                kanban_db.block_task(
                    connection, task_id, reason=reason, kind=kind,
                    expected_run_id=run_id,
                )
            finally:
                connection.close()
    except Exception as error:  # noqa: BLE001 - a failed park must not wedge the pass
        _defer(f"{board}/{task_id} block", error)


def recover_results(pool: Any) -> None:
    """Re-apply each durable result in isolation from its neighbours."""
    for record in pool.store.pending_results():
        try:
            pool.finalize(record)
            _settled(_context(record))
        except Exception as error:  # noqa: BLE001 - one bad result must not block the rest
            print(
                f"result recovery deferred: {type(error).__name__}: {error}",
                file=sys.stderr,
                flush=True,
            )


def _lease_failure_state(
    kanban_db: Any, connection: Any, binding: dict[str, Any], run_id: int
) -> str:
    """Classify one exhausted lease from authoritative Kanban evidence only."""
    task = kanban_db.get_task(connection, binding["task_id"])
    current = (
        canonical_run_id(_task_value(task, "current_run_id", None))
        if task is not None
        else None
    )
    if current != run_id:
        return "stale"
    if str(_task_value(task, "status", "") or "") != "running":
        return "finalized"
    if kanban_db.block_task(
        connection,
        binding["task_id"],
        reason=(
            "Distributed worker lease expired after "
            f"{binding['attempt']} fenced attempts"
        ),
        kind="transient",
        expected_run_id=run_id,
    ):
        return "finalized"
    # Kanban refused the exact-run park yet still reports the run as running, so
    # there is no terminal evidence. Keep the row for the next pass to retry.
    raise ProtocolError("Kanban refused the exact-run lease park")


def record_lease_failure(pool: Any, record: dict[str, Any]) -> None:
    """Record one exhausted lease against its exact Kanban run, idempotently.

    The row stays ``lease_failed`` unless this pass obtains authoritative
    evidence: the run moved on or vanished (``stale``), Kanban already parked it
    (``finalized``), or this call parked it (``finalized``). Anything else --
    including a refused write or an unreadable board -- leaves the row for the
    next pass, so no run is ever silently dropped.
    """
    from hermes_cli import kanban_db

    binding = pool._binding(record)
    run_id = canonical_run_id(binding["run_id"])
    if run_id is None:
        pool.store.finalize(binding, "stale")
        return
    with pool._kanban_lock, kanban_db.scoped_current_board(binding["board"]):
        connection = kanban_db.connect(board=binding["board"])
        try:
            state = _lease_failure_state(kanban_db, connection, binding, run_id)
        finally:
            connection.close()
    pool.store.finalize(binding, state)
    _settled(_context(record))


def recover_lease_failures(pool: Any) -> None:
    """Retry every released-but-unrecorded lease failure, each in isolation."""
    for record in pool.store.failed_leases():
        try:
            record_lease_failure(pool, record)
        except Exception as error:  # noqa: BLE001 - retried on the next pass
            _defer(f"{_context(record)} lease park", error)


def expire_leases(pool: Any) -> None:
    """Fence dead attempts, then drain every unconfirmed terminal record."""
    pool.store.expire_leases()
    recover_lease_failures(pool)


def _release_moved_run(pool: Any, kanban_db: Any, record: dict[str, Any]) -> None:
    """Release one durable assignment whose Kanban run is no longer ours."""
    with kanban_db.scoped_current_board(record["board"]):
        connection = kanban_db.connect(board=record["board"])
        try:
            task = kanban_db.get_task(connection, record["task_id"])
            current_run_id = (
                canonical_run_id(_task_value(task, "current_run_id", None))
                if task
                else None
            )
            current = str(current_run_id or "")
            status = str(_task_value(task, "status", "") or "") if task else ""
        finally:
            connection.close()
    if current != record["run_id"] or status != "running":
        pool.store.finalize(record, "stale")


def _adopt_task(
    pool: Any,
    kanban_db: Any,
    connection: Any,
    board: str,
    task: Any,
    ordinals: list[int],
    known: set[tuple[str, str, str]],
) -> None:
    """Adopt one orphaned pool-eligible run onto a free ordinal."""
    task_id = str(_task_value(task, "id", "") or "")
    raw_run_id = canonical_run_id(_task_value(task, "current_run_id", None))
    run_id = str(raw_run_id or "")
    assignee = str(_task_value(task, "assignee", "") or "")
    if (
        not task_id
        or not run_id
        or str(_task_value(task, "status", "")) != "running"
        or not assignee.startswith("cli-")
        or not distributed_workspace_eligible(task)
        or (board, task_id, run_id) in known
    ):
        return
    try:
        payload = pool.assignment_payload(kanban_db, connection, task, board)
    except RECOVERABLE_BOARD_ERRORS as error:
        _defer(f"{board}/{task_id} recovery payload", error)
        return
    except Exception as error:  # noqa: BLE001 - a real capability failure of this run
        _block_exact(
            kanban_db, board, task_id, raw_run_id,
            "Distributed assignment recovery failed: "
            f"{type(error).__name__}: {error}",
        )
        return
    binding = {
        "board": board, "task_id": task_id, "run_id": run_id,
        "worker_ordinal": ordinals[0], "attempt": 1,
    }
    try:
        pool.store.add(binding, payload)
    except ProtocolError as error:
        # A durable row already owns this identity or ordinal. Leave the ordinal
        # free and leave Kanban untouched; a later pass revisits it.
        _defer(f"{board}/{task_id} adoption", error)
        return
    known.add((board, task_id, run_id))
    ordinals.pop(0)


def _adopt_board(
    pool: Any,
    kanban_db: Any,
    board: str,
    ordinals: list[int],
    known: set[tuple[str, str, str]],
) -> None:
    """Adopt this board's orphaned runs without touching any other board."""
    with kanban_db.scoped_current_board(board):
        connection = kanban_db.connect(board=board)
        try:
            for task in kanban_db.list_tasks(connection):
                if not ordinals:
                    return
                _adopt_task(pool, kanban_db, connection, board, task, ordinals, known)
        finally:
            connection.close()


def reconcile(pool: Any) -> None:
    """Recover the narrow claim/assignment crash gaps without double execution."""
    from hermes_cli import kanban_db

    for record in pool.store.active_assignments():
        try:
            _release_moved_run(pool, kanban_db, record)
        except Exception as error:  # noqa: BLE001 - other assignments still reconcile
            _defer(f"{_context(record)} release", error)
    ordinals = pool.store.available_ordinals()
    if not ordinals:
        return
    known = pool.store.known_runs()
    try:
        boards = kanban_db.list_boards(include_archived=False)
    except RECOVERABLE_BOARD_ERRORS as error:
        _defer("board-registry", error)
        return
    _settled("board-registry")
    for raw_board in boards:
        board = cli_lane_dispatch._board_slug(raw_board)
        if not board or not ordinals:
            continue
        try:
            _adopt_board(pool, kanban_db, board, ordinals, known)
            _settled(f"{board} adoption")
        except Exception as error:  # noqa: BLE001 - later boards must still run
            _defer(f"{board} adoption", error)


def _materialize(
    pool: Any,
    kanban_db: Any,
    board: str,
    task_id: str,
    ordinal: int,
    known: set[tuple[str, str, str]],
) -> None:
    """Turn one freshly claimed run into exactly one durable assignment."""
    database_run_id: int | None = None
    try:
        with kanban_db.scoped_current_board(board):
            connection = kanban_db.connect(board=board)
            try:
                task = kanban_db.get_task(connection, task_id)
                if task is None:
                    return
                raw_run = canonical_run_id(
                    _task_value(task, "current_run_id", None)
                )
                if not distributed_workspace_eligible(task):
                    if raw_run is not None:
                        kanban_db.block_task(
                            connection, task_id,
                            reason=(
                                "Distributed claim fenced because an existing "
                                "workspace is owned by the local lane"
                            ),
                            kind="capability", expected_run_id=raw_run,
                        )
                    return
                if raw_run is None:
                    raise ProjectPolicyError("claimed task has no canonical run ID")
                database_run_id = raw_run
                payload = pool.assignment_payload(kanban_db, connection, task, board)
            finally:
                connection.close()
        run_id = str(database_run_id)
        if (board, task_id, run_id) in known:
            return
        pool.store.add(
            {
                "board": board, "task_id": task_id, "run_id": run_id,
                "worker_ordinal": ordinal, "attempt": 1,
            },
            payload,
        )
        known.add((board, task_id, run_id))
        _settled(f"{board}/{task_id} dispatch")
    except (ProtocolError, *RECOVERABLE_BOARD_ERRORS) as error:
        # Coordinator-side storage/identity conditions. Never convert these into
        # a Kanban capability block; reconcile() adopts the claim on a later pass.
        _defer(f"{board}/{task_id} dispatch", error)
    except Exception as error:  # noqa: BLE001 - a real capability failure of this run
        _block_exact(
            kanban_db, board, task_id, database_run_id,
            "Distributed assignment preparation failed: "
            f"{type(error).__name__}: {error}",
        )


def dispatch(pool: Any) -> None:
    """Claim up to the free ordinal count and materialize assignments."""
    from hermes_cli import kanban_db

    ordinals = pool.store.available_ordinals()
    if not ordinals:
        return
    try:
        # Only storage faults are absorbed here. An incompatible claim API must
        # still surface, because silently skipping the eligibility predicate is
        # how the pool would start claiming the local lane's owned workspaces.
        claimed = cli_lane_dispatch.claim_ready(
            set(), len(ordinals),
            lambda _board, task: distributed_workspace_eligible(task),
        )
    except RECOVERABLE_BOARD_ERRORS as error:
        _defer("claim-ready", error)
        return
    _settled("claim-ready")
    known = pool.store.known_runs()
    for ordinal, (board, task_id) in zip(ordinals, claimed, strict=False):
        _materialize(pool, kanban_db, board, task_id, ordinal, known)