Gitea-Tools/mcp_fleet_retirement.py

"""CAS-protected retirement planning for stale worker registrations (#980).

#978 (merged PR #979) made the fleet observable: every registered namespace
worker, its instance attribution, its heartbeat freshness, and a structured
classification. It deliberately stopped there — the snapshot is read-only and
the control plane still had no sanctioned way to retire registry rows whose
owning process is conclusively gone.

This module is the *decision layer* for that retirement. It is pure: callers
supply registry rows, a clock, and a PID probe; nothing here opens SQLite,
scans process tables, or mutates state. The transactional apply lives in
:meth:`mcp_worker_identity.WorkerRegistry.retire_stale_workers`, which calls
back into these same pure functions so plan and apply can never disagree about
what "the registry looks like" or "which rows are eligible".

Why a separate token
--------------------

``mcp_fleet_snapshot._consistency_token`` seeds its digest with ``snapshot_at``
at second precision, so ``registry_revision`` changes on every call even when
no registry row changed. A compare-and-swap gated on it can never pass — a
dry-run/apply cycle spanning more than one second aborts unconditionally. That
token is still useful as an observation stamp, so it is left exactly as it is;
#980 gets its own :func:`registry_fingerprint`, derived *only* from canonical
retirement-relevant row content:

* identical registry contents observed at any two times produce the same token,
* row order never affects the token (serialized rows are sorted),
* any create/delete/identity/liveness/ownership/registration-state change to a
  retirement-relevant field changes the token.

Fail-closed posture
-------------------

A worker is retired only when the control plane *conclusively* establishes it
is a stale orphan. Missing evidence is never read as permission: an unprobeable
PID, an unparsable heartbeat, a row that shares identity evidence with a live
or unprobeable worker, a foreign or absent repository binding, or a worker that
still owns an active workflow lease all preserve the row.

Trusted launcher identity (``inst-…`` provenance) is deliberately *not* part of
the conjunction. #980 places trusted ``client_instance_id`` propagation out of
scope and lists "backfilling trusted identity for legacy workers" as a non-goal;
requiring it here would preserve every legacy row forever and make the feature
inert. What *is* required is that the registry fields the conjunction reads are
actually present — see :data:`REQUIRED_IDENTITY_FIELDS`.
"""

from __future__ import annotations

import hashlib
from datetime import datetime
from typing import Any, Callable, Iterable, Mapping, Sequence

import mcp_fleet_snapshot as fleet
import mcp_worker_identity as mwi

# --- Outcomes -------------------------------------------------------------

OUTCOME_PLANNED = "planned"
OUTCOME_APPLIED = "applied"
OUTCOME_REGISTRY_MOVED = "registry_revision_moved"
OUTCOME_CANDIDATES_MOVED = "candidate_set_moved"
OUTCOME_ALREADY_RETIRED = "already_retired"
OUTCOME_NOTHING_REQUESTED = "nothing_requested"

# --- Reason codes ---------------------------------------------------------

#: The only reason code that authorizes retirement.
REASON_ELIGIBLE = "eligible_stale_orphan"

REASON_ALREADY_TERMINAL = "already_terminal_registration"
REASON_AMBIGUOUS_OWNERSHIP = "ambiguous_ownership_state"
REASON_CONFLICTING_IDENTITY = "conflicting_identity_evidence"
REASON_FOREIGN_REPOSITORY = "repository_binding_ambiguous"
REASON_HEARTBEAT_FRESH = "heartbeat_not_expired"
REASON_INCOMPLETE_IDENTITY = "incomplete_registry_identity"
REASON_NOT_IN_PLAN = "not_in_current_plan"
REASON_PID_ALIVE = "pid_alive"
REASON_PID_UNKNOWN = "pid_liveness_unknown"
REASON_PROTECTED_OWNER = "protected_active_workflow_owner"
REASON_ROW_CHANGED = "row_changed_since_plan"
REASON_ROW_MISSING = "registration_missing"
REASON_UNPARSABLE_HEARTBEAT = "unparsable_heartbeat"
REASON_WORKER_LIVE = "worker_live"

#: Registry columns that must carry a usable value before the eligibility
#: conjunction can even be evaluated. Absence is ambiguity, not permission.
REQUIRED_IDENTITY_FIELDS: tuple[str, ...] = (
    "worker_identity",
    "client_instance_id",
    "session_id",
    "generation_id",
    "status",
    "started_at",
    "last_heartbeat_at",
    "heartbeat_ttl_seconds",
    "pid",
)

#: Canonical retirement-relevant content. Ordering here is fixed and part of
#: the token contract; adding a field changes every fingerprint, so a change
#: here is a deliberate contract revision.
#:
#: Deliberately excluded: ``token_fingerprint`` (credential-adjacent, never a
#: retirement input), the four ``*_revision`` columns (revision drift is an
#: independent restart concern and is not part of the eligibility conjunction),
#: and the ``retired_*`` bookkeeping columns this feature adds.
FINGERPRINT_FIELDS: tuple[str, ...] = (
    "worker_identity",
    "client_name",
    "client_instance_id",
    "session_id",
    "generation_id",
    "role",
    "profile",
    "namespace",
    "remote",
    "repository_binding",
    "pid",
    "process_identity",
    "transport",
    "started_at",
    "last_heartbeat_at",
    "heartbeat_ttl_seconds",
    "fencing_epoch",
    "status",
    "fleet_run_id",
    "authenticated_account",
    "instance_id_provenance",
)

_FINGERPRINT_VERSION = "registryfp-v1"
_CANDIDATE_VERSION = "candidatefp-v1"
_UNIT = "\x1f"
_RECORD = "\x1e"


def _canon(value: Any) -> str:
    """Stable text for one field value, independent of Python/SQLite typing.

    ``900`` and ``900.0`` are the same TTL and must hash the same; a value that
    round-trips through SQLite as REAL must not produce a different token than
    the same value supplied by a caller as ``int``.
    """
    if value is None:
        return ""
    if isinstance(value, bool):
        return "true" if value else "false"
    if isinstance(value, float):
        if value != value or value in (float("inf"), float("-inf")):
            return repr(value)
        if value.is_integer():
            return str(int(value))
        return repr(value)
    if isinstance(value, int):
        return str(value)
    return str(value)


def _serialize_row(row: Mapping[str, Any]) -> str:
    return _UNIT.join(f"{name}={_canon(row.get(name))}" for name in FINGERPRINT_FIELDS)


def _digest(version: str, serialized: Sequence[str], prefix: str) -> str:
    ordered = sorted(serialized)
    material = _RECORD.join([version, str(len(ordered)), *ordered])
    return f"{prefix}-{hashlib.sha256(material.encode('utf-8')).hexdigest()[:32]}"


def registry_fingerprint(rows: Iterable[Mapping[str, Any]]) -> str:
    """Content-derived compare-and-swap token for the worker registry (#980).

    Derived exclusively from :data:`FINGERPRINT_FIELDS` across every row. It
    contains no ``snapshot_at``, wall-clock, request, or report-generation
    time, so two observations of an unchanged registry always agree, and the
    serialized rows are sorted so iteration order cannot perturb the digest.
    """
    return _digest(
        _FINGERPRINT_VERSION,
        [_serialize_row(row) for row in rows],
        "registryfp",
    )


def candidate_fingerprint(candidate_rows: Iterable[Mapping[str, Any]]) -> str:
    """Exact-candidate-set token over the selected rows' canonical content.

    A matching :func:`registry_fingerprint` already implies these rows are
    unchanged; this second token additionally pins *which* rows the operator
    approved, so an apply can never widen or narrow the approved set.
    """
    return _digest(
        _CANDIDATE_VERSION,
        [_serialize_row(row) for row in candidate_rows],
        "candidatefp",
    )


def _probe_pid(
    pid: Any, pid_alive_probe: Callable[[int | None], bool | None] | None
) -> bool | None:
    if pid_alive_probe is None or pid is None:
        return None
    try:
        result = pid_alive_probe(pid)
    except Exception:
        return None
    return None if result is None else bool(result)


def _evidence(
    row: Mapping[str, Any],
    snapshot_row: Mapping[str, Any],
    pid_alive: bool | None,
) -> dict[str, Any]:
    liveness = snapshot_row.get("liveness") or {}
    return {
        "worker_identity": row.get("worker_identity"),
        "client_type": snapshot_row.get("client_type"),
        "client_instance_id": row.get("client_instance_id"),
        "fleet_run_id": row.get("fleet_run_id"),
        "namespace": row.get("namespace"),
        "profile": row.get("profile"),
        "declared_role": row.get("role"),
        "session_id": row.get("session_id"),
        "generation_id": row.get("generation_id"),
        "fencing_epoch": row.get("fencing_epoch"),
        "process_identity": snapshot_row.get("process_identity"),
        "pid": row.get("pid"),
        "pid_alive": pid_alive,
        "repository_binding": row.get("repository_binding"),
        "foreign_repository": bool(snapshot_row.get("foreign_repository")),
        "status": row.get("status"),
        "started_at": row.get("started_at"),
        "last_heartbeat_at": row.get("last_heartbeat_at"),
        "heartbeat_ttl_seconds": row.get("heartbeat_ttl_seconds"),
        "heartbeat_age_seconds": liveness.get("heartbeat_age_seconds"),
        "heartbeat_fresh": liveness.get("heartbeat_fresh"),
        "live": bool(snapshot_row.get("live")),
        "ownership_state": snapshot_row.get("ownership_state"),
        "instance_id_provenance": snapshot_row.get("instance_id_provenance"),
        "instance_identity_trusted": bool(
            snapshot_row.get("instance_identity_trusted")
        ),
    }


def _conflict_keys(
    row: Mapping[str, Any], snapshot_row: Mapping[str, Any]
) -> list[tuple[str, str]]:
    keys: list[tuple[str, str]] = []
    for name, value in (
        ("session_id", row.get("session_id")),
        ("generation_id", row.get("generation_id")),
        ("process_identity", snapshot_row.get("process_identity")),
        ("pid", row.get("pid")),
    ):
        text = _canon(value)
        if text:
            keys.append((name, text))
    return keys


def _missing_identity_fields(row: Mapping[str, Any]) -> list[str]:
    missing: list[str] = []
    for name in REQUIRED_IDENTITY_FIELDS:
        value = row.get(name)
        if value is None or (isinstance(value, str) and not value.strip()):
            missing.append(name)
    return missing


def plan_stale_worker_retirement(
    rows: Iterable[Mapping[str, Any]],
    *,
    now: datetime | None = None,
    pid_alive_probe: Callable[[int | None], bool | None] | None = None,
    canonical_repository: str | None = None,
    protected_worker_identities: Iterable[str] | None = None,
    protected_session_ids: Iterable[str] | None = None,
    protected_pids: Iterable[Any] | None = None,
) -> dict[str, Any]:
    """Decide, without mutating anything, which registrations may be retired.

    Every row lands in exactly one of ``candidates`` (eligible) or
    ``preserved`` (with the reason code that stopped it), so the output
    explains the whole registry rather than only the interesting part.
    """
    all_rows = [dict(row) for row in rows]
    protected_ids = {str(w) for w in (protected_worker_identities or []) if w}
    protected_sessions = {str(s) for s in (protected_session_ids or []) if s}
    protected_pid_set = {_canon(p) for p in (protected_pids or []) if p is not None}

    snapshots: dict[int, dict[str, Any]] = {}
    pid_alive_by_index: dict[int, bool | None] = {}
    for index, row in enumerate(all_rows):
        pid_alive = _probe_pid(row.get("pid"), pid_alive_probe)
        pid_alive_by_index[index] = pid_alive
        snapshots[index] = fleet.build_worker_snapshot_row(
            row,
            now=now,
            pid_alive_probe=(lambda _pid, _value=pid_alive: _value),
            canonical_repository=canonical_repository,
        )

    # Identity evidence owned by a worker that is live, or whose liveness could
    # not be established, is ambiguous: anything sharing it is preserved.
    ambiguous_keys: set[tuple[str, str]] = set()
    identity_counts: dict[str, int] = {}
    for index, row in enumerate(all_rows):
        identity = _canon(row.get("worker_identity"))
        if identity:
            identity_counts[identity] = identity_counts.get(identity, 0) + 1
        snapshot_row = snapshots[index]
        liveness = snapshot_row.get("liveness") or {}
        unresolved = (
            pid_alive_by_index[index] is None
            or liveness.get("heartbeat_fresh") is None
        )
        if snapshot_row.get("live") or (
            str(row.get("status") or "") == mwi.STATUS_ACTIVE and unresolved
        ):
            ambiguous_keys.update(_conflict_keys(row, snapshot_row))

    candidates: list[dict[str, Any]] = []
    candidate_rows: list[Mapping[str, Any]] = []
    preserved: list[dict[str, Any]] = []

    for index, row in enumerate(all_rows):
        snapshot_row = snapshots[index]
        pid_alive = pid_alive_by_index[index]
        liveness = snapshot_row.get("liveness") or {}
        evidence = _evidence(row, snapshot_row, pid_alive)
        blocked: tuple[str, str] | None = None

        identity = _canon(row.get("worker_identity"))
        missing = _missing_identity_fields(row)
        shared = sorted(
            f"{name}={value}"
            for name, value in _conflict_keys(row, snapshot_row)
            if (name, value) in ambiguous_keys
        )
        binding = (row.get("repository_binding") or "").strip()
        protected_hits: list[str] = []
        if identity and identity in protected_ids:
            protected_hits.append(f"worker_identity={identity}")
        if _canon(row.get("session_id")) in protected_sessions:
            protected_hits.append(f"session_id={_canon(row.get('session_id'))}")
        if _canon(row.get("pid")) in protected_pid_set:
            protected_hits.append(f"pid={_canon(row.get('pid'))}")

        if identity and identity_counts.get(identity, 0) > 1:
            blocked = (
                REASON_CONFLICTING_IDENTITY,
                f"worker identity {identity!r} appears on more than one registry row",
            )
        elif str(row.get("status") or "") != mwi.STATUS_ACTIVE:
            blocked = (
                REASON_ALREADY_TERMINAL,
                f"registration status is {row.get('status')!r}; nothing to retire",
            )
        elif missing:
            blocked = (
                REASON_INCOMPLETE_IDENTITY,
                "registry row is missing field(s) the retirement conjunction "
                f"reads: {missing}",
            )
        elif mwi._parse_ts(row.get("last_heartbeat_at")) is None:
            blocked = (
                REASON_UNPARSABLE_HEARTBEAT,
                "last_heartbeat_at is not a parsable UTC stamp; liveness is unknown",
            )
        elif snapshot_row.get("live"):
            blocked = (REASON_WORKER_LIVE, "worker is live and must not be retired")
        elif pid_alive is None:
            blocked = (
                REASON_PID_UNKNOWN,
                f"pid {row.get('pid')!r} could not be probed; liveness is unproven",
            )
        elif pid_alive:
            blocked = (
                REASON_PID_ALIVE,
                f"recorded pid {row.get('pid')!r} is still running",
            )
        elif liveness.get("heartbeat_fresh") is not False:
            blocked = (
                REASON_HEARTBEAT_FRESH,
                "heartbeat has not expired under the canonical TTL policy",
            )
        elif snapshot_row.get("ownership_state") != "stale":
            blocked = (
                REASON_AMBIGUOUS_OWNERSHIP,
                "ownership_state is "
                f"{snapshot_row.get('ownership_state')!r}, not 'stale'",
            )
        elif not binding or snapshot_row.get("foreign_repository"):
            blocked = (
                REASON_FOREIGN_REPOSITORY,
                "repository binding is absent or does not match the canonical "
                "repository; retirement scope is ambiguous",
            )
        elif shared:
            blocked = (
                REASON_CONFLICTING_IDENTITY,
                "identity evidence is shared with a live or unprobeable worker: "
                f"{shared}",
            )
        elif protected_hits:
            blocked = (
                REASON_PROTECTED_OWNER,
                "worker still owns active workflow state requiring separate "
                f"reconciliation: {protected_hits}",
            )

        if blocked is not None:
            preserved.append(
                {
                    "worker_identity": row.get("worker_identity"),
                    "reason_code": blocked[0],
                    "detail": blocked[1],
                    "evidence": evidence,
                }
            )
            continue

        candidates.append(
            {
                "worker_identity": row.get("worker_identity"),
                "reason_code": REASON_ELIGIBLE,
                "detail": (
                    "dead pid, expired heartbeat, stale ownership, unambiguous "
                    "identity, canonical repository binding, no active workflow "
                    "ownership"
                ),
                "evidence": evidence,
            }
        )
        candidate_rows.append(row)

    counts: dict[str, int] = {}
    for entry in preserved:
        counts[entry["reason_code"]] = counts.get(entry["reason_code"], 0) + 1

    return {
        "success": True,
        "read_only": True,
        "mutation_performed": False,
        "outcome": OUTCOME_PLANNED,
        "registry_fingerprint": registry_fingerprint(all_rows),
        "candidate_fingerprint": candidate_fingerprint(candidate_rows),
        "assessed_count": len(all_rows),
        "candidate_count": len(candidates),
        "preserved_count": len(preserved),
        "candidates": candidates,
        "candidate_worker_identities": [c["worker_identity"] for c in candidates],
        "preserved": preserved,
        "preserved_reason_counts": counts,
        "protected_inputs": {
            "worker_identities": sorted(protected_ids),
            "session_ids": sorted(protected_sessions),
            "pids": sorted(protected_pid_set),
        },
        "canonical_repository": canonical_repository,
    }


def summarize_plan(plan: Mapping[str, Any]) -> dict[str, Any]:
    """Compact, log-safe view of a plan or apply result."""
    return {
        "outcome": plan.get("outcome"),
        "registry_fingerprint": plan.get("registry_fingerprint"),
        "candidate_fingerprint": plan.get("candidate_fingerprint"),
        "assessed_count": plan.get("assessed_count"),
        "candidate_count": plan.get("candidate_count"),
        "retired_count": plan.get("retired_count"),
        "preserved_count": plan.get("preserved_count"),
        "mutation_performed": plan.get("mutation_performed"),
    }