diff --git a/.changeset/coalesce-streaming-snapshot-storm.md b/.changeset/coalesce-streaming-snapshot-storm.md new file mode 100644 index 0000000000..0b69a06dbc --- /dev/null +++ b/.changeset/coalesce-streaming-snapshot-storm.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-client': patch +--- + +Stop the per-chunk re-render storm on high-frequency streams. A tool call whose arguments arrive as thousands of `TOOL_CALL_ARGS` deltas in one burst (e.g. Code Mode emitting a large generated program) used to patch the snapshot atom and notify every framework `useChat` hook once per chunk, saturating the main thread. Message updates made while a streamed chunk is being processed now coalesce to roughly one snapshot notification per animation frame (`requestAnimationFrame`, falling back to a macrotask off-DOM), so React/Vue/Solid/Svelte re-render a handful of times instead of thousands. Discrete, API-driven changes (`setMessages`, `clear`, `reload`, status/subscription updates) still notify synchronously, and `getSnapshot()` flushes any pending streamed messages so it is never stale. In development, the devtools bridge's per-event full-snapshot emit is likewise coalesced to one per microtask instead of firing after every streamed event. diff --git a/packages/ai-client/src/chat-client.ts b/packages/ai-client/src/chat-client.ts index 3085d2f526..ccf4cc620d 100644 --- a/packages/ai-client/src/chat-client.ts +++ b/packages/ai-client/src/chat-client.ts @@ -62,6 +62,7 @@ import { freezeSnapshotMessages, freezeSnapshotPart, patchAtom, + scheduleFrame, subscribeAtom, } from './snapshot-atom' import type { Atom } from './snapshot-atom' @@ -809,12 +810,7 @@ export class ChatClient< // its live handle before anyone reads the messages. this.syncSubagentHandles() this.persistor?.notifyMessagesChanged(messages) - patchAtom(this.snapshotAtom, { - messages: freezeSnapshotMessages( - this.frozenMessages, - messages as Array>, - ), - }) + this.onChunkMessagesChange(messages as Array>) this.callbacksRef.current.onMessagesChange(messages) // The bridge only snapshots on status changes, and the panel draws a // child card from the snapshot, so without this the card would update @@ -1833,19 +1829,73 @@ export class ChatClient< /** * Subscribe to UI snapshot changes. Does not fire with the current value; - * read {@link getSnapshot} first. + * read {@link getSnapshot} first. Notifies synchronously on every atom + * change — except the per-chunk message storm, which is coalesced to roughly + * one notification per animation frame (see {@link onChunkMessagesChange}). * * Named separately from {@link subscribe} (the live connection loop). */ subscribeSnapshot = (listener: () => void): (() => void) => subscribeAtom(this.snapshotAtom, listener) + /** Set only while a streamed chunk is being processed (see processIncomingChunk). */ + private processingChunk = false + /** Latest frozen messages not yet published to the atom while coalescing. */ + private pendingFrozenMessages: ReadonlyArray> | null = null + private cancelMessagesFrame: (() => void) | null = null + + /** + * Publish a `messages` change to the snapshot atom. While a chunk is being + * processed (`processingChunk`), a high-frequency stream — e.g. Code Mode + * delivering a program as thousands of `TOOL_CALL_ARGS` deltas in one burst — + * would otherwise patch the atom and notify every framework hook once per + * chunk. Coalesce those to one atom patch (hence one re-render) per frame. + * Discrete changes (setMessages, clear, reload) run outside chunk processing + * and publish synchronously, so their observers still see them at once. + * {@link getSnapshot} flushes any pending messages, so it is never stale. + */ + private onChunkMessagesChange( + messages: ReadonlyArray>, + ): void { + this.pendingFrozenMessages = freezeSnapshotMessages( + this.frozenMessages, + messages as Array>, + ) + if (!this.processingChunk) { + this.flushPendingMessages() + return + } + if (!this.cancelMessagesFrame) { + this.cancelMessagesFrame = scheduleFrame(() => { + this.cancelMessagesFrame = null + this.flushPendingMessages() + }) + } + } + + /** Patch the pending messages into the atom now (fires listeners once). */ + private flushPendingMessages(): void { + if (this.cancelMessagesFrame) { + this.cancelMessagesFrame() + this.cancelMessagesFrame = null + } + if (this.pendingFrozenMessages) { + const messages = this.pendingFrozenMessages + this.pendingFrozenMessages = null + patchAtom(this.snapshotAtom, { messages }) + } + } + /** * Current UI snapshot. The snapshot, `messages`, `queue`, and their entries * are frozen copies. `interruptState` is the interrupt manager's own state. + * Flushes any coalesced streaming messages first, so a synchronous read is + * never behind what the processor has produced. */ - getSnapshot = (): ChatClientSnapshot => - this.snapshotAtom.get() + getSnapshot = (): ChatClientSnapshot => { + if (this.pendingFrozenMessages) this.flushPendingMessages() + return this.snapshotAtom.get() + } private setIsLoading(isLoading: boolean): void { patchAtom(this.snapshotAtom, { isLoading }) @@ -2276,9 +2326,14 @@ export class ChatClient< } this.handingOverPendingToolCalls = chunk.type === 'RUN_FINISHED' && chunk.outcome?.type !== 'interrupt' + // Mark the window in which the processor may fire onMessagesChange per + // chunk, so those message updates coalesce (see onChunkMessagesChange) + // while discrete, API-driven message changes keep notifying synchronously. + this.processingChunk = true try { this.processor.processChunk(chunk) } finally { + this.processingChunk = false this.handingOverPendingToolCalls = false } this.syncSubagentHandles() diff --git a/packages/ai-client/src/devtools-noop.ts b/packages/ai-client/src/devtools-noop.ts index b66019f62d..2f9df1692f 100644 --- a/packages/ai-client/src/devtools-noop.ts +++ b/packages/ai-client/src/devtools-noop.ts @@ -60,6 +60,7 @@ export class NoOpChatDevtoolsBridge { emitRegistered(): void {} emitUpdated(): void {} emitSnapshot(): void {} + scheduleSnapshot(): void {} emitToolsRegistered(): void {} emitRunLifecycle( _eventType: unknown, diff --git a/packages/ai-client/src/devtools.ts b/packages/ai-client/src/devtools.ts index b01a568e99..f61bf583f9 100644 --- a/packages/ai-client/src/devtools.ts +++ b/packages/ai-client/src/devtools.ts @@ -923,6 +923,23 @@ export class ChatDevtoolsBridge extends ClientDevtoolsBridge { + this.snapshotScheduled = false + this.emitSnapshot() + }) + } + // --- Stream / run context API ------------------------------------------- setCurrentStreamId(streamId: string | null): void { @@ -1981,7 +1998,7 @@ class ChatDevtoolsAwareEventEmitter extends DefaultChatClientEventEmitter { private afterEmit(streamId?: string): void { if (streamId) this.helper.recordStreamId(streamId) - this.helper.emitSnapshot() + this.helper.scheduleSnapshot() } // -- methods with run context -------------------------------------------- diff --git a/packages/ai-client/src/snapshot-atom.ts b/packages/ai-client/src/snapshot-atom.ts index e691397392..b938123076 100644 --- a/packages/ai-client/src/snapshot-atom.ts +++ b/packages/ai-client/src/snapshot-atom.ts @@ -49,6 +49,20 @@ export function notifyListeners>( } } +/** + * Run `cb` on the next animation frame, falling back to a macrotask off-DOM + * (SSR/node). Returns a canceller. Used to coalesce a burst of snapshot + * changes into roughly one notification per frame instead of one per chunk. + */ +export function scheduleFrame(cb: () => void): () => void { + if (typeof globalThis.requestAnimationFrame === 'function') { + const handle = globalThis.requestAnimationFrame(cb) + return () => globalThis.cancelAnimationFrame(handle) + } + const handle = setTimeout(cb, 0) + return () => clearTimeout(handle) +} + /** Deep-freeze plain JSON. Class instances stay as they are. */ function freezePlainData(value: unknown): unknown { if (Array.isArray(value)) { diff --git a/packages/ai-client/tests/chat-client.test.ts b/packages/ai-client/tests/chat-client.test.ts index d0de6d518d..cb63afabe1 100644 --- a/packages/ai-client/tests/chat-client.test.ts +++ b/packages/ai-client/tests/chat-client.test.ts @@ -20,6 +20,11 @@ import type { UIMessage, } from '../src/types' +// subscribeSnapshot coalesces notifications to the next frame (rAF, or a +// setTimeout(0) macrotask off-DOM, which is this node test env). Awaiting a +// macrotask lets the scheduled notification fire. +const flushFrame = () => new Promise((resolve) => setTimeout(resolve, 0)) + describe('ChatClient', () => { const persistedMessage: UIMessage = { id: 'persisted-1', @@ -96,6 +101,8 @@ describe('ChatClient', () => { const listener = vi.fn() const stop = client.subscribeSnapshot(listener) + // A single, discrete change notifies synchronously (the leading edge of + // the coalescer), like an uncoalesced subscription. client.setMessagesManually([]) expect(listener).toHaveBeenCalledTimes(1) expect(client.getSnapshot().messages).toEqual([]) @@ -111,6 +118,51 @@ describe('ChatClient', () => { expect(listener).toHaveBeenCalledTimes(1) }) + it('coalesces the per-chunk message storm during streaming', async () => { + // One delta per character — a 40-chunk burst. Each delta fires + // onMessagesChange; without coalescing that is 40 snapshot notifications. + const text = 'The quick brown fox jumps over the lazy!' + const client = new ChatClient({ + threadId: 'thread-1', + connection: createMockConnectionAdapter({ + chunks: createTextChunks(text), + }), + }) + const listener = vi.fn() + const stop = client.subscribeSnapshot(listener) + + await client.sendMessage('go') + await flushFrame() + + // The message updates collapsed to far fewer notifications than chunks, + // and the final text is complete and readable synchronously. + expect(listener.mock.calls.length).toBeLessThan(text.length) + const assistant = client.getSnapshot().messages.at(-1) + expect(assistant?.role).toBe('assistant') + expect( + assistant?.parts.map((p) => ('content' in p ? p.content : '')).join(''), + ).toBe(text) + stop() + }) + + it('publishes a discrete (non-streaming) message change synchronously', () => { + const client = new ChatClient({ + threadId: 'thread-1', + connection: createMockConnectionAdapter(), + }) + const listener = vi.fn() + const stop = client.subscribeSnapshot(listener) + + client.setMessagesManually([ + { id: 'm-1', role: 'user', parts: [{ type: 'text', content: 'hi' }] }, + ]) + + // Not inside chunk processing, so it notifies at once (no coalescing). + expect(listener).toHaveBeenCalledTimes(1) + expect(client.getSnapshot().messages[0]?.id).toBe('m-1') + stop() + }) + it('isolates nested message parts from callers and old snapshots', () => { const adapter = createMockConnectionAdapter() const source = { @@ -285,11 +337,16 @@ describe('ChatClient', () => { }) await client.sendMessage('Hello') + await flushFrame() stop() const messages = client.getSnapshot().messages expect(messages.length).toBeGreaterThan(2) - expect(seen.length).toBeGreaterThan(1) + // Notifications are coalesced (leading edge + one trailing flush per + // frame), so the exact count is not asserted; what matters is that every + // snapshot the listener observed kept the first message's identity (the + // freeze cache reuses unchanged messages). + expect(seen.length).toBeGreaterThan(0) for (const message of seen) { expect(message).toBe(first) } diff --git a/packages/ai-client/tests/devtools.test.ts b/packages/ai-client/tests/devtools.test.ts index eb2eae7e27..8db77e45e1 100644 --- a/packages/ai-client/tests/devtools.test.ts +++ b/packages/ai-client/tests/devtools.test.ts @@ -878,6 +878,9 @@ describe('ChatClient devtools bridge', () => { const duplicateClient = createClient({ threadId: 'shared-thread', }) + // The surviving bridge's mount snapshot is coalesced to a microtask; let it + // flush before clearing so only post-send emissions are asserted below. + await Promise.resolve() vi.clearAllMocks() await firstClient.sendMessage('start') @@ -1044,11 +1047,13 @@ describe('ChatClient devtools bridge', () => { expect(eventClientMock.unsubscribe).toHaveBeenCalledTimes(2) }) - it('emits a snapshot when messages are set manually', () => { + it('emits a snapshot when messages are set manually', async () => { const client = createClient() vi.clearAllMocks() client.setMessagesManually([userMessage]) + // The event-driven snapshot is coalesced to a microtask; let it flush. + await Promise.resolve() expect(aiEventClient.emit).toHaveBeenCalledWith( 'hook:state-snapshot', @@ -1062,12 +1067,14 @@ describe('ChatClient devtools bridge', () => { client.dispose() }) - it('emits a snapshot when messages are cleared', () => { + it('emits a snapshot when messages are cleared', async () => { const client = createClient() client.setMessagesManually([userMessage]) vi.clearAllMocks() client.clear() + // The event-driven snapshot is coalesced to a microtask; let it flush. + await Promise.resolve() expect(aiEventClient.emit).toHaveBeenCalledWith( 'hook:state-snapshot',