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

#!/usr/bin/env python3
"""Reopen only due, coordinator-owned SCM publication retries."""

from __future__ import annotations

import json
import time
from typing import Any

import supervisor_state
from publication_retry import PublicationRetryError
from publication_retry_fence import guarded_unblock, read as read_fence


INITIAL_BACKOFF_SECONDS = 300


def _value(task: Any, name: str) -> Any:
    return task.get(name) if isinstance(task, dict) else getattr(task, name, None)


def _due(board: str, task_id: str, now: int) -> tuple[dict[str, Any], str] | None:
    """Return a validated unissued receipt only after one durable backoff."""
    with supervisor_state._connect(board) as connection:
        row = connection.execute(
            "SELECT source_run_id,source_ordinal,issued_run_id,resolved_run_id,retry_after,last_reissued_run_id "
            "FROM publication_retries WHERE board=? AND child_task_id=?", (board, task_id)
        ).fetchone()
        if row is None or row[2] or row[3]:
            return None
        if int(row[4]) == 0:
            connection.execute(
                "UPDATE publication_retries SET retry_after=? WHERE board=? AND child_task_id=? AND retry_after=0 "
                "AND issued_run_id='' AND resolved_run_id=''", (now + INITIAL_BACKOFF_SECONDS, board, task_id)
            )
            return None
        if int(row[4]) > now:
            return None
        expected_run = str(row[5] or row[0])
    try:
        receipt = supervisor_state.publication_retry(board, task_id, "")
    except (OSError, ValueError, PublicationRetryError):
        return None
    if receipt is None or str(receipt.get("source", {}).get("run_id", "")) != str(row[0]):
        return None
    return receipt, expected_run


def _same_evidence(left: Any, right: Any, *, source_run: bool) -> bool:
    """Compare immutable receipt evidence; a corrected title may differ on reissue."""
    if not isinstance(left, dict) or not isinstance(right, dict):
        return False
    keys = ("source", "baseline_sha", "head", "body", "structured")
    if any(left.get(key) != right.get(key) for key in keys):
        return False
    return not source_run or (left.get("title") == right.get("title") and left.get("result_digest") == right.get("result_digest"))


def _pool_owned(pool: Any, board: str, task_id: str, run_id: str, receipt: dict[str, Any]) -> bool:
    """Require a terminal exact run and no competing live pool assignment."""
    if any(record.get("board") == board and record.get("task_id") == task_id for record in pool.store.active_assignments()):
        return False
    record = pool.store.terminal_record(board, task_id, run_id)
    if record is None or record.get("state") != "finalized" or int(record.get("worker_ordinal", -1)) != int(receipt["source"]["worker_ordinal"]):
        return False
    result, payload = record.get("result"), record.get("payload")
    structured = result.get("structured") if isinstance(result, dict) else None
    if (not isinstance(structured, dict) or structured.get("status") != "blocked"
            or result.get("capacity_failure") is not True or result.get("scm_submission") is not None):
        return False
    source_run = run_id == str(receipt["source"]["run_id"])
    evidence = result.get("scm_resume") if source_run and isinstance(result, dict) else payload.get("scm_resume") if isinstance(payload, dict) else None
    return _same_evidence(evidence, receipt, source_run=source_run)


def _reopen(kanban_db: Any, board: str, task_id: str, expected_run: str, receipt: dict[str, Any]) -> bool:
    """Unblock an unchanged transient task only when its private parent chain holds."""
    child = supervisor_state.get_child(board, task_id)
    if child is None:
        return False
    try:
        if supervisor_state.publication_retry(board, task_id, "") != receipt:
            return False
    except (OSError, ValueError, PublicationRetryError):
        return False
    fence = read_fence(board, task_id, expected_run)
    if fence is None:
        return False
    with kanban_db.scoped_current_board(board):
        connection = kanban_db.connect(board=board)
        try:
            task = kanban_db.get_task(connection, task_id)
            parents = kanban_db.parent_ids(connection, task_id)
            if (
                task is None or _value(task, "status") != "blocked" or _value(task, "block_kind") != "transient"
                or _value(task, "current_run_id") is not None
                or child["root_task_id"] not in {str(value) for value in parents}
            ):
                return False
            return guarded_unblock(
                kanban_db, connection, task_id, child["root_task_id"], expected_run, fence
            )
        finally:
            connection.close()


def reopen_due(pool: Any, kanban_db: Any, boards: list[str], *, now: int | None = None) -> int:
    """Schedule due receipts only on their free source ordinal, once per pass."""
    current = int(time.time()) if now is None else int(now)
    reopened = 0
    available = set(pool.store.available_ordinals())
    for board in boards:
        with kanban_db.scoped_current_board(board):
            connection = kanban_db.connect(board=board)
            try:
                tasks = list(kanban_db.list_tasks(connection))
            finally:
                connection.close()
        for task in tasks:
            task_id = str(_value(task, "id") or "")
            if (
                not task_id or _value(task, "status") != "blocked"
                or _value(task, "block_kind") != "transient" or _value(task, "current_run_id") is not None
            ):
                continue
            candidate = _due(board, task_id, current)
            if candidate is None:
                continue
            receipt, expected_run = candidate
            ordinal = int(receipt["source"]["worker_ordinal"])
            if ordinal not in available or not _pool_owned(pool, board, task_id, expected_run, receipt):
                continue
            if _reopen(kanban_db, board, task_id, expected_run, receipt):
                available.remove(ordinal)
                reopened += 1
    return reopened