Gitea-Tools/maintenance_drain.py

"""Graceful MCP maintenance-drain mode (#659).

Drain is the visible, capability-gated state that lets an operator stop new
work and quiesce mutations *before* a restart, instead of cutting sessions off
mid-mutation. This module owns the pure decision layer:

* the drain state vocabulary and its normalization;
* the allowlist of safety operations that must keep working while draining
  (heartbeat, release/abandon, checkpoint, and drain exit itself — the exact
  calls an in-flight session needs to finish and hand off);
* the mutation-gate classification consumed by the MCP preflight chokepoint;
* the assignment-stop classification consumed by the allocator;
* the observable status payload sessions read to see the drain (AC4).

Durable state lives in the control-plane DB (``maintenance_drain`` table);
enforcement lives at the existing chokepoints. Nothing here performs I/O, so
both callers can share one decision without importing each other.

Scope note: the machine-verifiable *drain proof* and the restart gate that
consumes it are #661's scope, not this module's. Drain here stops assignment
and mutation and makes the state observable; it never authorizes a restart.
"""

from __future__ import annotations

from typing import Any, Mapping

# ── State vocabulary ──────────────────────────────────────────────────────────

STATE_INACTIVE = "inactive"
STATE_DRAINING = "draining"
DRAIN_STATES = frozenset({STATE_INACTIVE, STATE_DRAINING})

# Typed blocker code surfaced to clients (never a bare string at call sites).
BLOCKER_DRAIN_ACTIVE = "maintenance_drain_active"

# Reason code for the allocator's assignment stop.
REASON_ASSIGNMENT_STOPPED = "maintenance_drain_assignment_stopped"

DRAIN_SCHEMA_VERSION = 6


class MaintenanceDrainError(RuntimeError):
    """Raised when a mutation is refused because drain is active (fail closed)."""

    def __init__(self, message: str, *, decision: Mapping[str, Any] | None = None):
        super().__init__(message)
        self.decision = dict(decision or {})
        self.reason_code = BLOCKER_DRAIN_ACTIVE


# ── Safety allowlist ──────────────────────────────────────────────────────────

# Mutations that stay permitted while draining. Every entry is a *quiesce*
# operation: it either proves an in-flight task is still alive, hands its claim
# back, records the durable state a restart needs, or ends the drain. Nothing
# that creates new work, new branches, new PRs, or new review/merge verdicts is
# on this list — that is the whole point of the drain.
ALLOWLISTED_DRAIN_TASKS: frozenset[str] = frozenset(
    {
        # Liveness of work already in flight.
        "heartbeat_issue_lock",
        "heartbeat_reviewer_pr_lease",
        "post_heartbeat",
        # Handing claims back so nothing is stranded across the restart.
        "release_workflow_lease",
        "release_reviewer_pr_lease",
        "release_merger_pr_lease",
        "abandon_workflow_lease",
        # Durable recovery state (#660) must be writable *during* drain.
        "write_session_checkpoint",
        "checkpoint_session",
        # The drain controls themselves — exit must never be self-blocked.
        "enter_maintenance_drain",
        "exit_maintenance_drain",
    }
)


def normalize_task(task: str | None) -> str:
    """Normalize a task name, tolerating the ``gitea_`` tool-name prefix."""
    name = str(task or "").strip()
    if name.startswith("gitea_"):
        name = name[len("gitea_") :]
    return name


def is_allowlisted_task(task: str | None) -> bool:
    """Is *task* a safety operation permitted while draining?"""
    return normalize_task(task) in ALLOWLISTED_DRAIN_TASKS


def normalize_state(state: str | None) -> str:
    """Normalize a drain state; blank means inactive, unknown fails closed.

    Blank normalizes to ``inactive`` (no drain record = not draining), but an
    unrecognized non-blank value raises: silently treating ``"drainig"`` as
    inactive would disable the gate.
    """
    value = str(state or "").strip().lower()
    if not value:
        return STATE_INACTIVE
    if value not in DRAIN_STATES:
        raise MaintenanceDrainError(
            f"unknown maintenance-drain state {value!r}; expected one of "
            f"{sorted(DRAIN_STATES)} (fail closed)"
        )
    return value


def is_draining(record: Mapping[str, Any] | None) -> bool:
    """Is the given drain record (or None) an active drain?"""
    if not record:
        return False
    return normalize_state(record.get("state")) == STATE_DRAINING


# ── Decisions ─────────────────────────────────────────────────────────────────


def classify_mutation(
    task: str | None,
    record: Mapping[str, Any] | None,
) -> dict[str, Any]:
    """Decide whether *task* may mutate under the given drain record.

    Returns a decision dict with ``allowed``/``deferred`` and, when refused, a
    typed ``reason_code`` plus the one exact next action the caller may take.
    Deferred (not failed): the operation is legal again after drain exits, so
    the caller is told to wait rather than to retry a different way.
    """
    task_norm = normalize_task(task)
    draining = is_draining(record)

    if not draining:
        return {
            "allowed": True,
            "deferred": False,
            "drain_state": STATE_INACTIVE,
            "task": task_norm,
            "allowlisted": is_allowlisted_task(task_norm),
            "reason_code": None,
            "reasons": [],
            "exact_safe_next_action": None,
        }

    if is_allowlisted_task(task_norm):
        return {
            "allowed": True,
            "deferred": False,
            "drain_state": STATE_DRAINING,
            "task": task_norm,
            "allowlisted": True,
            "reason_code": None,
            "reasons": [
                f"task '{task_norm}' is an allowlisted drain safety operation; "
                "permitted so in-flight work can finish and hand off"
            ],
            "exact_safe_next_action": None,
        }

    return {
        "allowed": False,
        "deferred": True,
        "drain_state": STATE_DRAINING,
        "task": task_norm,
        "allowlisted": False,
        "reason_code": BLOCKER_DRAIN_ACTIVE,
        "reasons": [format_drain_reason(task_norm, record)],
        "exact_safe_next_action": (
            "Wait for maintenance drain to exit (or have an authorized "
            "controller call gitea_exit_maintenance_drain), then retry this "
            "mutation. Reads and gitea_maintenance_drain_status stay available."
        ),
    }


def classify_assignment(record: Mapping[str, Any] | None) -> dict[str, Any]:
    """Decide whether the allocator may assign new work (AC2)."""
    if not is_draining(record):
        return {
            "assignment_allowed": True,
            "drain_state": STATE_INACTIVE,
            "reason_code": None,
            "reasons": [],
        }
    return {
        "assignment_allowed": False,
        "drain_state": STATE_DRAINING,
        "reason_code": REASON_ASSIGNMENT_STOPPED,
        "reasons": [
            "maintenance drain is active: new work assignment is stopped and "
            "no lease was created (fail closed, #659)" + _scope_suffix(record)
        ],
    }


def format_drain_reason(task: str | None, record: Mapping[str, Any] | None) -> str:
    """Human-readable refusal line for a drained mutation."""
    task_norm = normalize_task(task) or "(unnamed task)"
    return (
        f"maintenance drain is active: mutation '{task_norm}' is deferred; only "
        "allowlisted drain safety operations "
        f"({', '.join(sorted(ALLOWLISTED_DRAIN_TASKS))}) and reads are permitted "
        "(fail closed, #659)" + _scope_suffix(record)
    )


def format_drain_block_error(decision: Mapping[str, Any]) -> str:
    """Format the typed error message raised at the mutation chokepoint."""
    reasons = list(decision.get("reasons") or [])
    head = reasons[0] if reasons else "maintenance drain is active (fail closed)"
    action = decision.get("exact_safe_next_action")
    return f"{head}. Exact safe next action: {action}" if action else head


def _scope_suffix(record: Mapping[str, Any] | None) -> str:
    """Append the drain's scope/reason/owner facts when the record carries them."""
    if not record:
        return ""
    bits: list[str] = []
    scope = "/".join(
        str(record.get(key) or "") for key in ("remote", "org", "repo")
    ).strip("/")
    if scope:
        bits.append(f"scope {scope}")
    if record.get("reason"):
        bits.append(f"reason: {record['reason']}")
    if record.get("requested_by"):
        bits.append(f"entered by {record['requested_by']}")
    if record.get("entered_at"):
        bits.append(f"at {record['entered_at']}")
    return f" ({'; '.join(bits)})" if bits else ""


# ── Observability (AC4) ───────────────────────────────────────────────────────


def status_payload(
    record: Mapping[str, Any] | None,
    *,
    remote: str = "",
    org: str = "",
    repo: str = "",
) -> dict[str, Any]:
    """Build the session-observable drain status payload.

    Always answers, including when no drain record exists: an absent record is
    a definitive "not draining", not an unknown.
    """
    draining = is_draining(record)
    rec: Mapping[str, Any] = record or {}
    return {
        "drain_state": STATE_DRAINING if draining else STATE_INACTIVE,
        "draining": draining,
        "remote": str(rec.get("remote") or "") or remote,
        "org": str(rec.get("org") or "") or org,
        "repo": str(rec.get("repo") or "") or repo,
        "reason": str(rec.get("reason") or ""),
        "requested_by": str(rec.get("requested_by") or ""),
        "requested_by_profile": str(rec.get("requested_by_profile") or ""),
        "session_id": str(rec.get("session_id") or ""),
        "entered_at": str(rec.get("entered_at") or ""),
        "exited_at": str(rec.get("exited_at") or ""),
        "assignment_stopped": draining,
        "mutations_deferred": draining,
        "allowlisted_tasks": sorted(ALLOWLISTED_DRAIN_TASKS),
        "reads_permitted": True,
        "record_present": bool(record),
        "schema_version": DRAIN_SCHEMA_VERSION,
        "drain_proof_scope": (
            "drain proof and the restart gate that consumes it are #661 scope; "
            "this status never authorizes a restart"
        ),
        "safe_next_action": (
            "Wait for drain to exit before retrying deferred mutations; "
            "allowlisted safety operations and reads remain available."
            if draining
            else "None; maintenance drain is not active."
        ),
    }