multi-agent-mux/.agents/skills/lib_py/delegation/engine.py

"""The single state-transition authority.

Every job-status write goes through `transition`. Before phase 2a there were two
write paths: `update_job_status` merged fields unconditionally, and
`pick_pending` hand-rolled a correct pending->running swap plus its own audit
mirror. Neither was wrong on its own; having two was the risk, because a fix or
a rule added to one silently misses the other.

It lives in `lib_py.delegation`, not `lib_py.loop`, because it writes the
registry and the registry writers are here. A frozen loop can bind
`lib_py.delegation` to a different tree than the parent package, so delegation
importing a sibling would mix trees -- the rule this package states in its own
__init__, and which caught this during implementation.

Lock discipline (the reason this is a module and not a helper):

    with registry_lock:          decide, write the record once, QUEUE notifications
    after the lock:              flush notifications, publish, anything slow

Nothing that can block on another process or the network may run while the lock
is held. `pick_pending` already mirrored its audit entries after releasing the
lock; this generalises that shape to every transition.

`MAM_ENGINE_CAS=0` restores the previous unconditional-merge behaviour without a
redeploy. The legacy bodies are preserved by their callers for that reason.
"""
from __future__ import annotations

import logging
import os
from typing import Any, Dict, List, Optional

from lib_py.delegation import mqtt_common

logger = logging.getLogger("loop.engine")

TERMINAL_STATUSES = ("completed", "error", "cancelled")
NON_TERMINAL_STATUSES = ("pending", "running")

# request -> the status it asks for. `merge`/`progress` ask for nothing.
REQUEST_TO_STATUS = {
    "claim": "running",
    "started": "running",
    "progress": None,
    "permission_required": None,
    "completed": "completed",
    "error": "error",
    "cancelled": "cancelled",
    "deadline_timeout": "error",
    "merge": None,
}

# A caller that merges `status=<value>` is making the request that owns that
# value; mapping it here is what stops a raw field write from sidestepping the
# CAS table. `pending` has no owning request (nothing transitions *to* pending),
# so it stays an ordinary metadata merge, exactly as before.
STATUS_TO_REQUEST = {
    "running": "started",
    "completed": "completed",
    "error": "error",
    "cancelled": "cancelled",
}


def cas_enabled() -> bool:
    """The rollback switch. Read per call so a test or operator can flip it."""
    return os.environ.get("MAM_ENGINE_CAS", "1") != "0"


class TransitionResult(object):
    __slots__ = ("committed", "from_status", "to_status", "reason", "record", "notifications")

    def __init__(self, committed, from_status, to_status, reason, record, notifications):
        self.committed = committed
        self.from_status = from_status
        self.to_status = to_status
        self.reason = reason
        self.record = record
        self.notifications = notifications

    def __repr__(self):  # pragma: no cover - debugging aid
        return "TransitionResult(committed={!r}, {!r}->{!r}, {!r})".format(
            self.committed, self.from_status, self.to_status, self.reason)


class _Decision(object):
    __slots__ = ("commit", "to_status", "reason", "conflict")

    def __init__(self, commit, to_status, reason, conflict=False):
        self.commit = commit
        self.to_status = to_status
        self.reason = reason
        self.conflict = conflict


def _decide(record: Dict[str, Any], request: str, expected_status, expected_attempt) -> _Decision:
    """Pure CAS decision. No I/O, so it is safe to call under the lock.

    Implements the target table in docs/contracts/issue-5/state-transitions.md §3.
    """
    current = record.get("status")

    # Stale attempt or generation loses before anything else is considered.
    if expected_attempt is not None:
        actual = record.get("attempt_id", record.get("iteration"))
        if actual is not None and actual != expected_attempt:
            return _Decision(False, current,
                             "stale attempt: expected {!r}, record has {!r}".format(
                                 expected_attempt, actual))

    if expected_status is not None and current != expected_status:
        return _Decision(False, current,
                         "precondition failed: expected status {!r}, found {!r}".format(
                             expected_status, current))

    requested = REQUEST_TO_STATUS.get(request, None)

    if current in TERMINAL_STATUSES:
        # Idempotent re-receipt of the same terminal is a no-op ack.
        if requested == current:
            return _Decision(False, current, "idempotent: already {}".format(current))
        if requested is None:
            return _Decision(False, current,
                             "late {} after terminal {}; refusing".format(request, current),
                             conflict=True)
        return _Decision(False, current,
                         "conflicting {} after terminal {}; refusing".format(requested, current),
                         conflict=True)

    if requested is None:
        # progress / permission_required / metadata-only merge: the engine does
        # not set a status, so any status in `fields` is applied as-is.
        return _Decision(True, None, "metadata updated ({})".format(request))

    return _Decision(True, requested, "{} -> {}".format(current, requested))


def _notification(job_id, kind, **fields) -> Dict[str, Any]:
    payload = {"job_id": job_id, "kind": kind}
    payload.update(fields)
    return payload


def transition(job_id, request="merge", *, registry_dir=None, expected_status=None,
               expected_attempt=None, fields=None, source="unknown") -> TransitionResult:
    """Decide and commit one state transition. The only writer of job status.

    Returns a result whose `notifications` have ALREADY been flushed; the field
    is kept so callers and tests can see what the transition emitted.
    """
    registry_dir = registry_dir or mqtt_common.DEFAULT_REGISTRY_DIR
    fields = dict(fields or {})

    # ---- short transaction: decide and write, never do I/O that can block ----
    with mqtt_common.registry_lock(registry_dir):
        record = mqtt_common.load_job(job_id, registry_dir)
        from_status = record.get("status")
        decision = _decide(record, request, expected_status, expected_attempt)

        notifications: List[Dict[str, Any]] = []
        if decision.commit:
            record.update(fields)
            if decision.to_status is not None:
                record["status"] = decision.to_status
            record["updated_at"] = mqtt_common._utcnow()
            # A monotonic counter so a late flush can still be ordered correctly
            # once the audit mirror happens outside the lock.
            record["transition_seq"] = int(record.get("transition_seq", 0)) + 1
            mqtt_common._atomic_write_record(job_id, registry_dir, record)

            to_status = record.get("status")
            if to_status != from_status:
                notifications.append(_notification(
                    job_id, "status_changed", **{
                        "from": from_status, "to": to_status,
                        "by": source, "timestamp": record["updated_at"],
                        "transition_seq": record["transition_seq"],
                    }))
        else:
            to_status = from_status
            if decision.conflict:
                notifications.append(_notification(
                    job_id, "conflict_refused",
                    attempted=request, current=from_status,
                    reason=decision.reason, by=source,
                    timestamp=mqtt_common._utcnow()))
    # ---- lock released; everything below may block ----

    _flush(notifications, registry_dir)
    if not decision.commit and decision.conflict:
        logger.warning("refused %s for job %s: %s", request, job_id, decision.reason)

    return TransitionResult(decision.commit, from_status, to_status,
                            decision.reason, record, notifications)


def _flush(notifications, registry_dir) -> None:
    """Mirror queued notifications into the audit log, outside the lock.

    Best effort by design: the registry record is already durable, and an audit
    write must never be able to fail a transition that has committed.
    """
    for note in notifications:
        job_id = note["job_id"]
        kind = note["kind"]
        try:
            if kind == "status_changed":
                mqtt_common.update_logged_status(
                    job_id, note["to"], updated_at=note.get("timestamp"))
                mqtt_common.append_event(job_id, {
                    "event": "status_changed",
                    "from": note.get("from"),
                    "to": note.get("to"),
                    "by": note.get("by"),
                    "timestamp": note.get("timestamp"),
                    "transition_seq": note.get("transition_seq"),
                })
            elif kind == "conflict_refused":
                mqtt_common.append_event(job_id, {
                    "event": "conflict_refused",
                    "attempted": note.get("attempted"),
                    "current": note.get("current"),
                    "reason": note.get("reason"),
                    "by": note.get("by"),
                    "timestamp": note.get("timestamp"),
                })
        except Exception as exc:  # pragma: no cover - best effort
            logger.warning("audit mirror failed for %s: %s", job_id, exc)