BioMysteryBench Claude Harness (Reproduce)

claude_infer.py

"""Reproduce Claude Opus 4.6 on BioMysteryBench-full (99 problems).

═══ SCORES ═══
avg pass@1 over 3 independent runs:
  Overall          54.21%
  Human-solvable   64.91%
  Human-difficult  18.84%
Setting: Anthropic native /v1/messages + extended thinking (budget=4096) +
bash + persistent IPython kernel, max_turns=128, concurrency=4.

═══ USAGE ═══

  export BMB_FULL_DIR=/path/to/BioMysteryBench-full
  export BIOMYSTERY_ENV=/path/to/miniconda3/envs/biomystery
  export BMB_STAGE_ROOT=/tmp/bmb_stage_claude        # per-run isolation
  export JUDGE_BASE_URL=https://api.openai.com/v1
  export JUDGE_API_KEY=$OPENAI_API_KEY
  export JUDGE_MODEL=gpt-5
  python claude_infer.py \\
    --base_url https://api.anthropic.com --api_key $ANTHROPIC_API_KEY \\
    --model claude-opus-4-6 --out ./claude_results \\
    --concurrency 4 --max_turns 128

═══ TOOLS GIVEN TO CLAUDE ═══

  bash    — fresh sub-shell per call, conda env on PATH, 1800s timeout
  python  — persistent IPython kernel (one per trial); variables/imports/
            DataFrames persist across calls; `!cmd` shell escape works

═══ SYSTEM PROMPT (handed to Claude each trial — full template in SYSTEM_TEMPLATE) ═══

  "You are a bioinformatics research assistant working on a real
  BioMysteryBench problem. Analyze the provided data files and answer
  the question. Data files at {workdir}. Files: {file_listing}.
  Allowed network domains: {allowed_domains}. You have two tools:
  `bash` (stateless) and `python` (persistent IPython kernel). You may
  call tools up to {max_turns} times. When done, return your answer
  as plain text wrapped in <answer>...</answer>."

═══ DEPENDENCIES (names only) ═══

  runner pip: httpx jupyter_client openai ipykernel
  conda CLI : samtools bcftools bedtools bwa minimap2 hisat2 STAR bismark
              salmon kallisto featureCounts blastn blastp tblastn makeblastdb
              seqkit seqtk fastqc cutadapt multiqc bamCoverage computeMatrix
              pairtools cooler macs2 meme fimo tomtom kraken2 centrifuge
              megahit prodigal humann picard gatk snpEff SnpSift foldseek
              TMalign mmseqs haplogrep esearch efetch elink prefetch
              fasterq-dump plink plink2 hmmscan jellyfish
              bigWigAverageOverBed pdftotext
  conda Py  : Bio pysam pyfaidx pybedtools pyBigWig pybiomart scanpy anndata
              scrublet leidenalg gemmi pyopenms pyteomics pyensembl mygene
              gseapy numpy pandas scipy sklearn matplotlib seaborn statsmodels
              networkx umap igraph requests requests_cache tqdm pyarrow
              jupyter_client ipykernel gensim cooler
  conda R   : DESeq2 edgeR limma SingleCellExperiment SummarizedExperiment
              GenomicRanges Rsamtools VariantAnnotation GenomicFeatures
              GenomicAlignments rtracklayer biomaRt clusterProfiler enrichplot
              DOSE GOSemSim ggplot2 apeglm ashr BiocManager
  system    : java>=17 curl wget jq awk sed gzip tar bzip2 unzip
"""
from __future__ import annotations

import argparse
import asyncio
import csv
import json
import os
import re
import subprocess
import sys
import time
import zipfile
from pathlib import Path
from statistics import mean

import httpx
import jupyter_client
from openai import AsyncOpenAI


# ============================================================================
# Persistent IPython kernel — same long-running interpreter across many tool
# calls in a single BMB trial. Variables, imports, in-memory DataFrames all
# stay alive between turns.
# ============================================================================

class PersistentPython:
    """One long-lived IPython kernel per BMB trial. Sync API; wrap with
    asyncio.to_thread() from coroutines."""

    def __init__(self, workdir=None, kernel_name="python3", startup_timeout=180):
        self.km = jupyter_client.KernelManager(kernel_name=kernel_name)
        kwargs = {}
        if workdir is not None:
            kwargs["cwd"] = str(workdir)
        self.km.start_kernel(**kwargs)
        self.kc = self.km.client()
        self.kc.start_channels()
        self.kc.wait_for_ready(timeout=startup_timeout)

    def _drain_iopub(self, budget=0.3):
        deadline = time.time() + budget
        while time.time() < deadline:
            try:
                self.kc.get_iopub_msg(timeout=0.05)
            except Exception:
                return

    def run(self, code, timeout=600):
        """Execute one cell. Returns {rc, stdout, stderr, result, elapsed, timed_out}."""
        self._drain_iopub()
        t0 = time.time()
        msg_id = self.kc.execute(code, store_history=False)
        stdout, stderr_stream = [], []
        result_text = ""
        error_text = ""
        timed_out = False
        deadline = time.time() + timeout
        while True:
            try:
                remaining = max(0.1, deadline - time.time())
                msg = self.kc.get_iopub_msg(timeout=remaining)
            except Exception:
                timed_out = True
                try:
                    self.km.interrupt_kernel()
                except Exception:
                    pass
                break
            if msg.get("parent_header", {}).get("msg_id") != msg_id:
                continue
            mtype = msg["header"]["msg_type"]
            c = msg["content"]
            if mtype == "stream":
                (stdout if c["name"] == "stdout" else stderr_stream).append(c["text"])
            elif mtype == "execute_result":
                result_text = c.get("data", {}).get("text/plain", "")
            elif mtype == "display_data":
                txt = c.get("data", {}).get("text/plain", "")
                if txt:
                    stdout.append(str(txt) + "\n")
            elif mtype == "error":
                tb = c.get("traceback", [])
                ansi = re.compile(r"\x1b\[[0-9;]*[mGKHF]")
                error_text = "\n".join(ansi.sub("", t) for t in tb)
            elif mtype == "status" and c.get("execution_state") == "idle":
                break
        return {
            "rc": 0 if not error_text and not timed_out else 1,
            "stdout": "".join(stdout),
            "stderr": error_text or "".join(stderr_stream),
            "result": result_text,
            "elapsed": time.time() - t0,
            "timed_out": timed_out,
        }

    def shutdown(self):
        try: self.kc.stop_channels()
        except Exception: pass
        try: self.km.shutdown_kernel(now=False)
        except Exception:
            try: self.km.shutdown_kernel(now=True)
            except Exception: pass


# ============================================================================
# Dataset + conda env setup
# ============================================================================

FULL = Path(os.environ.get("BMB_FULL_DIR", ""))
if not FULL or not FULL.exists():
    print("ERROR: set BMB_FULL_DIR to the unpacked BioMysteryBench-full directory", file=sys.stderr)
    sys.exit(2)
PROBLEMS_CSV = FULL / "problems.csv"
DATA_DIR = FULL / "data"

BIOMYSTERY_ENV = Path(os.environ.get("BIOMYSTERY_ENV",
                                     Path.home() / "miniconda3/envs/biomystery"))
ENV_BIN = BIOMYSTERY_ENV / "bin"
CONDA_BIN = BIOMYSTERY_ENV.parent.parent / "bin"

# CA bundle override: conda-pack tarball restores may leave libcurl.so with an
# unreplaced placeholder for the CA cert path; HTTPS via curl / libcurl-backed
# R packages then fails. Point env vars at the cert file shipped in the env.
_CA_BUNDLE = BIOMYSTERY_ENV / "ssl" / "cacert.pem"
ENV_VARS = {**os.environ,
            "PATH": f"{ENV_BIN}:{CONDA_BIN}:{os.environ.get('PATH','')}",
            "CONDA_DEFAULT_ENV": BIOMYSTERY_ENV.name,
            "CONDA_PREFIX": str(BIOMYSTERY_ENV),
            "PYTHONWARNINGS": "ignore"}
if _CA_BUNDLE.exists():
    ENV_VARS["CURL_CA_BUNDLE"] = str(_CA_BUNDLE)
    ENV_VARS["SSL_CERT_FILE"] = str(_CA_BUNDLE)
    ENV_VARS["REQUESTS_CA_BUNDLE"] = str(_CA_BUNDLE)

JUDGE_BASE = os.environ.get("JUDGE_BASE_URL", "https://api.openai.com/v1")
JUDGE_KEY = os.environ.get("JUDGE_API_KEY", "")
JUDGE_MODEL = os.environ.get("JUDGE_MODEL", "gpt-5")


# ============================================================================
# Agent prompts & tool schemas
# ============================================================================

SYSTEM_TEMPLATE = """You are a bioinformatics research assistant working on a real BioMysteryBench problem. Analyze the provided data files and answer the question.

The data files for THIS problem are at:
  {workdir}

Files available:
{file_listing}

Allowed network domains (you can install packages and download from these): {allowed_domains}

## Tool

You have one tool: `bash`. It runs in a fresh sub-shell with the conda env on PATH:
samtools, bcftools, bwa, minimap2, STAR, blast+, salmon, kallisto, subread,
scanpy, macs2, meme, deeptools, bismark, cooler, kraken2, megahit, biopython,
gemmi, pysam, pyensembl, mygene, DESeq2 + Bioconductor, entrez-direct, sra-tools,
bedtools, plink, pyBigWig, pyopenms, limma, fastqc, multiqc, picard, gatk (GATK4),
snpEff, SnpSift, foldseek, TMalign, mmseqs, humann, haplogrep,
R, Python (numpy/pandas/scipy/sklearn). You can `pip install` / `conda install` additional packages.

You have TWO tools — pick whichever fits the step:

  - `bash` — run any shell command in a fresh sub-shell. Stateless: each call is a new shell, so file system writes persist but environment variables / cwd / Python-level state do NOT carry over. Each call has a 1800s timeout.
  - `python` — execute Python code in a **persistent IPython kernel** that lives for the whole trial. Variables, imports, big DataFrames you load with `pd.read_csv(...)` etc. all stay alive across calls. Use this when you want to inspect a dataset across multiple steps without re-reading it every turn. Each call has a 1800s timeout. Shell escape inside is allowed via `!cmd`. Working directory is the problem's data folder.

You may call tools up to {max_turns} times in total. After each call you'll see captured stdout / stderr / value.

## Answer format

When you have the answer, return it as plain text content (no tool call) wrapped in <answer>...</answer>. Be concise — match the format the question asks for (e.g. species name, gene symbol, list of sample IDs).
"""

TOOL_RESP = "<tool_response>\nrc={rc} elapsed={el:.1f}s\nSTDOUT (last 8KB):\n{stdout}\n\nSTDERR (last 2KB):\n{stderr}\n</tool_response>"

JUDGE_PROMPT = """You are grading a BioMysteryBench question.

Question:
{question}

Rubric (contains the ground-truth answer + grading instructions):
{rubric}

Model's predicted answer:
{predicted}

Decide whether the predicted answer satisfies the rubric. Output exactly:
VERDICT: CORRECT
or
VERDICT: INCORRECT

Be strict — the rubric says "all or nothing, no partial credit"."""


BASH_TOOL_CLAUDE = {
    "name": "bash",
    "description": "Run a single shell command in the per-problem workdir. STATELESS — each call is a fresh sub-shell. Returns rc/stdout/stderr.",
    "input_schema": {
        "type": "object",
        "properties": {
            "command": {"type": "string", "description": "Shell command (≤2000 chars)."},
        },
        "required": ["command"],
    },
}

PYTHON_TOOL_CLAUDE = {
    "name": "python",
    "description": "Execute Python code in a PERSISTENT IPython kernel (one per trial). Variables, imports, DataFrames etc. persist across calls. Shell escape with `!cmd`. Returns stdout/stderr/value.",
    "input_schema": {
        "type": "object",
        "properties": {
            "code": {"type": "string", "description": "Python source to execute. Multi-line supported."},
        },
        "required": ["code"],
    },
}


# ============================================================================
# Tool execution helpers
# ============================================================================

def parse_answer(text):
    m = re.search(r"<answer>(.*?)</answer>", text, re.DOTALL)
    return m.group(1).strip() if m else None


def run_bash(cmd, workdir, timeout=1800):
    t0 = time.time()
    try:
        p = subprocess.run(["bash", "-c", cmd], capture_output=True, text=True,
                           timeout=timeout, env=ENV_VARS, cwd=workdir)
        return p.returncode, p.stdout, p.stderr, time.time() - t0
    except subprocess.TimeoutExpired as e:
        return -1, (e.stdout or b"").decode("utf-8", "ignore"), (e.stderr or b"").decode("utf-8", "ignore"), time.time() - t0
    except Exception as e:
        return -2, "", str(e), time.time() - t0


def truncate_head_tail(s, head=1024, tail=3072):
    """For bash output: small outputs whole; big ones keep head + tail with a marker."""
    if not s:
        return s
    if len(s) <= head + tail + 80:
        return s
    omitted = len(s) - head - tail
    return f"{s[:head]}\n...[{omitted} bytes truncated in the middle]...\n{s[-tail:]}"


def stage_problem(qid, stage_root):
    """Unzip data/{qid}.zip into a workdir."""
    workdir = stage_root / qid
    workdir.mkdir(parents=True, exist_ok=True)
    zp = DATA_DIR / f"{qid}.zip"
    if not zp.exists():
        return workdir
    if not any(workdir.iterdir()):
        with zipfile.ZipFile(zp) as z:
            z.extractall(workdir)
    return workdir


def list_files(workdir, max_n=30):
    files = []
    for p in sorted(workdir.rglob("*")):
        if p.is_file():
            files.append(f"  - {p.relative_to(workdir)}  ({p.stat().st_size:,} bytes)")
        if len(files) >= max_n:
            files.append("  ... (more files truncated; full listing available via `ls -R`)")
            break
    return "\n".join(files) or "  (no files)"


ANSWER_DEADLINE_WINDOW = 4


def maybe_inject_answer_urgency(messages, turn, max_turns):
    """Last few turns: nag model to emit <answer> rather than hitting the cutoff
    with no answer. Anthropic rejects two consecutive user messages, so merge
    into the existing trailing user message when present."""
    remaining = max_turns - turn
    if remaining > ANSWER_DEADLINE_WINDOW:
        return
    msg = (f"⚠️ {remaining} turn(s) left before the hard cutoff. "
           "STOP investigating and emit your final <answer>...</answer> on THIS response, "
           "based on whatever evidence you've gathered so far. A best-guess answer is "
           "strictly better than no answer.")
    if messages and messages[-1].get("role") == "user":
        last = messages[-1]
        content = last.get("content", "")
        if isinstance(content, list):
            content.append({"type": "text", "text": msg})
        else:
            last["content"] = (content or "") + "\n\n" + msg
    else:
        messages.append({"role": "user", "content": msg})


def _extract_arg(args_obj, keys):
    if isinstance(args_obj, str) and args_obj.strip():
        return args_obj
    if isinstance(args_obj, dict):
        for k in keys:
            v = args_obj.get(k)
            if isinstance(v, str) and v.strip():
                return v
    return ""


def _execute_tool(tool_name, raw_args, pk, workdir):
    """Dispatch one tool call. `pk` lazy-init'd. Returns (rc, tool_out_text, elapsed_s, pk)."""
    try:
        args = json.loads(raw_args) if isinstance(raw_args, str) else (raw_args or {})
    except Exception:
        args = raw_args
    if tool_name == "python":
        code = _extract_arg(args, ("code", "command", "cmd", "input", "python", "source"))
        if not code:
            return 1, "ERROR: python tool needs `code` string argument.", 0.0, pk
        if pk is None:
            pk = PersistentPython(workdir=str(workdir))
        r = pk.run(code, timeout=1800)
        out_t = truncate_head_tail(r["stdout"], head=1024, tail=3072)
        err_t = truncate_head_tail(r["stderr"], head=512, tail=1536)
        result_t = r["result"][:500] if r.get("result") else ""
        msg = f"<tool_response tool=python>\nrc={r['rc']} elapsed={r['elapsed']:.1f}s timed_out={r['timed_out']}\n"
        if result_t:
            msg += f"VALUE: {result_t}\n"
        msg += f"STDOUT:\n{out_t}\n"
        if err_t:
            msg += f"\nSTDERR / TRACEBACK:\n{err_t}\n"
        msg += "</tool_response>"
        return r["rc"], msg, r["elapsed"], pk
    # bash (default)
    cmd = _extract_arg(args, ("command", "cmd", "code", "input", "bash", "shell"))
    if not cmd:
        return 1, "ERROR: bash tool needs a `command` string argument.", 0.0, pk
    rc, out, err, el = run_bash(cmd, workdir=str(workdir))
    out_t = truncate_head_tail(out, head=1024, tail=3072)
    err_t = truncate_head_tail(err, head=512, tail=1536)
    return rc, TOOL_RESP.format(rc=rc, el=el, stdout=out_t, stderr=err_t), el, pk


# ============================================================================
# Claude agent loop — Anthropic native /v1/messages with extended thinking
# ============================================================================

async def run_one_claude(http_client, base_url, api_key, model, problem, workdir,
                         max_turns, thinking_budget, claude_max_tokens):
    """Agent loop. Anthropic requires the most-recent assistant turn's thinking
    blocks to be sent back on the next call (signature verified against the
    tool_use_id). Older thinking blocks are dropped to keep input under the
    200K-token context cap."""
    sysprompt = SYSTEM_TEMPLATE.format(
        workdir=str(workdir),
        file_listing=list_files(workdir),
        allowed_domains=problem["allowed_domains"],
        max_turns=max_turns,
    )
    messages = [{"role": "user", "content": problem["question"]}]
    transcript = []
    url = base_url.rstrip("/") + "/messages"
    headers = {
        "x-api-key": api_key,
        "anthropic-version": "2023-06-01",
        "content-type": "application/json",
    }
    pk = None  # lazy-init persistent python kernel for this trial

    for turn in range(max_turns):
        maybe_inject_answer_urgency(messages, turn, max_turns)
        body = {
            "model": model,
            "max_tokens": claude_max_tokens,
            "thinking": {"type": "enabled", "budget_tokens": thinking_budget},
            "temperature": 1,
            "system": sysprompt,
            "messages": messages,
            "tools": [BASH_TOOL_CLAUDE, PYTHON_TOOL_CLAUDE],
        }
        # Retry policy: transient = httpx network exc / 429 / 5xx ⇒ retry with
        # backoff. Anything else (auth, 4xx schema) ⇒ stop. No fixed attempt
        # cap; keep retrying up to WALL_BUDGET_S so a single LLM call survives
        # gateway outages.
        WALL_BUDGET_S = 1800
        BACKOFF_MAX = 120
        deadline = asyncio.get_event_loop().time() + WALL_BUDGET_S
        last_err = None
        attempt = 0
        data = None
        while True:
            try:
                r = await http_client.post(url, headers=headers, json=body, timeout=1800.0)
                if r.status_code == 429 or r.status_code >= 500:
                    raise RuntimeError(f"HTTP {r.status_code}: {r.text[:300]}")
                r.raise_for_status()
                data = r.json()
                last_err = None
                break
            except Exception as e:
                last_err = e
                transient = isinstance(e, (httpx.TimeoutException, httpx.NetworkError,
                                           httpx.RemoteProtocolError, httpx.ReadError,
                                           httpx.WriteError, httpx.ConnectError,
                                           httpx.ProxyError))
                if not transient:
                    msg = str(e).lower()
                    transient = (any(c in msg for c in ("524", "502", "503", "504", "429"))
                                 or any(k in msg for k in ("timeout", "timed out", "connection",
                                                            "reset", "disconnect", "broken", "eof",
                                                            "readerror", "remoteprotocol")))
                if not transient:
                    break
                if asyncio.get_event_loop().time() >= deadline:
                    break
                await asyncio.sleep(min(5 * (2 ** attempt), BACKOFF_MAX))
                attempt += 1

        if last_err is not None:
            err_type = type(last_err).__name__
            err_msg = str(last_err)[:300]
            err_body = ""
            try:
                err_body = getattr(getattr(last_err, "response", None), "text", "")[:300]
            except Exception:
                pass
            transcript.append({"turn": turn, "error": f"{err_type}: {err_msg}",
                               "response_body": err_body})
            if pk is not None:
                try: pk.shutdown()
                except Exception: pass
            return None, transcript

        blocks = data.get("content", [])
        text = "\n".join(b.get("text", "") for b in blocks if b.get("type") == "text")
        tool_uses = [b for b in blocks if b.get("type") == "tool_use"]
        n_think = sum(1 for b in blocks if b.get("type") == "thinking")

        transcript.append({"turn": turn, "assistant": text[:2000],
                           "n_tool_calls": len(tool_uses), "n_thinking_blocks": n_think,
                           "stop_reason": data.get("stop_reason")})

        ans = parse_answer(text)
        if ans is not None and not tool_uses:
            if pk is not None:
                try: pk.shutdown()
                except Exception: pass
            return ans, transcript

        # Drop thinking blocks from OLDER assistant messages but keep the most
        # recent one intact — Anthropic verifies its signature on next turn.
        for prev in messages:
            if prev.get("role") == "assistant" and isinstance(prev.get("content"), list):
                prev["content"] = [b for b in prev["content"] if b.get("type") != "thinking"]
        messages.append({"role": "assistant", "content": blocks})

        if not tool_uses:
            messages.append({"role": "user", "content":
                "You produced no tool_use and no <answer>. Either call the bash tool "
                "to investigate further, or reply with your final answer wrapped in <answer>...</answer>."})
            continue

        result_blocks = []
        for tu in tool_uses:
            tool_name = tu.get("name", "bash")
            inp = tu.get("input") or {}
            rc, tool_out, el, pk = await asyncio.to_thread(_execute_tool, tool_name, inp, pk, workdir)
            result_blocks.append({
                "type": "tool_result",
                "tool_use_id": tu["id"],
                "content": tool_out,
            })
            cmd_log = inp.get("command") or inp.get("code") or inp.get("cmd") or ""
            if isinstance(cmd_log, str):
                cmd_log = cmd_log[:500]
            transcript.append({"turn": turn, "tool": tool_name, "cmd_or_code": cmd_log,
                               "rc": rc, "elapsed": round(el, 1),
                               "stdout_tail": tool_out[-500:]})
        messages.append({"role": "user", "content": result_blocks})

    if pk is not None:
        try: pk.shutdown()
        except Exception: pass
    return None, transcript


# ============================================================================
# Judge + driver
# ============================================================================

async def judge(judge_client, q, rubric, predicted):
    if not predicted:
        return "INCORRECT", "no answer"
    prompt = JUDGE_PROMPT.format(question=q, rubric=rubric, predicted=predicted)
    try:
        r = await judge_client.chat.completions.create(
            model=JUDGE_MODEL,
            messages=[{"role": "user", "content": prompt}],
            max_completion_tokens=4096,
            reasoning_effort="low",
        )
        c = r.choices[0].message.content or ""
    except Exception as e:
        return "NOT_ATTEMPTED", f"judge error: {e}"
    m = re.search(r"VERDICT:\s*(CORRECT|INCORRECT)", c, re.IGNORECASE)
    return (m.group(1).upper() if m else "NOT_ATTEMPTED"), c[-300:]


def load_problems():
    return list(csv.DictReader(open(PROBLEMS_CSV)))


async def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", default="claude-opus-4-6",
                    help="Anthropic model id, e.g. claude-opus-4-6")
    ap.add_argument("--base_url", default="https://api.anthropic.com",
                    help="Anthropic /v1/messages base URL (origin, no /v1 suffix needed)")
    ap.add_argument("--api_key", required=True,
                    help="Anthropic API key (sent as x-api-key)")
    ap.add_argument("--out", required=True, help="output directory")
    ap.add_argument("--n_trials", type=int, default=1)
    ap.add_argument("--max_turns", type=int, default=128)
    ap.add_argument("--concurrency", type=int, default=4,
                    help="parallel trials; Anthropic rate-limits aggressively, 4 is safe")
    ap.add_argument("--limit", type=int, default=0, help="0=all; otherwise first N problems")
    ap.add_argument("--skip_existing", action="store_true", default=True)
    ap.add_argument("--thinking_budget", type=int, default=4096,
                    help="tokens reserved for extended-thinking blocks per turn")
    ap.add_argument("--claude_max_tokens", type=int, default=8192,
                    help="hard cap on total output tokens per turn")
    args = ap.parse_args()

    if not JUDGE_KEY:
        print("ERROR: set JUDGE_API_KEY env var", file=sys.stderr)
        sys.exit(2)

    problems = load_problems()
    if args.limit:
        problems = problems[:args.limit]
    print(f"[INFO] {args.model} | {len(problems)} problems × {args.n_trials} trials", flush=True)
    print(f"[INFO] Anthropic native /v1/messages | "
          f"thinking_budget={args.thinking_budget} max_tokens={args.claude_max_tokens}", flush=True)

    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)
    (out / "trials").mkdir(exist_ok=True)
    stage_root = Path(os.environ.get("BMB_STAGE_ROOT", "/tmp/bmb_full_stage"))
    stage_root.mkdir(parents=True, exist_ok=True)

    client = httpx.AsyncClient(timeout=1800.0)
    judge_client = AsyncOpenAI(base_url=JUDGE_BASE, api_key=JUDGE_KEY,
                               timeout=600.0, max_retries=4)
    sem = asyncio.Semaphore(args.concurrency)

    async def task(problem, trial):
        qid = problem["id"]
        trial_path = out / "trials" / f"{qid}__t{trial}.json"
        if args.skip_existing and trial_path.exists():
            try:
                existing = json.loads(trial_path.read_text())
                print(f"  [{qid} t{trial}] [cached] {existing.get('verdict')}", flush=True)
                return existing
            except Exception:
                pass
        try:
            wd = stage_problem(qid, stage_root)
            async with sem:
                t0 = time.time()
                ans, transcript = await run_one_claude(
                    client, args.base_url, args.api_key, args.model, problem, wd,
                    args.max_turns, args.thinking_budget, args.claude_max_tokens)
                verdict, judge_detail = await judge(judge_client, problem["question"],
                                                    problem["answer_rubric"], ans)
                rec = {
                    "qid": qid, "trial": trial,
                    "human_solvable": problem["human_solvable"],
                    "predicted": ans, "verdict": verdict,
                    "turns": len(transcript),
                    "duration_s": round(time.time() - t0, 1),
                    "judge_detail": judge_detail,
                }
                trial_path.write_text(json.dumps({**rec, "transcript": transcript},
                                                 ensure_ascii=False, indent=2))
                print(f"  [{qid} t{trial}] verdict={verdict} pred={(ans or '')[:60]!r} "
                      f"({rec['turns']}t, {rec['duration_s']}s)", flush=True)
                return rec
        except Exception as e:
            import traceback
            err_msg = f"{type(e).__name__}: {str(e)[:200]}"
            print(f"  [{qid} t{trial}] TASK_FAILED {err_msg}", flush=True)
            traceback.print_exc()
            return {"qid": qid, "trial": trial,
                    "human_solvable": problem["human_solvable"],
                    "predicted": None, "verdict": "NOT_ATTEMPTED",
                    "turns": 0, "duration_s": 0,
                    "judge_detail": f"task crashed: {err_msg}"}

    tasks = [task(p, t) for p in problems for t in range(args.n_trials)]
    results = await asyncio.gather(*tasks, return_exceptions=True)
    results = [r for r in results if isinstance(r, dict)]

    from collections import defaultdict
    by_q = defaultdict(list)
    for r in results:
        by_q[r["qid"]].append(r["verdict"] == "CORRECT")
    per_q_acc = {qid: mean(v) for qid, v in by_q.items()}
    solv_ids = {p["id"] for p in problems if p["human_solvable"] == "yes"}
    diff_ids = {p["id"] for p in problems if p["human_solvable"] != "yes"}
    summary = {
        "model": args.model, "n_trials": args.n_trials, "n_problems": len(problems),
        "overall_acc": mean(per_q_acc.values()),
        "human_solvable_acc": mean(per_q_acc[i] for i in solv_ids) if solv_ids else None,
        "human_difficult_acc": mean(per_q_acc[i] for i in diff_ids) if diff_ids else None,
        "per_problem_acc": per_q_acc,
        "pass_at_n": {qid: int(any(v)) for qid, v in by_q.items()},
    }
    (out / "summary.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2))
    with open(out / "report.txt", "w") as f:
        f.write(f"=== BioMysteryBench full — {args.model} (n_trials={args.n_trials}) ===\n")
        f.write(f"Overall acc:           {summary['overall_acc']*100:6.2f}%\n")
        if summary['human_solvable_acc'] is not None:
            f.write(f"Human-solvable ({len(solv_ids)}):  {summary['human_solvable_acc']*100:6.2f}%\n")
        if summary['human_difficult_acc'] is not None:
            f.write(f"Human-difficult ({len(diff_ids)}): {summary['human_difficult_acc']*100:6.2f}%\n")
    sv = summary.get('human_solvable_acc') or 0
    dv = summary.get('human_difficult_acc') or 0
    print(f"\n=== {args.model}: {summary['overall_acc']*100:.2f}% overall "
          f"({sv*100:.1f}% solvable / {dv*100:.1f}% difficult) ===")
    await client.aclose()


if __name__ == "__main__":
    asyncio.run(main())