multi-agent-mux/tests/test_issue50_onboard.py

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

- U: the decisions (lib_py.delegation.wrapper wait-terminal / onboard-exit).
- S: create_session.sh in the sandbox (mock herdr whose fake agent answers
  the onboarding prompt through the real publish_event.py): completed -> 0
  and a one-line result; error -> 44; never ended -> 45; the session is kept
  and the row says why; worktree_prepare.sh passes 44/45 through.
"""
from __future__ import annotations

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

import pytest
import yaml

REPO_ROOT = Path(__file__).resolve().parent.parent
SKILLS = REPO_ROOT / ".agents" / "skills"
if str(SKILLS) not in sys.path:
    sys.path.insert(0, str(SKILLS))

from lib_py.delegation import registry, wrapper as decide  # noqa: E402


# ---------------------------------------------------------------------------
# U - decisions
# ---------------------------------------------------------------------------

@pytest.mark.parametrize("status,code", [("completed", 0), ("error", 44), ("cancelled", 44),
                                         ("timeout", 45), ("missing", 45), ("", 45), ("garbage", 45)])
def test_u1_onboard_exit_table(status, code):
    assert decide.onboard_exit(status) == code


def _wait_terminal(reg, job, timeout=1, watch_pid=None):
    cmd = [sys.executable, "-m", "lib_py.delegation.wrapper", "wait-terminal", "--job", job,
           "--registry-dir", str(reg), "--timeout", str(timeout)]
    if watch_pid is not None:
        cmd += ["--watch-pid", str(watch_pid)]
    return subprocess.run(cmd, capture_output=True, text=True, cwd=str(SKILLS),
                          env=dict(os.environ, PYTHONPATH=str(SKILLS)), timeout=60)


@pytest.mark.parametrize("terminal", ["completed", "error", "cancelled"])
def test_u2_wait_terminal_prints_the_terminal_status_at_once(tmp_path, terminal):
    jid = registry.register_job(prompt="onboarding", agent="claude-code", agent_session="herdr:s",
                                registry_dir=str(tmp_path))
    rec = json.loads((tmp_path / f"{jid}.json").read_text())
    rec["status"] = terminal
    (tmp_path / f"{jid}.json").write_text(json.dumps(rec))
    t0 = time.time()
    res = _wait_terminal(tmp_path, jid, timeout=30)
    assert res.stdout.strip() == terminal and res.returncode == 0, res.stderr
    assert time.time() - t0 < 10


def test_u2b_wait_terminal_times_out_on_a_job_that_never_ends(tmp_path):
    jid = registry.register_job(prompt="onboarding", agent="claude-code", agent_session="herdr:s",
                                registry_dir=str(tmp_path))
    res = _wait_terminal(tmp_path, jid, timeout=1)
    assert res.stdout.strip() == "timeout", res.stdout + res.stderr


def test_u2c_wait_terminal_reports_a_missing_record(tmp_path):
    res = _wait_terminal(tmp_path, "nothere", timeout=30)
    assert res.stdout.strip() == "missing"


def test_u2d_wait_terminal_stops_with_its_caller(tmp_path):
    jid = registry.register_job(prompt="onboarding", agent="claude-code", agent_session="herdr:s",
                                registry_dir=str(tmp_path))
    gone = subprocess.Popen(["true"])
    gone.wait()
    res = _wait_terminal(tmp_path, jid, timeout=30, watch_pid=gone.pid)
    assert res.stdout == "" and "is gone" in res.stderr, res.stdout + res.stderr


# ---------------------------------------------------------------------------
# S - create_session.sh
# ---------------------------------------------------------------------------

def _create(sandbox, *extra, env=None):
    script = sandbox / ".agents" / "skills" / "multi-agent-mux-create" / "scripts" / "create_session.sh"
    cmd = ["bash", str(script), "--workspace", str(sandbox), "--agent", "claude", "--role", "creator",
           "--session", "t50-creator-claude", "--onboard", *extra]
    return subprocess.run(cmd, capture_output=True, text=True, cwd=str(sandbox), env=env, timeout=300)


def _row(sandbox):
    (row,) = yaml.safe_load((sandbox / ".mam" / "agent-sessions.yaml").read_text())["herdr_sessions"]
    return row


def _job(sandbox, jid):
    return json.loads((sandbox / ".mam" / "jobs" / f"{jid}.json").read_text())


def _mock_state(mock_herdr):
    return json.loads(Path(mock_herdr).read_text())


def test_s1_create_ends_with_the_onboarding_completed(mam_sandbox, mock_herdr, mock_agents):
    res = _create(mam_sandbox)
    assert res.returncode == 0, res.stdout + res.stderr
    row = _row(mam_sandbox)
    jid = row["delegate_job_id"]
    assert f"onboarding completed: job {jid} (" in res.stdout, res.stdout
    assert _job(mam_sandbox, jid)["status"] == "completed"
    assert not str(row.get("last_visible_status", "")).startswith("onboarding-")
    assert "onboarding" not in res.stderr.lower() or "WARNING" not in res.stderr


def test_s2_an_onboarding_that_ends_with_error_exits_44_and_keeps_the_session(mam_sandbox, mock_herdr, mock_agents,
                                                                               monkeypatch):
    monkeypatch.setenv("MAM_MOCK_ONBOARD_HANDSHAKE", "0")   # no agent answer: no started -> error
    res = _create(mam_sandbox)
    assert res.returncode == 44, res.stdout + res.stderr
    row = _row(mam_sandbox)
    jid = row["delegate_job_id"]
    assert jid, row
    assert f"onboarding error for 't50-creator-claude': job {jid}; the session is kept" in res.stderr, res.stderr
    assert _job(mam_sandbox, jid)["status"] == "error"
    assert row["status"] == "running" and row["last_visible_status"] == "onboarding-error"
    state = _mock_state(mock_herdr)
    assert "t50-creator-claude" in state["agents"], "the session was rolled back"
    assert not any(c[:1] == ["kill-session"] for c in state["calls"]), "the EXIT trap ran"
    assert "onboarding completed" not in res.stdout


def test_s3_an_onboarding_that_never_ends_exits_45_and_keeps_the_session(mam_sandbox, mock_herdr, mock_agents,
                                                                          monkeypatch):
    monkeypatch.setenv("MAM_MOCK_ONBOARD_HANDSHAKE", "started")   # started, never completed
    t0 = time.time()
    res = _create(mam_sandbox)
    assert res.returncode == 45, res.stdout + res.stderr
    row = _row(mam_sandbox)
    jid = row["delegate_job_id"]
    assert f"onboarding timeout for 't50-creator-claude': job {jid}; the session is kept" in res.stderr, res.stderr
    assert _job(mam_sandbox, jid)["status"] == "running"
    assert row["status"] == "running" and row["last_visible_status"] == "onboarding-timeout"
    assert "t50-creator-claude" in _mock_state(mock_herdr)["agents"]
    assert time.time() - t0 < 120, "MAM_ONBOARD_WAIT_S=3 bounds the wait in the sandbox"


def test_s3b_the_wait_knob_follows_the_integer_rule(mam_sandbox, mock_herdr, mock_agents, monkeypatch):
    monkeypatch.setenv("MAM_MOCK_ONBOARD_HANDSHAKE", "started")
    monkeypatch.setenv("MAM_ONBOARD_WAIT_S", "0")   # below the floor of 1: the default (600) with a warning
    monkeypatch.setenv("MAM_ONBOARD_ACK_S", "2")
    script = mam_sandbox / ".agents" / "skills" / "multi-agent-mux-create" / "scripts" / "create_session.sh"
    src = script.read_text()
    # bound the run: the default would wait 600 s; cap the python wait by a stub of the knob call
    assert 'ONBOARD_WAIT_S=$(mam_int_knob MAM_ONBOARD_WAIT_S 600 1)' in src
    script.write_text(src.replace('ONBOARD_WAIT_S=$(mam_int_knob MAM_ONBOARD_WAIT_S 600 1)',
                                  'ONBOARD_WAIT_S=$(mam_int_knob MAM_ONBOARD_WAIT_S 600 1); echo "KNOB=$ONBOARD_WAIT_S"; ONBOARD_WAIT_S=1'))
    res = _create(mam_sandbox)
    assert "KNOB=600" in res.stdout, res.stdout
    assert res.returncode == 45


def test_s4_the_source_clears_the_trap_before_the_named_exits():
    src = (SKILLS / "multi-agent-mux-create" / "scripts" / "create_session.sh").read_text()
    tail = src.split('_mark_row_note "onboarding-$ONBOARD_STATUS"')[1].split("esac")[0]
    assert tail.index("trap - EXIT") < tail.index("exit 44") < tail.index("exit 45")
    assert "delegate_wait_terminal" in src and "delegate_onboard_exit" in src


# ---------------------------------------------------------------------------
# S5 - worktree_prepare.sh passes 44 / 45 through with its own message
# ---------------------------------------------------------------------------

def _git(repo, *args):
    subprocess.run(["git", "-c", "user.email=t@example.com", "-c", "user.name=t", *args],
                   cwd=str(repo), check=True, capture_output=True, text=True)


@pytest.mark.parametrize("rc,word", [(44, "error"), (45, "timeout")])
def test_s5_worktree_prepare_passes_create_exit_through(mam_sandbox, rc, word):
    main = mam_sandbox
    _git(main, "init", "-q")
    (main / "README").write_text("x\n")
    _git(main, "add", "README")
    _git(main, "commit", "-q", "-m", "init")
    stub = main / ".agents" / "skills" / "multi-agent-mux-create" / "scripts" / "create_session.sh"
    stub.write_text(f'#!/usr/bin/env bash\necho "stub create: $*" >&2\nexit {rc}\n')
    prep = main / ".agents" / "skills" / "multi-agent-mux-create" / "scripts" / "worktree_prepare.sh"
    wt = main.parent / f"{main.name}-wt{rc}"
    res = subprocess.run(["bash", str(prep), "--path", str(wt), "--branch", f"b-{rc}",
                          "--creator-agent", "claude", "--onboard"],
                         capture_output=True, text=True, cwd=str(main), timeout=300)
    try:
        assert res.returncode == rc, res.stdout + res.stderr
        assert f"onboarding {word} for 'creator-claude-b-{rc}' (create_session.sh rc={rc}); session and worktree kept" in res.stderr, res.stderr
        assert "worktree_remove.sh" not in res.stderr
        assert "--onboard" in res.stderr and "stub create" in res.stderr
        assert (wt / "README").exists(), "the worktree is kept"
    finally:
        subprocess.run(["git", "worktree", "remove", "--force", str(wt)], cwd=str(main), capture_output=True)