diff --git a/.changeset/fix-use-chat-messages-ref-staleness.md b/.changeset/fix-use-chat-messages-ref-staleness.md new file mode 100644 index 00000000..0ac05de7 --- /dev/null +++ b/.changeset/fix-use-chat-messages-ref-staleness.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-react': patch +--- + +Update messagesRef synchronously during render instead of in useEffect to prevent stale messages when ChatClient is recreated diff --git a/packages/typescript/ai-react/src/use-chat.ts b/packages/typescript/ai-react/src/use-chat.ts index b49ac051..9fbaf030 100644 --- a/packages/typescript/ai-react/src/use-chat.ts +++ b/packages/typescript/ai-react/src/use-chat.ts @@ -29,10 +29,9 @@ export function useChat = any>( ) const isFirstMountRef = useRef(true) - // Update ref whenever messages change - useEffect(() => { - messagesRef.current = messages - }, [messages]) + // Update ref synchronously during render so it's always current when useMemo runs. + // A useEffect here would be async and messagesRef could be stale on client recreation. + messagesRef.current = messages // Track current options in a ref to avoid recreating client when options change const optionsRef = useRef>(options) diff --git a/packages/typescript/ai-react/tests/use-chat.test.ts b/packages/typescript/ai-react/tests/use-chat.test.ts index 282c2ae5..f231988b 100644 --- a/packages/typescript/ai-react/tests/use-chat.test.ts +++ b/packages/typescript/ai-react/tests/use-chat.test.ts @@ -1,7 +1,9 @@ import type { ModelMessage } from '@tanstack/ai' -import { waitFor } from '@testing-library/react' +import { act, renderHook, waitFor } from '@testing-library/react' +import { useState } from 'react' import { describe, expect, it, vi } from 'vitest' -import type { UIMessage } from '../src/types' +import type { UIMessage, UseChatOptions } from '../src/types' +import { useChat } from '../src/use-chat' import { createMockConnectionAdapter, createTextChunks, @@ -835,6 +837,64 @@ describe('useChat', () => { }) }) + describe('client recreation', () => { + it('should pass existing messages to new client when id changes in a batched update', async () => { + const connectSpy = vi.fn() + const chunks = createTextChunks('Reply') + const adapter = createMockConnectionAdapter({ + chunks, + onConnect: connectSpy, + }) + + // Control id via state so setMessages and setId are both React + // state updates that get batched into a single render. + const { result } = renderHook(() => { + const [id, setId] = useState('client-A') + const chat = useChat({ connection: adapter, id }) + return { ...chat, switchId: setId } + }) + + const messages: Array = [ + { + id: 'msg-1', + role: 'user', + parts: [{ type: 'text', content: 'Hello' }], + createdAt: new Date(), + }, + { + id: 'msg-2', + role: 'assistant', + parts: [{ type: 'text', content: 'Hi there!' }], + createdAt: new Date(), + }, + ] + + // Batch: set messages AND change id in one render cycle. + // With the useEffect ref pattern, the new ChatClient is created + // with stale (empty) initialMessages because messagesRef hasn't + // been updated yet. + act(() => { + result.current.setMessages(messages) + result.current.switchId('client-B') + }) + + // Send a message through the new client. If the client lost the + // previous messages, the adapter only receives the new message. + await act(async () => { + await result.current.sendMessage('Follow-up') + }) + + await waitFor(() => { + expect(connectSpy).toHaveBeenCalled() + }) + + // The adapter should receive the previous conversation + new message. + const sentMessages = connectSpy.mock.calls[0]![0] as Array + const userMessages = sentMessages.filter((m: any) => m.role === 'user') + expect(userMessages.length).toBeGreaterThanOrEqual(2) + }) + }) + describe('unmount behavior', () => { it('should not update state after unmount', async () => { const chunks = createTextChunks('Response')