fix(react): backport throttled chat snapshots to v6

.changeset/calm-chats-throttle.md

@@ -0,0 +1,5 @@
+---
+'@ai-sdk/react': patch
+---
+
+Fix `useChat` throttling so unrelated React renders cannot publish message snapshots ahead of the configured throttle cadence.

examples/ai-e2e-next/app/api/chat/throttle/route.ts

@@ -3,16 +3,28 @@ import { createUIMessageStreamResponse, simulateReadableStream } from 'ai';
 export async function POST(req: Request) {
   return createUIMessageStreamResponse({
     stream: simulateReadableStream({
-      initialDelayInMs: 0, // Delay before the first chunk
-      chunkDelayInMs: 0, // Delay between chunks
+      initialDelayInMs: 0,
+      chunkDelayInMs: 0,
       chunks: [
         {
           type: 'start',
         },
         {
           type: 'start-step',
         },
-        ...Array(5000).fill({ type: 'text', value: 'T\n' }),
+        {
+          type: 'text-start',
+          id: 'text-1',
+        },
+        ...Array(500).fill({
+          type: 'text-delta',
+          id: 'text-1',
+          delta: 'T\n',
+        }),
+        {
+          type: 'text-end',
+          id: 'text-1',
+        },
         {
           type: 'finish-step',
         },

examples/ai-e2e-next/app/chat/throttle/page.tsx

@@ -1,36 +1,165 @@
 'use client';
 
-import ChatInput from '@/components/chat-input';
 import { useChat } from '@ai-sdk/react';
 import { DefaultChatTransport } from 'ai';
-import { useLayoutEffect, useRef } from 'react';
+import { useEffect, useRef, useState } from 'react';
+
+const THROTTLE_MS = 50;
+const EXPECTED_ASSISTANT_CHARACTERS = 1000;
+
+type Result = {
+  assistantCharacterCountAtReady: number;
+  durationInMs: number;
+  maximumExpectedSnapshotChanges: number;
+  renderCount: number;
+  snapshotChangeCount: number;
+};
 
 export default function Chat() {
   const renderCount = useRef(0);
-  useLayoutEffect(() => {
-    console.log(`component rendered #${++renderCount.current}`);
-  });
+  renderCount.current += 1;
 
-  const { messages, status, sendMessage } = useChat({
+  const { error, messages, status, sendMessage } = useChat({
     transport: new DefaultChatTransport({ api: '/api/chat/throttle' }),
-    experimental_throttle: 50,
+    experimental_throttle: THROTTLE_MS,
   });
 
+  const previousMessages = useRef(messages);
+  const snapshotChangeCount = useRef(0);
+  if (previousMessages.current !== messages) {
+    previousMessages.current = messages;
+    snapshotChangeCount.current += 1;
+  }
+
+  const [result, setResult] = useState<Result>();
+  const [hasMounted, setHasMounted] = useState(false);
+  const [, forceUnrelatedRender] = useState(0);
+  const startedAt = useRef<number>();
+
+  useEffect(() => {
+    setHasMounted(true);
+  }, []);
+
+  const assistantMessages = messages.filter(
+    message => message.role === 'assistant',
+  );
+  const latestAssistantMessage =
+    assistantMessages[assistantMessages.length - 1];
+  const assistantText = (latestAssistantMessage?.parts ?? [])
+    .filter(part => part.type === 'text')
+    .map(part => part.text)
+    .join('');
+
+  useEffect(() => {
+    if (status !== 'submitted' && status !== 'streaming') {
+      return;
+    }
+
+    // Re-render independently from useChat while the response streams. Before
+    // the fix for #6166, these renders read the always-current messages array
+    // from getSnapshot and bypass the throttled subscription.
+    const interval = setInterval(() => {
+      forceUnrelatedRender(count => count + 1);
+    }, 0);
+
+    return () => clearInterval(interval);
+  }, [status]);
+
+  useEffect(() => {
+    if (startedAt.current == null || status !== 'ready') {
+      return;
+    }
+
+    const durationInMs = performance.now() - startedAt.current;
+    setResult({
+      assistantCharacterCountAtReady: assistantText.length,
+      durationInMs,
+      // Account for the leading update, trailing update, user message, and
+      // timer/commit boundary variance around the throttle window.
+      maximumExpectedSnapshotChanges: Math.ceil(durationInMs / THROTTLE_MS) + 4,
+      renderCount: renderCount.current,
+      snapshotChangeCount: snapshotChangeCount.current,
+    });
+    startedAt.current = undefined;
+  }, [assistantText.length, status]);
+
+  const runReproduction = () => {
+    previousMessages.current = messages;
+    renderCount.current = 0;
+    snapshotChangeCount.current = 0;
+    setResult(undefined);
+    startedAt.current = performance.now();
+    sendMessage({ text: 'Run the throttle snapshot reproduction' });
+  };
+
+  const passed =
+    result != null &&
+    result.snapshotChangeCount <= result.maximumExpectedSnapshotChanges &&
+    result.assistantCharacterCountAtReady === EXPECTED_ASSISTANT_CHARACTERS;
+
   return (
     <div className="flex flex-col w-full max-w-md py-24 mx-auto stretch">
       <h4 className="pb-4 text-xl font-bold text-gray-900 md:text-xl">
-        useChat throttle example
+        useChat throttle snapshot reproduction
       </h4>
-      {messages.map(m => (
-        <div key={m.id} className="whitespace-pre-wrap">
-          {m.role === 'user' ? 'User: ' : 'AI: '}
-          {m.parts
-            .map(part => (part.type === 'text' ? part.text : ''))
-            .join('')}
+      <p className="pb-4 text-sm text-gray-700">
+        Streams 500 chunks with a 50ms throttle while an unrelated timer
+        re-renders the component. Message snapshots should only change within
+        the throttle cadence.
+      </p>
+
+      <dl className="grid grid-cols-2 gap-2 pb-4 text-sm">
+        <dt>Status</dt>
+        <dd data-testid="status">{status}</dd>
+        <dt>React renders</dt>
+        <dd data-testid="render-count">
+          {hasMounted ? renderCount.current : '—'}
+        </dd>
+        <dt>Message snapshot changes</dt>
+        <dd data-testid="snapshot-change-count">
+          {snapshotChangeCount.current}
+        </dd>
+        <dt>Assistant characters</dt>
+        <dd data-testid="assistant-character-count">{assistantText.length}</dd>
+      </dl>
+
+      {result != null && (
+        <div
+          className={`mb-4 rounded border p-3 ${
+            passed
+              ? 'border-green-600 bg-green-50 text-green-900'
+              : 'border-red-600 bg-red-50 text-red-900'
+          }`}
+          data-testid="result"
+        >
+          <strong>{passed ? 'PASS' : 'FAIL'}</strong>: observed{' '}
+          {result.snapshotChangeCount} message snapshot changes in{' '}
+          {Math.round(result.durationInMs)}ms; expected at most{' '}
+          {result.maximumExpectedSnapshotChanges} with a {THROTTLE_MS}ms
+          throttle. Assistant characters when status became ready:{' '}
+          {result.assistantCharacterCountAtReady}/
+          {EXPECTED_ASSISTANT_CHARACTERS}. Total renders: {result.renderCount}.
         </div>
-      ))}
+      )}
+
+      {error != null && (
+        <pre
+          className="mb-4 whitespace-pre-wrap text-red-700"
+          data-testid="error"
+        >
+          {error.message}
+        </pre>
+      )}
 
-      <ChatInput status={status} onSubmit={text => sendMessage({ text })} />
+      <button
+        className="rounded bg-black px-4 py-2 text-white disabled:opacity-50"
+        data-testid="run-reproduction"
+        disabled={status !== 'ready' || startedAt.current != null}
+        onClick={runReproduction}
+        type="button"
+      >
+        Run reproduction
+      </button>
     </div>
   );
 }

packages/react/src/use-chat.ts

@@ -128,24 +128,81 @@ export function useChat<UI_MESSAGE extends UIMessage = UIMessage>({
     chatRef.current = 'chat' in options ? options.chat : new Chat(chatOptions);
   }
 
+  const chat = chatRef.current;
+  const messagesSnapshotRef = useRef({
+    chat,
+    messages: chat.messages,
+  });
+
+  if (messagesSnapshotRef.current.chat !== chat) {
+    messagesSnapshotRef.current = { chat, messages: chat.messages };
+  }
+
   const subscribeToMessages = useCallback(
-    (update: () => void) =>
-      chatRef.current['~registerMessagesCallback'](update, throttleWaitMs),
-    // `chatRef.current.id` is required to trigger re-subscription when the chat ID changes
-    // eslint-disable-next-line react-hooks/exhaustive-deps
-    [throttleWaitMs, chatRef.current.id],
+    (update: () => void) => {
+      let isSubscribed = true;
+
+      const updateMessages = () => {
+        if (!isSubscribed || messagesSnapshotRef.current.chat !== chat) {
+          return;
+        }
+
+        messagesSnapshotRef.current = { chat, messages: chat.messages };
+        update();
+      };
+
+      const unsubscribe = chat['~registerMessagesCallback'](
+        updateMessages,
+        throttleWaitMs,
+      );
+
+      // Synchronize changes that may have happened between render and
+      // subscription. useSyncExternalStore checks the snapshot after
+      // subscribing and schedules a render when it changed.
+      messagesSnapshotRef.current = { chat, messages: chat.messages };
+
+      return () => {
+        isSubscribed = false;
+        unsubscribe();
+      };
+    },
+    [chat, throttleWaitMs],
+  );
+
+  const getMessagesSnapshot = useCallback(
+    () => messagesSnapshotRef.current.messages,
+    [],
   );
 
   const messages = useSyncExternalStore(
     subscribeToMessages,
-    () => chatRef.current.messages,
-    () => chatRef.current.messages,
+    getMessagesSnapshot,
+    getMessagesSnapshot,
   );
 
+  const subscribeToStatus = useCallback(
+    (update: () => void) =>
+      chat['~registerStatusCallback'](() => {
+        if (messagesSnapshotRef.current.chat !== chat) {
+          return;
+        }
+
+        if (chat.status === 'ready' || chat.status === 'error') {
+          // Publish the latest messages before the terminal status can render.
+          messagesSnapshotRef.current = { chat, messages: chat.messages };
+        }
+
+        update();
+      }),
+    [chat],
+  );
+
+  const getStatusSnapshot = useCallback(() => chat.status, [chat]);
+
   const status = useSyncExternalStore(
-    chatRef.current['~registerStatusCallback'],
-    () => chatRef.current.status,
-    () => chatRef.current.status,
+    subscribeToStatus,
+    getStatusSnapshot,
+    getStatusSnapshot,
   );
 
   const error = useSyncExternalStore(