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 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 "
                "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,
            # 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",
]