src/agent_harness/repair_experiment.py

"""Frozen E03 end-to-end repair experiment over held-out GitLab Runner tasks."""

from __future__ import annotations

from contextlib import contextmanager
from dataclasses import asdict
from hashlib import sha256
import json
import os
from pathlib import Path
import re
import shlex
import subprocess
import sys
import tarfile
import tempfile
import time
from typing import Any, Iterator, Sequence

from .components import Candidate
from .context_packing import (
    EVIDENCE_TOKEN_BUDGET,
    candidates_from_records,
    load_ranking_records,
    pack_snippets,
    ranking_path,
)
from .interactive_experiment import exclusive_loaded
from .lm_studio import LMStudioClient
from .pilot import research_code_revision, retrieval_metrics
from .repository import GitSnapshot
from .specs import (
    HarnessSpec,
    TaskSpec,
    load_experiments,
    load_harnesses,
    load_models,
    load_task_split,
    load_tasks,
)
from .syntax_index import parse_go_file
from .telemetry import EventWriter, RunIdentity, run_directory
from .tokenization import QwenTokenCounter


REPAIR_SYSTEM = """You are repairing a held-out issue in a large Go repository.
Use only the issue and source evidence supplied by the harness. Return only a valid unified
diff that can be applied with `git apply`. Modify only evidence paths, do not add tests, do not
use Markdown fences, and do not include explanation outside the diff. Make the smallest correct
production-code change. If evidence is insufficient, still make the best evidence-grounded patch."""


class RepairExperimentError(RuntimeError):
    """Raised when E03 infrastructure cannot preserve the frozen protocol."""


class PatchOutputError(RepairExperimentError):
    """Raised when a model patch violates the frozen output interface."""


def repair_prompt(task: TaskSpec, evidence: str) -> str:
    return f"ISSUE:\n{task.statement}\n\nSOURCE EVIDENCE:\n{evidence}"


def extract_unified_diff(response: dict[str, Any]) -> str:
    """Extract a model diff without repairing or otherwise changing its content."""

    try:
        content = response["choices"][0]["message"]["content"]
    except (KeyError, IndexError, TypeError) as exc:
        raise PatchOutputError("chat completion has no assistant content") from exc
    if not isinstance(content, str) or not content.strip():
        raise PatchOutputError("assistant returned no patch text")
    stripped = content.strip()
    fence = re.search(r"```(?:diff|patch)?\s*\n(.*?)```", stripped, re.DOTALL)
    if fence:
        stripped = fence.group(1).strip()
    starts = [value for value in (stripped.find("diff --git "), stripped.find("--- a/")) if value >= 0]
    if not starts:
        raise PatchOutputError("assistant did not return a unified diff")
    patch = stripped[min(starts) :].rstrip() + "\n"
    if "+++ b/" not in patch or "@@" not in patch:
        raise PatchOutputError("assistant patch lacks unified-diff file or hunk headers")
    return patch


def modified_paths(patch: str) -> tuple[str, ...]:
    paths = re.findall(r"(?m)^\+\+\+ b/(.+)$", patch)
    unique = tuple(dict.fromkeys(item.strip() for item in paths))
    if not unique:
        raise PatchOutputError("assistant patch has no modified repository path")
    for value in unique:
        path = Path(value)
        if path.is_absolute() or ".." in path.parts or value == "/dev/null":
            raise PatchOutputError(f"assistant patch contains unsafe path: {value}")
    return unique


def validate_patch_scope(patch: str, allowed_paths: Sequence[str]) -> tuple[str, ...]:
    paths = modified_paths(patch)
    unknown = set(paths) - set(allowed_paths)
    if unknown:
        raise PatchOutputError(
            f"assistant modified paths outside the supplied evidence: {sorted(unknown)}"
        )
    return paths


def _iterative_ranking_path(root: Path, task_id: str) -> Path:
    matches = sorted(
        (root / "results" / "staging" / "E02").glob(
            f"*/H010/{task_id}/refined_ranking.json"
        )
    )
    if len(matches) != 1:
        raise RepairExperimentError(
            f"expected one confirmatory H010 refined ranking for {task_id}, found {len(matches)}"
        )
    return matches[0]


def oracle_hunk_ranges(patch: str, wanted_path: str) -> tuple[tuple[int, int], ...]:
    """Read only old-file hunk locations for an oracle symbol absent at the base commit."""

    current_path: str | None = None
    ranges: list[tuple[int, int]] = []
    for line in patch.splitlines():
        if line.startswith("+++ b/"):
            current_path = line.removeprefix("+++ b/").strip()
            continue
        if current_path != wanted_path or not line.startswith("@@"):
            continue
        match = re.match(r"@@ -(\d+)(?:,(\d+))? \+\d+(?:,\d+)? @@", line)
        if match:
            start = int(match.group(1))
            count = int(match.group(2) or "1")
            ranges.append((start, max(count, 1)))
    return tuple(ranges)


def _oracle_function_candidates(
    root: Path, snapshot: GitSnapshot, task: TaskSpec
) -> tuple[Candidate, ...]:
    candidates: list[Candidate] = []
    by_path: dict[str, set[str]] = {}
    for identifier in task.gold_symbols:
        path, separator, name = identifier.partition("::")
        if not separator:
            raise RepairExperimentError(f"malformed oracle symbol: {identifier}")
        by_path.setdefault(path, set()).add(name)
    for path, names in by_path.items():
        source = snapshot.read_file(task.base_commit, path)
        lines = source.text.splitlines()
        matched: set[str] = set()
        for symbol in parse_go_file(path, source.text):
            if symbol.name not in names:
                continue
            matched.add(symbol.name)
            start = max(symbol.line_start, 1)
            end = min(symbol.line_end, len(lines))
            candidates.append(
                Candidate(
                    path=path,
                    line_start=start,
                    line_end=end,
                    text="\n".join(lines[start - 1 : end]),
                    source="oracle_function",
                    score=1.0,
                    symbol=symbol.name,
                )
            )
        missing = names - matched
        if missing:
            patch = (root / "tasks" / task.gold_patch).read_text(encoding="utf-8")
            ranges = oracle_hunk_ranges(patch, path)
            if not ranges:
                raise RepairExperimentError(
                    f"oracle symbols absent at base and no patch hunk found: {sorted(missing)}"
                )
            for start, count in ranges:
                window_start = max(start - 20, 1)
                window_end = min(start + count + 20, len(lines))
                candidates.append(
                    Candidate(
                        path=path,
                        line_start=window_start,
                        line_end=window_end,
                        text="\n".join(lines[window_start - 1 : window_end]),
                        source="oracle_hunk_location",
                        score=1.0,
                        symbol=", ".join(sorted(missing)),
                    )
                )
    return tuple(candidates)


def repair_context(
    root: Path,
    task: TaskSpec,
    harness: HarnessSpec,
    snapshot: GitSnapshot,
    tokenizer: QwenTokenCounter,
) -> tuple[str, tuple[str, ...], int, str]:
    if harness.harness_id == "H019":
        candidates = _oracle_function_candidates(root, snapshot, task)
        source = "oracle_function"
    else:
        path = (
            _iterative_ranking_path(root, task.task_id)
            if harness.harness_id == "H010"
            else ranking_path(root, "E01", harness.harness_id, task.task_id)
        )
        records = load_ranking_records(path)
        candidates = candidates_from_records(snapshot, task.base_commit, records)
        source = str(path.relative_to(root))
    text, allowed, used = pack_snippets(tokenizer, candidates, EVIDENCE_TOKEN_BUDGET)
    if not allowed:
        raise RepairExperimentError(f"{harness.harness_id}/{task.task_id} packed no evidence")
    return text, allowed, used, source


@contextmanager
def isolated_source_tree(repository: Path, commit: str) -> Iterator[Path]:
    """Materialize a commit without mutating the user's repository checkout."""

    with tempfile.TemporaryDirectory(prefix="agent-harness-e03-") as temporary:
        base = Path(temporary)
        archive = base / "source.tar"
        tree = base / "tree"
        tree.mkdir()
        with archive.open("wb") as handle:
            result = subprocess.run(
                ["git", "archive", "--format=tar", commit],
                cwd=repository,
                stdout=handle,
                stderr=subprocess.PIPE,
                check=False,
                timeout=120,
            )
        if result.returncode != 0:
            raise RepairExperimentError(
                "git archive failed: " + result.stderr.decode("utf-8", errors="replace")
            )
        with tarfile.open(archive, "r") as handle:
            handle.extractall(tree, filter="data")
        yield tree


@contextmanager
def isolated_git_tree(
    repository: Path,
    commit: str,
    repository_url: str,
) -> Iterator[Path]:
    """Create an isolated shared clone with exact Git metadata for test suites."""

    with tempfile.TemporaryDirectory(prefix="agent-harness-study2-") as temporary:
        tree = Path(temporary) / "tree"
        clone = subprocess.run(
            [
                "git",
                "clone",
                "--quiet",
                "--shared",
                "--no-checkout",
                str(repository.resolve()),
                str(tree),
            ],
            text=True,
            capture_output=True,
            check=False,
            timeout=180,
        )
        if clone.returncode != 0:
            raise RepairExperimentError(
                "isolated local clone failed: " + (clone.stderr.strip() or clone.stdout.strip())
            )
        for arguments in (
            ["checkout", "--quiet", "--detach", commit],
            ["remote", "set-url", "origin", repository_url],
        ):
            result = subprocess.run(
                ["git", *arguments],
                cwd=tree,
                text=True,
                capture_output=True,
                check=False,
                timeout=180,
            )
            if result.returncode != 0:
                raise RepairExperimentError(
                    f"isolated git {' '.join(arguments)} failed: "
                    + (result.stderr.strip() or result.stdout.strip())
                )
        yield tree


def _apply_patch(tree: Path, patch_path: Path) -> dict[str, Any]:
    started = time.monotonic()
    result = subprocess.run(
        ["git", "apply", "--whitespace=nowarn", str(patch_path)],
        cwd=tree,
        text=True,
        capture_output=True,
        check=False,
        timeout=120,
    )
    return {
        "returncode": result.returncode,
        "stdout": result.stdout,
        "stderr": result.stderr,
        "elapsed_seconds": time.monotonic() - started,
    }


def run_test_command(tree: Path, command: str, timeout_seconds: int = 600) -> dict[str, Any]:
    arguments = shlex.split(command)
    is_go = arguments[:2] == ["go", "test"]
    is_python = arguments[:3] == ["python", "-m", "pytest"]
    if not (is_go or is_python):
        raise RepairExperimentError(
            f"only frozen `go test` or `python -m pytest` commands are permitted: {command}"
        )
    if is_python:
        arguments[0] = sys.executable
    environment = None
    if is_go:
        environment = dict(os.environ)
        cache_root = Path(tempfile.gettempdir()) / "agent-harness-go-runtime"
        build_cache = cache_root / "build"
        module_cache = cache_root / "modules"
        build_cache.mkdir(parents=True, exist_ok=True)
        module_cache.mkdir(parents=True, exist_ok=True)
        environment["GOCACHE"] = str(build_cache)
        environment["GOMODCACHE"] = str(module_cache)
    started = time.monotonic()
    try:
        result = subprocess.run(
            arguments,
            cwd=tree,
            text=True,
            capture_output=True,
            check=False,
            timeout=timeout_seconds,
            env=environment,
        )
        return {
            "command": command,
            "returncode": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr,
            "elapsed_seconds": time.monotonic() - started,
            "timed_out": False,
        }
    except subprocess.TimeoutExpired as exc:
        return {
            "command": command,
            "returncode": None,
            "stdout": exc.stdout or "",
            "stderr": exc.stderr or "",
            "elapsed_seconds": time.monotonic() - started,
            "timed_out": True,
        }


def validate_generated_patch(
    root: Path,
    repository: Path,
    task: TaskSpec,
    patch: str,
    preserve_git_metadata: bool = False,
) -> dict[str, Any]:
    context = (
        isolated_git_tree(repository, task.base_commit, task.repository_url)
        if preserve_git_metadata
        else isolated_source_tree(repository, task.base_commit)
    )
    with context as tree:
        test_patch = (root / "tasks" / task.test_patch).resolve()
        hidden_apply = _apply_patch(tree, test_patch)
        if hidden_apply["returncode"] != 0:
            raise RepairExperimentError(
                f"frozen hidden test patch failed to apply for {task.task_id}: {hidden_apply['stderr']}"
            )
        model_patch = tree.parent / "model.patch"
        model_patch.write_text(patch, encoding="utf-8")
        model_apply = _apply_patch(tree, model_patch)
        tests: list[dict[str, Any]] = []
        if model_apply["returncode"] == 0:
            for command in dict.fromkeys((*task.fail_to_pass_tests, *task.pass_to_pass_tests)):
                tests.append(run_test_command(tree, command))
        by_command = {item["command"]: item for item in tests}
        fail_to_pass = all(
            by_command.get(command, {}).get("returncode") == 0
            for command in task.fail_to_pass_tests
        )
        pass_to_pass = all(
            by_command.get(command, {}).get("returncode") == 0
            for command in task.pass_to_pass_tests
        )
        resolved = model_apply["returncode"] == 0 and fail_to_pass and pass_to_pass
        return {
            "hidden_test_patch_apply": hidden_apply,
            "model_patch_apply": model_apply,
            "tests": tests,
            "fail_to_pass": fail_to_pass,
            "pass_to_pass": pass_to_pass,
            "resolved_at_1": resolved,
            "failure_stage": (
                "resolved"
                if resolved
                else "patch_apply"
                if model_apply["returncode"] != 0
                else "tests"
            ),
        }


def _identity(
    experiment: Any,
    task: TaskSpec,
    harness: HarnessSpec,
    model: Any,
    revision: str,
) -> RunIdentity:
    return RunIdentity(
        experiment_id=experiment.experiment_id,
        task_id=task.task_id,
        harness_id=harness.harness_id,
        harness_hash=harness.config_hash,
        model_id=model.model_id,
        model_key=model.expected_inference_key,
        model_config_hash=model.config_hash,
        context_budget=experiment.context_budgets[0],
        seed=experiment.seeds[0],
        repetition=0,
        repository_sha=task.base_commit,
        code_revision=revision,
    )


def run_repair_experiment(
    root: Path,
    repository: Path,
    task_filter: set[str] | None = None,
    harness_filter: set[str] | None = None,
) -> dict[str, Any]:
    revision = research_code_revision(root)
    experiment = load_experiments(root)["E03"]
    model = load_models(root)[experiment.model_ids[0]]
    harness_catalog = load_harnesses(root)
    task_catalog = load_tasks(root)
    split = load_task_split(root / "tasks" / "splits" / f"{experiment.task_split}.txt")
    tasks = [
        task_catalog[item]
        for item in split
        if task_filter is None or item in task_filter
    ]
    harnesses = [
        harness_catalog[item]
        for item in experiment.harness_ids
        if harness_filter is None or item in harness_filter
    ]
    if not tasks or not harnesses:
        raise RepairExperimentError("filters selected no E03 cells")
    if any(task.validation_status != "end_to_end_ready" for task in tasks):
        raise RepairExperimentError("E03 split contains a task that is not end-to-end ready")

    client = LMStudioClient(model, timeout_seconds=experiment.timeout_seconds)
    discovery, resolved = client.resolve()
    loaded = exclusive_loaded(discovery.native_models)
    if loaded != (model.expected_inference_key,):
        raise RepairExperimentError(f"E03 requires exclusive Qwen residency: {loaded}")
    snapshot = GitSnapshot(repository)
    tokenizer = QwenTokenCounter()
    rows: list[dict[str, Any]] = []
    for task in tasks:
        snapshot.verify_commit(task.base_commit)
        for harness in harnesses:
            identity = _identity(experiment, task, harness, model, revision)
            directory = run_directory(root / "results", identity)
            if directory.exists():
                final_path = directory / "final_metrics.json"
                if not final_path.exists():
                    raise RepairExperimentError(f"incomplete existing run directory: {directory}")
                rows.append(json.loads(final_path.read_text(encoding="utf-8")))
                continue
            evidence, allowed_paths, context_tokens, evidence_source = repair_context(
                root, task, harness, snapshot, tokenizer
            )
            prompt = repair_prompt(task, evidence)
            started = time.monotonic()
            response = client.chat_completions(
                resolved.inference_key,
                [
                    {"role": "system", "content": REPAIR_SYSTEM},
                    {"role": "user", "content": prompt},
                ],
                max_tokens=model.max_tokens,
            )
            model_elapsed = time.monotonic() - started
            protocol_violation: str | None = None
            patch = ""
            paths: tuple[str, ...] = ()
            validation: dict[str, Any]
            try:
                patch = extract_unified_diff(response)
                paths = validate_patch_scope(patch, allowed_paths)
                validation = validate_generated_patch(root, repository, task, patch)
            except PatchOutputError as exc:
                protocol_violation = str(exc)
                validation = {
                    "hidden_test_patch_apply": None,
                    "model_patch_apply": None,
                    "tests": [],
                    "fail_to_pass": False,
                    "pass_to_pass": False,
                    "resolved_at_1": False,
                    "failure_stage": "protocol_violation",
                }
            localization = retrieval_metrics(paths, task.gold_files)
            final = {
                "run_id": identity.run_id,
                "experiment_id": "E03",
                "task_id": task.task_id,
                "harness_id": harness.harness_id,
                "resolved_at_1": validation["resolved_at_1"],
                "failure_stage": validation["failure_stage"],
                "patch_applied": bool(
                    validation["model_patch_apply"]
                    and validation["model_patch_apply"].get("returncode") == 0
                ),
                "fail_to_pass": validation["fail_to_pass"],
                "pass_to_pass": validation["pass_to_pass"],
                "modified_files": paths,
                "localization_metrics": localization,
                "protocol_violation": protocol_violation,
                "evidence_source": evidence_source,
                "context_tokens": context_tokens,
                "prompt_sha256": sha256(prompt.encode()).hexdigest(),
                "patch_sha256": sha256(patch.encode()).hexdigest() if patch else None,
                "elapsed_seconds": model_elapsed + sum(
                    item.get("elapsed_seconds", 0.0) for item in validation["tests"]
                ),
                "model_elapsed_seconds": model_elapsed,
                "usage": response.get("usage", {}),
                "test_results": validation["tests"],
            }
            with EventWriter(
                root / "results", identity, asdict(harness), resolved.to_dict()
            ) as writer:
                writer.emit("run_started", {"confirmatory": True, "evidence_source": evidence_source})
                writer.emit("model_call", {"elapsed_seconds": model_elapsed, "usage": response.get("usage", {})})
                writer.write_artifact("prompt.txt", prompt)
                writer.write_artifact("model_response.json", json.dumps(response, indent=2) + "\n")
                writer.write_artifact("model.patch", patch)
                writer.write_artifact("validation.json", json.dumps(validation, indent=2) + "\n")
                writer.write_artifact("final_metrics.json", json.dumps(final, indent=2) + "\n")
                for test in validation["tests"]:
                    writer.emit("test_run", test)
                writer.emit(
                    "run_finished",
                    {
                        "status": "completed_with_protocol_violation" if protocol_violation else "completed",
                        "resolved_at_1": validation["resolved_at_1"],
                        "failure_stage": validation["failure_stage"],
                    },
                )
            rows.append(final)
    return {
        "experiment_id": "E03",
        "code_revision": revision,
        "run_count": len(rows),
        "resolved_count": sum(bool(item["resolved_at_1"]) for item in rows),
        "rows": rows,
    }