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

"""Control-plane binding of a session row (Issue #21, item 3).

The binding is the four absolute paths a session's tools must agree on:
WORKSPACE_ROOT, AGENT_SESSIONS_YAML, MAM_JOBS_DIR, DELEGATE_JOB_LOGS_DIR.
create/resume persist it on the row (`control_plane`) and forward it to the
pane only when it differs from what `--workspace` alone would imply, through
the leading `K=V` tokens lib.sh's new-session already turns into `--env`.

CLI (used by lib.sh):
  python -m lib_py.control_plane json <workspace> [<row control_plane JSON>]
      -> {4 keys, "forward": bool, "whitespace": bool, "source": "row"|"env"}
  python -m lib_py.control_plane prefix <json>
      -> "WORKSPACE_ROOT=... AGENT_SESSIONS_YAML=... ... " or "" (no forward)
"""

from __future__ import annotations

import json
import os
import shlex
import sys
from typing import Dict, Optional

KEYS = ("workspace_root", "agent_sessions_yaml", "mam_jobs_dir", "delegate_job_logs_dir")
ENV = {"workspace_root": "WORKSPACE_ROOT", "agent_sessions_yaml": "AGENT_SESSIONS_YAML",
       "mam_jobs_dir": "MAM_JOBS_DIR", "delegate_job_logs_dir": "DELEGATE_JOB_LOGS_DIR"}


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


def defaults(workspace: str) -> Dict[str, str]:
    """D(ws): the binding `--workspace` implies on its own."""
    ws = _abs(workspace)
    mam = os.path.join(ws, ".mam")
    return {"workspace_root": ws, "agent_sessions_yaml": os.path.join(mam, "agent-sessions.yaml"),
            "mam_jobs_dir": os.path.join(mam, "jobs"), "delegate_job_logs_dir": os.path.join(mam, "delegate_job_logs")}


def from_env(env=None) -> Dict[str, str]:
    env = os.environ if env is None else env
    b = defaults(env.get("WORKSPACE_ROOT") or os.getcwd())
    for k, var in ENV.items():
        v = (env.get(var) or "").strip()
        if v:
            b[k] = _abs(v)
    return b


def describe(workspace: str, row: Optional[dict] = None) -> dict:
    """The binding to use: the row's persisted one if complete (D11), else the
    caller's env. A row binding that differs from the env is reported on stderr."""
    env_b = from_env()
    source = "env"
    b = env_b
    if isinstance(row, dict) and all(isinstance(row.get(k), str) and row.get(k) for k in KEYS):
        b = {k: row[k] for k in KEYS}
        source = "row"
        if b != env_b:
            print("WARNING: using the control-plane binding recorded in the row "
                  f"({json.dumps(b, sort_keys=True)}) instead of the caller's environment", file=sys.stderr)
    return {**b, "forward": b != defaults(workspace),
            "whitespace": any(any(c.isspace() for c in v) for v in b.values()), "source": source}


def prefix(desc: dict) -> str:
    if not desc.get("forward"):
        return ""
    return "".join(f"{ENV[k]}={shlex.quote(desc[k])} " for k in KEYS)


def main(argv=None) -> int:
    argv = list(sys.argv[1:] if argv is None else argv)
    if len(argv) >= 2 and argv[0] == "json":
        row = None
        if len(argv) >= 3 and argv[2].strip():
            try:
                row = json.loads(argv[2])
            except json.JSONDecodeError:
                row = None
        print(json.dumps(describe(argv[1], row)))
        return 0
    if len(argv) == 2 and argv[0] == "prefix":
        desc = json.loads(argv[1])
        if desc.get("forward") and desc.get("whitespace"):
            return 3  # lib.sh splits --env flags unquoted: refuse rather than break silently
        sys.stdout.write(prefix(desc))
        return 0
    print(__doc__, file=sys.stderr)
    return 2


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