band-aider-pub/tests/band_aider/test_trace_summary.py

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

"""In-memory trace summary aggregation contracts."""

from __future__ import annotations

import io
import time

from band_aider.reporting import trace_summary
from band_aider.ui.dto.human_intervention import (
    HumanInterventionRequestDTO,
    HumanInterventionResponseDTO,
)
from band_aider.ui.dto.loop_status import LoopStatusDTO
from band_aider.ui.dto.model_escalation import ModelEscalationDTO
from band_aider.ui.dto.run_completion import RunCompletionDTO
from band_aider.ui.dto.supreme_status import SupremeStatusDTO
from band_aider.ui.dto.token_budget_stats import TokenBudgetStatsDTO
from band_aider.ui.trace_publisher import TRACE_FORMAT, TracePublisher


class _NullPublisher:
    """Discard events; only the trace stream matters in these contracts."""

    def publish(self, event: object) -> None:
        del event


def _loop_status(current_loop: str) -> LoopStatusDTO:
    """Build one loop transition event for the named loop."""

    return LoopStatusDTO(
        previous_loop=None,
        current_loop=current_loop,
        next_loop=None,
        current_item=None,
        work_focus=None,
    )


def _escalation(trigger_name: str) -> ModelEscalationDTO:
    """Build one model escalation event for the named trigger."""

    return ModelEscalationDTO(
        trigger_name=trigger_name,
        occurred_at=time.monotonic(),
        prompt_preview="prompt",
        reason="reason",
        request_md5="0" * 32,
        local_model="local",
        active_model="active",
        config_index=0,
        config_count=1,
        state="escalated",
    )


def _synthetic_trace() -> io.StringIO:
    """Record one representative run into an in-memory trace stream.

    Loop and token-budget status travel nested inside SupremeStatusDTO in
    real runs (see band_aider.ui.status), so most events here take that
    shape; one bare LoopStatusDTO proves the flat form still parses.
    """

    stream = io.StringIO()
    publisher = TracePublisher(_NullPublisher(), stream, origin="memory")
    publisher.publish(SupremeStatusDTO(loop=_loop_status("BootstrapLoop")))
    publisher.publish(SupremeStatusDTO(loop=_loop_status("BootstrapLoop")))
    publisher.publish(SupremeStatusDTO(loop=_loop_status("StabilizationLoop")))
    publisher.publish(_loop_status("AlphaLoop"))
    publisher.publish(_escalation("FenceCrashingClient"))
    publisher.publish(_escalation("FenceCrashingClient"))
    publisher.publish(HumanInterventionRequestDTO(
        request_id="r1",
        check_name="StabilizationDecision",
        original_prompt="prompt",
        repeated_response_tail="tail",
        request_md5="0" * 32,
        response_md5="1" * 32,
        local_model="local",
        escalated_model="active",
    ))
    publisher.publish(HumanInterventionResponseDTO(
        request_id="r1", guidance="try smaller edits", action="guide",
    ))
    publisher.publish(SupremeStatusDTO(token_budget=TokenBudgetStatsDTO(
        max_prompt_tokens=1000,
        last_ai_token_used=10,
        total_tokens=123,
        bucket_window_seconds=300,
        buckets=(),
        baseline_tokens=100,
        escalated_tokens=23,
    )))
    publisher.publish(RunCompletionDTO(
        final_summary="all done",
        report_path="band-aider reports/run.md",
        total_tokens_spent=123,
        largest_query_tokens=50,
        total_seconds=9.5,
        average_seconds_per_cot=4.75,
        cot_count=2,
        files_modified=3,
        files_added=1,
        llm_call_count=8,
    ))
    stream.seek(0)
    return stream


def test_trace_summary_aggregates_events_loops_escalations_and_completion() -> None:
    """One synthetic run folds into counts, transitions, and terminal facts."""

    stream = _synthetic_trace()
    summary = trace_summary.summarize_stream(stream)

    assert summary["format"] == TRACE_FORMAT
    assert summary["origin"] == "memory"
    assert summary["total_events"] == 10
    assert summary["duration_seconds"] >= 0
    assert summary["event_counts"]["SupremeStatusDTO"] == 4
    assert summary["event_counts"]["LoopStatusDTO"] == 1
    assert summary["loop_sequence"] == [
        "BootstrapLoop",
        "StabilizationLoop",
        "AlphaLoop",
    ]
    assert summary["loop_transitions"] == 2
    assert summary["escalation_count"] == 2
    assert summary["escalation_triggers"] == {"FenceCrashingClient": 2}
    assert summary["intervention_requests"] == 1
    assert summary["intervention_responses"] == 1
    assert summary["run_completion"]["final_summary"] == "all done"
    assert summary["run_completion"]["total_tokens_spent"] == 123
    assert summary["run_completion"]["llm_call_count"] == 8
    assert summary["final_token_budget"]["total_tokens"] == 123
    assert summary["final_token_budget"]["escalated_tokens"] == 23

    text = trace_summary.render_text(summary)
    assert "BootstrapLoop -> StabilizationLoop -> AlphaLoop" in text
    assert "escalations: 2" in text
    assert "run completion:" in text


def test_trace_summary_tolerates_empty_and_truncated_traces() -> None:
    """Missing and half-written traces still yield a coherent digest."""

    empty = trace_summary.summarize_stream(io.StringIO(""))
    assert empty["total_events"] == 0
    assert empty["loop_sequence"] == []
    assert empty["run_completion"] is None
    assert "no loop transitions" in trace_summary.render_text(empty)

    truncated = io.StringIO(
        '{"seq": 0, "ts": "2026-10-28T05:40:00.000+00:00", '
        '"event": "trace_start", "data": {"format": "band-aider-trace/1"}}\n'
        '{"seq": 1, "ts": "2026-10-28T05:40:01.000+00:00", '
        '"event": "LoopStatusDTO", "data": {"current_loop": "AlphaLoop"}}\n'
        '{"seq": 2, "ts": "2026-10-28T05:40:02.00'  # process died mid-write
    )
    partial = trace_summary.summarize_stream(truncated)
    assert partial["total_records"] == 2
    assert partial["total_events"] == 1
    assert partial["loop_sequence"] == ["AlphaLoop"]