src/agent_harness/telemetry.py

"""Append-only run telemetry and deterministic artifact layout."""

from __future__ import annotations

from dataclasses import asdict, dataclass
from datetime import datetime, timezone
from hashlib import sha256
import json
import os
from pathlib import Path
import time
from typing import Any


def utc_now() -> str:
    return datetime.now(timezone.utc).isoformat()


@dataclass(frozen=True, slots=True)
class RunIdentity:
    experiment_id: str
    task_id: str
    harness_id: str
    harness_hash: str
    model_id: str
    model_key: str
    model_config_hash: str
    context_budget: int
    seed: int
    repetition: int
    repository_sha: str
    code_revision: str

    @property
    def run_id(self) -> str:
        payload = json.dumps(asdict(self), sort_keys=True, separators=(",", ":"))
        return sha256(payload.encode("utf-8")).hexdigest()[:20]


def run_directory(results_root: Path, identity: RunIdentity) -> Path:
    return (
        results_root
        / "raw"
        / identity.experiment_id
        / identity.harness_id
        / identity.task_id
        / identity.run_id
    )


def load_completed_or_archive_incomplete(
    results_root: Path, identity: RunIdentity
) -> dict[str, Any] | None:
    """Resume a completed cell or retain an interrupted attempt outside raw/."""

    directory = run_directory(results_root, identity)
    if not directory.exists():
        return None
    final_path = directory / "final_metrics.json"
    if final_path.exists():
        value = json.loads(final_path.read_text(encoding="utf-8"))
        if not isinstance(value, dict):
            raise ValueError(f"completed metrics are not a JSON object: {final_path}")
        return value
    archive_root = (
        results_root
        / "infrastructure_attempts"
        / identity.experiment_id
        / identity.harness_id
        / identity.task_id
    )
    archive_root.mkdir(parents=True, exist_ok=True)
    destination = archive_root / f"{identity.run_id}-{time.time_ns()}"
    directory.replace(destination)
    with (destination / "archive_record.json").open("x", encoding="utf-8") as handle:
        json.dump(
            {
                "schema_version": 1,
                "reason": "incomplete cell retained before identical-identity retry",
                "identity": asdict(identity),
                "archived_at": utc_now(),
                "original_directory": str(directory),
            },
            handle,
            indent=2,
            sort_keys=True,
        )
        handle.write("\n")
    return None


class EventWriter:
    """Writes one durable JSON object per event and never overwrites a run."""

    def __init__(
        self,
        results_root: Path,
        identity: RunIdentity,
        resolved_harness: dict[str, Any],
        resolved_model: dict[str, Any],
    ):
        self.identity = identity
        self.directory = run_directory(results_root, identity)
        self.directory.mkdir(parents=True, exist_ok=False)
        self._sequence = 0
        self._stream = (self.directory / "trajectory.jsonl").open("x", encoding="utf-8")
        self._write_json_exclusive(
            self.directory / "run_manifest.json",
            {
                "schema_version": 1,
                "created_at": utc_now(),
                "run_id": identity.run_id,
                "identity": asdict(identity),
                "resolved_harness": resolved_harness,
                "resolved_model": resolved_model,
            },
        )

    @staticmethod
    def _write_json_exclusive(path: Path, value: dict[str, Any]) -> None:
        with path.open("x", encoding="utf-8") as handle:
            json.dump(value, handle, indent=2, sort_keys=True)
            handle.write("\n")

    def emit(self, event_type: str, payload: dict[str, Any]) -> None:
        event = {
            "schema_version": 1,
            "run_id": self.identity.run_id,
            "sequence": self._sequence,
            "recorded_at": utc_now(),
            "event_type": event_type,
            "payload": payload,
        }
        self._stream.write(json.dumps(event, sort_keys=True, separators=(",", ":")) + "\n")
        self._stream.flush()
        os.fsync(self._stream.fileno())
        self._sequence += 1

    def write_artifact(self, filename: str, content: str) -> Path:
        if Path(filename).name != filename:
            raise ValueError("artifact filename must not contain directories")
        path = self.directory / filename
        with path.open("x", encoding="utf-8") as handle:
            handle.write(content)
        return path

    def close(self) -> None:
        if not self._stream.closed:
            self._stream.close()

    def __enter__(self) -> "EventWriter":
        return self

    def __exit__(self, exc_type: Any, exc: Any, traceback: Any) -> None:
        self.close()


ALLOWED_EVENT_TYPES = {
    "run_started",
    "model_call",
    "tool_call",
    "retrieval_candidate",
    "file_read",
    "edit",
    "test_run",
    "resource_sample",
    "run_finished",
}