tatlock/src/core/tracing_router.py

"""
Trace viewer router.

Serves the trace viewer UI and trace files when tracing is enabled.
Only available when DEBUG=true.
"""
from pathlib import Path

from fastapi import APIRouter, HTTPException
from fastapi.responses import HTMLResponse, JSONResponse

from src.core.config import config
from src.core.logging_config import get_logger

logger = get_logger(__name__)

router = APIRouter(prefix="/traces", tags=["traces"])

TRACES_DIR = Path("logs/traces")
VIEWER_PATH = TRACES_DIR / "viewer.html"


def tracing_enabled() -> bool:
    """Check if tracing is enabled."""
    return config.DEBUG


@router.get("", response_class=HTMLResponse)
async def get_trace_viewer():
    """
    Serve the trace viewer UI.

    Returns the standalone HTML viewer for browsing traces.
    """
    if not tracing_enabled():
        raise HTTPException(status_code=404, detail="Tracing not enabled")

    if not VIEWER_PATH.exists():
        raise HTTPException(status_code=404, detail="Viewer not found")

    return HTMLResponse(content=VIEWER_PATH.read_text())


@router.get("/list")
async def list_traces(
    limit: int = 50,
    since_minutes: int | None = None,
    status: str | None = None,
    search: str | None = None,
):
    """
    List available trace files.

    Returns most recent traces first, with basic metadata.

    Args:
        limit: Maximum number of traces to return (default 50)
        since_minutes: Only return traces from the last N minutes
        status: Filter by status (completed, error, streaming)
        search: Search in request preview text
    """
    if not tracing_enabled():
        raise HTTPException(status_code=404, detail="Tracing not enabled")

    if not TRACES_DIR.exists():
        return {"traces": [], "total": 0}

    import json
    from datetime import datetime, timezone, timedelta

    # Calculate cutoff time if filtering by time
    cutoff_time = None
    if since_minutes:
        cutoff_time = datetime.now(timezone.utc) - timedelta(minutes=since_minutes)

    # Get all trace files, sorted by modification time (newest first)
    trace_files = sorted(
        TRACES_DIR.glob("trace_*.json"),
        key=lambda p: p.stat().st_mtime,
        reverse=True,
    )

    traces = []
    for path in trace_files:
        if len(traces) >= limit:
            break

        try:
            with open(path) as f:
                data = json.load(f)

            # Parse timestamp for filtering
            trace_timestamp = data.get("timestamp")
            if cutoff_time and trace_timestamp:
                try:
                    ts = datetime.fromisoformat(trace_timestamp.replace('Z', '+00:00'))
                    if ts < cutoff_time:
                        continue
                except (ValueError, TypeError):
                    pass

            # Filter by status
            trace_status = data.get("status", "")
            if status and trace_status != status:
                continue

            # Filter by search text
            request_preview = data.get("request", {}).get("input_preview", "")
            if search and search.lower() not in request_preview.lower():
                continue

            traces.append({
                "trace_id": data.get("trace_id"),
                "timestamp": trace_timestamp,
                "user": data.get("user"),
                "status": trace_status,
                "total_duration_ms": data.get("total_duration_ms"),
                "span_count": len(data.get("spans", [])),
                "request_preview": request_preview[:100],
            })
        except Exception as e:
            logger.warning("trace_list_parse_error", path=str(path), error=str(e))

    return {"traces": traces, "total": len(traces)}


@router.get("/{trace_id}")
async def get_trace(trace_id: str):
    """
    Get a specific trace by ID.

    Returns the full trace JSON.
    """
    if not tracing_enabled():
        raise HTTPException(status_code=404, detail="Tracing not enabled")

    # Sanitize trace_id to prevent path traversal
    if not trace_id.startswith("trace_") or "/" in trace_id or "\\" in trace_id:
        raise HTTPException(status_code=400, detail="Invalid trace ID")

    trace_path = TRACES_DIR / f"{trace_id}.json"

    if not trace_path.exists():
        raise HTTPException(status_code=404, detail="Trace not found")

    try:
        import json
        with open(trace_path) as f:
            data = json.load(f)
        return JSONResponse(content=data)
    except Exception as e:
        logger.error("trace_read_error", trace_id=trace_id, error=str(e))
        raise HTTPException(status_code=500, detail="Failed to read trace")