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 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 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.
    """
    report = SchedulerReport()
    now = datetime.now(timezone.utc)
    with session_scope(engine) as session:
        # Workflows currently in a state that needs a worker.
        rows = session.execute(
            text(
                "SELECT workflow_id, current_state, current_tier "
                "FROM workflows "
                "WHERE current_state IN "
                "  ('ANALYZING', 'IMPLEMENTING', 'REVIEWING', "
                "   'CONFLICT_RESOLVING', 'ESCALATING')"
            )
        ).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
            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
                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.
            existing = session.execute(
                text(
                    "SELECT attempt_id FROM workflow_attempts "
                    "WHERE workflow_id = :wf_id AND role = :role "
                    "  AND status IN ('pending', 'in_progress')"
                ),
                {"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 {role} attempt")
                )
                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
                logger.warning(
                    "prefetch failed for workflow %s role %s: %s",
                    r.workflow_id, role, exc,
                )
                report.workflows_skipped += 1
                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 {
        "ANALYZING": "estimator",
        "IMPLEMENTING": "implementer",
        "REVIEWING": "reviewer",
        "CONFLICT_RESOLVING": "conflict_resolver",
    }.get(state)


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,
            "payload": 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."""
    # 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

    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,
            "payload": json.dumps(input_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.
    return result.lastrowid


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