Xray Lens: Request Trace

"use client";

import * as React from "react";

import { XrayLens, XrayLensSpot } from "@/components/ui/xray-lens";

// One trace, in the order it ran. Each span is drawn over the part of the answer it produced.
const TRACE = [
  { id: "embed", ms: 18 },
  { id: "retrieve", ms: 42 },
  { id: "rerank", ms: 31 },
  { id: "generate", ms: 812 },
  { id: "cite", ms: 12 },
  { id: "suggest", ms: 64 },
] as const;

type SpanId = (typeof TRACE)[number]["id"];
type At = { start: number; ms: number };
type Timing = Record<SpanId, At>;

const BASE = Object.fromEntries(TRACE.map((s) => [s.id, s.ms])) as Record<SpanId, number>;

function lay(durations: Record<SpanId, number>) {
  let t = 0;
  const timing = {} as Timing;
  for (const { id } of TRACE) {
    timing[id] = { start: t, ms: durations[id] };
    t += durations[id];
  }
  return { timing, total: t };
}

/** The trace re-measured every couple of seconds, as a live dashboard would show it. */
function useLiveTrace() {
  const [durations, setDurations] = React.useState(BASE);
  React.useEffect(() => {
    const id = window.setInterval(() => {
      setDurations(
        Object.fromEntries(
          TRACE.map((s) => [s.id, Math.round(s.ms * (0.88 + Math.random() * 0.24))]),
        ) as Record<SpanId, number>,
      );
    }, 1600);
    return () => window.clearInterval(id);
  }, []);
  return lay(durations);
}

const spanBox = "outline-1 -outline-offset-1 outline-emerald-400/40 bg-emerald-400/[0.06]";

/** Where a span sits on the whole request's timeline. */
function Bar({ at, total }: { at: At; total: number }) {
  return (
    <span className="relative h-1 min-w-8 flex-1 rounded-full bg-white/10">
      <span
        className="absolute inset-y-0 min-w-[3px] rounded-full bg-emerald-400 transition-[left,width] duration-500"
        style={{ left: `${(at.start / total) * 100}%`, width: `${(at.ms / total) * 100}%` }}
      />
    </span>
  );
}

/** A span in a box: its name and time, its place on the timeline, and a note. */
function Span({ name, at, total, note }: { name: string; at: At; total: number; note?: string }) {
  return (
    <span className="flex h-full min-w-0 flex-col justify-center gap-1.5 px-2.5 font-mono text-[10.5px] leading-none">
      <span className="flex items-center justify-between gap-3 whitespace-nowrap">
        <span className="text-emerald-300">{name}</span>
        <span className="text-zinc-100 tabular-nums">{at.ms}ms</span>
      </span>
      <span className="flex">
        <Bar at={at} total={total} />
      </span>
      {note && <span className="truncate text-zinc-500">{note}</span>}
    </span>
  );
}

/** A span on one line, for the section labels. */
function SpanRow({ name, at, total }: { name: string; at: At; total: number }) {
  return (
    <span className="flex h-full items-center gap-2.5 font-mono text-[10px] leading-none whitespace-nowrap">
      <span className="text-emerald-300">{name}</span>
      <Bar at={at} total={total} />
      <span className="text-zinc-100 tabular-nums">{at.ms}ms</span>
    </span>
  );
}

// The answer as text and citations. The surface and the xray both render it from
// here, so they wrap the same way.
const ANSWER: (string | number)[] = [
  "Create the new key and deploy it alongside the old one. Once traffic has moved over, revoke the old key",
  1,
  ". Both keys stay valid during the overlap, so no request fails",
  2,
  ".",
];
const WORDS = ANSWER.filter((part) => typeof part === "string").join(" ").split(/\s+/).length;

const answer = "text-[15px] leading-relaxed text-pretty";
const cite =
  "mx-0.5 inline-grid size-[18px] place-items-center rounded-full align-[2px] text-[10px] font-semibold";

/** The answer again, each word colored by when it streamed in: green first, violet last. */
function Streamed() {
  let n = 0;
  return (
    <p className={answer}>
      {ANSWER.map((part, i) =>
        typeof part === "number" ? (
          <span key={i} className={`${cite} text-transparent`}>
            {part}
          </span>
        ) : (
          part.split(/(\s+)/).map((word, j) =>
            /\s/.test(word) || !word ? (
              word
            ) : (
              <span key={`${i}-${j}`} style={{ color: `oklch(0.82 0.14 ${160 + (n++ / WORDS) * 140})` }}>
                {word}
              </span>
            ),
          )
        ),
      )}
    </p>
  );
}

const SOURCES = [
  { host: "docs.acme.dev", title: "Rotating API keys", score: "0.91" },
  { host: "acme.dev/blog", title: "Zero-downtime deploys", score: "0.88" },
  { host: "status.acme.dev", title: "Incident 4127", score: "0.64" },
];

function Answer({ timing, total }: { timing: Timing; total: number }) {
  return (
    <div className="flex flex-col gap-5 p-5 sm:p-6">
      <XrayLensSpot
        label="embed · query"
        xrayClassName={spanBox}
        xray={<Span name="embed" at={timing.embed} total={total} note="trace 7f3a2c · 1536 dims" />}
        className="flex h-12 items-center gap-2.5 rounded-xl bg-zinc-100 px-3.5 text-sm text-zinc-900 dark:bg-zinc-900 dark:text-zinc-100"
      >
        <svg aria-hidden viewBox="0 0 20 20" className="size-4 shrink-0 fill-none stroke-zinc-400 stroke-2">
          <circle cx="9" cy="9" r="6" />
          <path d="M13.5 13.5 17 17" strokeLinecap="round" />
        </svg>
        <span className="truncate">How do I rotate API keys without downtime?</span>
      </XrayLensSpot>

      <div>
        <XrayLensSpot
          label="generate"
          xray={<SpanRow name="generate" at={timing.generate} total={total} />}
          className="text-xs leading-4 font-medium tracking-wider text-zinc-500 uppercase"
        >
          Answer
        </XrayLensSpot>
        <XrayLensSpot
          label={`stream · 214 tokens · first at ${Math.round(timing.generate.ms * 0.22)}ms`}
          xray={<Streamed />}
          className="mt-2"
        >
          <p className={`${answer} text-zinc-700 dark:text-zinc-300`}>
            {ANSWER.map((part, i) =>
              typeof part === "number" ? (
                <XrayLensSpot
                  key={i}
                  as="span"
                  label={`cite · doc_${part} · ${timing.cite.ms}ms`}
                  xrayClassName="outline-1 -outline-offset-1 outline-amber-300/80 bg-amber-300/15"
                  xray={
                    <span className="grid size-full place-items-center font-mono text-[9px] text-amber-200">
                      d{part}
                    </span>
                  }
                  className={`${cite} bg-zinc-100 text-zinc-600 dark:bg-zinc-800 dark:text-zinc-300`}
                >
                  {part}
                </XrayLensSpot>
              ) : (
                part
              ),
            )}
          </p>
        </XrayLensSpot>
      </div>

      <div>
        <XrayLensSpot
          label="retrieve"
          xray={<SpanRow name="retrieve" at={timing.retrieve} total={total} />}
          className="text-xs leading-4 font-medium tracking-wider text-zinc-500 uppercase"
        >
          Sources
        </XrayLensSpot>
        <div className="mt-2 grid gap-2 sm:grid-cols-3">
          {SOURCES.map((source, i) => (
            <XrayLensSpot
              key={source.host}
              label={`${i === 2 ? "rerank" : "retrieve"} · doc_${i + 1}`}
              xrayClassName={spanBox}
              xray={
                <Span
                  name={i === 2 ? "rerank" : `doc_${i + 1}`}
                  at={i === 2 ? timing.rerank : timing.retrieve}
                  total={total}
                  note={`score ${source.score}${i === 2 ? " · dropped" : ""}`}
                />
              }
              className="rounded-lg p-2.5 ring-1 ring-zinc-950/10 dark:ring-white/10"
            >
              <p className="truncate text-[11px] text-zinc-500">{source.host}</p>
              <p className="mt-0.5 truncate text-sm font-medium text-zinc-900 dark:text-zinc-100">
                {source.title}
              </p>
            </XrayLensSpot>
          ))}
        </div>
      </div>

      <XrayLensSpot
        label="suggest · 2 follow-ups"
        xrayClassName={spanBox}
        xray={<SpanRow name="suggest" at={timing.suggest} total={total} />}
        className="flex flex-wrap gap-2 self-start"
      >
        {["Revoke a leaked key", "Scope a key to one project"].map((q) => (
          <button
            key={q}
            type="button"
            className="h-8 rounded-full px-3 text-xs text-zinc-600 ring-1 ring-zinc-950/10 transition-colors hover:bg-zinc-50 dark:text-zinc-400 dark:ring-white/10 dark:hover:bg-zinc-900"
          >
            {q}
          </button>
        ))}
      </XrayLensSpot>
    </div>
  );
}

export function XrayLensTraceDemo() {
  const { timing, total } = useLiveTrace();
  return (
    <section className="w-full max-w-2xl">
      <div className="mb-8 text-center">
        <p className="font-mono text-xs tracking-widest text-zinc-500 uppercase">Observability</p>
        <h3 className="mt-3 text-3xl font-semibold tracking-tight text-balance text-zinc-900 dark:text-zinc-50">
          Every answer, traced to the millisecond.
        </h3>
      </div>
      <XrayLens
        rest={false}
        size="28%"
        autoPlay
        zoom={1.08}
        lensLabel="Trace lens"
        labels={["Answer", "Trace"]}
        toggleClassName="top-full right-auto bottom-auto left-1/2 mt-4 -translate-x-1/2"
        surface={<Answer timing={timing} total={total} />}
        xray={
          <div className="absolute inset-0 bg-[linear-gradient(90deg,rgb(52_211_153/0.07)_1px,transparent_1px)] bg-size-[24px_100%]" />
        }
        xrayClassName="bg-[#050807] text-zinc-100"
        className="rounded-2xl bg-white shadow-[0_1px_2px_rgb(0_0_0/0.05),0_16px_40px_-20px_rgb(0_0_0/0.3)] ring-1 ring-zinc-950/5 dark:bg-zinc-950 dark:ring-white/10"
      />
      <p className="mt-20 text-center text-sm text-zinc-500">
        The lens tours each span. Point at any part of the answer to take over.
      </p>
    </section>
  );
}