Biomanus_upload/run_bioagent_bench.py

#!/usr/bin/env python3
import argparse
import csv
import gzip
import json
import os
import re
import shutil
import sys
from datetime import datetime
from pathlib import Path
from time import perf_counter

PROJECT_ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(PROJECT_ROOT))

DATASET_ROOT = Path(os.getenv("BIOAGENT_BENCH_DATASET_ROOT", PROJECT_ROOT.parent / "bioagent-bench" / "dataset"))
METADATA_PATH = Path(os.getenv("BIOAGENT_BENCH_METADATA", PROJECT_ROOT.parent / "bioagent-bench" / "src" / "task_metadata.json"))
DEFAULT_OUTPUT_ROOT = PROJECT_ROOT / "bioagent-bench-runs"
DEFAULT_MCP_CONFIG = PROJECT_ROOT / "mcp_config_shim.yaml"
DEFAULT_EXECUTION_ENV_PREFIX = Path(os.getenv("BIOMNI_EXECUTION_ENV_PREFIX", sys.prefix))
FORBIDDEN_VISIBLE_DIRS = {"biomni_data", "__pycache__"}


TASK_OUTPUTS = {
    "alzheimer-mouse": ["pathway_comparison.csv"],
    "comparative-genomics": ["cluster_annotation_mapping.csv"],
    "cystic-fibrosis": ["cf_variants.csv"],
    "deseq": ["up_regulated_genes.csv"],
    "evolution": ["variants_shared.csv", "gene_annotations.csv"],
    "giab": ["predicted.vcf.gz"],
    "metagenomics": ["phylum_relative_abundances.csv"],
    "single-cell": ["all_clusters_de_genes.csv"],
    "transcript-quant": ["truth.tsv"],
    "viral-metagenomics": ["taxonomy.csv"],
}

TASK_SCHEMA_RULES = {
    "alzheimer-mouse": {
        "pathway_comparison.csv": {
            "format": "csv",
            "required_columns": ["pathway", "5xFAD_pvalue", "3xTG_AD_pvalue", "PS3O1S_pvalue"],
        }
    },
    "comparative-genomics": {
        "cluster_annotation_mapping.csv": {
            "format": "csv",
            "required_columns": ["cluster_number", "consensus_annotation"],
        }
    },
    "cystic-fibrosis": {
        "cf_variants.csv": {
            "format": "csv",
            "required_columns": [
                "chromosome",
                "position",
                "variant_id",
                "reference",
                "alternate",
                "gene_name",
                "gene_id",
                "annotation",
                "impact",
                "transcript_id",
                "hgvs_c",
                "hgvs_p",
                "clinical_significance",
                "diseases",
                "review_status",
                "rs_id",
            ],
        }
    },
    "deseq": {
        "up_regulated_genes.csv": {
            "format": "csv",
            "required_columns": ["gene_id", "log2FoldChange", "pvalue", "padj"],
        }
    },
    "evolution": {
        "variants_shared.csv": {
            "format": "csv",
            "required_columns": ["CHROM", "POS", "REF", "ALT", "GENE", "IMPACT", "EFFECT", "STATUS"],
        },
        "gene_annotations.csv": {
            "format": "csv",
            "required_columns": ["CHROM", "POS", "REF", "ALT", "GENE", "IMPACT", "EFFECT", "STATUS"],
        },
    },
    "giab": {
        "predicted.vcf.gz": {
            "format": "vcf.gz",
        }
    },
    "metagenomics": {
        "phylum_relative_abundances.csv": {
            "format": "csv",
            "required_columns": ["OTU", "Kingdom", "Phylum", "JP4D", "JC1A"],
        }
    },
    "single-cell": {
        "all_clusters_de_genes.csv": {
            "format": "csv",
            "required_columns": [
                "cluster_id",
                "predicted_cell_type",
                "gene_name",
                "logfoldchanges",
                "pvals",
                "pvals_adj",
                "direction",
                "abs_logfc",
            ],
        }
    },
    "transcript-quant": {
        "truth.tsv": {
            "format": "tsv_no_header",
        }
    },
    "viral-metagenomics": {
        "taxonomy.csv": {
            "format": "csv",
            "required_columns": ["contig_count", "domain", "species"],
        }
    },
}

TASK_EXTRA_INSTRUCTIONS = {
    "cystic-fibrosis": (
        "For the final row, "
        "variant_id should be the ClinVar VCF ID for the matching record, and rs_id should be the "
        "numeric dbSNP identifier from ClinVar INFO when available, without adding an extra 'rs' prefix."
    ),
    "giab": (
        "The final deliverable must be a bgzip-compatible .vcf.gz file. "
        "If you also generate index or benchmark helper files, keep them in the same run directory."
    ),
    "transcript-quant": (
        "The final deliverable must be a two-column tab-separated file with no header line and no extra "
        "commentary around the table. Each line should be: transcript_id<TAB>count."
    ),
}


def load_task_metadata(metadata_path: Path) -> list[dict]:
    return json.loads(metadata_path.read_text(encoding="utf-8"))


def build_agent_kwargs(args: argparse.Namespace) -> dict:
    provider = os.getenv("BIOMNI_LLM_PROVIDER", "").strip().lower()
    kwargs = {
        "expected_data_lake_files": [],
        "rewrite_user_query": args.rewrite_user_query,
        "dynamic_mcp_registration": args.dynamic_mcp_registration,
        "use_graph_retriever": args.use_graph_retriever,
        "use_tool_retriever": args.use_tool_retriever,
        "timeout_seconds": args.timeout_seconds,
        "mcp_server_top_k": args.mcp_server_top_k,
        "mcp_tool_top_k": args.mcp_tool_top_k,
    }
    if provider == "deepseek":
        kwargs.update(
            {
                "llm": os.getenv("DEEPSEEK_MODEL_NAME", "deepseek-chat"),
                "source": "Custom",
                "base_url": os.getenv("DEEPSEEK_BASE_URL", "https://api.deepseek.com/v1"),
                "api_key": os.getenv("DEEPSEEK_API_KEY"),
            }
        )
    elif args.llm:
        kwargs["llm"] = args.llm
        if args.source:
            kwargs["source"] = args.source
        if args.base_url:
            kwargs["base_url"] = args.base_url
        if args.api_key:
            kwargs["api_key"] = args.api_key
    return kwargs


def list_visible_files(directory: Path, limit: int = 24) -> list[str]:
    if not directory.exists():
        return []
    files = []
    for path in sorted(directory.rglob("*")):
        if path.is_file():
            rel_path = path.relative_to(directory)
            rel_parts = rel_path.parts
            if rel_parts and rel_parts[0] in FORBIDDEN_VISIBLE_DIRS:
                continue
            files.append(str(rel_path))
            if len(files) >= limit:
                break
    return files


def build_benchmark_policy(task_meta: dict, task_dir: Path, run_dir: Path) -> str:
    data_dir = task_dir / "data"
    ref_dir = task_dir / "reference"
    results_dir = task_dir / "results"

    lines = [
        "Benchmark data policy:",
        f"- Allowed input data directory: {data_dir}",
        f"- Allowed reference directory: {ref_dir if ref_dir.exists() else '<none>'}",
        f"- Allowed scratch/output directory: {run_dir}",
        f"- Forbidden truth/results directory: {results_dir}",
        f"- Forbidden sibling benchmark task directories: {DATASET_ROOT}/<any task other than {task_meta['task_id']}>",
        f"- Forbidden generated Biomni cache/runtime directories inside benchmark inputs: {data_dir}/biomni_data and {ref_dir}/biomni_data",
        "- Do not inspect previous bioagent-bench-runs as data sources.",
        "- Do not download external databases or install new packages during the benchmark run.",
        "- You may use installed command-line tools, Python/R packages, and MCP servers as executors, but their inputs must come from the allowed paths above.",
    ]
    return "\n".join(lines)


def build_execution_guard(task_meta: dict, task_dir: Path, run_dir: Path) -> dict:
    task_id = task_meta["task_id"]
    dataset_root = re.escape(str(DATASET_ROOT))
    task_id_re = re.escape(task_id)
    task_dir_re = re.escape(str(task_dir))
    run_root_re = re.escape(str(run_dir.parent))
    run_name_re = re.escape(run_dir.name)

    return {
        "enabled": True,
        "allowed_roots": [
            str(task_dir / "data"),
            str(task_dir / "reference"),
            str(run_dir),
        ],
        "forbidden_patterns": [
            rf"{dataset_root}/(?!{task_id_re}(?:/|$|[\s'\"<>]))[^\s'\"<>]+",
            rf"{task_dir_re}/results(?:/|$|[^\s'\"<>]*)",
            rf"{task_dir_re}/(?:data|reference)/biomni_data(?:/|$|[^\s'\"<>]*)",
            rf"{run_root_re}/(?!{run_name_re}(?:/|$|[\s'\"<>]))[^\s'\"<>]+",
            rf"os\\.walk\\(['\"]{dataset_root}['\"]\\)",
            rf"Path\\(['\"]{dataset_root}['\"]\\)\\.rglob",
        ],
        "forbidden_substrings": [
            "pip install",
            "conda install",
            "mamba install",
            "install.packages(",
            "BiocManager::install",
            "http://",
            "https://",
        ],
        "forbidden_commands": [
            "wget ",
            "curl ",
            "aws s3 cp",
            "gsutil cp",
        ],
    }


def build_benchmark_task_context(task_meta: dict, output_paths: list[Path]) -> dict:
    task_id = task_meta["task_id"]
    return {
        "task_id": task_id,
        "task_name": task_meta.get("name", task_id),
        "description": task_meta.get("description", ""),
        "task_prompt": task_meta.get("task_prompt", ""),
        "extra_instruction": TASK_EXTRA_INSTRUCTIONS.get(task_id, ""),
        "required_outputs": [path.name for path in output_paths],
    }


def build_delivery_guardrails(task_id: str, output_paths: list[Path]) -> str:
    rules = TASK_SCHEMA_RULES.get(task_id, {})
    lines = [
        "Before writing the final <solution>, validate the deliverable yourself against these fairness-preserving checks:",
        "1. The final file names must match the required output paths exactly.",
        "2. The final file schema must match the requested columns/format exactly.",
        "3. Do not export a background universe or broad intermediate table when the prompt asks for a filtered/shared/significant final result set.",
        "4. Do not switch to a different reference coordinate system, taxonomy database, or condition contrast without explicitly proving it is still the task's provided one.",
        "5. If a tool path fails, do not silently change the biological question, reference space, or output definition just to produce a file.",
    ]
    for output_path in output_paths:
        spec = rules.get(output_path.name, {})
        required_columns = spec.get("required_columns", [])
        if required_columns:
            lines.append(f"- {output_path.name} required columns: {', '.join(required_columns)}")
        for warning in spec.get("warnings", []):
            lines.append(f"- {output_path.name}: {warning}")
    return "\n".join(lines)


def _safe_float(value):
    try:
        return float(value)
    except (TypeError, ValueError):
        return None


def _safe_int(value):
    try:
        return int(str(value).strip())
    except (TypeError, ValueError):
        return None


def _read_text_preview(path: Path, limit: int = 2000) -> str:
    if path.suffix == ".gz":
        with gzip.open(path, "rt", encoding="utf-8", errors="ignore") as handle:
            return handle.read(limit)
    return path.read_text(encoding="utf-8", errors="ignore")[:limit]


def _find_reference_contigs(task_dir: Path) -> set[str]:
    contigs: set[str] = set()
    for path in sorted((task_dir / "reference").glob("*")):
        if not path.is_file():
            continue
        suffixes = "".join(path.suffixes).lower()
        if not any(token in suffixes for token in (".fa", ".fasta", ".fna", ".fa.gz", ".fasta.gz", ".fna.gz")):
            continue
        try:
            if path.suffix == ".gz":
                handle = gzip.open(path, "rt", encoding="utf-8", errors="ignore")
            else:
                handle = path.open("r", encoding="utf-8", errors="ignore")
            with handle:
                for line in handle:
                    if line.startswith(">"):
                        contigs.add(line[1:].strip().split()[0])
                    if len(contigs) >= 5000:
                        return contigs
        except OSError:
            continue
    return contigs


def validate_output_file(task_id: str, task_dir: Path, path: Path) -> dict:
    spec = TASK_SCHEMA_RULES.get(task_id, {}).get(path.name, {})
    result = {
        "file": str(path),
        "exists": path.exists(),
        "errors": [],
        "warnings": [],
        "summary": {},
    }
    if not path.exists():
        result["errors"].append("missing_output_file")
        return result

    fmt = spec.get("format")
    if fmt in {"csv", "tsv_no_header"}:
        delimiter = "\t" if fmt == "tsv_no_header" else ","
        with path.open("r", encoding="utf-8", errors="ignore", newline="") as handle:
            rows = list(csv.reader(handle, delimiter=delimiter))
        result["summary"]["row_count"] = max(0, len(rows) - (0 if fmt == "tsv_no_header" else 1))
        if fmt == "tsv_no_header":
            if rows and len(rows[0]) != 2:
                result["errors"].append("expected_two_columns_without_header")
        else:
            header = rows[0] if rows else []
            result["summary"]["header"] = header
            required_columns = spec.get("required_columns", [])
            missing_columns = [col for col in required_columns if col not in header]
            if missing_columns:
                result["errors"].append(f"missing_required_columns:{','.join(missing_columns)}")

            if task_id == "alzheimer-mouse" and result["summary"]["row_count"] > 150:
                result["warnings"].append("appears_to_export_large_pathway_universe")
            if task_id == "comparative-genomics" and result["summary"]["row_count"] > 500:
                result["warnings"].append("appears_to_export_unfiltered_cluster_universe")
            if task_id == "deseq":
                try:
                    dict_rows = list(csv.DictReader(path.open("r", encoding="utf-8", errors="ignore")))
                    non_positive = sum(
                        1
                        for row in dict_rows
                        if (_safe_float(row.get("log2FoldChange")) is not None and _safe_float(row.get("log2FoldChange")) <= 0)
                    )
                    if non_positive:
                        result["warnings"].append(f"contains_{non_positive}_non_upregulated_rows")
                except OSError:
                    pass
            if task_id == "metagenomics":
                try:
                    dict_rows = list(csv.DictReader(path.open("r", encoding="utf-8", errors="ignore")))
                    kingdoms = sorted({(row.get("Kingdom") or "").strip() for row in dict_rows if row.get("Kingdom")})
                    result["summary"]["kingdoms"] = kingdoms
                    if any(k and k != "Bacteria" for k in kingdoms):
                        result["warnings"].append("contains_non_bacterial_rows")
                    sums = {}
                    for sample in ("JP4D", "JC1A"):
                        vals = [_safe_float(row.get(sample)) for row in dict_rows]
                        vals = [v for v in vals if v is not None]
                        if vals:
                            sums[sample] = round(sum(vals), 4)
                            if not (99.0 <= sums[sample] <= 101.0):
                                result["warnings"].append(f"{sample}_relative_abundance_sum_not_near_100")
                    result["summary"]["sample_sums"] = sums
                except OSError:
                    pass
            if task_id == "single-cell":
                try:
                    dict_rows = list(csv.DictReader(path.open("r", encoding="utf-8", errors="ignore")))
                    bad_direction = 0
                    for row in dict_rows[:5000]:
                        direction = (row.get("direction") or "").strip().lower()
                        logfc = _safe_float(row.get("logfoldchanges"))
                        if logfc is None or direction not in {"up", "down"}:
                            continue
                        if (direction == "up" and logfc < 0) or (direction == "down" and logfc > 0):
                            bad_direction += 1
                    if bad_direction:
                        result["warnings"].append(f"direction_logfc_mismatch_rows:{bad_direction}")
                except OSError:
                    pass
            if task_id == "viral-metagenomics":
                try:
                    dict_rows = list(csv.DictReader(path.open("r", encoding="utf-8", errors="ignore")))
                    negative_counts = sum(
                        1 for row in dict_rows if (_safe_int(row.get("contig_count")) is not None and _safe_int(row.get("contig_count")) < 0)
                    )
                    if negative_counts:
                        result["errors"].append("negative_contig_count")
                    if len(dict_rows) > 25:
                        result["warnings"].append("appears_to_export_overly_broad_taxonomic_summary")
                except OSError:
                    pass

    elif fmt == "vcf.gz":
        preview = _read_text_preview(path)
        result["summary"]["preview"] = preview[:400]
        if not preview.startswith("##") and "#CHROM" not in preview:
            result["errors"].append("vcf_header_not_detected")

    if task_id in {"evolution", "giab"} and path.exists():
        reference_contigs = _find_reference_contigs(task_dir)
        if reference_contigs:
            observed_contigs = set()
            try:
                if path.suffix == ".gz":
                    handle = gzip.open(path, "rt", encoding="utf-8", errors="ignore")
                    is_vcf = True
                else:
                    handle = path.open("r", encoding="utf-8", errors="ignore")
                    is_vcf = False
                with handle:
                    for line in handle:
                        if not line.strip():
                            continue
                        if is_vcf and line.startswith("#"):
                            continue
                        if path.suffix != ".gz" and line.lower().startswith("chrom,"):
                            continue
                        observed_contigs.add(line.split("\t", 1)[0] if is_vcf else line.split(",", 1)[0])
                        if len(observed_contigs) >= 100:
                            break
            except OSError:
                observed_contigs = set()

            if observed_contigs and observed_contigs.isdisjoint(reference_contigs):
                result["warnings"].append("observed_coordinate_system_not_in_reference_headers")
            result["summary"]["observed_contig_examples"] = sorted(list(observed_contigs))[:10]

    return result


def validate_outputs(task_id: str, task_dir: Path, output_paths: list[Path]) -> dict:
    file_reports = [validate_output_file(task_id, task_dir, path) for path in output_paths]
    fatal = [err for report in file_reports for err in report["errors"]]
    warnings = [warning for report in file_reports for warning in report["warnings"]]
    return {
        "passed": not fatal,
        "file_reports": file_reports,
        "fatal_errors": fatal,
        "warnings": warnings,
    }


def build_query(
    task_meta: dict,
    task_dir: Path,
    run_dir: Path,
    output_paths: list[Path],
) -> str:
    data_dir = task_dir / "data"
    ref_dir = task_dir / "reference"
    output_lines = [f"- {path.name}: {path}" for path in output_paths]
    data_lines = [f"- {name}" for name in list_visible_files(data_dir)]
    ref_lines = [f"- {name}" for name in list_visible_files(ref_dir)] if ref_dir.exists() else []
    extra = TASK_EXTRA_INSTRUCTIONS.get(task_meta["task_id"], "")
    policy = build_benchmark_policy(task_meta, task_dir, run_dir)
    return f"""
You are running a bioagent-bench task with local files already prepared.

Task ID: {task_meta["task_id"]}
Task name: {task_meta["name"]}
Benchmark prompt:
{task_meta["task_prompt"]}
Data background:
{task_meta["description"]}
Constraints:
1. Use only the benchmark inputs and references explicitly listed below.
2. Do not inspect or use any files under benchmark truth/results directories, sibling task directories, generated biomni_data caches, or previous run outputs.
3. Save the required final deliverables exactly to the paths listed below.
4. Save any intermediate scripts, logs, and scratch outputs inside this run directory: {run_dir}
5. Keep final deliverables in the same schema/format requested by the benchmark prompt.
6. Return a concise final summary after writing the required files.
7. The runner, Python REPL, MCP servers, Rscript, and CLI subprocesses are bound to this conda environment: {os.environ.get("CONDA_PREFIX", DEFAULT_EXECUTION_ENV_PREFIX)}. Do not switch to another conda environment.

{policy}

Input data directory:
{data_dir}
Visible input files:
{chr(10).join(data_lines) if data_lines else "- <empty>"}

Reference data directory:
{ref_dir if ref_dir.exists() else "<none>"}
Visible reference files:
{chr(10).join(ref_lines) if ref_lines else "- <none>"}

Required final output paths:
{chr(10).join(output_lines)}

""".strip()


def save_json(path: Path, payload: dict) -> None:
    path.write_text(json.dumps(payload, ensure_ascii=False, indent=2, default=str), encoding="utf-8")


def count_tokens(text: str | None) -> int:
    if not isinstance(text, str) or not text.strip():
        return 0
    try:
        import tiktoken

        return len(tiktoken.get_encoding("cl100k_base").encode(text))
    except Exception:
        return max(1, len(re.findall(r"\S+", text)))


def build_token_usage_payload(agent) -> dict[str, int]:
    usage = dict(getattr(agent, "last_token_usage", {}) or {})
    payload: dict[str, int] = {}
    for key in ("prompt_tokens", "completion_tokens", "total_tokens", "llm_call_count"):
        value = usage.get(key, 0)
        try:
            payload[key] = int(value)
        except (TypeError, ValueError):
            payload[key] = 0
    return payload


def select_tasks_for_shard(task_ids: list[str], shard_index: int | None, shard_count: int | None) -> list[str]:
    if shard_index is None and shard_count is None:
        return task_ids
    if shard_index is None or shard_count is None:
        raise SystemExit("Provide both --shard-index and --shard-count together.")
    if shard_count <= 0:
        raise SystemExit("--shard-count must be > 0.")
    if shard_index < 0 or shard_index >= shard_count:
        raise SystemExit("--shard-index must satisfy 0 <= shard_index < shard_count.")
    return [task_id for idx, task_id in enumerate(task_ids) if idx % shard_count == shard_index]


def _build_bound_env(env_prefix: Path) -> dict[str, str]:
    """Return an environment that resolves Python/CLI tools from one conda env."""
    env = os.environ.copy()
    env_bin = env_prefix / "bin"
    path_parts = []
    for part in env.get("PATH", "").split(os.pathsep):
        if not part:
            continue
        # Avoid silently falling back to another conda env such as bioenv_cli.
        if "/miniconda3/envs/" in part and Path(part).resolve() != env_bin.resolve():
            continue
        if part not in path_parts:
            path_parts.append(part)

    env["PATH"] = os.pathsep.join([str(env_bin), *path_parts])
    env["CONDA_PREFIX"] = str(env_prefix)
    env["CONDA_DEFAULT_ENV"] = env_prefix.name
    env["CONDA_SHLVL"] = "1"
    env["PYTHONNOUSERSITE"] = "1"
    env["BIOMNI_EXECUTION_ENV_PREFIX"] = str(env_prefix)
    env["BIOMNI_EXECUTION_PYTHON"] = str(env_bin / "python")
    env.pop("VIRTUAL_ENV", None)
    return env


def bind_process_to_execution_env(env_prefix: Path, *, reexec: bool = True) -> dict[str, str]:
    """Bind this benchmark runner to biomni_e1 and optionally re-exec into its Python."""
    env_prefix = env_prefix.expanduser().resolve()
    env_python = env_prefix / "bin" / "python"
    if not env_python.exists():
        raise SystemExit(f"Execution environment Python not found: {env_python}")

    bound_env = _build_bound_env(env_prefix)
    current_python = Path(sys.executable).resolve()
    if reexec and current_python != env_python.resolve():
        os.execve(str(env_python), [str(env_python), *sys.argv], bound_env)

    os.environ.clear()
    os.environ.update(bound_env)
    return bound_env


def _rewrite_csv_header(path: Path, header_map: dict[str, str]) -> bool:
    if not path.exists() or path.stat().st_size == 0:
        return False
    with path.open("r", encoding="utf-8", newline="") as handle:
        rows = list(csv.reader(handle))
    if not rows:
        return False
    original = rows[0]
    rewritten = [header_map.get(col.strip(), header_map.get(col.strip().lower(), col.strip())) for col in original]
    if rewritten == original:
        return False
    with path.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.writer(handle)
        writer.writerow(rewritten)
        writer.writerows(rows[1:])
    return True


def _normalize_cystic_fibrosis_csv(path: Path) -> dict:
    report = {"file": str(path), "actions": []}
    if not path.exists() or path.stat().st_size == 0:
        report["actions"].append("missing_or_empty")
        return report

    with path.open("r", encoding="utf-8", newline="") as handle:
        reader = csv.DictReader(handle)
        rows = list(reader)
        fieldnames = reader.fieldnames or []

    if "rs_id" in fieldnames:
        changed = False
        for row in rows:
            value = (row.get("rs_id") or "").strip()
            if value.lower().startswith("rs") and value[2:].isdigit():
                row["rs_id"] = value[2:]
                changed = True
        if changed:
            with path.open("w", encoding="utf-8", newline="") as handle:
                writer = csv.DictWriter(handle, fieldnames=fieldnames)
                writer.writeheader()
                writer.writerows(rows)
            report["actions"].append("stripped_rs_prefix")
    return report


def _normalize_transcript_quant_tsv(path: Path) -> dict:
    report = {"file": str(path), "actions": []}
    if not path.exists() or path.stat().st_size == 0:
        report["actions"].append("missing_or_empty")
        return report

    original_lines = path.read_text(encoding="utf-8").splitlines()
    cleaned = [line.strip() for line in original_lines if line.strip()]
    if cleaned and cleaned[0].lower().replace(" ", "") in {"transcript_id\tcount", "transcript_id,count"}:
        cleaned = cleaned[1:]
        report["actions"].append("removed_header")

    normalized = []
    for line in cleaned:
        parts = [part.strip() for part in line.replace(",", "\t").split("\t") if part.strip()]
        if len(parts) >= 2:
            normalized.append(f"{parts[0]}\t{parts[1]}")

    if normalized != original_lines:
        path.write_text("\n".join(normalized) + ("\n" if normalized else ""), encoding="utf-8")
        report["actions"].append("normalized_two_column_tsv")
    return report


def postprocess_outputs(task_id: str, output_paths: list[Path]) -> list[dict]:
    """Apply schema-only cleanup that does not read benchmark truth files."""
    reports = []
    for path in output_paths:
        report = {"file": str(path), "actions": []}
        if not path.exists():
            report["actions"].append("missing")
            reports.append(report)
            continue

        if task_id == "alzheimer-mouse" and path.name == "pathway_comparison.csv":
            if _rewrite_csv_header(path, {"pathway": "Pathway"}):
                report["actions"].append("canonicalized_pathway_header")
        elif task_id == "cystic-fibrosis" and path.name == "cf_variants.csv":
            report = _normalize_cystic_fibrosis_csv(path)
        elif task_id == "transcript-quant" and path.name == "truth.tsv":
            report = _normalize_transcript_quant_tsv(path)

        reports.append(report)
    return reports


def run_task(task_meta: dict, args: argparse.Namespace, output_root: Path) -> dict:
    from biomni.agent import A1

    task_id = task_meta["task_id"]
    task_dir = DATASET_ROOT / task_id
    timestamp = datetime.utcnow().strftime("%Y%m%d_%H%M%S")
    run_dir = output_root / f"{task_id}_{timestamp}"
    run_dir.mkdir(parents=True, exist_ok=True)

    output_filenames = TASK_OUTPUTS[task_id]
    output_paths = [run_dir / name for name in output_filenames]
    query = build_query(task_meta, task_dir, run_dir, output_paths)

    agent_kwargs = build_agent_kwargs(args)
    agent_kwargs["path"] = str(run_dir / "agent_runtime")
    agent_kwargs["execution_env_prefix"] = str(Path(args.execution_env_prefix).expanduser().resolve())
    agent_kwargs["benchmark_guard"] = build_execution_guard(task_meta, task_dir, run_dir)
    agent_kwargs["benchmark_task_context"] = build_benchmark_task_context(task_meta, output_paths)

    metadata = {
        "task_id": task_id,
        "task_name": task_meta["name"],
        "run_dir": str(run_dir),
        "dataset_dir": str(task_dir),
        "data_dir": str(task_dir / "data"),
        "reference_dir": str(task_dir / "reference"),
        "agent_runtime_dir": agent_kwargs["path"],
        "output_paths": [str(path) for path in output_paths],
        "agent_kwargs": agent_kwargs,
        "query": query,
        "benchmark_policy": build_benchmark_policy(task_meta, task_dir, run_dir),
        "benchmark_execution_guard": agent_kwargs["benchmark_guard"],
        "benchmark_task_context": agent_kwargs["benchmark_task_context"],
        "timestamp_utc": timestamp,
        "runtime_environment": {
            "execution_env_prefix": str(Path(args.execution_env_prefix).expanduser().resolve()),
            "execution_python": os.environ.get("BIOMNI_EXECUTION_PYTHON", sys.executable),
            "conda_default_env": os.environ.get("CONDA_DEFAULT_ENV"),
            "conda_prefix": os.environ.get("CONDA_PREFIX"),
            "path_head": os.environ.get("PATH", "").split(os.pathsep)[:5],
        },
    }
    save_json(run_dir / "run_metadata.json", metadata)
    (run_dir / "task_query.txt").write_text(query, encoding="utf-8")

    agent = A1(**agent_kwargs)
    if args.use_mcp and args.mcp_graph and Path(args.mcp_graph).exists():
        agent.attach_prebuilt_mcp_graph(str(args.mcp_graph), executable_only=args.executable_mcp_only)
    elif args.use_mcp and args.mcp_config and Path(args.mcp_config).exists():
        agent.attach_mcp_catalog(str(args.mcp_config))

    run_started = perf_counter()
    log_entries, answer = agent.go(query)
    total_runtime_seconds = perf_counter() - run_started
    token_usage = build_token_usage_payload(agent)
    planning_context_text = getattr(agent, "last_planning_context_text", None)
    planning_context_tokens = count_tokens(planning_context_text)
    (run_dir / "final_answer.txt").write_text(str(answer), encoding="utf-8")
    save_json(run_dir / "execution_log.json", {"log_entries": log_entries})
    (run_dir / "execution_log.txt").write_text("\n\n".join(str(entry) for entry in log_entries), encoding="utf-8")
    save_json(
        run_dir / "retrieval_plan.json",
        {
            "query_context": getattr(agent, "query_context", {}),
            "mcp_enabled": args.use_mcp,
            "graph_enabled": args.use_graph_retriever,
            "mcp_graph_route": getattr(agent, "last_graph_route", {}),
            "internal_tool_graph_route": getattr(agent, "last_internal_tool_route", {}),
            "planning_context_text": planning_context_text,
            "planning_context_tokens": planning_context_tokens,
            "planning_context_chars": len(planning_context_text) if isinstance(planning_context_text, str) else 0,
            "planning_latency_seconds": getattr(agent, "last_retrieval_latency_seconds", None),
            "total_runtime_seconds": total_runtime_seconds,
            "token_usage": token_usage,
            "selected_resources": getattr(agent, "last_selected_resources", {}),
            "selected_resource_names": getattr(agent, "last_selected_resources_names", {}),
            "registered_tool_count": len(agent.tool_registry.tools)
            if hasattr(getattr(agent, "tool_registry", None), "tools")
            else None,
            "registered_tool_names": [
                tool.get("name")
                for tool in getattr(getattr(agent, "tool_registry", None), "tools", [])
                if isinstance(tool, dict) and tool.get("name")
            ],
        },
    )

    postprocess_report = postprocess_outputs(task_id, output_paths)
    validation_report = validate_outputs(task_id, task_dir, output_paths)
    save_json(
        run_dir / "output_validation.json",
        {
            "postprocess": postprocess_report,
            "validator": validation_report,
        },
    )

    output_status = []
    for path in output_paths:
        output_status.append(
            {
                "path": str(path),
                "exists": path.exists(),
                "size_bytes": path.stat().st_size if path.exists() else 0,
            }
        )

    result = {
        "task_id": task_id,
        "run_dir": str(run_dir),
        "final_answer_path": str(run_dir / "final_answer.txt"),
        "metadata_path": str(run_dir / "run_metadata.json"),
        "query_path": str(run_dir / "task_query.txt"),
        "retrieval_plan_path": str(run_dir / "retrieval_plan.json"),
        "output_validation_path": str(run_dir / "output_validation.json"),
        "outputs": output_status,
        "planning_latency_seconds": getattr(agent, "last_retrieval_latency_seconds", None),
        "total_runtime_seconds": total_runtime_seconds,
        "planning_context_tokens": planning_context_tokens,
        "planning_context_chars": len(planning_context_text) if isinstance(planning_context_text, str) else 0,
        "validation_passed": validation_report["passed"],
        "validation_warning_count": len(validation_report["warnings"]),
        **token_usage,
    }
    save_json(run_dir / "run_summary.json", result)
    metadata["post_run_metrics"] = {
        "planning_latency_seconds": getattr(agent, "last_retrieval_latency_seconds", None),
        "total_runtime_seconds": total_runtime_seconds,
        "planning_context_tokens": planning_context_tokens,
        "planning_context_chars": len(planning_context_text) if isinstance(planning_context_text, str) else 0,
        "validation_passed": validation_report["passed"],
        "validation_warning_count": len(validation_report["warnings"]),
        "validation_fatal_errors": validation_report["fatal_errors"],
        **token_usage,
    }
    save_json(run_dir / "run_metadata.json", metadata)
    return result


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Run bioagent-bench tasks with Biomanus.")
    parser.add_argument("--task", action="append", help="Task ID to run. Can be provided multiple times.")
    parser.add_argument("--all", action="store_true", help="Run all supported tasks.")
    parser.add_argument("--metadata", default=str(METADATA_PATH))
    parser.add_argument("--dataset-root", default=str(DATASET_ROOT))
    parser.add_argument("--output-root", default=str(DEFAULT_OUTPUT_ROOT))
    parser.add_argument("--mcp-config", default=str(DEFAULT_MCP_CONFIG))
    parser.add_argument(
        "--mcp-graph",
        default=None,
        help="Prebuilt MCP graph directory containing server_catalog.json. Takes precedence over --mcp-config.",
    )
    parser.add_argument("--executable-mcp-only", action="store_true", help="When using --mcp-graph, skip entries without commands.")
    parser.add_argument("--llm", default=None)
    parser.add_argument("--source", default=None)
    parser.add_argument("--base-url", default=None)
    parser.add_argument("--api-key", default=None)
    parser.add_argument("--timeout-seconds", type=int, default=1200)
    parser.add_argument("--mcp-server-top-k", type=int, default=20)
    parser.add_argument("--mcp-tool-top-k", type=int, default=12)
    parser.add_argument(
        "--execution-env-prefix",
        default=str(DEFAULT_EXECUTION_ENV_PREFIX),
        help="Conda environment prefix used for the runner, Python REPL, MCP servers, and CLI tools.",
    )
    parser.add_argument(
        "--no-env-reexec",
        dest="env_reexec",
        action="store_false",
        help="Do not re-exec the benchmark runner with --execution-env-prefix/bin/python.",
    )
    parser.set_defaults(env_reexec=True)
    parser.add_argument("--rewrite-user-query", action="store_true", default=True)
    parser.add_argument("--no-rewrite-user-query", dest="rewrite_user_query", action="store_false")
    parser.add_argument("--dynamic-mcp-registration", action="store_true", default=True)
    parser.add_argument("--no-dynamic-mcp-registration", dest="dynamic_mcp_registration", action="store_false")
    parser.add_argument("--use-graph-retriever", action="store_true", default=True)
    parser.add_argument("--no-graph-retriever", dest="use_graph_retriever", action="store_false")
    parser.add_argument("--use-mcp", action="store_true", default=True)
    parser.add_argument("--no-mcp", dest="use_mcp", action="store_false")
    parser.add_argument("--use-tool-retriever", action="store_true", default=True)
    parser.add_argument("--no-use-tool-retriever", dest="use_tool_retriever", action="store_false")
    parser.add_argument("--shard-index", type=int, default=None, help="0-based shard index over the selected task list.")
    parser.add_argument("--shard-count", type=int, default=None, help="Total number of shards over the selected task list.")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    execution_env = bind_process_to_execution_env(Path(args.execution_env_prefix), reexec=args.env_reexec)
    global DATASET_ROOT
    DATASET_ROOT = Path(args.dataset_root)
    metadata = load_task_metadata(Path(args.metadata))
    meta_by_task = {item["task_id"]: item for item in metadata if item["task_id"] in TASK_OUTPUTS}

    if args.all:
        task_ids = list(meta_by_task)
    else:
        task_ids = args.task or []
    task_ids = select_tasks_for_shard(task_ids, args.shard_index, args.shard_count)
    if not task_ids:
        if args.shard_index is not None:
            print(
                f"No tasks assigned to shard {args.shard_index}/{args.shard_count} under the current selection; exiting."
            )
            return 0
        raise SystemExit("Provide --task <task_id> or use --all.")

    unsupported = [task_id for task_id in task_ids if task_id not in meta_by_task]
    if unsupported:
        raise SystemExit(f"Unsupported task IDs: {unsupported}")

    output_root = Path(args.output_root)
    output_root.mkdir(parents=True, exist_ok=True)

    print(
        "Execution environment: "
        f"{execution_env['CONDA_DEFAULT_ENV']} ({execution_env['CONDA_PREFIX']}); "
        f"python={sys.executable}"
    )

    summaries = []
    for task_id in task_ids:
        print(f"Running task: {task_id}")
        summary = run_task(meta_by_task[task_id], args, output_root)
        summaries.append(summary)
        print(json.dumps(summary, ensure_ascii=False, indent=2))

    batch_summary = {
        "timestamp_utc": datetime.utcnow().strftime("%Y%m%d_%H%M%S"),
        "shard_index": args.shard_index,
        "shard_count": args.shard_count,
        "tasks": summaries,
    }
    save_json(output_root / "latest_batch_summary.json", batch_summary)
    print(f"Batch summary: {output_root / 'latest_batch_summary.json'}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())