fix(ai-client): coalesce the per-chunk snapshot storm while streaming

.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.

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<UIMessage<TTools>>,
-            ),
-          })
+          this.onChunkMessagesChange(messages as Array<UIMessage<TTools>>)
           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<UIMessage<TTools>> | 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<UIMessage<TTools>>,
+  ): void {
+    this.pendingFrozenMessages = freezeSnapshotMessages(
+      this.frozenMessages,
+      messages as Array<UIMessage<TTools>>,
+    )
+    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<TTools, TInterrupts> =>
-    this.snapshotAtom.get()
+  getSnapshot = (): ChatClientSnapshot<TTools, TInterrupts> => {
+    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()

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,

packages/ai-client/src/devtools.ts

@@ -923,6 +923,23 @@ export class ChatDevtoolsBridge extends ClientDevtoolsBridge<AIDevtoolsChatSnaps
     )
   }
 
+  private snapshotScheduled = false
+
+  /**
+   * Coalesced {@link emitSnapshot}: at most one full-state emit per microtask,
+   * however many events arrive. A high-frequency stream fires a per-chunk event
+   * (text/tool-call), each of which would otherwise rebuild and emit the whole
+   * snapshot; this caps that to one rebuild per microtask.
+   */
+  scheduleSnapshot(): void {
+    if (this.snapshotScheduled) return
+    this.snapshotScheduled = true
+    queueMicrotask(() => {
+      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 --------------------------------------------

packages/ai-client/src/snapshot-atom.ts

@@ -49,6 +49,20 @@ export function notifyListeners<TArgs extends Array<unknown>>(
   }
 }
 
+/**
+ * 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)) {

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)
       }

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',