rein-aharness/rein_aharness/metrics.py

"""Write per-run metrics to compatibility or durable external storage.

Legacy grant-absent runs retain the kaizen-agentic ADR-004 repository layout.
Accepted grant runs use private external state and a projection descriptor so
metrics cannot dirty the validated checkout.

  .kaizen/metrics/<agent>/
    executions.jsonl   # append-only
    summary.json       # regenerated on write
"""

from __future__ import annotations

import fcntl
import hashlib
import json
import os
import uuid
from dataclasses import asdict, dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Any


@dataclass
class ExecutionRecord:
    timestamp: str
    agent: str
    success: bool
    execution_time_s: float = 0.0
    session_id: str | None = None
    quality_score: float | None = None
    primary_metric: dict[str, Any] | None = None
    metadata: dict[str, Any] = field(default_factory=dict)
    # Helix / harness correlation (ADR-004 optional fields)
    repo: str | None = None
    tokens: int | None = None
    committed: bool | None = None
    head_after: str | None = None
    reason: str | None = None
    harness: str = "rein-aharness"

    def to_json_line(self) -> str:
        data = asdict(self)
        # Drop Nones for a compact record; required fields always present.
        compact = {k: v for k, v in data.items() if v is not None}
        return json.dumps(compact, sort_keys=True)


def metrics_dir(project_root: Path, agent: str) -> Path:
    return Path(project_root) / ".kaizen" / "metrics" / agent


def external_metrics_dir(
    project_root: Path | str,
    agent: str,
    *,
    state_dir: Path | None = None,
) -> Path:
    """Return the private durable metrics directory for a granted run."""
    root = Path(project_root).expanduser().resolve()
    base = state_dir.expanduser().resolve() if state_dir else _state_dir()
    repo_id = hashlib.sha256(str(root).encode("utf-8")).hexdigest()[:32]
    agent_id = hashlib.sha256(str(agent).encode("utf-8")).hexdigest()[:32]
    return base / "execution-metrics" / repo_id / agent_id


def _utc_now() -> str:
    return datetime.now(timezone.utc).replace(microsecond=0).isoformat().replace(
        "+00:00", "Z"
    )


def _load_executions(path: Path) -> list[dict[str, Any]]:
    if not path.is_file():
        return []
    records: list[dict[str, Any]] = []
    for line in path.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line:
            continue
        try:
            records.append(json.loads(line))
        except json.JSONDecodeError:
            continue
    return records


def _trend(values: list[float]) -> str:
    if len(values) < 4:
        return "stable"
    mid = len(values) // 2
    early = sum(values[:mid]) / max(1, mid)
    late = sum(values[mid:]) / max(1, len(values) - mid)
    if late - early > 0.05:
        return "up"
    if early - late > 0.05:
        return "down"
    return "stable"


def regenerate_summary(agent: str, executions: list[dict[str, Any]]) -> dict[str, Any]:
    count = len(executions)
    successes = [e for e in executions if e.get("success")]
    success_rate = (len(successes) / count) if count else 0.0
    times = [float(e["execution_time_s"]) for e in executions if "execution_time_s" in e]
    qualities = [
        float(e["quality_score"])
        for e in executions
        if isinstance(e.get("quality_score"), (int, float))
    ]
    last_ts = executions[-1].get("timestamp") if executions else None
    return {
        "agent": agent,
        "execution_count": count,
        "success_rate": round(success_rate, 3),
        "avg_quality_score": (
            round(sum(qualities) / len(qualities), 3) if qualities else None
        ),
        "avg_execution_time_s": (
            round(sum(times) / len(times), 3) if times else None
        ),
        "last_execution": last_ts,
        "trend": {
            "success_rate": _trend(
                [1.0 if e.get("success") else 0.0 for e in executions]
            ),
            "quality_score": _trend(qualities) if qualities else "stable",
        },
    }


def record_execution(
    project_root: Path | str,
    agent: str,
    *,
    success: bool,
    execution_time_s: float = 0.0,
    tokens: int | None = None,
    committed: bool | None = None,
    head_after: str | None = None,
    reason: str | None = None,
    metadata: dict[str, Any] | None = None,
    session_id: str | None = None,
) -> Path:
    """Append one execution record and regenerate summary.json.

    Returns the path to executions.jsonl. Never raises for missing parent
    dirs (creates them). Callers that must not write should skip this.
    """
    root = Path(project_root)
    directory = metrics_dir(root, agent)
    directory.mkdir(parents=True, exist_ok=True)

    record = _execution_record(
        root,
        agent,
        success=success,
        execution_time_s=execution_time_s,
        tokens=tokens,
        committed=committed,
        head_after=head_after,
        reason=reason,
        metadata=metadata,
        session_id=session_id,
    )

    executions_path = directory / "executions.jsonl"
    with executions_path.open("a", encoding="utf-8") as fh:
        fh.write(record.to_json_line() + "\n")

    all_records = _load_executions(executions_path)
    summary = regenerate_summary(agent, all_records)
    (directory / "summary.json").write_text(
        json.dumps(summary, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    return executions_path


def record_external_execution(
    project_root: Path | str,
    agent: str,
    *,
    success: bool,
    execution_time_s: float = 0.0,
    tokens: int | None = None,
    committed: bool | None = None,
    head_after: str | None = None,
    reason: str | None = None,
    metadata: dict[str, Any] | None = None,
    session_id: str | None = None,
    state_dir: Path | None = None,
) -> Path:
    """Durably record granted-run metrics without dirtying the target checkout."""
    root = Path(project_root).expanduser().resolve()
    projection_target = _projection_target(agent)
    directory = external_metrics_dir(root, agent, state_dir=state_dir)
    _ensure_private_directory(directory)
    lock_path = directory / ".lock"
    lock_fd = os.open(lock_path, os.O_RDWR | os.O_CREAT, 0o600)
    try:
        os.fchmod(lock_fd, 0o600)
        fcntl.flock(lock_fd, fcntl.LOCK_EX)
        record = _execution_record(
            root,
            agent,
            success=success,
            execution_time_s=execution_time_s,
            tokens=tokens,
            committed=committed,
            head_after=head_after,
            reason=reason,
            metadata=metadata,
            session_id=session_id,
        )
        executions_path = directory / "executions.jsonl"
        existing = (
            executions_path.read_text(encoding="utf-8")
            if executions_path.exists()
            else ""
        )
        if existing and not existing.endswith("\n"):
            raise OSError("external metrics ledger has an incomplete record")
        records = _load_external_executions(executions_path) if existing else []
        already_recorded = session_id is not None and any(
            item.get("session_id") == session_id for item in records
        )
        if not already_recorded:
            record_line = record.to_json_line()
            _atomic_write_text(executions_path, existing + record_line + "\n")
            records.append(json.loads(record_line))
        _atomic_write_text(
            directory / "summary.json",
            json.dumps(regenerate_summary(agent, records), indent=2, sort_keys=True)
            + "\n",
        )
        repo_id = directory.parent.name
        _atomic_write_text(
            directory / "projection.json",
            json.dumps(
                {
                    "agent": agent,
                    "repository_id": repo_id,
                    "repository_name": root.name,
                    "target_relative_directory": projection_target,
                },
                indent=2,
                sort_keys=True,
            )
            + "\n",
        )
        _fsync_directory(directory)
        return executions_path
    finally:
        try:
            fcntl.flock(lock_fd, fcntl.LOCK_UN)
        finally:
            os.close(lock_fd)


def _projection_target(agent: str) -> str:
    """Return a checkout-relative projection target for a safe agent identity."""
    if (
        not agent
        or agent in {".", ".."}
        or "/" in agent
        or "\\" in agent
        or "\x00" in agent
        or len(agent) > 200
    ):
        raise OSError("agent identity cannot form a safe metrics projection")
    return f".kaizen/metrics/{agent}"


def _execution_record(
    root: Path,
    agent: str,
    *,
    success: bool,
    execution_time_s: float,
    tokens: int | None,
    committed: bool | None,
    head_after: str | None,
    reason: str | None,
    metadata: dict[str, Any] | None,
    session_id: str | None,
) -> ExecutionRecord:
    return ExecutionRecord(
        timestamp=_utc_now(),
        agent=agent,
        success=success,
        execution_time_s=float(execution_time_s),
        session_id=session_id,
        metadata=metadata or {},
        repo=root.name,
        tokens=tokens,
        committed=committed,
        head_after=head_after,
        reason=reason,
    )


def _load_external_executions(path: Path) -> list[dict[str, Any]]:
    records: list[dict[str, Any]] = []
    for line in path.read_text(encoding="utf-8").splitlines():
        try:
            value = json.loads(line)
        except json.JSONDecodeError as exc:
            raise OSError("external metrics ledger contains invalid JSON") from exc
        if not isinstance(value, dict):
            raise OSError("external metrics ledger contains a non-object record")
        records.append(value)
    return records


def _state_dir() -> Path:
    explicit = os.environ.get("REIN_AHARNESS_STATE_DIR", "").strip()
    if explicit:
        return Path(explicit).expanduser().resolve()
    xdg = os.environ.get("XDG_STATE_HOME", "").strip()
    if xdg:
        return (Path(xdg).expanduser() / "rein-aharness").resolve()
    return (Path.home() / ".local" / "state" / "rein-aharness").resolve()


def _ensure_private_directory(path: Path) -> None:
    path.mkdir(parents=True, exist_ok=True)
    current = path
    while current.name and current != current.parent:
        os.chmod(current, 0o700)
        if current.name == "execution-metrics":
            break
        current = current.parent


def _atomic_write_text(path: Path, value: str) -> None:
    temporary = path.parent / f".{path.name}.{uuid.uuid4().hex}.tmp"
    fd = os.open(temporary, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
    try:
        with os.fdopen(fd, "w", encoding="utf-8") as handle:
            fd = -1
            handle.write(value)
            handle.flush()
            os.fsync(handle.fileno())
        os.replace(temporary, path)
        os.chmod(path, 0o600)
        _fsync_directory(path.parent)
    except BaseException:
        if fd >= 0:
            os.close(fd)
        try:
            temporary.unlink()
        except FileNotFoundError:
            pass
        raise


def _fsync_directory(path: Path) -> None:
    fd = os.open(path, os.O_RDONLY)
    try:
        os.fsync(fd)
    finally:
        os.close(fd)