From c9fdf689cb01fbb90ece0c09cb40cde8cef1c9d5 Mon Sep 17 00:00:00 2001 From: Jack Herrington Date: Fri, 30 Oct 2026 14:59:03 -0700 Subject: [PATCH] fix(ai-client): coalesce the per-chunk snapshot storm while streaming A large tool-call stream re-rendered every useChat hook once per chunk. A Code Mode program arrives as thousands of TOOL_CALL_ARGS deltas in one burst, so the client published each delta to the snapshot atom synchronously and the tab locked up (the SSE read runs on the same thread and could abort). Message publishes made while the client processes a chunk now coalesce to about one per animation frame (setTimeout fallback off-DOM). Discrete changes (setMessages, clear, reload, status) stay synchronous, and getSnapshot() flushes any pending streamed messages so it is never stale. In development, the devtools bridge's per-event snapshot emit is likewise coalesced to one per microtask. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../coalesce-streaming-snapshot-storm.md | 5 ++ packages/ai-client/src/chat-client.ts | 73 ++++++++++++++++--- packages/ai-client/src/devtools-noop.ts | 1 + packages/ai-client/src/devtools.ts | 19 ++++- packages/ai-client/src/snapshot-atom.ts | 14 ++++ packages/ai-client/tests/chat-client.test.ts | 59 ++++++++++++++- packages/ai-client/tests/devtools.test.ts | 11 ++- 7 files changed, 169 insertions(+), 13 deletions(-) create mode 100644 .changeset/coalesce-streaming-snapshot-storm.md 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',