"""Neutral okstra team CLI for pane-backed worker dispatch. Under cmux this is every lead's door onto cmux surfaces, because okstra owns the panes there rather than the host. Outside cmux a worker owns no pane at all — it runs as a cli-wrapper subprocess — so there is no pane for either closing command to act on. Which backend a run uses is read from its run manifest. `reclaim` is the round boundary and `teardown` is the end of the run: the first closes nothing but the finished dispatches' panes, the second takes every recorded pane and writes off whatever never finished. """ from __future__ import annotations import argparse import json import subprocess import sys from pathlib import Path from typing import Any, Mapping, Sequence from . import cmux from .adapters.dispatch import provider_worker_wrappers from .adapters.dispatch.cmux import dispatch_port_for_terminal_backend from .application.dispatch_assignments import dispatch_assignments from .dispatch_state import ( TEARDOWN_BEFORE_TERMINAL_REASON, TERMINAL_WORKER_STATUSES, mutate_team_state, ) from .dispatch_core import ( BACKEND_CLI_WRAPPER, BACKEND_CMUX_PANE, DispatchError, DispatchPlan, await_dispatches, dispatch_plan, ) from .ports.worker_dispatch import WorkerDispatchRequest from .registry.host_registry import default_host_registry from .registry.provider_registry import default_provider_registry from .session import observe_lead_session from .json_boundary import load_owned_object _SUPPORTED_WRAPPERS = provider_worker_wrappers(default_provider_registry()) def main(argv: Sequence[str] | None = None) -> int: parser = _parser() args = parser.parse_args(argv) try: if args.command == "dispatch": return _dispatch(args) if args.command == "await": return _await(args) if args.command == "teardown": return _teardown(args) if args.command == "reclaim": return _reclaim(args) except DispatchError as exc: print(f"okstra team: {exc}", file=sys.stderr) return 2 parser.print_help() return 2 def _parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(prog="okstra team") parser.add_argument("--workspace-root", required=True) parser.add_argument("--okstra-bin", required=True) sub = parser.add_subparsers(dest="command", required=True) _add_dispatch_parser(sub) _add_await_parser(sub) _add_teardown_parser(sub) _add_reclaim_parser(sub) return parser def _add_dispatch_parser(sub) -> None: parser = sub.add_parser("dispatch", help="dispatch pane-backed workers") _add_run_args(parser) parser.add_argument("--workers", default="") parser.add_argument("--jobs-file", default="") parser.add_argument("--dry-run", action="store_true") parser.add_argument("--idle-timeout-seconds", type=int, default=None) parser.add_argument("--dispatch-kind", default="initial") def _add_await_parser(sub) -> None: parser = sub.add_parser("await", help="wait for pane-backed workers") _add_run_args(parser) parser.add_argument("--poll-interval-seconds", type=int, default=5) parser.add_argument("--timeout-seconds", type=int, default=None) parser.add_argument("--heartbeat-seconds", type=int, default=30) parser.add_argument("--json", action="store_true") def _add_teardown_parser(sub) -> None: parser = sub.add_parser( "teardown", help="close every recorded pane at the end of the run" ) _add_run_args(parser) parser.add_argument("--dry-run", action="store_true") parser.add_argument("--json", action="store_true") def _add_reclaim_parser(sub) -> None: parser = sub.add_parser( "reclaim", help="close the finished dispatches' panes at a round boundary", ) _add_run_args(parser) parser.add_argument("--dry-run", action="store_true") parser.add_argument("--json", action="store_true") def _add_run_args(parser) -> None: parser.add_argument("--project-root", required=True) parser.add_argument("--run-manifest", required=True) def _dispatch(args) -> int: if args.jobs_file and args.workers: raise DispatchError("--jobs-file and --workers cannot be combined") manifest = _load_manifest(args.project_root, args.run_manifest) _observe_lead_session_from_manifest(Path(args.project_root).resolve(), manifest) terminal_backend = _manifest_backend(manifest) lead_runtime = _require_string(manifest, "leadRuntime") fallback_port = default_host_registry().resolve( lead_runtime ).worker_dispatch() port = dispatch_port_for_terminal_backend( terminal_backend, fallback_port, supported_worker_wrappers=_SUPPORTED_WRAPPERS, unsupported_worker_label=f"{lead_runtime} lead", required_lead_runtime=None if _is_cmux_run(manifest) else lead_runtime, ) request = WorkerDispatchRequest( project_root=Path(args.project_root), run_manifest=Path(args.run_manifest), workspace_root=Path(args.workspace_root), okstra_bin=Path(args.okstra_bin), requested_workers=tuple(_parse_workers(args.workers)), idle_timeout_seconds=args.idle_timeout_seconds, dispatch_kind=args.dispatch_kind, jobs_file=Path(args.jobs_file) if args.jobs_file else None, ) plan = dispatch_assignments(request, port) backend_plan = plan.backend_plan if backend_plan is None: raise DispatchError("team dispatch did not produce a backend plan") if args.dry_run: _print_json(backend_plan.to_payload(dry_run=True)) return 0 result = dispatch_plan(backend_plan, wait=False) _print_json(backend_plan.to_payload(dry_run=False)) return result def _await(args) -> int: manifest = _load_manifest(args.project_root, args.run_manifest) _validate_team_manifest(manifest) _observe_lead_session_from_manifest(Path(args.project_root).resolve(), manifest) plan = _plan_for_existing(Path(args.project_root), Path(args.workspace_root), Path(args.run_manifest), manifest) code = await_dispatches( plan, poll_interval_seconds=args.poll_interval_seconds, timeout_seconds=args.timeout_seconds, heartbeat_seconds=0 if args.json else args.heartbeat_seconds, ) if args.json: _print_json(_await_payload(plan, code == 0)) else: print("ALL_WORKERS_DONE" if code == 0 else "POLL_TIMEOUT") return code def _teardown(args) -> int: manifest = _load_manifest(args.project_root, args.run_manifest) _validate_team_manifest(manifest) project_root = Path(args.project_root).resolve() team_state_path = _resolve_project_path(project_root, _require_string(manifest, "teamStatePath")) team_state = _load_json(team_state_path, "team-state") panes = _reclaimable_panes(manifest, team_state) if args.dry_run: _emit_panes(args.json, panes) return 0 _close_panes(manifest, panes) _mark_teardown_errors(team_state_path) _emit_panes(args.json, panes) return 0 def _reclaim(args) -> int: """Close the finished dispatches' surfaces at a round boundary. Two things teardown does are wrong here. It closes every recorded surface, which mid-round would kill the workers still running; and it writes off every non-terminal dispatch as an error, which would drop a worker out of the retry path it has not reached yet. So this shares the closing and the reporting and nothing else. """ manifest = _load_manifest(args.project_root, args.run_manifest) _validate_team_manifest(manifest) project_root = Path(args.project_root).resolve() team_state_path = _resolve_project_path( project_root, _require_string(manifest, "teamStatePath") ) team_state = _load_json(team_state_path, "team-state") panes = _reclaimable_panes(manifest, team_state, finished_only=True) if args.dry_run: _emit_panes(args.json, panes) return 0 _close_panes(manifest, panes) _emit_panes(args.json, panes) return 0 def _close_panes(manifest: Mapping[str, Any], panes: list[dict[str, str]]) -> None: """Close the given surfaces, then give the lead its width back. The width recovery belongs here rather than to teardown alone. cmux hands the freed width to a neighbour it picks, and that neighbour is not always the lead — so a round boundary that closed panes and stopped there leaves the lead squeezed for exactly the stretch the user spends reading it. """ from .adapters.runtime.assembly import port_for, runtime_chain from .domain.worker_runtime import RuntimeHandle, SURFACE_CMUX_PANE chain = runtime_chain(_manifest_backend(manifest)) # cli-wrapper chains have no cmux port; leftover paneIds must not KeyError. if panes and any(port.surface == SURFACE_CMUX_PANE for port in chain): closer = port_for(chain, SURFACE_CMUX_PANE) for pane in panes: closer.close( RuntimeHandle(SURFACE_CMUX_PANE, pane["paneId"], False) ) try: chain[0].restore_lead() except (OSError, subprocess.SubprocessError, RuntimeError) as exc: print(f"okstra team: could not restore the lead's width: {exc}", file=sys.stderr) def _observe_lead_session_from_manifest( project_root: Path, manifest: Mapping[str, Any] ) -> None: """Collect this run's lead session generations at a phase boundary. Resume and compaction split a lead across several session files, so the one id written at prepare time can name a session holding no record of the run. Dispatch and await are the boundaries a lead cannot route around, which is why the observation hangs off them instead of a call the lead has to make. `--dry-run` observes too, because the hook sits ahead of the dry-run return on purpose: moving it behind would lose the generation a lead was on when a dispatch failed. What gets appended is not a rehearsal either — that lead really did cross this boundary, whatever the dispatch went on to do. A manifest missing `teamStatePath` leaves nothing to observe into; that is the command's own failure to report, not this side effect's. """ try: team_state_path = _resolve_project_path( project_root, _require_string(manifest, "teamStatePath") ) except DispatchError: return observe_lead_session(project_root, team_state_path) def _plan_for_existing( project_root: Path, workspace_root: Path, run_manifest_path: Path, manifest: Mapping[str, Any] ) -> DispatchPlan: root = project_root.resolve() manifest_path = _resolve_project_path(root, str(run_manifest_path)) return DispatchPlan( project_root=root, workspace_root=workspace_root.resolve(), manifest_path=manifest_path, team_state_path=_resolve_project_path(root, _require_string(manifest, "teamStatePath")), lead_events_path=_resolve_project_path(root, _require_string(manifest, "leadEventsPath")), manifest=manifest, jobs=(), default_backend=_manifest_backend(manifest), ) def _await_payload(plan: DispatchPlan, completed: bool) -> dict[str, Any]: team_state = _load_json(plan.team_state_path, "team-state") dispatches = [d for d in team_state.get("workerDispatches", []) if isinstance(d, dict)] timed_out = [d for d in dispatches if d.get("status") == "timeout"] return { "completed": completed, "timedOut": len(timed_out), "workers": dispatches, "teamStatePath": str(plan.team_state_path), } def _reclaimable_panes( manifest: Mapping[str, Any], team_state: Mapping[str, Any], *, finished_only: bool = False, ) -> list[dict[str, str]]: """Everything this run owns and may close. `finished_only` is what separates a round boundary from the end of the run. Teardown closes every recorded surface because no dispatch is expected to continue past it. Mid-round the live workers' surfaces must survive, so reclaim asks for the finished ones only — closing an `in-progress` surface kills that worker and the round has no result to show for it. The recorded ids are the only candidates. There is no per-pane tag API to sweep with, and scanning by title would be worse than nothing: cmux labels its own agent surfaces with the same glyph the harness uses for a teammate pane, so a sweep could close the lead. Only surfaces okstra created are recorded, so only those can be closed. A cli-wrapper run records no surface at all, which is why it reclaims nothing. """ seen: set[str] = set() panes: list[dict[str, str]] = [] for record in team_state.get("workerDispatches", []): if not isinstance(record, dict): continue if finished_only and record.get("status") not in TERMINAL_WORKER_STATUSES: continue _append_pane(panes, seen, str(record.get("paneId", "")), "worker") if _is_cmux_run(manifest): return _still_open_surfaces(panes) return panes def _still_open_surfaces(panes: list[dict[str, str]]) -> list[dict[str, str]]: """The recorded surfaces cmux still shows. Nothing prunes `workerDispatches` — a repeat of one `dispatchId` replaces that row in place, and no row is ever dropped — so a surface closed at an earlier round boundary stays recorded for the rest of the run. Taking the ledger as the residual set makes the run-end cleanup gate offer to close panes that left the screen rounds ago, on a workspace holding none. An unreachable cmux keeps the ledger rather than reporting an empty set: a wedged app is not evidence that the surfaces are gone, closing one that already went away is a silent no-op, and skipping a live one strands it on the user's screen for the rest of the session. """ workspace = cmux.resolve_lead_workspace() if not workspace: return panes try: open_ids = cmux.open_surface_ids(workspace) except (RuntimeError, OSError, subprocess.SubprocessError): return panes return [pane for pane in panes if pane["paneId"] in open_ids] def _append_pane(panes: list[dict[str, str]], seen: set[str], pane_id: str, kind: str) -> None: if pane_id and pane_id not in seen: panes.append({"paneId": pane_id, "kind": kind}) seen.add(pane_id) def _mark_teardown_errors(team_state_path: Path) -> None: def mark(payload: dict[str, Any]) -> bool: changed = False for record in payload.get("workerDispatches", []): if isinstance(record, dict) and record.get("status") not in TERMINAL_WORKER_STATUSES: record["status"] = "error" record["reason"] = TEARDOWN_BEFORE_TERMINAL_REASON changed = True return changed mutate_team_state(team_state_path, mark) def _emit_panes(as_json: bool, panes: list[dict[str, str]]) -> None: if as_json: _print_json({"panes": panes}) return for pane in panes: print(f"{pane['paneId']}\t{pane['kind']}") def _manifest_backend(manifest: Mapping[str, Any]) -> str: """The backend prepare recorded for this run. Deliberately not a flag on this command: the manifest already answers it, and a flag would be a second answer free to disagree. A manifest with no recorded backend reads as cli-wrapper, the pane-less path. """ return str(manifest.get("terminalBackend") or "") or BACKEND_CLI_WRAPPER def _is_cmux_run(manifest: Mapping[str, Any]) -> bool: return _manifest_backend(manifest) == BACKEND_CMUX_PANE def _validate_team_manifest(manifest: Mapping[str, Any]) -> None: if _is_cmux_run(manifest): return runtime = _require_string(manifest, "leadRuntime") descriptor = default_host_registry().resolve(runtime).descriptor if descriptor.launch_mode == "team": return raise DispatchError( f"leadRuntime={descriptor.id} does not use the okstra team lifecycle" ) def _load_manifest(project_root: str, run_manifest: str) -> dict[str, Any]: root = Path(project_root).resolve() path = _resolve_project_path(root, run_manifest) return _load_json(path, "run manifest") def _load_json(path: Path, label: str) -> dict[str, Any]: if not path.is_file(): raise DispatchError(f"{label} not found: {path}") payload = load_owned_object(path, artifact=label) if not isinstance(payload, dict): raise DispatchError(f"{label} must be a JSON object: {path}") return payload def _resolve_project_path(project_root: Path, value: str | Path) -> Path: path = Path(value) return path if path.is_absolute() else project_root / path def _require_string(payload: Mapping[str, Any], key: str) -> str: value = payload.get(key) if not isinstance(value, str) or not value.strip(): raise DispatchError(f"required string missing: {key}") return value.strip() def _parse_workers(raw: str) -> list[str]: return [item.strip() for item in raw.split(",") if item.strip()] def _print_json(payload: Mapping[str, Any]) -> None: print(json.dumps(payload, ensure_ascii=False, indent=2)) if __name__ == "__main__": raise SystemExit(main(sys.argv[1:]))