src/agent_harness/study5_experiment.py

"""Manifest-driven prospective Study 5 harness-factor experiments."""

from __future__ import annotations

from hashlib import sha256
import json
from pathlib import Path
import time
from typing import Any

from .lm_studio_embeddings import LMStudioEmbeddingClient
from .lm_studio_management import LMStudioResidencyManager, LMStudioServer
from .pilot import research_code_revision
from .protocol_experiment import (
    ProtocolExperimentError,
    _build_task_retrieval,
    _repository_for_task,
    run_protocol_cell,
)
from .repository import GitSnapshot
from .retrieval import SQLiteEmbeddingCache
from .specs import (
    load_edit_interfaces,
    load_embeddings,
    load_experiments,
    load_harnesses,
    load_models,
    load_repositories,
    load_tasks,
)
from .study2_experiment import _RuntimeLease


class Study5ExperimentError(RuntimeError):
    """Raised when a Study 5 manifest or runtime violates its frozen design."""


def _manifest_hash(value: dict[str, Any]) -> str:
    payload = dict(value)
    expected = payload.pop("design_sha256", None)
    observed = sha256(
        json.dumps(payload, sort_keys=True, separators=(",", ":")).encode("utf-8")
    ).hexdigest()
    if expected != observed:
        raise Study5ExperimentError(
            f"Study 5 manifest hash mismatch: expected {expected}, observed {observed}"
        )
    return observed


def _write_progress(
    root: Path,
    experiment_id: str,
    revision: str,
    manifest_hash: str,
    planned: int,
    rows: list[dict[str, Any]],
) -> Path:
    path = root / "results" / "reports" / f"{experiment_id}_progress.json"
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(
        json.dumps(
            {
                "schema_version": 1,
                "experiment_id": experiment_id,
                "code_revision": revision,
                "manifest_sha256": manifest_hash,
                "planned_cells": planned,
                "completed_cells": len(rows),
                "accepted_edit_cells": sum(bool(row["accepted_edit_cell"]) for row in rows),
                "applicable_patch_cells": sum(bool(row["applicable_final_patch"]) for row in rows),
                "resolved_cells": sum(bool(row["resolved_at_1"]) for row in rows),
                "rows": rows,
            },
            indent=2,
            sort_keys=True,
        )
        + "\n",
        encoding="utf-8",
    )
    return path


def run_study5_experiment(
    root: Path,
    experiment_id: str,
    task_filter: set[str] | None = None,
    harness_filter: set[str] | None = None,
    interface_filter: set[str] | None = None,
    model_filter: set[str] | None = None,
    stop_server_when_complete: bool = True,
) -> dict[str, Any]:
    if experiment_id not in {"E13", "E14", "E15", "E16"}:
        raise Study5ExperimentError("Study 5 runner requires E13, E14, E15, or E16")
    revision = research_code_revision(root)
    experiment = load_experiments(root)[experiment_id]
    manifest_path = root / "configs" / "study5" / f"{experiment_id}_cells.json"
    manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
    manifest_hash = _manifest_hash(manifest)
    if manifest.get("experiment_id") != experiment_id or not manifest.get("outcome_blind"):
        raise Study5ExperimentError("Study 5 manifest identity/freeze flag mismatch")
    all_cells = manifest.get("cells")
    if not isinstance(all_cells, list) or len(all_cells) != int(manifest["planned_cells"]):
        raise Study5ExperimentError("Study 5 manifest cell count mismatch")

    tasks = load_tasks(root)
    harnesses = load_harnesses(root)
    interfaces = load_edit_interfaces(root)
    models = load_models(root)
    repositories = load_repositories(root)
    embedding = load_embeddings(root)[experiment.embedding_id]
    cells = [
        item
        for item in all_cells
        if (task_filter is None or item["task_id"] in task_filter)
        and (harness_filter is None or item["harness_id"] in harness_filter)
        and (interface_filter is None or item["interface_id"] in interface_filter)
        and (model_filter is None or item["model_id"] in model_filter)
    ]
    if not cells:
        raise Study5ExperimentError("Study 5 filters selected an empty execution block")

    identities: set[tuple[str, str, str, str]] = set()
    for item in cells:
        task = tasks[item["task_id"]]
        harness = harnesses[item["harness_id"]]
        interface = interfaces[item["interface_id"]]
        model = models[item["model_id"]]
        identity = (task.task_id, harness.harness_id, interface.interface_id, model.model_id)
        if identity in identities:
            raise Study5ExperimentError(f"duplicate Study 5 cell: {identity}")
        identities.add(identity)
        expected = (
            task.base_commit,
            harness.config_hash,
            interface.config_hash,
            model.config_hash,
        )
        observed = (
            item["repository_sha"],
            item["harness_hash"],
            item["interface_hash"],
            item["model_hash"],
        )
        if observed != expected:
            raise Study5ExperimentError(f"frozen configuration drift for {identity}")
        if task.validation_status != "end_to_end_ready":
            raise Study5ExperimentError(f"{task.task_id} is not end-to-end ready")

    server = LMStudioServer(port=1234)
    first_model = models[cells[0]["model_id"]]
    residency = LMStudioResidencyManager(
        first_model.base_url,
        first_model.api_token_env,
        timeout_seconds=experiment.timeout_seconds,
    )
    embedding_client = LMStudioEmbeddingClient(
        embedding, timeout_seconds=experiment.timeout_seconds
    )
    cache_path = root / "indexes" / "embeddings" / f"{embedding.config_hash}.sqlite3"
    rows: list[dict[str, Any]] = []
    task_summaries: list[dict[str, Any]] = []
    runtime = _RuntimeLease(server, residency, stop_server_when_complete)
    grouped: dict[str, list[dict[str, Any]]] = {}
    for cell in cells:
        grouped.setdefault(str(cell["task_id"]), []).append(cell)

    with runtime as server_state, SQLiteEmbeddingCache(cache_path, embedding) as cache:
        for task_id, task_cells in grouped.items():
            task = tasks[task_id]
            repository_spec = _repository_for_task(repositories, task)
            repository = (root / repository_spec.local_path).resolve()
            snapshot = GitSnapshot(repository)
            snapshot.verify_commit(task.base_commit)
            index_transition = residency.ensure_exclusive(
                embedding.model_key, embedding.loaded_context_length
            )
            embedding_client.resolve()
            index_started = time.monotonic()
            retrieval = _build_task_retrieval(snapshot, task, embedding, embedding_client, cache)
            index_elapsed = time.monotonic() - index_started
            task_rows: list[dict[str, Any]] = []
            for cell in sorted(task_cells, key=lambda item: int(item["order"])):
                row = run_protocol_cell(
                    root,
                    repository,
                    experiment,
                    task,
                    interfaces[cell["interface_id"]],
                    models[cell["model_id"]],
                    residency,
                    server,
                    revision,
                    retrieval_harness=harnesses[cell["harness_id"]],
                    retrieval=retrieval,
                    embedding=embedding,
                    seed=int(cell["seed"]),
                    context_budget=int(cell["context_budget"]),
                )
                rows.append(row)
                task_rows.append(row)
                _write_progress(
                    root, experiment_id, revision, manifest_hash, len(cells), rows
                )
            task_summaries.append(
                {
                    "task_id": task_id,
                    "repository_id": repository_spec.repository_id,
                    "language": task.language,
                    "cells": len(task_rows),
                    "accepted_edits": sum(bool(row["accepted_edit_cell"]) for row in task_rows),
                    "applicable_patches": sum(bool(row["applicable_final_patch"]) for row in task_rows),
                    "resolved": sum(bool(row["resolved_at_1"]) for row in task_rows),
                    "embedding_index_transition": index_transition.to_dict(),
                    "index_elapsed_seconds": index_elapsed,
                    "dense_index_stats": retrieval.dense_index_stats,
                }
            )

    if len(rows) != len(cells):
        raise Study5ExperimentError(f"finalized {len(rows)}/{len(cells)} selected cells")
    report = {
        "schema_version": 1,
        "experiment_id": experiment_id,
        "code_revision": revision,
        "manifest_sha256": manifest_hash,
        "planned_cells": len(cells),
        "run_count": len(rows),
        "accepted_edit_count": sum(bool(row["accepted_edit_cell"]) for row in rows),
        "applicable_patch_count": sum(bool(row["applicable_final_patch"]) for row in rows),
        "resolved_count": sum(bool(row["resolved_at_1"]) for row in rows),
        "server_lifecycle": server_state,
        "server_stop": runtime.stop_state,
        "final_residency_transition": runtime.final_transition,
        "cleanup_errors": runtime.cleanup_errors,
        "task_summaries": task_summaries,
        "rows": rows,
    }
    report_path = (
        root
        / "results"
        / "reports"
        / f"{experiment_id}_{revision[:12]}_{int(time.time())}.json"
    )
    report_path.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    return {**report, "report_path": str(report_path)}