engine/ui_trace.py

"""Trace rendering for the Gradio web UI."""
from __future__ import annotations

from .agent import Step
from .trace_collector import TraceEvent

_TOOL_ICON = {
    "write_file": "✏️", "read_file": "📖", "list_files": "📂",
    "run_python": "▶️", "run_tests": "🧪", "check_app": "🌐",
}


def merge_step_metadata(events: list[TraceEvent], raw_history: list) -> list[TraceEvent]:
    """Attach LiteForge timing/token stats to tool_call events."""
    if not raw_history:
        return events
    calls = [e for e in events if e.kind == "tool_call"]
    merged: list[TraceEvent] = []
    call_idx = 0
    for ev in events:
        if ev.kind != "tool_call" or call_idx >= len(raw_history):
            merged.append(ev)
            continue
        step = raw_history[call_idx]
        call_idx += 1
        merged.append(TraceEvent(
            kind=ev.kind, name=ev.name, detail=ev.detail, step=ev.step,
            duration_ms=getattr(step, "duration_ms", None),
            tokens=getattr(step, "total_tokens", None),
        ))
    return merged


def format_trace_md(
    events: list[TraceEvent],
    *,
    steps: list[Step] | None = None,
    max_detail: int = 500,
    idle: str = "_waiting for the model…_",
) -> str:
    """Render trace events as markdown with expandable tool I/O."""
    if not events and not steps:
        return idle
    if not events and steps:
        return _steps_only_md(steps)

    lines: list[str] = []
    step_no = 0
    i = 0
    while i < len(events):
        ev = events[i]
        if ev.kind == "tool_call":
            icon = _TOOL_ICON.get(ev.name, "🔧")
            meta = _meta_badge(ev)
            summary = f"`{step_no}` &nbsp; {icon} **{ev.name}**{meta}"
            detail = _truncate(ev.detail, max_detail)
            block = f"<details><summary>{summary}</summary>\n\n```json\n{detail}\n```\n</details>"
            if i + 1 < len(events) and events[i + 1].kind == "tool_result":
                result = _truncate(events[i + 1].detail, max_detail)
                block += f"\n\n↳ result:\n\n```json\n{result}\n```"
                i += 1
            lines.append(block)
            step_no += 1
        elif ev.kind == "tier_escalation":
            lines.append(f"⬆️ **escalated** → `{ev.name}`: {ev.detail}")
        elif ev.kind == "final":
            lines.append("✅ **final answer**")
        elif ev.kind == "error":
            lines.append(f"⚠️ **error**: {_truncate(ev.detail, max_detail)}")
        i += 1
    return "\n\n".join(lines) if lines else idle


def format_fanout_trace_md(results) -> str:
    """Per-subagent expandable traces for fan-out mode."""
    if not results:
        return "_no subagents_"
    blocks = []
    for r in results:
        events = getattr(r, "trace_events", None) or []
        inner = format_trace_md(events, steps=r.steps, idle="_no steps yet_")
        verdict = "✓ verified" if r.verified else ("⚠️ error" if r.error else "· unverified")
        blocks.append(
            f"<details><summary>`{r.index + 1}` **subagent** ({r.model}): "
            f"{len(r.steps)} steps · {verdict}</summary>\n\n{inner}\n</details>"
        )
    return "\n\n".join(blocks)


def _steps_only_md(steps: list[Step]) -> str:
    lines = []
    for s in steps:
        kind = s.kind
        if kind.startswith("tool_call:"):
            tool = kind.split(":", 1)[1]
            icon = _TOOL_ICON.get(tool, "🔧")
            meta = ""
            if s.total_tokens:
                meta = f" · {s.total_tokens} tok"
            lines.append(f"`{s.number}` &nbsp; {icon} **{tool}**{meta}")
        elif kind == "response":
            lines.append("✅ **final answer**")
        else:
            lines.append(f"• {kind}")
    return "\n\n".join(lines) if lines else "_waiting for the model…_"


def _meta_badge(ev: TraceEvent) -> str:
    parts = []
    if ev.duration_ms is not None:
        parts.append(f"{ev.duration_ms}ms")
    if ev.tokens is not None:
        parts.append(f"{ev.tokens} tok")
    return f" <span class='trace-meta'>({', '.join(parts)})</span>" if parts else ""


def _truncate(text: str, limit: int) -> str:
    text = text.strip()
    if len(text) <= limit:
        return text
    return text[:limit] + f"\n… ({len(text)} chars total)"