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

"""Same-cwd creation/resume ordering for post-discovered agents (Issue #63).

While a post-discovered agent (agy, codex, hermes, opencode) in a workspace is
 still being created (marker held by a live create/resume process) or is running
without a recorded conversation id (live pane pid), a second create or resume of
the same agent type in that workspace waits until the first records its id, up
to MAM_ID_ORDER_WAIT_S (default 600 s; 0 = do not wait). A marker whose owner
died before the row was registered keeps blocking until herdr answers that its
session is gone (the pane may still discover the workspace id).
"""
from __future__ import annotations

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

from lib_py import ints
from lib_py.agents.autopin import claimed_elsewhere
from lib_py.agents.base import POST_DISCOVERED
from lib_py.agents.registry import agent_of_row, get_adapter
from lib_py.agents.sanitize import sanitize_herdr_agent_name
from lib_py.agents.session_ids import canonical_cwd, claiming_rows, own_id
from lib_py.state import load_state_json

DEFAULT_WAIT_S = 600
PROBE_TIMEOUT_S = 30.0


def wait_s(env: Mapping[str, str] = os.environ) -> int:
    raw = env.get("MAM_ID_ORDER_WAIT_S")
    if raw is None:
        return DEFAULT_WAIT_S
    n = ints.positive_int(raw, 0)
    return DEFAULT_WAIT_S if n is None else n


def pid_alive(pid: int) -> bool:
    if not isinstance(pid, int) or pid <= 0:
        return False
    try:
        os.kill(pid, 0)
        return True
    except ProcessLookupError:
        return False
    except PermissionError:
        return True
    except OSError:
        return False


def herdr_session_state(server: str, session: str, timeout: float) -> Optional[bool]:
    """Is `session` live on herdr server `server`: True, False, or None when
    herdr gave no authoritative answer. False only from a successful, parsed
    `agent list` without the name; a failed (server down included), timed-out
    or unparseable lookup, or no time left, is None. Never starts a server
    (MAM_NO_AUTOSTART)."""
    if timeout <= 0:
        return None
    skills_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
    env = dict(os.environ, MAM_NO_AUTOSTART="1", HERDR_SESSION_NAME=server or "default")
    script = 'source "$1" && _herdr agent list'
    try:
        res = subprocess.run(
            ["bash", "-c", script, "_", os.path.join(skills_dir, "lib.sh")],
            env=env, stdin=subprocess.DEVNULL, capture_output=True, text=True,
            timeout=min(timeout, PROBE_TIMEOUT_S),
        )
    except (OSError, subprocess.SubprocessError):
        return None
    if res.returncode != 0:
        return None
    try:
        agents = json.loads(res.stdout)["result"]["agents"]
    except (ValueError, KeyError, TypeError):
        return None
    if not isinstance(agents, list):
        return None
    names = {session, sanitize_herdr_agent_name(session)}
    return any(isinstance(a, dict) and a.get("name") in names for a in agents)


def order_dir(cwd: str, agent: str, env: Mapping[str, str] = os.environ,
              home: Optional[str] = None) -> str:
    h = home or env.get("HOME") or env.get("HOME_DIR") or os.path.expanduser("~")
    real_cwd = canonical_cwd(cwd or ".")
    digest = hashlib.sha256(real_cwd.encode("utf-8")).hexdigest()[:16]
    return os.path.join(h, ".local", "state", "mam", "id-order", f"{digest}-{agent}")


def _parse_pos_int(val: Any) -> Optional[int]:
    if isinstance(val, bool):
        return None
    if isinstance(val, int):
        return val if val > 0 else None
    s = str(val or "").strip()
    return int(s) if s.isdigit() and int(s) > 0 else None


def blockers(state: Dict[str, Any], cwd: str, agent: str, me: str,
             markers: List[Dict[str, Any]],
             is_alive: Callable[[int], bool] = pid_alive) -> List[str]:
    target_cwd = canonical_cwd(cwd)
    found: List[str] = []
    rows_by_name: Dict[str, Dict[str, Any]] = {}

    for s in state.get("herdr_sessions") or []:
        if not isinstance(s, dict):
            continue
        rname = str(s.get("name") or "")
        if rname:
            rows_by_name[rname] = s
        if not rname or rname == me:
            continue
        s_cwd = (s.get("pane") or {}).get("cwd") or ""
        if not s_cwd or canonical_cwd(s_cwd) != target_cwd:
            continue
        if agent_of_row(s) != agent or s.get("status") != "running":
            continue
        if own_id(s, agent):
            continue
        pane_pid = _parse_pos_int((s.get("pane") or {}).get("pid"))
        if pane_pid is None or is_alive(pane_pid):
            if rname not in found:
                found.append(rname)

    for m in markers or []:
        if not isinstance(m, dict):
            continue
        m_sess = str(m.get("session") or "")
        if not m_sess or m_sess in found:
            continue
        m_cwd = str(m.get("cwd") or "")
        if m_cwd and canonical_cwd(m_cwd) != target_cwd:
            continue
        m_agent = str(m.get("agent") or "")
        if m_agent and m_agent != agent:
            continue
        row = rows_by_name.get(m_sess)
        if row is None:
            found.append(m_sess)
            continue
        status = str(row.get("status") or "")
        cur_own = own_id(row, agent)
        row_epoch = int(row.get("herdr_session_epoch") or 0)
        end_epoch = int(row.get("stopped_at_epoch") or row.get("terminated_at_epoch") or 0)
        m_at = int(m.get("at") or 0)
        prev_own = str(m.get("prev_own_id") or "")
        prev_epoch = int(m.get("prev_epoch") or 0)
        prev_end_epoch = int(m.get("prev_end_epoch") or 0)

        pinned_for_attempt = bool(cur_own) and (
            cur_own != prev_own
            or (status == "running" and row_epoch > prev_epoch)
        )
        if pinned_for_attempt:
            continue
        if status == "running":
            found.append(m_sess)
        else:
            ended_after_attempt = (
                end_epoch > max(m_at, prev_end_epoch)
                and (prev_epoch == 0 or row_epoch > prev_epoch)
            )
            if not ended_after_attempt:
                found.append(m_sess)

    return found


def _registered_for_attempt(m: Dict[str, Any], row: Optional[Dict[str, Any]]) -> bool:
    if not isinstance(row, dict) or row.get("status") != "running":
        return False
    return int(row.get("herdr_session_epoch") or 0) > int(m.get("prev_epoch") or 0)


def _read_and_prune_markers(dpath: str, me: str, my_pid: int, state: Dict[str, Any],
                            is_alive: Callable[[int], bool] = pid_alive,
                            pane_state: Callable[[str, str], Optional[bool]] = (
                                lambda server, session: None),
                            ) -> List[Dict[str, Any]]:
    """Markers that still count. A live owner keeps its marker (another live
    attempt on `me` included). A dead owner's marker goes once its row is
    registered for the attempt (the row decides from then on) or herdr answers
    that its session is gone; until then it is kept as an orphan."""
    rows_by_name = {
        str(s.get("name") or ""): s
        for s in (state.get("herdr_sessions") or []) if isinstance(s, dict)
    }
    out: List[Dict[str, Any]] = []
    try:
        names = sorted(os.listdir(dpath))
    except OSError:
        return out
    for fname in names:
        if not fname.endswith(".json") or fname.startswith("."):
            continue
        fpath = os.path.join(dpath, fname)
        try:
            with open(fpath, encoding="utf-8") as f:
                data = json.load(f)
        except (OSError, ValueError):
            try:
                os.remove(fpath)
            except OSError:
                pass
            continue
        if not isinstance(data, dict):
            continue
        owner_pid = _parse_pos_int(data.get("owner_pid"))
        sess = str(data.get("session") or "")
        if owner_pid is None or owner_pid == my_pid:
            drop = True
        elif is_alive(owner_pid):
            drop = False
        elif sess == me or not sess:
            drop = True
        elif _registered_for_attempt(data, rows_by_name.get(sess)):
            drop = True
        else:
            live = pane_state(str(data.get("herdr_server") or ""), sess)
            drop = live is False
            data["_orphan"] = "alive" if live else "unknown"
        if drop:
            try:
                os.remove(fpath)
            except OSError:
                pass
            continue
        out.append(data)
    return out


def write_marker(dpath: str, session: str, owner_pid: int, cwd: str, agent: str,
                 prev_row: Optional[Dict[str, Any]] = None, herdr_server: str = "") -> str:
    os.makedirs(dpath, exist_ok=True)
    marker_path = os.path.join(dpath, f"{session}.json")
    payload: Dict[str, Any] = {
        "session": session,
        "owner_pid": int(owner_pid),
        "cwd": canonical_cwd(cwd),
        "agent": agent,
        "herdr_server": herdr_server,
        "at": int(time.time()),
    }
    if isinstance(prev_row, dict):
        payload["prev_own_id"] = own_id(prev_row, agent)
        payload["prev_epoch"] = int(prev_row.get("herdr_session_epoch") or 0)
        payload["prev_end_epoch"] = int(
            prev_row.get("stopped_at_epoch") or prev_row.get("terminated_at_epoch") or 0
        )
    fd, tmp = tempfile.mkstemp(dir=dpath, prefix=".marker.", suffix=".tmp")
    with os.fdopen(fd, "w", encoding="utf-8") as f:
        json.dump(payload, f)
    os.replace(tmp, marker_path)
    return marker_path


def wait(cwd: str, agent: str, session: str, owner_pid: int,
         uuid: str = "", yaml_path: Optional[str] = None,
         home: Optional[str] = None, wait_s_val: Optional[int] = None,
         poll: float = 0.25,
         is_alive: Callable[[int], bool] = pid_alive,
         herdr_server: str = "",
         pane_state: Optional[Callable[[str, str, float], Optional[bool]]] = None) -> int:
    adapter = get_adapter(agent)
    if adapter is None or adapter.id_lifecycle != POST_DISCOVERED:
        return 0

    cap = wait_s(os.environ) if wait_s_val is None else int(wait_s_val)
    dpath = order_dir(cwd, agent, os.environ, home=home)
    os.makedirs(dpath, exist_ok=True)
    lock_file = os.path.join(dpath, "lock")
    deadline = time.monotonic() + cap
    announced = False
    row_ref = {"name": session}
    probe = pane_state or herdr_session_state

    answers: Dict[str, bool] = {}  # last definite answer per session, for the message

    def bounded_probe(server: str, sess: str) -> Optional[bool]:
        live = probe(server, sess, deadline - time.monotonic())
        if live is not None:
            answers[sess] = live
        return live

    while True:
        with open(lock_file, "a", encoding="utf-8") as lf:
            fcntl.flock(lf, fcntl.LOCK_EX)
            state = load_state_json(yaml_path)
            if uuid:
                claimed = set(claimed_elsewhere(row_ref, state))
                if uuid in claimed:
                    others = ", ".join(claiming_rows(row_ref, uuid, state))
                    print(
                        f"ERROR: {session}'s recorded id {uuid} is also recorded by {others}; "
                        f"refusing to resume on a conversation another session claims (Issue #63).",
                        file=sys.stderr, flush=True,
                    )
                    return 1
            markers = _read_and_prune_markers(dpath, session, owner_pid, state,
                                              is_alive=is_alive, pane_state=bounded_probe)
            blks = blockers(state, cwd, agent, session, markers, is_alive=is_alive)
            if not blks:
                prev_row = next(
                    (
                        s for s in (state.get("herdr_sessions") or [])
                        if isinstance(s, dict) and s.get("name") == session
                    ),
                    None,
                )
                already_recorded = bool(uuid) and own_id(prev_row, agent) == uuid
                if not already_recorded and owner_pid and owner_pid > 0:
                    write_marker(dpath, session, owner_pid, cwd, agent, prev_row=prev_row,
                                 herdr_server=herdr_server)
                return 0

        names = ", ".join(blks)
        if time.monotonic() >= deadline:
            print(
                f"ERROR: {names} in {cwd} still has no recorded {agent} id after {cap}s; "
                f"not starting {session} (Issue #63). Wait for its first turn to finish, "
                f"or stop it with multi-agent-mux-stop.",
                file=sys.stderr, flush=True,
            )
            for m in markers:
                if m.get("_orphan") and m.get("session") in blks:
                    srv = m.get("herdr_server") or "default"
                    alive = m["_orphan"] == "alive" or answers.get(m["session"]) is True
                    seen = ("its herdr session is alive" if alive
                            else "herdr could not tell whether its session is gone")
                    if not alive and cap == 0:
                        seen += " (MAM_ID_ORDER_WAIT_S=0 leaves no time to ask herdr)"
                    print(
                        f"ERROR: {m['session']}'s create/resume exited before recording it, "
                        f"and {seen}; close it with: "
                        f"HERDR_SESSION_NAME={srv} herdr kill-session -t {m['session']}",
                        file=sys.stderr, flush=True,
                    )
            return 1
        if not announced:
            print(
                f"id-order: waiting for {names} ({agent}, no recorded id yet) in {cwd} "
                f"- up to {cap}s (MAM_ID_ORDER_WAIT_S)",
                file=sys.stderr, flush=True,
            )
            announced = True
        remaining = deadline - time.monotonic()
        if remaining > 0:
            time.sleep(min(poll, remaining))


def main(argv: Optional[List[str]] = None) -> int:
    p = argparse.ArgumentParser(prog="python -m lib_py.agents.id_order")
    sub = p.add_subparsers(dest="cmd", required=True)
    s_wait = sub.add_parser("wait")
    s_wait.add_argument("--cwd", required=True)
    s_wait.add_argument("--agent", required=True)
    s_wait.add_argument("--session", required=True)
    s_wait.add_argument("--owner-pid", type=int, default=0)
    s_wait.add_argument("--uuid", default="")
    s_wait.add_argument("--herdr-server", default="")
    a = p.parse_args(argv)
    return wait(
        cwd=a.cwd,
        agent=a.agent,
        session=a.session,
        owner_pid=a.owner_pid,
        uuid=a.uuid or "",
        herdr_server=a.herdr_server or "",
    )


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