multi-agent-mux/tests/test_issue45_restore.py

"""Issue #45 against a real herdr server: a restarted MAM server must not relaunch
saved agents (A), and resume must not leave the session's restored empty shell
behind (B).

Everything runs on a throwaway herdr session under a short throwaway HOME, so
the server's socket, session.json and config never touch the user's herdr (the
running `multi-agent-mux` server included). The agent is a fake `claude`
script that logs its argv; the user's herdr config sets `[terminal]
default_shell` to a shell that puts the fake first on PATH, so herdr's own
restore command (`claude --resume <id>`) reaches the fake too. That is the
positive control: with restore on, the fake logs `--resume <id>`; through the
MAM bootstrap it must not. No real claude runs and no API is called.

A crash is `kill -9` of the server after session.json has the pane's
agent_session (herdr writes it a few seconds after a change).
"""
from __future__ import annotations

import json
import os
import shutil
import signal
import subprocess
import sys
import tempfile
import time
import uuid
from pathlib import Path

import pytest

HERDR = shutil.which("herdr")
pytestmark = pytest.mark.skipif(not HERDR, reason="real herdr is not installed")

REPO_ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(REPO_ROOT / ".agents" / "skills"))
from lib_py import herdr_config  # noqa: E402

# Inherited from the pane this suite may run in: they would point herdr at that
# pane's server and workspace.
SCOPE_ENV = ("HERDR_WORKSPACE_ID", "HERDR_PANE_ID", "HERDR_TAB_ID", "HERDR_SOCKET_PATH",
             "HERDR_SESSION", "HERDR_ENV", "HERDR_CONFIG_PATH", "XDG_CONFIG_HOME")


class Rig:
    def __init__(self, sandbox: Path, root: Path, session: str):
        self.sandbox, self.root, self.session = sandbox, root, session
        self.bin = root / "bin"
        self.argv_log = root / "argv.log"
        self.sdir = root / ".config" / "herdr" / "sessions" / session
        self.user_cfg = root / ".config" / "herdr" / "config.toml"

    # -- herdr ---------------------------------------------------------------
    def env(self, **extra) -> dict:
        env = {k: v for k, v in os.environ.items() if k not in SCOPE_ENV}
        env["PATH"] = f"{self.bin}:{env['PATH']}"
        env.update(extra)
        return env

    def hd(self, *args) -> dict:
        out = subprocess.run([HERDR, "--session", self.session, *args], env=self.env(),
                             capture_output=True, text=True, timeout=30)
        assert out.returncode == 0, f"herdr {args}: {out.stdout}{out.stderr}"
        return json.loads(out.stdout).get("result") or {} if out.stdout.strip() else {}

    def panes(self) -> list:
        return self.hd("pane", "list").get("panes") or []

    def server_pids(self) -> list:
        out = subprocess.run(["pgrep", "-f", "--", f"--session {self.session} server"],
                             capture_output=True, text=True).stdout
        return [int(p) for p in out.split()]

    def _wait_socket(self):
        sock = self.sdir / "herdr.sock"
        for _ in range(80):
            if sock.exists():
                return
            time.sleep(0.25)
        raise AssertionError("herdr server socket never appeared")

    def restores(self) -> int:
        log = self.sdir / "herdr-server.log"
        return log.read_text().count("session restore evaluated") if log.exists() else 0

    def start_raw(self, config: Path | None = None):
        """A server as the user would start it: restore on unless `config` says otherwise."""
        env = self.env(**({"HERDR_CONFIG_PATH": str(config)} if config else {}))
        subprocess.Popen([HERDR, "--session", self.session, "server"], env=env, start_new_session=True,
                         stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        self._wait_socket()

    def run_lib(self, body: str, **extra) -> subprocess.CompletedProcess:
        lib = self.sandbox / ".agents" / "skills" / "lib.sh"
        return subprocess.run(["bash", "-c", f'source "{lib}"\n{body}'], cwd=str(self.sandbox),
                              env=self.env(HERDR_SESSION_NAME=self.session, **extra),
                              capture_output=True, text=True, timeout=120)

    def start_via_shim(self):
        """The only way MAM starts a server: the shim's headless bootstrap."""
        res = self.run_lib("herdr has-session -t nobody-here || true")
        assert res.returncode == 0, res.stderr
        assert "restore-on-restart stays enabled" not in res.stderr, res.stderr
        try:
            self._wait_socket()
        except AssertionError:
            raise AssertionError(f"the shim did not start the server: {res.stderr}")

    def wait_restored(self, before: int):
        for _ in range(60):
            if self.restores() > before:
                time.sleep(3)  # let herdr spawn whatever it is going to spawn
                return
            time.sleep(0.25)
        raise AssertionError("the server never evaluated its saved session")

    # -- panes ---------------------------------------------------------------
    def wait_shell(self, pane_id: str):
        for _ in range(40):
            info = self.hd("pane", "process-info", "--pane", pane_id)["process_info"]
            if info.get("shell_pid"):
                return
            time.sleep(0.25)

    def new_pane(self, cwd: Path, label: str) -> str:
        cwd.mkdir(parents=True, exist_ok=True)
        pane = self.hd("workspace", "create", "--cwd", str(cwd), "--label", label, "--no-focus")["root_pane"]["pane_id"]
        self.wait_shell(pane)
        return pane

    def run_agent(self, pane_id: str, session_id: str):
        """The fake claude, plus the session report herdr's claude integration would send."""
        self.hd("pane", "run", pane_id, f"claude --session-id {session_id}")
        self.hd("pane", "report-agent", "--source", "herdr:claude", "--agent", "claude",
                "--state", "idle", "--agent-session-id", session_id, pane_id)

    def crash(self, *session_ids: str):
        """kill -9 the server once session.json holds every id."""
        saved = self.sdir / "session.json"
        for _ in range(80):
            text = saved.read_text() if saved.exists() else ""
            if all(s in text for s in session_ids):
                break
            time.sleep(0.25)
        else:
            raise AssertionError(f"session.json never recorded {session_ids}")
        for pid in self.server_pids():
            os.kill(pid, signal.SIGKILL)
        self.kill_fakes()
        time.sleep(0.5)
        (self.sdir / "herdr.sock").unlink(missing_ok=True)

    def kill_fakes(self):
        subprocess.run(["pkill", "-9", "-f", f"{self.bin}/claude"], capture_output=True)

    def resumed(self, session_id: str) -> bool:
        return self.argv_log.exists() and f"--resume {session_id}" in self.argv_log.read_text()


def _process_env(pid: int) -> list:
    """A process's environment as KEY=VALUE strings: /proc on Linux; on macOS
    `ps -E` appends it to the command line, so the entries are matched there."""
    proc = Path(f"/proc/{pid}/environ")
    if proc.exists():
        return [e.decode(errors="replace") for e in proc.read_bytes().split(b"\0") if e]
    out = subprocess.run(["ps", "-E", "-ww", "-o", "command=", "-p", str(pid)],
                         capture_output=True, text=True, timeout=10).stdout
    return out.split()


@pytest.fixture
def rig(mam_sandbox, monkeypatch):
    # A pytest tmp_path is too long for a unix socket path under it.
    # Resolved: on macOS /tmp is a link to /private/tmp and herdr reports real paths.
    root = Path(tempfile.mkdtemp(prefix="m45", dir="/tmp")).resolve()
    session = f"mamt45-{os.getpid()}-{uuid.uuid4().hex[:4]}"
    for k in SCOPE_ENV:
        monkeypatch.delenv(k, raising=False)
    monkeypatch.setenv("HOME", str(root))
    monkeypatch.setenv("HOME_DIR", str(root))
    r = Rig(mam_sandbox, root, session)
    r.bin.mkdir()
    fake = r.bin / "claude"
    fake.write_text(f'#!/usr/bin/env bash\n'
                    'if [ "${1:-}" = auth ]; then echo \'{"loggedIn": true}\'; exit 0; fi\n'
                    f'echo "$*" >> "{r.argv_log}"\n'
                    'echo "Welcome to Claude Code (fake)"\nwhile :; do sleep 1; done\n')
    fake.chmod(0o755)
    shell = root / "shell.sh"
    shell.write_text(f'#!/usr/bin/env bash\nexport PATH="{r.bin}:$PATH"\nexec bash --norc -i\n')
    shell.chmod(0o755)
    r.user_cfg.parent.mkdir(parents=True)
    r.user_cfg.write_text(f'[terminal]\ndefault_shell = "{shell}"\n')
    try:
        yield r
    finally:
        subprocess.run([HERDR, "session", "stop", session], env=r.env(), capture_output=True, timeout=30)
        for pid in r.server_pids():
            try:
                os.kill(pid, signal.SIGKILL)
            except ProcessLookupError:
                pass
        r.kill_fakes()
        subprocess.run(["pkill", "-9", "-f", str(root)], capture_output=True)
        subprocess.run([HERDR, "session", "delete", session], env=r.env(), capture_output=True, timeout=30)
        shutil.rmtree(root, ignore_errors=True)


def _pane_with(rig: Rig, session_id: str) -> dict | None:
    return next((p for p in rig.panes() if (p.get("agent_session") or {}).get("value") == session_id), None)


# --- A ---------------------------------------------------------------------

@pytest.mark.parametrize("user_session_table", [
    "",
    # an existing [session] table: appending a second one would break the TOML
    # and herdr would silently restore again
    "[session]\nresume_agents_on_restore = true\n",
], ids=["plain-user-config", "user-config-with-session-table"])
def test_a_mam_server_does_not_relaunch_saved_agents(rig, user_session_table):
    rig.user_cfg.write_text(rig.user_cfg.read_text() + user_session_table)
    sid = str(uuid.uuid4())

    rig.start_raw()
    rig.run_agent(rig.new_pane(rig.root / "wsA", "A"), sid)
    rig.crash(sid)

    # MAM brings the server back: nothing may be relaunched.
    before = rig.restores()
    rig.start_via_shim()
    rig.wait_restored(before)
    assert not rig.resumed(sid), rig.argv_log.read_text()
    pane = _pane_with(rig, sid)
    assert pane is not None and not pane.get("agent"), rig.panes()
    info = rig.hd("pane", "process-info", "--pane", pane["pane_id"])["process_info"]
    assert info["foreground_process_group_id"] == info["shell_pid"], info
    # and it got there through HERDR_CONFIG_PATH
    want = f"HERDR_CONFIG_PATH={rig.user_cfg.parent / herdr_config.OUTPUT_NAME}"
    assert want in _process_env(rig.server_pids()[0]), want

    # Positive control: the same saved layout under a plain server relaunches
    # the agent, so the "not resumed" above is not vacuous.
    rig.crash(sid)
    before = rig.restores()
    rig.start_raw()
    rig.wait_restored(before)
    assert rig.resumed(sid), rig.argv_log.read_text() if rig.argv_log.exists() else "no argv log"


# --- B ---------------------------------------------------------------------

def _close(rig: Rig, session_id: str) -> subprocess.CompletedProcess:
    env = rig.env(HERDR_SESSION_NAME=rig.session, PYTHONPATH=str(rig.sandbox / ".agents" / "skills"))
    return subprocess.run([sys.executable, "-m", "lib_py.herdr_restore", "close-restored-shells",
                           "--session-id", session_id, "--herdr", HERDR],
                          env=env, capture_output=True, text=True, timeout=60)


def test_b_only_the_resumed_sessions_restored_shell_is_closed(rig):
    ids = {k: str(uuid.uuid4()) for k in ("X", "Y", "Z", "X2")}
    rig.start_raw()
    for k in ids:
        rig.run_agent(rig.new_pane(rig.root / f"ws{k}", k), ids[k])
    rig.new_pane(rig.root / "wsU", "U")  # a user's plain shell: no agent_session
    rig.crash(*ids.values())
    shutil.rmtree(rig.root / "wsY")  # a removed worktree: restored into $HOME

    before = rig.restores()
    rig.start_via_shim()
    rig.wait_restored(before)
    rig.run_agent(_pane_with(rig, ids["Z"])["pane_id"], ids["Z"])  # running again
    rig.hd("pane", "run", _pane_with(rig, ids["X2"])["pane_id"], "sleep 60")  # busy shell
    time.sleep(1)
    assert _pane_with(rig, ids["Y"])["cwd"] == str(rig.root)
    n = len(rig.panes())

    for k in ("Z", "X2"):
        res = _close(rig, ids[k])
        assert res.returncode == 0 and "closed" not in res.stderr, res.stderr
    res = _close(rig, "")
    assert res.returncode == 0 and "closed" not in res.stderr, res.stderr
    assert len(rig.panes()) == n

    x = _pane_with(rig, ids["X"])["pane_id"]
    res = _close(rig, ids["X"])
    assert res.returncode == 0, res.stderr
    assert f"closed restored pane {x}" in res.stderr
    assert _pane_with(rig, ids["X"]) is None and len(rig.panes()) == n - 1

    res = _close(rig, ids["Y"])
    assert res.returncode == 0 and "closed restored pane" in res.stderr, res.stderr
    assert _pane_with(rig, ids["Y"]) is None and len(rig.panes()) == n - 2
    # the user's shell, the running agent and the busy shell are all still there
    assert any(p.get("cwd") == str(rig.root / "wsU") for p in rig.panes())
    assert _pane_with(rig, ids["Z"]) and _pane_with(rig, ids["X2"])


def test_b_default_server_is_never_touched(rig):
    env = rig.env(HERDR_SESSION_NAME="default", PYTHONPATH=str(rig.sandbox / ".agents" / "skills"))
    # a herdr that would fail loudly if called at all
    res = subprocess.run([sys.executable, "-m", "lib_py.herdr_restore", "close-restored-shells",
                          "--session-id", "some-id", "--herdr", "/nonexistent/herdr"],
                         env=env, capture_output=True, text=True, timeout=30)
    assert res.returncode == 0 and res.stderr == "", res.stderr


def _script(rig: Rig, rel: str, *args: str, timeout: int = 180) -> subprocess.CompletedProcess:
    """Output goes through files: a server the script bootstraps outlives it."""
    path = rig.sandbox / ".agents" / "skills" / rel
    out, err = rig.root / "script.out", rig.root / "script.err"
    with open(out, "w") as o, open(err, "w") as e:
        try:
            p = subprocess.run(["bash", str(path), *args], cwd=str(rig.sandbox), env=rig.env(),
                               stdout=o, stderr=e, text=True, timeout=timeout)
        except subprocess.TimeoutExpired:
            raise AssertionError(f"{rel} timed out:\n{out.read_text()}\n{err.read_text()}")
    return subprocess.CompletedProcess(p.args, p.returncode, out.read_text(), err.read_text())


def test_b_resume_leaves_no_restored_shell_behind(rig):
    """The real create -> crash -> resume path. The row create writes stays
    `running` through the crash, as it does in production; resume runs from a
    worktree whose control plane is the sandbox, so it spawns with a CP_PREFIX
    and the shim cannot reuse the restored shell (it splits a new pane)."""
    ws = rig.root / "wsX"
    ws.mkdir()
    name = "creator-claude-m45"
    res = _script(rig, "multi-agent-mux-create/scripts/create_session.sh", "--workspace", str(ws),
                  "--agent", "claude", "--role", "creator", "--session", name,
                  "--herdr-session", rig.session, "--no-onboard", timeout=60)
    assert res.returncode == 0, res.stdout + res.stderr
    rows = json.loads(rig.run_lib("load_state_json").stdout)["herdr_sessions"]
    row = next(r for r in rows if r["name"] == name)
    sid = row["claude_session_id_own"]
    assert row["status"] == "running" and sid
    pane = next(p for p in rig.panes() if p.get("cwd") == str(ws))
    # what herdr's claude integration reports for a real claude (checked on a live server)
    rig.hd("pane", "report-agent", "--source", "herdr:claude", "--agent", "claude",
           "--state", "idle", "--agent-session-id", sid, pane["pane_id"])
    rig.crash(sid)

    res = _script(rig, "multi-agent-mux-resume/scripts/resume_session.sh", "--workspace", str(ws),
                  "--agent", "claude", "--session", name)
    assert res.returncode == 0, res.stdout + res.stderr
    assert not rig.resumed(sid)  # herdr did not relaunch it ...
    # ... MAM did, with its own command: create's launch plus resume's, both flagged
    launches = rig.argv_log.read_text().splitlines()
    assert len(launches) == 2 and all("--dangerously-skip-permissions" in l and sid in l for l in launches), launches
    in_ws = [p for p in rig.panes() if p.get("cwd") == str(ws)]
    assert len(in_ws) == 1, in_ws
    info = rig.hd("pane", "process-info", "--pane", in_ws[0]["pane_id"])["process_info"]
    assert info["foreground_process_group_id"] != info["shell_pid"], info  # the agent, not an empty shell
    assert "closed restored pane" in res.stderr, res.stderr