scripts/codex_route_ab.py

#!/usr/bin/env python3
"""Codex CLI route A/B: same prompt, same model/effort, API-key vs ChatGPT-OAuth auth.

Question: does the Codex (ChatGPT OAuth) backend spend fewer reasoning tokens than the
OpenAI API for an identical request made by the *same* client (the official Codex CLI)?

Each route gets its own throwaway CODEX_HOME (mode 0700, deleted at exit), so the user's
~/.codex is never read or written. Credentials reach the CLI only via stdin to
`codex login`; they are never put in argv, printed, or kept after the run. Provider
environment variables are stripped from `codex exec`, so the OAuth arm cannot silently
fall back to an API key. Rounds interleave the routes to spread time-of-day effects.

Requirements: Python >= 3.9 (stdlib only), Codex CLI on PATH (`codex --version`).

Usage (both routes, 4 interleaved rounds, Luna medium):
    export OPENAI_API_KEY=...                       # API arm
    python3 codex_route_ab.py --model gpt-6-luna --effort medium --rounds 4 \\
        --oauth-codex-auth ~/.codex/auth.json       # OAuth arm: copy of a Codex CLI login
        # or: --oauth-access-token-env VAR (Codex *agent identity* JWT, not a ChatGPT token)
    python3 codex_route_ab.py --routes oauth ...    # one arm only

The OAuth arm copies a ChatGPT-mode Codex CLI auth.json into the throwaway home. Codex
refreshes tokens when `last_refresh` is ~8 days old, and a refresh rotates the refresh
token (which would strand the original login), so the script refuses logins older than
7 days: run any `codex` command normally first to refresh your real login.

Outputs (default ./codex-route-ab-<UTC timestamp>/): per-run raw `codex exec --json`
event logs (<route>-r<N>.jsonl), final answers, results.jsonl and summary.json with
codex version, prompt sha256, per-run usage (input, cached, output, reasoning output),
wall time, tool/command count and answer check. Event logs contain model output only.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import os
import re
import shutil
import statistics
import subprocess
import sys
import tempfile
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime, timezone
from pathlib import Path

DEFAULT_PROMPT = (
    "Find all pairs of integers (x, y) with 1 <= x <= y <= 100 such that x^2 + y^2 + 1 "
    "is divisible by x*y.\nReason carefully and prove that your list is complete (do not "
    "just test a few values).\nFinish with a line \"ANSWER:\" followed by the pairs in "
    "increasing order of y.\n")
DEFAULT_EXPECTED = "(1,1),(1,2),(2,5),(5,13),(13,34),(34,89)"
SECRET_ENV = ("OPENAI_API_KEY", "CODEX_API_KEY", "CODEX_ACCESS_TOKEN", "OPENAI_ORG_ID",
              "OPENAI_ORGANIZATION", "OPENAI_PROJECT", "OPENAI_BASE_URL")


def clean_env(home: Path) -> dict:
    env = {k: v for k, v in os.environ.items() if k not in SECRET_ENV}
    env["CODEX_HOME"] = str(home)
    return env


def run(cmd, env, stdin=None, timeout=60):
    return subprocess.run(cmd, env=env, input=stdin, capture_output=True, text=True,
                          timeout=timeout)


def login(route: str, home: Path, args) -> str:
    """Authenticate one isolated CODEX_HOME; returns a non-secret description."""
    env = clean_env(home)
    if route == "api":
        key = os.environ.get(args.api_key_env)
        if not key:
            sys.exit(f"{args.api_key_env} is not set (API arm)")
        proc = run(["codex", "login", "--with-api-key"], env, stdin=key + "\n")
        method = f"api key from ${args.api_key_env}"
    elif args.oauth_codex_auth:
        source = Path(args.oauth_codex_auth).expanduser()
        data = json.loads(source.read_text())
        if data.get("auth_mode") not in (None, "chatgpt") or not (data.get("tokens") or {}).get("access_token"):
            sys.exit("oauth: --oauth-codex-auth is not a ChatGPT-mode Codex login")
        refreshed = data.get("last_refresh")
        if refreshed:
            age = datetime.now(timezone.utc) - datetime.fromisoformat(refreshed.replace("Z", "+00:00"))
            if age.total_seconds() > 7 * 86400:
                sys.exit("oauth: login last refreshed > 7 days ago; refresh it with normal codex "
                         "use first so the copy does not rotate its refresh token")
        target = home / "auth.json"
        shutil.copyfile(source, target)
        os.chmod(target, 0o600)
        proc = run(["codex", "login", "status"], env)
        method = "copy of a Codex CLI auth.json"
    else:
        token = os.environ.get(args.oauth_access_token_env)
        if not token:
            sys.exit(f"{args.oauth_access_token_env} is not set (OAuth arm)")
        proc = run(["codex", "login", "--with-access-token"], env, stdin=token + "\n")
        method = f"access token from ${args.oauth_access_token_env}"
    if proc.returncode != 0:  # never echo CLI output: it may contain masked key fragments
        sys.exit(f"{route}: codex login failed (exit {proc.returncode})")
    status = run(["codex", "login", "status"], env)
    if status.returncode != 0:
        sys.exit(f"{route}: codex login status failed (exit {status.returncode})")
    said = (status.stdout + status.stderr).lower()  # status prints to stderr; never echoed
    kind = "api-key" if "using an api key" in said else "chatgpt" if "chatgpt" in said else "unknown"
    want = "api-key" if route == "api" else "chatgpt"
    if kind != want:
        sys.exit(f"{route}: expected {want} auth, codex reports {kind}")
    return method


def normalise_pairs(text: str) -> str | None:
    lines = [l for l in text.splitlines() if "ANSWER" in l.upper()]
    if not lines:
        return None
    pairs = re.findall(r"\(\s*(\d+)\s*,\s*(\d+)\s*\)", lines[-1])
    return ",".join(f"({a},{b})" for a, b in pairs) or None


def exec_once(route, home, workdir, args, prompt, out_dir, index):
    cmd = ["codex", "exec", "--json", "--ephemeral", "--ignore-user-config",
           "--ignore-rules", "--skip-git-repo-check", "-s", "read-only",
           "-C", str(workdir), "-m", args.model,
           "-c", f'model_reasoning_effort="{args.effort}"']
    for override in args.config:
        cmd += ["-c", override]
    cmd.append("-")
    started_at = datetime.now(timezone.utc).isoformat(timespec="seconds")
    started = time.monotonic()
    proc = subprocess.run(cmd, env=clean_env(home), input=prompt, capture_output=True,
                          text=True, timeout=args.timeout)
    seconds = round(time.monotonic() - started, 1)
    raw = out_dir / f"{route}-r{index}.jsonl"
    raw.write_text(proc.stdout)
    usage, message, items, errors = {}, "", {}, []
    for line in proc.stdout.splitlines():
        try:
            event = json.loads(line)
        except json.JSONDecodeError:
            continue
        kind = event.get("type", "")
        if isinstance(event.get("usage"), dict):  # turn.completed; sum across turns
            for k, v in event["usage"].items():
                if isinstance(v, (int, float)):
                    usage[k] = usage.get(k, 0) + v
        item = event.get("item") or {}
        if kind == "item.completed" and isinstance(item, dict):
            items[item.get("type", "?")] = items.get(item.get("type", "?"), 0) + 1
            if item.get("type") in ("agent_message", "assistant_message"):
                message = item.get("text") or message
        if kind in ("error", "turn.failed"):
            errors.append(str(event.get("message") or event.get("error"))[:300])
    (out_dir / f"{route}-r{index}.answer.txt").write_text(message)
    got = normalise_pairs(message)
    return {"route": route, "round": index, "started_at": started_at,
            "exit": proc.returncode, "seconds": seconds,
            "usage": usage, "items": items, "errors": errors,
            "answer": got, "correct": (got == args.expected) if args.expected else None}


def summarise(rows, routes):
    out = {}
    for route in routes:
        ok = [r for r in rows if r["route"] == route and r["exit"] == 0 and r["usage"]]
        def med(key):
            values = [r["usage"].get(key) for r in ok if r["usage"].get(key) is not None]
            return statistics.median(values) if values else None
        out[route] = {"runs": len([r for r in rows if r["route"] == route]), "ok": len(ok),
                      "correct": sum(bool(r["correct"]) for r in ok),
                      "median_reasoning_output_tokens": med("reasoning_output_tokens"),
                      "median_output_tokens": med("output_tokens"),
                      "median_input_tokens": med("input_tokens"),
                      "median_seconds": statistics.median([r["seconds"] for r in ok]) if ok else None,
                      "reasoning_output_tokens": [r["usage"].get("reasoning_output_tokens") for r in ok]}
    return out


def main():
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--model", required=True)
    p.add_argument("--effort", default="medium")
    p.add_argument("--rounds", type=int, default=4)
    p.add_argument("--routes", default="api,oauth", help="comma list of api,oauth")
    p.add_argument("--prompt-file", help="default: built-in number-theory prompt")
    p.add_argument("--expected", default=None,
                   help=f"normalised ANSWER pairs; default for built-in prompt: {DEFAULT_EXPECTED}")
    p.add_argument("--api-key-env", default="OPENAI_API_KEY")
    p.add_argument("--oauth-codex-auth", help="path to a Codex CLI auth.json (copied into the temp home)")
    p.add_argument("--oauth-access-token-env", default="CODEX_ACCESS_TOKEN")
    p.add_argument("-c", "--config", action="append", default=[],
                   help="extra `codex exec -c key=value` overrides, applied to both routes")
    p.add_argument("--timeout", type=int, default=1200)
    p.add_argument("--out", help="output directory")
    p.add_argument("--parallel", type=int, default=1,
                   help="concurrent workers (each with its own homes); jobs stay interleaved")
    p.add_argument("--login-only", action="store_true",
                   help="authenticate each route's temp home, verify auth kind, run nothing")
    args = p.parse_args()
    routes = [r.strip() for r in args.routes.split(",") if r.strip()]
    if not routes or set(routes) - {"api", "oauth"}:
        sys.exit("--routes must be a comma list of api,oauth")
    prompt = Path(args.prompt_file).read_text() if args.prompt_file else DEFAULT_PROMPT
    if args.expected is None and not args.prompt_file:
        args.expected = DEFAULT_EXPECTED
    stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
    out_dir = Path(args.out or f"codex-route-ab-{stamp}")
    out_dir.mkdir(parents=True, exist_ok=False)
    version = run(["codex", "--version"], {k: v for k, v in os.environ.items()
                                            if k not in SECRET_ENV}).stdout.strip()
    meta = {"started_at": stamp, "codex_version": version, "model": args.model,
            "effort": args.effort, "rounds": args.rounds, "routes": routes,
            "prompt_sha256": hashlib.sha256(prompt.encode()).hexdigest(),
            "expected": args.expected, "extra_config": args.config, "parallel": args.parallel,
            "exec_flags": "--json --ephemeral --ignore-user-config --ignore-rules "
                          "--skip-git-repo-check -s read-only"}
    # Exact command lines (temp paths shown as placeholders; credentials only via stdin).
    meta["commands"] = {
        "api_login": "printf '%s\\n' \"$" + args.api_key_env + "\" | CODEX_HOME=$HOME_API codex login --with-api-key",
        "oauth_login": ("cp " + str(args.oauth_codex_auth) + " $HOME_OAUTH/auth.json  # chmod 600; then codex login status"
                        if args.oauth_codex_auth else
                        "printf '%s\\n' \"$" + args.oauth_access_token_env + "\" | CODEX_HOME=$HOME_OAUTH codex login --with-access-token"),
        "exec": "env -u " + " -u ".join(SECRET_ENV) + " CODEX_HOME=$HOME_<ROUTE> codex exec --json --ephemeral "
                "--ignore-user-config --ignore-rules --skip-git-repo-check -s read-only -C $WORKDIR "
                f"-m {args.model} -c 'model_reasoning_effort=\"{args.effort}\"'"
                + "".join(f" -c '{c}'" for c in args.config) + " - < prompt.txt"}
    (out_dir / "prompt.txt").write_text(prompt)
    # Codex refuses to install its helper binaries under /tmp, so keep the throwaway
    # homes under ~/.cache (still private and deleted at exit).
    cache = Path(os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache"))
    cache.mkdir(parents=True, exist_ok=True)
    temp_root = Path(tempfile.mkdtemp(prefix="codex-route-ab-", dir=cache))
    os.chmod(temp_root, 0o700)
    rows = []
    try:
        if args.parallel < 1:
            sys.exit("--parallel must be >= 1")
        # One private home + workdir per (worker, route): concurrent codex processes never
        # share CODEX_HOME state.
        slots = []
        for worker in range(args.parallel):
            homes, workdirs = {}, {}
            for route in routes:
                homes[route] = temp_root / f"home-{route}-w{worker}"
                homes[route].mkdir(mode=0o700)
                workdirs[route] = temp_root / f"work-{route}-w{worker}"
                workdirs[route].mkdir(mode=0o700)
                meta[f"{route}_auth"] = login(route, homes[route], args)
            slots.append((homes, workdirs))
        print(json.dumps(meta), file=sys.stderr)
        if args.login_only:
            print("login-only: all homes authenticated as expected; no model calls", file=sys.stderr)
            return
        jobs = [(index, route) for index in range(1, args.rounds + 1)
                for route in (routes if index % 2 else list(reversed(routes)))]
        free, lock = list(range(args.parallel)), threading.Lock()
        results = open(out_dir / "results.jsonl", "w")

        def work(job):
            index, route = job
            with lock:
                worker = free.pop()
            try:
                homes, workdirs = slots[worker]
                row = exec_once(route, homes[route], workdirs[route], args, prompt, out_dir, index)
            finally:
                with lock:
                    free.append(worker)
            with lock:
                rows.append(row)
                results.write(json.dumps(row) + "\n")
                results.flush()
                u = row["usage"]
                print(f"r{index} {route:5} exit={row['exit']} {row['seconds']:7.1f}s "
                      f"reasoning={u.get('reasoning_output_tokens')} output={u.get('output_tokens')} "
                      f"input={u.get('input_tokens')} correct={row['correct']} items={row['items']}",
                      file=sys.stderr)

        try:
            with ThreadPoolExecutor(max_workers=args.parallel) as pool:
                list(pool.map(work, jobs))
        finally:
            results.close()
    finally:
        shutil.rmtree(temp_root, ignore_errors=True)  # removes the temporary credentials
    summary = {**meta, "summary": summarise(rows, routes)}
    (out_dir / "summary.json").write_text(json.dumps(summary, indent=2) + "\n")
    print(json.dumps(summary["summary"], indent=2))
    print(f"results: {out_dir}", file=sys.stderr)


if __name__ == "__main__":
    main()