okstra 0.159.0 → 0.161.0

package/runtime/python/okstra_ctl/cmux.py

@@ -0,0 +1,531 @@
+"""cmux command helpers for okstra-owned worker surfaces.
+
+The tmux sibling of this module is `tmux.py`. cmux is a GUI app, so there is no
+detached-server equivalent here — every surface okstra creates lands in the
+workspace the user is already looking at.
+"""
+from __future__ import annotations
+
+import json
+import os
+import shlex
+import shutil
+import subprocess
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Sequence
+
+PING_OK = "PONG"
+
+# cmux answers over a local unix socket, so a probe that has not returned in a
+# few seconds means the app is wedged or gone — not that it is being slow.
+PROBE_TIMEOUT_SECONDS = 5
+
+# A login shell sources the user's whole rc chain, which can be slow on a
+# developer machine, so this gets far more room than a socket probe.
+LOGIN_SHELL_TIMEOUT_SECONDS = 15
+
+# cmux plants a per-surface directory of CLI shims on PATH; every entry under it
+# re-enters cmux's own agent wrapper instead of the real CLI.
+SHIM_DIR_MARKER = "cmux-cli-shims"
+
+# Workers fill two columns and then grow downward. A third column would cost
+# width, and width is the dimension a horizontal split cannot give back.
+GRID_COLUMNS = 2
+
+# Measured against a Claude Code worker: 67 columns renders losslessly, 33 drops
+# content off the right edge. Below this floor okstra stacks a tab instead of
+# splitting, because an unreadable pane costs the lead its view of that worker.
+WORKER_MIN_COLUMNS = 60
+
+# What the lead keeps for itself once workers arrive.
+LEAD_TARGET_COLUMNS = 80
+
+# Sidebar entries are keyed by source so tools do not overwrite each other's.
+SIDEBAR_SOURCE = "okstra"
+
+
+def cmux_cli_path() -> str:
+    """Absolute path to the cmux CLI, or "" when cmux is not installed.
+
+    The bundled CLI wins over PATH: cmux documents the /usr/local/bin/cmux
+    symlink as a manual step, so a working install may leave PATH untouched.
+    """
+    bundled = os.environ.get("CMUX_BUNDLED_CLI_PATH", "")
+    if bundled and os.access(bundled, os.X_OK):
+        return bundled
+    return shutil.which("cmux") or ""
+
+
+def run_cmux(
+    args: Sequence[str], *, timeout: int = PROBE_TIMEOUT_SECONDS
+) -> subprocess.CompletedProcess[str]:
+    cli = cmux_cli_path()
+    if not cli:
+        raise FileNotFoundError(
+            "cmux CLI not found: CMUX_BUNDLED_CLI_PATH unset and no cmux on PATH"
+        )
+    return subprocess.run(
+        [cli, *args],
+        capture_output=True,
+        text=True,
+        timeout=timeout,
+        check=False,
+    )
+
+
+def cmux_available() -> bool:
+    """True only when okstra can drive cmux *and* knows where the lead sits.
+
+    Being able to drive cmux is not sufficient. Workers attach to the lead's
+    workspace, so a run whose lead location is unresolvable would open panes on
+    a screen nobody is watching; degrading to the blocking wrapper is the better
+    failure. Nested tmux is exactly that case — see `resolve_lead_workspace`.
+    """
+    if not cmux_cli_path():
+        return False
+    # Cheapest discriminator first: no workspace in the environment means this
+    # is not a cmux-hosted session, and the probes below would only confirm that
+    # the app happens to be installed.
+    if not _lead_workspace_env():
+        return False
+    if not _ping_answers():
+        return False
+    return bool(resolve_lead_workspace())
+
+
+def resolve_lead_workspace() -> str:
+    """The lead's workspace UUID, or "" when it cannot be resolved.
+
+    The UUID is read from the environment rather than from `identify`, whose
+    refs are positional and shift as surfaces open and close. `identify` is
+    consulted only to prove the environment is current: a tmux server freezes
+    the CMUX_* block it was launched with, so those variables can outlive the
+    surface they name.
+    """
+    workspace = _lead_workspace_env()
+    if not workspace:
+        return ""
+    if not identify_caller():
+        return ""
+    return workspace
+
+
+def _lead_workspace_env() -> str:
+    return os.environ.get("CMUX_WORKSPACE_ID", "").strip()
+
+
+@dataclass(frozen=True)
+class PaneGeometry:
+    """One pane as `pane.list` reports it.
+
+    `columns` and `rows` come straight from cmux rather than being derived from
+    the container frame, so a display with a different cell size needs no
+    conversion here.
+    """
+
+    pane_id: str
+    surface_ids: tuple[str, ...]
+    columns: int
+    rows: int
+    x: int
+    y: int
+    cell_width_points: int
+    ref: str = ""
+    selected_surface_id: str = ""
+
+
+@dataclass(frozen=True)
+class Placement:
+    """Where the next worker goes: split `pane_id`, or stack a tab into it."""
+
+    pane_id: str
+    direction: str
+    stack_as_tab: bool
+
+
+def plan_worker_placement(
+    panes: Sequence[PaneGeometry], *, lead_pane_id: str, min_columns: int
+) -> Placement:
+    """Pick the next worker slot from the workspace's current geometry.
+
+    Stateless by design: okstra records surface UUIDs, never a layout, so a
+    resumed or crashed run cannot carry a layout model that no longer matches
+    the screen. Every dispatch re-reads the panes and derives the next slot.
+    """
+    workers = [pane for pane in panes if pane.pane_id != lead_pane_id]
+    if not workers:
+        return Placement(pane_id=lead_pane_id, direction="right", stack_as_tab=False)
+
+    columns = _panes_by_column(workers)
+    if len(columns) < GRID_COLUMNS:
+        return _widen_the_grid(workers, min_columns=min_columns)
+    return _extend_the_shortest_column(columns)
+
+
+def lead_shrink_points(lead: PaneGeometry, *, target_columns: int) -> int:
+    """How far to push the lead's right border, in the points `pane.resize` takes.
+
+    The API's `amount` is points, not cells — measured at this pane's own
+    `cell_width_points`, so passing a column count shrinks by an eighth of the
+    intent on a typical display.
+    """
+    surplus = lead.columns - target_columns
+    if surplus <= 0:
+        return 0
+    return surplus * lead.cell_width_points
+
+
+def _widen_the_grid(
+    workers: Sequence[PaneGeometry], *, min_columns: int
+) -> Placement:
+    rightmost = max(workers, key=lambda pane: pane.x)
+    if rightmost.columns // GRID_COLUMNS >= min_columns:
+        return Placement(pane_id=rightmost.pane_id, direction="right", stack_as_tab=False)
+    roomiest = min(workers, key=lambda pane: (len(pane.surface_ids), pane.x))
+    return Placement(pane_id=roomiest.pane_id, direction="", stack_as_tab=True)
+
+
+def _extend_the_shortest_column(
+    columns: dict[int, list[PaneGeometry]]
+) -> Placement:
+    shortest = min(columns.values(), key=lambda group: (len(group), group[0].x))
+    bottom = max(shortest, key=lambda pane: pane.y)
+    return Placement(pane_id=bottom.pane_id, direction="down", stack_as_tab=False)
+
+
+def _panes_by_column(
+    workers: Sequence[PaneGeometry],
+) -> dict[int, list[PaneGeometry]]:
+    columns: dict[int, list[PaneGeometry]] = {}
+    for pane in workers:
+        columns.setdefault(pane.x, []).append(pane)
+    return columns
+
+
+def shim_free_login_path() -> str:
+    """The user's login PATH with cmux's per-surface CLI shims removed.
+
+    A cmux pane execs its command through `login … bash --noprofile --norc`, so
+    none of the user's shell rc runs and PATH is cmux's own short list. What that
+    list does hold is a shim directory bound to the new surface, where `claude`
+    and `codex` are wrappers into cmux's agent lifecycle rather than the real
+    CLIs — and where a provider cmux does not know has no entry at all, which is
+    a plain exit 127 inside the worker wrapper.
+
+    Restoring the login shell's PATH and dropping the shim entries gives the
+    wrapper what it would see in an ordinary terminal, which is what its own
+    `command -v <cli>` check expects. Nothing here is provider-specific, so
+    adding a provider does not touch this path.
+    """
+    entries = _login_shell_path().split(os.pathsep)
+    return os.pathsep.join(
+        entry for entry in entries if entry and SHIM_DIR_MARKER not in entry
+    )
+
+
+def worker_command_line(
+    *, cwd: Path, argv: Sequence[str], path_value: str
+) -> str:
+    """The single shell line a cmux pane execs to run one worker.
+
+    Every part is shell-quoted. PATH entries routinely contain spaces — macOS
+    ships `/Applications/VMware Fusion.app/Contents/Public` on any machine with
+    VMware — and an unquoted assignment stops at the first space, then hands the
+    remainder to the shell as a command name.
+
+    The `cd` is part of the command because neither `new-split` nor
+    `respawn-pane` accepts a working directory the way `tmux split-window -c`
+    does.
+    """
+    return (
+        f"cd {shlex.quote(str(cwd))} && "
+        f"PATH={shlex.quote(path_value)} exec {shlex.join(argv)}"
+    )
+
+
+def list_panes(workspace: str) -> list[PaneGeometry]:
+    payload = rpc("pane.list", {"workspace_id": workspace})
+    return [_pane_geometry(entry) for entry in payload.get("panes", [])]
+
+
+def spawn_worker_surface(
+    *, workspace: str, cwd: Path, command: Sequence[str], title: str
+) -> str:
+    """Start one worker beside the lead and return its surface UUID.
+
+    The UUID is what okstra records and later closes by. Positional refs cannot
+    serve that purpose: cmux renumbers them as surfaces open and close, so a
+    close by ref can land on a pane okstra never created.
+    """
+    panes = list_panes(workspace)
+    lead = _lead_pane(panes)
+    placement = plan_worker_placement(
+        panes, lead_pane_id=lead.pane_id, min_columns=WORKER_MIN_COLUMNS
+    )
+    target = _pane_by_id(panes, placement.pane_id)
+    surface_uuid = _open_worker_surface(workspace, placement, target)
+    run_cmux(["rename-tab", "--surface", surface_uuid, "--title", title])
+    _exec_worker(surface_uuid, cwd=cwd, command=command)
+    _shrink_lead_pane(workspace)
+    return surface_uuid
+
+
+def close_surface(surface_uuid: str) -> None:
+    """Close an okstra-created surface, killing whatever still runs inside it."""
+    try:
+        run_cmux(["close-surface", "--surface", surface_uuid])
+    except (OSError, subprocess.SubprocessError):
+        return
+
+
+def capture_surface(surface_uuid: str, *, last_lines: int = 200) -> str:
+    """What the worker's screen shows — for the lead to look at, never to parse.
+
+    cmux hands back a rendered grid: wrapped to the pane's width with a finite
+    history, so a path or a count read off it can be silently truncated.
+    """
+    try:
+        result = run_cmux(
+            [
+                "capture-pane",
+                "--surface",
+                surface_uuid,
+                "--scrollback",
+                "--lines",
+                str(last_lines),
+            ]
+        )
+    except (OSError, subprocess.SubprocessError):
+        return ""
+    return result.stdout if result.returncode == 0 else ""
+
+
+def sidebar_log(workspace: str, message: str, *, level: str = "info") -> None:
+    """Append one line to the workspace sidebar's log."""
+    _sidebar_call(
+        workspace,
+        ["log", "--workspace", workspace, "--level", level,
+         "--source", SIDEBAR_SOURCE, "--", message],
+    )
+
+
+def sidebar_notify(workspace: str, *, title: str, body: str) -> None:
+    """Raise a notification, which badges the tab even from another workspace."""
+    _sidebar_call(
+        workspace,
+        ["notify", "--workspace", workspace, "--title", title, "--body", body],
+    )
+
+
+def sidebar_progress(workspace: str, done: int, total: int, *, label: str) -> None:
+    """Show `done`/`total` as the sidebar's progress bar.
+
+    An empty roster reports nothing rather than zero percent: no workers is not
+    the same claim as no progress.
+    """
+    if total <= 0:
+        return
+    fraction = min(1.0, done / total)
+    _sidebar_call(
+        workspace,
+        ["set-progress", str(round(fraction, 2)), "--workspace", workspace,
+         "--label", label],
+    )
+
+
+def _sidebar_call(workspace: str, args: Sequence[str]) -> None:
+    """Best-effort sidebar write.
+
+    The sidebar is a view of the run, not part of it, so a wedged or closed cmux
+    must never take a dispatch down. The workspace is always explicit: letting
+    cmux choose would drop one run's noise into whichever workspace happens to
+    be focused, which on a machine running two okstra tasks is the other one.
+    """
+    if not workspace:
+        return
+    try:
+        run_cmux(args)
+    except (OSError, subprocess.SubprocessError):
+        return
+
+
+def rpc(method: str, params: dict[str, Any]) -> dict[str, Any]:
+    """Call a cmux rpc method.
+
+    The rpc surface is used rather than the same-named CLI verbs wherever both
+    exist, because the CLI swallows failures — `resize-pane` reports success on
+    a rejected resize, while `pane.resize` returns the reason.
+    """
+    result = run_cmux(["rpc", method, json.dumps(params)])
+    if result.returncode != 0:
+        raise RuntimeError(result.stderr.strip() or f"cmux {method} failed")
+    try:
+        payload = json.loads(result.stdout)
+    except ValueError as exc:
+        raise RuntimeError(f"cmux {method} returned non-JSON output") from exc
+    return payload if isinstance(payload, dict) else {}
+
+
+def _open_worker_surface(
+    workspace: str, placement: Placement, target: PaneGeometry
+) -> str:
+    before = _surface_ids(workspace)
+    _create_surface(workspace, placement, target)
+    new_ids = _surface_ids(workspace) - before
+    if len(new_ids) != 1:
+        raise RuntimeError(
+            f"cmux opened {len(new_ids)} surfaces where exactly one was expected"
+        )
+    return new_ids.pop()
+
+
+def _surface_ids(workspace: str) -> set[str]:
+    """Every surface UUID in the workspace.
+
+    The new surface is identified by diffing this set rather than by translating
+    the `OK surface:N` echo, because `list-pane-surfaces` reports only the
+    focused pane unless given a `--pane`, and the new pane is not focused.
+    """
+    return {
+        surface_id
+        for pane in list_panes(workspace)
+        for surface_id in pane.surface_ids
+    }
+
+
+def _create_surface(
+    workspace: str, placement: Placement, target: PaneGeometry
+) -> None:
+    if placement.stack_as_tab:
+        created = run_cmux(
+            ["new-surface", "--workspace", workspace, "--pane", target.pane_id]
+        )
+    else:
+        created = run_cmux(
+            [
+                "new-split",
+                placement.direction,
+                "--workspace",
+                workspace,
+                "--surface",
+                target.selected_surface_id or target.surface_ids[0],
+            ]
+        )
+    if created.returncode != 0:
+        raise RuntimeError(created.stderr.strip() or "cmux could not open a pane")
+
+
+def _exec_worker(surface_uuid: str, *, cwd: Path, command: Sequence[str]) -> None:
+    line = worker_command_line(
+        cwd=cwd, argv=command, path_value=shim_free_login_path()
+    )
+    started = run_cmux(["respawn-pane", "--surface", surface_uuid, "--command", line])
+    if started.returncode != 0:
+        raise RuntimeError(started.stderr.strip() or "cmux could not start the worker")
+
+
+def _shrink_lead_pane(workspace: str) -> None:
+    """Give the lead's width to the workers by pushing its right border left.
+
+    The lead cannot shrink itself: `pane.resize` moves the named pane's border,
+    so asking the leftmost pane to move `left` fails with no adjacent border and
+    `right` widens it. The neighbour on its right carries the request instead.
+    """
+    panes = list_panes(workspace)
+    lead = _lead_pane(panes)
+    amount = lead_shrink_points(lead, target_columns=LEAD_TARGET_COLUMNS)
+    if amount <= 0:
+        return
+    neighbours = [pane for pane in panes if pane.x > lead.x]
+    if not neighbours:
+        return
+    rpc(
+        "pane.resize",
+        {
+            "workspace_id": workspace,
+            "pane_id": min(neighbours, key=lambda pane: pane.x).pane_id,
+            "direction": "left",
+            "amount": amount,
+        },
+    )
+
+
+def _lead_pane(panes: Sequence[PaneGeometry]) -> PaneGeometry:
+    caller_pane_ref = identify_caller().get("pane_ref", "")
+    for pane in panes:
+        if caller_pane_ref and pane.ref == caller_pane_ref:
+            return pane
+    raise RuntimeError("cmux could not locate the lead's pane in its workspace")
+
+
+def _pane_by_id(panes: Sequence[PaneGeometry], pane_id: str) -> PaneGeometry:
+    for pane in panes:
+        if pane.pane_id == pane_id:
+            return pane
+    raise RuntimeError(f"cmux pane {pane_id} disappeared while placing a worker")
+
+
+def _pane_geometry(entry: dict[str, Any]) -> PaneGeometry:
+    frame = entry.get("pixel_frame") or {}
+    return PaneGeometry(
+        pane_id=str(entry.get("id", "")),
+        surface_ids=tuple(str(s) for s in entry.get("surface_ids") or ()),
+        columns=int(entry.get("columns", 0)),
+        rows=int(entry.get("rows", 0)),
+        x=int(frame.get("x", 0)),
+        y=int(frame.get("y", 0)),
+        cell_width_points=int(entry.get("cell_width_points", 0)),
+        ref=str(entry.get("ref", "")),
+        selected_surface_id=str(entry.get("selected_surface_id", "")),
+    )
+
+
+def _login_shell_path() -> str:
+    """PATH as the user's own login shell resolves it.
+
+    Falls back to the inherited PATH: a worker started with a shim-filtered
+    inherited PATH is still better than one started with cmux's bare list.
+    """
+    shell = os.environ.get("SHELL", "") or "/bin/sh"
+    try:
+        result = subprocess.run(
+            [shell, "-lc", 'printf %s "$PATH"'],
+            capture_output=True,
+            text=True,
+            timeout=LOGIN_SHELL_TIMEOUT_SECONDS,
+            check=False,
+        )
+    except (OSError, subprocess.SubprocessError):
+        return os.environ.get("PATH", "")
+    if result.returncode != 0 or not result.stdout.strip():
+        return os.environ.get("PATH", "")
+    return result.stdout.strip()
+
+
+def identify_caller() -> dict[str, Any]:
+    """The caller's ref bundle, or {} when cmux cannot resolve it.
+
+    `identify` exits 0 even for an id it cannot resolve, so the exit code says
+    nothing; a null `caller` is the only signal that the lookup failed.
+    """
+    try:
+        result = run_cmux(["identify"])
+    except (OSError, subprocess.SubprocessError):
+        return {}
+    if result.returncode != 0:
+        return {}
+    try:
+        payload = json.loads(result.stdout)
+    except ValueError:
+        return {}
+    caller = payload.get("caller")
+    return caller if isinstance(caller, dict) else {}
+
+
+def _ping_answers() -> bool:
+    try:
+        result = run_cmux(["ping"])
+    except (OSError, subprocess.SubprocessError):
+        return False
+    return result.returncode == 0 and result.stdout.strip() == PING_OK

package/runtime/python/okstra_ctl/codex_dispatch.py

@@ -283,7 +283,7 @@ def _build_selected_worker_dispatches(
     idle_timeout_seconds: int,
     selected_workers: Sequence[str],
     worker_requests: Sequence[InitialPromptWorkerRequest],
-    prompt_paths: Sequence[Path],
+    prompt_paths: Mapping[str, Path],
 ) -> tuple[WorkerJob, ...]:
     return tuple(
         _build_worker_dispatch(
@@ -298,11 +298,11 @@ def _build_selected_worker_dispatches(
             model_execution_value=worker_request.model,
             prompt_path=prompt_path,
         )
-        for worker_id, worker_request, prompt_path in zip(
-            selected_workers,
-            worker_requests,
-            prompt_paths,
-            strict=True,
+        for worker_id, worker_request, prompt_path in (
+            (worker_id, worker_request, prompt_paths[worker_id])
+            for worker_id, worker_request in zip(
+                selected_workers, worker_requests, strict=True
+            )
         )
     )