@assistant-ui/react 0.15.24 → 0.15.25

package/src/unstable/useMessageStallDetection.test.tsx

@@ -0,0 +1,40 @@
+// @vitest-environment jsdom
+
+import { renderToString } from "react-dom/server";
+import { afterEach, describe, expect, it, vi } from "vitest";
+
+const { message } = vi.hoisted(() => ({
+  message: {
+    status: { type: "running" as const },
+    content: [{ type: "text" as const, text: "streaming" }],
+  },
+}));
+
+vi.mock("@assistant-ui/store", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("@assistant-ui/store")>()),
+  useAuiState: (selector: (state: { message: typeof message }) => unknown) =>
+    selector({ message }),
+}));
+
+vi.mock("@assistant-ui/store/internal", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("@assistant-ui/store/internal")>()),
+  useShallowSelector: <T,>(selector: T) => selector,
+}));
+
+import { unstable_useMessageStallDetection } from "./useMessageStallDetection";
+
+describe("unstable_useMessageStallDetection", () => {
+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+
+  it("does not read the wall clock during server render", () => {
+    const nowSpy = vi.spyOn(Date, "now");
+    const Probe = () => (
+      <span>{String(unstable_useMessageStallDetection().stalled)}</span>
+    );
+
+    expect(renderToString(<Probe />)).toBe("<span>false</span>");
+    expect(nowSpy).not.toHaveBeenCalled();
+  });
+});

package/src/unstable/useMessageStallDetection.ts

@@ -57,9 +57,9 @@ export function unstable_useMessageStallDetection(
   );
 
   const running = activity[0] === true;
-  const lastActivityRef = useRef(Date.now());
+  const lastActivityRef = useRef(0);
   const [stalled, setStalled] = useState(false);
-  const [, setTick] = useState(0);
+  const [now, setNow] = useState(0);
 
   useEffect(() => {
     if (!running) return undefined;
@@ -73,26 +73,30 @@ export function unstable_useMessageStallDetection(
       return undefined;
     }
 
+    const stall = () => {
+      setNow(Date.now());
+      setStalled(true);
+    };
     const sinceActivity = Date.now() - lastActivityRef.current;
     if (sinceActivity >= thresholdMs) {
-      setStalled(true);
+      stall();
       return undefined;
     }
 
     setStalled(false);
-    const id = setTimeout(() => setStalled(true), thresholdMs - sinceActivity);
+    const id = setTimeout(stall, thresholdMs - sinceActivity);
     return () => clearTimeout(id);
   }, [running, activity, thresholdMs]);
 
   useEffect(() => {
     if (!stalled) return undefined;
-    const id = setInterval(() => setTick((t) => t + 1), 1000);
+    const id = setInterval(() => setNow(Date.now()), 1000);
     return () => clearInterval(id);
   }, [stalled]);
 
   if (!stalled) return { stalled: false, stalledForMs: 0 };
   return {
     stalled: true,
-    stalledForMs: Math.max(0, Date.now() - lastActivityRef.current),
+    stalledForMs: Math.max(0, now - lastActivityRef.current),
   };
 }

package/src/utils/smooth/useSmooth.test.tsx

@@ -1,5 +1,6 @@
 /** @vitest-environment jsdom */
 import { act, renderHook } from "@testing-library/react";
+import { renderToString } from "react-dom/server";
 import { afterEach, describe, expect, it, vi } from "vitest";
 import type {
   MessagePartState,
@@ -68,7 +69,7 @@ const driveAndCount = (minCommitMs: number) => {
     .spyOn(globalThis, "cancelAnimationFrame")
     .mockImplementation(() => {});
   let now = 1000;
-  const nowSpy = vi.spyOn(Date, "now").mockImplementation(() => now);
+  const nowSpy = vi.spyOn(performance, "now").mockImplementation(() => now);
 
   const state = runningState("0123456789");
   const { result, unmount } = renderHook(() =>
@@ -105,6 +106,16 @@ describe("useSmooth", () => {
     Reflect.deleteProperty(window, "matchMedia");
   });
 
+  it("does not read the wall clock during server render", () => {
+    const nowSpy = vi.spyOn(Date, "now");
+    const Probe = () => (
+      <span>{useSmooth(runningState("streaming"), true).text}</span>
+    );
+
+    expect(renderToString(<Probe />)).toBe("<span></span>");
+    expect(nowSpy).not.toHaveBeenCalled();
+  });
+
   it("returns the input state unchanged when disabled", () => {
     const state = textState("hello");
     const { result } = renderHook(() => useSmooth(state, false));
@@ -189,8 +200,8 @@ describe("useSmooth", () => {
       return raf.length;
     });
     vi.spyOn(globalThis, "cancelAnimationFrame").mockImplementation(() => {});
-    let now = 1_000_000;
-    vi.spyOn(Date, "now").mockImplementation(() => now);
+    let now = 0;
+    vi.spyOn(performance, "now").mockImplementation(() => now);
 
     const frame = () => {
       if (raf.length === 0) return;
@@ -246,7 +257,7 @@ describe("useSmooth", () => {
     });
     vi.spyOn(globalThis, "cancelAnimationFrame").mockImplementation(() => {});
     let now = 1000;
-    vi.spyOn(Date, "now").mockImplementation(() => now);
+    vi.spyOn(performance, "now").mockImplementation(() => now);
 
     const frame = () => {
       if (raf.length === 0) return;
@@ -297,7 +308,7 @@ describe("useSmooth", () => {
     });
     vi.spyOn(globalThis, "cancelAnimationFrame").mockImplementation(() => {});
     let now = 1000;
-    vi.spyOn(Date, "now").mockImplementation(() => now);
+    vi.spyOn(performance, "now").mockImplementation(() => now);
 
     const frame = () => {
       if (raf.length === 0) return;

package/src/utils/smooth/useSmooth.ts

@@ -49,8 +49,8 @@ const DEFAULT_MAX_CHAR_INTERVAL_MS = 5;
 
 class TextStreamAnimator {
   private animationFrameId: number | null = null;
-  private lastUpdateTime: number = Date.now();
-  public lastCommitTime: number = 0;
+  private lastUpdateTime: number = 0;
+  public lastCommitTime: number = -Infinity;
 
   public targetText: string = "";
   public drainMs: number = DEFAULT_DRAIN_MS;
@@ -68,7 +68,7 @@ class TextStreamAnimator {
 
   start() {
     if (this.animationFrameId !== null) return;
-    this.lastUpdateTime = Date.now();
+    this.lastUpdateTime = performance.now();
     this.animate();
   }
 
@@ -80,7 +80,7 @@ class TextStreamAnimator {
   }
 
   private animate = () => {
-    const currentTime = Date.now();
+    const currentTime = performance.now();
     const deltaTime = currentTime - this.lastUpdateTime;
     let timeToConsume = deltaTime;
 
@@ -214,8 +214,8 @@ export const useSmooth = (
     }
   }, [smoothStatusStore, enabled, text, displayedText, state.status]);
 
-  const [animatorRef] = useState<TextStreamAnimator>(
-    new TextStreamAnimator(displayedText, setText),
+  const [animatorRef] = useState(
+    () => new TextStreamAnimator(displayedText, setText),
   );
 
   useEffect(() => {
@@ -240,7 +240,7 @@ export const useSmooth = (
       if (state.status.type === "running") {
         animatorRef.currentText = "";
         animatorRef.targetText = text;
-        animatorRef.lastCommitTime = 0;
+        animatorRef.lastCommitTime = -Infinity;
         animatorRef.start();
       } else {
         animatorRef.currentText = text;