tai/tai_pkg/codex.py

"""Codex CLI backend."""

import json
import os
import shutil
import subprocess
import tempfile

from tai_pkg.clirun import Cli
from tai_pkg.errors import BackendError
from tai_pkg.muse import Muse
from tai_pkg.opencode import Opencode
from tai_pkg.text import Text

__all__ = ["Codex"]


class Codex:
    """Codex."""

    @staticmethod
    def codex_binary():
        return os.environ.get("CODEX_BIN") or shutil.which("codex")

    @staticmethod
    def codex_jsonl_texts(stdout):
        texts = []
        last_error = ""
        for line in (stdout or "").splitlines():
            line = line.strip()
            if not line.startswith("{"):
                continue
            try:
                event = json.loads(line)
            except ValueError:
                continue
            if not isinstance(event, dict):
                continue
            etype = event.get("type")
            if etype == "item.completed":
                item = event.get("item") or {}
                if item.get("type") == "agent_message" and item.get("text"):
                    texts.append(item["text"])
            elif etype in ("turn.failed", "error"):
                msg = event.get("message") or event.get("error") or ""
                if isinstance(msg, dict):
                    msg = msg.get("message") or json.dumps(msg)
                if msg:
                    last_error = str(msg)[:300]
        return texts, last_error

    @staticmethod
    def codex_parse_output(stdout, outfile_text="", used_schema=False):
        if isinstance(outfile_text, str) and outfile_text.strip():
            if used_schema:
                try:
                    answer = json.loads(outfile_text)
                except ValueError:
                    answer = None
                if isinstance(answer, dict) and ("content" in answer or "tool_calls" in answer):
                    return answer
            return Opencode.opencode_parse(outfile_text)
        texts, last_error = Codex.codex_jsonl_texts(stdout)
        if texts:
            return Opencode.opencode_parse("".join(texts))
        if last_error:
            raise BackendError("codex: %s" % last_error)
        return Opencode.opencode_parse(stdout)

    @staticmethod
    def codex_run(prompt, timeout):
        binary = Codex.codex_binary()
        if not binary:
            raise BackendError("codex CLI not found (CODEX_BIN or PATH)")
        schema = json.dumps(Muse.muse_output_schema())
        with tempfile.TemporaryDirectory(prefix="tai-codex-") as folder:
            schema_path = os.path.join(folder, "schema.json")
            out_path = os.path.join(folder, "last.txt")
            with open(schema_path, "w", encoding="utf-8") as handle:
                handle.write(schema)
            base = [binary, "exec", "--json", "--ephemeral", "--sandbox", "read-only", "--skip-git-repo-check"]
            if len(prompt.encode("utf-8", "replace")) > 120000:
                variants = [
                    (base + ["--output-schema", schema_path, "-o", out_path, "-"], prompt, True),
                    (base + ["-o", out_path, "-"], prompt, False),
                ]
            else:
                variants = [
                    (base + ["--output-schema", schema_path, "-o", out_path, prompt], None, True),
                    (base + ["-o", out_path, prompt], None, False),
                    (base + [prompt], None, False),
                ]
            last_error = ""
            for argv, prompt_input, used_schema in variants:
                try:
                    done = subprocess.run(argv, input=prompt_input, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True, cwd=folder, timeout=timeout)
                except subprocess.TimeoutExpired:
                    raise BackendError("codex run timed out after %ds" % timeout)
                except OSError as exc:
                    raise BackendError("codex run failed: %s" % Text.short_error(exc))
                if done.returncode != 0 and Cli._cli_unknown_flag(done.stderr) and (argv, prompt_input, used_schema) != variants[-1]:
                    last_error = (done.stderr or "").strip()[:200]
                    continue
                outfile_text = ""
                if "-o" in argv:
                    try:
                        with open(out_path, "r", encoding="utf-8", errors="replace") as handle:
                            outfile_text = handle.read()
                    except OSError:
                        outfile_text = ""
                if done.returncode != 0 and not (outfile_text.strip() or (done.stdout or "").strip()):
                    blob = ((done.stderr or "") + "\n" + (done.stdout or "")).strip()
                    if Cli._cli_auth_failure(blob):
                        raise BackendError("codex not authenticated (run codex login): %s" % blob[:200])
                    raise BackendError("codex exit %d: %s" % (done.returncode, blob[:300] or last_error or "no output"))
                if not (outfile_text.strip() or (done.stdout or "").strip()):
                    blob = (done.stderr or "").strip()
                    raise BackendError("codex returned no output: %s" % (blob[:300] or last_error or "empty stdout"))
                return done.stdout or "", outfile_text, used_schema
        raise BackendError("codex failed: %s" % (last_error or "no usable output"))

    @staticmethod
    def codex_health():
        binary = Codex.codex_binary()
        if not binary:
            return "missing: codex CLI not found"
        try:
            done = subprocess.run([binary, "login", "status"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True, timeout=15)
        except (OSError, subprocess.SubprocessError) as exc:
            return "fail: %s" % Text.short_error(exc)[:60]
        detail = Cli._cli_version_detail(done, "codex")
        if done.returncode != 0:
            return "not logged in: %s" % detail
        return "ok, subscription (%s)" % detail