scripts/validate_confirmatory_tasks.py

"""Validate held-out task patches against clean archived Git snapshots."""

from __future__ import annotations

import argparse
from io import BytesIO
import json
from pathlib import Path
import shlex
import subprocess
import tarfile
import tempfile
import time
from typing import Any, Iterable

from agent_harness.specs import TaskSpec, load_tasks


def run(
    arguments: list[str],
    cwd: Path,
    timeout: int = 900,
) -> dict[str, Any]:
    started = time.monotonic()
    result = subprocess.run(
        arguments,
        cwd=cwd,
        check=False,
        capture_output=True,
        text=True,
        timeout=timeout,
    )
    combined = result.stdout + result.stderr
    return {
        "command": shlex.join(arguments),
        "exit_code": result.returncode,
        "elapsed_seconds": time.monotonic() - started,
        "output_tail": combined[-12000:],
    }


def extract_snapshot(repository: Path, commit: str, destination: Path) -> None:
    archive = subprocess.run(
        ["git", "archive", "--format=tar", commit],
        cwd=repository,
        check=True,
        capture_output=True,
        timeout=120,
    ).stdout
    with tarfile.open(fileobj=BytesIO(archive), mode="r:") as handle:
        for member in handle.getmembers():
            path = Path(member.name)
            if path.is_absolute() or ".." in path.parts:
                raise RuntimeError(f"Unsafe git archive member: {member.name}")
        handle.extractall(destination)


def apply_patch(checkout: Path, patch: Path) -> dict[str, Any]:
    return run(
        ["git", "apply", "--whitespace=nowarn", str(patch.resolve())],
        checkout,
        timeout=120,
    )


def run_commands(checkout: Path, commands: Iterable[str]) -> list[dict[str, Any]]:
    results: list[dict[str, Any]] = []
    for command in commands:
        arguments = shlex.split(command)
        if not arguments or arguments[0] != "go" or arguments[1:2] != ["test"]:
            raise RuntimeError(f"Only explicit go test validation is allowed: {command}")
        results.append(run(arguments, checkout))
    return results


def all_pass(results: Iterable[dict[str, Any]]) -> bool:
    values = list(results)
    return bool(values) and all(item["exit_code"] == 0 for item in values)


def validate_task(root: Path, repository: Path, task: TaskSpec) -> dict[str, Any]:
    if not task.gold_patch or not task.test_patch:
        raise RuntimeError(f"{task.task_id} has no source/test patch pair")
    source_patch = root / "tasks" / task.gold_patch
    test_patch = root / "tasks" / task.test_patch
    started = time.monotonic()
    with tempfile.TemporaryDirectory(prefix=f"{task.task_id.lower()}-") as raw_checkout:
        checkout = Path(raw_checkout)
        extract_snapshot(repository, task.base_commit, checkout)

        original_tests = run_commands(checkout, task.pass_to_pass_tests)
        test_apply = apply_patch(checkout, test_patch)
        test_only_results = (
            run_commands(checkout, task.fail_to_pass_tests)
            if test_apply["exit_code"] == 0
            else []
        )
        source_apply = apply_patch(checkout, source_patch)
        gold_results = (
            run_commands(checkout, task.pass_to_pass_tests)
            if source_apply["exit_code"] == 0
            else []
        )

    original_pass = all_pass(original_tests)
    test_exposes_bug = bool(test_only_results) and any(
        item["exit_code"] != 0 for item in test_only_results
    )
    gold_pass = all_pass(gold_results)
    valid = (
        original_pass
        and test_apply["exit_code"] == 0
        and test_exposes_bug
        and source_apply["exit_code"] == 0
        and gold_pass
    )
    return {
        "schema_version": 1,
        "task_id": task.task_id,
        "base_commit": task.base_commit,
        "gold_commit": task.gold_commit,
        "valid_end_to_end": valid,
        "checks": {
            "original_suite_passes": original_pass,
            "test_patch_applies": test_apply["exit_code"] == 0,
            "test_patch_exposes_bug": test_exposes_bug,
            "source_patch_applies": source_apply["exit_code"] == 0,
            "gold_suite_passes": gold_pass,
        },
        "original_tests": original_tests,
        "test_patch_apply": test_apply,
        "test_only_results": test_only_results,
        "source_patch_apply": source_apply,
        "gold_results": gold_results,
        "elapsed_seconds": time.monotonic() - started,
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--root", type=Path, default=Path(__file__).resolve().parents[1])
    parser.add_argument("--repository", type=Path, default=Path("data/repos/gitlab-runner"))
    parser.add_argument("--task", action="append")
    args = parser.parse_args()
    root = args.root.resolve()
    repository = args.repository.resolve()
    selected = set(args.task or [])
    tasks = [
        task
        for task_id, task in load_tasks(root).items()
        if task_id.startswith("TASK_CR_") and (not selected or task_id in selected)
    ]
    if not tasks:
        raise SystemExit("No confirmatory tasks selected")
    output_dir = root / "tasks" / "validation"
    output_dir.mkdir(parents=True, exist_ok=True)
    go_version = subprocess.run(
        ["go", "version"], check=True, capture_output=True, text=True, timeout=30
    ).stdout.strip()
    summary: list[dict[str, Any]] = []
    for task in tasks:
        print(f"VALIDATING {task.task_id}", flush=True)
        result = validate_task(root, repository, task)
        result["go_version"] = go_version
        (output_dir / f"{task.task_id}.json").write_text(
            json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8"
        )
        summary.append(
            {
                "task_id": task.task_id,
                "valid_end_to_end": result["valid_end_to_end"],
                "elapsed_seconds": result["elapsed_seconds"],
                "checks": result["checks"],
            }
        )
        print(json.dumps(summary[-1], sort_keys=True), flush=True)
    (output_dir / "summary.json").write_text(
        json.dumps(
            {
                "schema_version": 1,
                "go_version": go_version,
                "tasks": summary,
                "valid_task_ids": [item["task_id"] for item in summary if item["valid_end_to_end"]],
            },
            indent=2,
            sort_keys=True,
        )
        + "\n",
        encoding="utf-8",
    )


if __name__ == "__main__":
    main()