codex_workflow.py

#!/usr/bin/env python3
"""codex_workflow.py — a deterministic orchestrator for OpenAI Codex.

Mirrors Claude Code's Workflow tool. The model never decides how many agents
to spawn — THIS script does, deterministically. Each agent() is one
`codex exec` subprocess; parallel()/pipeline() control concurrency.

Primitives (Workflow-tool parity):
  agent(prompt, schema=, role=, isolation='worktree', retries=) -> str|dict|None
  parallel(thunks)          -> list   barrier; concurrent; failed thunk -> None
  pipeline(items, *stages)  -> list   per-item chains, no barrier
  log(msg)                  -> None   progress line (mirrors Workflow log())
  phase(title)              -> None   name the current progress group
  tokens_used()             -> int    output-token meter (parsed from codex --json)
  budget                    -> obj    .total / .spent() / .remaining()
  collected_diffs()         -> list   diffs from isolation='worktree' writers
  start_run/save_result/write_ledger  durable run trail under .codex/ultracode/runs/

Env knobs:
  CODEX_WF_CONCURRENCY  max concurrent codex exec processes (default min(16, cores-2))
  CODEX_WF_MODEL        model for agents (default gpt-5.5)
  CODEX_WF_EFFORT       reasoning effort (default medium)
  CODEX_WF_CWD          base working dir / git repo agents run in (default $PWD)
  CODEX_WF_TIMEOUT      per-agent timeout seconds (default 1800)
  CODEX_WF_BUDGET       output-token ceiling; unset = unlimited
  CODEX_WF_EST_TOKENS   per-agent reservation used to bound concurrent overshoot (default 8000)
  CODEX_WF_ARGS         JSON value exposed verbatim as the module global `args`
"""
import atexit
import copy
import json
import os
import re
import subprocess
import sys
import tempfile
import threading
import time
import uuid
from concurrent.futures import ThreadPoolExecutor

MODEL = os.environ.get("CODEX_WF_MODEL", "gpt-5.5")
EFFORT = os.environ.get("CODEX_WF_EFFORT", "medium")
CONCURRENCY = int(os.environ.get("CODEX_WF_CONCURRENCY", str(min(16, max(1, (os.cpu_count() or 4) - 2)))))
CWD = os.environ.get("CODEX_WF_CWD", os.getcwd())
TIMEOUT = int(os.environ.get("CODEX_WF_TIMEOUT", "1800"))
# Budget is an OUTPUT-token ceiling. Unset (env absent/empty) == unlimited == None, mirroring
# Workflow's budget.total being null when the user set no target (NOT 0, which is a real ceiling).
_b = os.environ.get("CODEX_WF_BUDGET")
BUDGET = int(_b) if _b else None
EST_TOKENS = int(os.environ.get("CODEX_WF_EST_TOKENS", "8000"))
AGENT_CAP = 1000          # lifetime spawn backstop (Workflow caps a run at 1000 agents)
MAX_BATCH = 4096          # a single parallel()/pipeline() accepts at most this many items

# `args` — the Workflow-tool global, ported as a JSON env value the script can branch on.
try:
    args = json.loads(os.environ["CODEX_WF_ARGS"]) if os.environ.get("CODEX_WF_ARGS") else None
except (ValueError, KeyError):
    args = None

_pool = ThreadPoolExecutor(max_workers=CONCURRENCY)
_tok_lock = threading.Lock()
_tokens = [0]             # realized output tokens
_reserved = [0]           # in-flight reservations (bounds concurrent overshoot)
_spawned = [0]            # lifetime agent count
_collected = []           # {branch, diff} from worktree writers
_collected_lock = threading.Lock()
_phase = [None]
_warned_no_jsonschema = [False]

try:
    import jsonschema as _jsonschema
except ImportError:
    _jsonschema = None


class BudgetExceeded(Exception):
    pass


class AgentCapExceeded(Exception):
    pass


# Terminal failures are NOT retried and NOT swallowed to None — they propagate so a fan-out stops
# cleanly (a thrown ceiling is a stop signal, not a transient error). parallel() still catches them.
_TERMINAL = (BudgetExceeded, AgentCapExceeded)


def log(msg):
    """Progress narration, mirrors the Workflow tool's log(). Goes to stderr."""
    prefix = f"[{_phase[0]}] " if _phase[0] else ""
    print(f"[wf] {prefix}{msg}", file=sys.stderr, flush=True)


def phase(title):
    """Start a named progress group; subsequent log() lines are tagged with it (mirrors phase())."""
    _phase[0] = title


def tokens_used():
    """Approximate OUTPUT tokens across all agents this run (parsed from codex --json usage)."""
    with _tok_lock:
        return _tokens[0]


class _Budget:
    """Mirrors Workflow's budget object. total is None when unset (unlimited)."""
    @property
    def total(self):
        return BUDGET

    def spent(self):
        return tokens_used()

    def remaining(self):
        return float("inf") if BUDGET is None else max(0, BUDGET - tokens_used())


budget = _Budget()


def _reserve():
    """Admission control under the lock: enforce the lifetime cap and the budget ceiling against
    realized + in-flight spend, then reserve an estimate so concurrent agents can't all pass a
    stale gate and overshoot without bound (Workflow's ceiling is hard; this is its approximation)."""
    with _tok_lock:
        if _spawned[0] >= AGENT_CAP:
            raise AgentCapExceeded(f"lifetime agent cap {AGENT_CAP} reached")
        if BUDGET is not None and _tokens[0] + _reserved[0] >= BUDGET:
            raise BudgetExceeded(f"output-token budget {BUDGET} reached ({_tokens[0]} used)")
        _spawned[0] += 1
        _reserved[0] += EST_TOKENS


def _settle(actual):
    with _tok_lock:
        _reserved[0] -= EST_TOKENS
        _tokens[0] += max(0, actual)


def _iter_turn_usages(node):
    """Yield every `usage` dict attached to a turn.completed event in a codex --json stream,
    robust to event wrapping. turn.completed carries the authoritative per-turn token counts."""
    if isinstance(node, dict):
        t = node.get("type")
        if isinstance(t, str) and t.endswith("turn.completed") and isinstance(node.get("usage"), dict):
            yield node["usage"]
        for v in node.values():
            yield from _iter_turn_usages(v)
    elif isinstance(node, list):
        for v in node:
            yield from _iter_turn_usages(v)


def _output_tokens(usage):
    for k in ("output_tokens", "output", "completion_tokens"):
        if isinstance(usage.get(k), int):
            return usage[k]
    return None


def _parse_output_tokens(text):
    """OUTPUT tokens from a codex --json stream: sum turn.completed usages. Falls back to any
    usage dict, then to the human footer ('tokens used N', a TOTAL not output — last resort)."""
    total, found = 0, False
    for line in text.splitlines():
        line = line.strip()
        if not line or line[0] not in "{[":
            continue
        try:
            ev = json.loads(line)
        except ValueError:
            continue
        for usage in _iter_turn_usages(ev):
            ot = _output_tokens(usage)
            if ot is not None:
                total += ot
                found = True
    if found:
        return total
    m = re.findall(r"tokens used[\s:]*([\d,]+)", text)
    return int(m[-1].replace(",", "")) if m else 0


def _make_nullable(sub):
    """Widen a property's type to admit null, so an OpenAI-strict schema (which must list every
    property as required) can still let the model OMIT an optional field by emitting null."""
    if not isinstance(sub, dict):
        return
    t = sub.get("type")
    if isinstance(t, list):
        if "null" not in t:
            t.append("null")
    elif isinstance(t, str):
        sub["type"] = [t, "null"]


def _strict(node):
    """OpenAI structured output is strict: every object needs additionalProperties:false and ALL
    its properties listed in required. We honor that WITHOUT destroying optionality — a property
    the caller left out of `required` is widened to nullable (may be null) instead of forced."""
    if isinstance(node, dict):
        if node.get("type") == "object" and "properties" in node:
            props = node["properties"]
            required = set(node.get("required", []))
            for key, sub in props.items():
                if key not in required:
                    _make_nullable(sub)
            node["required"] = list(props.keys())
            node.setdefault("additionalProperties", False)
        for v in node.values():
            _strict(v)
    elif isinstance(node, list):
        for v in node:
            _strict(v)
    return node


def _extract_json(text):
    """Balanced-brace extraction of the first complete JSON object/array in text
    (robust to prose around it — unlike a greedy {.*} regex which merges blobs)."""
    start = next((i for i, c in enumerate(text) if c in "{["), None)
    if start is None:
        return text
    open_c, close_c = text[start], "}" if text[start] == "{" else "]"
    depth, instr, esc = 0, False, False
    for i in range(start, len(text)):
        c = text[i]
        if instr:
            esc = (c == "\\" and not esc)
            if c == '"' and not esc:
                instr = False
        elif c == '"':
            instr = True
        elif c == open_c:
            depth += 1
        elif c == close_c:
            depth -= 1
            if depth == 0:
                return text[start:i + 1]
    return text[start:]


def _run_once(prompt, schema, cwd, sandbox, model, effort):
    _reserve()
    actual = 0
    try:
        with tempfile.TemporaryDirectory() as td:
            out = os.path.join(td, "out.txt")
            logf = os.path.join(td, "log.txt")
            # -a never and -s are GLOBAL flags — must precede `exec` (codex rejects -a after exec).
            # --json gives us the structured event stream (turn.completed.usage) for OUTPUT-token
            # accounting instead of scraping a human footer. mcp_servers={} disables MCP: faster
            # startup AND restores --output-schema, which codex silently ignores when an MCP server
            # is active (openai/codex#15451).
            cmd = [
                "codex", "-a", "never", "-s", sandbox,
                "exec", "--json", "--skip-git-repo-check",
                "-c", f"model_reasoning_effort={effort or EFFORT}",
                "-c", "notify=[]",
                "-c", "mcp_servers={}",
                "-m", model or MODEL,
                "--cd", cwd,
                "-o", out,
            ]
            strict_schema = None
            if schema is not None:
                strict_schema = _strict(copy.deepcopy(schema))
                sf = os.path.join(td, "schema.json")
                with open(sf, "w") as f:
                    json.dump(strict_schema, f)
                cmd += ["--output-schema", sf]
            cmd.append(prompt)
            # Child output -> FILE, never PIPE: codex emits enough text that pipe buffers (~64KB)
            # fill under concurrency and subprocess.run deadlocks. The result comes from -o anyway.
            with open(logf, "w") as lf:
                subprocess.run(cmd, check=True, stdout=lf, stderr=lf, timeout=TIMEOUT)
            actual = _parse_output_tokens(open(logf).read())
            text = open(out).read().strip()
            if schema is None:
                return text
            # --output-schema yields pure JSON; parse directly, fall back to balanced extraction.
            try:
                obj = json.loads(text)
            except json.JSONDecodeError:
                obj = json.loads(_extract_json(text))
            if _jsonschema is not None:
                # validate against the SAME strictened schema codex was constrained to, so the
                # check reflects the actual generation contract (nullable optionals included).
                _jsonschema.validate(obj, strict_schema)  # raises -> caller's retry loop re-runs
            elif not _warned_no_jsonschema[0]:
                _warned_no_jsonschema[0] = True
                log("WARNING: jsonschema not installed — structured output is returned UNVALIDATED")
            return obj
    finally:
        _settle(actual)
        # Don't let a spend cap silently vanish: if a budget is set and we couldn't read usage
        # (parse miss, timeout, crash before the footer), say so instead of counting 0.
        if BUDGET is not None and actual == 0:
            log("WARNING: token usage unreadable for an agent; budget undercounts this run")


# Writer agents edit files only: in exec mode under an on-request/cloud policy, shell
# commands (tests/builds) are auto-REJECTED ("approval not supported in exec mode"), so
# verification is the orchestrator's job after collecting diffs.
_WRITER_RULE = ("Make your changes by editing files only. Do NOT run shell commands "
                "(tests, builds, installers) — they are rejected in this mode; the "
                "orchestrator runs verification after collecting your diff.\n\n")

# Default (schema-less) agents return raw text that IS the return value — discourage chatty prose.
_RETURN_RULE = ("Your final message IS the return value of this task. Output only the requested "
                "result as raw data — no preamble, explanation, or sign-off.\n\n")


def agent(prompt, schema=None, cwd=None, sandbox="read-only", model=None, effort=None,
          role=None, isolation=None, retries=1, label=None):
    """One codex exec subagent.
      schema     JSON Schema dict -> returns a validated dict (else returns text). Optional fields
                 (any property not in `required`) stay optional — the model may emit null.
      role       short role framing prepended to the prompt (explorer/skeptic/etc).
      isolation  'worktree' -> run in a fresh git worktree (sandbox forced workspace-write); the
                 worktree+branch are auto-removed after the diff is captured (no leak, no manual
                 cleanup), the diff is recorded in collected_diffs(), and the agent's bare result
                 is returned — same shape as a normal call.
      retries    extra attempts on nonzero exit / parse / validation failure (default 1).
      label      display label used in failure logs.
    Returns None after exhausting retries (mirrors Workflow's agent() returning null on terminal
    failure); a budget/cap ceiling propagates instead, to stop the fan-out."""
    parts = []
    if role:
        parts.append(f"You are acting as: {role}.")
    if schema is None and isolation != "worktree":
        parts.append(_RETURN_RULE.strip())
    parts.append(prompt)
    full = "\n\n".join(parts)
    last = None
    for _ in range(retries + 1):
        try:
            if isolation == "worktree":
                return _agent_worktree(full, schema, model, effort)
            return _run_once(full, schema, cwd or CWD, sandbox, model, effort)
        except _TERMINAL:
            raise
        except Exception as e:
            last = e
    log(f"agent failed after {retries + 1} attempts{f' [{label}]' if label else ''}: {last}")
    return None


def _remove_worktree(base, wt, branch):
    subprocess.run(["git", "-C", base, "worktree", "remove", "--force", wt],
                   stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    subprocess.run(["git", "-C", base, "branch", "-D", branch],
                   stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)


def _agent_worktree(prompt, schema, model, effort, base=None):
    base = base or CWD
    wt = tempfile.mkdtemp(prefix="cxwt-")
    branch = "cxwf/" + os.path.basename(wt)
    subprocess.run(["git", "-C", base, "worktree", "add", "-q", "-b", branch, wt, "HEAD"],
                   check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    try:
        result = _run_once(_WRITER_RULE + prompt, schema, wt, "workspace-write", model, effort)
        subprocess.run(["git", "-C", wt, "add", "-A"], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        diff = subprocess.run(["git", "-C", wt, "diff", "--cached"], capture_output=True, text=True).stdout
    except Exception:
        _remove_worktree(base, wt, branch)
        raise
    # Success: the orchestrator applies the captured diff TEXT, not the live tree, so tear the
    # worktree+branch down now (auto-clean, including the unchanged case — empty diff records
    # nothing). This is the parity behavior: caller never manages teardown of successful worktrees.
    _remove_worktree(base, wt, branch)
    if diff.strip():
        with _collected_lock:
            _collected.append({"branch": branch, "diff": diff})
    return result


def collected_diffs():
    """Diffs produced by isolation='worktree' writers this run, in completion order.
    Each is {'branch', 'diff'}; apply/verify these in the orchestrator."""
    with _collected_lock:
        return list(_collected)


def _cleanup_leaked_worktrees():
    """atexit backstop: prune any cxwf/ worktrees the success/failure paths somehow missed
    (e.g. a hard interrupt between `worktree add` and teardown)."""
    out = subprocess.run(["git", "-C", CWD, "worktree", "list", "--porcelain"],
                         capture_output=True, text=True).stdout
    for line in out.splitlines():
        if line.startswith("worktree ") and os.path.basename(line[9:]).startswith("cxwt-"):
            wt = line[9:]
            _remove_worktree(CWD, wt, "cxwf/" + os.path.basename(wt))


atexit.register(_cleanup_leaked_worktrees)


def parallel(thunks):
    """Run thunks (0-arg callables) concurrently, capped at CONCURRENCY.
    Promise.allSettled semantics: a thunk that raises yields None, never aborts the batch."""
    thunks = list(thunks)
    if len(thunks) > MAX_BATCH:
        raise ValueError(f"parallel() accepts at most {MAX_BATCH} items, got {len(thunks)}")
    futs = [_pool.submit(t) for t in thunks]
    results = []
    for f in futs:
        try:
            results.append(f.result())
        except Exception as e:
            results.append(None)
            log(f"agent failed: {e}")
    return results


def pipeline(items, *stages):
    """Each item flows through all stages independently (no barrier between stages).
    Stage signature: stage(prev_result, original_item, index). A stage that raises
    drops that item to None and skips its remaining stages."""
    items = list(items)
    if len(items) > MAX_BATCH:
        raise ValueError(f"pipeline() accepts at most {MAX_BATCH} items, got {len(items)}")

    def run_chain(item, idx):
        val = item
        for stage in stages:
            try:
                val = stage(val, item, idx)
            except Exception as e:
                log(f"item {idx} dropped: {e}")
                return None
        return val
    return parallel([(lambda it=it, i=i: run_chain(it, i)) for i, it in enumerate(items)])


# --- Durable run trail (the start_run/save_result/write_ledger API the ultracode skill expects) ---

_RUNS_ROOT = os.path.join(CWD, ".codex", "ultracode", "runs")
_current_run = [None]


def start_run(label=None):
    """Mint a run id and create .codex/ultracode/runs/<id>/ (with results/). Returns (id, path)."""
    run_id = f"{int(time.time())}-{uuid.uuid4().hex[:6]}"
    path = os.path.join(_RUNS_ROOT, run_id)
    os.makedirs(os.path.join(path, "results"), exist_ok=True)
    with open(os.path.join(path, "run.json"), "w") as f:
        json.dump({"id": run_id, "label": label, "started": int(time.time())}, f, indent=2)
    _current_run[0] = path
    return run_id, path


def save_result(name, obj):
    """Persist a structured result to results/<name>.json under the current run."""
    if _current_run[0] is None:
        start_run()
    with open(os.path.join(_current_run[0], "results", f"{name}.json"), "w") as f:
        json.dump(obj, f, indent=2)


def write_ledger(sections):
    """Write ledger.md under the current run. `sections` is {heading: body} or a markdown string."""
    if _current_run[0] is None:
        start_run()
    if isinstance(sections, dict):
        body = "\n\n".join(f"## {h}\n\n{b}" for h, b in sections.items())
    else:
        body = str(sections)
    with open(os.path.join(_current_run[0], "ledger.md"), "w") as f:
        f.write(body + "\n")