diff --git a/.changeset/chat-client-stream-speed.md b/.changeset/chat-client-stream-speed.md new file mode 100644 index 000000000..53d76d8eb --- /dev/null +++ b/.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. diff --git a/docs/chat/streaming.md b/docs/chat/streaming.md index 3254bf6e8..cd36fb82a 100644 --- a/docs/chat/streaming.md +++ b/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: diff --git a/packages/ai-client/src/chat-client.ts b/packages/ai-client/src/chat-client.ts index fc0299c69..475e62a3c 100644 --- a/packages/ai-client/src/chat-client.ts +++ b/packages/ai-client/src/chat-client.ts @@ -84,6 +84,20 @@ interface InternalQueuedMessage extends QueuedMessage { body?: Record } +const STREAM_PROCESSING_BUDGET_MS = 8 + +type SchedulerWithYield = { + yield?: () => Promise +} + +function yieldToHost(): Promise { + 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> @@ -430,6 +444,8 @@ export class ChatClient< private continuationPending = false private subscriptionAbortController: AbortController | null = null private processingResolve: (() => void) | null = null + private chunkProcessingTime = 0 + private chunkProcessingYield: Promise | null = null /** * `connect()` adapters push the full HTTP body into the subscribe queue, then * wait until that queue is idle. After `send()` returns, every chunk from this @@ -1692,14 +1708,42 @@ export class ChatClient< }) } - /** - * Consume chunks from the connection subscription. - */ private async consumeSubscription(signal: AbortSignal): Promise { - 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, + signal: AbortSignal, + beforeProcess?: (chunk: StreamChunk) => void, + ): Promise { 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 + } + } + } } } @@ -1722,9 +1766,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 @@ -1753,18 +1794,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. @@ -1832,10 +1875,7 @@ export class ChatClient< } } - private async processIncomingChunk( - chunk: StreamChunk, - options?: { defer?: boolean }, - ): Promise { + private processIncomingChunk(chunk: StreamChunk): void { chunk = restoreInboundChunk(chunk) if ( chunk.type === 'RUN_ERROR' && @@ -1867,15 +1907,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) } diff --git a/packages/ai-client/tests/chat-client-hidden-tab-yield.test.ts b/packages/ai-client/tests/chat-client-hidden-tab-yield.test.ts deleted file mode 100644 index b44eb0a1e..000000000 --- a/packages/ai-client/tests/chat-client-hidden-tab-yield.test.ts +++ /dev/null @@ -1,47 +0,0 @@ -import { afterEach, describe, expect, it, vi } from 'vitest' -import { ChatClient } from '../src/chat-client' -import { createMockConnectionAdapter, createTextChunks } from './test-utils' - -afterEach(() => { - vi.unstubAllGlobals() -}) - -/** - * Live `processIncomingChunk` either awaits `setTimeout(0)` (visible / Node) - * or skips it (hidden). A marker timer scheduled *before* sendMessage has - * not run yet if sendMessage never awaited a macrotask. - */ -async function streamWithMacrotaskMarker(documentHidden?: boolean) { - if (documentHidden !== undefined) { - vi.stubGlobal('document', { hidden: documentHidden }) - } - - const chunks = createTextChunks('ab') - const client = new ChatClient({ - connection: createMockConnectionAdapter({ chunks }), - }) - let macrotaskRan = false - setTimeout(() => { - macrotaskRan = true - }, 0) - await client.sendMessage('Hi') - return macrotaskRan -} - -describe('ChatClient live yield', () => { - it('does not await a macrotask after live chunks when the page is hidden', async () => { - const macrotaskRan = await streamWithMacrotaskMarker(true) - expect(macrotaskRan).toBe(false) - }) - - it('awaits a macrotask after live chunks when the page is visible', async () => { - const macrotaskRan = await streamWithMacrotaskMarker(false) - expect(macrotaskRan).toBe(true) - }) - - it('awaits a macrotask after live chunks when document is missing', async () => { - expect(typeof document).toBe('undefined') - const macrotaskRan = await streamWithMacrotaskMarker() - expect(macrotaskRan).toBe(true) - }) -}) diff --git a/packages/ai-client/tests/chat-client-stream-processing.test.ts b/packages/ai-client/tests/chat-client-stream-processing.test.ts new file mode 100644 index 000000000..9e42a9c08 --- /dev/null +++ b/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((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() + } + }) +}) diff --git a/packages/ai-client/tests/chat-client.test.ts b/packages/ai-client/tests/chat-client.test.ts index 44a8ce4ce..a546eaa38 100644 --- a/packages/ai-client/tests/chat-client.test.ts +++ b/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([ { diff --git a/packages/ai-client/tests/resume-snapshot.test.ts b/packages/ai-client/tests/resume-snapshot.test.ts index 7f02d8dc9..b0cd1b85a 100644 --- a/packages/ai-client/tests/resume-snapshot.test.ts +++ b/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 = [] + 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 diff --git a/packages/ai-client/tests/test-utils.ts b/packages/ai-client/tests/test-utils.ts index 4e127b312..dec8accbd 100644 --- a/packages/ai-client/tests/test-utils.ts +++ b/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 diff --git a/testing/e2e/src/routeTree.gen.ts b/testing/e2e/src/routeTree.gen.ts index 70cd62134..ceb3861ab 100644 --- a/testing/e2e/src/routeTree.gen.ts +++ b/testing/e2e/src/routeTree.gen.ts @@ -27,6 +27,7 @@ import { Route as DevtoolsRouteARouteImport } from './routes/devtools-route-a' import { Route as DevtoolsMemoryRouteImport } from './routes/devtools-memory' import { Route as DevtoolsGenerationHooksRouteImport } from './routes/devtools-generation-hooks' import { Route as DevtoolsChatRouteImport } from './routes/devtools-chat' +import { Route as ChatClientStreamProcessingRouteImport } from './routes/chat-client-stream-processing' import { Route as ChatClientDefaultBridgeRouteImport } from './routes/chat-client-default-bridge' import { Route as ByokRouteImport } from './routes/byok' import { Route as IndexRouteImport } from './routes/index' @@ -188,6 +189,12 @@ const DevtoolsChatRoute = DevtoolsChatRouteImport.update({ path: '/devtools-chat', getParentRoute: () => rootRouteImport, } as any) +const ChatClientStreamProcessingRoute = + ChatClientStreamProcessingRouteImport.update({ + id: '/chat-client-stream-processing', + path: '/chat-client-stream-processing', + getParentRoute: () => rootRouteImport, + } as any) const ChatClientDefaultBridgeRoute = ChatClientDefaultBridgeRouteImport.update({ id: '/chat-client-default-bridge', path: '/chat-client-default-bridge', @@ -550,6 +557,7 @@ export interface FileRoutesByFullPath { '/': typeof IndexRoute '/byok': typeof ByokRoute '/chat-client-default-bridge': typeof ChatClientDefaultBridgeRoute + '/chat-client-stream-processing': typeof ChatClientStreamProcessingRoute '/devtools-chat': typeof DevtoolsChatRoute '/devtools-generation-hooks': typeof DevtoolsGenerationHooksRoute '/devtools-memory': typeof DevtoolsMemoryRoute @@ -638,6 +646,7 @@ export interface FileRoutesByTo { '/': typeof IndexRoute '/byok': typeof ByokRoute '/chat-client-default-bridge': typeof ChatClientDefaultBridgeRoute + '/chat-client-stream-processing': typeof ChatClientStreamProcessingRoute '/devtools-chat': typeof DevtoolsChatRoute '/devtools-generation-hooks': typeof DevtoolsGenerationHooksRoute '/devtools-memory': typeof DevtoolsMemoryRoute @@ -727,6 +736,7 @@ export interface FileRoutesById { '/': typeof IndexRoute '/byok': typeof ByokRoute '/chat-client-default-bridge': typeof ChatClientDefaultBridgeRoute + '/chat-client-stream-processing': typeof ChatClientStreamProcessingRoute '/devtools-chat': typeof DevtoolsChatRoute '/devtools-generation-hooks': typeof DevtoolsGenerationHooksRoute '/devtools-memory': typeof DevtoolsMemoryRoute @@ -817,6 +827,7 @@ export interface FileRouteTypes { | '/' | '/byok' | '/chat-client-default-bridge' + | '/chat-client-stream-processing' | '/devtools-chat' | '/devtools-generation-hooks' | '/devtools-memory' @@ -905,6 +916,7 @@ export interface FileRouteTypes { | '/' | '/byok' | '/chat-client-default-bridge' + | '/chat-client-stream-processing' | '/devtools-chat' | '/devtools-generation-hooks' | '/devtools-memory' @@ -993,6 +1005,7 @@ export interface FileRouteTypes { | '/' | '/byok' | '/chat-client-default-bridge' + | '/chat-client-stream-processing' | '/devtools-chat' | '/devtools-generation-hooks' | '/devtools-memory' @@ -1082,6 +1095,7 @@ export interface RootRouteChildren { IndexRoute: typeof IndexRoute ByokRoute: typeof ByokRoute ChatClientDefaultBridgeRoute: typeof ChatClientDefaultBridgeRoute + ChatClientStreamProcessingRoute: typeof ChatClientStreamProcessingRoute DevtoolsChatRoute: typeof DevtoolsChatRoute DevtoolsGenerationHooksRoute: typeof DevtoolsGenerationHooksRoute DevtoolsMemoryRoute: typeof DevtoolsMemoryRoute @@ -1290,6 +1304,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof DevtoolsChatRouteImport parentRoute: typeof rootRouteImport } + '/chat-client-stream-processing': { + id: '/chat-client-stream-processing' + path: '/chat-client-stream-processing' + fullPath: '/chat-client-stream-processing' + preLoaderRoute: typeof ChatClientStreamProcessingRouteImport + parentRoute: typeof rootRouteImport + } '/chat-client-default-bridge': { id: '/chat-client-default-bridge' path: '/chat-client-default-bridge' @@ -1831,6 +1852,7 @@ const rootRouteChildren: RootRouteChildren = { IndexRoute: IndexRoute, ByokRoute: ByokRoute, ChatClientDefaultBridgeRoute: ChatClientDefaultBridgeRoute, + ChatClientStreamProcessingRoute: ChatClientStreamProcessingRoute, DevtoolsChatRoute: DevtoolsChatRoute, DevtoolsGenerationHooksRoute: DevtoolsGenerationHooksRoute, DevtoolsMemoryRoute: DevtoolsMemoryRoute, diff --git a/testing/e2e/src/routes/chat-client-stream-processing.tsx b/testing/e2e/src/routes/chat-client-stream-processing.tsx new file mode 100644 index 000000000..960974d31 --- /dev/null +++ b/testing/e2e/src/routes/chat-client-stream-processing.tsx @@ -0,0 +1,264 @@ +import { useEffect, useMemo, useRef, useState } from 'react' +import { createFileRoute } from '@tanstack/react-router' +import { EventType } from '@tanstack/ai' +import { useChat } from '@tanstack/ai-react' +import type { StreamChunk } from '@tanstack/ai' +import type { + RunAgentInputContext, + SubscribeConnectionAdapter, +} from '@tanstack/ai-react' + +export const Route = createFileRoute('/chat-client-stream-processing')({ + component: ChatClientStreamProcessingPage, +}) + +const CONTENT_CHUNK_COUNT = 2_000 +const CHUNK_WORK_MS = 0.075 + +type DrainResult = { + contentChunkCount: number + firstContentBeforeUserBlockingTask: boolean + longTaskObserverSupported: boolean + longTaskCount: number + ordered: boolean + userBlockingTaskBeforeRunFinished: boolean +} + +type TaskScheduler = { + postTask: ( + callback: () => void, + options: { priority: 'user-blocking' }, + ) => Promise +} + +declare const scheduler: TaskScheduler + +function createChunks(run: RunAgentInputContext): Array { + const messageId = 'buffered-message' + const timestamp = Date.now() + const chunks: Array = [ + { + type: EventType.RUN_STARTED, + runId: run.runId, + threadId: run.threadId, + timestamp, + }, + { + type: EventType.TEXT_MESSAGE_START, + messageId, + role: 'assistant', + timestamp, + }, + ] + + for (let index = 0; index < CONTENT_CHUNK_COUNT; index++) { + chunks.push({ + type: EventType.TEXT_MESSAGE_CONTENT, + messageId, + delta: String(index), + timestamp, + }) + } + + chunks.push( + { + type: EventType.TEXT_MESSAGE_END, + messageId, + timestamp, + }, + { + type: EventType.RUN_FINISHED, + runId: run.runId, + threadId: run.threadId, + timestamp, + }, + ) + + return chunks +} + +function createBufferedConnection(): SubscribeConnectionAdapter { + const queue: Array = [] + let wake: (() => void) | undefined + + return { + subscribe(signal) { + return (async function* () { + while (!signal?.aborted) { + if (queue.length === 0) { + await new Promise((resolve) => { + const onAbort = () => resolve() + wake = () => { + signal?.removeEventListener('abort', onAbort) + resolve() + } + signal?.addEventListener('abort', onAbort, { once: true }) + }) + wake = undefined + } + + let chunk = queue.shift() + while (chunk !== undefined) { + yield chunk + chunk = queue.shift() + } + } + })() + }, + send(_messages, _data, _signal, run) { + if (run === undefined) { + return Promise.reject(new Error('Missing run context')) + } + queue.push(...createChunks(run)) + wake?.() + return Promise.resolve() + }, + } +} + +function occupyMainThread(): void { + const deadline = performance.now() + CHUNK_WORK_MS + while (performance.now() < deadline) { + // Keep each chunk cheap while making an uninterrupted drain a long task. + } +} + +function ChatClientStreamProcessingPage() { + const connection = useMemo(createBufferedConnection, []) + const userBlockingTaskRan = useRef(false) + const nextContentIndex = useRef(0) + const nextChunkIndex = useRef(0) + const ordered = useRef(true) + const firstContentBeforeUserBlockingTask = useRef( + undefined, + ) + const userBlockingTaskBeforeRunFinished = useRef(false) + const longTasks = useRef>([]) + const [complete, setComplete] = useState(false) + const [error, setError] = useState() + const [hydrated, setHydrated] = useState(false) + const [result, setResult] = useState() + const { isLoading, messages, sendMessage } = useChat({ + connection, + onChunk(chunk) { + occupyMainThread() + + const chunkIndex = nextChunkIndex.current++ + if (chunkIndex === 0) { + ordered.current = + ordered.current && chunk.type === EventType.RUN_STARTED + } else if (chunkIndex === 1) { + ordered.current = + ordered.current && chunk.type === EventType.TEXT_MESSAGE_START + } else if (chunkIndex <= CONTENT_CHUNK_COUNT + 1) { + ordered.current = + ordered.current && chunk.type === EventType.TEXT_MESSAGE_CONTENT + } else if (chunkIndex === CONTENT_CHUNK_COUNT + 2) { + ordered.current = + ordered.current && chunk.type === EventType.TEXT_MESSAGE_END + } else if (chunkIndex === CONTENT_CHUNK_COUNT + 3) { + ordered.current = + ordered.current && chunk.type === EventType.RUN_FINISHED + } else { + ordered.current = false + } + + if (chunk.type === EventType.TEXT_MESSAGE_CONTENT) { + firstContentBeforeUserBlockingTask.current ??= + !userBlockingTaskRan.current + if (chunk.delta !== String(nextContentIndex.current)) { + ordered.current = false + } + nextContentIndex.current++ + } else if (chunk.type === EventType.RUN_FINISHED) { + userBlockingTaskBeforeRunFinished.current = userBlockingTaskRan.current + } + }, + }) + const assistantText = messages + .filter((message) => message.role === 'assistant') + .flatMap((message) => + message.parts.flatMap((part) => + part.type === 'text' ? [part.content] : [], + ), + ) + .join('') + + useEffect(() => { + setHydrated(true) + }, []) + + const run = async () => { + setComplete(false) + setError(undefined) + setResult(undefined) + userBlockingTaskRan.current = false + nextContentIndex.current = 0 + nextChunkIndex.current = 0 + ordered.current = true + firstContentBeforeUserBlockingTask.current = undefined + userBlockingTaskBeforeRunFinished.current = false + longTasks.current = [] + + const observer = new PerformanceObserver((list) => { + longTasks.current.push(...list.getEntries()) + }) + observer.observe({ type: 'longtask', buffered: true }) + const startedAt = performance.now() + void scheduler.postTask( + () => { + userBlockingTaskRan.current = true + }, + { priority: 'user-blocking' }, + ) + + try { + await sendMessage('Drain buffered chunks') + const finishedAt = performance.now() + await new Promise((resolve) => + requestAnimationFrame(() => setTimeout(resolve, 0)), + ) + setResult({ + contentChunkCount: nextContentIndex.current, + firstContentBeforeUserBlockingTask: + firstContentBeforeUserBlockingTask.current ?? false, + longTaskObserverSupported: + PerformanceObserver.supportedEntryTypes.includes('longtask'), + longTaskCount: longTasks.current.filter( + (entry) => + entry.startTime < finishedAt && + entry.startTime + entry.duration > startedAt, + ).length, + ordered: + ordered.current && nextChunkIndex.current === CONTENT_CHUNK_COUNT + 4, + userBlockingTaskBeforeRunFinished: + userBlockingTaskBeforeRunFinished.current, + }) + } catch (runError) { + setError(String(runError)) + } finally { + observer.disconnect() + setComplete(true) + } + } + + return ( +
+ + {String(isLoading)} + {String(complete)} + {assistantText} + {result !== undefined && ( + {JSON.stringify(result)} + )} + {error !== undefined && {error}} +
+ ) +} diff --git a/testing/e2e/tests/chat-client-stream-processing.spec.ts b/testing/e2e/tests/chat-client-stream-processing.spec.ts new file mode 100644 index 000000000..7d9633012 --- /dev/null +++ b/testing/e2e/tests/chat-client-stream-processing.spec.ts @@ -0,0 +1,25 @@ +import { expect, test } from '@playwright/test' + +test('ChatClient time-slices a large buffered stream', async ({ page }) => { + await page.goto('/chat-client-stream-processing') + + await page.getByTestId('run').click() + await expect(page.getByTestId('complete')).toHaveText('true') + + const result = JSON.parse( + (await page.getByTestId('result').textContent()) ?? '{}', + ) + expect(result).toEqual({ + contentChunkCount: 2_000, + firstContentBeforeUserBlockingTask: true, + longTaskObserverSupported: true, + longTaskCount: 0, + ordered: true, + userBlockingTaskBeforeRunFinished: true, + }) + await expect(page.getByTestId('assistant-text')).toHaveText( + Array.from({ length: 2_000 }, (_, index) => String(index)).join(''), + ) + await expect(page.getByTestId('loading')).toHaveText('false') + await expect(page.getByTestId('error')).toHaveCount(0) +})