runner.py

#!/usr/bin/env python3
# DEPRECATED: This file is being replaced by `ostk bench` (the kernel-native runner).
# See docs/OSTK_BENCH_SPEC.md for the migration plan.
# runner.py will be removed once ostk bench supports Docker execution natively.
"""needle-bench runner — agent loop for benchmark evaluation."""

import argparse
import json
import os
import subprocess
import sys
import time
import urllib.request

ANTHROPIC_ENDPOINT = "https://api.anthropic.com/v1/messages"
GOOGLE_ENDPOINT = "https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent"
OPENROUTER_ENDPOINT = "https://openrouter.ai/api/v1/chat/completions"

# OpenRouter model name mapping (needle-bench short name → OpenRouter path)
# Every model in run_missing.sh must have an entry here.
OPENROUTER_MODELS = {
    # Anthropic
    "claude-haiku-3-5-20251001": "anthropic/claude-haiku-4.5",
    "claude-haiku-3-5-20241022": "anthropic/claude-haiku-4.5",
    "claude-haiku-4-5": "anthropic/claude-haiku-4.5",
    "claude-sonnet-4-6": "anthropic/claude-sonnet-4.6",
    "claude-opus-4-6": "anthropic/claude-opus-4.6",
    # Google
    "gemini-2.5-flash": "google/gemini-2.5-flash",
    "gemini-2.5-pro": "google/gemini-2.5-pro-preview",
    "gemini-3-flash-preview": "google/gemini-3-flash-preview",
    "gemini-3.1-pro-preview": "google/gemini-3.1-pro-preview",
    # OpenAI
    "gpt-4.1": "openai/gpt-4.1",
    "gpt-5-codex": "openai/gpt-5-codex",
    "o3": "openai/o3",
    "o4-mini": "openai/o4-mini",
    # xAI (Grok)
    "grok-3": "x-ai/grok-3",
    "grok-3-fast": "x-ai/grok-3-fast",
    "grok-3-mini": "x-ai/grok-3-mini",
    "grok-4": "x-ai/grok-4",
    "grok-4-fast": "x-ai/grok-4-fast",
    "grok-4.1-fast": "x-ai/grok-4.1-fast",
    "grok-4.20": "x-ai/grok-4.20",
    "grok-code-fast-1": "x-ai/grok-code-fast-1",
    # Mistral
    "codestral-2508": "mistralai/codestral-2508",
    "devstral-2512": "mistralai/devstral-2512",
    "devstral-medium": "mistralai/devstral-medium",
    "devstral-small-latest": "mistralai/devstral-small-latest",
    # DeepSeek
    "deepseek-r1": "deepseek/deepseek-r1",
    "deepseek-r1-0528": "deepseek/deepseek-r1-0528",
    "deepseek-v3.2": "deepseek/deepseek-v3.2",
    # Qwen
    "qwen3-coder": "qwen/qwen3-coder",
    "qwen3-coder-flash": "qwen/qwen3-coder-flash",
    "qwen3-coder-plus": "qwen/qwen3-coder-plus",
    # Moonshot AI (Kimi)
    "kimi-k2.5": "moonshotai/kimi-k2.5",
    "kimi-k2-0905": "moonshotai/kimi-k2-0905",
    "kimi-k2-thinking": "moonshotai/kimi-k2-thinking",
    # Meta
    "llama-4-maverick": "meta-llama/llama-4-maverick",
}

# Reverse mapping: OpenRouter path → canonical short name
_OPENROUTER_REVERSE = {v: k for k, v in OPENROUTER_MODELS.items()}


def _canonical_agent_name(model: str) -> str:
    """Return the short canonical name for a model, regardless of how it was specified.

    If the user passed an OpenRouter path like 'anthropic/claude-opus-4.6',
    map it back to the short name 'claude-opus-4-6'.
    """
    if model in _OPENROUTER_REVERSE:
        return _OPENROUTER_REVERSE[model]
    # If it looks like an OpenRouter path but isn't in our map, take the part after '/'
    if "/" in model:
        return model.split("/", 1)[1].replace(".", "-")
    return model


# Estimated cost per 1M tokens (input, output) — USD.
# Aligned with consolidate_scores.py RATE_CARD.
MODEL_PRICING = {
    # Anthropic
    "claude-opus-4-6": (5.0, 25.0),
    "claude-sonnet-4-6": (3.0, 15.0),
    "claude-haiku-4-5": (1.0, 5.0),
    # Google
    "gemini-2.5-flash": (0.30, 2.50),
    "gemini-2.5-pro": (1.25, 10.0),
    "gemini-3-flash-preview": (0.50, 3.00),
    "gemini-3.1-pro-preview": (2.00, 12.00),
    # OpenAI
    "gpt-4.1": (2.00, 8.00),
    "gpt-5-codex": (1.25, 10.00),
    "o3": (2.00, 8.00),
    "o4-mini": (1.10, 4.40),
    # xAI (Grok)
    "grok-3": (3.00, 15.00),
    "grok-3-fast": (0.20, 0.50),
    "grok-3-mini": (0.30, 0.50),
    "grok-4": (3.00, 15.00),
    "grok-4-fast": (0.20, 0.50),
    "grok-4.1-fast": (0.20, 0.50),
    "grok-4.20": (2.00, 6.00),
    "grok-code-fast-1": (0.20, 1.50),
    # Mistral
    "codestral-2508": (0.30, 0.90),
    "devstral-2512": (0.40, 2.00),
    "devstral-medium": (0.40, 2.00),
    "devstral-small-latest": (0.10, 0.30),
    # DeepSeek
    "deepseek-r1": (0.70, 2.50),
    "deepseek-r1-0528": (0.45, 2.15),
    "deepseek-v3.2": (0.26, 0.38),
    # Qwen
    "qwen3-coder": (0.22, 1.00),
    "qwen3-coder-flash": (0.20, 0.97),
    "qwen3-coder-plus": (0.65, 3.25),
    # Moonshot AI (Kimi)
    "kimi-k2.5": (0.40, 1.99),
    # Meta
    "llama-4-maverick": (0.15, 0.60),
    # defaults
    "_default_input": 3.0,
    "_default_output": 15.0,
}


def _model_cost_usd(model: str, input_tokens: int, output_tokens: int) -> float:
    """Estimate cost in USD for a given token usage."""
    # Try exact match, then prefix match
    pricing = None
    for key, val in MODEL_PRICING.items():
        if key.startswith("_"):
            continue
        if model == key or model.startswith(key):
            pricing = val
            break
    if pricing is None:
        pricing = (MODEL_PRICING["_default_input"], MODEL_PRICING["_default_output"])
    cost_in = input_tokens * pricing[0] / 1_000_000
    cost_out = output_tokens * pricing[1] / 1_000_000
    return round(cost_in + cost_out, 6)


# ── Token / cache-split accounting (task #11) ────────────────────────────
# Cache billing multipliers relative to the base input rate, matching the
# downstream RATE CARD in consolidate_harden.py / emit_cells.py (→2062 SCHEMA
# PARITY): cache reads bill at 0.1x, 5-minute cache writes at 1.25x, 1-hour
# cache writes at 2.0x. Unsplit cache-creation is priced as 5m (the Anthropic
# default TTL). Providers without a cache-creation concept report 0 for those
# buckets, so their billed-equivalent reduces to fresh input.
CACHE_READ_MULTIPLIER = 0.10
CACHE_CREATE_5M_MULTIPLIER = 1.25
CACHE_CREATE_1H_MULTIPLIER = 2.00


def extract_usage(resp: dict) -> dict:
    """Normalize per-turn token usage across providers.

    Returns the four canonical Anthropic usage fields plus the 5m/1h
    cache-creation breakdown, defaulting every field to 0 so that older
    responses or providers that omit cache accounting never raise. Anthropic
    populates all of these directly (the 5m/1h split lives under the nested
    ``cache_creation`` object); Google/OpenRouter surface the cache-read field
    when available (see call_google / call_openrouter).
    """
    u = (resp or {}).get("usage", {}) or {}
    creation = u.get("cache_creation") or {}
    return {
        "input_tokens": u.get("input_tokens", 0) or 0,
        "output_tokens": u.get("output_tokens", 0) or 0,
        "cache_read_input_tokens": u.get("cache_read_input_tokens", 0) or 0,
        "cache_creation_input_tokens": u.get("cache_creation_input_tokens", 0) or 0,
        "cache_creation_5m_tokens": creation.get("ephemeral_5m_input_tokens", 0) or 0,
        "cache_creation_1h_tokens": creation.get("ephemeral_1h_input_tokens", 0) or 0,
    }


def billed_equivalent_input_tokens(fresh: int, cache_read: int,
                                   cache_create_total: int,
                                   cc_5m: int = 0, cc_1h: int = 0) -> float:
    """Input-token count weighted by the cache billing multipliers.

    Mirrors consolidate_harden.bucket_cost's input side: unsplit cache-create
    (total minus the known 5m/1h buckets) is priced at the 5m rate.
    """
    cc_unsplit = max(0, cache_create_total - cc_5m - cc_1h)
    return (fresh
            + cache_read * CACHE_READ_MULTIPLIER
            + cc_5m * CACHE_CREATE_5M_MULTIPLIER
            + cc_1h * CACHE_CREATE_1H_MULTIPLIER
            + cc_unsplit * CACHE_CREATE_5M_MULTIPLIER)


class PostRecorder:
    """Writes structured POST (power-on self-test) records per benchmark run.

    Output: runs/<model>/<benchmark>/post.jsonl
    Events: post.start, agent.bash, agent.edit, post.end
    """

    def __init__(self, runs_dir: str, bench_name: str, model: str):
        run_dir = os.path.join(runs_dir, bench_name)
        os.makedirs(run_dir, exist_ok=True)
        self._path = os.path.join(run_dir, "post.jsonl")
        self._f = open(self._path, "w")
        self._bench = bench_name
        self._model = model

    def _emit(self, record: dict):
        self._f.write(json.dumps(record) + "\n")
        self._f.flush()

    def start(self, initial_test_output: str, prompt: str):
        self._emit({
            "event": "post.start",
            "benchmark": self._bench,
            "model": self._model,
            "initial_test_output": initial_test_output[:2000],
            "prompt": prompt,
            "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        })

    def bash(self, cmd: str, output: str, turn: int):
        self._emit({
            "event": "agent.bash",
            "cmd": cmd[:500],
            "output_preview": output[:500],
            "turn": turn,
        })

    def read(self, path: str, turn: int):
        self._emit({
            "event": "agent.read",
            "path": path,
            "turn": turn,
        })

    def edit(self, path: str, old_str: str, new_str: str, turn: int):
        self._emit({
            "event": "agent.edit",
            "path": path,
            "old_str_preview": old_str[:200],
            "new_str_preview": new_str[:200],
            "turn": turn,
        })

    def end(self, resolved: bool, final_test_output: str, turns: int):
        self._emit({
            "event": "post.end",
            "resolved": resolved,
            "final_test_output": final_test_output[:2000],
            "turns": turns,
            "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        })

    def close(self):
        self._f.close()

    @property
    def path(self) -> str:
        return self._path


class MetricsRecorder:
    """Tracks and records token consumption per benchmark run.

    Output: runs/<model>/<benchmark>/metrics.jsonl
    Events: token.usage (per turn), run.complete (summary)
    """

    def __init__(self, runs_dir: str, bench_name: str, model: str):
        run_dir = os.path.join(runs_dir, bench_name)
        os.makedirs(run_dir, exist_ok=True)
        self._path = os.path.join(run_dir, "metrics.jsonl")
        self._f = open(self._path, "w")
        self._model = model
        self._cum_in = 0
        self._cum_out = 0
        self._cum_cache_read = 0
        self._cum_cache_creation = 0
        self._cum_cc_5m = 0
        self._cum_cc_1h = 0

    def _emit(self, record: dict):
        self._f.write(json.dumps(record) + "\n")
        self._f.flush()

    def record_turn(self, turn: int, input_tokens: int, output_tokens: int,
                    cache_read_input_tokens: int = 0,
                    cache_creation_input_tokens: int = 0,
                    cache_creation_5m_tokens: int = 0,
                    cache_creation_1h_tokens: int = 0):
        self._cum_in += input_tokens
        self._cum_out += output_tokens
        self._cum_cache_read += cache_read_input_tokens
        self._cum_cache_creation += cache_creation_input_tokens
        self._cum_cc_5m += cache_creation_5m_tokens
        self._cum_cc_1h += cache_creation_1h_tokens
        self._emit({
            "event": "token.usage",
            "turn": turn,
            "input_tokens": input_tokens,
            "output_tokens": output_tokens,
            "cache_read_input_tokens": cache_read_input_tokens,
            "cache_creation_input_tokens": cache_creation_input_tokens,
            "cache_creation_5m_tokens": cache_creation_5m_tokens,
            "cache_creation_1h_tokens": cache_creation_1h_tokens,
            "cumulative_input": self._cum_in,
            "cumulative_output": self._cum_out,
            "cumulative_cache_read": self._cum_cache_read,
            "cumulative_cache_creation": self._cum_cache_creation,
        })

    def complete(self):
        total = self._cum_in + self._cum_out
        cost = _model_cost_usd(self._model, self._cum_in, self._cum_out)
        ratio = (round(self._cum_cache_read / self._cum_cache_creation, 4)
                 if self._cum_cache_creation else None)
        billed = billed_equivalent_input_tokens(
            self._cum_in, self._cum_cache_read, self._cum_cache_creation,
            self._cum_cc_5m, self._cum_cc_1h)
        self._emit({
            "event": "run.complete",
            "total_input_tokens": self._cum_in,
            "total_output_tokens": self._cum_out,
            "total_cache_read_input_tokens": self._cum_cache_read,
            "total_cache_creation_input_tokens": self._cum_cache_creation,
            "cache_read_to_create_ratio": ratio,
            "total_billed_equivalent_input_tokens": round(billed, 2),
            "total_tokens": total,
            "estimated_cost_usd": cost,
        })
        return self._cum_in, self._cum_out, total, cost

    @property
    def cum_cache_read(self) -> int:
        return self._cum_cache_read

    @property
    def cum_cache_creation(self) -> int:
        return self._cum_cache_creation

    @property
    def cum_cache_creation_5m(self) -> int:
        return self._cum_cc_5m

    @property
    def cum_cache_creation_1h(self) -> int:
        return self._cum_cc_1h

    def close(self):
        self._f.close()

    @property
    def path(self) -> str:
        return self._path

# ── Unified bench prompts (SPEC-bench-v2 §1) ────────────────────────────
# Every arm (native, kernel, kernel-cpu) receives exactly these prompts.
# No silent context, no BOOT injection, no extra instructions.
BENCH_SYSTEM_PROMPT = (
    "This is a needle-bench benchmark. There is a bug hidden in this codebase. "
    "Find it and fix it. Run test.sh to verify your fix."
)
BENCH_USER_PROMPT = "find the needle. run test.sh to verify your fix."

# Legacy aliases — kept so existing callers don't break during migration
DEFAULT_SYSTEM_PROMPT = BENCH_SYSTEM_PROMPT
SYSTEM_PROMPT = BENCH_SYSTEM_PROMPT
DEFAULT_INSTANCE = BENCH_USER_PROMPT

# ---------------------------------------------------------------------------
# Tool definitions — keyed by the API tool name the LLM sees.
# ---------------------------------------------------------------------------
ALL_TOOLS = {
    "bash": {"name": "bash", "description": "Run a bash command",
             "input_schema": {"type": "object", "properties": {"command": {"type": "string", "description": "The bash command to run"}}, "required": ["command"]}},
    "read": {"name": "read", "description": "Read a file from the workspace.",
             "input_schema": {"type": "object", "properties": {"path": {"type": "string", "description": "Absolute path to read"}}, "required": ["path"]}},
    "edit": {"name": "edit", "description": "Edit a file by replacing old_str with new_str",
             "input_schema": {"type": "object", "properties": {"path": {"type": "string", "description": "Path to the file"}, "old_str": {"type": "string", "description": "Exact string to find"}, "new_str": {"type": "string", "description": "Replacement string"}}, "required": ["path", "old_str", "new_str"]}},
}

# Mapping: Agentfile canonical TOOL names → API tool names the LLM sees.
AGENTFILE_TO_API = {
    # v2.0 canonical names
    "shell":      "bash",
    "file:read":  "read",
    "file:edit":  "edit",
    "file:write": "edit",   # alias
    # v1.x compat
    "sh_run":     "bash",
    "ss":         "edit",
    "ss_session": "read",
    "bash":       "bash",
    "read":       "read",
    "edit":       "edit",
}

# Fallback: if no TOOL directives or none resolve, use all tools.
DEFAULT_TOOLS = list(ALL_TOOLS.values())


def resolve_tools(agentfile_tools):
    """Map Agentfile TOOL directives to API tool definitions.

    Returns a list of tool dicts suitable for the messages API.
    Unknown tool names (e.g. 'spawn', 'interact') are silently skipped —
    they are haystack-level tools, not runner-level tools.
    """
    if not agentfile_tools:
        return DEFAULT_TOOLS
    api_names = set()
    for tool in agentfile_tools:
        api_name = AGENTFILE_TO_API.get(tool)
        if api_name:
            api_names.add(api_name)
    if not api_names:
        return DEFAULT_TOOLS
    return [ALL_TOOLS[name] for name in ALL_TOOLS if name in api_names]


def _resolve_var(value):
    """Resolve ${VAR:-default} syntax, returning the raw string.

    If value matches ${VAR:-default}, check os.environ for VAR and fall back
    to the default.  Otherwise return the value unchanged.
    """
    import re
    m = re.match(r'^\$\{(\w+):-([^}]*)\}$', value)
    if m:
        env_name, default = m.group(1), m.group(2)
        return os.environ.get(env_name, default)
    m = re.match(r'^\$\{(\w+)\}$', value)
    if m:
        return os.environ.get(m.group(1), value)
    return value


def parse_agentfile(path):
    cfg = {"tools": [], "limits": {}, "prompt": None, "from_image": None, "boot": None}
    with open(path) as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            parts = line.split(None, 1)
            directive, rest = parts[0], parts[1] if len(parts) > 1 else ""
            if directive == "FROM":
                cfg["from_image"] = _resolve_var(rest)
            elif directive == "TOOL":
                cfg["tools"].append(rest)
            elif directive == "LIMIT":
                k, v = rest.split(None, 1)
                cfg["limits"][k] = int(_resolve_var(v))
            elif directive == "PROMPT":
                cfg["prompt"] = rest
            elif directive == "BOOT":
                cfg["boot"] = rest
    return cfg


def solution_files(bench_dir):
    patch = os.path.join(bench_dir, ".bench", "solution.patch")
    files = []
    if os.path.exists(patch):
        with open(patch) as f:
            for line in f:
                if line.startswith("+++ b/"):
                    files.append(line[6:].strip())
    return files


def docker_exec(container, cmd):
    # Inject scoped secrets via env — resolved from host env, never in image
    env_flags = []
    gh_token = os.environ.get("GH_NEEDLE_BENCH_PROOF", "")
    if gh_token:
        env_flags = ["-e", f"GH_NEEDLE_BENCH_PROOF={gh_token}"]
    try:
        r = subprocess.run(["docker", "exec"] + env_flags + [container, "bash", "-c", cmd],
                           capture_output=True, text=True, timeout=120)
        return r.returncode, r.stdout[-4000:] if len(r.stdout) > 4000 else r.stdout, r.stderr[-2000:] if len(r.stderr) > 2000 else r.stderr
    except subprocess.TimeoutExpired:
        return 124, "timeout: command exceeded 120s limit", ""


def do_edit(container, path, old_str, new_str):
    script = (
        "import sys\n"
        f"p = {path!r}\n"
        "with open(p) as f: c = f.read()\n"
        f"o = {old_str!r}\n"
        "if o not in c:\n"
        "    print('old_str not found in ' + p, file=sys.stderr); sys.exit(1)\n"
        f"c = c.replace(o, {new_str!r}, 1)\n"
        "with open(p, 'w') as f: f.write(c)\n"
        "print('OK')\n"
    )
    r = subprocess.run(["docker", "exec", container, "python3", "-c", script],
                       capture_output=True, text=True, timeout=30)
    return r.returncode, r.stdout.strip(), r.stderr.strip()


def call_anthropic(model, messages, api_key, system_prompt=None, tools=None):
    body = json.dumps({
        "model": model, "max_tokens": 4096, "system": system_prompt or SYSTEM_PROMPT,
        "tools": tools if tools is not None else DEFAULT_TOOLS, "messages": messages,
    }).encode()
    req = urllib.request.Request(ANTHROPIC_ENDPOINT, data=body, headers={
        "Content-Type": "application/json",
        "x-api-key": api_key,
        "anthropic-version": "2023-06-01",
    })
    with urllib.request.urlopen(req, timeout=120) as resp:
        return json.loads(resp.read())


def call_google(model, messages, api_key, system_prompt=None, tools=None):
    _tools = tools if tools is not None else DEFAULT_TOOLS
    _sys_prompt = system_prompt or SYSTEM_PROMPT
    contents = []
    for m in messages:
        role = "user" if m["role"] == "user" else "model"
        if isinstance(m["content"], str):
            contents.append({"role": role, "parts": [{"text": m["content"]}]})
        elif isinstance(m["content"], list):
            parts = []
            for block in m["content"]:
                if block.get("type") == "text":
                    parts.append({"text": block["text"]})
                elif block.get("type") == "tool_result":
                    parts.append({"text": f"[tool_result id={block.get('tool_use_id','')}] {block.get('content','')}"})
                elif block.get("type") == "tool_use":
                    parts.append({"functionCall": {"name": block["name"], "args": block.get("input", {})}})
            if parts:
                contents.append({"role": role, "parts": parts})

    google_tools = [{"function_declarations": [
        {"name": t["name"], "description": t["description"], "parameters": t["input_schema"]} for t in _tools
    ]}]
    body = json.dumps({
        "contents": contents, "tools": google_tools,
        "systemInstruction": {"parts": [{"text": _sys_prompt}]},
        "generationConfig": {"maxOutputTokens": 4096},
    }).encode()
    url = GOOGLE_ENDPOINT.format(model=model) + f"?key={api_key}"
    req = urllib.request.Request(url, data=body, headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=120) as resp:
        data = json.loads(resp.read())

    # Translate Google response to Anthropic-like format
    result = {"stop_reason": "end_turn", "content": [], "usage": {"input_tokens": 0, "output_tokens": 0}}
    usage = data.get("usageMetadata", {})
    result["usage"]["input_tokens"] = usage.get("promptTokenCount", 0)
    result["usage"]["output_tokens"] = usage.get("candidatesTokenCount", 0)
    # Gemini context-cache reads (no cache-creation concept on this API).
    result["usage"]["cache_read_input_tokens"] = usage.get("cachedContentTokenCount", 0) or 0
    result["usage"]["cache_creation_input_tokens"] = 0
    for candidate in data.get("candidates", []):
        for part in candidate.get("content", {}).get("parts", []):
            if "text" in part:
                result["content"].append({"type": "text", "text": part["text"]})
            elif "functionCall" in part:
                fc = part["functionCall"]
                result["content"].append({"type": "tool_use", "id": f"google_{int(time.time()*1000)}", "name": fc["name"], "input": fc.get("args", {})})
                result["stop_reason"] = "tool_use"
    return result


def _anthropic_messages_to_openai(messages):
    """Convert Anthropic-format messages to OpenAI-format for OpenRouter."""
    oai_messages = []
    for m in messages:
        role = m["role"]
        content = m["content"]
        if isinstance(content, str):
            oai_messages.append({"role": role, "content": content})
        elif isinstance(content, list):
            # Anthropic list content: text blocks, tool_use blocks, tool_result blocks
            if role == "assistant":
                # Check if there are tool_use blocks
                tool_calls = []
                text_parts = []
                for block in content:
                    btype = block.get("type")
                    if btype == "text":
                        text_parts.append(block.get("text", ""))
                    elif btype == "tool_use":
                        tool_calls.append({
                            "id": block.get("id", ""),
                            "type": "function",
                            "function": {
                                "name": block.get("name", ""),
                                "arguments": json.dumps(block.get("input", {})),
                            },
                        })
                msg = {"role": "assistant", "content": " ".join(text_parts) or None}
                if tool_calls:
                    msg["tool_calls"] = tool_calls
                oai_messages.append(msg)
            elif role == "user":
                # tool_result blocks become tool messages.
                # The runner appends a synthetic "_test" tool_result after each edit
                # (tool_use_id = original_id + "_test") that has no matching tool_call.
                # OpenAI/OpenRouter rejects orphan tool messages, so we merge _test
                # content into the preceding real tool result instead.
                for block in content:
                    btype = block.get("type")
                    if btype == "tool_result":
                        tid = block.get("tool_use_id", "")
                        result_content = block.get("content", "")
                        if isinstance(result_content, list):
                            result_content = " ".join(
                                b.get("text", "") for b in result_content if b.get("type") == "text"
                            )
                        if tid.endswith("_test") and oai_messages and oai_messages[-1].get("role") == "tool":
                            # Merge test.sh output into the preceding tool message
                            oai_messages[-1]["content"] += "\n" + result_content
                        else:
                            oai_messages.append({
                                "role": "tool",
                                "tool_call_id": tid,
                                "content": result_content,
                            })
                    elif btype == "text":
                        oai_messages.append({"role": "user", "content": block.get("text", "")})
    return oai_messages


def call_openrouter(model, messages, api_key, system_prompt=None, tools=None):
    """Call any model via OpenRouter's OpenAI-compatible API with tool support."""
    _tools = tools if tools is not None else DEFAULT_TOOLS
    or_model = OPENROUTER_MODELS.get(model, model)

    # Convert tool defs (Anthropic format) to OpenAI function format
    oai_tools = [
        {
            "type": "function",
            "function": {
                "name": t["name"],
                "description": t["description"],
                "parameters": t["input_schema"],
            },
        }
        for t in _tools
    ]

    oai_messages = _anthropic_messages_to_openai(messages)
    _sys = system_prompt or SYSTEM_PROMPT
    oai_messages = [{"role": "system", "content": _sys}] + oai_messages

    # GPT models need tool_choice=required to actually invoke tools.
    # Other models work with "auto" or may reject "required".
    tc = "required" if "gpt" in or_model.lower() else "auto"
    payload = json.dumps({
        "model": or_model,
        "messages": oai_messages,
        "tools": oai_tools,
        "tool_choice": tc,
        "max_tokens": 4096,
    }).encode()
    req = urllib.request.Request(
        OPENROUTER_ENDPOINT,
        data=payload,
        headers={
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
            "HTTP-Referer": "https://needle-bench.cc",
            "X-Title": "needle-bench",
        },
    )
    with urllib.request.urlopen(req, timeout=120) as resp:
        data = json.loads(resp.read())

    choice = data["choices"][0]
    msg = choice["message"]
    usage = data.get("usage", {})
    finish_reason = choice.get("finish_reason", "end_turn")

    # Translate OpenAI response back to Anthropic format
    content_blocks = []
    if msg.get("content"):
        content_blocks.append({"type": "text", "text": msg["content"]})
    for tc in msg.get("tool_calls") or []:
        fn = tc.get("function", {})
        try:
            args = json.loads(fn.get("arguments", "{}"))
        except json.JSONDecodeError:
            args = {"command": fn.get("arguments", "")}
        content_blocks.append({
            "type": "tool_use",
            "id": tc.get("id", f"call_{int(time.time()*1000)}"),
            "name": fn.get("name", ""),
            "input": args,
        })

    stop_reason = "tool_use" if msg.get("tool_calls") else "end_turn"

    # OpenAI/OpenRouter report cache reads under prompt_tokens_details.cached_tokens.
    # There is no separate cache-creation count in this shape, so it stays 0.
    _cached = (usage.get("prompt_tokens_details") or {}).get("cached_tokens", 0) or 0

    return {
        "content": content_blocks,
        "stop_reason": stop_reason,
        "usage": {
            "input_tokens": usage.get("prompt_tokens", 0),
            "output_tokens": usage.get("completion_tokens", 0),
            "cache_read_input_tokens": _cached,
            "cache_creation_input_tokens": 0,
            "cost_usd": usage.get("cost", 0),  # real cost from OpenRouter
        },
    }


def call_model(model, messages, provider, system_prompt=None, tools=None):
    if provider == "anthropic":
        return call_anthropic(model, messages, os.environ["ANTHROPIC_API_KEY"], system_prompt=system_prompt, tools=tools)
    elif provider == "google":
        return call_google(model, messages, os.environ["GOOGLE_API_KEY"], system_prompt=system_prompt, tools=tools)
    elif provider == "openrouter":
        return call_openrouter(model, messages, os.environ["OPENROUTER_API_KEY"], system_prompt=system_prompt, tools=tools)
    else:
        raise ValueError(f"Unknown provider: {provider}")


def detect_provider(model):
    # Prefer OpenRouter if ANTHROPIC_API_KEY is absent but OPENROUTER_API_KEY is set
    if model.startswith("gemini") and os.environ.get("GOOGLE_API_KEY"):
        return "google"
    if model.startswith("kimi"):
        return "openrouter"
    if not os.environ.get("ANTHROPIC_API_KEY") and os.environ.get("OPENROUTER_API_KEY"):
        return "openrouter"
    if model.startswith("gemini"):
        return "google"
    return "anthropic"


def _load_difficulty_json():
    """Load difficulty.json from the project root, returning (tiers, benchmarks) or (None, None)."""
    proj_root = os.path.dirname(os.path.abspath(__file__))
    diff_path = os.path.join(proj_root, "difficulty.json")
    if not os.path.exists(diff_path):
        return None, None
    with open(diff_path) as f:
        data = json.load(f)
    return data.get("tiers", {}), data.get("benchmarks", {})


def _resolve_difficulty_limits(bench_name):
    """Look up difficulty-tier limits for a benchmark.

    Returns a dict with turns/tokens/wall_clock if the benchmark is in
    difficulty.json, or None to signal the caller should fall back to
    per-Agentfile limits.
    """
    tiers, benchmarks = _load_difficulty_json()
    if tiers is None or benchmarks is None:
        return None
    tier_name = benchmarks.get(bench_name)
    if tier_name is None:
        return None
    return tiers.get(tier_name)


def run_benchmark(model, bench_name, bench_dir, provider):
    # Use canonical name for file paths/scoring, original for API calls
    api_model = model  # preserve full name (e.g. deepseek/deepseek-chat) for OpenRouter
    model = _canonical_agent_name(model)  # short name for score files

    # Try the unified Agentfile.bench first, fall back to per-benchmark Agentfile
    proj_root = os.path.dirname(os.path.abspath(__file__))
    bench_agentfile = os.path.join(proj_root, "Agentfile.bench")
    per_bench_agentfile = os.path.join(bench_dir, "Agentfile")
    if os.path.exists(bench_agentfile):
        cfg = parse_agentfile(bench_agentfile)
    else:
        cfg = parse_agentfile(per_bench_agentfile)
    tools = resolve_tools(cfg["tools"])
    sol_files = solution_files(bench_dir)

    # Resolve limits: difficulty.json tier > Agentfile LIMIT > defaults
    diff_limits = _resolve_difficulty_limits(bench_name)
    if diff_limits is not None:
        max_turns = diff_limits.get("turns", 30)
        max_tokens = diff_limits.get("tokens", 150000)
        max_wall = diff_limits.get("wall_clock", 600)
    else:
        limits = cfg["limits"]
        max_turns = limits.get("turns", 20)
        max_tokens = limits.get("tokens", 100000)
        max_wall = limits.get("wall_clock", 300)
    has_prompt = cfg["prompt"] is not None

    runs_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "runs", model)
    os.makedirs(runs_dir, exist_ok=True)
    log_path = os.path.join(runs_dir, f"{bench_name}.jsonl")

    # POST and metrics recorders write to runs/<model>/<benchmark>/
    post = PostRecorder(runs_dir, bench_name, model)
    metrics = MetricsRecorder(runs_dir, bench_name, model)

    # Build image
    subprocess.run(["docker", "build", "-t", f"needle-bench-{bench_name}", bench_dir],
                   capture_output=True, check=True)

    # Start container
    ts = str(int(time.time()))
    container = f"nb-{model.replace('/', '-')}-{bench_name}-{ts}"
    subprocess.run(["docker", "run", "-d", "--name", container, f"needle-bench-{bench_name}", "sleep", "3600"],
                   capture_output=True, check=True)

    # Detect WORKDIR from the running container
    _, wdir_out, _ = docker_exec(container, "pwd")
    workdir = wdir_out.strip() or "/workspace"

    # Snapshot workspace before agent starts so diff works later.
    docker_exec(container, f"cp -a {workdir} {workdir}.orig")
    docker_exec(container, f"cd {workdir} && git init -q && git add -A && git commit -q -m baseline")

    log_f = open(log_path, "w")
    start_time = time.time()
    total_tokens_in = 0
    total_tokens_out = 0
    total_cache_read = 0
    total_cache_creation = 0
    total_cc_5m = 0
    total_cc_1h = 0
    total_cost_usd = 0.0
    turn_events = []
    # v2.0 metrics: accumulate tool-call counters for read_tool_ratio / tool_calls_per_turn
    total_tool_calls = 0
    total_read_calls = 0   # file:read tool
    total_cat_calls = 0    # cat-via-bash

    def emit(event):
        log_f.write(json.dumps(event) + "\n")
        log_f.flush()

    emit({"event": "run.start", "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
          "benchmark": bench_name, "model": model, "has_prompt": has_prompt, "solution_files": sol_files})

    # SPEC-bench-v2 §1: identical prompt for all arms, no Agentfile override
    instance_prompt = BENCH_USER_PROMPT
    system_prompt = BENCH_SYSTEM_PROMPT

    # BOOT context injection removed (SPEC-bench-v2 §1): all arms get identical prompts.
    # No silent context, no boot output prepended.

    # POST start: capture initial test output before agent touches anything
    _irc, _istdout, _istderr = docker_exec(container, "bash test.sh")
    initial_test_output = _istdout + ("\n" + _istderr if _istderr else "")
    post.start(initial_test_output, instance_prompt)

    messages = [{"role": "user", "content": instance_prompt}]
    final_test_exit = 1
    final_test_output = ""

    try:
        for turn in range(1, max_turns + 1):
            elapsed = time.time() - start_time
            if elapsed >= max_wall:
                break
            if total_tokens_in + total_tokens_out >= max_tokens:
                break

            resp = call_model(api_model, messages, provider, system_prompt=system_prompt, tools=tools)
            usage = extract_usage(resp)
            tokens_in = usage["input_tokens"]
            tokens_out = usage["output_tokens"]
            cache_read = usage["cache_read_input_tokens"]
            cache_creation = usage["cache_creation_input_tokens"]
            cc_5m = usage["cache_creation_5m_tokens"]
            cc_1h = usage["cache_creation_1h_tokens"]
            turn_cost_usd = resp.get("usage", {}).get("cost_usd", 0)
            total_tokens_in += tokens_in
            total_tokens_out += tokens_out
            total_cache_read += cache_read
            total_cache_creation += cache_creation
            total_cc_5m += cc_5m
            total_cc_1h += cc_1h
            total_cost_usd += turn_cost_usd

            # AC2: record per-turn token usage (incl. cache split, task #11)
            metrics.record_turn(turn, tokens_in, tokens_out, cache_read, cache_creation,
                                cc_5m, cc_1h)

            content_blocks = resp.get("content", [])
            stop_reason = resp.get("stop_reason", "end_turn")

            files_edited = []
            files_read = []
            test_exit = None

            # Build assistant message
            messages.append({"role": "assistant", "content": content_blocks})

            tool_results = []
            for block in content_blocks:
                if block.get("type") != "tool_use":
                    continue
                name = block["name"]
                inp = block.get("input", {})
                tool_id = block.get("id", "")
                total_tool_calls += 1

                if name == "bash":
                    cmd = inp.get("command", "")
                    # v2.0: detect cat-via-bash for read_tool_ratio
                    if cmd.strip().startswith("cat "):
                        total_cat_calls += 1
                    # Track file reads from cat/less/head commands
                    # Bug 3 fix: resolve relative paths against /workspace
                    for token in cmd.split():
                        if token.startswith("/dev"):
                            continue
                        if token.startswith("/"):
                            files_read.append(token)
                        elif "." in token and not token.startswith("-"):
                            # Looks like a relative file path (e.g. app.py, src/main.rs)
                            files_read.append("/workspace/" + token)
                    rc, stdout, stderr = docker_exec(container, cmd)
                    output = stdout
                    if stderr:
                        output += ("\n" if output else "") + stderr
                    tool_results.append({"type": "tool_result", "tool_use_id": tool_id,
                                         "content": output if output else f"(exit {rc})"})
                    # AC1: record bash call
                    post.bash(cmd, output, turn)
                elif name == "read":
                    total_read_calls += 1
                    path = inp.get("path", "")
                    files_read.append(path)
                    rc, stdout, stderr = docker_exec(container, f"cat {path!r}")
                    output = stdout
                    if rc != 0:
                        output = f"ERROR: {stderr}" if stderr else f"(exit {rc})"
                    tool_results.append({"type": "tool_result", "tool_use_id": tool_id,
                                         "content": output if output else "(empty file)"})
                    post.read(path, turn)
                elif name == "edit":
                    path = inp.get("path", "")
                    old_str = inp.get("old_str", "")
                    new_str = inp.get("new_str", "")
                    rc, stdout, stderr = do_edit(container, path, old_str, new_str)
                    if rc == 0:
                        files_edited.append(path)
                    result_text = stdout if rc == 0 else f"ERROR: {stderr}"
                    tool_results.append({"type": "tool_result", "tool_use_id": tool_id, "content": result_text})
                    # AC1: record edit call
                    post.edit(path, old_str, new_str, turn)

                    # Run test.sh after every edit
                    if rc == 0:
                        trc, tstdout, tstderr = docker_exec(container, "bash test.sh")
                        test_exit = trc
                        test_output = tstdout
                        if tstderr:
                            test_output += ("\n" if test_output else "") + tstderr
                        final_test_output = test_output
                        # Send test output as a user message after the tool results,
                        # not as a tool_result (Anthropic rejects orphan tool_use_ids)
                        # and not appended to edit result (confuses the agent)
                        pass  # test output sent after tool_results block below

            turn_event = {"event": "turn", "turn": turn, "files_edited": files_edited,
                          "files_read": files_read, "tokens_in": tokens_in, "tokens_out": tokens_out,
                          "cache_read_input_tokens": cache_read,
                          "cache_creation_input_tokens": cache_creation,
                          "test_exit": test_exit}
            emit(turn_event)
            turn_events.append(turn_event)

            if tool_results:
                messages.append({"role": "user", "content": tool_results})

            # Send test output as a separate user message (not a tool_result)
            if final_test_output and test_exit is not None:
                messages.append({"role": "user", "content": f"[test.sh exit={test_exit}]\n{final_test_output}"})

            # Check if test passed
            if test_exit == 0:
                final_test_exit = 0
                break

            # Model stopped producing tool calls
            if stop_reason != "tool_use":
                break

    finally:
        # Final test run
        if final_test_exit != 0:
            trc, fout, ferr = docker_exec(container, "bash test.sh")
            final_test_exit = trc
            final_test_output = fout + ("\n" + ferr if ferr else "")

        # Count correct lines vs solution.patch
        correct_lines = 0
        patch_path = os.path.join(bench_dir, ".bench", "solution.patch")
        if os.path.exists(patch_path):
            with open(patch_path) as f:
                patch_adds = [l[1:].strip() for l in f if l.startswith("+") and not l.startswith("+++")]
            # Get agent's diff
            drc, diff_out, _ = docker_exec(container, f"cd {workdir} && git diff 2>/dev/null || diff -ruN {workdir}.orig {workdir} 2>/dev/null")
            agent_adds = [l[1:].strip() for l in diff_out.splitlines() if l.startswith("+") and not l.startswith("+++")]
            for line in patch_adds:
                if line and line in agent_adds:
                    correct_lines += 1

        wall_clock = time.time() - start_time
        emit({"event": "run.end", "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
              "test_exit": final_test_exit, "total_turns": len(turn_events), "correct_lines": correct_lines})
        log_f.close()

        # AC1: write POST end record
        resolved_flag = (final_test_exit == 0)
        post.end(resolved_flag, final_test_output, len(turn_events))
        post.close()

        # AC2: write metrics summary + print token summary
        cum_in, cum_out, total_tok, cost = metrics.complete()
        metrics.close()
        tok_fmt = f"{total_tok:,}"
        print(f"tokens: {tok_fmt} total (${cost:.4f})", file=sys.stderr)

        # Capture container logs before cleanup (SPEC-bench-v2 §6)
        try:
            log_result = subprocess.run(
                ["docker", "logs", container],
                capture_output=True, text=True, timeout=30,
            )
            container_log_path = os.path.join(runs_dir, f"{bench_name}.log")
            with open(container_log_path, "w") as clf:
                clf.write(log_result.stdout)
                if log_result.stderr:
                    clf.write("\n--- stderr ---\n")
                    clf.write(log_result.stderr)
        except Exception:
            pass  # best-effort — don't fail the run

        # Cleanup container
        subprocess.run(["docker", "rm", "-f", container], capture_output=True)

    # Compute scores
    resolved = final_test_exit == 0
    total_turns = len(turn_events)
    token_cost = total_tokens_in + total_tokens_out

    # turns_to_discovery: first turn that edited a solution file
    turns_to_discovery = max_turns
    for te in turn_events:
        for f in te.get("files_edited", []) + te.get("files_read", []):
            basename = f.lstrip("/").replace("workspace/", "", 1) if f.startswith("/workspace/") else (f.lstrip("/").replace("app/", "", 1) if f.startswith("/app/") else f)
            if any(basename == sf or f.endswith(sf) for sf in sol_files):
                turns_to_discovery = te["turn"]
                break
        if turns_to_discovery != max_turns:
            break

    # turns_to_fix: first turn where test passed
    turns_to_fix = max_turns
    for te in turn_events:
        if te.get("test_exit") == 0:
            turns_to_fix = te["turn"]
            break

    # signal_to_noise
    productive = 0
    for te in turn_events:
        edited = te.get("files_edited", [])
        read = te.get("files_read", [])
        is_productive = False
        for f in edited + read:
            basename = f.lstrip("/").replace("workspace/", "", 1) if f.startswith("/workspace/") else (f.lstrip("/").replace("app/", "", 1) if f.startswith("/app/") else f)
            if any(basename == sf or f.endswith(sf) for sf in sol_files):
                is_productive = True
                break
        if is_productive:
            productive += 1
    signal_to_noise = productive / total_turns if total_turns > 0 else 0.0

    # false_positives: files edited not in solution
    all_edited = set()
    for te in turn_events:
        for f in te.get("files_edited", []):
            all_edited.add(f)
    false_pos = 0
    for f in all_edited:
        basename = f.lstrip("/").replace("workspace/", "", 1) if f.startswith("/workspace/") else (f.lstrip("/").replace("app/", "", 1) if f.startswith("/app/") else f)
        if not any(basename == sf or f.endswith(sf) for sf in sol_files):
            if "test" not in f.lower():
                false_pos += 1

    # tokens_per_correct_line (v1.0 compat — null instead of "Infinity" for JSON consumers)
    tpcl = None if correct_lines == 0 else token_cost / correct_lines

    # v2.0: estimated_cost_usd — prefer real OpenRouter cost, fall back to MODEL_PRICING estimate
    estimated_cost_usd = total_cost_usd if total_cost_usd > 0 else cost

    # task #11: per-arm cache-split aggregates. Field names mirror the canonical
    # bucket schema consumed downstream by consolidate_harden.py / emit_cells.py
    # (→2062 SCHEMA PARITY): fresh + cache_read + cache_create buckets, priced
    # via the shared multipliers. All additive — existing fields are untouched.
    billed_tokens = total_tokens_in + total_cache_read + total_cache_creation
    cache_read_to_create_ratio = (
        round(total_cache_read / total_cache_creation, 4)
        if total_cache_creation else None
    )
    total_billed_equivalent_tokens = round(
        billed_equivalent_input_tokens(
            total_tokens_in, total_cache_read, total_cache_creation,
            total_cc_5m, total_cc_1h) + total_tokens_out, 2)

    # v2.0: dollars_per_correct_line
    if correct_lines > 0:
        dollars_per_correct_line = round(estimated_cost_usd / correct_lines, 6)
    else:
        dollars_per_correct_line = None

    # v2.0: difficulty_tier
    _tiers, _benchmarks = _load_difficulty_json()
    difficulty_tier = (_benchmarks or {}).get(bench_name, "medium")

    # v2.0: tool_calls_per_turn
    tool_calls_per_turn = round(total_tool_calls / max(total_turns, 1), 2)

    # v2.0: read_tool_ratio — file:read vs cat-via-bash
    _total_reads = total_read_calls + total_cat_calls
    read_tool_ratio = round(total_read_calls / _total_reads, 2) if _total_reads > 0 else None

    # recovery_events and recovery_rate (simplified: count reverts)
    recovery_events = 0
    for te in turn_events:
        if len(te.get("files_edited", [])) > 0 and te.get("test_exit") is not None and te["test_exit"] != 0:
            # Check if a previous turn also edited the same file (possible revert)
            for prev in turn_events:
                if prev["turn"] >= te["turn"]:
                    break
                if set(prev.get("files_edited", [])) & set(te.get("files_edited", [])):
                    recovery_events += 1
                    break
    recovery_rate = 1.0 if recovery_events == 0 else (1.0 if resolved else 0.0)

    blind_discovery = resolved and not has_prompt

    # Capture commit hash of the benchmark repo for audit trail
    try:
        import subprocess as _sp
        _bench_sha = _sp.check_output(
            ["git", "-C", bench_dir, "rev-parse", "--short", "HEAD"],
            stderr=_sp.DEVNULL, text=True
        ).strip()
    except Exception:
        _bench_sha = "unknown"

    score = {
        "benchmark": bench_name, "agent": model,
        "commit": _bench_sha,
        "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "difficulty_tier": difficulty_tier,
        "resolved": resolved, "turns_to_discovery": turns_to_discovery,
        "turns_to_fix": turns_to_fix, "signal_to_noise": round(signal_to_noise, 3),
        "false_positives": false_pos, "token_cost": token_cost,
        "input_tokens": total_tokens_in,
        "output_tokens": total_tokens_out,
        "fresh_input_tokens": total_tokens_in,
        "cache_read_tokens": total_cache_read,
        "cache_create_tokens": total_cache_creation,
        "cache_create_5m_tokens": total_cc_5m,
        "cache_create_1h_tokens": total_cc_1h,
        "billed_tokens": billed_tokens,
        "cache_read_to_create_ratio": cache_read_to_create_ratio,
        "total_billed_equivalent_tokens": total_billed_equivalent_tokens,
        "estimated_cost_usd": estimated_cost_usd,
        "dollars_per_correct_line": dollars_per_correct_line,
        "tokens_per_correct_line": tpcl,
        "recovery_events": recovery_events, "recovery_rate": round(recovery_rate, 3),
        "wall_clock": round(wall_clock, 1), "blind_discovery": blind_discovery,
        "read_tool_ratio": read_tool_ratio,
        "tool_calls_per_turn": tool_calls_per_turn,
    }

    score_path = os.path.join(runs_dir, f"{bench_name}.score.json")
    with open(score_path, "w") as f:
        json.dump(score, f, indent=2)

    # Also write score into the per-run subdir so all artifacts are co-located
    run_score_path = os.path.join(runs_dir, bench_name, "score.json")
    with open(run_score_path, "w") as f:
        json.dump(score, f, indent=2)

    print(json.dumps(score, indent=2))
    return score


def list_benchmarks():
    """Return benchmark names from difficulty.json (canonical list).

    Falls back to filesystem scan only if difficulty.json doesn't exist.
    """
    proj_root = os.path.dirname(os.path.abspath(__file__))
    diff_path = os.path.join(proj_root, "difficulty.json")
    if os.path.exists(diff_path):
        with open(diff_path) as f:
            data = json.load(f)
        return sorted(data.get("benchmarks", {}).keys())

    # Fallback: filesystem scan
    bench_dir = os.path.join(proj_root, "benchmarks")
    names = []
    for name in sorted(os.listdir(bench_dir)):
        if name.startswith("_"):
            continue
        if os.path.isdir(os.path.join(bench_dir, name)):
            names.append(name)
    return names


def main():
    parser = argparse.ArgumentParser(description="needle-bench runner")
    parser.add_argument("--model", help="Model name (e.g. claude-haiku-3-5-20241022)")
    parser.add_argument("--benchmark", help="Benchmark name")
    parser.add_argument("--all", action="store_true", help="Run all benchmarks")
    parser.add_argument("--list", action="store_true", help="List benchmarks")
    parser.add_argument("--provider", help="API provider (anthropic or google)")
    args = parser.parse_args()

    if args.list:
        for name in list_benchmarks():
            print(f"  {name}")
        return

    if not args.model:
        print("ERROR: --model required", file=sys.stderr)
        sys.exit(1)

    provider = args.provider or detect_provider(args.model)

    if provider == "anthropic" and not os.environ.get("ANTHROPIC_API_KEY"):
        print("ERROR: ANTHROPIC_API_KEY not set (tip: set --provider openrouter)", file=sys.stderr)
        sys.exit(1)
    if provider == "google" and not os.environ.get("GOOGLE_API_KEY"):
        print("ERROR: GOOGLE_API_KEY not set", file=sys.stderr)
        sys.exit(1)
    if provider == "openrouter" and not os.environ.get("OPENROUTER_API_KEY"):
        print("ERROR: OPENROUTER_API_KEY not set", file=sys.stderr)
        sys.exit(1)

    base = os.path.dirname(os.path.abspath(__file__))
    benchmarks = list_benchmarks() if args.all else [args.benchmark] if args.benchmark else []
    if not benchmarks:
        print("ERROR: --benchmark or --all required", file=sys.stderr)
        sys.exit(1)

    for bench in benchmarks:
        bench_dir = os.path.join(base, "benchmarks", bench)
        if not os.path.isdir(bench_dir):
            print(f"ERROR: benchmark not found: {bench}", file=sys.stderr)
            continue
        print(f"\n=== {bench} ({args.model}) ===", file=sys.stderr)
        try:
            run_benchmark(args.model, bench, bench_dir, provider)
        except Exception as e:
            print(f"ERROR running {bench}: {e}", file=sys.stderr)


if __name__ == "__main__":
    main()