feat(mcp): implement graceful maintenance-drain mode (Closes #659)

Add durable per-repo drain state, allocator assignment stop, mutation deferral with a safety allowlist, observable status, and capability-gated enter/exit tools. Drain proof/restart gate remain #661. Co-Authored-By: Claude Opus 4.8 (1M context) <[email protected]>
2026-07-25 17:02:12 -04:00
parent 76f293eb28
commit e91b94db56
7 changed files with 1056 additions and 1 deletions
@@ -0,0 +1,281 @@
+"""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."
+        ),
+    }