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
import uuid
from typing import Any, Dict, List, Optional

from lib_py.delegation import mqtt_common, notifier

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,
    # Operator-authority requests. The wrapper drives multi-iteration delegation
    # (`submit --type loop|discuss`) by re-dispatching a finished job to the next
    # session, and by finalising it as an error when the iteration budget is
    # exhausted. Both legitimately move a terminal record, which is why they are
    # separate requests gated on authority rather than an exception in the table.
    "redispatch": "pending",
    "finalize": "error",
    # Bookkeeping: records that an owed notification was delivered. It cannot
    # change status, and it is *about* a terminal, so refusing it in a terminal
    # state would make an owed notification impossible to ever settle.
    "notification_settled": None,
}

# Requests that carry no status at all and are therefore safe in any state,
# including a terminal one. Keep this list short and obviously status-free.
BOOKKEEPING_REQUESTS = ("notification_settled",)

# The only fields a bookkeeping request may write.
BOOKKEEPING_FIELDS = ("pending_notification",)

# Distinguishes "no expectation" from "expect the field to be None".
_UNSET = object()

# An opaque identity for one dispatch attempt, minted when an attempt begins and
# stored in the record itself. It lives here, not only in the shadow store,
# because a precondition can only be atomic with the commit if the transaction
# can read it under its own lock. Comparing the shadow store before the call
# leaves a window in which the job is completed, re-dispatched and claimed again,
# after which stale evidence satisfies a status check and completes a new
# attempt. An opaque token is used rather than the iteration counter because
# iteration already means the loop round, and conflating the two produces a guard
# that is either vacuous or wrong.
DISPATCH_TOKEN_FIELD = "dispatch_token"

# The requests that begin a new attempt, and therefore mint a new token.
ATTEMPT_STARTING_REQUESTS = ("claim", "redispatch")

# Only these sources may move a record out of a terminal state. An agent can
# never walk back its own terminal; that is the whole point of the CAS table.
OPERATOR_SOURCES = ("registry_cli", "operator", "delegate_wrapper")
OPERATOR_REQUESTS = ("redispatch", "finalize")

# 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. Every status has an owning request, `pending` included: the wrapper
# resets a finished job to pending to hand it to the next iteration.
STATUS_TO_REQUEST = {
    "pending": "redispatch",
    "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", "idempotent")

    def __init__(self, committed, from_status, to_status, reason, record, notifications,
                 idempotent=False):
        self.committed = committed
        # Not every uncommitted transition is a refusal. Re-receiving a terminal
        # that already holds is a no-op ACK in the contract table, and a caller
        # that treats it as failure turns a success path into an error.
        self.idempotent = idempotent
        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", "idempotent")

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


def _decide(record: Dict[str, Any], request: str, expected_status, expected_attempt,
            source="unknown", fields=None, expected_notification=_UNSET,
            expected_dispatch_token=None) -> _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")
    fields = fields or {}

    # An unrecognised request must never fall through to the permissive branch.
    if request not in REQUEST_TO_STATUS:
        return _Decision(False, current, "unknown request: {!r}".format(request))

    # A status may only be set by the request that owns it. Letting one ride in
    # as an ordinary field would let `progress` commit a terminal, which is the
    # decision table being bypassed from inside its own boundary.
    if "status" in fields:
        return _Decision(False, current,
                         "status must be requested, not merged as a field "
                         "(got status={!r} on request {!r})".format(fields["status"], request))

    # A supplied precondition must be *proved*, not merely un-contradicted. If a
    # caller names an expected attempt and the record cannot show one, the
    # compare-and-swap has nothing to compare and must fail closed.
    if expected_attempt is not None:
        actual = record.get("attempt_id", record.get("iteration"))
        if actual is None:
            return _Decision(False, current,
                             "expected attempt {!r} but the record carries none".format(
                                 expected_attempt))
        if actual != expected_attempt:
            return _Decision(False, current,
                             "stale attempt: expected {!r}, record has {!r}".format(
                                 expected_attempt, actual))

    requested = REQUEST_TO_STATUS.get(request, None)

    if expected_dispatch_token is not None:
        actual_token = record.get(DISPATCH_TOKEN_FIELD)
        if actual_token != expected_dispatch_token:
            return _Decision(False, current,
                             "dispatch token changed: the attempt this decision was "
                             "made for is no longer the current one")
    elif expected_attempt is None and source not in OPERATOR_SOURCES and requested in TERMINAL_STATUSES:
        actual_token = record.get(DISPATCH_TOKEN_FIELD)
        rec_attempt = record.get("attempt_id", record.get("iteration", 1))
        if actual_token is not None and rec_attempt > 1:
            return _Decision(False, current,
                             "dispatch token required: redispatched job carries a dispatch token "
                             "but none was provided")

    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))

    # Authority first, so a refusal says *why* rather than blaming the state.
    if request in OPERATOR_REQUESTS and source not in OPERATOR_SOURCES:
        return _Decision(False, current,
                         "{} requires operator authority (source={!r})".format(request, source))

    if request in BOOKKEEPING_REQUESTS:
        # Bookkeeping is allowed in a terminal state precisely because it cannot
        # change anything of consequence. That has to be enforced here, at the
        # boundary, and not merely by what today's one caller happens to pass:
        # otherwise a status-free request becomes an unrestricted merge on a
        # record nothing else is allowed to touch.
        extra = sorted(set(fields) - set(BOOKKEEPING_FIELDS))
        if extra:
            return _Decision(False, current,
                             "bookkeeping may not merge {}".format(extra),
                             conflict=True)
        expected = expected_notification
        if expected is not _UNSET:
            current_intent = record.get("pending_notification")
            if current_intent != expected:
                # A newer intent replaced the one being settled while it was in
                # flight. Clearing it now would drop a notification nobody sent.
                return _Decision(False, current,
                                 "notification intent changed while it was being sent")
        return _Decision(True, None, "bookkeeping ({})".format(request))

    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),
                             idempotent=True)

        operator = source in OPERATOR_SOURCES
        if operator and request == "redispatch":
            # Re-dispatching a finished job to the next iteration. The attempt
            # must move forward, so a replayed redispatch cannot resurrect it.
            new_attempt = fields.get("attempt_id")
            if new_attempt is not None:
                old_attempt = record.get("attempt_id", record.get("iteration"))
                if old_attempt is not None and int(new_attempt) <= int(old_attempt):
                    return _Decision(False, current,
                                     "redispatch attempt_id {!r} does not advance past {!r}".format(
                                         new_attempt, old_attempt), conflict=True)
                return _Decision(True, "pending", "{} -> pending (redispatch)".format(current))

            new_iteration = fields.get("iteration")
            old_iteration = record.get("iteration")
            if new_iteration is None:
                return _Decision(False, current,
                                 "redispatch requires an iteration", conflict=True)
            if old_iteration is not None and int(new_iteration) <= int(old_iteration):
                return _Decision(False, current,
                                 "redispatch iteration {!r} does not advance past {!r}".format(
                                     new_iteration, old_iteration), conflict=True)
            return _Decision(True, "pending", "{} -> pending (redispatch)".format(current))

        if operator and request in ("finalize", "error"):
            return _Decision(True, "error", "{} -> error (finalize)".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. No status is set
        # here: a status-bearing field was already refused above.
        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",
               notify=None, expected_notification=_UNSET,
               expected_dispatch_token=None) -> 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,
                           source=source, fields=fields,
                           expected_notification=expected_notification,
                           expected_dispatch_token=expected_dispatch_token)

        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) or 0) + 1
            if request in ATTEMPT_STARTING_REQUESTS:
                # A new attempt begins here, so anything decided about the
                # previous one stops being committable from this moment on.
                record[DISPATCH_TOKEN_FIELD] = uuid.uuid4().hex
                if "attempt_id" in fields:
                    record["attempt_id"] = fields["attempt_id"]
                elif "iteration" in fields:
                    record["attempt_id"] = fields["iteration"]
                else:
                    record.setdefault("attempt_id", record.get("iteration", 1))
            elif record.get(DISPATCH_TOKEN_FIELD) is None:
                record[DISPATCH_TOKEN_FIELD] = uuid.uuid4().hex
                record.setdefault("attempt_id", record.get("iteration", 1))
            if notify is not None:
                # One atomic commit records both the new state and the fact that
                # the world has not been told yet. If the process dies before the
                # publish, the record still says a notification is owed.
                record[notifier.FIELD] = notifier.intent_for(
                    notify, record["transition_seq"], record["updated_at"],
                    detail=str(fields.get("detail", "")))
            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,
                            idempotent=decision.idempotent)


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"),
                    transition_seq=note.get("transition_seq"))
                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)