"""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())