File size: 7,278 Bytes
3b0d6b3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 | import json
import logging
import os
import re
import shlex
import shutil
import subprocess
import tempfile
import time
from config import CODEX_CLI_PATH, get_pricing
from models import RunStatus
from adapters.base import AgentAdapter, RunContext, RunResult
logger = logging.getLogger(__name__)
# Codex prints token usage like: "tokens used\n1.019"
_TOKENS_USED_RE = re.compile(r"tokens\s+used\s*\n\s*([0-9.]+)", re.IGNORECASE)
# Broader alternatives
_TOKEN_ALT_RE = re.compile(r"(?:total\s+)?tokens[:\s]+(\d[\d,]*)", re.I)
_COST_ALT_RE = re.compile(r"(?:cost|price|total)[:\s]*\$([0-9.]+)", re.I)
# Default model assumed when Codex doesn't report one
_DEFAULT_MODEL = "gpt-4o"
class CodexAdapter(AgentAdapter):
def name(self) -> str:
return "Codex"
def is_available(self) -> bool:
return shutil.which(CODEX_CLI_PATH) is not None
async def run(self, ctx: RunContext, sandbox) -> RunResult:
"""Run Codex CLI on the host against a local clone of the task repo.
Codex uses ChatGPT Plus OAuth (no API key needed).
Runs non-interactively via `codex exec -s full-auto`.
"""
start = time.time()
tmpdir = tempfile.mkdtemp(prefix="arena-codex-")
try:
# Clone repo to host temp dir.
repo_url = ctx.repo_url or ""
repo_ref = ctx.repo_ref or "main"
if repo_url:
clone_result = subprocess.run(
["git", "clone", "--depth=1", "-b", repo_ref, repo_url, tmpdir],
capture_output=True, text=True, timeout=120,
)
if clone_result.returncode != 0:
return RunResult(
status=RunStatus.error,
duration_s=round(time.time() - start, 2),
raw_output=clone_result.stderr[:5000],
)
work_dir = tmpdir
else:
work_dir = ctx.work_dir
# Git checkpoint for diff.
subprocess.run(
["git", "tag", "arena-checkpoint"],
cwd=work_dir, capture_output=True, timeout=10,
)
# Use model from agent config, fall back to default.
model = self.config.get("model", _DEFAULT_MODEL)
prompt = (
f"Fix the following issue in the codebase:\n\n"
f"{ctx.task_prompt}\n\n"
f"After fixing, verify your changes pass: {ctx.success_cmd}"
)
cmd = [
CODEX_CLI_PATH,
"exec",
"-s", "danger-full-access",
"--ephemeral",
"--model", model,
prompt,
]
try:
proc = subprocess.run(
cmd,
cwd=work_dir,
capture_output=True,
text=True,
timeout=ctx.time_limit_s,
)
returncode = proc.returncode
stdout = proc.stdout
stderr = proc.stderr
except subprocess.TimeoutExpired:
returncode = 124
stdout = ""
stderr = f"Command timed out after {ctx.time_limit_s}s"
duration = time.time() - start
# Collect diff, exclude __pycache__ and .pyc.
def _get_diff(rev):
r = subprocess.run(
["git", "diff", rev, "--", ".", ":!**/__pycache__/", ":!**/*.pyc"],
cwd=work_dir, capture_output=True, text=True, timeout=30,
)
return r.stdout if r.returncode == 0 else ""
diff_text = _get_diff("arena-checkpoint")
if not diff_text.strip():
diff_text = _get_diff("arena-checkpoint..HEAD")
# Parse token usage from output.
tokens_input = 0
tokens_output = 0
cost_usd = 0.0
combined = stdout + "\n" + stderr
# Log raw output for diagnostics
logger.info("CODEX RAW OUTPUT (first 2000 chars):\n%s", combined[:2000])
# Codex prints "tokens used\n1.019" (in thousands)
m = _TOKENS_USED_RE.search(combined)
if m:
total_k = float(m.group(1))
total_tokens = int(total_k * 1000)
# Rough split: 70% input, 30% output
tokens_input = int(total_tokens * 0.7)
tokens_output = total_tokens - tokens_input
# Fallback: try alternative token/cost patterns
if tokens_input == 0 and tokens_output == 0:
# Filter out metadata lines (e.g., "Repo-map: using 1024 tokens")
content_lines = [
line for line in combined.split("\n")
if not re.search(r"repo.?map|using \d+ tokens.*refresh", line, re.I)
]
filtered = "\n".join(content_lines)
alt_tokens = _TOKEN_ALT_RE.findall(filtered)
if alt_tokens:
total = sum(int(t.replace(",", "")) for t in alt_tokens)
if total >= 100: # Sanity check: real LLM usage > 100 tokens
tokens_input = int(total * 0.7)
tokens_output = total - tokens_input
if cost_usd == 0.0:
alt_costs = _COST_ALT_RE.findall(combined)
if alt_costs:
cost_usd = sum(float(c) for c in alt_costs)
# Estimate cost from tokens using dynamic pricing lookup.
if cost_usd == 0.0 and (tokens_input or tokens_output):
rate_in, _, _, rate_out = get_pricing(model)
cost_usd = (tokens_input * rate_in + tokens_output * rate_out) / 1_000_000
# Last resort: estimate from raw output length (better than $0)
if cost_usd == 0.0 and (tokens_input + tokens_output) < 50 and len(combined) > 100:
estimated_output_tokens = len(combined) // 4
rate_in, _, _, rate_out = get_pricing(model)
tokens_input = estimated_output_tokens * 3
tokens_output = estimated_output_tokens
cost_usd = (tokens_input * rate_in + tokens_output * rate_out) / 1_000_000
logger.warning("Cost estimated from output length (%d chars) — no token data found", len(combined))
trajectory = [{"tool_name": "Bash", "input_preview": f"codex exec --model {model}"}]
if duration >= ctx.time_limit_s:
status = RunStatus.timeout
elif returncode == 0:
status = RunStatus.success
else:
status = RunStatus.failure
return RunResult(
status=status,
duration_s=round(duration, 2),
tokens_input=tokens_input,
tokens_output=tokens_output,
cost_usd=round(cost_usd, 6),
diff_text=diff_text,
trajectory=trajectory,
raw_output=(stdout + stderr)[:50_000],
)
finally:
shutil.rmtree(tmpdir, ignore_errors=True)
|