tracer-cloud--opensre/surfaces/cli/ui/renderer/diagnose.py

"""Live Markdown renderer for diagnose-node token streams."""

from __future__ import annotations

import sys
import time
from collections.abc import Callable, Mapping
from typing import Any

from rich.console import Console
from rich.live import Live
from rich.markdown import Markdown
from rich.spinner import Spinner
from rich.text import Text

from core.domain.stream import StreamEvent
from platform.analytics.cli import capture_investigation_lifecycle_event
from platform.analytics.events import Event
from surfaces.cli.ui.renderer.constants import (
    _BOLD,
    _DIAGNOSE_LIVE_REFRESH,
    _DIAGNOSE_NODE,
    _DIAGNOSE_RENDER_INTERVAL_S,
    _DIAGNOSE_SPINNER_COLOR,
    _DIAGNOSE_SPINNER_NAME,
    _DIM,
    _GREEN,
    _RESET,
    _WHITE,
    _render_source,
)
from surfaces.interactive_shell.ui.output import (
    ProgressTracker,
    _repl_progress_active,
    get_output_format,
    set_live_console,
    stop_display,
    unregister_live_console,
)


class _DiagnoseStreamRenderer:
    """Owns the diagnose-node live-streaming state machine.

    Encapsulates the buffer of incoming token deltas, the lazy Rich Console
    + Live region, and the throttled Markdown re-parse cadence. Exists so
    :class:`StreamRenderer` keeps a single responsibility (event dispatch
    + node lifecycle + final report) while diagnose-specific streaming
    concerns live in one focused place.

    Lifecycle: :meth:`start` → :meth:`append_chunk` (per token-delta event)
    → :meth:`finish`. The same instance can be reused across multiple
    investigation runs — :meth:`start` resets all state.
    """

    def __init__(
        self,
        console: Console | None = None,
        tracker: ProgressTracker | None = None,
        *,
        local: bool = False,
        state_provider: Callable[[], Mapping[str, object]] | None = None,
    ) -> None:
        self.buffer: list[str] = []
        self._live: Live | None = None
        self._started: float = 0.0
        # Last time we re-rendered ``Markdown(buffer)`` into the Live region.
        # Throttled to ``_DIAGNOSE_RENDER_INTERVAL_S`` so long streams don't
        # incur O(n²) parsing.
        self._last_render: float = 0.0
        self._console: Console | None = console
        self._tracker: ProgressTracker | None = tracker
        self._local = local
        self._state_provider = state_provider

    @property
    def streamed(self) -> bool:
        """True if any chunks were buffered during the run.

        Callers (specifically :meth:`StreamRenderer._print_report`) use this
        to decide whether the final ``Root Cause`` summary should be
        suppressed — it would duplicate text the user just watched stream.
        """
        return bool(self.buffer)

    def start(self) -> None:
        """Reset state and open the Live region (rich) or print a placeholder (text)."""
        self.buffer = []
        self._started = time.monotonic()
        # 0.0 sentinel forces the first chunk past the throttle gate so the
        # user sees something rendered as soon as tokens arrive.
        self._last_render = 0.0

        if _repl_progress_active():
            return

        if get_output_format() != "rich":
            sys.stdout.write(f"  … {_DIAGNOSE_NODE}\n")
            sys.stdout.flush()
            return

        if self._console is None:
            self._console = Console(highlight=False)
        spinner = Spinner(
            _DIAGNOSE_SPINNER_NAME,
            text=Text(
                f"{_DIAGNOSE_NODE}  reasoning…",
                style=f"bold {_DIAGNOSE_SPINNER_COLOR}",
            ),
            style=f"bold {_DIAGNOSE_SPINNER_COLOR}",
        )
        self._live = Live(
            spinner,
            console=self._console,
            refresh_per_second=_DIAGNOSE_LIVE_REFRESH,
            transient=False,
        )

        # Shrink the gap: stop previous display immediately before starting new one
        if self._tracker is not None:
            self._tracker.stop()
        else:
            stop_display()

        # Register console globally so that print_above_renderable fallbacks
        # correctly print above this live region during the diagnose phase.
        set_live_console(self._console)
        self._live.start()

    def append_chunk(self, event: StreamEvent) -> None:
        """Append a token delta to the buffer; refresh the Live region (throttled).

        The chunk's ``content`` shape varies by provider: OpenAI emits a
        plain string; some Anthropic SDK paths emit a list of content blocks.
        :func:`_flatten_chunk_content` handles both — calling ``str()`` on
        the list shape would render its Python repr instead of reasoning.
        """
        chunk = event.data.get("data", {}).get("chunk", {})
        content = chunk.get("content", "") if isinstance(chunk, dict) else ""
        if not content:
            return
        text = _flatten_chunk_content(content)
        if not text:
            return
        self.buffer.append(text)
        if len(self.buffer) == 1:
            latency_ms = (time.monotonic() - self._started) * 1000
            state = dict(self._state_provider()) if self._state_provider is not None else None
            capture_investigation_lifecycle_event(
                Event.INVESTIGATION_FIRST_HYPOTHESIS_RENDERED,
                {
                    "latency_ms": int(latency_ms),
                    "stage": _DIAGNOSE_NODE,
                    "source": _render_source(local=self._local),
                },
                state=state,
            )
        if self._live is None:
            if _repl_progress_active() and self._tracker is not None:
                preview = "".join(self.buffer)
                if len(preview) > 80:
                    preview = "…" + preview[-77:]
                self._tracker.update_subtext(_DIAGNOSE_NODE, preview, duration=30.0)
            return
        # Throttle Markdown re-parse to once per refresh window; the final
        # flush in :meth:`finish` guarantees the latest buffer is rendered
        # before the Live region closes.
        now = time.monotonic()
        if now - self._last_render >= _DIAGNOSE_RENDER_INTERVAL_S:
            self._live.update(Markdown("".join(self.buffer)))
            self._last_render = now

    def finish(self, message: str | None = None) -> None:
        """Close the Live region (or text-mode flush) and print the resolved-dot line.

        ``message`` is appended dim-styled to the resolution line — typically
        a validity-score summary built by ``_build_node_message``.
        """
        elapsed = time.monotonic() - self._started

        if self._live is not None:
            # Final flush: any chunks pending in the last throttle window
            # render here so the user sees the complete reasoning.
            if self.buffer:
                self._live.update(Markdown("".join(self.buffer)))
            try:
                self._live.stop()
            finally:
                self._live = None
                # Unregister only if we own it (safeguard against subsequent activations)
                unregister_live_console(self._console)
            sys.stdout.write(
                f"  {_GREEN}●{_RESET}  {_BOLD}{_WHITE}{_DIAGNOSE_NODE}{_RESET}"
                f"  {_DIM}{elapsed:.1f}s{_RESET}"
            )
            if message:
                sys.stdout.write(f"  {_DIM}{message}{_RESET}")
            sys.stdout.write("\n")
            sys.stdout.flush()
        else:
            if self.buffer:
                for line in "".join(self.buffer).strip().splitlines():
                    print(f"  {line}")
            tail = f"  ● {_DIAGNOSE_NODE}  {elapsed:.1f}s"
            if message:
                tail += f"  {message}"
            print(tail)


def _flatten_chunk_content(content: Any) -> str:
    """Resolve a chat-model chunk's ``content`` to plain text.

    OpenAI emits a string. Anthropic-style adapters may emit a list of content
    blocks where each block may be an object with ``.text`` or a dict
    with a ``"text"`` key. Non-text blocks (tool-use, image) are skipped.
    """
    if isinstance(content, str):
        return content
    if not isinstance(content, list):
        return ""
    parts: list[str] = []
    for block in content:
        if isinstance(block, dict):
            text_value = block.get("text")
            if isinstance(text_value, str):
                parts.append(text_value)
            continue
        text_value = getattr(block, "text", None)
        if isinstance(text_value, str):
            parts.append(text_value)
    return "".join(parts)