diff --git a/.changeset/calm-chats-throttle.md b/.changeset/calm-chats-throttle.md new file mode 100644 index 000000000000..055b7fecde84 --- /dev/null +++ b/.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. diff --git a/examples/ai-e2e-next/app/api/chat/throttle/route.ts b/examples/ai-e2e-next/app/api/chat/throttle/route.ts index 641c7a70f161..a74a0979b0c7 100644 --- a/examples/ai-e2e-next/app/api/chat/throttle/route.ts +++ b/examples/ai-e2e-next/app/api/chat/throttle/route.ts @@ -3,8 +3,8 @@ 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', @@ -12,7 +12,19 @@ export async function POST(req: Request) { { 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', }, diff --git a/examples/ai-e2e-next/app/chat/throttle/page.tsx b/examples/ai-e2e-next/app/chat/throttle/page.tsx index 3f50ba3d0650..3b0f4d153c95 100644 --- a/examples/ai-e2e-next/app/chat/throttle/page.tsx +++ b/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(); + const [hasMounted, setHasMounted] = useState(false); + const [, forceUnrelatedRender] = useState(0); + const startedAt = useRef(); + + 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 (

- useChat throttle example + useChat throttle snapshot reproduction

- {messages.map(m => ( -
- {m.role === 'user' ? 'User: ' : 'AI: '} - {m.parts - .map(part => (part.type === 'text' ? part.text : '')) - .join('')} +

+ Streams 500 chunks with a 50ms throttle while an unrelated timer + re-renders the component. Message snapshots should only change within + the throttle cadence. +

+ +
+
Status
+
{status}
+
React renders
+
+ {hasMounted ? renderCount.current : '—'} +
+
Message snapshot changes
+
+ {snapshotChangeCount.current} +
+
Assistant characters
+
{assistantText.length}
+
+ + {result != null && ( +
+ {passed ? 'PASS' : 'FAIL'}: 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}.
- ))} + )} + + {error != null && ( + 
+          {error.message}
+
+ )} - sendMessage({ text })} /> +
); } diff --git a/packages/react/src/use-chat.ts b/packages/react/src/use-chat.ts index e01082e92520..6c8c72c1f2f3 100644 --- a/packages/react/src/use-chat.ts +++ b/packages/react/src/use-chat.ts @@ -128,24 +128,81 @@ export function useChat({ 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( diff --git a/packages/react/src/use-chat.ui.test.tsx b/packages/react/src/use-chat.ui.test.tsx index c32aeffad0b0..bf56810a853c 100644 --- a/packages/react/src/use-chat.ui.test.tsx +++ b/packages/react/src/use-chat.ui.test.tsx @@ -5,7 +5,13 @@ import { TestResponseController, } from '@ai-sdk/test-server/with-vitest'; import { mockId } from '@ai-sdk/provider-utils/test'; -import { cleanup, screen, waitFor, render } from '@testing-library/react'; +import { + cleanup, + fireEvent, + screen, + waitFor, + render, +} from '@testing-library/react'; import userEvent from '@testing-library/user-event'; import { DefaultChatTransport, @@ -2283,15 +2289,28 @@ describe('stop', () => { describe('experimental_throttle', () => { const throttleMs = 50; + beforeEach(() => { + vi.useFakeTimers(); + vi.setSystemTime(new Date('2024-01-01T00:00:00Z')); + }); + + afterEach(() => { + vi.useRealTimers(); + }); + setupTestComponent(() => { - const { messages, sendMessage, status } = useChat({ + const [id, setId] = useState('first-id'); + const { error, messages, sendMessage, status, stop } = useChat({ + id, experimental_throttle: throttleMs, generateId: mockId(), }); + const [, forceUnrelatedRender] = useState(0); return (
{status.toString()}
+ {error != null &&
{error.message}
} {messages.map((m, idx) => (
{m.role === 'user' ? 'User: ' : 'AI: '} @@ -2306,6 +2325,12 @@ describe('experimental_throttle', () => { sendMessage({ parts: [{ text: 'hi', type: 'text' }] }); }} /> +
); }); @@ -2318,11 +2343,9 @@ describe('experimental_throttle', () => { controller, }; - await userEvent.click(screen.getByTestId('do-send')); + fireEvent.click(screen.getByTestId('do-send')); expect(screen.getByTestId('message-0')).toHaveTextContent('User: hi'); - vi.useFakeTimers(); - controller.write(formatChunk({ type: 'text-start', id: '0' })); controller.write( formatChunk({ type: 'text-delta', id: '0', delta: 'Hel' }), @@ -2353,8 +2376,146 @@ describe('experimental_throttle', () => { expect(screen.getByTestId('message-1')).toHaveTextContent( 'AI: Hello There', ); + }); - vi.useRealTimers(); + it('should not publish a new message snapshot during an unrelated render', async () => { + const controller = new TestResponseController(); + + server.urls['/api/chat'].response = { + type: 'controlled-stream', + controller, + }; + + fireEvent.click(screen.getByTestId('do-send')); + + controller.write(formatChunk({ type: 'text-start', id: '0' })); + controller.write( + formatChunk({ type: 'text-delta', id: '0', delta: 'Hel' }), + ); + await act(async () => { + await vi.advanceTimersByTimeAsync(throttleMs + 10); + }); + + expect(screen.getByTestId('message-1')).toHaveTextContent('AI: Hel'); + + controller.write(formatChunk({ type: 'text-delta', id: '0', delta: 'lo' })); + fireEvent.click(screen.getByTestId('force-unrelated-render')); + + expect(screen.getByTestId('message-1')).toHaveTextContent('AI: Hel'); + + await act(async () => { + await vi.advanceTimersByTimeAsync(throttleMs + 10); + }); + + expect(screen.getByTestId('message-1')).toHaveTextContent('AI: Hello'); + }); + + it('should publish the final message snapshot with ready status', async () => { + const controller = new TestResponseController(); + + server.urls['/api/chat'].response = { + type: 'controlled-stream', + controller, + }; + + fireEvent.click(screen.getByTestId('do-send')); + controller.write(formatChunk({ type: 'text-start', id: '0' })); + controller.write( + formatChunk({ type: 'text-delta', id: '0', delta: 'Hello' }), + ); + controller.write(formatChunk({ type: 'text-end', id: '0' })); + controller.close(); + + await act(async () => { + await vi.advanceTimersByTimeAsync(0); + }); + + expect(screen.getByTestId('status')).toHaveTextContent('ready'); + expect(screen.getByTestId('message-1')).toHaveTextContent('AI: Hello'); + }); + + it('should publish the latest message snapshot with error status', async () => { + const controller = new TestResponseController(); + + server.urls['/api/chat'].response = { + type: 'controlled-stream', + controller, + }; + + fireEvent.click(screen.getByTestId('do-send')); + controller.write(formatChunk({ type: 'text-start', id: '0' })); + controller.write( + formatChunk({ type: 'text-delta', id: '0', delta: 'Hello' }), + ); + controller.write( + formatChunk({ type: 'error', errorText: 'stream failed' }), + ); + controller.close(); + + await act(async () => { + await vi.advanceTimersByTimeAsync(0); + }); + + expect(screen.getByTestId('status')).toHaveTextContent('error'); + expect(screen.getByTestId('error')).toHaveTextContent('stream failed'); + expect(screen.getByTestId('message-1')).toHaveTextContent('AI: Hello'); + }); + + it('should publish the latest message snapshot when an abort becomes ready', async () => { + const controller = new TestResponseController(); + + server.urls['/api/chat'].response = { + type: 'controlled-stream', + controller, + }; + + fireEvent.click(screen.getByTestId('do-send')); + await act(async () => { + await controller.write(formatChunk({ type: 'text-start', id: '0' })); + await controller.write( + formatChunk({ type: 'text-delta', id: '0', delta: 'Hello' }), + ); + }); + + expect(screen.queryByTestId('message-1')).not.toBeInTheDocument(); + + await act(async () => { + fireEvent.click(screen.getByTestId('stop')); + await vi.advanceTimersByTimeAsync(0); + }); + + expect(screen.getByTestId('status')).toHaveTextContent('ready'); + expect(screen.getByTestId('message-1')).toHaveTextContent('AI: Hello'); + }); + + it('should ignore a delayed publication after changing chats', async () => { + const controller = new TestResponseController(); + + server.urls['/api/chat'].response = { + type: 'controlled-stream', + controller, + }; + + fireEvent.click(screen.getByTestId('do-send')); + await act(async () => { + await controller.write(formatChunk({ type: 'text-start', id: '0' })); + await controller.write( + formatChunk({ type: 'text-delta', id: '0', delta: 'Hello' }), + ); + }); + + expect(screen.getByTestId('status')).toHaveTextContent('streaming'); + expect(screen.queryByTestId('message-1')).not.toBeInTheDocument(); + + fireEvent.click(screen.getByTestId('change-chat')); + expect(screen.queryByTestId('message-0')).not.toBeInTheDocument(); + + await act(async () => { + await vi.advanceTimersByTimeAsync(throttleMs + 10); + }); + + expect(screen.queryByTestId('message-0')).not.toBeInTheDocument(); + controller.close(); }); });