multi-agent-mux/tests/test_issue30_stall_watch_stop.py

"""Issue #30 review (PR #41, N1): the wrapper never waits forever for its
stall watcher, and the watcher never outlives its wrapper.

The wrapper stops the watcher with SIGTERM the moment its subscriber ends.
A TERM that lands inside the watcher's fork-exec window is swallowed (bash
defers terminating signals in the forked child until the exec), so with an
unbounded `wait` the wrapper hung until someone killed the watcher by hand.
Both tests run on macOS and Linux: no /proc, no live broker.
"""

from __future__ import annotations

import json
import os
import re
import subprocess
import sys
import time
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parent.parent
SKILLS = REPO_ROOT / ".agents" / "skills"
WRAPPER = SKILLS / "multi-agent-mux-delegate-job" / "multi-agent-mux-delegate-job"
PUBLISH = SKILLS / "multi-agent-mux-delegate-job" / "scripts" / "publish_event.py"

UNREACHABLE = {"MQTT_BROKER": "127.0.0.1", "MQTT_HOST": "127.0.0.1", "MQTT_PORT": "65534", "MQTT_TLS": "0"}


def _env(home: Path, **extra) -> dict:
    env = {k: v for k, v in os.environ.items() if not k.startswith("HERDR_")}
    env.update(UNREACHABLE, DELEGATE_JOB_PYTHON=sys.executable, AGENT_PYTHON_BIN=sys.executable,
               PYTHONPATH=str(SKILLS), MAM_JOB_ID="", HOME=str(home), WORKSPACE_ROOT=str(home),
               DELEGATE_JOB_LOGS_DIR=str(home / ".mam" / "logs"))
    env.update({k: str(v) for k, v in extra.items()})
    return env


def _record(reg: Path, jid: str) -> dict:
    return json.loads((reg / f"{jid}.json").read_text())


def _publish(env: dict, reg: Path, jid: str, event: str, *extra: str) -> None:
    res = subprocess.run([sys.executable, str(PUBLISH), "--registry-dir", str(reg), "--job", jid, "--event", event,
                          "--expect-dispatch-token", _record(reg, jid)["dispatch_token"], "--attempts", "1", *extra],
                         capture_output=True, text=True, env=env, timeout=120)
    assert res.returncode in (0, 2), res.stderr


def _alive(pid: int) -> bool:
    try:
        os.kill(pid, 0)
    except ProcessLookupError:
        return False
    return True


def _until(pred, limit=20.0):
    end = time.monotonic() + limit
    while time.monotonic() < end:
        if pred():
            return True
        time.sleep(0.2)
    return pred()


def test_stop_stall_watch_is_bounded_when_the_watcher_ignores_term(tmp_path):
    """The wrapper's stop routine, run as is against a stand-in that ignores
    SIGTERM (what a swallowed TERM looks like from the wrapper's side): it
    returns within its 5 s budget, says so, and the stand-in is dead."""
    fn = subprocess.run(["sed", "-n", "/^stop_stall_watch()/,/^}/p", str(WRAPPER)],
                        capture_output=True, text=True, check=True).stdout
    assert fn.startswith("stop_stall_watch()") and fn.rstrip().endswith("}")
    script = (fn + '\nJOB_ID=t; REGISTRY_DIR="$1"\n'
              '(trap "" TERM; exec sleep 300) & STALL_PID=$!\n'
              'sleep 0.5; t0=$SECONDS; stop_stall_watch; echo "took=$((SECONDS - t0))"\n'
              'kill -0 "$!" 2>/dev/null && echo "stand-in alive" || echo "stand-in gone"\n')
    t0 = time.monotonic()
    res = subprocess.run(["bash", "-c", script, "stop", str(tmp_path)], capture_output=True, text=True, timeout=30)
    took = time.monotonic() - t0
    assert res.returncode == 0, res.stdout + res.stderr
    assert "stand-in gone" in res.stdout, res.stdout
    assert "stall watch for t ignored TERM; killed" in res.stderr, res.stderr
    assert 4 < took < 12, took


def test_the_wrapper_returns_when_its_watcher_ignores_term(tmp_path, mock_herdr):
    """The Issue #14 ack scenario that hung: broker unreachable, the RPC fails,
    the registry ack wins, the subscriber ends at once and the wrapper stops
    the watcher. Here the watcher is a stand-in whose shim ignores SIGTERM from
    its first line: the wrapper must still return, on its usual
    broker_unavailable path, and leave no watcher behind."""
    reg = tmp_path / "jobs"
    reg.mkdir()
    pidfile = tmp_path / "watcher.pid"
    shim = tmp_path / "py-shim"
    shim.write_text(
        "#!/bin/sh\n"
        'case "$*" in *lib_py.stall_watch*)\n'
        "  trap '' TERM\n"
        f'  exec "{sys.executable}" -c "import os, time; '
        f"open({str(pidfile)!r}, 'w').write(str(os.getpid())); time.sleep(300)\" ;;\n"
        "esac\n"
        f'exec "{sys.executable}" "$@"\n')
    shim.chmod(0o755)

    state_file = Path(mock_herdr)
    state = json.loads(state_file.read_text())
    state["agents"]["ack-agent"] = {"agent": "codex", "agent_status": "idle", "prompt_wait_result": "transport_error"}
    state["prompt_wait_result"] = "transport_error"
    state_file.write_text(json.dumps(state, indent=2))

    env = _env(tmp_path, DELEGATE_JOB_PYTHON=str(shim), SKS_IDLE_TIMEOUT_MS=500, MAM_INJECT_ACK_GRACE_S=8)
    proc = subprocess.Popen(["bash", str(WRAPPER), "submit", "--agent", "codex", "--agent-session", "herdr:ack-agent",
                             "--prompt", "issue 30 stop", "--registry-dir", str(reg), "--workdir", str(tmp_path),
                             "--timeout", "6", "--idle-timeout", "2"],
                            cwd=str(tmp_path), env=env, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    try:
        assert _until(lambda: list(reg.glob("*.json")), 10), "job was not registered"
        jid = list(reg.glob("*.json"))[0].stem
        _publish(env, reg, jid, "started", "--attempt-id", "1")
        t0 = time.monotonic()
        stdout, stderr = proc.communicate(timeout=60)
        took = time.monotonic() - t0
    finally:
        if proc.poll() is None:
            proc.kill()
    assert proc.returncode == 3, f"expected the broker_unavailable path (rc=3), got {proc.returncode}\n{stderr}\n{stdout}"
    # ack grace (8 s) + the 5 s TERM budget + slack; an unbounded wait never returns.
    assert took < 40, took
    if pidfile.exists():  # the stand-in got past its shim (a TERM in the shim's own start-up can still end it)
        assert _until(lambda: not _alive(int(pidfile.read_text())), 5), "the watcher outlived the wrapper"
        assert re.search(r"stall watch for \S+ ignored TERM", stderr), stderr


def test_the_watcher_stops_once_its_wrapper_is_gone(mam_sandbox, mock_herdr):
    """`--watch-pid`: a watcher whose wrapper vanished (killed, crashed) exits
    with rc 0 on its own instead of polling herdr until the job closes."""
    reg = mam_sandbox / ".mam" / "jobs"
    env = _env(mam_sandbox, MAM_STALL_IDLE_S=600)
    res = subprocess.run(["bash", str(WRAPPER), "submit", "--agent", "human", "--agent-session", "herdr:img",
                          "--dry-run", "--registry-dir", str(reg), "--prompt", "work"],
                         capture_output=True, text=True, env=env, cwd=str(mam_sandbox), timeout=120)
    assert res.returncode == 0, res.stdout + res.stderr
    jid = re.search(r"registered job:\s*(\S+)", res.stdout).group(1)
    _publish(env, reg, jid, "started")
    assert _record(reg, jid)["status"] == "running"

    state_file = Path(mock_herdr)
    state = json.loads(state_file.read_text())
    state["wait_blocks_s"] = 0.5
    state.setdefault("agents", {})["creator-codex-t"] = {
        "name": "creator-codex-t", "agent": "codex", "status": "running", "agent_status": "idle",
        "pane_id": "w1:p1", "workspace_id": "w1", "pid": 4242, "buffer": "› Ask Codex to do anything"}
    state_file.write_text(json.dumps(state, indent=2))

    stand_in = subprocess.Popen(["sleep", "300"])
    w = subprocess.Popen([sys.executable, "-m", "lib_py.stall_watch", "--job", jid, "--registry-dir", str(reg),
                          "--target", "creator-codex-t", "--agent", "codex", "--watch-pid", str(stand_in.pid)],
                         env=env, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
    try:
        time.sleep(3)
        assert w.poll() is None, "the watcher left while its wrapper was alive:\n" + w.stdout.read()
        stand_in.kill()
        stand_in.wait()
        out, _ = w.communicate(timeout=30)
    finally:
        if w.poll() is None:
            w.kill()
    assert w.returncode == 0, out
    assert f"wrapper pid {stand_in.pid} is gone" in out, out
    assert _record(reg, jid)["status"] == "running"  # notify-only: it closed nothing


# ---------------------------------------------------------------------------
# codex R1 - the real wrapper killed by TERM leaves no watcher behind
# ---------------------------------------------------------------------------

def test_a_terminated_wrapper_leaves_no_watcher_behind(mam_sandbox, mock_herdr):
    """The wrapper's own functions: `start_stall_watch` then `wait`, and TERM to
    the wrapper only. The orphaned watcher must end on its own (rc 0), since
    no trap can run for a wrapper that is simply gone."""
    reg = mam_sandbox / ".mam" / "jobs"
    env = _env(mam_sandbox, MAM_STALL_IDLE_S=600)
    res = subprocess.run(["bash", str(WRAPPER), "submit", "--agent", "human", "--agent-session", "herdr:img",
                          "--dry-run", "--registry-dir", str(reg), "--prompt", "work"],
                         capture_output=True, text=True, env=env, cwd=str(mam_sandbox), timeout=120)
    assert res.returncode == 0, res.stdout + res.stderr
    jid = re.search(r"registered job:\s*(\S+)", res.stdout).group(1)
    _publish(env, reg, jid, "started")

    state_file = Path(mock_herdr)
    state = json.loads(state_file.read_text())
    state["wait_blocks_s"] = 0.5
    state.setdefault("agents", {})["creator-codex-t"] = {
        "name": "creator-codex-t", "agent": "codex", "status": "running", "agent_status": "idle",
        "pane_id": "w1:p1", "workspace_id": "w1", "pid": 4242, "buffer": "› Ask Codex to do anything"}
    state_file.write_text(json.dumps(state, indent=2))

    script = ('MAM_DELEGATE_JOB_NO_MAIN=1 source "$1"; JOB_ID="$2"; REGISTRY_DIR="$3"; PY="$4"\n'
              'start_stall_watch herdr:creator-codex-t codex\n'
              'echo "WATCHER=$STALL_PID"; wait\n')
    proc = subprocess.Popen(["bash", "-c", script, "probe", str(WRAPPER), jid, str(reg), sys.executable],
                            env=env, cwd=str(mam_sandbox), stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
    wpid = None
    try:
        for line in proc.stdout:
            if line.startswith("WATCHER="):
                wpid = int(line.split("=", 1)[1])
                break
        assert wpid and _alive(wpid), "the watcher did not start"
        time.sleep(2)
        proc.terminate()  # TERM to the wrapper only
        proc.wait(timeout=10)
        assert _until(lambda: not _alive(wpid), 30), "the watcher outlived its wrapper"
    finally:
        if proc.poll() is None:
            proc.kill()
        if wpid and _alive(wpid):
            os.kill(wpid, 9)
    out = (reg / f"{jid}.stall.out").read_text()
    assert f"wrapper pid {proc.pid} is gone" in out, out
    assert _record(reg, jid)["status"] == "running"


# ---------------------------------------------------------------------------
# codex R2 - a job that closed during a blocking wait is not nudged
# ---------------------------------------------------------------------------

class _Src:
    """A source whose facts always say "idle, alone, empty input"."""

    def __init__(self):
        self.prompts = []

    def collect(self, sig):
        sig.state_change_seq, sig.agent_pid, sig.background_children, sig.input_empty = 1, 4242, [], True
        return sig

    def prompt(self, text):
        self.prompts.append(text)

    def blocked_text(self, pane_id):
        return ""

    def agent_get(self):
        return {"pane_id": "w1:p1"}


def _watch_in_process(tmp_path, monkeypatch, waits, notify):
    from lib_py import stall_watch
    from lib_py.loop import watcher as w
    import argparse
    import uuid
    jid = uuid.uuid4().hex[:8]
    rec = tmp_path / f"{jid}.json"
    rec.write_text(json.dumps({"job_id": jid, "status": "running", "dispatch_token": "tok", "attempt_id": 1}))
    calls = []

    def fake_wait_once(target, until=(), timeout_ms=None, session=None, run=None):
        calls.append(until)
        return waits[min(len(calls), len(waits)) - 1](rec)

    monkeypatch.setattr(stall_watch.watcher, "wait_once", fake_wait_once)
    monkeypatch.setenv("MAM_STALL_IDLE_S", "1")
    args = argparse.Namespace(job=jid, registry_dir=str(tmp_path), target="creator-codex-t", agent="codex",
                              herdr_session="", expect_dispatch_token="tok", attempt_id=1, watch_pid=None)
    src = _Src()
    sw = stall_watch.StallWatch(args, source=src)
    notes = []
    monkeypatch.setattr(sw, "_notify", lambda detail: notes.append(detail) or notify)
    rc = sw.run()
    return rc, src.prompts, notes, rec, w


def _idle(rec):
    from lib_py.loop import watcher as w
    return w.WaitOutcome(w.MATCHED, agent={"agent_status": "idle", "state_change_seq": 1, "pane_id": "w1:p1"})


def _completed_during_the_wait(rec):
    from lib_py.loop import watcher as w
    time.sleep(1.2)  # past S while "blocked in herdr"
    d = json.loads(rec.read_text())
    d["status"] = "completed"
    rec.write_text(json.dumps(d))
    return w.WaitOutcome(w.TIMEOUT, error_code="timeout")


def _redispatched_during_the_wait(rec):
    from lib_py.loop import watcher as w
    time.sleep(1.2)
    d = json.loads(rec.read_text())
    if d.get("ticks"):  # a watcher that ignores the new token is stopped by the close on the next tick
        d["status"] = "completed"
    d["ticks"] = d.get("ticks", 0) + 1
    d["dispatch_token"], d["attempt_id"] = "tok-2", 2
    rec.write_text(json.dumps(d))
    return w.WaitOutcome(w.TIMEOUT, error_code="timeout")


def test_a_job_completed_during_the_wait_is_not_nudged(tmp_path, monkeypatch):
    rc, prompts, notes, _, _ = _watch_in_process(tmp_path, monkeypatch, [_idle, _completed_during_the_wait], True)
    assert rc == 0 and prompts == [] and notes == [], (prompts, notes)


def test_a_job_redispatched_during_the_wait_is_not_nudged(tmp_path, monkeypatch):
    rc, prompts, notes, rec, _ = _watch_in_process(tmp_path, monkeypatch, [_idle, _redispatched_during_the_wait], True)
    assert rc == 0 and prompts == [] and notes == [], (prompts, notes)
    assert json.loads(rec.read_text())["status"] == "running"  # it left on the token, and closed nothing


def test_a_refused_idle_notice_withholds_the_nudge(tmp_path, monkeypatch):
    """The registry (CAS) or the wire refused the `stalled:idle` progress: no
    prompt goes to the pane, and the job is then left alone (nudged once)."""
    def _idle_then_close(rec):
        from lib_py.loop import watcher as w
        time.sleep(1.2)
        d = json.loads(rec.read_text())
        # the second idle tick nudges (or not); the third finds the job closed and ends the run
        if d.get("ticks"):
            d["status"] = "completed"
        d["ticks"] = d.get("ticks", 0) + 1
        rec.write_text(json.dumps(d))
        return w.WaitOutcome(w.TIMEOUT, error_code="timeout")
    rc, prompts, notes, _, _ = _watch_in_process(tmp_path, monkeypatch, [_idle, _idle_then_close], False)
    assert rc == 0 and notes == ["stalled:idle"] and prompts == [], (prompts, notes)