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

"""Decisions behind the worktree helpers (Issue #21 phase 2, item 4).

worktree_prepare.sh / worktree_remove.sh keep the orchestration in shell and
ask this module the questions that need parsing: is a path a linked worktree
of the control root (and is it locked), which session rows ran under it, which
of their jobs are still open, and whether its submodules hold anything a
forced removal would destroy. Standard library only.

CLI (JSON on stdout, rc 0; the session state comes from MAM_STATE_JSON):
  python -m lib_py.worktree is-linked <main> <path>
  python -m lib_py.worktree is-locked <main> <path>
  python -m lib_py.worktree rows <path>
  python -m lib_py.worktree open-jobs <path> [--exclude id,id]
  python -m lib_py.worktree submodules <worktree>
"""

from __future__ import annotations

import glob
import json
import os
import subprocess
import sys
from typing import Dict, Iterable, List, Optional

OPEN_STATUSES = ("pending", "running")


def _git(cwd: str, *args: str) -> subprocess.CompletedProcess:
    return subprocess.run(["git", "-C", str(cwd), *args], capture_output=True, text=True)


def _real(p: str) -> str:
    return os.path.realpath(p)


def worktree_blocks(main: str) -> List[Dict[str, object]]:
    """`git worktree list --porcelain`, one dict per worktree."""
    res = _git(main, "worktree", "list", "--porcelain")
    blocks: List[Dict[str, object]] = []
    cur: Dict[str, object] = {}
    for line in res.stdout.splitlines():
        if not line.strip():
            if cur:
                blocks.append(cur)
                cur = {}
            continue
        key, _, val = line.partition(" ")
        cur[key] = val if val else True
    if cur:
        blocks.append(cur)
    return blocks


def _block(main: str, path: str) -> Optional[Dict[str, object]]:
    rp, rm = _real(path), _real(main)
    if rp == rm:
        return None
    for b in worktree_blocks(main):
        w = b.get("worktree")
        if isinstance(w, str) and _real(w) == rp:
            return b
    return None


def is_linked_worktree(main: str, path: str) -> bool:
    return _block(main, path) is not None


def is_locked(main: str, path: str) -> bool:
    b = _block(main, path)
    return bool(b and "locked" in b)


def rows_under(state: dict, path: str) -> List[dict]:
    """Session rows whose pane.cwd is `path` or below it."""
    rp = _real(path).rstrip(os.sep)
    out = []
    for s in state.get("herdr_sessions") or []:
        cwd = (s.get("pane") or {}).get("cwd") or ""
        if not cwd:
            continue
        rc = _real(cwd)
        if rc == rp or rc.startswith(rp + os.sep):
            out.append(s)
    return out


def registry_dirs(rows: Iterable[dict], env=None) -> List[str]:
    env = os.environ if env is None else env
    cands = [env.get("MAM_JOBS_DIR") or ""]
    if env.get("WORKSPACE_ROOT"):
        cands.append(os.path.join(env["WORKSPACE_ROOT"], ".mam", "jobs"))
    cands += [((s.get("control_plane") or {}).get("mam_jobs_dir") or "") for s in rows]
    seen, out = set(), []
    for c in cands:
        if c and os.path.isdir(c) and _real(c) not in seen:
            seen.add(_real(c))
            out.append(c)
    return out


def open_jobs(rows: List[dict], dirs: Iterable[str], exclude_ids: Iterable[str] = ()) -> List[dict]:
    """Pending/running jobs addressed to one of the rows (herdr:<name>)."""
    names = {f"herdr:{s.get('name')}" for s in rows if s.get("name")}
    skip = {j for j in exclude_ids if j}
    out, seen = [], set()
    for d in dirs:
        for p in sorted(glob.glob(os.path.join(d, "*.json"))):
            try:
                with open(p, encoding="utf-8") as fh:
                    rec = json.load(fh)
            except (OSError, ValueError):
                continue
            jid = rec.get("job_id") if isinstance(rec, dict) else None
            if not jid or jid in skip or jid in seen:
                continue
            if rec.get("agent_session") in names and rec.get("status") in OPEN_STATUSES:
                seen.add(jid)
                out.append({"job_id": jid, "agent_session": rec.get("agent_session"),
                            "status": rec.get("status"), "registry_dir": d})
    return out


def submodule_findings(wt: str) -> Dict[str, object]:
    """trace: the worktree git dir holds modules/ (git then refuses a plain
    `worktree remove`, even after a failed init); dirty / unpushed / stashes:
    what a forced removal would destroy - it deletes the worktree's own
    submodule repositories. `HEAD` is part of the unpushed check: a submodule
    is on a detached HEAD after `submodule update`, so commits made there are
    on no branch and `--branches --not --remotes` alone misses them. Stashes
    live in refs/stash (and its reflog), reachable from neither."""
    gd = _git(wt, "rev-parse", "--absolute-git-dir")
    if gd.returncode != 0 or not gd.stdout.strip():
        return {"trace": False, "dirty": "", "unpushed": "", "stashes": "",
                "error": gd.stderr.strip() or "rev-parse failed"}
    res: Dict[str, object] = {"trace": os.path.isdir(os.path.join(gd.stdout.strip(), "modules")),
                              "dirty": "", "unpushed": "", "stashes": "", "error": ""}
    if not res["trace"]:
        return res
    checks = (("dirty", "git status --porcelain"),
              ("unpushed", "git log --oneline HEAD --branches --not --remotes"),
              ("stashes", "git stash list"))
    errors = []
    for key, cmd in checks:
        r = _git(wt, "submodule", "foreach", "--recursive", "--quiet", cmd)
        if r.returncode != 0:
            errors.append(r.stderr.strip() or f"{key} check failed (rc {r.returncode})")
        else:
            res[key] = r.stdout.strip()
    res["error"] = "; ".join(errors)
    return res


def _state() -> dict:
    try:
        return json.loads(os.environ.get("MAM_STATE_JSON") or "{}") or {}
    except ValueError:
        return {}


def main(argv=None) -> int:
    argv = list(sys.argv[1:] if argv is None else argv)
    cmd = argv[0] if argv else ""
    if cmd == "is-linked" and len(argv) == 3:
        print(json.dumps(is_linked_worktree(argv[1], argv[2])))
    elif cmd == "is-locked" and len(argv) == 3:
        print(json.dumps(is_locked(argv[1], argv[2])))
    elif cmd == "rows" and len(argv) == 2:
        print(json.dumps(rows_under(_state(), argv[1])))
    elif cmd == "open-jobs" and len(argv) in (2, 4) and (len(argv) == 2 or argv[2] == "--exclude"):
        rows = rows_under(_state(), argv[1])
        exclude = argv[3].replace(",", " ").split() if len(argv) == 4 else []
        print(json.dumps(open_jobs(rows, registry_dirs(rows), exclude)))
    elif cmd == "submodules" and len(argv) == 2:
        print(json.dumps(submodule_findings(argv[1])))
    else:
        print(__doc__, file=sys.stderr)
        return 2
    return 0


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