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>>
@@ -416,6 +430,8 @@ export class ChatClient<
   private continuationPending = false
   private subscriptionAbortController: AbortController | null = null
   private processingResolve: (() => void) | null = null
+  private chunkProcessingTime = 0
+  private chunkProcessingYield: Promise<void> | null = null
   private errorReportedGeneration: number | null = null
   private streamGeneration = 0
   private continuationGeneration = 0
@@ -1668,14 +1684,42 @@ export class ChatClient<
       })
   }
 
-  /**
-   * Consume chunks from the connection subscription.
-   */
   private async consumeSubscription(signal: AbortSignal): Promise<void> {
-    const stream = this.connection.subscribe(signal)
+    await this.consumeChunks(this.connection.subscribe(signal), signal)
+  }
+
+  /** Consume chunks in order against the client-wide processing budget. */
+  private async consumeChunks(
+    stream: AsyncIterable<StreamChunk>,
+    signal: AbortSignal,
+    beforeProcess?: (chunk: StreamChunk) => void,
+  ): Promise<void> {
     for await (const chunk of stream) {
       if (signal.aborted) break
-      await this.processIncomingChunk(chunk)
+      const pendingYield = this.chunkProcessingYield
+      if (pendingYield) {
+        await pendingYield
+        if (signal.aborted) break
+      }
+      const startedAt = performance.now()
+      beforeProcess?.(chunk)
+      this.processIncomingChunk(chunk)
+      this.chunkProcessingTime += performance.now() - startedAt
+      if (
+        this.chunkProcessingTime >= STREAM_PROCESSING_BUDGET_MS &&
+        (typeof document === 'undefined' || !document.hidden)
+      ) {
+        this.chunkProcessingTime = 0
+        const processingYield = yieldToHost()
+        this.chunkProcessingYield = processingYield
+        try {
+          await processingYield
+        } finally {
+          if (this.chunkProcessingYield === processingYield) {
+            this.chunkProcessingYield = null
+          }
+        }
+      }
     }
   }
 
@@ -1698,9 +1742,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
@@ -1729,18 +1770,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.
@@ -1808,10 +1851,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' &&
@@ -1843,15 +1883,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,140 @@
+import { afterEach, describe, expect, it, vi } from 'vitest'
+import { ChatClient } from '../src/chat-client'
+import { createMockConnectionAdapter, createTextChunks } from './test-utils'
+import type { StreamChunk } from '@tanstack/ai/client'
+
+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()
+  })
+
+  it('shares the processing budget across live and joined streams', async () => {
+    let releaseYield!: () => void
+    const schedulerYield = vi.fn(
+      () =>
+        new Promise<void>((resolve) => {
+          releaseYield = resolve
+        }),
+    )
+    vi.stubGlobal('scheduler', { yield: schedulerYield })
+    let time = 0
+    vi.spyOn(performance, 'now').mockImplementation(() => (time += 5))
+    const processed = vi.fn()
+    const chunk = (name: string): StreamChunk => ({
+      type: 'CUSTOM',
+      name,
+      timestamp: Date.now(),
+      value: null,
+    })
+    const client = new ChatClient({
+      threadId: 't1',
+      connection: {
+        subscribe: async function* () {
+          yield chunk('live-1')
+          yield chunk('live-2')
+        },
+        send: () => Promise.resolve(),
+        joinRun: async function* () {
+          yield chunk('joined')
+        },
+      },
+      initialResumeSnapshot: {
+        resumeState: { threadId: 't1', runId: 'r1' },
+      },
+      onChunk: processed,
+    })
+
+    client.subscribe()
+    client.attach()
+    try {
+      await vi.waitFor(() => expect(schedulerYield).toHaveBeenCalledTimes(1))
+      expect(processed).toHaveBeenCalledTimes(2)
+
+      releaseYield()
+      await vi.waitFor(() => expect(processed).toHaveBeenCalledTimes(3))
+      expect(schedulerYield).toHaveBeenCalledTimes(1)
+    } finally {
+      client.dispose()
+    }
+  })
+})

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

@@ -1636,6 +1636,22 @@ describe('ChatClient', () => {
       expect(client.getConnectionStatus()).toBe('error')
     })
 
+    it('should expose connectionStatus error when subscribe throws', async () => {
+      const connection = {
+        subscribe() {
+          throw new Error('subscription failed')
+        },
+        send: async () => {},
+      }
+      const client = new ChatClient({ connection })
+
+      expect(() => client.subscribe()).not.toThrow()
+      await vi.waitFor(() => {
+        expect(client.getIsSubscribed()).toBe(false)
+        expect(client.getConnectionStatus()).toBe('error')
+      })
+    })
+
     it('should remain pending without terminal run events', async () => {
       const adapter = createSubscribeAdapter([
         {