fix(ai-client): avoid per-chunk stream scheduling delays

.changeset/chat-client-stream-speed.md

@@ -0,0 +1,5 @@
+---
+'@tanstack/ai-client': patch
+---
+
+Process live chat chunks without waiting for a separate macrotask after each chunk.

docs/chat/streaming.md

@@ -77,6 +77,8 @@ messages.forEach((message) => {
 });
 ```
 
+Across every framework integration, the shared `ChatClient` processes ready live chunks in order without inserting a task between every chunk. After a bounded amount of chunk-processing work, it yields to keep the main thread responsive before continuing.
+
 ## Stream Events (AG-UI Protocol)
 
 TanStack AI implements the [AG-UI Protocol](https://docs.ag-ui.com/introduction) for streaming. Stream events contain different types of data:

packages/ai-client/src/chat-client.ts

@@ -84,6 +84,20 @@ interface InternalQueuedMessage extends QueuedMessage {
   body?: Record<string, any>
 }
 
+const STREAM_PROCESSING_BUDGET_MS = 8
+
+type SchedulerWithYield = {
+  yield?: () => Promise<void>
+}
+
+function yieldToHost(): Promise<void> {
+  const { scheduler } = globalThis as typeof globalThis & {
+    scheduler?: SchedulerWithYield
+  }
+  if (scheduler?.yield) return scheduler.yield()
+  return new Promise((resolve) => setTimeout(resolve, 0))
+}
+
 function assertUniqueInterruptDefinitions(
   interrupts:
     | ReadonlyArray<InterruptDefinition<any, any, any, any>>
@@ -1632,14 +1646,30 @@ export class ChatClient<
       })
   }
 
-  /**
-   * Consume chunks from the connection subscription.
-   */
-  private async consumeSubscription(signal: AbortSignal): Promise<void> {
-    const stream = this.connection.subscribe(signal)
+  private consumeSubscription(signal: AbortSignal): Promise<void> {
+    return this.consumeChunks(this.connection.subscribe(signal), signal)
+  }
+
+  /** Consume chunks in order, yielding after bounded processing work. */
+  private async consumeChunks(
+    stream: AsyncIterable<StreamChunk>,
+    signal: AbortSignal,
+    beforeProcess?: (chunk: StreamChunk) => void,
+  ): Promise<void> {
+    let processingTime = 0
     for await (const chunk of stream) {
       if (signal.aborted) break
-      await this.processIncomingChunk(chunk)
+      beforeProcess?.(chunk)
+      const startedAt = performance.now()
+      this.processIncomingChunk(chunk)
+      processingTime += performance.now() - startedAt
+      if (
+        processingTime >= STREAM_PROCESSING_BUDGET_MS &&
+        (typeof document === 'undefined' || !document.hidden)
+      ) {
+        await yieldToHost()
+        processingTime = 0
+      }
     }
   }
 
@@ -1662,9 +1692,6 @@ export class ChatClient<
    * give up after {@link REJOIN_CONNECT_DEADLINE_MS} if no chunk arrives and
    * clear the dead pointer so it does not retry on the next load.
    *
-   * Replay chunks are processed WITHOUT the per-chunk yield the live path uses,
-   * so the buffered prefix snaps in and only the genuinely-live tail streams at
-   * network speed — a reload looks like the run continued, not like it re-typed.
    */
   private resumeInFlightRun(runId: string): void {
     const joinRun = this.connection.joinRun
@@ -1692,18 +1719,20 @@ export class ChatClient<
         if (!attached) controller.abort()
       }, REJOIN_CONNECT_DEADLINE_MS)
       try {
-        for await (const chunk of joinRun(runId, controller.signal)) {
-          if (controller.signal.aborted) break
-          if (!attached) {
-            attached = true
-            clearTimeout(connectTimer)
-          }
-          if (!rebuilt && REJOIN_REBUILD_TRIGGERS.has(chunk.type)) {
-            rebuilt = true
-            this.dropTrailingInFlightAssistant()
-          }
-          await this.processIncomingChunk(chunk, { defer: false })
-        }
+        await this.consumeChunks(
+          joinRun(runId, controller.signal),
+          controller.signal,
+          (chunk) => {
+            if (!attached) {
+              attached = true
+              clearTimeout(connectTimer)
+            }
+            if (!rebuilt && REJOIN_REBUILD_TRIGGERS.has(chunk.type)) {
+              rebuilt = true
+              this.dropTrailingInFlightAssistant()
+            }
+          },
+        )
         // Same contract as `streamResponse`: client tools may finish (and
         // queue a resume) while `isLoading` is still true. Wait for them
         // before teardown so `drainPostStreamActions` below sees the queue.
@@ -1771,10 +1800,7 @@ export class ChatClient<
     }
   }
 
-  private async processIncomingChunk(
-    chunk: StreamChunk,
-    options?: { defer?: boolean },
-  ): Promise<void> {
+  private processIncomingChunk(chunk: StreamChunk): void {
     chunk = restoreInboundChunk(chunk)
     if (
       chunk.type === 'RUN_ERROR' &&
@@ -1806,15 +1832,6 @@ export class ChatClient<
     this.processor.processChunk(chunk)
     this.updateRunLifecycle(chunk)
     this.observeInterruptState(chunk)
-    // Live path: yield a macrotask so the UI can paint. Skip when the page is
-    // hidden. Browsers clamp setTimeout there, and that wait paces stream pull.
-    // Replay passes defer: false so a backlog applies in one batch.
-    if (
-      options?.defer !== false &&
-      (typeof document === 'undefined' || !document.hidden)
-    ) {
-      await new Promise((resolve) => setTimeout(resolve, 0))
-    }
     this.resolveJoinedRun(chunk)
   }
 

packages/ai-client/tests/chat-client-stream-processing.test.ts

@@ -0,0 +1,89 @@
+import { afterEach, describe, expect, it, vi } from 'vitest'
+import { ChatClient } from '../src/chat-client'
+import { createMockConnectionAdapter, createTextChunks } from './test-utils'
+
+afterEach(() => {
+  vi.restoreAllMocks()
+  vi.unstubAllGlobals()
+})
+
+describe('ChatClient stream processing', () => {
+  it('does not wait for a macrotask after each live chunk', async () => {
+    vi.spyOn(performance, 'now').mockReturnValue(0)
+    const client = new ChatClient({
+      connection: createMockConnectionAdapter({
+        chunks: createTextChunks('ab'),
+      }),
+    })
+    let macrotaskRan = false
+    setTimeout(() => {
+      macrotaskRan = true
+    }, 0)
+
+    await client.sendMessage('Hi')
+
+    expect(macrotaskRan).toBe(false)
+  })
+
+  it('falls back to a timer after a full processing slice', async () => {
+    vi.stubGlobal('scheduler', {})
+    let time = 0
+    vi.spyOn(performance, 'now').mockImplementation(() => (time += 9))
+    const client = new ChatClient({
+      connection: createMockConnectionAdapter({
+        chunks: createTextChunks('ab'),
+      }),
+    })
+    let macrotaskRan = false
+    setTimeout(() => {
+      macrotaskRan = true
+    }, 0)
+
+    await client.sendMessage('Hi')
+
+    expect(macrotaskRan).toBe(true)
+  })
+
+  it('uses the scheduler after a full processing slice', async () => {
+    const schedulerYield = vi.fn(() => Promise.resolve())
+    vi.stubGlobal('scheduler', { yield: schedulerYield })
+    let time = 0
+    vi.spyOn(performance, 'now').mockImplementation(() => (time += 9))
+    const client = new ChatClient({
+      connection: createMockConnectionAdapter({
+        chunks: createTextChunks('ab'),
+      }),
+    })
+    let macrotaskRan = false
+    setTimeout(() => {
+      macrotaskRan = true
+    }, 0)
+
+    await client.sendMessage('Hi')
+
+    expect(schedulerYield).toHaveBeenCalled()
+    expect(macrotaskRan).toBe(false)
+  })
+
+  it('does not yield in a hidden document', async () => {
+    vi.stubGlobal('document', { hidden: true })
+    const schedulerYield = vi.fn(() => Promise.resolve())
+    vi.stubGlobal('scheduler', { yield: schedulerYield })
+    let time = 0
+    vi.spyOn(performance, 'now').mockImplementation(() => (time += 9))
+    const client = new ChatClient({
+      connection: createMockConnectionAdapter({
+        chunks: createTextChunks('ab'),
+      }),
+    })
+    let macrotaskRan = false
+    setTimeout(() => {
+      macrotaskRan = true
+    }, 0)
+
+    await client.sendMessage('Hi')
+
+    expect(macrotaskRan).toBe(false)
+    expect(schedulerYield).not.toHaveBeenCalled()
+  })
+})

packages/ai-client/tests/resume-snapshot.test.ts

@@ -247,6 +247,48 @@ describe('ChatClient auto-rejoin after reload', () => {
     void client
   })
 
+  it('yields after a full replay processing slice', async () => {
+    const schedulerYield = vi.fn(() => Promise.resolve())
+    vi.stubGlobal('scheduler', { yield: schedulerYield })
+    let time = 0
+    const now = vi
+      .spyOn(performance, 'now')
+      .mockImplementation(() => (time += 9))
+    const joinRun = vi.fn(async function* () {
+      for (const chunk of runChunks('r1', 't1')) {
+        yield chunk
+      }
+    })
+    const connection: ResumableConnectConnectionAdapter = {
+      connect: async function* () {},
+      joinRun,
+    }
+    let latest: Array<UIMessage> = []
+    const client = mountedChatClient({
+      threadId: 't1',
+      connection,
+      initialResumeSnapshot: {
+        resumeState: { threadId: 't1', runId: 'r1' },
+      },
+      onMessagesChange: (messages) => {
+        latest = messages
+      },
+    })
+
+    try {
+      await vi.waitFor(() => {
+        const assistant = latest.find((message) => message.role === 'assistant')
+        const text = assistant?.parts.find((part) => part.type === 'text')
+        expect(text && 'content' in text && text.content).toBe('world')
+      })
+      expect(schedulerYield).toHaveBeenCalled()
+    } finally {
+      client.dispose()
+      now.mockRestore()
+      vi.unstubAllGlobals()
+    }
+  })
+
   it('persistence:true hydrates history AND tails a live run from the server on mount', async () => {
     // Server-authoritative: the client caches no transcript and no run pointer.
     // On mount it calls connection.hydrate(threadId), which returns the stored

packages/ai-client/tests/test-utils.ts

@@ -158,11 +158,8 @@ export function createMockConnectionAdapter(
 /**
  * Subscribe/send adapter that tests can push chunks into at any time.
  *
- * `ChatClient.processIncomingChunk` yields a `setTimeout(0)` after each chunk
- * so React can paint. A test that pushes the next batch during that gap would
- * lose the wake on a naive mock (the generator is not parked, so `wake()` is
- * a no-op, then the generator parks on a new waiter and the chunk sits
- * forever). This helper:
+ * A push between reading the queue and parking the waiter would lose the wake
+ * on a naive mock. This helper:
  * - rechecks the queue after every yielded batch
  * - rechecks again after parking the waiter, so a push in that window still
  *   wakes