From e91b94db567c25128a70ff0ef7f9d4461f67d52a Mon Sep 17 00:00:00 2001 From: Jason Walker <913443@dadeschools.net> Date: Sat, 25 Jul 2026 17:01:42 -0400 Subject: [PATCH] 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) --- allocator_service.py | 41 +++++ control_plane_db.py | 195 +++++++++++++++++++++- docs/mcp-maintenance-drain.md | 45 +++++ gitea_mcp_server.py | 240 +++++++++++++++++++++++++++ maintenance_drain.py | 281 ++++++++++++++++++++++++++++++++ task_capability_map.py | 18 ++ tests/test_maintenance_drain.py | 237 +++++++++++++++++++++++++++ 7 files changed, 1056 insertions(+), 1 deletion(-) create mode 100644 docs/mcp-maintenance-drain.md create mode 100644 maintenance_drain.py create mode 100644 tests/test_maintenance_drain.py diff --git a/allocator_service.py b/allocator_service.py index 9f32fd9..c47a71a 100644 --- a/allocator_service.py +++ b/allocator_service.py @@ -25,6 +25,7 @@ import uuid from dataclasses import dataclass, field from typing import Any, Mapping, Sequence +import maintenance_drain from control_plane_db import ( ControlPlaneDB, ControlPlaneError, @@ -885,6 +886,46 @@ def allocate_next_work( "allocation_mode": (allocation_mode or "").strip() or None, } + # #659 AC2: while maintenance drain is active, no new work is assigned — + # for dry-run and apply alike, so a preview can never be read as evidence + # that work was assignable during the drain. Checked before session + # registration so a drained allocator leaves no new state behind. + try: + drain_record = db.read_maintenance_drain(remote=remote, org=org, repo=repo) + except Exception as exc: # noqa: BLE001 — unreadable drain state fails closed + return { + "success": False, + "outcome": OUTCOME_NO_SAFE, + "reasons": [ + f"maintenance-drain state lookup failed: {exc} (fail closed, #659)" + ], + "skipped": [], + "assignment": None, + "substrate": "control_plane_db", + } + + drain_decision = maintenance_drain.classify_assignment(drain_record) + if not drain_decision["assignment_allowed"]: + return { + "success": True, + "outcome": OUTCOME_WAIT, + "apply": apply, + "role": role_norm, + "allocation_mode": mode, + "remote": remote, + "org": org, + "repo": repo, + "selected": None, + "reasons": list(drain_decision["reasons"]), + "reason_code": drain_decision["reason_code"], + "skipped": [], + "assignment": None, + "substrate": "control_plane_db", + "maintenance_drain": maintenance_drain.status_payload( + drain_record, remote=remote, org=org, repo=repo + ), + } + session_id = (session_id or "").strip() or f"alloc-{uuid.uuid4().hex[:12]}" try: db.upsert_session( diff --git a/control_plane_db.py b/control_plane_db.py index a1f90d2..cd0ae77 100644 --- a/control_plane_db.py +++ b/control_plane_db.py @@ -31,8 +31,9 @@ from typing import Any, Iterator, Sequence import dependency_graph import gitea_audit +import maintenance_drain -SCHEMA_VERSION = 5 +SCHEMA_VERSION = 6 # Assignable work kinds only — raw monitoring incidents are never work items. WORK_KINDS = frozenset({"issue", "pr"}) @@ -239,6 +240,31 @@ CREATE INDEX IF NOT EXISTS idx_session_checkpoints_session CREATE INDEX IF NOT EXISTS idx_session_checkpoints_work ON session_checkpoints(remote, org, repo, work_kind, work_number); +-- Graceful maintenance-drain state (#659). One current row per repository +-- scope — drain is a *state*, not a history, so entering and exiting update +-- the same row and every transition is audited to ``events``. Creating the +-- table is the v5->v6 migration: additive, idempotent, and it never touches +-- prior tables. ``state`` is CHECK-constrained so an unknown value can never +-- be written and later read as "not draining". +CREATE TABLE IF NOT EXISTS maintenance_drain ( + drain_id TEXT PRIMARY KEY, + remote TEXT NOT NULL, + org TEXT NOT NULL, + repo TEXT NOT NULL, + state TEXT NOT NULL DEFAULT 'inactive' + CHECK (state IN ('inactive', 'draining')), + reason TEXT NOT NULL DEFAULT '', + requested_by TEXT NOT NULL DEFAULT '', + requested_by_profile TEXT NOT NULL DEFAULT '', + session_id TEXT NOT NULL DEFAULT '', + entered_at TEXT NOT NULL DEFAULT '', + exited_at TEXT NOT NULL DEFAULT '', + drain_schema_version INTEGER NOT NULL DEFAULT 6, + created_at TEXT NOT NULL, + updated_at TEXT NOT NULL, + UNIQUE (remote, org, repo) +); + -- Model usage, token cost, latency, and performance events (#651) CREATE TABLE IF NOT EXISTS usage_events ( usage_id INTEGER PRIMARY KEY AUTOINCREMENT, @@ -3025,3 +3051,170 @@ class ControlPlaneDB: "live_lease_id": None if live_lease_id is None else str(live_lease_id), "reconcile_action": "reconcile_required" if stale else "safe_to_resume", } + + # ── Maintenance drain (#659) ───────────────────────────────────────────── + + @staticmethod + def _maintenance_drain_row(row: sqlite3.Row | None) -> dict[str, Any] | None: + """Convert a ``maintenance_drain`` row to a plain record.""" + if row is None: + return None + return {key: row[key] for key in row.keys()} + + def read_maintenance_drain( + self, *, remote: str, org: str, repo: str + ) -> dict[str, Any] | None: + """Return the current drain record for a scope, or None if never set. + + None and a stored ``inactive`` row mean the same thing to callers — + ``maintenance_drain.is_draining`` treats both as not draining — so the + read never has to invent a record to answer the gate. + """ + with self._tx(immediate=False) as conn: + row = conn.execute( + """ + SELECT * FROM maintenance_drain + WHERE remote = ? AND org = ? AND repo = ? + """, + (str(remote or ""), str(org or ""), str(repo or "")), + ).fetchone() + return self._maintenance_drain_row(row) + + def set_maintenance_drain( + self, + *, + remote: str, + org: str, + repo: str, + state: str, + reason: str = "", + requested_by: str = "", + requested_by_profile: str = "", + session_id: str = "", + ) -> dict[str, Any]: + """Enter or exit maintenance drain for one repository scope (AC1). + + The state transition is audited to ``events`` — entering and exiting + are exactly the moments an operator has to be able to reconstruct + later. Re-entering an already-draining scope is idempotent: it refreshes + the reason/owner metadata, keeps the original ``entered_at``, and + records no duplicate transition event. + + Capability authorization happens above this layer (the drain tasks + carry a non-``gitea.*`` permission in the task capability map); the DB + records who asked and why, and never grants the right itself. + """ + state_norm = maintenance_drain.normalize_state(state) + raw = { + "reason": str(reason or ""), + "requested_by": str(requested_by or ""), + "requested_by_profile": str(requested_by_profile or ""), + "session_id": str(session_id or ""), + } + clean = gitea_audit.redact(raw) + remote_s, org_s, repo_s = str(remote or ""), str(org or ""), str(repo or "") + now_s = _ts() + + with self._tx() as conn: + existing = conn.execute( + """ + SELECT * FROM maintenance_drain + WHERE remote = ? AND org = ? AND repo = ? + """, + (remote_s, org_s, repo_s), + ).fetchone() + + prior_state = ( + maintenance_drain.normalize_state(existing["state"]) + if existing is not None + else maintenance_drain.STATE_INACTIVE + ) + transitioned = prior_state != state_norm + + prior_entered = ( + str(existing["entered_at"] or "") if existing is not None else "" + ) + prior_exited = ( + str(existing["exited_at"] or "") if existing is not None else "" + ) + if state_norm == maintenance_drain.STATE_DRAINING: + # A re-entry keeps the original entry time (the drain never + # stopped); a fresh entry stamps now and clears the old exit. + entered_at = prior_entered if (not transitioned and prior_entered) else now_s + exited_at = "" + else: + entered_at = prior_entered + exited_at = now_s if (transitioned or not prior_exited) else prior_exited + + if existing is None: + drain_id = uuid.uuid4().hex + conn.execute( + """ + INSERT INTO maintenance_drain( + drain_id, remote, org, repo, state, reason, + requested_by, requested_by_profile, session_id, + entered_at, exited_at, drain_schema_version, + created_at, updated_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + drain_id, remote_s, org_s, repo_s, state_norm, + clean["reason"], clean["requested_by"], + clean["requested_by_profile"], clean["session_id"], + entered_at, exited_at, + maintenance_drain.DRAIN_SCHEMA_VERSION, now_s, now_s, + ), + ) + else: + drain_id = str(existing["drain_id"]) + conn.execute( + """ + UPDATE maintenance_drain + SET state = ?, reason = ?, requested_by = ?, + requested_by_profile = ?, session_id = ?, + entered_at = ?, exited_at = ?, + drain_schema_version = ?, updated_at = ? + WHERE drain_id = ? + """, + ( + state_norm, clean["reason"], clean["requested_by"], + clean["requested_by_profile"], clean["session_id"], + entered_at, exited_at, + maintenance_drain.DRAIN_SCHEMA_VERSION, now_s, drain_id, + ), + ) + + if transitioned: + event_type = ( + "maintenance_drain_enter" + if state_norm == maintenance_drain.STATE_DRAINING + else "maintenance_drain_exit" + ) + conn.execute( + """ + INSERT INTO events(work_item_id, event_type, message, created_at) + VALUES (NULL, ?, ?, ?) + """, + ( + event_type, + f"drain {drain_id} scope {remote_s}/{org_s}/{repo_s} " + f"{prior_state} -> {state_norm} by " + f"{clean['requested_by'] or '(unknown)'} " + f"({clean['requested_by_profile'] or 'no profile'}); " + f"reason: {clean['reason'] or '(none)'}", + now_s, + ), + ) + + row = conn.execute( + "SELECT * FROM maintenance_drain WHERE drain_id = ?", (drain_id,) + ).fetchone() + + record = self._maintenance_drain_row(row) or {} + return { + "record": record, + "drain_id": drain_id, + "state": state_norm, + "prior_state": prior_state, + "transitioned": transitioned, + } diff --git a/docs/mcp-maintenance-drain.md b/docs/mcp-maintenance-drain.md new file mode 100644 index 0000000..f5131b7 --- /dev/null +++ b/docs/mcp-maintenance-drain.md @@ -0,0 +1,45 @@ +# MCP maintenance-drain mode (#659) + +Graceful **maintenance drain** stops new work assignment and defers non-allowlisted +mutations so sessions can finish critical handoffs and checkpoint before a +restart. It is **not** a restart authorization: the drain *proof* and apply gate +remain #661. + +## State + +Per repository scope (`remote`/`org`/`repo`) in the control-plane DB table +`maintenance_drain` (schema v6): + +| State | Meaning | +|-------|---------| +| `inactive` | Normal operation (also: no row) | +| `draining` | Assignment stopped; non-allowlisted mutations deferred | + +Enter/exit transitions are audited as `maintenance_drain_enter` / +`maintenance_drain_exit` events. + +## Tools + +| Tool | Permission | Effect | +|------|------------|--------| +| `gitea_maintenance_drain_status` | `gitea.read` | Observe drain (every session) | +| `gitea_enter_maintenance_drain` | `runtime.maintenance_drain` | Enter drain (capability-gated) | +| `gitea_exit_maintenance_drain` | `runtime.maintenance_drain` | Exit drain | + +`runtime.maintenance_drain` is intentionally **not** a `gitea.*` op, so ordinary +author profiles cannot enter drain by accident. + +## Enforcement + +1. **Allocator** (`allocate_next_work`): while draining, returns `outcome=wait` + with `reason_code=maintenance_drain_assignment_stopped` for dry-run and apply. +2. **Mutation preflight** (`verify_preflight_purity`): non-allowlisted mutation + tasks raise `MaintenanceDrainError` with a typed next action. +3. **Allowlist** (safety only): heartbeats, lease release/abandon, session + checkpoints, enter/exit drain. Reads always work. + +## Restart relationship + +Drain mode prepares the blast radius. Restart apply still requires a clean +`DrainProof` (#661) or authorized break-glass. Status payloads never claim +restart permission. diff --git a/gitea_mcp_server.py b/gitea_mcp_server.py index aab8249..90a54df 100644 --- a/gitea_mcp_server.py +++ b/gitea_mcp_server.py @@ -1472,6 +1472,9 @@ def verify_preflight_purity( # contaminated by manual MCP daemon process killing (reconciler-exempt). _enforce_runtime_recovery_contamination_gate(task, remote) + # #659 AC3: defer non-allowlisted mutations while maintenance drain is active. + _enforce_maintenance_drain_gate(task, remote=remote, org=org, repo=repo) + ctx = _resolve_namespace_mutation_context(worktree_path) workspace = ctx["workspace_path"] canonical_root = ctx["canonical_repo_root"] @@ -2004,6 +2007,55 @@ def _enforce_runtime_recovery_contamination_gate( ) +def _enforce_maintenance_drain_gate( + task: str | None, + remote: str | None = None, + org: str | None = None, + repo: str | None = None, +) -> None: + """#659 AC3: defer non-allowlisted mutations while drain is active. + + The single mutation chokepoint already used by every gated task, so drain + coverage cannot drift per-tool. Allowlisted safety operations (heartbeat, + release/abandon, checkpoint, drain exit) pass through so an in-flight + session can still finish and hand off; everything else is deferred with a + typed blocker. Unreadable drain state fails closed — a drain that cannot be + read is not evidence that no drain is running. + """ + if _preflight_in_test_mode() and not os.environ.get( + "GITEA_TEST_FORCE_MAINTENANCE_DRAIN" + ): + return + if maintenance_drain.is_allowlisted_task(task): + return + + try: + _h, o, r = _resolve(remote, None, org, repo) + except Exception: # noqa: BLE001 — scope resolution is best-effort here + o, r = (org or ""), (repo or "") + + db, errs = _control_plane_db_or_error() + if db is None: + raise RuntimeError( + "maintenance-drain state could not be read: " + f"{'; '.join(errs) or 'control-plane DB unavailable'} (fail closed, #659)" + ) + try: + record = db.read_maintenance_drain(remote=remote or "", org=o, repo=r) + except Exception as exc: # noqa: BLE001 + raise RuntimeError( + f"maintenance-drain state could not be read: {_redact(str(exc))} " + "(fail closed, #659)" + ) from exc + + decision = maintenance_drain.classify_mutation(task, record) + if not decision["allowed"]: + raise maintenance_drain.MaintenanceDrainError( + maintenance_drain.format_drain_block_error(decision), + decision=decision, + ) + + def _enforce_stable_branch_contamination_gate( task: str | None, remote: str | None = None, @@ -2064,6 +2116,7 @@ import allocator_service # noqa: E402 import allocator_dependencies # noqa: E402 import dependency_graph # noqa: E402 # #784 durable dependency edges import control_plane_db # noqa: E402 +import maintenance_drain # noqa: E402 # #659 graceful maintenance-drain mode import lease_lifecycle # noqa: E402 import lease_policy # noqa: E402 import workflow_dashboard # noqa: E402 # #605 live queue/lease dashboard @@ -22559,6 +22612,193 @@ def gitea_workflow_dashboard( return payload +@mcp.tool() +def gitea_maintenance_drain_status( + remote: str = "dadeschools", + host: str | None = None, + org: str | None = None, + repo: str | None = None, +) -> dict: + """Read-only: current maintenance-drain state for a repository scope (#659 AC4). + + Every session must be able to observe drain so it can stop creating new work + and finish only allowlisted safety operations. Never mutates; never restarts. + """ + read_block = _profile_operation_gate("gitea.read") + if read_block: + return { + "success": False, + "read_only": True, + "reasons": read_block, + "permission_report": _permission_block_report("gitea.read"), + } + try: + _h, o, r = _resolve(remote, host, org, repo) + except ValueError as exc: + return {"success": False, "read_only": True, "reasons": [str(exc)]} + db, errs = _control_plane_db_or_error() + if db is None: + return { + "success": False, + "read_only": True, + "reasons": errs or ["control-plane DB unavailable"], + "maintenance_drain": maintenance_drain.status_payload( + None, remote=remote, org=o, repo=r + ), + } + try: + record = db.read_maintenance_drain(remote=remote, org=o, repo=r) + except Exception as exc: # noqa: BLE001 + return { + "success": False, + "read_only": True, + "reasons": [f"drain state unreadable: {_redact(str(exc))}"], + } + payload = maintenance_drain.status_payload( + record, remote=remote, org=o, repo=r + ) + return {"success": True, "read_only": True, "maintenance_drain": payload} + + +@mcp.tool() +def gitea_enter_maintenance_drain( + reason: str = "", + remote: str = "dadeschools", + host: str | None = None, + org: str | None = None, + repo: str | None = None, + session_id: str | None = None, +) -> dict: + """Enter graceful maintenance-drain mode for a repository scope (#659 AC1). + + Stops new assignment and defers non-allowlisted mutations until exit. Requires + ``runtime.maintenance_drain`` (controller/lifecycle capability — not granted + by ordinary Gitea author profiles). Audited in the control-plane event log. + """ + cap_block = _profile_operation_gate("runtime.maintenance_drain") + if cap_block: + return { + "success": False, + "performed": False, + "reasons": cap_block, + "permission_report": _permission_block_report( + "runtime.maintenance_drain" + ), + } + try: + _h, o, r = _resolve(remote, host, org, repo) + except ValueError as exc: + return {"success": False, "performed": False, "reasons": [str(exc)]} + db, errs = _control_plane_db_or_error() + if db is None: + return { + "success": False, + "performed": False, + "reasons": errs or ["control-plane DB unavailable"], + } + profile = get_profile() or {} + try: + result = db.set_maintenance_drain( + remote=remote, + org=o, + repo=r, + state=maintenance_drain.STATE_DRAINING, + reason=reason or "operator-entered maintenance drain", + requested_by=str( + (profile.get("identity") or {}).get("username") + or profile.get("expected_username") + or "" + ), + requested_by_profile=str(profile.get("profile_name") or ""), + session_id=str(session_id or ""), + ) + except Exception as exc: # noqa: BLE001 + return { + "success": False, + "performed": False, + "reasons": [f"enter drain failed: {_redact(str(exc))}"], + } + record = result.get("record") or {} + return { + "success": True, + "performed": True, + "transitioned": bool(result.get("transitioned")), + "state": result.get("state"), + "prior_state": result.get("prior_state"), + "drain_id": result.get("drain_id"), + "maintenance_drain": maintenance_drain.status_payload( + record, remote=remote, org=o, repo=r + ), + } + + +@mcp.tool() +def gitea_exit_maintenance_drain( + reason: str = "", + remote: str = "dadeschools", + host: str | None = None, + org: str | None = None, + repo: str | None = None, + session_id: str | None = None, +) -> dict: + """Exit graceful maintenance-drain mode (#659 AC1). Restores assignment and mutations.""" + cap_block = _profile_operation_gate("runtime.maintenance_drain") + if cap_block: + return { + "success": False, + "performed": False, + "reasons": cap_block, + "permission_report": _permission_block_report( + "runtime.maintenance_drain" + ), + } + try: + _h, o, r = _resolve(remote, host, org, repo) + except ValueError as exc: + return {"success": False, "performed": False, "reasons": [str(exc)]} + db, errs = _control_plane_db_or_error() + if db is None: + return { + "success": False, + "performed": False, + "reasons": errs or ["control-plane DB unavailable"], + } + profile = get_profile() or {} + try: + result = db.set_maintenance_drain( + remote=remote, + org=o, + repo=r, + state=maintenance_drain.STATE_INACTIVE, + reason=reason or "operator-exited maintenance drain", + requested_by=str( + (profile.get("identity") or {}).get("username") + or profile.get("expected_username") + or "" + ), + requested_by_profile=str(profile.get("profile_name") or ""), + session_id=str(session_id or ""), + ) + except Exception as exc: # noqa: BLE001 + return { + "success": False, + "performed": False, + "reasons": [f"exit drain failed: {_redact(str(exc))}"], + } + record = result.get("record") or {} + return { + "success": True, + "performed": True, + "transitioned": bool(result.get("transitioned")), + "state": result.get("state"), + "prior_state": result.get("prior_state"), + "drain_id": result.get("drain_id"), + "maintenance_drain": maintenance_drain.status_payload( + record, remote=remote, org=o, repo=r + ), + } + + @mcp.tool() def gitea_request_mcp_restart( remote: str = "dadeschools", diff --git a/maintenance_drain.py b/maintenance_drain.py new file mode 100644 index 0000000..584111e --- /dev/null +++ b/maintenance_drain.py @@ -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." + ), + } diff --git a/task_capability_map.py b/task_capability_map.py index 878cf0a..c2576c6 100644 --- a/task_capability_map.py +++ b/task_capability_map.py @@ -397,6 +397,24 @@ TASK_CAPABILITY_MAP: dict[str, dict[str, str]] = { "role": "controller", }, + # #659 maintenance drain. Same reasoning as the lifecycle controls above: + # entering/exiting drain quiesces a whole namespace, so it carries a + # non-``gitea.*`` permission that no configured Gitea profile satisfies by + # accident (AC1 — capability-gated and audited). Reading drain state is + # ordinary read authority: every session must be able to see the drain (AC4). + "enter_maintenance_drain": { + "permission": "runtime.maintenance_drain", + "role": "controller", + }, + "exit_maintenance_drain": { + "permission": "runtime.maintenance_drain", + "role": "controller", + }, + "maintenance_drain_status": { + "permission": "gitea.read", + "role": "author", + }, + # #601 first-class lease lifecycle — inspect/list need read; mutations gate on # ownership in the control-plane DB (not a separate Gitea write permission). "list_workflow_leases": { diff --git a/tests/test_maintenance_drain.py b/tests/test_maintenance_drain.py new file mode 100644 index 0000000..b3aac32 --- /dev/null +++ b/tests/test_maintenance_drain.py @@ -0,0 +1,237 @@ +"""Tests for graceful MCP maintenance-drain mode (#659). + +Acceptance coverage: + +1. Enter/exit is durable and audited (DB substrate). +2. New work assignment stops during drain (allocator WAIT). +3. Mutations deferred except allowlisted safety ops. +4. Sessions can observe drain state. +5. Fail-closed on unreadable drain state. +""" + +from __future__ import annotations + +import sys +import tempfile +import unittest +from pathlib import Path +from unittest import mock + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +import maintenance_drain +from control_plane_db import ControlPlaneDB +from allocator_service import WorkCandidate, allocate_next_work, OUTCOME_WAIT + + +class TestDrainDecisions(unittest.TestCase): + def test_inactive_allows_mutations_and_assignment(self): + decision = maintenance_drain.classify_mutation("create_pr", None) + self.assertTrue(decision["allowed"]) + self.assertFalse(decision["deferred"]) + assign = maintenance_drain.classify_assignment(None) + self.assertTrue(assign["assignment_allowed"]) + + def test_draining_defers_non_allowlisted_mutation(self): + record = {"state": "draining", "remote": "prgs", "org": "o", "repo": "r"} + decision = maintenance_drain.classify_mutation("create_pr", record) + self.assertFalse(decision["allowed"]) + self.assertTrue(decision["deferred"]) + self.assertEqual(decision["reason_code"], maintenance_drain.BLOCKER_DRAIN_ACTIVE) + self.assertIn("create_pr", decision["reasons"][0]) + + def test_allowlisted_safety_ops_pass_during_drain(self): + record = {"state": "draining"} + for task in ( + "heartbeat_issue_lock", + "gitea_release_reviewer_pr_lease", + "write_session_checkpoint", + "exit_maintenance_drain", + ): + with self.subTest(task=task): + decision = maintenance_drain.classify_mutation(task, record) + self.assertTrue(decision["allowed"], decision) + + def test_assignment_stopped_during_drain(self): + record = {"state": "draining", "reason": "upgrade"} + decision = maintenance_drain.classify_assignment(record) + self.assertFalse(decision["assignment_allowed"]) + self.assertEqual( + decision["reason_code"], maintenance_drain.REASON_ASSIGNMENT_STOPPED + ) + + def test_unknown_state_fails_closed(self): + with self.assertRaises(maintenance_drain.MaintenanceDrainError): + maintenance_drain.normalize_state("drainig") + + def test_status_payload_always_answers(self): + inactive = maintenance_drain.status_payload(None, remote="prgs", org="o", repo="r") + self.assertFalse(inactive["draining"]) + self.assertTrue(inactive["reads_permitted"]) + active = maintenance_drain.status_payload( + {"state": "draining", "reason": "reboot", "requested_by": "ops"}, + remote="prgs", + org="o", + repo="r", + ) + self.assertTrue(active["draining"]) + self.assertTrue(active["assignment_stopped"]) + self.assertTrue(active["mutations_deferred"]) + self.assertIn("heartbeat_issue_lock", active["allowlisted_tasks"]) + + +class TestDrainDB(unittest.TestCase): + def setUp(self): + self._tmp = tempfile.TemporaryDirectory() + self.addCleanup(self._tmp.cleanup) + self.db = ControlPlaneDB(db_path=str(Path(self._tmp.name) / "cp.sqlite3")) + + def test_enter_exit_idempotent_and_audited(self): + first = self.db.set_maintenance_drain( + remote="prgs", + org="Scaled-Tech-Consulting", + repo="Gitea-Tools", + state="draining", + reason="planned restart", + requested_by="sysadmin", + requested_by_profile="prgs-controller", + session_id="s1", + ) + self.assertTrue(first["transitioned"]) + self.assertEqual(first["state"], "draining") + self.assertTrue(maintenance_drain.is_draining(first["record"])) + + again = self.db.set_maintenance_drain( + remote="prgs", + org="Scaled-Tech-Consulting", + repo="Gitea-Tools", + state="draining", + reason="still draining", + requested_by="sysadmin", + requested_by_profile="prgs-controller", + session_id="s1", + ) + self.assertFalse(again["transitioned"]) + self.assertEqual(again["record"]["entered_at"], first["record"]["entered_at"]) + + exited = self.db.set_maintenance_drain( + remote="prgs", + org="Scaled-Tech-Consulting", + repo="Gitea-Tools", + state="inactive", + reason="done", + requested_by="sysadmin", + requested_by_profile="prgs-controller", + session_id="s1", + ) + self.assertTrue(exited["transitioned"]) + self.assertFalse(maintenance_drain.is_draining(exited["record"])) + self.assertTrue(exited["record"]["exited_at"]) + + # Events recorded for transitions only (enter + exit). + with self.db._tx(immediate=False) as conn: + rows = conn.execute( + "SELECT event_type FROM events WHERE event_type LIKE 'maintenance_drain_%' " + "ORDER BY event_id" + ).fetchall() + types = [r[0] for r in rows] + self.assertEqual(types, ["maintenance_drain_enter", "maintenance_drain_exit"]) + + def test_read_missing_is_none_not_error(self): + self.assertIsNone( + self.db.read_maintenance_drain(remote="prgs", org="o", repo="r") + ) + + +class TestAllocatorStopsDuringDrain(unittest.TestCase): + def setUp(self): + self._tmp = tempfile.TemporaryDirectory() + self.addCleanup(self._tmp.cleanup) + self.db = ControlPlaneDB(db_path=str(Path(self._tmp.name) / "cp.sqlite3")) + + def test_allocate_returns_wait_while_draining(self): + self.db.set_maintenance_drain( + remote="prgs", + org="Scaled-Tech-Consulting", + repo="Gitea-Tools", + state="draining", + reason="test", + requested_by="tester", + ) + candidates = [ + WorkCandidate( + kind="issue", + number=659, + title="drain", + labels=("status:ready",), + priority=20, + ) + ] + result = allocate_next_work( + self.db, + role="author", + session_id="test-session", + remote="prgs", + org="Scaled-Tech-Consulting", + repo="Gitea-Tools", + apply=False, + candidates=candidates, + username="jcwalker3", + profile_name="prgs-author", + ) + self.assertEqual(result["outcome"], OUTCOME_WAIT) + self.assertIsNone(result.get("selected")) + self.assertEqual( + result.get("reason_code"), + maintenance_drain.REASON_ASSIGNMENT_STOPPED, + ) + self.assertTrue(result["maintenance_drain"]["draining"]) + + def test_allocate_works_when_inactive(self): + candidates = [ + WorkCandidate( + kind="issue", + number=659, + title="drain", + labels=("status:ready",), + priority=20, + ) + ] + result = allocate_next_work( + self.db, + role="author", + session_id="test-session-2", + remote="prgs", + org="Scaled-Tech-Consulting", + repo="Gitea-Tools", + apply=False, + candidates=candidates, + username="jcwalker3", + profile_name="prgs-author", + ) + self.assertNotEqual( + result.get("reason_code"), + maintenance_drain.REASON_ASSIGNMENT_STOPPED, + ) + + +class TestCapabilityMap(unittest.TestCase): + def test_drain_tasks_mapped(self): + import task_capability_map as tcm + + self.assertEqual( + tcm.required_permission("enter_maintenance_drain"), + "runtime.maintenance_drain", + ) + self.assertEqual( + tcm.required_permission("exit_maintenance_drain"), + "runtime.maintenance_drain", + ) + self.assertEqual( + tcm.required_permission("maintenance_drain_status"), + "gitea.read", + ) + + +if __name__ == "__main__": + unittest.main()