GEWU/scripts/prepare_four_file_run.py

#!/usr/bin/env python3
"""Prepare a private clean worktree and data copy for a four-file benchmark run."""

from __future__ import annotations

import argparse
import json
import shutil
import subprocess
import sys
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

if __package__ in {None, ""}:
    sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from scripts.build_benchmark_manifest import create_manifest


def git(repo: Path, *args: str) -> str:
    result = subprocess.run(
        ["git", "-C", str(repo), *args],
        check=True,
        capture_output=True,
        text=True,
    )
    return result.stdout.strip()


def private_tree(path: Path) -> None:
    for item in path.rglob("*"):
        item.chmod(item.stat().st_mode & ~0o077)
    path.chmod(path.stat().st_mode & ~0o077)


def prepare(
    *,
    repo: Path,
    source_data: Path,
    sources: list[Path],
    reference: Path,
    extractions: Path,
    model_id: str,
    search_mode: str,
    budget_tokens: int,
    output_root: Path,
    commit: str,
) -> dict[str, Any]:
    repo, source_data, output_root = repo.resolve(), source_data.resolve(), output_root.resolve()
    if output_root.exists():
        raise FileExistsError(f"Run output already exists: {output_root}")
    if output_root == repo or repo in output_root.parents or output_root in repo.parents:
        raise ValueError("Run output must be isolated from the source repository")
    dirty = git(repo, "status", "--porcelain", "--untracked-files=all")
    if dirty:
        raise ValueError("Source checkout must be clean before creating a benchmark worktree")
    resolved_commit = git(repo, "rev-parse", f"{commit}^{{commit}}")
    output_root.mkdir(parents=True, mode=0o700)
    worktree = output_root / "checkout"
    data_copy = output_root / "data"
    try:
        subprocess.run(
            ["git", "-C", str(repo), "worktree", "add", "--detach", str(worktree), resolved_commit],
            check=True,
            capture_output=True,
            text=True,
        )
        subprocess.run(
            ["git", "-C", str(worktree), "submodule", "update", "--init", "--recursive"],
            check=True,
            capture_output=True,
            text=True,
        )
        shutil.copytree(source_data, data_copy, copy_function=shutil.copy2)
        private_tree(data_copy)
        if git(worktree, "rev-parse", "HEAD") != resolved_commit:
            raise ValueError("Prepared worktree commit differs from the requested revision")
        if git(worktree, "status", "--porcelain", "--untracked-files=all"):
            raise ValueError("Prepared worktree is not clean after submodule initialization")
        manifest = create_manifest(
            repo=worktree,
            pack=worktree / "agents" / "gewu",
            sources=[path.resolve() for path in sources],
            reference=reference.resolve(),
            extractions=data_copy / "extractions",
            model_id=model_id,
            search_mode=search_mode,
            budget_tokens=budget_tokens,
        )
        manifest["run_setup"] = {
            "source_checkout": str(repo),
            "source_data": str(source_data),
            "data_copy": str(data_copy),
            "worktree_path": str(worktree),
            "requested_commit": resolved_commit,
        }
        manifest_path = output_root / "benchmark-manifest.json"
        with manifest_path.open("x", encoding="utf-8") as stream:
            stream.write(json.dumps(manifest, ensure_ascii=False, indent=2, sort_keys=True) + "\n")
        state = {
            "schema_version": 1,
            "run_id": output_root.name,
            "status": "prepared",
            "stage": "isolated-worktree-and-data-copy-ready",
            "created_utc": datetime.now(UTC).isoformat(),
            "backend_commit": resolved_commit,
            "checkout": str(worktree),
            "data": str(data_copy),
            "manifest": str(manifest_path),
            "owner_decisions_written": 0,
            "canonical_data_mutated": False,
        }
        (output_root / "run-state.json").write_text(
            json.dumps(state, ensure_ascii=False, indent=2, sort_keys=True) + "\n",
            encoding="utf-8",
        )
        private_tree(output_root)
        return state
    except Exception as error:
        failure = {
            "schema_version": 1,
            "status": "preparation_failed",
            "error_type": type(error).__name__,
            "message": str(error),
            "created_utc": datetime.now(UTC).isoformat(),
        }
        (output_root / "preparation-failure.json").write_text(
            json.dumps(failure, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
        )
        private_tree(output_root)
        raise


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--repo", type=Path, required=True)
    parser.add_argument("--source-data", type=Path, required=True)
    parser.add_argument("--source", type=Path, action="append", required=True)
    parser.add_argument("--reference", type=Path, required=True)
    parser.add_argument("--model-id", required=True)
    parser.add_argument("--search-mode", required=True)
    parser.add_argument("--budget-tokens", type=int, required=True)
    parser.add_argument("--output-root", type=Path, required=True)
    parser.add_argument("--commit", default="HEAD")
    args = parser.parse_args()
    result = prepare(
        repo=args.repo,
        source_data=args.source_data,
        sources=args.source,
        reference=args.reference,
        extractions=args.source_data / "extractions",
        model_id=args.model_id,
        search_mode=args.search_mode,
        budget_tokens=args.budget_tokens,
        output_root=args.output_root,
        commit=args.commit,
    )
    print(
        json.dumps(
            {
                "status": result["status"],
                "run_id": result["run_id"],
                "output": str(args.output_root.resolve()),
            }
        )
    )


if __name__ == "__main__":
    main()