feat(ui): smooth streaming of agent message tokens

packages/ui/src/features/sessions/components/session-update/AgentMessage.tsx

@@ -3,6 +3,7 @@ import { Box, Code, IconButton } from "@radix-ui/themes";
 import { memo, useCallback, useMemo, useState } from "react";
 import type { Components } from "react-markdown";
 import { HighlightedCode } from "../../../../primitives/HighlightedCode";
+import { useSmoothText } from "../../../../primitives/hooks/useSmoothText";
 import { Tooltip } from "../../../../primitives/Tooltip";
 import { usePendingScrollStore } from "../../../code-editor/pendingScrollStore";
 import { MarkdownRenderer } from "../../../editor/components/MarkdownRenderer";
@@ -141,6 +142,9 @@ export const AgentMessage = memo(function AgentMessage({
   content,
 }: AgentMessageProps) {
   const [copied, setCopied] = useState(false);
+  // Reveal streamed tokens at a steady rate so bursty chunks read as smooth
+  // typing. Copy always uses the full `content`, not the partial reveal.
+  const displayedContent = useSmoothText(content);
 
   const handleCopy = useCallback(() => {
     navigator.clipboard.writeText(content);
@@ -151,7 +155,7 @@ export const AgentMessage = memo(function AgentMessage({
   return (
     <Box className="group/msg relative py-1 pl-3 text-[13px] [&>*:last-child]:mb-0 [&_p]:leading-[1.9]">
       <MarkdownRenderer
-        content={content}
+        content={displayedContent}
         componentsOverride={agentComponents}
       />
       <Box className="absolute top-1 left-full ml-2 opacity-0 transition-opacity group-hover/msg:opacity-100">

packages/ui/src/primitives/hooks/useSmoothText.test.ts

@@ -0,0 +1,113 @@
+import {
+  nextRevealLength,
+  useSmoothText,
+} from "@posthog/ui/primitives/hooks/useSmoothText";
+import { act, renderHook } from "@testing-library/react";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+
+describe("nextRevealLength", () => {
+  it.each<[string, number, number, number, number, number]>([
+    // label                                     current  target  elapsedMs  rate  expected
+    ["caught up: returns the target", 10, 10, 16, 120, 10],
+    ["caught up past target: clamps to target", 12, 10, 16, 120, 10],
+    // 120 chars/sec over 100ms = 12 chars.
+    ["advances proportionally to elapsed/rate", 0, 100, 100, 120, 12],
+    ["never overshoots the target", 95, 100, 1000, 120, 100],
+    // Tiny elapsed time would round to zero; keep forward progress.
+    ["always advances at least one when behind", 0, 100, 0, 120, 1],
+    ["snaps when lag is too large to ease", 0, 5000, 16, 120, 5000],
+  ])("%s", (_label, current, target, elapsedMs, rate, expected) => {
+    expect(nextRevealLength(current, target, elapsedMs, rate)).toBe(expected);
+  });
+});
+
+describe("useSmoothText", () => {
+  let now: number;
+  let rafCallbacks: Array<(t: number) => void>;
+
+  beforeEach(() => {
+    now = 0;
+    rafCallbacks = [];
+    vi.stubGlobal(
+      "requestAnimationFrame",
+      (cb: (t: number) => void): number => {
+        rafCallbacks.push(cb);
+        return rafCallbacks.length;
+      },
+    );
+    vi.stubGlobal("cancelAnimationFrame", () => {});
+    vi.stubGlobal("matchMedia", () => ({ matches: false }));
+  });
+
+  afterEach(() => {
+    vi.unstubAllGlobals();
+  });
+
+  // Advance one animation frame by the given time delta.
+  const flushFrame = (deltaMs: number) => {
+    now += deltaMs;
+    const callbacks = rafCallbacks;
+    rafCallbacks = [];
+    act(() => {
+      for (const cb of callbacks) cb(now);
+    });
+  };
+
+  it("shows existing text immediately on mount (no replay)", () => {
+    const { result } = renderHook(() => useSmoothText("already here"));
+    expect(result.current).toBe("already here");
+  });
+
+  it("reveals appended text gradually instead of all at once", () => {
+    const { result, rerender } = renderHook(
+      ({ text }) => useSmoothText(text, 100),
+      { initialProps: { text: "" } },
+    );
+
+    // A burst of 50 characters arrives at once.
+    rerender({ text: "x".repeat(50) });
+
+    // First frame establishes the clock and makes minimal forward progress.
+    flushFrame(0);
+    expect(result.current.length).toBe(1);
+
+    // 100ms at 100 chars/sec reveals ~10 more chars, nowhere near all 50.
+    flushFrame(100);
+    expect(result.current.length).toBe(11);
+    expect(result.current.length).toBeLessThan(50);
+
+    // Given enough time it catches up fully.
+    flushFrame(1000);
+    expect(result.current).toBe("x".repeat(50));
+  });
+
+  it("snaps when the target is replaced rather than appended", () => {
+    const { result, rerender } = renderHook(
+      ({ text }) => useSmoothText(text, 100),
+      { initialProps: { text: "" } },
+    );
+
+    rerender({ text: "hello world" });
+    flushFrame(0);
+    flushFrame(20); // partway through revealing "hello world"
+    expect(result.current.length).toBeLessThan("hello world".length);
+
+    // A completely different string is not a prefix -> show it all at once.
+    rerender({ text: "totally different" });
+    expect(result.current).toBe("totally different");
+  });
+
+  it("snaps immediately when reduced motion is preferred", () => {
+    vi.stubGlobal("matchMedia", (query: string) => ({
+      matches: query.includes("reduce"),
+    }));
+
+    const { result, rerender } = renderHook(
+      ({ text }) => useSmoothText(text, 100),
+      { initialProps: { text: "" } },
+    );
+
+    rerender({ text: "x".repeat(50) });
+    expect(result.current).toBe("x".repeat(50));
+  });
+});

packages/ui/src/primitives/hooks/useSmoothText.ts

@@ -0,0 +1,102 @@
+import { useEffect, useRef, useState } from "react";
+
+// Default reveal rate. Streamed LLM tokens arrive in bursty chunks; revealing
+// at a steady character rate makes the text feel like it is being typed rather
+// than appearing in jumps. ~120 chars/sec roughly matches a fast token stream
+// while staying smooth. See https://upstash.com/blog/smooth-streaming.
+const DEFAULT_CHARS_PER_SECOND = 120;
+
+// Once the displayed text falls this far behind the target we stop easing and
+// snap, so a large catch-up (e.g. a long buffered chunk) never feels sluggish.
+const MAX_LAG_CHARS = 600;
+
+/**
+ * Given the current reveal length, the target length and how much time elapsed
+ * since the last frame, return the next reveal length. Pure so the easing is
+ * unit-testable without timers. Never exceeds `target` and never goes backwards
+ * (callers handle resets when the target is replaced rather than appended).
+ */
+export function nextRevealLength(
+  current: number,
+  target: number,
+  elapsedMs: number,
+  charsPerSecond: number,
+): number {
+  if (current >= target) return target;
+  if (target - current > MAX_LAG_CHARS) return target;
+  const step = Math.ceil((charsPerSecond * elapsedMs) / 1000);
+  return Math.min(target, current + Math.max(step, 1));
+}
+
+function prefersReducedMotion(): boolean {
+  return (
+    typeof window !== "undefined" &&
+    typeof window.matchMedia === "function" &&
+    window.matchMedia("(prefers-reduced-motion: reduce)").matches
+  );
+}
+
+/**
+ * Smoothly reveals a growing string a few characters per frame instead of
+ * jumping whenever the source updates. Intended for streamed assistant text:
+ * tokens are appended to `target` in bursts, and this hook animates the visible
+ * prefix up to it at a steady rate.
+ *
+ * The text the source already had on mount is shown immediately, so re-rendered
+ * history and completed messages never replay a typewriter effect. Only growth
+ * that happens while mounted is animated. If the target is replaced (no longer a
+ * prefix of what we were showing) or motion is reduced, it snaps.
+ */
+export function useSmoothText(
+  target: string,
+  charsPerSecond = DEFAULT_CHARS_PER_SECOND,
+): string {
+  // Start fully revealed: history/completed messages should not animate.
+  const [revealLength, setRevealLength] = useState(target.length);
+  const revealRef = useRef(revealLength);
+  revealRef.current = revealLength;
+
+  // Track the previous target so we can detect replacement vs. append.
+  const prevTargetRef = useRef(target);
+
+  useEffect(() => {
+    const prevTarget = prevTargetRef.current;
+    prevTargetRef.current = target;
+
+    // Replacement (not an extension of what we were revealing) or reduced
+    // motion: show everything immediately.
+    if (!target.startsWith(prevTarget) || prefersReducedMotion()) {
+      setRevealLength(target.length);
+      return;
+    }
+
+    if (revealRef.current >= target.length) {
+      setRevealLength(target.length);
+      return;
+    }
+
+    let frame = 0;
+    let lastTime: number | null = null;
+
+    const tick = (now: number) => {
+      const elapsed = lastTime === null ? 0 : now - lastTime;
+      lastTime = now;
+      const next = nextRevealLength(
+        revealRef.current,
+        target.length,
+        elapsed,
+        charsPerSecond,
+      );
+      revealRef.current = next;
+      setRevealLength(next);
+      if (next < target.length) {
+        frame = requestAnimationFrame(tick);
+      }
+    };
+
+    frame = requestAnimationFrame(tick);
+    return () => cancelAnimationFrame(frame);
+  }, [target, charsPerSecond]);
+
+  return target.slice(0, revealLength);
+}