multi-agent-mux/tests/_mam_state_guard.py

"""Issue #20 - helpers for the session guard over real control planes.

A pytest process can reach two kinds of `.mam`: this checkout's own, and the
one the agent pane's control-plane binding named before conftest's
pytest_configure dropped it (a worktree creator's pane carries
WORKSPACE_ROOT, AGENT_SESSIONS_YAML, MAM_JOBS_DIR and DELEGATE_JOB_LOGS_DIR,
all at the live main checkout). The session guard snapshots their job
registries and audit-log roots and, at the end, reports the new entries no
live delegation explains. Pure functions, unit-tested in
test_issue20_state_isolation.py (S4, S5); the name does not start with
`test_`, so pytest never collects it.
"""
from __future__ import annotations

import json
import os
import re

SESSION_PREFIXES = ("herdr:", "remote:")
# Top-level items of `herdr_sessions:` as atomic_dump_yaml writes them
# (`- name: <session>` at column 0); indented `- name:` lines belong to nested
# lists such as mcp_attachments. A regex keeps the guard free of PyYAML.
_SESSION_LINE = re.compile(r"""^- name:\s*['"]?([^'"\s]+)""")


def _listdir(path):
    try:
        return os.listdir(path)
    except OSError:
        return []


def _load(path):
    try:
        with open(path, encoding="utf-8") as fh:
            data = json.load(fh)
    except (OSError, ValueError):
        return {}
    return data if isinstance(data, dict) else {}


def watch_roots(repo_root, recorded):
    """`.mam` directories to watch, realpath-deduplicated, in order: this
    checkout's own (always, so a leak that creates it is still seen), then
    those the recorded binding names - `$WORKSPACE_ROOT/.mam` and the parents
    of MAM_JOBS_DIR / DELEGATE_JOB_LOGS_DIR - when they exist."""
    roots = [os.path.realpath(os.path.join(str(repo_root), ".mam"))]
    candidates = []
    if recorded.get("WORKSPACE_ROOT"):
        candidates.append(os.path.join(recorded["WORKSPACE_ROOT"], ".mam"))
    for var in ("MAM_JOBS_DIR", "DELEGATE_JOB_LOGS_DIR"):
        if recorded.get(var):
            candidates.append(os.path.dirname(os.path.abspath(recorded[var])))
    for path in candidates:
        real = os.path.realpath(path)
        if os.path.isdir(real) and real not in roots:
            roots.append(real)
    return roots


def state_dirs(roots, recorded):
    """(registry dirs, audit-log dirs): `<root>/jobs` and
    `<root>/delegate_job_logs` of every root, plus a recorded MAM_JOBS_DIR /
    DELEGATE_JOB_LOGS_DIR that is not one of those."""
    jobs = [os.path.join(r, "jobs") for r in roots]
    logs = [os.path.join(r, "delegate_job_logs") for r in roots]
    for var, dirs in (("MAM_JOBS_DIR", jobs), ("DELEGATE_JOB_LOGS_DIR", logs)):
        if recorded.get(var):
            real = os.path.realpath(recorded[var])
            if os.path.isdir(real) and real not in dirs:
                dirs.append(real)
    return jobs, logs


def snapshot(roots, recorded):
    """Registry records (`*.json`) and audit-log entries per watched dir; a
    missing dir counts as empty."""
    jobs, logs = state_dirs(roots, recorded)
    return {
        "jobs": {d: {n for n in _listdir(d) if n.endswith(".json")} for d in jobs},
        "logs": {d: set(_listdir(d)) for d in logs},
    }


def session_names(yaml_path):
    """Session names registered in an agent-sessions.yaml; empty if unreadable."""
    try:
        with open(yaml_path, encoding="utf-8", errors="replace") as fh:
            return {m.group(1) for m in map(_SESSION_LINE.match, fh) if m}
    except OSError:
        return set()


def _bare(session):
    for prefix in SESSION_PREFIXES:
        if session.startswith(prefix):
            return session[len(prefix):]
    return session


def leaks(roots, before, recorded):
    """New entries since `before` that no live delegation explains.

    A real delegation writes its registry record before its audit log and
    targets a registered session (registry.register_job), so a new entry is a
    leak when
      (a) it is an audit log whose job has no record in any watched registry, or
      (b) its record's agent_session is missing or names no session in the
          watched agent-sessions.yaml files.
    Returns [{"path", "reason", "agent_session", "prompt"}] sorted by path."""
    now = snapshot(roots, recorded)
    yamls = [os.path.join(r, "agent-sessions.yaml") for r in roots]
    if recorded.get("AGENT_SESSIONS_YAML"):
        yamls.append(recorded["AGENT_SESSIONS_YAML"])
    known = set()
    for path in yamls:
        known |= session_names(path)
    records = {}
    for d, names in now["jobs"].items():
        for name in names:
            records.setdefault(name[:-len(".json")], os.path.join(d, name))

    found = []

    def judge(path, meta, reason=None):
        session = meta.get("agent_session")
        if reason is None:
            if not session:
                reason = "no agent_session"
            elif _bare(session) not in known:
                reason = "agent_session names no registered session"
            else:
                return
        found.append({"path": path, "reason": reason, "agent_session": session,
                      "prompt": str(meta.get("prompt", ""))[:60]})

    for d, names in now["jobs"].items():
        for name in names - before["jobs"].get(d, set()):
            path = os.path.join(d, name)
            judge(path, _load(path))
    for d, names in now["logs"].items():
        for name in names - before["logs"].get(d, set()):
            path = os.path.join(d, name)
            if name in records:
                judge(path, _load(records[name]))
            else:
                judge(path, _load(os.path.join(path, "meta.json")), "audit log without a registry record")
    return sorted(found, key=lambda f: f["path"])


def describe(found):
    """Failure message for the session guard."""
    lines = ["Issue #20: {} entr{} made during this session in a real control plane "
             "(left in place; check and remove by hand):".format(len(found), "y" if len(found) == 1 else "ies")]
    lines += ["  {path}  [{reason}]  agent_session={agent_session!r}  prompt={prompt!r}".format(**f)
              for f in found]
    return "\n".join(lines)