band-aider-pub/tests/ui/test_trace_publisher.py

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

"""In-memory JSONL trace publisher contracts."""

from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime
from enum import Enum
import io
import json
from pathlib import Path

from band_aider.ui.dto.log_message import LogMessageDTO
from band_aider.ui.dto.loop_status import LoopStatusDTO
from band_aider.ui.dto.run_completion import RunCompletionDTO
from band_aider.ui.dto.token_budget_stats import TokenUsageBreakdownDTO
from band_aider.ui.trace_publisher import TRACE_FORMAT, TracePublisher


class _MemoryPublisher:
    """Retain published events in local memory instead of a live UI."""

    def __init__(self) -> None:
        self.events: list[object] = []

    def publish(self, event: object) -> None:
        self.events.append(event)


def _read_records(stream: io.StringIO) -> list[dict]:
    """Parse every JSONL record collected in one in-memory stream."""

    return [json.loads(line) for line in stream.getvalue().splitlines()]


def test_trace_publisher_writes_one_ordered_jsonl_record_per_event() -> None:
    """Each published event becomes one timestamped record after the header."""

    stream = io.StringIO()
    publisher = TracePublisher(_MemoryPublisher(), stream, origin="memory")

    publisher.publish(LoopStatusDTO(previous_loop=None, current_loop="AlphaLoop",
                                    next_loop="BetaLoop", current_item=None,
                                    work_focus=None))
    publisher.publish(LogMessageDTO(created=1.5, level="INFO",
                                    logger_name="loop", message="entered"))

    records = _read_records(stream)
    assert [record["seq"] for record in records] == [0, 1, 2]
    assert records[0]["event"] == "trace_start"
    assert records[0]["data"]["format"] == TRACE_FORMAT
    assert records[0]["data"]["origin"] == "memory"
    assert [record["event"] for record in records[1:]] == [
        "LoopStatusDTO",
        "LogMessageDTO",
    ]
    assert records[1]["data"]["current_loop"] == "AlphaLoop"
    assert records[2]["data"]["message"] == "entered"
    for record in records:
        # Every timestamp must round-trip as a timezone-aware ISO-8601 instant.
        assert datetime.fromisoformat(record["ts"]).tzinfo is not None


def test_trace_publisher_expands_nested_dataclasses_tuples_enums_and_paths() -> None:
    """Real DTO payloads survive JSON round-tripping without custom encoders."""

    class Phase(Enum):
        """Local enum proving enum values serialize as plain data."""

        ALPHA = "alpha"

    @dataclass(frozen=True)
    class Nested:
        """Local carrier mixing the non-trivial payload shapes at once."""

        phase: Phase
        target: Path
        pair: tuple[int, ...]

    stream = io.StringIO()
    publisher = TracePublisher(_MemoryPublisher(), stream)
    breakdown = TokenUsageBreakdownDTO(
        started_at=10,
        model_name="m",
        endpoint="e",
        level="base",
        tokens_spent=7,
    )
    publisher.publish(RunCompletionDTO(
        final_summary="done",
        report_path=None,
        total_tokens_spent=7,
        largest_query_tokens=7,
        total_seconds=1.0,
        average_seconds_per_cot=1.0,
        cot_count=1,
        files_modified=1,
        files_added=0,
        token_breakdown=(breakdown,),
    ))
    publisher.publish(Nested(phase=Phase.ALPHA, target=Path("src/x.py"), pair=(1, 2)))

    completion = _read_records(stream)[1]["data"]
    assert completion["token_breakdown"] == [
        {
            "started_at": 10,
            "model_name": "m",
            "endpoint": "e",
            "level": "base",
            "tokens_spent": 7,
            "cache_hit_tokens": 0,
        }
    ]
    nested = _read_records(stream)[2]["data"]
    assert nested == {"phase": "alpha", "target": "src/x.py", "pair": [1, 2]}


def test_trace_publisher_forwards_every_event_to_the_inner_publisher() -> None:
    """The wrapped publisher observes identical events in identical order."""

    inner = _MemoryPublisher()
    publisher = TracePublisher(inner, io.StringIO())
    first = LogMessageDTO(created=0.0, level="INFO", logger_name="a", message="one")
    second = LoopStatusDTO(previous_loop=None, current_loop="BetaLoop",
                           next_loop=None, current_item=None, work_focus=None)

    publisher.publish(first)
    publisher.publish(second)

    assert inner.events == [first, second]


def test_trace_publisher_redacts_secret_named_fields_but_keeps_token_telemetry() -> None:
    """Exact secret field names hide while token-budget fields stay visible."""

    @dataclass(frozen=True)
    class SecretCarrier:
        """Local carrier proving redaction is name-exact, not substring-based."""

        api_key: str
        total_tokens: int

    stream = io.StringIO()
    publisher = TracePublisher(_MemoryPublisher(), stream)

    publisher.publish(SecretCarrier(api_key="sk-live-secret", total_tokens=42))

    data = _read_records(stream)[1]["data"]
    assert data == {"api_key": "<redacted>", "total_tokens": 42}


def test_trace_publisher_never_breaks_the_run_it_observes() -> None:
    """Unserializable payloads and a dead stream degrade silently, not loudly."""

    class Unpleasant:
        """Local object whose repr fails, the worst case for the fallback path."""

        def __repr__(self) -> str:
            raise RuntimeError("no repr for you")

    inner = _MemoryPublisher()
    stream = io.StringIO()
    publisher = TracePublisher(inner, stream)

    publisher.publish(Unpleasant())
    # The unserializable event was forwarded but never recorded.
    assert len(_read_records(stream)) == 1

    stream.close()
    publisher.publish(LogMessageDTO(created=0.0, level="INFO",
                                    logger_name="a", message="after close"))
    publisher.close()

    assert len(inner.events) == 2


def test_trace_publisher_close_is_idempotent() -> None:
    """Closing twice leaves the publisher settled rather than raising."""

    publisher = TracePublisher(_MemoryPublisher(), io.StringIO())

    publisher.close()
    publisher.close()