multi-agent-mux/.agents/skills/lib_py/agents/autopin.py

"""Auto-pin of a post-discovered session id on a verified local `completed` (Issue #40).

Class A agents (agy, opencode, codex, hermes) write their conversation id to
disk only after their first turn, so create records the row without one. The
side that verified a job's `completed` - the delegate-job wrapper, create
--onboard, and the loop's wait_for_job as a recovery path - calls
handle_completed(). It captures the id through the adapter
(capture_active_session_id) and pins it with pin_row() under atomic_dump_yaml.
reconcile's drift C calls the same capture() and pin_row(), so the two paths
reach one decision.

One (job_id, dispatch_token, attempt_id) is handled once: a job-scoped flock
and a marker in the job's directory make a second call (wrapper + loop
recovery) a no-op. Remote and human jobs have no local disk or pane and are
skipped; the agent's own publish_event.py never writes the control plane
(#15/#20/#21). A failed pin never changes the job's result: the reason is
written on the row and returned.

CLI: python -m lib_py.agents.autopin --registry-dir DIR --job ID
         [--yaml AGENT_SESSIONS_YAML] [--expect-dispatch-token TOKEN]
"""

import argparse
import fcntl
import json
import os
import subprocess
import sys
import tempfile
import time
from typing import Any, Callable, Dict, List, Optional

from lib_py.agents.base import DiscoveryContext, POST_DISCOVERED, capture_result
from lib_py.agents.registry import _ADAPTERS, agent_of_row, get_adapter

OWN_KEYS = tuple(a.own_key for a in _ADAPTERS.values())
MARKER = "autopin.json"


def claimed_elsewhere(row: Dict[str, Any], state: Dict[str, Any]) -> List[str]:
    """Every id another row records: an id already some session's is never this one's."""
    ids: List[str] = []
    for other in state.get("herdr_sessions") or []:
        if other is row or other.get("name") == row.get("name"):
            continue
        for key in OWN_KEYS:
            v = other.get(key)
            if v and v not in ids:
                ids.append(v)
    return ids


def capture(row: Dict[str, Any], agent: str, state: Dict[str, Any], home_dir: Optional[str] = None) -> Dict[str, Any]:
    """The adapter's decision for this row (read-only). Besides the capture
    statuses: 'pre-assigned' and 'already-set' (nothing to capture), 'skipped'."""
    adapter = get_adapter(agent)
    if adapter is None:
        return capture_result("skipped", reason=f"unknown agent {agent!r}")
    if adapter.id_lifecycle != POST_DISCOVERED:
        return capture_result("pre-assigned")
    if row.get(adapter.own_key):
        return capture_result("already-set", session_id=row[adapter.own_key])
    cwd = (row.get("pane") or {}).get("cwd") or ""
    if not cwd:
        return capture_result("no-candidate", reason="the row has no pane cwd")
    row_eval = dict(row)
    row_eval["_sibling_claimed_uuids"] = claimed_elsewhere(row, state)
    ctx = DiscoveryContext(workspace=cwd, agent_name=agent, home_dir=home_dir,
                           epoch=row.get("herdr_session_epoch") or 0, row=row_eval, mode="discover")
    return adapter.capture_active_session_id(ctx)


def pin_row(row: Dict[str, Any], agent: str, uuid: str) -> None:
    """Record a verified id: own key, source and verified flag, in one change."""
    row[get_adapter(agent).own_key] = uuid
    row["session_id_source"] = "discovered"
    row["session_id_verified"] = True
    row["last_visible_status"] = "pinned"


def _skills_dir() -> str:
    return os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))


def resume_dry_run(skills_dir: str, cwd: str, agent: str, name: str, env: Optional[Dict[str, str]] = None) -> str:
    """last_visible_status after a pin: 'resume_verified' or why the dry-run failed."""
    cmd = ["bash", os.path.join(skills_dir, "multi-agent-mux-resume", "scripts", "resume_session.sh"),
           "--workspace", cwd, "--agent", agent, "--session", name, "--dry-run"]
    res = subprocess.run(cmd, capture_output=True, text=True, env=env)
    if res.returncode == 0:
        return "resume_verified"
    return f"resume dry-run failed: {res.stderr.strip() or res.stdout.strip()}"


def tui_viewport(lib_sh: str, name: str, agent: str, cwd: str, env: Optional[Dict[str, str]] = None) -> int:
    """lib.sh verify_tui_viewport: 0 match, 1 mismatch, other = no evidence."""
    script = 'source "$1" && verify_tui_viewport "$2" "$3" "$4"'
    return subprocess.run(["bash", "-c", script, "_", lib_sh, name, agent, cwd], env=env).returncode


def _load_record(registry_dir: str, job_id: str) -> Optional[Dict[str, Any]]:
    try:
        with open(os.path.join(registry_dir, f"{job_id}.json"), encoding="utf-8") as f:
            rec = json.load(f)
        return rec if isinstance(rec, dict) else None
    except (OSError, ValueError):
        return None


_MUTATION: Optional[Callable[[Dict[str, Any]], bool]] = None


def _atomic_update(yaml_path: str, mutate: Callable[[Dict[str, Any]], bool]) -> None:
    """Run `mutate(d)` as an atomic_dump_yaml mutation in this process; it
    returns False to leave the registry untouched."""
    from lib_py import atomic_yaml
    global _MUTATION
    saved = {k: os.environ.get(k) for k in ("YAML_PATH", "AGENT_SESSIONS_MUTATION")}
    _MUTATION = mutate
    os.environ["YAML_PATH"] = yaml_path
    # this very module (it is __main__ under `python -m`), not a second copy
    os.environ["AGENT_SESSIONS_MUTATION"] = (
        "import sys as _sys\n"
        f"if not _sys.modules[{__name__!r}]._MUTATION(d):\n"
        "    raise SystemExit(0)\n")
    try:
        atomic_yaml.atomic_dump_yaml_main()
    except SystemExit as exc:
        if exc.code not in (0, None):
            raise RuntimeError(str(exc.code))
    finally:
        _MUTATION = None
        for k, v in saved.items():
            if v is None:
                os.environ.pop(k, None)
            else:
                os.environ[k] = v



def _subprocess_env(yaml_path: str, registry_dir: str) -> Dict[str, str]:
    """The control plane the lib.sh checks run against: the registry's own."""
    env = dict(os.environ)
    env.setdefault("AGENT_SESSIONS_YAML", yaml_path)
    env.setdefault("MAM_JOBS_DIR", registry_dir)
    mam = os.path.dirname(yaml_path)
    if os.path.basename(mam) == ".mam":
        env.setdefault("WORKSPACE_ROOT", os.path.dirname(mam))
    return env


def _pin_session(name: str, job_id: str, yaml_path: str, home_dir: Optional[str],
                 viewport: Callable[[str, str, str], int], resume: Callable[[str, str, str], str]) -> Dict[str, Any]:
    out: Dict[str, Any] = {}

    def mutate(d: Dict[str, Any]) -> bool:
        row = next((s for s in d.get("herdr_sessions") or [] if s.get("name") == name), None)
        if row is None:
            out.update(capture_result("skipped", reason=f"no session row {name!r}"))
            return False
        if row.get("status") != "running":
            out.update(capture_result("skipped", reason=f"session {name!r} is {row.get('status')}"))
            return False
        agent = agent_of_row(row)
        res = capture(row, agent or "", d, home_dir=home_dir)
        cwd = (row.get("pane") or {}).get("cwd") or ""
        if res["status"] == "captured":
            rc = viewport(name, agent, cwd)
            if rc == 1:
                res = capture_result("viewport-mismatch", reason="the pane does not show this workspace",
                                     candidates=res["candidates"])
            else:
                pin_row(row, agent, res["session_id"])
                row["last_visible_status"] = resume(name, agent, cwd)
                res = dict(res, status="pinned", degraded=(rc != 0))
                out.update(res)
                return True
        out.update(res)
        if res["status"] in ("pre-assigned", "already-set", "skipped"):
            return False
        row["last_visible_status"] = f"autopin {res['status']} (job {job_id}): {res['reason']}"
        return True

    _atomic_update(yaml_path, mutate)
    return out


def handle_completed(registry_dir: str, job_id: str, yaml_path: Optional[str] = None,
                     expect_token: Optional[str] = None, home_dir: Optional[str] = None,
                     viewport: Optional[Callable[[str, str, str], int]] = None,
                     resume: Optional[Callable[[str, str, str], str]] = None) -> Dict[str, Any]:
    """Pin the job's session id if the job completed locally and the row still
    lacks a post-discovered id. Returns the result (also kept in the marker)."""
    rec = _load_record(registry_dir, job_id)
    if rec is None:
        return capture_result("skipped", reason=f"no job record {job_id}")
    if rec.get("status") != "completed":
        return capture_result("skipped", reason=f"job {job_id} is {rec.get('status')}")
    token = rec.get("dispatch_token") or ""
    if expect_token and token != expect_token:
        return capture_result("skipped", reason="stale dispatch token")
    session = str(rec.get("agent_session") or "")
    if (rec.get("transport") == "remote" or session.startswith("remote:")
            or rec.get("agent") == "human" or session in ("", "human", "herdr:human")):
        return capture_result("skipped", reason="not a local agent session")
    name = session[len("herdr:"):] if session.startswith("herdr:") else session
    yaml_path = yaml_path or os.environ.get("AGENT_SESSIONS_YAML") or os.path.join(
        os.path.dirname(os.path.abspath(registry_dir)), "agent-sessions.yaml")
    if not (os.path.exists(yaml_path) or os.path.exists(os.path.splitext(yaml_path)[0] + ".db")):
        return capture_result("skipped", reason=f"no session registry at {yaml_path}")

    key = f"{job_id}:{token}:{rec.get('attempt_id')}"
    job_dir = os.path.join(registry_dir, job_id)
    os.makedirs(job_dir, exist_ok=True)
    env = _subprocess_env(yaml_path, registry_dir)
    if viewport is None:
        lib_sh = os.environ.get("LIB_SH") or os.path.join(_skills_dir(), "lib.sh")
        viewport = lambda n, a, c: tui_viewport(lib_sh, n, a, c, env=env)  # noqa: E731
    if resume is None:
        resume = lambda n, a, c: resume_dry_run(_skills_dir(), c, a, n, env=env)  # noqa: E731

    with open(os.path.join(job_dir, "autopin.lock"), "w") as lock:
        fcntl.flock(lock, fcntl.LOCK_EX)
        marker = os.path.join(job_dir, MARKER)
        try:
            with open(marker, encoding="utf-8") as f:
                done = json.load(f)
            if done.get("key") == key:
                return dict(done.get("result") or {}, deduplicated=True)
        except (OSError, ValueError):
            pass
        try:
            result = _pin_session(name, job_id, yaml_path, home_dir, viewport, resume)
        except Exception as exc:
            result = capture_result("error", reason=f"{type(exc).__name__}: {exc}")
        result["session"] = name
        fd, tmp = tempfile.mkstemp(dir=job_dir, prefix=".autopin.", suffix=".tmp")
        with os.fdopen(fd, "w", encoding="utf-8") as f:
            json.dump({"key": key, "at": int(time.time()), "result": result}, f)
        os.replace(tmp, marker)
        return result


def main(argv: Optional[List[str]] = None) -> int:
    p = argparse.ArgumentParser(prog="python -m lib_py.agents.autopin")
    p.add_argument("--registry-dir", required=True)
    p.add_argument("--job", required=True)
    p.add_argument("--yaml", default=None, help="agent-sessions.yaml (default: beside the registry)")
    p.add_argument("--expect-dispatch-token", default=None)
    a = p.parse_args(argv)
    res = handle_completed(a.registry_dir, a.job, yaml_path=a.yaml, expect_token=a.expect_dispatch_token)
    if res.get("status") == "pinned":
        print(f"autopin: {res['session']} pinned {res['session_id']} (job {a.job})", file=sys.stderr)
    elif res.get("status") not in ("skipped", "pre-assigned", "already-set") and not res.get("deduplicated"):
        print(f"autopin: {res.get('session')} not pinned - {res['status']}: {res.get('reason')} (job {a.job})",
              file=sys.stderr)
    return 0


if __name__ == "__main__":
    sys.exit(main())