cleveragents-core/tools/controller/master/scheduler.py

"""Per-workflow scheduling — enqueue the next attempt after a transition.

When the state machine transitions a workflow to a state that needs
a worker (ANALYZING / IMPLEMENTING / REVIEWING / CONFLICT_RESOLVING),
the master inserts a fresh ``workflow_attempts`` row with
``status='pending'`` so a worker can dequeue it.

Per plan v9: the master pre-fetches the full input_payload at enqueue
time. This module ships the SCHEDULER (what role + tier + attempt
number, and how to escalate), with the prefetch as a pluggable callback
the tests inject (production wires it to the Forgejo prefetch helpers
in Phase 1d-3c).

States and what they enqueue:
- ANALYZING → role='estimator', tier=None
- IMPLEMENTING(tier) → role='implementer', tier=current_tier
- REVIEWING → role='reviewer', tier=None (reviewer tier is fixed)
- CONFLICT_RESOLVING → role='conflict_resolver', tier=current_tier
- ESCALATING → master transitions to IMPLEMENTING(tier+1) OR ABANDONED
  per the static escalation policy; if IMPLEMENTING, enqueue same as
  IMPLEMENTING above; if ABANDONED, no enqueue
- AWAITING_CI → no worker enqueue (master CI-poll thread handles it
  in Phase 1d-3c; this module skips it)
- MERGING → no worker enqueue (master-direct Forgejo merge call
  in Phase 1d-3c)
- Terminal states (MERGED / ABANDONED / STUCK / CREATED_PR) → no enqueue
"""

from __future__ import annotations

import json
import logging

from .._json_safe import safe_json_dumps as _safe_json_dumps
from collections.abc import Callable
from dataclasses import dataclass, field
from datetime import datetime, timezone

from sqlalchemy import text
from sqlalchemy.engine import Engine

from ..db.session import session_scope
from ..state_machine import KNOWN_STATES

logger = logging.getLogger(__name__)


MAX_TIER = 2
"""Plan v9: tiers go 0/1/2. Beyond max_tier escalation transitions to
ABANDONED via the static escalation policy."""


_PREFETCH_FAILURE_LIMIT = 10
"""Per-state cap on prefetch failures before a workflow is routed to
STUCK.

A persistently failing ``prefetch()`` — a deleted PR, a contract bug,
a sustained Forgejo outage — used to be retried every master tick
FOREVER: the only prefetch-loop guard (T4-4) covered nothing but the
issue-kind sub-case. This is the same unbounded-retry class as the
run-2 174x estimator loop. At the 30s default tick cadence, 10
failures ≈ 5 minutes — long enough to ride out a transient Forgejo
blip, bounded enough that a genuinely broken workflow surfaces to an
operator instead of spinning silently. A workflow STUCK by a Forgejo
outage that later recovers is recovered with operator_unstick."""


# Prefetch callback signature. Tests inject a fake. Production wires
# this to the Forgejo prefetch path that builds the V1 ContractInput
# matching the role.
#
# Inputs: workflow_id, role, tier
# Returns: (input_payload dict, input_version string)
PrefetchCallback = Callable[[int, str, int | None], tuple[dict, str]]


@dataclass
class ScheduledAttempt:
    """One attempt the scheduler enqueued."""

    workflow_id: int
    attempt_id: int
    role: str
    tier: int | None


@dataclass
class SchedulerReport:
    """Per-tick scheduler summary."""

    workflows_scheduled: int = 0
    scheduled: list[ScheduledAttempt] = field(default_factory=list)
    workflows_skipped: int = 0
    skip_reasons: list[tuple[int, str]] = field(default_factory=list)


def schedule_next_attempts(
    engine: Engine,
    *,
    prefetch: PrefetchCallback,
    max_tier: int = MAX_TIER,
) -> SchedulerReport:
    """Find workflows that NEED a worker attempt but don't have one
    pending/in-progress; enqueue exactly one fresh attempt for each.

    "Needs an attempt" means:
    - current_state ∈ {ANALYZING, IMPLEMENTING, REVIEWING, CONFLICT_RESOLVING}
    - No row in workflow_attempts for this workflow with status IN
      ('pending', 'in_progress') AND matching the role this state
      needs.

    The ESCALATING state is handled inline: master transitions to
    IMPLEMENTING(tier+1) or ABANDONED + schedules accordingly.

    PAUSED workflows are deliberately EXCLUDED so the scheduler doesn't
    enqueue work for a paused workflow between two reconciliation ticks
    (the operator removed the label; reconciliation will catch up).
    """
    report = SchedulerReport()
    now = datetime.now(timezone.utc)
    with session_scope(engine) as session:
        # Workflows currently in a state that needs a worker.
        # PAUSED intentionally absent — reconciliation manages it.
        rows = session.execute(
            text(
                "SELECT workflow_id, current_state, current_tier, "
                "       last_transition_at "
                "FROM workflows "
                "WHERE current_state IN "
                "  ('GROOMING', 'ANALYZING', 'IMPLEMENTING', 'REVIEWING', "
                "   'CONFLICT_RESOLVING', 'ESCALATING') "
                # Phase 1 grooming-plan defer block (kept post-revert
                # 2026-05-25): skip workflows with deferred_reason set.
                # Harmless when no grooming code sets this column
                # (deferred_reason stays NULL for all workflows, and
                # `IS NULL` is true). The column is Phase 0 schema and
                # the defer mechanism (close_issue/defer_issue) is
                # Phase 0 code; both will still be the eventual
                # consumers when the worker-shape grooming dispatch is
                # designed.
                "  AND deferred_reason IS NULL"
            )
        ).all()

        for r in rows:
            # ESCALATING is special: resolve it to IMPLEMENTING(tier+1)
            # or ABANDONED first, then schedule.
            current_state = r.current_state
            current_tier = r.current_tier
            # Window for the prefetch-failure cap: failures are counted
            # since the workflow last changed state. After an escalation
            # below it is bumped to ``now`` so the new state-occupancy
            # starts with a clean count.
            last_transition_at = r.last_transition_at
            if current_state == "ESCALATING":
                next_state, next_tier = _resolve_escalation(
                    current_tier,
                    max_tier=max_tier,
                )
                _commit_escalation(
                    session, r.workflow_id, current_state, next_state, next_tier, now
                )
                current_state = next_state
                current_tier = next_tier
                last_transition_at = now
                if current_state == "ABANDONED":
                    report.workflows_skipped += 1
                    report.skip_reasons.append(
                        (r.workflow_id, "ABANDONED via escalation")
                    )
                    continue

            # State-to-role mapping.
            role = _role_for_state(current_state)
            if role is None:
                report.workflows_skipped += 1
                report.skip_reasons.append(
                    (r.workflow_id, f"no role for state {current_state!r}")
                )
                continue

            # Already-pending check: skip if a pending/in-progress
            # attempt for this role exists.
            #
            # T4-1 fix (2026-05-19): ALSO skip if a complete attempt
            # for this role has finished but tick hasn't processed it
            # yet (finished_at > last_transition_at). Without this
            # check the scheduler enqueues a SECOND estimator/etc.
            # when the first completes between two scheduler runs.
            # That second attempt's outcome then fires from a state
            # the workflow has already moved past (e.g.,
            # estimator_done from IMPLEMENTING) → IllegalTransition
            # → STUCK. Observed live in trial-4 run 2026-05-19.
            existing = session.execute(
                text(
                    "SELECT a.attempt_id FROM workflow_attempts a "
                    " JOIN workflows w ON w.workflow_id = a.workflow_id "
                    "WHERE a.workflow_id = :wf_id AND a.role = :role "
                    "  AND ("
                    "    a.status IN ('pending', 'in_progress') "
                    "    OR (a.status = 'complete' "
                    "        AND a.finished_at IS NOT NULL "
                    "        AND a.finished_at > w.last_transition_at)"
                    "  )"
                ),
                {"wf_id": r.workflow_id, "role": role},
            ).first()
            if existing is not None:
                report.workflows_skipped += 1
                report.skip_reasons.append(
                    (
                        r.workflow_id,
                        f"already has pending/in-progress/unprocessed-complete {role} attempt",
                    )
                )
                continue

            # T4-4 fix (2026-05-19): issue workflows currently lack
            # an implementer/reviewer/conflict_resolver prefetch path
            # (those builders demand kind='pr'). State machine treats
            # issues identically to PRs at the ANALYZING→IMPLEMENTING
            # transition, but the per-role prefetchers reject. Without
            # this guard the scheduler retries every ~5s forever
            # (observed in trial-4 wf=6 #34 — log spam, never reaches
            # terminal state). Route the workflow to STUCK with a clear
            # reason so an operator can decide.
            wf_kind = _read_workflow_kind(session, r.workflow_id)
            if wf_kind == "issue" and role in (
                "implementer",
                "reviewer",
                "conflict_resolver",
            ):
                _transition_issue_to_stuck(
                    session,
                    r.workflow_id,
                    current_state,
                    now,
                    role=role,
                )
                report.workflows_skipped += 1
                report.skip_reasons.append(
                    (
                        r.workflow_id,
                        f"issue workflows have no {role!r} path; routed to STUCK",
                    )
                )
                continue
            # Prefetch + enqueue.
            try:
                input_payload, input_version = prefetch(
                    r.workflow_id,
                    role,
                    current_tier,
                )
            except Exception as exc:  # noqa: BLE001 — prefetch failures are common
                # A persistently failing prefetch used to retry every
                # master tick forever (T4-4 guarded only the issue-kind
                # sub-case). Journal each failure and, once a workflow
                # has burned _PREFETCH_FAILURE_LIMIT of them in its
                # current state, route it to STUCK so an operator sees
                # it instead of perpetual silent log spam.
                _record_prefetch_failure(
                    session, r.workflow_id, current_state, role, exc, now
                )
                failure_count = _count_prefetch_failures(
                    session, r.workflow_id, since=last_transition_at
                )
                logger.warning(
                    "prefetch failed for workflow %s role %s (%d/%d): %s",
                    r.workflow_id,
                    role,
                    failure_count,
                    _PREFETCH_FAILURE_LIMIT,
                    exc,
                )
                report.workflows_skipped += 1
                if failure_count >= _PREFETCH_FAILURE_LIMIT:
                    _transition_prefetch_failed_to_stuck(
                        session,
                        r.workflow_id,
                        current_state,
                        now,
                        role=role,
                        failure_count=failure_count,
                        last_error=str(exc),
                    )
                    report.skip_reasons.append(
                        (
                            r.workflow_id,
                            f"prefetch failed {failure_count}x for {role!r}; "
                            f"routed to STUCK",
                        )
                    )
                else:
                    report.skip_reasons.append(
                        (r.workflow_id, f"prefetch raised: {exc}")
                    )
                continue

            attempt_id = _insert_pending_attempt(
                session,
                r.workflow_id,
                role,
                current_tier,
                input_payload,
                input_version,
                now,
            )
            report.workflows_scheduled += 1
            report.scheduled.append(
                ScheduledAttempt(
                    workflow_id=r.workflow_id,
                    attempt_id=attempt_id,
                    role=role,
                    tier=current_tier,
                )
            )

    return report


# ─── helpers ──────────────────────────────────────────────────────────


def _role_for_state(state: str) -> str | None:
    """Map a current_state to the worker role that drives it."""
    return {
        "GROOMING": "grooming_stage_b",
        "ANALYZING": "estimator",
        "IMPLEMENTING": "implementer",
        "REVIEWING": "reviewer",
        "CONFLICT_RESOLVING": "conflict_resolver",
    }.get(state)


def _read_workflow_kind(session, workflow_id: int) -> str:
    """Return the workflow's ``kind`` column ('pr' or 'issue')."""
    row = session.execute(
        text("SELECT kind FROM workflows WHERE workflow_id = :wf_id"),
        {"wf_id": workflow_id},
    ).first()
    return row.kind if row else ""


def _transition_issue_to_stuck(
    session,
    workflow_id: int,
    from_state: str,
    now: datetime,
    *,
    role: str,
) -> None:
    """T4-4: issue workflows have no implementer/reviewer/conflict_resolver
    prefetch path today. Route them to STUCK with a clear reason so
    they don't loop forever in the scheduler's prefetch-fail spiral.

    Future: when issue-implementation support lands (via
    IssueImplementerOutputV1 + a separate prefetch builder), remove
    this short-circuit and route issues through the proper path.
    """
    reason = f"issue-workflow has no {role!r} path (T4-4)"
    session.execute(
        text(
            "UPDATE workflows SET "
            "  current_state = 'STUCK', "
            "  last_transition_at = :now, "
            "  entered_state_at = :now "
            "WHERE workflow_id = :wf_id "
            "  AND current_state = :from_state"
        ),
        {"now": now, "wf_id": workflow_id, "from_state": from_state},
    )
    session.execute(
        text(
            "INSERT INTO controller_events "
            "(workflow_id, ts, event_type, from_state, to_state, "
            " payload, forgejo_write_pending, replay_attempts) "
            "VALUES (:wf_id, :ts, 'issue-not-supported', :from_state, "
            "        'STUCK', :payload, 0, 0)"
        ),
        {
            "wf_id": workflow_id,
            "ts": now,
            "from_state": from_state,
            "payload": json.dumps({"reason": reason, "source": "scheduler"}),
        },
    )
    logger.info(
        "scheduler: workflow %s (issue) routed to STUCK — no %s path; "
        "see PENDING_FIXES T4-4",
        workflow_id,
        role,
    )


def _record_prefetch_failure(
    session,
    workflow_id: int,
    current_state: str,
    role: str,
    exc: BaseException,
    now: datetime,
) -> None:
    """Journal one prefetch failure as a ``prefetch-failed``
    controller_events row.

    The row is both the audit trail of a Forgejo/contract outage and
    the counter that ``_count_prefetch_failures`` reads to enforce
    ``_PREFETCH_FAILURE_LIMIT``. ``from_state``/``to_state`` are both
    the current state — a prefetch failure is not a transition.
    """
    session.execute(
        text(
            "INSERT INTO controller_events "
            "(workflow_id, ts, event_type, from_state, to_state, "
            " payload, forgejo_write_pending, replay_attempts) "
            "VALUES (:wf_id, :ts, 'prefetch-failed', :state, :state, "
            "        :payload, 0, 0)"
        ),
        {
            "wf_id": workflow_id,
            "ts": now,
            "state": current_state,
            # _safe_json_dumps (not raw json.dumps) for consistency with
            # _commit_escalation — preempts a future contributor adding
            # a datetime/Decimal field here and tripping a TypeError.
            "payload": _safe_json_dumps(
                {
                    "role": role,
                    "error": str(exc)[:500],
                    "source": "scheduler",
                }
            ),
        },
    )


def _count_prefetch_failures(session, workflow_id: int, *, since) -> int:
    """Count ``prefetch-failed`` events for this workflow since it last
    changed state.

    The window is ``ts >= since`` (``since`` = the workflow's
    ``last_transition_at``), so a workflow that re-enters a
    prefetch-needing state after a successful transition starts with a
    clean count. ``since`` is None only for a workflow with no recorded
    transition yet — count every failure in that case rather than zero,
    so the cap can never be silently disabled.
    """
    if since is None:
        row = session.execute(
            text(
                "SELECT COUNT(*) AS n FROM controller_events "
                "WHERE workflow_id = :wf_id "
                "  AND event_type = 'prefetch-failed'"
            ),
            {"wf_id": workflow_id},
        ).first()
    else:
        row = session.execute(
            text(
                "SELECT COUNT(*) AS n FROM controller_events "
                "WHERE workflow_id = :wf_id "
                "  AND event_type = 'prefetch-failed' "
                "  AND ts >= :since"
            ),
            {"wf_id": workflow_id, "since": since},
        ).first()
    return int(row.n) if row else 0


def _transition_prefetch_failed_to_stuck(
    session,
    workflow_id: int,
    from_state: str,
    now: datetime,
    *,
    role: str,
    failure_count: int,
    last_error: str,
) -> None:
    """Route a workflow whose prefetch has failed too many times to
    STUCK with an operator-readable reason.

    The UPDATE is guarded on ``current_state = from_state`` so a
    concurrent transition between this tick's SELECT and now is not
    clobbered; the STUCK event row is journaled only when the UPDATE
    actually moved the workflow (rowcount > 0), so the journal never
    records a transition the guard rejected (R-2 / R-8 / AUDIT-9).
    """
    reason = (
        f"prefetch failed {failure_count}x for role {role!r} in "
        f"{from_state} (limit {_PREFETCH_FAILURE_LIMIT}); "
        f"last error: {last_error[:300]}"
    )
    result = session.execute(
        text(
            "UPDATE workflows SET "
            "  current_state = 'STUCK', "
            "  last_transition_at = :now, "
            "  entered_state_at = :now "
            "WHERE workflow_id = :wf_id "
            "  AND current_state = :from_state"
        ),
        {"now": now, "wf_id": workflow_id, "from_state": from_state},
    )
    if not result.rowcount:
        # A concurrent transition moved the workflow between this tick's
        # SELECT and now — leave it to whoever owns it now; do NOT
        # journal a STUCK that did not happen.
        logger.info(
            "scheduler: workflow %s no longer in %s; skipping "
            "prefetch-failed STUCK transition",
            workflow_id,
            from_state,
        )
        return
    session.execute(
        text(
            "INSERT INTO controller_events "
            "(workflow_id, ts, event_type, from_state, to_state, "
            " payload, forgejo_write_pending, replay_attempts) "
            "VALUES (:wf_id, :ts, 'prefetch-failed-stuck', :from_state, "
            "        'STUCK', :payload, 0, 0)"
        ),
        {
            "wf_id": workflow_id,
            "ts": now,
            "from_state": from_state,
            "payload": _safe_json_dumps(
                {
                    "reason": reason,
                    "role": role,
                    "failure_count": failure_count,
                    "source": "scheduler",
                }
            ),
        },
    )
    logger.warning(
        "scheduler: workflow %s routed to STUCK — prefetch failed %dx "
        "for role %s (limit %d)",
        workflow_id,
        failure_count,
        role,
        _PREFETCH_FAILURE_LIMIT,
    )


def _resolve_escalation(
    current_tier: int | None,
    *,
    max_tier: int,
) -> tuple[str, int | None]:
    """Static escalation policy: min(current_tier+1, MAX_TIER).
    Beyond MAX_TIER → ABANDONED. v9 deterministic.
    """
    if current_tier is None:
        # ESCALATING from a no-tier state (shouldn't happen normally).
        # Treat as escalate from tier 0.
        current_tier = 0
    next_tier = current_tier + 1
    if next_tier > max_tier:
        return "ABANDONED", current_tier
    return "IMPLEMENTING", next_tier


def _commit_escalation(
    session,
    workflow_id: int,
    from_state: str,
    to_state: str,
    new_tier: int | None,
    now: datetime,
) -> None:
    session.execute(
        text(
            "UPDATE workflows SET "
            "  current_state = :to_state, "
            "  current_tier = :tier, "
            "  last_transition_at = :now, "
            "  entered_state_at = :now "
            "WHERE workflow_id = :wf_id"
        ),
        {
            "to_state": to_state,
            "tier": new_tier,
            "now": now,
            "wf_id": workflow_id,
        },
    )
    session.execute(
        text(
            "INSERT INTO controller_events "
            "(workflow_id, ts, event_type, from_state, to_state, "
            " payload, forgejo_write_pending, replay_attempts) "
            "VALUES (:wf_id, :ts, 'transition', :from_state, :to_state, "
            "        :payload, 0, 0)"
        ),
        {
            "wf_id": workflow_id,
            "ts": now,
            "from_state": from_state,
            "to_state": to_state,
            # Use safe_json_dumps consistently with the rest of the
            # scheduler — preempts a future contributor adding a
            # datetime/Decimal field here and tripping the raw
            # json.dumps with a TypeError.
            "payload": _safe_json_dumps(
                {
                    "event": (
                        "escalate_next_tier_available"
                        if to_state == "IMPLEMENTING"
                        else "escalate_max_tier_exhausted"
                    ),
                    "new_tier": new_tier,
                    "reason": "scheduler escalation",
                }
            ),
        },
    )


def _insert_pending_attempt(
    session,
    workflow_id: int,
    role: str,
    tier: int | None,
    input_payload: dict,
    input_version: str,
    now: datetime,
) -> int:
    """Insert a pending attempt; return its id.

    Phase 1k+ refinement: ``input_payload`` is patched in-place AFTER
    the INSERT with the real ``attempt_id`` (autoincrement PK) and
    ``attempt_number`` so the DB never holds the placeholder
    ``attempt_id=0`` / ``attempt_number=1`` values that
    ``prefetch.py`` initially writes. Post-mortem debugging then sees
    the real values instead of chasing ghost zeros.
    """
    # Compute next attempt_number = max+1 (or 1 if first).
    row = session.execute(
        text(
            "SELECT COALESCE(MAX(attempt_number), 0) + 1 AS next "
            "FROM workflow_attempts WHERE workflow_id = :wf_id"
        ),
        {"wf_id": workflow_id},
    ).first()
    next_n = row.next

    # Patch placeholders BEFORE the INSERT so the stored row has the
    # right attempt_number from the start. attempt_id is patched after
    # INSERT via UPDATE — we don't know the autoincrement until then.
    patched_payload = dict(input_payload)
    if "attempt_number" in patched_payload:
        patched_payload["attempt_number"] = next_n

    result = session.execute(
        text(
            "INSERT INTO workflow_attempts "
            "(workflow_id, attempt_number, role, tier, status, "
            " input_payload, input_version, created_at, pickup_count, "
            " lock_ttl_seconds, input_payload_truncated, strict_parse_retries) "
            "VALUES (:wf_id, :n, :role, :tier, 'pending', "
            "        :payload, :version, :now, 0, 600, 0, 0)"
        ),
        {
            "wf_id": workflow_id,
            "n": next_n,
            "role": role,
            "tier": tier,
            # Restricted encoder: datetime / Decimal / UUID / Path / set
            # only — anything else raises so prefetch output regressions
            # surface loudly. Allowlist enough that nested CI summary
            # fields (CISummary.observed_at is a datetime) still
            # serialize.
            "payload": _safe_json_dumps(patched_payload),
            "version": input_version,
            "now": now,
        },
    )
    # SQLAlchemy 2.0 + SQLite: ``lastrowid`` is the way to retrieve
    # the auto-incremented PK from a raw text() insert.
    attempt_id = result.lastrowid

    # Patch in the real attempt_id and re-serialize. One extra UPDATE
    # per attempt; cheap vs. the alternative of an audit-trail lie.
    if "attempt_id" in patched_payload:
        patched_payload["attempt_id"] = attempt_id
        session.execute(
            text(
                "UPDATE workflow_attempts SET input_payload = :payload "
                "WHERE attempt_id = :aid"
            ),
            {
                "payload": _safe_json_dumps(patched_payload),
                "aid": attempt_id,
            },
        )

    return attempt_id


__all__ = [
    "MAX_TIER",
    "PrefetchCallback",
    "ScheduledAttempt",
    "SchedulerReport",
    "schedule_next_attempts",
]