run_agent.py

#!/usr/bin/env python3
"""
Run AI agent to generate destination implementations for MCode benchmarks.
Reads config.toml for benchmark selection, runs claude headless inside docker container.

Prompts are generated on-the-fly from templates in templates/<version>/.
"""

import yaml
import subprocess
import time
import sys
import shutil
import argparse
import json
import os
import tomllib
from datetime import datetime
from pathlib import Path
from typing import Optional, Dict, List

from generate_prompts import generate_prompt, load_benchmark_config, DEFAULT_TEMPLATE_VERSION


def extract_session_id(output_file: Path) -> Optional[str]:
    """
    Extract session ID from Claude Code JSONL output.

    Claude Code outputs session_id in the system.init message:
    {"type":"system","subtype":"init",...,"session_id":"ef349e3a-f4cc-4f69-ac17-69599f2dd382",...}

    Returns:
        Session ID string (UUID format) or None if not found.
    """
    if not output_file.exists():
        return None

    with open(output_file, 'r') as f:
        for line in f:
            try:
                data = json.loads(line)
                # Claude Code native format
                if data.get("type") == "system" and data.get("subtype") == "init":
                    session_id = data.get("session_id")
                    if session_id:
                        return session_id
            except json.JSONDecodeError:
                continue

    return None


def process_agent_config(config: dict, session_id: str = None) -> dict:
    """
    Process a single agent config, replacing template placeholders.
    Can be used after loading from agents.yml and optionally overriding model.

    Args:
        config: Agent configuration dict
        session_id: Optional session ID for resume_template processing
    """
    config = config.copy()  # Don't modify original
    cli_version = config.get('cli_version', '')
    model = config.get('model', '')
    api_key_env = config.get('api_key_env', '')

    # Replace placeholders in install_cmd and cmd_template
    if 'install_cmd' in config:
        config['install_cmd'] = config['install_cmd'].format(
            cli_version=cli_version,
            model=model,
            api_key_env=f"${api_key_env}"
        )
    if 'cmd_template' in config:
        config['cmd_template'] = config['cmd_template'].format(
            cli_version=cli_version,
            model=model,
            api_key_env=api_key_env
        )
    # Process resume_template if present and session_id is provided
    if 'resume_template' in config:
        config['resume_template'] = config['resume_template'].format(
            cli_version=cli_version,
            model=model,
            api_key_env=api_key_env,
            session_id=session_id or "{session_id}"  # Keep placeholder if no session_id
        )

    return config


def load_agents_config(config_path: Path) -> Dict[str, dict]:
    """Load agent configurations from agents.yml and process templates."""
    if not config_path.exists():
        print(f"Error: agents.yml not found at {config_path}")
        sys.exit(1)

    with open(config_path) as f:
        agents = yaml.safe_load(f)

    # Process templates for each agent
    for agent_id in agents:
        agents[agent_id] = process_agent_config(agents[agent_id])

    return agents

# Load .env file if exists
def load_env():
    env_file = Path(__file__).parent / ".env"
    if env_file.exists():
        with open(env_file) as f:
            for line in f:
                line = line.strip()
                if line and not line.startswith('#') and '=' in line:
                    key, value = line.split('=', 1)
                    os.environ[key.strip()] = value.strip().strip('"').strip("'")


def load_config(config_path: Path) -> dict:
    """Load configuration from TOML file."""
    if not config_path.exists():
        return {}
    with open(config_path, 'rb') as f:
        return tomllib.load(f)


def discover_benchmarks(dataset_dir: Path) -> Dict[str, dict]:
    """Discover all benchmarks in the dataset directory."""
    benchmarks = {}
    for item in dataset_dir.iterdir():
        if item.is_dir():
            benchmark_file = item / "benchmark.yml"
            metadata_file = item / "metadata.yml"
            if benchmark_file.exists():
                with open(benchmark_file) as f:
                    config = yaml.safe_load(f)
                # Skip incomplete benchmarks
                if 'type' not in config.get('benchmark', {}):
                    continue
                # Get ID from metadata
                benchmark_id = "000"
                if metadata_file.exists():
                    with open(metadata_file) as f:
                        metadata = yaml.safe_load(f)
                        benchmark_id = metadata.get("id", "000")
                benchmarks[benchmark_id] = {
                    "name": config['benchmark']['name'],
                    "type": config['benchmark']['type'],
                    "dir": item,
                }
    return benchmarks


def run_cmd(cmd: list, **kwargs) -> subprocess.CompletedProcess:
    """Run a command and return the result."""
    return subprocess.run(cmd, **kwargs)


class AgentRunner:
    """Runs the AI agent for a single benchmark."""

    def __init__(self, benchmark_id: str, benchmark_info: dict, log_file: Path, timeout: int, agent_config: dict, workspace_dir: Path = None, template_version: str = None, iterations: int = 1):
        self.benchmark_id = benchmark_id
        self.benchmark_info = benchmark_info
        self.benchmark_dir = benchmark_info["dir"]
        self.log_file = log_file
        self.timeout = timeout  # Per-iteration timeout
        self.agent_config = agent_config
        self.temp_compose_file = None
        # Use custom workspace_dir if provided, otherwise default to benchmark_dir/workspace
        self.workspace_dir = workspace_dir if workspace_dir else self.benchmark_dir / "workspace"
        # Template version for prompt generation (None means use existing prompt.md)
        self.template_version = template_version
        # Number of iterations (for Ralph Wiggum experiments)
        self.iterations = iterations
        self.session_id = None  # Captured from first iteration

    def generate_prompt_file(self) -> bool:
        """Generate prompt.md on-the-fly from template."""
        prompt_path = self.workspace_dir / "prompt.md"

        # If no template version specified and prompt.md exists, use it
        if self.template_version is None and prompt_path.exists():
            print(f"  Using existing prompt.md")
            return True

        # Otherwise generate from template (default to v1)
        template_version = self.template_version or "v1"

        benchmark_config = load_benchmark_config(self.benchmark_dir)
        if benchmark_config is None:
            print(f"  ❌ Failed to load benchmark config from {self.benchmark_dir}")
            return False

        try:
            prompt_content = generate_prompt(benchmark_config, template_version)
            with open(prompt_path, 'w') as f:
                f.write(prompt_content)
            print(f"  Generated prompt.md (template: {template_version})")
            return True
        except Exception as e:
            print(f"  ❌ Failed to generate prompt: {e}")
            return False

    def initialize_workspace(self) -> bool:
        """Initialize workspace by cloning src and creating dst with git."""
        import subprocess

        print(f"  Initializing workspace from scratch...")

        self.workspace_dir.mkdir(parents=True, exist_ok=True)

        # Load metadata for provenance
        metadata_file = self.benchmark_dir / "metadata.yml"
        source_repo = None
        source_commit = None
        if metadata_file.exists():
            with open(metadata_file) as f:
                metadata = yaml.safe_load(f)
                provenance = metadata.get("provenance", {})
                source_repo = provenance.get("source_repo")
                source_commit = provenance.get("source_commit")

        temp_src = self.workspace_dir / "src"
        if source_repo and source_commit:
            print(f"    Cloning {source_repo}...")
            result = subprocess.run(
                ["git", "clone", "--recurse-submodules", source_repo, str(temp_src)],
                capture_output=True, text=True
            )
            if result.returncode != 0:
                print(f"    Clone failed: {result.stderr}")
                return False

            result = subprocess.run(
                ["git", "checkout", source_commit],
                cwd=temp_src, capture_output=True, text=True
            )
            if result.returncode != 0:
                print(f"    Checkout failed: {result.stderr}")
                return False

            subprocess.run(
                ["git", "submodule", "update", "--init", "--recursive"],
                cwd=temp_src, capture_output=True
            )
            print(f"    Cloned at {source_commit[:8]}")
        else:
            workspace_src = self.benchmark_dir / "workspace" / "src"
            if workspace_src.exists():
                shutil.copytree(workspace_src, temp_src)
                print(f"    Copied workspace/src (no provenance)")
            else:
                print(f"    No source_repo in metadata and no workspace/src")
                return False

        # Create dst/ with git repo
        temp_dst = self.workspace_dir / "dst"
        temp_dst.mkdir(exist_ok=True)
        subprocess.run(["git", "init"], cwd=temp_dst, capture_output=True)
        subprocess.run(
            ["git", "config", "user.email", "agent@benchmark.local"],
            cwd=temp_dst, capture_output=True
        )
        subprocess.run(
            ["git", "config", "user.name", "Benchmark Agent"],
            cwd=temp_dst, capture_output=True
        )

        return True

    def setup(self):
        """Prepare docker-compose file for running commands."""
        print(f"  Preparing docker environment...")

        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
        self.temp_compose_file = self.benchmark_dir / f"docker-compose-agent-{timestamp}.yml"

        # Copy docker-compose file
        with open(self.benchmark_dir / "docker-compose.dev.yml", "r") as f:
            compose_content = f.read()

        # Remove port mapping to avoid conflicts (not needed for run)
        compose_content = compose_content.replace(
            "    ports:\n      - \"3000:3000\"\n",
            ""
        )

        # Replace workspace mount with custom workspace path for isolation
        # Original: ./workspace:/workspace
        # New: /absolute/path/to/experiment/workspace:/workspace
        import re
        compose_content = re.sub(
            r'\./workspace:/workspace',
            f'{self.workspace_dir.absolute()}:/workspace',
            compose_content
        )

        with open(self.temp_compose_file, "w") as f:
            f.write(compose_content)

        return True

    def cleanup(self):
        """Clean up temporary files."""
        if self.temp_compose_file and self.temp_compose_file.exists():
            # Clean up any containers created by 'run'
            run_cmd(
                ["docker", "compose", "-f", str(self.temp_compose_file), "down", "--remove-orphans"],
                cwd=self.benchmark_dir,
                capture_output=True
            )
            self.temp_compose_file.unlink()

    def count_commits(self) -> int:
        """Count the number of commits in workspace/dst."""
        dst_dir = self.workspace_dir / "dst"
        if not (dst_dir / ".git").exists():
            return 0
        result = run_cmd(
            ["git", "rev-list", "--count", "HEAD"],
            cwd=dst_dir,
            capture_output=True,
            text=True
        )
        if result.returncode == 0:
            try:
                return int(result.stdout.strip())
            except ValueError:
                return 0
        return 0

    def run(self) -> bool:
        """Run the agent and return success status."""
        # Build environment variables to pass to container
        env_flags = []

        # Handle traditional API key authentication
        api_key_env = self.agent_config.get("api_key_env")
        if api_key_env:
            api_key = os.environ.get(api_key_env)
            if not api_key:
                print(f"  ❌ {api_key_env} not set")
                return False
            env_flags.extend(["-e", f"{api_key_env}={api_key}"])

        # Handle additional env_vars (e.g., for Vertex AI)
        env_vars = self.agent_config.get("env_vars", {})
        for key, value in env_vars.items():
            env_flags.extend(["-e", f"{key}={value}"])

        # For Vertex AI ADC, we need to mount the gcloud credentials
        # Also mount ~/.gemini for cached OAuth credentials (needed for Gemini 3 models)
        auth_type = self.agent_config.get("auth_type")
        extra_mounts = []
        if auth_type == "vertex_ai_adc":
            # Mount the Application Default Credentials from host
            adc_path = Path.home() / ".config" / "gcloud"
            if not adc_path.exists():
                print(f"  ❌ gcloud config not found at {adc_path}")
                print(f"  Run: gcloud auth application-default login")
                return False
            extra_mounts.extend(["-v", f"{adc_path}:/root/.config/gcloud:ro"])

            # Mount ~/.gemini for cached OAuth credentials (read-write needed for tmp files)
            gemini_path = Path.home() / ".gemini"
            if gemini_path.exists():
                extra_mounts.extend(["-v", f"{gemini_path}:/root/.gemini"])

        # Initialize workspace if it doesn't exist (standalone mode)
        if not self.workspace_dir.exists():
            if not self.initialize_workspace():
                return False

        # Generate prompt on-the-fly if template_version is specified
        if not self.generate_prompt_file():
            return False

        try:
            if not self.setup():
                return False

            iterations_str = f", iterations: {self.iterations}" if self.iterations > 1 else ""
            print(f"  Running {self.agent_config['name']} (timeout: {self.timeout}s{iterations_str})...")

            # Temp output file inside container (workspace is mounted)
            container_output = "/workspace/agent_output.jsonl"
            host_output = self.workspace_dir / "agent_output.jsonl"

            # Track results
            all_iteration_outputs = []
            final_success = False
            final_exit_code = -1
            final_timed_out = False
            install_cmd = self.agent_config['install_cmd']

            # Build the shell command for all iterations
            # For multi-iteration: run everything in a single container to preserve session
            if self.iterations > 1:
                # Build a shell script that runs all iterations in sequence
                first_cmd = self.agent_config['cmd_template']
                resume_cmd_template = self.agent_config.get('resume_template', first_cmd)

                # Shell script to run iterations inside container
                # 1. Run first iteration, save output
                # 2. Extract session_id from output
                # 3. Run remaining iterations with -r SESSION_ID
                iteration_script = f'''
{install_cmd}
cd /workspace

echo "=== ITERATION 1 ===" >&2
({first_cmd}) > {container_output}_1 2>&1
ITER1_EXIT=$?
echo "{{\\\"type\\\":\\\"mcode_iter_status\\\",\\\"iteration\\\":1,\\\"exit_code\\\":$ITER1_EXIT}}"

# Extract session_id from first iteration output
SESSION_ID=$(grep -o '"session_id":"[^"]*"' {container_output}_1 | head -1 | sed 's/"session_id":"//;s/"//g')
echo "=== SESSION_ID: $SESSION_ID ===" >&2

# Run remaining iterations with resume
for i in $(seq 2 {self.iterations}); do
    echo "=== ITERATION $i ===" >&2
    RESUME_CMD="{resume_cmd_template}"
    RESUME_CMD="${{RESUME_CMD//\\{{session_id\\}}/$SESSION_ID}}"
    eval "($RESUME_CMD)" > {container_output}_$i 2>&1
    ITER_EXIT=$?
    echo "{{\\\"type\\\":\\\"mcode_iter_status\\\",\\\"iteration\\\":$i,\\\"exit_code\\\":$ITER_EXIT}}"
done

# Combine all outputs
cat {container_output}_* > {container_output}
rm -f {container_output}_*
'''
                shell_cmd = iteration_script
            else:
                # Single iteration - simple command
                shell_cmd = f"{install_cmd} && cd /workspace && touch {container_output} && ({self.agent_config['cmd_template']}) > {container_output} 2>&1"

            agent_cmd = [
                "docker", "compose", "-f", str(self.temp_compose_file), "run",
                "--rm",  # Remove container after it exits
                "--build",  # Build image if needed
                *env_flags,  # Environment variables
                *extra_mounts,  # Mount gcloud and gemini credentials
                "-T",  # Disable pseudo-TTY
                "dev", "bash", "-c",
                shell_cmd
            ]

            print(f"    Running {self.iterations} iteration(s) in single container...")

            start_time = time.time()
            timed_out = False
            exit_code = -1

            try:
                # Run agent, Docker build logs go to stdout (not captured)
                # Total timeout = per-iteration timeout * number of iterations
                total_timeout = self.timeout * self.iterations
                result = run_cmd(
                    agent_cmd,
                    cwd=self.benchmark_dir,
                    timeout=total_timeout
                )
                total_elapsed = time.time() - start_time
                exit_code = result.returncode
                final_success = exit_code == 0

            except subprocess.TimeoutExpired:
                total_elapsed = time.time() - start_time
                timed_out = True
                final_success = False

            final_exit_code = exit_code
            final_timed_out = timed_out

            # Read combined output
            iteration_output = ""
            if host_output.exists():
                with open(host_output, 'r') as f:
                    iteration_output = f.read()

                # Extract session_id from output
                self.session_id = extract_session_id(host_output)

                host_output.unlink()  # Clean up temp file

            # For multi-iteration, we store a single combined output
            all_iteration_outputs.append({
                "iteration": "all",
                "elapsed_seconds": round(total_elapsed, 1),
                "exit_code": exit_code,
                "success": final_success,
                "timed_out": timed_out,
                "output": iteration_output
            })

            # Write final log file with metadata + all iteration outputs + status
            # Use append mode to preserve logs from previous iterations (in test_each_iteration mode)
            with open(self.log_file, 'a') as log:
                # Write metadata as first JSON line
                metadata = {
                    "type": "mcode_metadata",
                    "agent": self.agent_config['name'],
                    "cli_version": self.agent_config.get('cli_version', 'unknown'),
                    "model": self.agent_config.get('model', 'unknown'),
                    "benchmark": self.benchmark_info['name'],
                    "id": self.benchmark_id,
                    "timestamp": datetime.now().isoformat(),
                    "timeout": self.timeout,
                    "iterations": self.iterations,
                    "session_id": self.session_id
                }
                log.write(json.dumps(metadata) + "\n")

                # Write all iteration outputs (the actual agent JSONL content)
                for iter_data in all_iteration_outputs:
                    # Write iteration marker
                    iter_marker = {
                        "type": "mcode_iteration",
                        "iteration": iter_data["iteration"],
                        "elapsed_seconds": iter_data["elapsed_seconds"],
                        "exit_code": iter_data["exit_code"],
                        "success": iter_data["success"],
                        "timed_out": iter_data["timed_out"]
                    }
                    log.write(json.dumps(iter_marker) + "\n")
                    # Write the actual agent output for this iteration
                    log.write(iter_data["output"])

                # Count commits made by agent
                commits = self.count_commits()

                # Determine termination reason
                if final_success:
                    termination_reason = "success"
                elif final_timed_out:
                    termination_reason = "timeout"
                else:
                    termination_reason = "error"

                # Write final status
                status = {
                    "type": "mcode_status",
                    "exit_code": final_exit_code,
                    "elapsed_seconds": round(total_elapsed, 1),
                    "success": final_success,
                    "timed_out": final_timed_out,
                    "termination_reason": termination_reason,
                    "commits": commits,
                    "iterations_completed": len(all_iteration_outputs),
                    "iterations_requested": self.iterations
                }
                log.write(json.dumps(status) + "\n")

            if final_success:
                print(f"  ✅ Agent completed ({total_elapsed:.1f}s total)")
            else:
                print(f"  ❌ Agent failed ({total_elapsed:.1f}s total)")

            return final_success

        finally:
            self.cleanup()


def run_agent_for_benchmark(
    benchmark_name: str,
    benchmark_dir: Path,
    agent_config: dict,
    output_dir: Path,
    timeout: int = 3600,
    workspace_dir: Path = None,
    template_version: str = None,
    iterations: int = 1
) -> dict:
    """
    Run agent for a single benchmark. Designed to be called from run_experiment.py.

    Args:
        benchmark_name: Name of the benchmark
        benchmark_dir: Path to benchmark directory (e.g., dataset/jq-gojq)
        agent_config: Agent configuration dict (from agents.yml, already processed)
        output_dir: Directory to save agent_log.jsonl
        timeout: Timeout in seconds (per iteration)
        workspace_dir: Optional custom workspace path for isolation
        template_version: Template version for prompt generation (None = use existing prompt.md)
        iterations: Number of iterations to run (default: 1, for Ralph Wiggum experiments)

    Returns:
        dict with keys: success, elapsed_seconds, commits, termination_reason, iterations_completed
    """
    load_env()

    # Load benchmark metadata to get ID
    metadata_file = benchmark_dir / "metadata.yml"
    benchmark_id = "000"
    if metadata_file.exists():
        with open(metadata_file) as f:
            metadata = yaml.safe_load(f)
            benchmark_id = metadata.get("id", "000")

    # Create benchmark info dict
    benchmark_info = {
        "name": benchmark_name,
        "type": "unknown",  # Not strictly needed for agent run
        "dir": benchmark_dir,
    }

    # Output file
    log_file = output_dir / "agent_log.jsonl"

    # Run agent
    runner = AgentRunner(benchmark_id, benchmark_info, log_file, timeout, agent_config, workspace_dir, template_version, iterations)
    success = runner.run()

    # Read back the status from the log file to get full results
    result = {
        "success": success,
        "elapsed_seconds": 0,
        "commits": 0,
        "termination_reason": "unknown",
        "iterations_completed": 1,
        "iterations_requested": iterations
    }

    if log_file.exists():
        with open(log_file, 'r') as f:
            for line in f:
                try:
                    data = json.loads(line)
                    if data.get("type") == "mcode_status":
                        result["elapsed_seconds"] = data.get("elapsed_seconds", 0)
                        result["commits"] = data.get("commits", 0)
                        result["termination_reason"] = data.get("termination_reason", "unknown")
                        result["iterations_completed"] = data.get("iterations_completed", 1)
                        result["iterations_requested"] = data.get("iterations_requested", iterations)
                        break
                except json.JSONDecodeError:
                    continue

    return result


def main():
    load_env()

    # Setup paths
    root_dir = Path(__file__).parent
    agents_config_path = root_dir / "agents.yml"

    # Load agents config first to get available choices
    agents = load_agents_config(agents_config_path)

    parser = argparse.ArgumentParser(description='Run AI agent to generate dst implementations')
    parser.add_argument('--agent', type=str, default='claude-code',
                        choices=list(agents.keys()),
                        help=f'Agent to use (default: claude-code). Options: {", ".join(agents.keys())}')
    parser.add_argument('--config', type=str, default='config.toml',
                        help='Path to config file (default: config.toml)')
    parser.add_argument('--timeout', type=int, default=3600,
                        help='Timeout per benchmark in seconds (default: 3600 = 1 hour)')
    args = parser.parse_args()

    # Get agent configuration
    agent_config = agents[args.agent]

    dataset_dir = root_dir / "dataset"
    logs_dir = root_dir / "logs"
    config_path = root_dir / args.config

    # Load config
    config = load_config(config_path)
    benchmark_ids = config.get('benchmarks', {}).get('ids', [])

    # Create log directory (include agent name for easy identification)
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    log_session_dir = logs_dir / f"{args.agent}-{timestamp}"
    log_session_dir.mkdir(parents=True, exist_ok=True)

    # Discover benchmarks
    benchmarks = discover_benchmarks(dataset_dir)

    if not benchmarks:
        print("No benchmarks found in dataset/")
        sys.exit(1)

    # Filter by IDs from config
    if benchmark_ids:
        benchmarks = {k: v for k, v in benchmarks.items() if k in benchmark_ids}
        if not benchmarks:
            print(f"No benchmarks found matching IDs: {benchmark_ids}")
            sys.exit(1)

    # Sort by ID
    sorted_benchmarks = sorted(benchmarks.items(), key=lambda x: x[0])

    print("="*70)
    print("MCode Agent Runner")
    print(f"Agent: {agent_config['name']}")
    print(f"CLI Version: {agent_config.get('cli_version', 'unknown')}")
    print(f"Model: {agent_config.get('model', 'unknown')}")
    print(f"Running {len(sorted_benchmarks)} benchmark(s)")
    print(f"Timeout: {args.timeout}s per benchmark")
    print("="*70)

    results = []

    for benchmark_id, info in sorted_benchmarks:
        print(f"\n{'='*70}")
        print(f"[{benchmark_id}] {info['name']}")
        print("="*70)

        log_file = log_session_dir / f"{info['name']}.jsonl"
        runner = AgentRunner(benchmark_id, info, log_file, args.timeout, agent_config)
        success = runner.run()
        results.append({"id": benchmark_id, "name": info["name"], "success": success})

    # Print summary
    print(f"\n{'='*70}")
    print("SUMMARY")
    print("="*70)
    for r in results:
        status = "✅" if r["success"] else "❌"
        print(f"[{r['id']}] {r['name']}: {status}")
    print("="*70)

    print(f"\nLogs saved to: {log_session_dir}/")

    # Exit code
    all_success = all(r["success"] for r in results)
    sys.exit(0 if all_success else 1)


if __name__ == "__main__":
    main()