band-aider-pub/band_aider/reporting/trace_summary.py

# SPDX-FileCopyrightText: 2026 Art Sea LLC
#
# SPDX-License-Identifier: AGPL-3.0-only

"""Summarize a Band-Aider JSONL execution trace into a readable digest.

A trace file answers "what happened during this run?" with total
fidelity, but fidelity is not digestibility.  This module folds a trace
produced by band_aider.ui.trace_publisher into the handful of facts a
reviewer actually checks first: how long the run lasted, which loops it
moved through, where escalations and human interventions occurred, how
many tokens it spent, and how it ended.

Usage:
    python -m band_aider.reporting.trace_summary path/to/trace.jsonl
    python -m band_aider.reporting.trace_summary --json trace.jsonl
"""

from __future__ import annotations

import argparse
from collections import Counter
from datetime import datetime
import json
from pathlib import Path
from typing import Any, Iterable, TextIO

from band_aider.ui.trace_publisher import TRACE_START_EVENT

# Loop names appear in LoopStatusDTO.current_loop; the sequence of
# distinct values is the closest thing the trace has to a state machine.
# In real runs loop status travels nested inside SupremeStatusDTO.loop
# (see band_aider.ui.status.publish_loop_status), while escalations,
# intervention requests, and run completion are published top-level.
_LOOP_TRANSITION_EVENT = "LoopStatusDTO"
_SUPREME_STATUS_EVENT = "SupremeStatusDTO"
_ESCALATION_EVENT = "ModelEscalationDTO"
_INTERVENTION_REQUEST_EVENT = "HumanInterventionRequestDTO"
_INTERVENTION_RESPONSE_EVENT = "HumanInterventionResponseDTO"
_RUN_COMPLETION_EVENT = "RunCompletionDTO"
_TOKEN_BUDGET_EVENT = "TokenBudgetStatsDTO"

_RUN_COMPLETION_FIELDS = (
    "final_summary",
    "report_path",
    "total_tokens_spent",
    "total_seconds",
    "cot_count",
    "files_modified",
    "files_added",
    "llm_call_count",
)

_TOKEN_BUDGET_FIELDS = (
    "max_prompt_tokens",
    "total_tokens",
    "baseline_tokens",
    "escalated_tokens",
    "total_cache_hit_tokens",
)


def load_records(stream: TextIO) -> list[dict[str, Any]]:
    """Parse one JSONL trace, keeping every record that still parses.

    A killed process can leave a truncated final line behind; the trace
    remains useful without it, so unparseable lines are skipped rather
    than failing the whole summary.
    """

    records: list[dict[str, Any]] = []
    for line in stream:
        line = line.strip()
        if not line:
            continue
        try:
            record = json.loads(line)
        except json.JSONDecodeError:
            continue
        if isinstance(record, dict):
            records.append(record)
    return records


def _parse_ts(record: dict[str, Any]) -> datetime | None:
    """Return one record timestamp, or None when absent or malformed."""

    raw = record.get("ts")
    if not isinstance(raw, str):
        return None
    try:
        return datetime.fromisoformat(raw)
    except ValueError:
        return None


def _loop_sequence(records: Iterable[dict[str, Any]]) -> list[str]:
    """Collapse consecutive loop status events into distinct transitions."""

    sequence: list[str] = []
    for record in records:
        current: Any = None
        data = record.get("data")
        if not isinstance(data, dict):
            continue
        if record.get("event") == _LOOP_TRANSITION_EVENT:
            current = data.get("current_loop")
        elif record.get("event") == _SUPREME_STATUS_EVENT:
            loop = data.get("loop")
            if isinstance(loop, dict):
                current = loop.get("current_loop")
        if isinstance(current, str) and current and (
            not sequence or sequence[-1] != current
        ):
            sequence.append(current)
    return sequence


def _final_token_budget(events: list[dict[str, Any]]) -> dict[str, Any] | None:
    """Return the last token budget, top-level or nested in supreme status."""

    for record in reversed(events):
        data = record.get("data")
        if not isinstance(data, dict):
            continue
        if record.get("event") == _TOKEN_BUDGET_EVENT:
            return data
        if record.get("event") == _SUPREME_STATUS_EVENT:
            budget = data.get("token_budget")
            if isinstance(budget, dict):
                return budget
    return None


def summarize_records(records: list[dict[str, Any]]) -> dict[str, Any]:
    """Fold parsed trace records into one summary dictionary."""

    events = [r for r in records if r.get("event") != TRACE_START_EVENT]
    header = next(
        (r for r in records if r.get("event") == TRACE_START_EVENT),
        None,
    )
    header_data = header.get("data") if isinstance(header, dict) else None

    timestamps = [ts for ts in (_parse_ts(r) for r in events) if ts is not None]
    event_counts = Counter(
        str(r.get("event")) for r in events if r.get("event") is not None
    )

    escalations = [r for r in events if r.get("event") == _ESCALATION_EVENT]
    escalation_triggers = Counter(
        str(r.get("data", {}).get("trigger_name"))
        for r in escalations
        if isinstance(r.get("data"), dict)
    )

    completion = next(
        (r for r in reversed(events) if r.get("event") == _RUN_COMPLETION_EVENT),
        None,
    )
    completion_data = (
        completion.get("data")
        if isinstance(completion, dict) and isinstance(completion.get("data"), dict)
        else None
    )
    budget_data = _final_token_budget(events)

    duration_seconds = None
    if len(timestamps) >= 2:
        duration_seconds = round((max(timestamps) - min(timestamps)).total_seconds(), 3)

    loop_sequence = _loop_sequence(events)
    return {
        "format": (
            header_data.get("format") if isinstance(header_data, dict) else None
        ),
        "origin": (
            header_data.get("origin") if isinstance(header_data, dict) else None
        ),
        "total_records": len(records),
        "total_events": len(events),
        "first_event_at": min(timestamps).isoformat() if timestamps else None,
        "last_event_at": max(timestamps).isoformat() if timestamps else None,
        "duration_seconds": duration_seconds,
        "event_counts": dict(event_counts.most_common()),
        "loop_sequence": loop_sequence,
        "loop_transitions": max(0, len(loop_sequence) - 1),
        "escalation_count": len(escalations),
        "escalation_triggers": dict(escalation_triggers.most_common()),
        "intervention_requests": sum(
            1 for r in events if r.get("event") == _INTERVENTION_REQUEST_EVENT
        ),
        "intervention_responses": sum(
            1 for r in events if r.get("event") == _INTERVENTION_RESPONSE_EVENT
        ),
        "run_completion": (
            {k: completion_data.get(k) for k in _RUN_COMPLETION_FIELDS}
            if completion_data
            else None
        ),
        "final_token_budget": (
            {k: budget_data.get(k) for k in _TOKEN_BUDGET_FIELDS}
            if budget_data
            else None
        ),
    }


def summarize_stream(stream: TextIO) -> dict[str, Any]:
    """Parse and summarize one open trace stream."""

    return summarize_records(load_records(stream))


def summarize_path(trace_path: Path) -> dict[str, Any]:
    """Open, parse, and summarize one trace file."""

    with trace_path.open("r", encoding="utf-8") as stream:
        return summarize_stream(stream)


def render_text(summary: dict[str, Any]) -> str:
    """Render one summary dictionary as a compact human-readable digest."""

    lines = ["Band-Aider execution trace summary"]
    lines.append(f"  trace format: {summary['format'] or 'unknown'}")
    if summary["origin"]:
        lines.append(f"  trace origin: {summary['origin']}")
    lines.append(f"  records: {summary['total_records']} "
                 f"({summary['total_events']} events)")
    lines.append(f"  first event: {summary['first_event_at'] or 'n/a'}")
    lines.append(f"  last event:  {summary['last_event_at'] or 'n/a'}")
    lines.append(f"  duration:    {summary['duration_seconds'] or 0}s")

    lines.append("")
    lines.append("  event counts:")
    for name, count in summary["event_counts"].items():
        lines.append(f"    {name}: {count}")
    if not summary["event_counts"]:
        lines.append("    (none)")

    lines.append("")
    lines.append(f"  loop sequence ({summary['loop_transitions']} transitions):")
    if summary["loop_sequence"]:
        lines.append("    " + " -> ".join(summary["loop_sequence"]))
    else:
        lines.append("    (no loop transitions recorded)")

    lines.append("")
    lines.append(f"  escalations: {summary['escalation_count']}")
    for trigger, count in summary["escalation_triggers"].items():
        lines.append(f"    {trigger}: {count}")
    lines.append(
        f"  interventions: {summary['intervention_requests']} requests, "
        f"{summary['intervention_responses']} responses"
    )

    completion = summary["run_completion"]
    lines.append("")
    if completion:
        lines.append("  run completion:")
        for key, value in completion.items():
            lines.append(f"    {key}: {value}")
    else:
        lines.append("  run completion: (run did not report completion)")

    budget = summary["final_token_budget"]
    if budget:
        lines.append("")
        lines.append("  final token budget:")
        for key, value in budget.items():
            lines.append(f"    {key}: {value}")

    return "\n".join(lines)


def main(argv: list[str] | None = None) -> int:
    """Summarize one trace file as text (default) or JSON (--json)."""

    parser = argparse.ArgumentParser(
        prog="python -m band_aider.reporting.trace_summary",
        description="Summarize a Band-Aider JSONL execution trace.",
    )
    parser.add_argument("trace_file", type=Path, help="Path to the JSONL trace file.")
    parser.add_argument(
        "--json",
        action="store_true",
        help="Emit the raw summary dictionary as JSON instead of text.",
    )
    args = parser.parse_args(argv)

    summary = summarize_path(args.trace_file)
    if args.json:
        print(json.dumps(summary, indent=2, ensure_ascii=False))
    else:
        print(render_text(summary))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())