db.py

from __future__ import annotations

import json
import sqlite3
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

from .types import ExistingRun


def _json_dumps(value: Any) -> str:
    return json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=True)


def connect_database(db_path: Path) -> sqlite3.Connection:
    db_path.parent.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(str(db_path))
    conn.row_factory = sqlite3.Row
    return conn


def _column_exists(conn: sqlite3.Connection, table: str, column: str) -> bool:
    rows = conn.execute(f"PRAGMA table_info({table})").fetchall()
    return any(row[1] == column for row in rows)


def ensure_comparison_id_column(conn: sqlite3.Connection) -> None:
    """Idempotently add ``comparison_id`` to ``benchmark_runs`` for old DBs."""
    if not _column_exists(conn, "benchmark_runs", "comparison_id"):
        conn.execute("ALTER TABLE benchmark_runs ADD COLUMN comparison_id TEXT")
    conn.execute(
        "CREATE INDEX IF NOT EXISTS idx_benchmark_runs_comparison "
        "ON benchmark_runs(comparison_id)"
    )
    conn.commit()


def _ensure_column(conn: sqlite3.Connection, table: str, column: str, declaration: str) -> None:
    if not _column_exists(conn, table, column):
        conn.execute(f"ALTER TABLE {table} ADD COLUMN {column} {declaration}")


def ensure_metrics_columns(conn: sqlite3.Connection) -> None:
    for column, declaration in (
        ("trajectory_summary_json", "TEXT"),
        ("token_metrics_json", "TEXT"),
        ("cache_metrics_json", "TEXT"),
        ("performance_metrics_json", "TEXT"),
        ("trajectory_count", "INTEGER"),
        ("llm_call_count", "INTEGER"),
        ("total_prompt_tokens", "INTEGER"),
        ("total_completion_tokens", "INTEGER"),
        ("total_cache_read_input_tokens", "INTEGER"),
        ("total_cache_creation_input_tokens", "INTEGER"),
        ("mean_latency_ms", "REAL"),
        ("p95_latency_ms", "REAL"),
        ("throughput_per_second", "REAL"),
    ):
        _ensure_column(conn, "benchmark_runs", column, declaration)
    conn.commit()


def initialize_database(conn: sqlite3.Connection) -> None:
    conn.executescript(
        """
        PRAGMA journal_mode=WAL;

        CREATE TABLE IF NOT EXISTS run_groups (
            run_group_id TEXT PRIMARY KEY,
            created_at TEXT NOT NULL,
            finished_at TEXT,
            request_json TEXT NOT NULL,
            benchmarks_json TEXT NOT NULL,
            repo_meta_json TEXT NOT NULL,
            created_by TEXT NOT NULL
        );

        CREATE TABLE IF NOT EXISTS benchmark_runs (
            run_id TEXT PRIMARY KEY,
            run_group_id TEXT NOT NULL,
            benchmark_id TEXT NOT NULL,
            benchmark_directory TEXT NOT NULL,
            signature TEXT NOT NULL,
            status TEXT NOT NULL,
            attempt INTEGER NOT NULL,
            agent TEXT NOT NULL,
            provider TEXT NOT NULL,
            model TEXT NOT NULL,
            extra_config_json TEXT NOT NULL,
            started_at TEXT NOT NULL,
            ended_at TEXT,
            duration_seconds REAL,
            command_json TEXT NOT NULL,
            cwd TEXT NOT NULL,
            stdout_path TEXT NOT NULL,
            stderr_path TEXT NOT NULL,
            result_json_path TEXT,
            score REAL,
            unit TEXT,
            higher_is_better INTEGER,
            metrics_json TEXT NOT NULL,
            artifacts_json TEXT NOT NULL,
            error TEXT,
            high_score_label TEXT,
            high_score_value REAL,
            delta_to_high_score REAL,
            benchmark_version TEXT,
            benchmarks_commit TEXT,
            eliza_commit TEXT,
            eliza_version TEXT,
            created_at TEXT NOT NULL,
            comparison_id TEXT,
            FOREIGN KEY(run_group_id) REFERENCES run_groups(run_group_id)
        );

        CREATE TABLE IF NOT EXISTS benchmark_run_trajectories (
            run_id TEXT NOT NULL,
            trajectory_file TEXT NOT NULL,
            turn_index INTEGER NOT NULL,
            prompt_tokens INTEGER NOT NULL DEFAULT 0,
            completion_tokens INTEGER NOT NULL DEFAULT 0,
            total_tokens INTEGER NOT NULL DEFAULT 0,
            cached_tokens INTEGER NOT NULL DEFAULT 0,
            cache_creation_tokens INTEGER NOT NULL DEFAULT 0,
            latency_ms REAL,
            prompt_chars INTEGER NOT NULL DEFAULT 0,
            PRIMARY KEY(run_id, trajectory_file, turn_index),
            FOREIGN KEY(run_id) REFERENCES benchmark_runs(run_id)
        );

        CREATE INDEX IF NOT EXISTS idx_benchmark_runs_signature
            ON benchmark_runs(signature);
        CREATE INDEX IF NOT EXISTS idx_benchmark_runs_signature_status
            ON benchmark_runs(signature, status, ended_at);
        CREATE INDEX IF NOT EXISTS idx_benchmark_runs_group
            ON benchmark_runs(run_group_id, started_at);
        CREATE INDEX IF NOT EXISTS idx_benchmark_runs_lookup
            ON benchmark_runs(benchmark_id, provider, model, agent, started_at);
        CREATE INDEX IF NOT EXISTS idx_benchmark_run_trajectories_run
            ON benchmark_run_trajectories(run_id);
        """
    )
    conn.commit()
    _ensure_column(
        conn,
        "benchmark_run_trajectories",
        "total_tokens",
        "INTEGER NOT NULL DEFAULT 0",
    )
    conn.commit()
    ensure_comparison_id_column(conn)
    ensure_metrics_columns(conn)


def create_run_group(
    conn: sqlite3.Connection,
    *,
    run_group_id: str,
    created_at: str,
    request: dict[str, Any],
    benchmarks: list[str],
    repo_meta: dict[str, Any],
) -> None:
    conn.execute(
        """
        INSERT INTO run_groups (
            run_group_id,
            created_at,
            request_json,
            benchmarks_json,
            repo_meta_json,
            created_by
        ) VALUES (?, ?, ?, ?, ?, ?)
        """,
        (
            run_group_id,
            created_at,
            _json_dumps(request),
            _json_dumps(benchmarks),
            _json_dumps(repo_meta),
            "benchmarks.orchestrator",
        ),
    )
    conn.commit()


def finish_run_group(conn: sqlite3.Connection, *, run_group_id: str, finished_at: str) -> None:
    conn.execute(
        "UPDATE run_groups SET finished_at = ? WHERE run_group_id = ?",
        (finished_at, run_group_id),
    )
    conn.commit()


def get_latest_run_for_signature(conn: sqlite3.Connection, signature: str) -> ExistingRun | None:
    row = conn.execute(
        """
        SELECT run_id, signature, status, attempt
        FROM benchmark_runs
        WHERE signature = ?
        ORDER BY attempt DESC, started_at DESC
        LIMIT 1
        """,
        (signature,),
    ).fetchone()
    if row is None:
        return None
    return ExistingRun(
        run_id=str(row["run_id"]),
        signature=str(row["signature"]),
        status=str(row["status"]),
        attempt=int(row["attempt"]),
    )


def get_latest_succeeded_run_for_signature(conn: sqlite3.Connection, signature: str) -> ExistingRun | None:
    row = conn.execute(
        """
        SELECT run_id, signature, status, attempt
        FROM benchmark_runs
        WHERE signature = ? AND status = 'succeeded'
        ORDER BY attempt DESC, started_at DESC
        LIMIT 1
        """,
        (signature,),
    ).fetchone()
    if row is None:
        return None
    return ExistingRun(
        run_id=str(row["run_id"]),
        signature=str(row["signature"]),
        status=str(row["status"]),
        attempt=int(row["attempt"]),
    )


def next_attempt_for_signature(conn: sqlite3.Connection, signature: str) -> int:
    row = conn.execute(
        "SELECT MAX(attempt) AS max_attempt FROM benchmark_runs WHERE signature = ?",
        (signature,),
    ).fetchone()
    if row is None or row["max_attempt"] is None:
        return 1
    return int(row["max_attempt"]) + 1


def insert_run_start(
    conn: sqlite3.Connection,
    *,
    run_id: str,
    run_group_id: str,
    benchmark_id: str,
    benchmark_directory: str,
    signature: str,
    attempt: int,
    agent: str,
    provider: str,
    model: str,
    extra_config: dict[str, Any],
    started_at: str,
    command: list[str],
    cwd: str,
    stdout_path: str,
    stderr_path: str,
    benchmark_version: str | None,
    benchmarks_commit: str | None,
    eliza_commit: str | None,
    eliza_version: str | None,
) -> None:
    conn.execute(
        """
        INSERT INTO benchmark_runs (
            run_id,
            run_group_id,
            signature,
            benchmark_id,
            benchmark_directory,
            status,
            attempt,
            agent,
            provider,
            model,
            extra_config_json,
            started_at,
            command_json,
            cwd,
            stdout_path,
            stderr_path,
            result_json_path,
            score,
            unit,
            higher_is_better,
            metrics_json,
            artifacts_json,
            error,
            high_score_label,
            high_score_value,
            delta_to_high_score,
            benchmark_version,
            benchmarks_commit,
            eliza_commit,
            eliza_version,
            created_at
        ) VALUES (?, ?, ?, ?, ?, 'running', ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, NULL, NULL, NULL, NULL, '{}', '[]', NULL, NULL, NULL, NULL, ?, ?, ?, ?, ?)
        """,
        (
            run_id,
            run_group_id,
            signature,
            benchmark_id,
            benchmark_directory,
            attempt,
            agent,
            provider,
            model,
            _json_dumps(extra_config),
            started_at,
            _json_dumps(command),
            cwd,
            stdout_path,
            stderr_path,
            benchmark_version,
            benchmarks_commit,
            eliza_commit,
            eliza_version,
            started_at,
        ),
    )
    conn.commit()


def update_run_result(
    conn: sqlite3.Connection,
    *,
    run_id: str,
    status: str,
    ended_at: str,
    duration_seconds: float | None,
    score: float | None,
    unit: str | None,
    higher_is_better: bool | None,
    metrics: dict[str, Any],
    result_json_path: str | None,
    artifacts: list[str],
    error: str | None,
    high_score_label: str | None,
    high_score_value: float | None,
    delta_to_high_score: float | None,
    trajectory_summary: dict[str, Any] | None = None,
    token_metrics: dict[str, Any] | None = None,
    cache_metrics: dict[str, Any] | None = None,
    performance_metrics: dict[str, Any] | None = None,
) -> None:
    hib: int | None
    if higher_is_better is None:
        hib = None
    else:
        hib = 1 if higher_is_better else 0

    conn.execute(
        """
        UPDATE benchmark_runs
        SET
            status = ?,
            ended_at = ?,
            duration_seconds = ?,
            score = ?,
            unit = ?,
            higher_is_better = ?,
            metrics_json = ?,
            result_json_path = ?,
            artifacts_json = ?,
            error = ?,
            high_score_label = ?,
            high_score_value = ?,
            delta_to_high_score = ?,
            trajectory_summary_json = ?,
            token_metrics_json = ?,
            cache_metrics_json = ?,
            performance_metrics_json = ?,
            trajectory_count = ?,
            llm_call_count = ?,
            total_prompt_tokens = ?,
            total_completion_tokens = ?,
            total_cache_read_input_tokens = ?,
            total_cache_creation_input_tokens = ?,
            mean_latency_ms = ?,
            p95_latency_ms = ?,
            throughput_per_second = ?
        WHERE run_id = ?
        """,
        (
            status,
            ended_at,
            duration_seconds,
            score,
            unit,
            hib,
            _json_dumps(metrics),
            result_json_path,
            _json_dumps(artifacts),
            error,
            high_score_label,
            high_score_value,
            delta_to_high_score,
            _json_dumps(trajectory_summary or {}),
            _json_dumps(token_metrics or {}),
            _json_dumps(cache_metrics or {}),
            _json_dumps(performance_metrics or {}),
            _int_or_none((trajectory_summary or {}).get("files")),
            _int_or_none((token_metrics or {}).get("llm_call_count")),
            _int_or_none((token_metrics or {}).get("prompt_tokens")),
            _int_or_none((token_metrics or {}).get("completion_tokens")),
            _int_or_none((cache_metrics or {}).get("cache_read_input_tokens")),
            _int_or_none((cache_metrics or {}).get("cache_creation_input_tokens")),
            _float_or_none((performance_metrics or {}).get("mean_latency_ms")),
            _float_or_none((performance_metrics or {}).get("p95_latency_ms")),
            _float_or_none((performance_metrics or {}).get("throughput_per_second")),
            run_id,
        ),
    )
    conn.commit()


def _int_or_none(value: Any) -> int | None:
    if isinstance(value, bool):
        return None
    if isinstance(value, (int, float)):
        return int(value)
    return None


def _float_or_none(value: Any) -> float | None:
    if isinstance(value, bool):
        return None
    if isinstance(value, (int, float)):
        return float(value)
    return None


def replace_run_trajectories(
    conn: sqlite3.Connection,
    *,
    run_id: str,
    trajectories: list[dict[str, Any]],
) -> None:
    conn.execute("DELETE FROM benchmark_run_trajectories WHERE run_id = ?", (run_id,))
    if trajectories:
        conn.executemany(
            """
            INSERT INTO benchmark_run_trajectories (
                run_id,
                trajectory_file,
                turn_index,
                prompt_tokens,
                completion_tokens,
                total_tokens,
                cached_tokens,
                cache_creation_tokens,
                latency_ms,
                prompt_chars
            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """,
            [
                (
                    run_id,
                    str(row.get("trajectory_file") or row.get("file") or ""),
                    int(row.get("turn_index") or row.get("index") or 0),
                    int(row.get("prompt_tokens") or 0),
                    int(row.get("completion_tokens") or 0),
                    int(
                        row.get("total_tokens")
                        or (
                            int(row.get("prompt_tokens") or 0)
                            + int(row.get("completion_tokens") or 0)
                        )
                    ),
                    int(row.get("cached_tokens") or 0),
                    int(row.get("cache_creation_tokens") or 0),
                    _float_or_none(row.get("latency_ms")),
                    int(row.get("prompt_chars") or 0),
                )
                for row in trajectories
            ],
        )
    conn.commit()


def tag_run_with_comparison(
    conn: sqlite3.Connection,
    *,
    run_id: str,
    comparison_id: str,
) -> None:
    conn.execute(
        "UPDATE benchmark_runs SET comparison_id = ? WHERE run_id = ?",
        (comparison_id, run_id),
    )
    conn.commit()


def list_runs_for_comparison(
    conn: sqlite3.Connection,
    *,
    comparison_id: str,
) -> list[dict[str, Any]]:
    rows = conn.execute(
        """
        SELECT run_id, run_group_id, benchmark_id, status, agent, provider, model,
               score, unit, higher_is_better, metrics_json, started_at, ended_at,
               duration_seconds, error, trajectory_summary_json,
               token_metrics_json, cache_metrics_json, performance_metrics_json
        FROM benchmark_runs
        WHERE comparison_id = ?
        ORDER BY started_at ASC, run_id ASC
        """,
        (comparison_id,),
    ).fetchall()
    out: list[dict[str, Any]] = []
    for row in rows:
        record = dict(row)
        for json_key in (
            "metrics_json",
            "trajectory_summary_json",
            "token_metrics_json",
            "cache_metrics_json",
            "performance_metrics_json",
        ):
            raw = record.pop(json_key, None)
            out_key = json_key.removesuffix("_json")
            if isinstance(raw, str):
                try:
                    record[out_key] = json.loads(raw)
                except json.JSONDecodeError:
                    record[out_key] = {}
            else:
                record[out_key] = {}
        hib = record.get("higher_is_better")
        record["higher_is_better"] = None if hib is None else bool(hib)
        out.append(record)
    return out


def list_runs(
    conn: sqlite3.Connection,
    *,
    limit: int | None = 5000,
) -> list[dict[str, Any]]:
    query = """
        SELECT
            run_id,
            run_group_id,
            signature,
            benchmark_id,
            benchmark_directory,
            status,
            attempt,
            agent,
            provider,
            model,
            extra_config_json,
            started_at,
            ended_at,
            duration_seconds,
            command_json,
            cwd,
            stdout_path,
            stderr_path,
            result_json_path,
            score,
            unit,
            higher_is_better,
            metrics_json,
            trajectory_summary_json,
            token_metrics_json,
            cache_metrics_json,
            performance_metrics_json,
            trajectory_count,
            llm_call_count,
            total_prompt_tokens,
            total_completion_tokens,
            total_cache_read_input_tokens,
            total_cache_creation_input_tokens,
            mean_latency_ms,
            p95_latency_ms,
            throughput_per_second,
            artifacts_json,
            error,
            high_score_label,
            high_score_value,
            delta_to_high_score,
            benchmark_version,
            benchmarks_commit,
            eliza_commit,
            eliza_version
        FROM benchmark_runs
        ORDER BY started_at DESC, run_id DESC
        """
    params: tuple[Any, ...] = ()
    if limit is not None:
        query += " LIMIT ?"
        params = (limit,)
    rows = conn.execute(query, params).fetchall()

    out: list[dict[str, Any]] = []
    for row in rows:
        record = dict(row)
        for key in (
            "extra_config_json",
            "command_json",
            "metrics_json",
            "trajectory_summary_json",
            "token_metrics_json",
            "cache_metrics_json",
            "performance_metrics_json",
            "artifacts_json",
        ):
            raw = record.get(key)
            if isinstance(raw, str):
                try:
                    record[key.removesuffix("_json") if key.endswith("_json") else key] = json.loads(raw)
                except json.JSONDecodeError:
                    record[key.removesuffix("_json") if key.endswith("_json") else key] = raw
            if key in record:
                del record[key]
        hib = record.get("higher_is_better")
        if hib is None:
            record["higher_is_better"] = None
        else:
            record["higher_is_better"] = bool(hib)
        out.append(record)
    return out


def repair_nonzero_returncode_statuses(conn: sqlite3.Connection) -> int:
    """Mark legacy rows with nonzero process exits as failed.

    Older runner versions recorded a row as ``succeeded`` whenever a result
    JSON existed, even if the subprocess exited nonzero. Keep the artifacts and
    score, but make the status honest so latest snapshots/viewer summaries do
    not treat task-level process failures as clean wins.
    """

    rows = conn.execute(
        """
        SELECT run_id, metrics_json, error
        FROM benchmark_runs
        WHERE status = 'succeeded'
        """
    ).fetchall()
    repaired = 0
    for row in rows:
        raw = row["metrics_json"]
        if not isinstance(raw, str):
            continue
        try:
            metrics = json.loads(raw)
        except json.JSONDecodeError:
            continue
        if not isinstance(metrics, dict):
            continue
        return_code = metrics.get("return_code", metrics.get("returncode"))
        if return_code in (None, 0):
            return_code = metrics.get("nonzero_return_code_with_result")
        if return_code in (None, 0):
            return_code = metrics.get("exit_code")
        if not isinstance(return_code, (int, float)) or int(return_code) == 0:
            continue
        error = row["error"] or (
            "Command produced a result JSON but exited with "
            f"return code {int(return_code)}"
        )
        conn.execute(
            """
            UPDATE benchmark_runs
            SET status = 'failed', error = ?
            WHERE run_id = ?
            """,
            (error, row["run_id"]),
        )
        repaired += 1
    if repaired:
        conn.commit()
    return repaired


def repair_nonpublishable_success_statuses(conn: sqlite3.Connection) -> int:
    """Mark legacy "successful" rows that never executed as failed.

    This is intentionally narrow. Some deterministic/calibration benchmarks do
    not emit LLM telemetry, but rows with zero evaluated samples in public
    benchmark result shapes are empty artifacts rather than scored attempts.
    """

    rows = conn.execute(
        """
        SELECT run_id, benchmark_id, metrics_json, token_metrics_json, trajectory_summary_json, error
        FROM benchmark_runs
        WHERE status = 'succeeded'
        """
    ).fetchall()
    repaired = 0
    for row in rows:
        raw_metrics = row["metrics_json"]
        if not isinstance(raw_metrics, str):
            continue
        try:
            metrics = json.loads(raw_metrics)
        except json.JSONDecodeError:
            continue
        if not isinstance(metrics, dict):
            continue
        reason = _nonpublishable_success_reason(str(row["benchmark_id"]), metrics)
        if reason is None:
            continue
        error = row["error"] or reason
        conn.execute(
            """
            UPDATE benchmark_runs
            SET status = 'failed', error = ?
            WHERE run_id = ?
            """,
            (error, row["run_id"]),
        )
        repaired += 1
    if repaired:
        conn.commit()
    return repaired


def _metric_number(metrics: dict[str, Any], key: str) -> float | None:
    value = metrics.get(key)
    if isinstance(value, bool):
        return None
    if isinstance(value, (int, float)):
        return float(value)
    return None


def _nonpublishable_success_reason(
    benchmark_id: str,
    metrics: dict[str, Any],
) -> str | None:
    if benchmark_id == "solana":
        messages = metrics.get("messages")
        cumulative_rewards = metrics.get("cumulative_rewards")
        final_programs = metrics.get("final_programs")
        if final_programs is None and isinstance(metrics.get("programs_discovered"), dict):
            final_programs = len(metrics["programs_discovered"])
        empty_rollout = (
            isinstance(messages, list)
            and len(messages) == 0
            and isinstance(cumulative_rewards, list)
            and len(cumulative_rewards) == 0
            and (final_programs in (None, 0, 0.0))
        )
        return "Solana benchmark produced an empty rollout artifact" if empty_rollout else None

    positive_sample_keys: dict[str, tuple[str, ...]] = {
        "abliteration-robustness": ("n",),
        "humaneval": ("n",),
        "lifeops_bench": ("seeds",),
        "mmlu": ("n",),
    }
    keys = positive_sample_keys.get(benchmark_id)
    if keys is None:
        return None
    for key in keys:
        value = _metric_number(metrics, key)
        if value is None or value <= 0:
            return f"{benchmark_id} benchmark produced a zero-sample success artifact ({key}={value!r})"
    if benchmark_id == "lifeops_bench":
        scenario_count = _metric_number(metrics, "scenario_count")
        if scenario_count is not None and scenario_count <= 0:
            return (
                "lifeops_bench benchmark produced a zero-scenario success "
                f"artifact (scenario_count={scenario_count!r})"
            )
    return None


def list_run_groups(conn: sqlite3.Connection, *, limit: int = 2000) -> list[dict[str, Any]]:
    rows = conn.execute(
        """
        SELECT run_group_id, created_at, finished_at, request_json, benchmarks_json, repo_meta_json
        FROM run_groups
        ORDER BY created_at DESC, run_group_id DESC
        LIMIT ?
        """,
        (limit,),
    ).fetchall()
    out: list[dict[str, Any]] = []
    for row in rows:
        record = dict(row)
        for key in ("request_json", "benchmarks_json", "repo_meta_json"):
            raw = record.get(key)
            if isinstance(raw, str):
                try:
                    record[key.removesuffix("_json")] = json.loads(raw)
                except json.JSONDecodeError:
                    record[key.removesuffix("_json")] = raw
            if key in record:
                del record[key]
        out.append(record)
    return out


def summarize_latest_scores(conn: sqlite3.Connection) -> list[dict[str, Any]]:
    rows = conn.execute(
        """
        WITH latest AS (
            SELECT
                benchmark_id,
                agent,
                MAX(started_at) AS max_started_at
            FROM benchmark_runs
            WHERE status = 'succeeded'
              AND score IS NOT NULL
            GROUP BY benchmark_id, agent
        )
        SELECT
            r.benchmark_id,
            r.run_id,
            r.run_group_id,
            r.started_at,
            r.score,
            r.unit,
            r.agent,
            r.provider,
            r.model,
            r.high_score_label,
            r.high_score_value,
            r.delta_to_high_score,
            r.token_metrics_json,
            r.llm_call_count,
            r.total_prompt_tokens,
            r.total_completion_tokens,
            r.total_cache_read_input_tokens,
            r.total_cache_creation_input_tokens
        FROM benchmark_runs r
        JOIN latest l
          ON r.benchmark_id = l.benchmark_id
         AND r.agent = l.agent
         AND r.started_at = l.max_started_at
        WHERE r.status = 'succeeded'
          AND r.score IS NOT NULL
        ORDER BY r.benchmark_id ASC, r.agent ASC
        """
    ).fetchall()
    out: list[dict[str, Any]] = []
    for row in rows:
        record = dict(row)
        raw_token_metrics = record.pop("token_metrics_json", None)
        token_metrics: dict[str, Any] = {}
        if isinstance(raw_token_metrics, str):
            try:
                decoded = json.loads(raw_token_metrics)
            except json.JSONDecodeError:
                decoded = {}
            if isinstance(decoded, dict):
                token_metrics = decoded

        prompt_tokens = _int_or_none(
            token_metrics.get("input_tokens", token_metrics.get("prompt_tokens"))
        )
        if prompt_tokens is None:
            prompt_tokens = _int_or_none(record.get("total_prompt_tokens"))
        completion_tokens = _int_or_none(
            token_metrics.get("output_tokens", token_metrics.get("completion_tokens"))
        )
        if completion_tokens is None:
            completion_tokens = _int_or_none(record.get("total_completion_tokens"))
        total_tokens = _int_or_none(token_metrics.get("total_tokens"))
        if total_tokens is None and prompt_tokens is not None and completion_tokens is not None:
            total_tokens = prompt_tokens + completion_tokens
        cached_tokens = _int_or_none(
            token_metrics.get("cached_tokens", token_metrics.get("cache_read_input_tokens"))
        )
        if cached_tokens is None:
            cached_tokens = _int_or_none(record.get("total_cache_read_input_tokens"))
        calls = _int_or_none(token_metrics.get("llm_call_count", token_metrics.get("call_count")))
        if calls is None:
            calls = _int_or_none(record.get("llm_call_count"))

        record["input_tokens"] = prompt_tokens if prompt_tokens is not None else 0
        record["output_tokens"] = completion_tokens if completion_tokens is not None else 0
        record["total_tokens"] = total_tokens if total_tokens is not None else 0
        record["cached_tokens"] = cached_tokens if cached_tokens is not None else 0
        record["cache_creation_input_tokens"] = _int_or_none(
            token_metrics.get("cache_creation_input_tokens")
        )
        if record["cache_creation_input_tokens"] is None:
            record["cache_creation_input_tokens"] = _int_or_none(
                record.get("total_cache_creation_input_tokens")
            ) or 0
        record["llm_call_count"] = calls if calls is not None else 0
        record["call_count"] = _int_or_none(token_metrics.get("call_count"))
        if record["call_count"] is None:
            record["call_count"] = record["llm_call_count"]
        record["token_metrics"] = {
            "input_tokens": record["input_tokens"],
            "output_tokens": record["output_tokens"],
            "total_tokens": record["total_tokens"],
            "cached_tokens": record["cached_tokens"],
            "cache_creation_input_tokens": record["cache_creation_input_tokens"],
            "llm_call_count": record["llm_call_count"],
            "call_count": record["call_count"],
        }
        out.append(record)
    return out


def recover_stale_running_runs(
    conn: sqlite3.Connection,
    *,
    stale_before: str,
    ended_at: str,
) -> list[str]:
    rows = conn.execute(
        """
        SELECT run_id, run_group_id, started_at
        FROM benchmark_runs
        WHERE status = 'running'
          AND started_at < ?
        ORDER BY started_at ASC
        """,
        (stale_before,),
    ).fetchall()
    if not rows:
        return []

    recovered_ids: list[str] = []
    touched_groups: set[str] = set()
    metrics_json = _json_dumps({"reason": "orchestrator_interrupted"})

    ended_dt = datetime.fromisoformat(ended_at)
    if ended_dt.tzinfo is None:
        ended_dt = ended_dt.replace(tzinfo=UTC)

    for row in rows:
        run_id = str(row["run_id"])
        run_group_id = str(row["run_group_id"])
        started_raw = str(row["started_at"])

        duration_seconds: float | None = None
        try:
            started_dt = datetime.fromisoformat(started_raw)
            if started_dt.tzinfo is None:
                started_dt = started_dt.replace(tzinfo=UTC)
            duration_seconds = max(0.0, (ended_dt - started_dt).total_seconds())
        except ValueError:
            duration_seconds = None

        conn.execute(
            """
            UPDATE benchmark_runs
            SET
                status = 'failed',
                ended_at = ?,
                duration_seconds = ?,
                metrics_json = ?,
                error = ?,
                result_json_path = NULL
            WHERE run_id = ?
            """,
            (
                ended_at,
                duration_seconds,
                metrics_json,
                "Recovered stale running run from interrupted orchestrator process",
                run_id,
            ),
        )
        recovered_ids.append(run_id)
        touched_groups.add(run_group_id)

    for run_group_id in sorted(touched_groups):
        still_running = conn.execute(
            """
            SELECT 1
            FROM benchmark_runs
            WHERE run_group_id = ? AND status = 'running'
            LIMIT 1
            """,
            (run_group_id,),
        ).fetchone()
        if still_running is None:
            conn.execute(
                """
                UPDATE run_groups
                SET finished_at = COALESCE(finished_at, ?)
                WHERE run_group_id = ?
                """,
                (ended_at, run_group_id),
            )

    conn.commit()
    return recovered_ids