run_agent.py

#!/usr/bin/env python3
"""Host-side launcher for the Claude Code ProgramBench agent image."""

from __future__ import annotations

import argparse
import shutil
import subprocess
import tempfile
import time
from pathlib import Path

from runtime_utils import DEFAULT_INSTANCE_ID, package_submission


AGENT_IMAGE_ORG = "local-programbench"
AGENT_IMAGE_TAG = "agent"


def agent_image_ref(instance_id: str) -> str:
    image_name = instance_id.replace("__", "_1776_")
    return f"{AGENT_IMAGE_ORG}/{image_name}:{AGENT_IMAGE_TAG}"


def run_agent_container(
    *,
    instance_id: str,
    run_dir: Path,
    settings_file: Path | None,
    prompt_file: Path,
    network: str,
    keep_container: bool,
    keep_workdir: Path | None,
    container_name: str | None,
    force: bool,
) -> Path:
    temp_ctx = tempfile.TemporaryDirectory(prefix=f"programbench-agent-{instance_id}-") if keep_workdir is None else None
    tmp = Path(temp_ctx.name) if temp_ctx is not None else keep_workdir
    assert tmp is not None
    try:
        host_out_dir = Path(tmp) / "out"
        host_out_dir.mkdir(parents=True, exist_ok=True)
        name = container_name or f"programbench-agent-{instance_id.replace('__', '-').replace('.', '-')}-{int(time.time())}"
        cmd = [
            "docker",
            "run",
            "--name",
            name,
            "-w",
            "/workspace",
            "-e",
            "IS_SANDBOX=1",
            "--network",
            network,
            "-v",
            f"{host_out_dir}:/workspace/out",
        ]
        if not keep_container:
            cmd.insert(2, "--rm")
        if settings_file is not None:
            cmd.extend(["-v", f"{settings_file.resolve()}:/settings.json:ro"])
        cmd.extend(["-v", f"{prompt_file.resolve()}:/prompt.md:ro"])
        cmd.extend(
            [
                agent_image_ref(instance_id),
                "--instance-id",
                instance_id,
                "--out-dir",
                "/workspace/out",
            ]
        )
        if settings_file is not None:
            cmd.extend(["--settings-file", "/settings.json"])
        cmd.extend(["--prompt-file", "/prompt.md"])
        try:
            subprocess.run(cmd, check=True)
        except subprocess.CalledProcessError:
            if keep_container:
                print(f"container kept for debugging: {name}")
                print(f"host output dir: {host_out_dir}")
                print(f"enter with: docker start -ai {name}")
            raise
        archive = package_submission(host_out_dir, run_dir, instance_id, force=force)

        artifacts = run_dir / instance_id / "agent-output"
        if artifacts.exists():
            shutil.rmtree(artifacts)
        shutil.copytree(host_out_dir, artifacts)
        return archive
    finally:
        if temp_ctx is not None:
            temp_ctx.cleanup()


def open_agent_shell(
    *,
    instance_id: str,
    settings_file: Path | None,
    prompt_file: Path | None,
    workdir: Path,
    network: str,
    container_name: str | None,
) -> None:
    host_out_dir = workdir / "out"
    host_out_dir.mkdir(parents=True, exist_ok=True)
    name = container_name or f"programbench-agent-shell-{instance_id.replace('__', '-').replace('.', '-')}"
    cmd = [
        "docker",
        "run",
        "--rm",
        "-it",
        "--name",
        name,
        "-w",
        "/workspace",
        "-e",
        "IS_SANDBOX=1",
        "--network",
        network,
        "--entrypoint",
        "bash",
        "-v",
        f"{host_out_dir.resolve()}:/workspace/out",
    ]
    if settings_file is not None:
        cmd.extend(["-v", f"{settings_file.resolve()}:/settings.json:ro"])
    if prompt_file is not None:
        cmd.extend(["-v", f"{prompt_file.resolve()}:/prompt.md:ro"])
    cmd.append(agent_image_ref(instance_id))
    subprocess.run(cmd, check=True)


def main() -> None:
    parser = argparse.ArgumentParser(description="Run Claude Code and package its ProgramBench submission")
    subparsers = parser.add_subparsers(dest="command", required=True)

    run = subparsers.add_parser("run", help="Run Claude Code and package /out")
    run.add_argument("--instance-id", default=DEFAULT_INSTANCE_ID)
    run.add_argument("--run-dir", type=Path, required=True)
    run.add_argument("--settings-file", type=Path, required=True)
    run.add_argument("--prompt-file", type=Path, required=True)
    run.add_argument("--network", default="bridge")
    run.add_argument("--keep-container", action="store_true")
    run.add_argument("--keep-workdir", type=Path)
    run.add_argument("--container-name")
    run.add_argument("--force", action="store_true")

    shell = subparsers.add_parser("shell", help="Open an interactive shell in the Claude Code agent image")
    shell.add_argument("--instance-id", default=DEFAULT_INSTANCE_ID)
    shell.add_argument("--settings-file", type=Path)
    shell.add_argument("--prompt-file", type=Path)
    shell.add_argument("--workdir", type=Path, default=Path("debug-agent"))
    shell.add_argument("--network", default="bridge")
    shell.add_argument("--container-name")
    args = parser.parse_args()

    if args.command == "shell":
        open_agent_shell(
            instance_id=args.instance_id,
            settings_file=args.settings_file,
            prompt_file=args.prompt_file,
            workdir=args.workdir,
            network=args.network,
            container_name=args.container_name,
        )
        return

    print(
        run_agent_container(
            instance_id=args.instance_id,
            run_dir=args.run_dir,
            settings_file=args.settings_file,
            prompt_file=args.prompt_file,
            network=args.network,
            keep_container=args.keep_container,
            keep_workdir=args.keep_workdir,
            container_name=args.container_name,
            force=args.force,
        )
    )


if __name__ == "__main__":
    main()