Run tty

scripts/run_codex_tty.py

@@ -0,0 +1,245 @@
+#!/usr/bin/env -S uv run --script
+# /// script
+# dependencies = ["pyte"]
+# ///
+import argparse
+import os
+import queue
+import select
+import signal
+import subprocess
+import sys
+import threading
+import time
+from pathlib import Path
+
+try:
+    import pyte
+except ImportError as exc:
+    raise SystemExit(
+        "pyte is required. Install with: pip install pyte"
+    ) from exc
+
+
+ENTER = b"\r"
+STEER_OVERRIDES = [
+    "--config",
+    "features.collaboration_modes=true",
+    "--config",
+    "tui.experimental_mode=plan",
+]
+TYPE_DELAY_SEC = 0.02
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Run N Codex TTY instances, send a prompt, and capture ASCII screenshots."
+    )
+    parser.add_argument("prompt", help="Prompt to send to each Codex instance.")
+    parser.add_argument(
+        "-n",
+        "--instances",
+        type=int,
+        default=1,
+        help="Number of Codex instances to launch.",
+    )
+    parser.add_argument(
+        "--out-dir",
+        default="screens",
+        help="Directory to write ASCII screenshots.",
+    )
+    parser.add_argument(
+        "--cols",
+        type=int,
+        default=120,
+        help="TTY columns for each instance.",
+    )
+    parser.add_argument(
+        "--rows",
+        type=int,
+        default=1200,
+        help="TTY rows for each instance.",
+    )
+    return parser.parse_args()
+
+
+def build_codex(repo_root: Path) -> Path:
+    codex_rs = repo_root / "codex-rs"
+    result = subprocess.run(
+        ["cargo", "build", "-p", "codex-cli"],
+        cwd=codex_rs,
+        check=False,
+    )
+    if result.returncode != 0:
+        raise SystemExit("cargo build failed")
+    codex_bin = codex_rs / "target" / "debug" / "codex"
+    if not codex_bin.exists():
+        raise SystemExit(f"codex binary not found at {codex_bin}")
+    return codex_bin
+
+
+def set_winsize(fd: int, rows: int, cols: int) -> None:
+    import fcntl
+    import struct
+    import termios
+
+    winsize = struct.pack("HHHH", rows, cols, 0, 0)
+    fcntl.ioctl(fd, termios.TIOCSWINSZ, winsize)
+
+
+def run_instance(
+    index: int,
+    codex_bin: Path,
+    repo_root: Path,
+    prompt: str,
+    out_dir: Path,
+    rows: int,
+    cols: int,
+    result_queue: "queue.Queue[tuple[int, Path]]",
+) -> None:
+    master_fd, slave_fd = os.openpty()
+    set_winsize(slave_fd, rows, cols)
+
+    env = os.environ.copy()
+    env["COLUMNS"] = str(cols)
+    env["LINES"] = str(rows)
+
+    proc = subprocess.Popen(
+        [str(codex_bin), *STEER_OVERRIDES],
+        stdin=slave_fd,
+        stdout=slave_fd,
+        stderr=slave_fd,
+        env=env,
+        cwd=repo_root,
+        preexec_fn=os.setsid,
+        close_fds=True,
+    )
+    os.close(slave_fd)
+
+    screen = pyte.Screen(cols, rows)
+    stream = pyte.Stream(screen)
+
+    last_output = time.monotonic()
+    sent_prompt = False
+    start_time = time.monotonic()
+    prompt_attempt_deadline = start_time + 5.0
+    prompt_sent_at: float | None = None
+    response_started = False
+    last_enter_retry = 0.0
+    buffer = b""
+    cpr_probe = b""
+
+    try:
+        while True:
+            if not sent_prompt and time.monotonic() - start_time >= 1.0:
+                prompt_line = next((line for line in screen.display if "›" in line), "")
+                has_prompt = bool(prompt_line)
+                if has_prompt or time.monotonic() >= prompt_attempt_deadline:
+                    write_with_delay(master_fd, prompt.encode(), TYPE_DELAY_SEC)
+                    time.sleep(0.2)
+                    os.write(master_fd, ENTER)
+                    sent_prompt = True
+                    prompt_sent_at = time.monotonic()
+
+            timeout = 0.1
+            rlist, _, _ = select.select([master_fd], [], [], timeout)
+            if rlist:
+                data = os.read(master_fd, 4096)
+                if not data:
+                    break
+                last_output = time.monotonic()
+                cpr_probe = (cpr_probe + data)[-8:]
+                if b"\x1b[6n" in cpr_probe:
+                    os.write(master_fd, b"\x1b[1;1R")
+                buffer += data
+                while True:
+                    try:
+                        chunk = buffer.decode(errors="ignore")
+                    except UnicodeDecodeError:
+                        break
+                    stream.feed(chunk)
+                    if not response_started and any("Working" in line for line in screen.display):
+                        response_started = True
+                    buffer = b""
+                    break
+
+            if sent_prompt and not response_started:
+                now = time.monotonic()
+                if now - prompt_sent_at >= 2.0 and now - last_enter_retry >= 2.0:
+                    os.write(master_fd, ENTER)
+                    last_enter_retry = now
+
+            if response_started and time.monotonic() - last_output >= 10.0:
+                break
+
+            if proc.poll() is not None:
+                break
+    finally:
+        try:
+            os.killpg(proc.pid, signal.SIGTERM)
+        except:
+            pass
+        try:
+            proc.wait(timeout=5)
+        except subprocess.TimeoutExpired:
+            try:
+                os.killpg(proc.pid, signal.SIGKILL)
+            except:
+                pass
+
+        out_dir.mkdir(parents=True, exist_ok=True)
+        out_path = out_dir / f"codex_{index}.txt"
+        with out_path.open("w", encoding="utf-8") as f:
+            for line in screen.display:
+                f.write(line.rstrip() + "\n")
+        result_queue.put((index, out_path))
+
+        os.close(master_fd)
+
+
+def main() -> int:
+    args = parse_args()
+    if args.instances < 1:
+        raise SystemExit("--instances must be >= 1")
+
+    repo_root = Path(__file__).resolve().parents[1]
+    codex_bin = build_codex(repo_root)
+
+    out_dir = Path(args.out_dir)
+    result_queue: "queue.Queue[tuple[int, Path]]" = queue.Queue()
+    threads = []
+    for i in range(1, args.instances + 1):
+        thread = threading.Thread(
+            target=run_instance,
+            args=(
+                i,
+                codex_bin,
+                repo_root,
+                args.prompt,
+                out_dir,
+                args.rows,
+                args.cols,
+                result_queue,
+            ),
+            daemon=True,
+        )
+        thread.start()
+        threads.append(thread)
+
+    for thread in threads:
+        thread.join()
+
+    results = sorted(result_queue.get() for _ in range(args.instances))
+    for index, path in results:
+        print(f"instance {index}: {path}")
+    return 0
+
+
+def write_with_delay(fd: int, data: bytes, delay: float = 0.05) -> None:
+    for ch in data:
+        os.write(fd, bytes([ch]))
+        time.sleep(delay)
+
+
+if __name__ == "__main__":
+    sys.exit(main())