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

"""Issue #17 item 4: the delegate-job wrapper's decisions.

The bash wrapper (multi-agent-mux-delegate-job) spawns, waits, kills and
injects; everything it used to decide in bash is here, and it branches only on
what a subcommand prints or returns:

  sub-status <rc> --job J --registry-dir R   subscriber rc -> job status (was _map_sub_rc)
  subscriber-up <rc> [--print]               exit 0: a subscriber that ended before
                                             SUBSCRIBED still lets the job proceed
  exit-code <status>                         the wrapper's exit code (was _rc_for, with
                                             MAM_DELEGATE_FAIL_CLOSED=0 -> 0)
  wait-rc <rc>                               `wait`'s rc (either clock is a timeout, 2)
  stall-watch <session> <agent> [--print]    exit 0: start the stall watcher (Issue #30)

The two yes/no decisions answer by exit code (0 yes, 1 no); with --print they
print `yes`/`no` and exit 0, so a failure to decide (any other exit code, or
no answer) cannot pass for "no". The wrapper uses --print.
  finalize-begin / finalize-end ...          finalize_unfinished, split around the one
                                             herdr diagnostic probe bash still runs

Issue #50 (the registry, not herdr's agent_status, says whether a session is
free, and create --onboard waits for the onboarding to end):
  wait-terminal --job J --registry-dir R --timeout T [--watch-pid P]
                                             prints the job's terminal status, or
                                             `timeout` / `missing`
  onboard-exit <status>                      create's exit code for it: 0, 44, 45
  open-jobs --job J --agent-session S --registry-dir R
                                             diagnostic: the live open jobs of S
  wait-open-jobs ... --wait-s W --start-time T0 --timeout T [--watch-pid P]
                                             prints `clear` or `busy <ids>`

finalize-end answers on fd 3 (stdout when fd 3 is closed): its stdout carries
registry.py's close output to the wrapper's stdout, as before. It runs as a
child of the wrapper and must not outlive it (review R1): with --watch-pid it
checks that the wrapper is still alive throughout both waits and before each
side effect, and leaves without closing or publishing once it is gone - as the
same loops in bash ended with the wrapper.

This module imports only the standard library and runs on Python 3.9 (macOS
/usr/bin/python3). The wrapper runs it as `python -m lib_py.delegation.wrapper`,
which imports the package first: lib_py/delegation/__init__.py loads the
service and with it paho-mqtt, so the wrapper's $PY needs paho - as it always
did for the subscriber and publish_event.py. registry.py and publish_event.py
run as the processes the wrapper ran, with its interpreter.
"""
from __future__ import annotations

import argparse
import calendar
import json
import os
import subprocess
import sys
import time
from typing import List, Optional

# job_subscriber.py's RC_* (a test keeps them equal; importing it would load paho).
SUB_RC_COMPLETED, SUB_RC_ERROR, SUB_RC_NO_EVENT, SUB_RC_BROKER_UNAVAILABLE, SUB_RC_WALL_CLOCK = 0, 1, 2, 3, 4

TERMINAL = ("completed", "error", "cancelled")
EXIT_CODES = {"completed": 0, "error": 1, "timeout": 2, "broker_unavailable": 3, "cancelled": 5}
# Issue #50: create --onboard's exit code for how the onboarding job ended.
# 44 = the agent ended it with error (or it was cancelled), 45 = it did not end
# within MAM_ONBOARD_WAIT_S (or its record is gone). The session is kept either way.
ONBOARD_EXIT_CODES = {"completed": 0, "error": 44, "cancelled": 44, "timeout": 45, "missing": 45}
DEFAULT_ACK_TIMEOUT_SEC, DEFAULT_IDLE_TIMEOUT_SEC = 180, 120  # registry.py's defaults, as a last resort
DISPATCH_SIGNAL_FIELD = "dispatch_wait_at"
DISPATCH_SIGNAL_EVERY = 10


def _now() -> int:
    return int(time.time())  # bash's $(date +%s)


def disk_status(registry_dir: str, job_id: str) -> str:
    """The record's status, "" when it is missing or unreadable."""
    try:
        with open(os.path.join(registry_dir, f"{job_id}.json")) as fh:
            return str(json.load(fh).get("status", ""))
    except Exception:
        return ""


def sub_status(rc: int, registry_dir: str, job_id: str) -> str:
    if rc == SUB_RC_COMPLETED:
        return "completed"
    if rc == SUB_RC_ERROR:
        return "error"
    if rc == SUB_RC_BROKER_UNAVAILABLE:
        disk = disk_status(registry_dir, job_id)
        return disk if disk in TERMINAL else "broker_unavailable"
    return "timeout"  # SUB_RC_NO_EVENT, SUB_RC_WALL_CLOCK (and any other code)


def subscriber_up(rc: int) -> bool:
    return rc in (SUB_RC_COMPLETED, SUB_RC_BROKER_UNAVAILABLE)


def fail_closed() -> bool:
    return (os.environ.get("MAM_DELEGATE_FAIL_CLOSED") or "1") != "0"


def exit_code(status: str) -> int:
    return EXIT_CODES.get(status, 1) if fail_closed() else 0


def wait_rc(rc: int) -> int:
    return SUB_RC_NO_EVENT if rc == SUB_RC_WALL_CLOCK else rc


def _bare(session: str) -> str:
    return session[len("herdr:"):] if session.startswith("herdr:") else session


def stall_watch(session: str, agent: str) -> bool:
    """A remote or human job has no local pane to watch."""
    if (os.environ.get("MAM_STALL_WATCH") or "1") == "0":
        return False
    if session.startswith("remote:"):
        print("stall watch: remote job, no local pane - stall detection off", file=sys.stderr)
        return False
    if agent == "human" or _bare(session) == "human":
        print("stall watch: human job, no local pane - stall detection off", file=sys.stderr)
        return False
    return True


def recent_disk_heartbeat(registry_dir: str, job_id: str, window: float) -> bool:
    """ISSUE-14: publish_event appends every event to the job's events.log
    before its wire publish, so a heartbeat the subscriber never saw is
    visible here; true when the last one is at most `window` s old."""
    try:
        last = ""
        with open(os.path.join(registry_dir, f"{job_id}.events.log")) as fh:
            for line in fh:
                if line.strip():
                    last = line
        seen = calendar.timegm(time.strptime(json.loads(last)["timestamp"], "%Y-%m-%dT%H:%M:%SZ"))
    except Exception:
        return False
    return time.time() - seen <= window


class OwnerGone(Exception):
    """The wrapper that ran finalize-end is gone."""


def _alive(pid: int) -> bool:
    """A running process. An exited one its caller has not reaped yet is a
    zombie that still answers kill(0) (review R1, round 2): `ps -o stat=`
    (POSIX, no /proc) tells it apart, as tests/conftest.py::_pid_alive does."""
    try:
        os.kill(pid, 0)
    except ProcessLookupError:
        return False
    except PermissionError:
        return True
    try:
        state = subprocess.run(["ps", "-o", "stat=", "-p", str(pid)], capture_output=True, text=True,
                               timeout=10).stdout.strip()
    except (OSError, subprocess.SubprocessError):
        return True  # no ps answer: kill(0) said it exists
    return bool(state) and not state.startswith("Z")


def _check_owner(a: argparse.Namespace) -> None:
    if a.watch_pid and not _alive(a.watch_pid):
        raise OwnerGone(a.watch_pid)


def _sleep(a: argparse.Namespace, seconds: float) -> None:
    """time.sleep that looks at the wrapper every 0.5 s. The span is measured
    on the monotonic clock, as bash's `sleep N` is: a wall clock stepped back
    during the wait must not stretch it (review R1, codex). Deadlines stay on
    the wall clock (_now)."""
    end = time.monotonic() + seconds
    while True:
        _check_owner(a)
        left = end - time.monotonic()
        if left <= 0:
            return
        time.sleep(min(left, 0.5))


def _tail(path: str, n: int) -> List[str]:
    with open(path, errors="replace") as fh:
        return fh.read().splitlines()[-n:]


# ---------------------------------------------------------------------------
# Issue #50 A: create --onboard waits for the onboarding job to end
# ---------------------------------------------------------------------------

def onboard_exit(status: str) -> int:
    return ONBOARD_EXIT_CODES.get(status, 45)


def wait_terminal(a: argparse.Namespace) -> str:
    """The job's terminal status once it has one; `timeout` after --timeout s,
    `missing` when the record is gone. Reads the disk every second."""
    deadline = _now() + a.timeout
    record = os.path.join(a.registry_dir, f"{a.job}.json")
    while True:
        disk = disk_status(a.registry_dir, a.job)
        if disk in TERMINAL:
            return disk
        if not os.path.isfile(record):
            return "missing"
        if _now() >= deadline:
            return "timeout"
        _sleep(a, 1)


# ---------------------------------------------------------------------------
# Issue #50 C: the registry says whether a session is free
# ---------------------------------------------------------------------------

def _parse_ts(text: str) -> Optional[int]:
    try:
        return calendar.timegm(time.strptime(str(text), "%Y-%m-%dT%H:%M:%SZ"))
    except (TypeError, ValueError):
        return None


def _load_record(registry_dir: str, job_id: str) -> Optional[dict]:
    try:
        with open(os.path.join(registry_dir, f"{job_id}.json")) as fh:
            rec = json.load(fh)
    except Exception:
        return None
    return rec if isinstance(rec, dict) else None


def _int_field(rec: dict, key: str, default: int) -> int:
    try:
        return max(1, int(rec.get(key) or default))
    except (TypeError, ValueError):
        return default


def open_jobs(registry_dir: str, session: str, job_id: str, now: Optional[int] = None):
    """(live, stale) among the other non-terminal records of `session`.

    live:  [(created_at, job_id, status)] - open jobs the Issue #18 contract
           still calls alive: `running` with an event in events.log within
           idle_timeout_sec; `pending` with a dispatch signal (the waiter's
           `dispatch_wait_at`, else created_at) within ack_timeout_sec.
    stale: [(job_id, reason)] - non-terminal records nothing has touched for
           longer than that: a wrapper killed before or while dispatching, a
           zombie's unclosed record. They are reported, never waited for, so
           a dead record cannot block a session forever."""
    now = _now() if now is None else now
    sess = _bare(session)
    live, stale = [], []
    try:
        names = sorted(os.listdir(registry_dir))
    except OSError:
        return live, stale
    for name in names:
        if not name.endswith(".json"):
            continue
        other = name[:-len(".json")]
        if other == job_id:
            continue
        rec = _load_record(registry_dir, other)
        if rec is None or _bare(str(rec.get("agent_session") or "")) != sess:
            continue
        status = rec.get("status")
        if status == "running":
            idle = _int_field(rec, "idle_timeout_sec", DEFAULT_IDLE_TIMEOUT_SEC)
            if recent_disk_heartbeat(registry_dir, other, idle):
                live.append((str(rec.get("created_at") or ""), other, status))
            else:
                stale.append((other, f"running, no event for {idle}s"))
        elif status == "pending":
            ack = _int_field(rec, "ack_timeout_sec", DEFAULT_ACK_TIMEOUT_SEC)
            seen = _parse_ts(rec.get(DISPATCH_SIGNAL_FIELD) or rec.get("created_at"))
            if seen is not None and now - seen <= ack:
                live.append((str(rec.get("created_at") or ""), other, status))
            else:
                stale.append((other, f"pending, no dispatch signal for {ack}s"))
    return live, stale


def _signal_dispatch_wait(registry_dir: str, job_id: str) -> None:
    """The waiter's liveness: a merge of dispatch_wait_at into its own pending
    record through the registry's one transition boundary (status untouched,
    no event in events.log - that log is the agent's)."""
    from lib_py.delegation import mqtt_common  # loads paho, as the package import did already
    try:
        mqtt_common.update_job_status(job_id, registry_dir, _source="dispatch_wait",
                                      **{DISPATCH_SIGNAL_FIELD: mqtt_common._utcnow()})
    except Exception as exc:  # the record is gone or unwritable: keep waiting, say so once per tick
        print(f"dispatch wait: could not signal on {job_id}: {exc}", file=sys.stderr)


def wait_open_jobs(a: argparse.Namespace) -> str:
    """`clear` when no live open job of the session remains and this job is
    the earliest live pending one (FIFO by (created_at, job_id)); `busy <ids>`
    when the bound runs out first. The bound is min(--wait-s, this job's
    remaining wall clock). Every DISPATCH_SIGNAL_EVERY s the waiter signals on
    its own record so later waiters see it alive; a waiter whose wrapper is
    gone stops (OwnerGone) and its signal goes stale."""
    me = _load_record(a.registry_dir, a.job) or {}
    mine = (str(me.get("created_at") or ""), a.job)
    remaining = (a.start_time if a.start_time is not None else _now()) + a.timeout - _now()
    deadline = _now() + max(0, min(a.wait_s, remaining))
    reported = set()
    last_signal = None
    while True:
        now = _now()
        live, stale = open_jobs(a.registry_dir, a.agent_session, a.job, now)
        for other, reason in stale:
            if other not in reported:
                reported.add(other)
                print(f"stale open job {other} ignored ({reason})", file=sys.stderr)
        if not any(st == "running" for _c, _j, st in live):
            queue = [(c, j) for c, j, st in live if st == "pending"] + [mine]
            if min(queue) == mine:
                return "clear"
        if now >= deadline:
            return "busy " + " ".join(j for _c, j, _st in live)
        if last_signal is None or now - last_signal >= DISPATCH_SIGNAL_EVERY:
            _signal_dispatch_wait(a.registry_dir, a.job)
            last_signal = now
        _sleep(a, 1)


def finalize_begin(a: argparse.Namespace) -> str:
    """`keep` (FINAL_JOB_STATUS untouched), `done <final>`, or
    `go|probe <final> <evidence_until> <deadline>`."""
    if not fail_closed() or a.status == "completed":
        return "keep"
    stall_out = os.path.join(a.registry_dir, f"{a.job}.stall.out")
    if os.path.isfile(stall_out) and os.path.getsize(stall_out) > 0:
        # Issue #30: what the stall watcher saw goes with an unfinished job's close.
        print(f"stall watch log for {a.job} (last 3 lines):", file=sys.stderr)
        try:
            for line in _tail(stall_out, 3):
                print(line, file=sys.stderr)
        except OSError:
            pass
    disk = disk_status(a.registry_dir, a.job)
    if disk == "completed":
        return "done completed"
    final = disk if disk in TERMINAL else a.status
    now = _now()
    deadline = (a.start_time if a.start_time is not None else now) + a.timeout
    evidence_until = now + a.idle_timeout
    sess = _bare(a.session)
    # herdr status and the pane are a diagnostic only (ISSUE-14 / ISSUE-18):
    # bash runs that probe, bounded by the deadline, when asked to.
    probe = (a.status == "timeout" and a.agent != "human" and sess != "human" and sess != ""
             and not a.session.startswith("remote:"))
    return f"{'probe' if probe else 'go'} {final} {evidence_until} {deadline}"


def _run(a: argparse.Namespace, script: str, args: List[str], merge: bool) -> int:
    sys.stdout.flush()
    cmd = [a.python, os.path.join(a.scripts_dir, script), "--registry-dir", a.registry_dir, *args]
    return subprocess.call(cmd, stderr=subprocess.STDOUT if merge else subprocess.DEVNULL)


def finalize_end(a: argparse.Namespace) -> str:
    """The rest of finalize_unfinished; returns FINAL_JOB_STATUS."""
    job, reg, status, final = a.job, a.registry_dir, a.status, a.final
    clock = ""
    deadline = (a.start_time if a.start_time is not None else _now()) + a.timeout
    if status == "broker_unavailable":
        # Re-observe the disk at 3 s intervals until this dispatch's wall-clock
        # deadline, read off the clock on every pass.
        while True:
            left = deadline - _now()
            if left <= 0:
                break
            _sleep(a, min(left, 3))
            disk = disk_status(reg, job)
            if disk in TERMINAL:
                status = final = disk
                if disk == "completed":
                    return final
                break
        if _now() >= deadline:  # Issue #33 (review R3): a wait into the deadline is the wall clock
            clock = "wall"
    if status == "timeout":
        # ISSUE-14 / ISSUE-18: keep re-observing the disk while the contract
        # says the agent is alive - an event in events.log at most IDLE_TIMEOUT
        # (K x N) old. Entering this grace opens one window and every disk event
        # re-opens it; evidence never extends the wall clock.
        evidence_until = a.evidence_until
        while _now() < deadline:
            disk = disk_status(reg, job)
            if disk in TERMINAL:
                status = final = disk
                if disk == "completed":
                    return final
                break
            if recent_disk_heartbeat(reg, job, a.idle_timeout):
                evidence_until = _now() + a.idle_timeout
            elif _now() >= evidence_until:
                break
            if deadline - _now() <= 0:
                break
            _sleep(a, 1)
        if _now() >= deadline:
            clock = "wall"

    # Issue #33: the subscriber's rc says which clock ran out (4 wall clock,
    # 2 no event); a grace above that ends on the deadline is the wall clock too.
    clock = clock or "idle"
    if a.sub_rc == SUB_RC_WALL_CLOCK:
        clock = "wall"
    close_req = "wall_clock_timeout" if clock == "wall" and a.phase != "inject" else "deadline_timeout"
    detail = ""
    if a.phase in ("preflight", "idle_wait", "registry_wait"):
        close_req, detail = "cancelled", "session_busy"
        # Issue #34: a target the preflight could not see was never busy
        if a.phase == "preflight" and a.preflight_reason:
            detail = a.preflight_reason
    if close_req != "cancelled":
        print(f"finalize: {job} closed by {'wall_clock' if clock == 'wall' else 'no_event'} clock "
              f"(budget {a.timeout}s / idle {a.idle_timeout}s)", file=sys.stderr)
    tok = ["--expect-dispatch-token", a.dispatch_token] if a.dispatch_token else []
    _check_owner(a)
    close_args = ["close", "--job", job, "--request", close_req] + (["--detail", detail] if detail else []) + tok
    if _run(a, "registry.py", close_args, merge=True) == 4:
        # a dispatch token mismatch (stale attempt): no secondary event (CAS/ABA)
        print("finalize_unfinished: dispatch token mismatch (stale attempt); suppressing event emission",
              file=sys.stderr)
        return final
    if disk_status(reg, job) == "completed":
        return "completed"
    # direct path only: publish the error/cancelled event (loop/discuss forbidden: O-1)
    if a.type == "direct":
        _check_owner(a)
        if close_req == "cancelled":
            _run(a, "publish_event.py", ["--job", job, "--event", "cancelled", "--detail", detail] + tok, merge=False)
        elif a.sub_rc != SUB_RC_ERROR or a.phase == "inject":
            err_detail = "agent_blocked" if a.phase == "inject" else close_req
            _run(a, "publish_event.py", ["--job", job, "--event", "error", "--detail", err_detail] + tok, merge=False)
    return final


def _answer(text: str) -> None:
    try:
        os.write(3, (text + "\n").encode())
    except OSError:
        print(text)


def _parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(prog="python -m lib_py.delegation.wrapper")
    sub = p.add_subparsers(dest="cmd", required=True)
    s = sub.add_parser("sub-status")
    s.add_argument("rc", type=int)
    s.add_argument("--job", required=True)
    s.add_argument("--registry-dir", required=True)
    s = sub.add_parser("subscriber-up")
    s.add_argument("rc", type=int)
    s.add_argument("--print", action="store_true", help="print yes/no and exit 0")
    sub.add_parser("exit-code").add_argument("status")
    sub.add_parser("wait-rc").add_argument("rc", type=int)
    s = sub.add_parser("stall-watch")
    s.add_argument("session")
    s.add_argument("agent")
    s.add_argument("--print", action="store_true", help="print yes/no and exit 0")
    s = sub.add_parser("wait-terminal")
    s.add_argument("--job", required=True)
    s.add_argument("--registry-dir", required=True)
    s.add_argument("--timeout", type=int, required=True)
    s.add_argument("--watch-pid", type=int, default=None, help="the caller's pid; gone -> stop waiting")
    sub.add_parser("onboard-exit").add_argument("status")
    for name in ("open-jobs", "wait-open-jobs"):
        s = sub.add_parser(name)
        s.add_argument("--job", required=True)
        s.add_argument("--agent-session", required=True)
        s.add_argument("--registry-dir", required=True)
        if name == "wait-open-jobs":
            s.add_argument("--wait-s", type=int, required=True, help="MAM_DISPATCH_WAIT_S")
            s.add_argument("--start-time", type=int, default=None)
            s.add_argument("--timeout", type=int, required=True, help="the job's wall-clock budget")
            s.add_argument("--watch-pid", type=int, default=None, help="the wrapper's pid; gone -> stop waiting")
    for name in ("finalize-begin", "finalize-end"):
        s = sub.add_parser(name)
        s.add_argument("--job", required=True)
        s.add_argument("--registry-dir", required=True)
        s.add_argument("--status", required=True)
        s.add_argument("--start-time", type=int, default=None)
        s.add_argument("--timeout", type=int, required=True)
        s.add_argument("--idle-timeout", type=int, required=True)
        if name == "finalize-begin":
            s.add_argument("--agent", default="")
            s.add_argument("--session", default="")
        else:
            s.add_argument("--sub-rc", type=int, required=True)
            s.add_argument("--phase", default="")
            s.add_argument("--type", default="direct")
            s.add_argument("--final", required=True)
            s.add_argument("--evidence-until", type=int, default=0)
            s.add_argument("--dispatch-token", default="")
            s.add_argument("--preflight-reason", default="")
            s.add_argument("--python", default=sys.executable)
            s.add_argument("--scripts-dir", required=True)
            s.add_argument("--watch-pid", type=int, default=None, help="the wrapper's pid; gone -> no close, no event")
    return p


def _yes_no(answer: bool, printed: bool) -> int:
    if printed:
        print("yes" if answer else "no")
        return 0
    return 0 if answer else 1


def main(argv: Optional[List[str]] = None) -> int:
    a = _parser().parse_args(argv)
    if a.cmd == "sub-status":
        print(sub_status(a.rc, a.registry_dir, a.job))
    elif a.cmd == "subscriber-up":
        return _yes_no(subscriber_up(a.rc), a.print)
    elif a.cmd == "exit-code":
        print(exit_code(a.status))
    elif a.cmd == "wait-rc":
        print(wait_rc(a.rc))
    elif a.cmd == "stall-watch":
        return _yes_no(stall_watch(a.session, a.agent), a.print)
    elif a.cmd == "finalize-begin":
        print(finalize_begin(a))
    elif a.cmd == "wait-terminal":
        try:
            print(wait_terminal(a))
        except OwnerGone:
            print(f"wait-terminal: caller pid {a.watch_pid} is gone; stopped waiting on job {a.job}", file=sys.stderr)
    elif a.cmd == "onboard-exit":
        print(onboard_exit(a.status))
    elif a.cmd == "open-jobs":
        live, stale = open_jobs(a.registry_dir, a.agent_session, a.job)
        for created, job, status in live:
            print(f"{job}  {status:8s}  created {created}")
        for job, reason in stale:
            print(f"{job}  stale     {reason}")
    elif a.cmd == "wait-open-jobs":
        try:
            print(wait_open_jobs(a))
        except OwnerGone:
            print(f"wait-open-jobs: wrapper pid {a.watch_pid} is gone; stopped waiting on job {a.job}", file=sys.stderr)
    else:
        try:
            _answer(finalize_end(a))
        except OwnerGone:
            print(f"finalize: wrapper pid {a.watch_pid} is gone; job {a.job} left unclosed", file=sys.stderr)
    return 0


if __name__ == "__main__":
    sys.exit(main())