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>                         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>              exit 0: start the stall watcher (Issue #30)
  finalize-begin / finalize-end ...          finalize_unfinished, split around the one
                                             herdr diagnostic probe bash still runs

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.

Stdlib only and Python 3.9 (macOS /usr/bin/python3); 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}


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:
    try:
        os.kill(pid, 0)
    except ProcessLookupError:
        return False
    except PermissionError:
        return True
    return True


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.2 s."""
    end = time.time() + seconds
    while True:
        _check_owner(a)
        left = end - time.time()
        if left <= 0:
            return
        time.sleep(min(left, 0.2))


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


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"):
        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)
    sub.add_parser("subscriber-up").add_argument("rc", type=int)
    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")
    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 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 0 if subscriber_up(a.rc) else 1
    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 0 if stall_watch(a.session, a.agent) else 1
    elif a.cmd == "finalize-begin":
        print(finalize_begin(a))
    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())