@assistant-ui/react 0.14.18 → 0.14.20

package/src/tests/toolCallTiming.test.tsx

@@ -0,0 +1,221 @@
+// @vitest-environment jsdom
+
+import { act, render, screen } from "@testing-library/react";
+import type { FC, PropsWithChildren } from "react";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import { AssistantRuntimeProvider } from "../context";
+import * as MessagePrimitive from "../primitives/message";
+import * as ThreadPrimitive from "../primitives/thread";
+import { useLocalRuntime } from "../legacy-runtime/runtime-cores/local/useLocalRuntime";
+import type { ChatModelAdapter, ThreadMessageLike } from "../index";
+import {
+  useToolCallElapsed,
+  unstable_useMessageStallDetection,
+} from "../index";
+
+const noOpAdapter: ChatModelAdapter = {
+  async *run() {},
+};
+
+const toolMessage = (timing?: {
+  startedAt: number;
+  completedAt?: number;
+}): ThreadMessageLike => ({
+  role: "assistant",
+  content: [
+    {
+      type: "tool-call",
+      toolCallId: "tc-1",
+      toolName: "search",
+      args: {},
+      argsText: "{}",
+      ...(timing !== undefined && { timing }),
+    },
+  ],
+  status: { type: "running" },
+});
+
+const ElapsedProbe: FC = () => {
+  const elapsedMs = useToolCallElapsed();
+  return (
+    <span data-testid="elapsed">
+      {elapsedMs === undefined ? "none" : String(elapsedMs)}
+    </span>
+  );
+};
+
+const StallProbe: FC = () => {
+  const { stalled } = unstable_useMessageStallDetection({ thresholdMs: 2000 });
+  return <span data-testid="stalled">{String(stalled)}</span>;
+};
+
+const RuntimeProvider: FC<
+  PropsWithChildren<{ messages: ThreadMessageLike[] }>
+> = ({ messages, children }) => {
+  const runtime = useLocalRuntime(noOpAdapter, { initialMessages: messages });
+  return (
+    <AssistantRuntimeProvider runtime={runtime}>
+      {children}
+    </AssistantRuntimeProvider>
+  );
+};
+
+const renderHarness = async (messages: ThreadMessageLike[], probe: FC) => {
+  render(<Harness messages={messages} probe={probe} />);
+  // Flush the runtime's deferred initialization, which fake timers hold back.
+  await act(async () => {
+    await vi.advanceTimersByTimeAsync(0);
+  });
+};
+
+const Harness: FC<{ messages: ThreadMessageLike[]; probe: FC }> = ({
+  messages,
+  probe: Probe,
+}) => (
+  <RuntimeProvider messages={messages}>
+    <ThreadPrimitive.Messages
+      components={{
+        Message: () => (
+          <MessagePrimitive.Parts components={{ tools: { Fallback: Probe } }} />
+        ),
+      }}
+    />
+  </RuntimeProvider>
+);
+
+describe("useToolCallElapsed", () => {
+  beforeEach(() => {
+    vi.useFakeTimers();
+    vi.setSystemTime(new Date("2026-01-01T00:00:10.000Z"));
+  });
+
+  afterEach(() => {
+    vi.useRealTimers();
+  });
+
+  it("returns undefined when the part carries no timing", async () => {
+    await renderHarness([toolMessage()], ElapsedProbe);
+    expect(screen.getByTestId("elapsed").textContent).toBe("none");
+  });
+
+  it("returns the settled duration for completed calls", async () => {
+    const startedAt = Date.now() - 5000;
+    await renderHarness(
+      [toolMessage({ startedAt, completedAt: startedAt + 1500 })],
+      ElapsedProbe,
+    );
+    expect(screen.getByTestId("elapsed").textContent).toBe("1500");
+  });
+
+  it("ticks while the call is running", async () => {
+    const startedAt = Date.now() - 1000;
+    await renderHarness([toolMessage({ startedAt })], ElapsedProbe);
+
+    expect(screen.getByTestId("elapsed").textContent).toBe("1000");
+
+    await act(async () => {
+      await vi.advanceTimersByTimeAsync(3000);
+    });
+
+    expect(screen.getByTestId("elapsed").textContent).toBe("4000");
+  });
+});
+
+describe("useToolCallElapsed on terminated calls", () => {
+  beforeEach(() => {
+    vi.useFakeTimers();
+    vi.setSystemTime(new Date("2026-01-01T00:00:10.000Z"));
+  });
+
+  afterEach(() => {
+    vi.useRealTimers();
+  });
+
+  it("claims no duration when the call ends without a result (cancelled run)", async () => {
+    const startedAt = Date.now() - 1000;
+    const cancelled: ThreadMessageLike = {
+      role: "assistant",
+      content: [
+        {
+          type: "tool-call",
+          toolCallId: "tc-1",
+          toolName: "search",
+          args: {},
+          argsText: "{}",
+          timing: { startedAt },
+        },
+      ],
+      status: { type: "incomplete", reason: "cancelled" },
+    };
+    await renderHarness([cancelled], ElapsedProbe);
+
+    expect(screen.getByTestId("elapsed").textContent).toBe("none");
+
+    await act(async () => {
+      await vi.advanceTimersByTimeAsync(3000);
+    });
+
+    expect(screen.getByTestId("elapsed").textContent).toBe("none");
+  });
+});
+
+describe("useToolCallElapsed outside any scope", () => {
+  it("returns undefined without a provider instead of throwing", () => {
+    const Standalone: FC = () => {
+      const elapsedMs = useToolCallElapsed();
+      return (
+        <span data-testid="standalone">
+          {elapsedMs === undefined ? "none" : String(elapsedMs)}
+        </span>
+      );
+    };
+    render(<Standalone />);
+    expect(screen.getByTestId("standalone").textContent).toBe("none");
+  });
+});
+
+describe("unstable_useMessageStallDetection", () => {
+  beforeEach(() => {
+    vi.useFakeTimers();
+  });
+
+  afterEach(() => {
+    vi.useRealTimers();
+  });
+
+  it("reports a stall after the threshold with no content changes", async () => {
+    await renderHarness([toolMessage({ startedAt: 0 })], StallProbe);
+
+    expect(screen.getByTestId("stalled").textContent).toBe("false");
+
+    await act(async () => {
+      await vi.advanceTimersByTimeAsync(2500);
+    });
+
+    expect(screen.getByTestId("stalled").textContent).toBe("true");
+  });
+
+  it("never stalls on settled messages", async () => {
+    const settled: ThreadMessageLike = {
+      role: "assistant",
+      content: [
+        {
+          type: "tool-call",
+          toolCallId: "tc-1",
+          toolName: "search",
+          args: {},
+          argsText: "{}",
+          result: "done",
+        },
+      ],
+      status: { type: "complete", reason: "stop" },
+    };
+    await renderHarness([settled], StallProbe);
+
+    await act(async () => {
+      await vi.advanceTimersByTimeAsync(5000);
+    });
+
+    expect(screen.getByTestId("stalled").textContent).toBe("false");
+  });
+});

package/src/unstable/useMessageStallDetection.ts

@@ -0,0 +1,91 @@
+"use client";
+
+import { useEffect, useRef, useState } from "react";
+import { useAuiState } from "@assistant-ui/store";
+
+export type Unstable_MessageStallDetectionOptions = {
+  /**
+   * Milliseconds of unchanged message content before the message counts as
+   * stalled.
+   * @default 2000
+   */
+  thresholdMs?: number | undefined;
+};
+
+export type Unstable_MessageStallDetection = {
+  /** True while the message is running and its content has not changed for at least `thresholdMs`. */
+  stalled: boolean;
+  /** Milliseconds since the last observed content change. `0` while not stalled. */
+  stalledForMs: number;
+};
+
+/**
+ * @deprecated Under active development and might change without notice.
+ *
+ * Detects mid-run output stalls on the current message: while the message is
+ * running, watches a fingerprint of its content (part count plus text,
+ * argument, and result sizes) and reports a stall once the fingerprint stops
+ * changing for `thresholdMs`. Useful for re-surfacing a "still working"
+ * indicator during tool think-time or provider stalls, after the first
+ * tokens have already streamed.
+ *
+ * Must be used inside a message scope.
+ */
+export function unstable_useMessageStallDetection(
+  options?: Unstable_MessageStallDetectionOptions,
+): Unstable_MessageStallDetection {
+  const thresholdMs = options?.thresholdMs ?? 2000;
+
+  const fingerprint = useAuiState((s) => {
+    if (s.message.status?.type !== "running") return undefined;
+    let size = 0;
+    for (const part of s.message.content) {
+      if (part.type === "text" || part.type === "reasoning") {
+        size += part.text.length;
+      } else if (part.type === "tool-call") {
+        size += part.argsText.length + (part.result !== undefined ? 1 : 0);
+      }
+    }
+    return `${s.message.content.length}:${size}`;
+  });
+
+  const running = fingerprint !== undefined;
+  const lastActivityRef = useRef(Date.now());
+  const [stalled, setStalled] = useState(false);
+  const [, setTick] = useState(0);
+
+  useEffect(() => {
+    if (!running) return undefined;
+    lastActivityRef.current = Date.now();
+    return undefined;
+  }, [running, fingerprint]);
+
+  useEffect(() => {
+    if (!running) {
+      setStalled(false);
+      return undefined;
+    }
+
+    const sinceActivity = Date.now() - lastActivityRef.current;
+    if (sinceActivity >= thresholdMs) {
+      setStalled(true);
+      return undefined;
+    }
+
+    setStalled(false);
+    const id = setTimeout(() => setStalled(true), thresholdMs - sinceActivity);
+    return () => clearTimeout(id);
+  }, [running, fingerprint, thresholdMs]);
+
+  useEffect(() => {
+    if (!stalled) return undefined;
+    const id = setInterval(() => setTick((t) => t + 1), 1000);
+    return () => clearInterval(id);
+  }, [stalled]);
+
+  if (!stalled) return { stalled: false, stalledForMs: 0 };
+  return {
+    stalled: true,
+    stalledForMs: Math.max(0, Date.now() - lastActivityRef.current),
+  };
+}