multi-agent-mux/tests/test_issue72_admission.py

"""Issue #72 (decision 8, U2): while an update runs, the entries that start
something or write the registry refuse - create, resume, run_loop.sh, delegate
submit, stop (plain and --purge-conversation) and orc-onboard - and change
nothing. A marker its update left behind (`stale`) refuses nothing.

The marker is .mam/update-in-progress as the update skill publishes it (pid +
lstart, loop_lock.sh); `live` here is a sleeping helper process's.
"""
from __future__ import annotations

import hashlib
import json
import os
import subprocess
from pathlib import Path

import pytest
import yaml

UUID = "11111111-2222-3333-4444-555555555555"
ORC = "99999999-8888-7777-6666-555555555555"
ORC_OLD = "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee"


@pytest.fixture
def marker(mam_sandbox):
    procs = []

    def make(live: bool):
        p = mam_sandbox / ".mam" / "update-in-progress"
        if live:
            proc = subprocess.Popen(["sleep", "120"], start_new_session=True)
            procs.append(proc)
            ls = " ".join(subprocess.run(["ps", "-p", str(proc.pid), "-o", "lstart="],
                                         capture_output=True, text=True).stdout.split())
            p.write_text("pid=%d\nlstart=%s\nstarted=x\n" % (proc.pid, ls))
        else:
            dead = subprocess.Popen(["true"])
            dead.wait()
            p.write_text("pid=%d\nlstart=Thu Jan  1 00:00:00 1970\nstarted=x\n" % dead.pid)
        return p

    yield make
    for proc in procs:
        proc.kill()
        proc.wait()


def _seed(sandbox, rows, **top):
    lib = sandbox / ".agents" / "skills" / "lib.sh"
    mutation = "d['herdr_sessions'] = %r\n" % (rows,) + "".join("d[%r] = %r\n" % kv for kv in top.items())
    res = subprocess.run(["bash", "-c", 'source "$1" && atomic_dump_yaml "$AGENT_SESSIONS_YAML"', "_", str(lib)],
                         input=mutation, capture_output=True, text=True)
    assert res.returncode == 0, res.stderr


def _registry(sandbox):
    out = {}
    for f in ("agent-sessions.yaml", "agent-sessions.db"):
        p = sandbox / ".mam" / f
        out[f] = hashlib.sha256(p.read_bytes()).hexdigest() if p.exists() else None
    return out


def _script(sandbox, rel):
    return str(sandbox / ".agents" / "skills" / rel)


def _run(sandbox, argv, env=None, timeout=120):
    return subprocess.run(argv, capture_output=True, text=True, cwd=str(sandbox), timeout=timeout,
                          env=dict(os.environ, **(env or {})), stdin=subprocess.DEVNULL)


def _refused(res, sandbox):
    assert res.returncode != 0, res.stdout + res.stderr
    assert "an update of %s is in progress" % sandbox in res.stderr, res.stdout + res.stderr


def _row(name, status="running", agent="claude"):
    r = {"name": name, "status": status, "role": "creator", "herdr_session": "default", "herdr_server": "default",
         "pane": {"cwd": "", "cmd": agent}, "agent": agent}
    if agent == "claude":
        r["claude_session_id_own"] = UUID
    return r


def _agent(cwd, pane):
    return {"agent": "claude", "status": "running", "cwd": str(cwd), "workspace_id": "w1",
            "pid": 9999, "pane_id": pane, "command": "claude", "buffer": "Ready"}


def _set_mock(state_file, **kw):
    st = json.loads(Path(state_file).read_text())
    st.update(kw)
    Path(state_file).write_text(json.dumps(st))


def test_create_refuses_while_an_update_runs(mam_sandbox, mock_herdr, mock_agents, marker):
    marker(True)
    before = _registry(mam_sandbox)
    res = _run(mam_sandbox, ["bash", _script(mam_sandbox, "multi-agent-mux-create/scripts/create_session.sh"),
                             "--workspace", str(mam_sandbox), "--agent", "claude", "--role", "creator",
                             "--session", "gated-create", "--no-onboard"], timeout=300)
    _refused(res, mam_sandbox)
    assert _registry(mam_sandbox) == before
    assert json.loads(Path(mock_herdr).read_text())["agents"] == {}, "nothing was spawned"


def test_resume_refuses_on_the_spawn_and_the_alive_path(mam_sandbox, mock_herdr, mock_agents, marker):
    _seed(mam_sandbox, [dict(_row("gated-agy", status="stopped", agent="agy"), pane={"cwd": str(mam_sandbox), "cmd": "agy"})])
    marker(True)
    before = _registry(mam_sandbox)
    script = _script(mam_sandbox, "multi-agent-mux-resume/scripts/resume_session.sh")
    res = _run(mam_sandbox, ["bash", script, "--workspace", str(mam_sandbox), "--agent", "agy",
                             "--session", "gated-agy"], timeout=300)
    _refused(res, mam_sandbox)
    assert _registry(mam_sandbox) == before
    assert json.loads(Path(mock_herdr).read_text())["agents"] == {}, "nothing was spawned"

    _set_mock(mock_herdr, agents={"gated-agy": _agent(mam_sandbox, "w1:p9")})
    res = _run(mam_sandbox, ["bash", script, "--workspace", str(mam_sandbox), "--agent", "agy",
                             "--session", "gated-agy"], timeout=300)
    _refused(res, mam_sandbox)
    assert _registry(mam_sandbox) == before


def test_run_loop_refuses_before_its_loop_marker(mam_sandbox, marker):
    marker(True)
    loop_marker = mam_sandbox / ".mam" / "loop-guard-active"
    res = _run(mam_sandbox, ["bash", _script(mam_sandbox, "multi-agent-mux-loop/scripts/run_loop.sh"),
                             "--creator", "c", "--task", "t"],
               env={"MAM_LOOP_NO_FREEZE": "1", "MAM_REAL_ROOT": str(mam_sandbox)})
    _refused(res, mam_sandbox)
    assert not loop_marker.exists()


def test_delegate_submit_registers_no_job(mam_sandbox, mock_herdr, marker):
    marker(True)
    res = _run(mam_sandbox, ["bash", _script(mam_sandbox, "multi-agent-mux-delegate-job/multi-agent-mux-delegate-job"),
                             "submit", "--agent", "claude-code", "--agent-session", "herdr:anyone", "--prompt", "hi"])
    _refused(res, mam_sandbox)
    jobs = mam_sandbox / ".mam" / "jobs"
    assert not list(jobs.glob("*.json")) if jobs.exists() else True


@pytest.mark.parametrize("purge", [False, True])
def test_stop_refuses_and_changes_nothing(mam_sandbox, mock_herdr, marker, purge):
    _seed(mam_sandbox, [_row("gated-stop")])
    _set_mock(mock_herdr, agents={"gated-stop": _agent(mam_sandbox, "w1:p8")})
    marker(True)
    before = _registry(mam_sandbox)
    extra = ["--purge-conversation", "--yes"] if purge else []
    res = _run(mam_sandbox, ["bash", _script(mam_sandbox, "multi-agent-mux-stop/scripts/stop_session.sh"),
                             "--session", "gated-stop", "--agent", "claude", *extra])
    _refused(res, mam_sandbox)
    assert _registry(mam_sandbox) == before
    assert "gated-stop" in json.loads(Path(mock_herdr).read_text())["agents"], "the session was not touched"
    assert not (mam_sandbox / ".mam" / "purging-gated-stop").exists()


def test_orc_onboard_refuses_and_changes_nothing(mam_sandbox, marker):
    _seed(mam_sandbox, [], orchestrator_uuids=[ORC_OLD])
    marker(True)
    before = _registry(mam_sandbox)
    script = _script(mam_sandbox, "multi-agent-mux-orc-onboard/scripts/orc_onboard.sh")
    for args in (["--uuid", ORC], ["--remove", ORC_OLD]):
        res = _run(mam_sandbox, ["bash", script, *args])
        _refused(res, mam_sandbox)
    assert _registry(mam_sandbox) == before


def test_a_stale_marker_refuses_nothing(mam_sandbox, mock_herdr, marker):
    _seed(mam_sandbox, [_row("stale-stop")], orchestrator_uuids=[])
    _set_mock(mock_herdr, agents={"stale-stop": _agent(mam_sandbox, "w1:p7")})
    marker(False)
    res = _run(mam_sandbox, ["bash", _script(mam_sandbox, "multi-agent-mux-orc-onboard/scripts/orc_onboard.sh"),
                             "--uuid", ORC])
    assert res.returncode == 0, res.stdout + res.stderr
    assert "stale update marker" in res.stderr
    res = _run(mam_sandbox, ["bash", _script(mam_sandbox, "multi-agent-mux-stop/scripts/stop_session.sh"),
                             "--session", "stale-stop", "--agent", "claude"])
    assert res.returncode == 0, res.stdout + res.stderr
    d = yaml.safe_load((mam_sandbox / ".mam" / "agent-sessions.yaml").read_text())
    assert d["orchestrator_uuids"] == [ORC]
    assert {r["name"]: r["status"] for r in d["herdr_sessions"]} == {"stale-stop": "stopped"}


def test_the_admission_lock_is_released_after_the_record(mam_sandbox, marker):
    """An entry takes the lock shared and drops it once its record is written:
    a later exclusive taker (the update) gets it at once."""
    lib = mam_sandbox / ".agents" / "skills" / "lib.sh"
    res = subprocess.run(["bash", "-c", 'source "$1" && mam_update_admit && mam_update_admit_release && '
                          'exec 7>>"$WORKSPACE_ROOT/.mam/update-admission.lock" && '
                          'MAM_UPDATE_ADMISSION_WAIT_S=1 python3 -m lib_py.pane_lock acquire --fd 7 --wait-env MAM_UPDATE_ADMISSION_WAIT_S',
                          "_", str(lib)], capture_output=True, text=True, timeout=60)
    assert res.returncode == 0 and res.stdout.strip() == "locked", res.stdout + res.stderr