"""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." ), }