From 2aae92137c9bbeed6ac26f14f491ddbc643f0fed Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Wed, 19 Aug 2026 19:53:02 +0200 Subject: [PATCH] fix(ai-client): keep subscribe test chunks from being dropped Release CI on main failed on should clear all runs on RUN_ERROR without runId. The test mock lost the wake when ChatClient yielded setTimeout(0) between chunks. The RUN_ERROR chunk then never ran, so sessionGenerating stayed true. The new helper rechecks the queue after each batch and after it starts waiting. Tests now wait for the real state instead of a fixed timeout. --- packages/ai-client/tests/chat-client.test.ts | 323 +++++++------------ packages/ai-client/tests/test-utils.ts | 74 ++++- 2 files changed, 189 insertions(+), 208 deletions(-) diff --git a/packages/ai-client/tests/chat-client.test.ts b/packages/ai-client/tests/chat-client.test.ts index 93083dff65..dcb197af11 100644 --- a/packages/ai-client/tests/chat-client.test.ts +++ b/packages/ai-client/tests/chat-client.test.ts @@ -5,6 +5,7 @@ import { createApprovalToolCallChunks, createCustomEventChunks, createMockConnectionAdapter, + createPushableSubscribeConnection, createTextChunks, createThinkingChunks, createToolCallChunks, @@ -1969,41 +1970,31 @@ describe('ChatClient', () => { }) it('should stay true during concurrent runs until all finish', async () => { - const wake = { fn: null as (() => void) | null } - const chunks: Array = [] - const connection = { - subscribe: async function* (signal?: AbortSignal) { - while (!signal?.aborted) { - if (chunks.length > 0) { - const batch = chunks.splice(0) - for (const chunk of batch) { - yield chunk - } - } - await new Promise((resolve) => { - wake.fn = resolve - const onAbort = () => resolve() - signal?.addEventListener('abort', onAbort, { once: true }) - }) - } - }, - send: async () => { - wake.fn?.() - }, - } + const { connection, push } = createPushableSubscribeConnection() const generatingChanges: Array = [] + const finishedRunIds: Array = [] const client = new ChatClient({ connection, + onChunk: (chunk) => { + if ( + chunk.type === EventType.RUN_FINISHED && + 'runId' in chunk && + typeof chunk.runId === 'string' + ) { + finishedRunIds.push(chunk.runId) + } + }, onSessionGeneratingChange: (isGenerating) => { generatingChanges.push(isGenerating) }, }) client.subscribe() - await new Promise((resolve) => setTimeout(resolve, 10)) + await vi.waitFor(() => { + expect(client.getIsSubscribed()).toBe(true) + }) - // Simulate two concurrent runs starting - chunks.push( + push( { type: EventType.RUN_STARTED, runId: 'run-1', @@ -2019,13 +2010,12 @@ describe('ChatClient', () => { timestamp: Date.now(), }, ) - wake.fn?.() - await new Promise((resolve) => setTimeout(resolve, 20)) - expect(client.getSessionGenerating()).toBe(true) + await vi.waitFor(() => { + expect(client.getSessionGenerating()).toBe(true) + }) - // First run finishes — should still be generating because run-2 is active - chunks.push({ + push({ type: EventType.RUN_FINISHED, runId: 'run-1', threadId: 'thread-1', @@ -2033,13 +2023,13 @@ describe('ChatClient', () => { timestamp: Date.now(), finishReason: 'stop', }) - wake.fn?.() - await new Promise((resolve) => setTimeout(resolve, 20)) + await vi.waitFor(() => { + expect(finishedRunIds).toContain('run-1') + }) expect(client.getSessionGenerating()).toBe(true) - // Second run finishes — now should be false - chunks.push({ + push({ type: EventType.RUN_FINISHED, runId: 'run-2', threadId: 'thread-1', @@ -2047,39 +2037,17 @@ describe('ChatClient', () => { timestamp: Date.now(), finishReason: 'stop', }) - wake.fn?.() - await new Promise((resolve) => setTimeout(resolve, 20)) - expect(client.getSessionGenerating()).toBe(false) - // Only two transitions: false→true at start, true→false when all done + await vi.waitFor(() => { + expect(client.getSessionGenerating()).toBe(false) + }) expect(generatingChanges).toEqual([true, false]) client.unsubscribe() }) it('should process future live subscription chunks after persistence clear', async () => { - const wake = { fn: null as (() => void) | null } - const chunks: Array = [] - const connection = { - subscribe: async function* (signal?: AbortSignal) { - while (!signal?.aborted) { - if (chunks.length > 0) { - const batch = chunks.splice(0) - for (const chunk of batch) { - yield chunk - } - } - await new Promise((resolve) => { - wake.fn = resolve - const onAbort = () => resolve() - signal?.addEventListener('abort', onAbort, { once: true }) - }) - } - }, - send: async () => { - wake.fn?.() - }, - } + const { connection, push } = createPushableSubscribeConnection() const persistence = createPersistence() const client = new ChatClient({ connection, @@ -2093,8 +2061,7 @@ describe('ChatClient', () => { }) client.clear() - chunks.push(...createTextChunks('future live', 'future-live')) - wake.fn?.() + push(...createTextChunks('future live', 'future-live')) await vi.waitFor(() => { expect( @@ -2112,28 +2079,7 @@ describe('ChatClient', () => { }) it('should clear all runs on RUN_ERROR without runId', async () => { - const wake = { fn: null as (() => void) | null } - const chunks: Array = [] - const connection = { - subscribe: async function* (signal?: AbortSignal) { - while (!signal?.aborted) { - if (chunks.length > 0) { - const batch = chunks.splice(0) - for (const chunk of batch) { - yield chunk - } - } - await new Promise((resolve) => { - wake.fn = resolve - const onAbort = () => resolve() - signal?.addEventListener('abort', onAbort, { once: true }) - }) - } - }, - send: async () => { - wake.fn?.() - }, - } + const { connection, push } = createPushableSubscribeConnection() const generatingChanges: Array = [] const client = new ChatClient({ connection, @@ -2143,10 +2089,11 @@ describe('ChatClient', () => { }) client.subscribe() - await new Promise((resolve) => setTimeout(resolve, 10)) + await vi.waitFor(() => { + expect(client.getIsSubscribed()).toBe(true) + }) - // Two runs active - chunks.push( + push( { type: EventType.RUN_STARTED, runId: 'run-1', @@ -2162,23 +2109,22 @@ describe('ChatClient', () => { timestamp: Date.now(), }, ) - wake.fn?.() - await new Promise((resolve) => setTimeout(resolve, 20)) - expect(client.getSessionGenerating()).toBe(true) + await vi.waitFor(() => { + expect(client.getSessionGenerating()).toBe(true) + }) - // Session-level error without runId clears everything - chunks.push({ + push({ type: EventType.RUN_ERROR, message: 'session crashed', model: 'test', timestamp: Date.now(), error: { message: 'session crashed' }, }) - wake.fn?.() - await new Promise((resolve) => setTimeout(resolve, 20)) - expect(client.getSessionGenerating()).toBe(false) + await vi.waitFor(() => { + expect(client.getSessionGenerating()).toBe(false) + }) expect(generatingChanges).toEqual([true, false]) client.unsubscribe() @@ -3919,46 +3865,27 @@ describe('ChatClient', () => { describe('concurrent runs and reconnect correctness', () => { it('concurrent runs should not produce duplicate messages or corrupt content', async () => { - const wake = { fn: null as (() => void) | null } - const chunks: Array = [] - const connection = { - subscribe: async function* (signal?: AbortSignal) { - while (!signal?.aborted) { - if (chunks.length > 0) { - const batch = chunks.splice(0) - for (const chunk of batch) { - yield chunk - } - // Re-check: new chunks may have been pushed while yielding - // (the consumer's setTimeout(0) between chunks allows the test - // to push more before we reach the await below) - if (chunks.length > 0) continue - } - await new Promise((resolve) => { - wake.fn = resolve - const onAbort = () => resolve() - signal?.addEventListener('abort', onAbort, { once: true }) - }) - } - }, - send: async () => { - wake.fn?.() - }, - } - - const messagesSnapshots: Array> = [] + const { connection, push } = createPushableSubscribeConnection() + const finishedRunIds: Array = [] const client = new ChatClient({ connection, - onMessagesChange: (msgs) => { - messagesSnapshots.push(msgs.map((m) => ({ ...m }))) + onChunk: (chunk) => { + if ( + chunk.type === EventType.RUN_FINISHED && + 'runId' in chunk && + typeof chunk.runId === 'string' + ) { + finishedRunIds.push(chunk.runId) + } }, }) client.subscribe() - await new Promise((resolve) => setTimeout(resolve, 10)) + await vi.waitFor(() => { + expect(client.getIsSubscribed()).toBe(true) + }) - // Run A starts with text message - chunks.push( + push( { type: EventType.RUN_STARTED, runId: 'run-a', @@ -3972,20 +3899,23 @@ describe('ChatClient', () => { role: 'assistant', model: 'test', timestamp: Date.now(), - } as StreamChunk, + }, { type: EventType.TEXT_MESSAGE_CONTENT, messageId: 'msg-a', model: 'test', timestamp: Date.now(), delta: 'Story: ', - } as StreamChunk, + }, ) - wake.fn?.() - await new Promise((resolve) => setTimeout(resolve, 20)) - // Run B starts concurrently - chunks.push( + await vi.waitFor(() => { + expect( + client.getMessages().some((message) => message.id === 'msg-a'), + ).toBe(true) + }) + + push( { type: EventType.RUN_STARTED, runId: 'run-b', @@ -3999,20 +3929,23 @@ describe('ChatClient', () => { role: 'assistant', model: 'test', timestamp: Date.now(), - } as StreamChunk, + }, { type: EventType.TEXT_MESSAGE_CONTENT, messageId: 'msg-b', model: 'test', timestamp: Date.now(), delta: 'Hi!', - } as StreamChunk, + }, ) - wake.fn?.() - await new Promise((resolve) => setTimeout(resolve, 20)) - // Run B finishes — Run A should still be active - chunks.push({ + await vi.waitFor(() => { + expect( + client.getMessages().some((message) => message.id === 'msg-b'), + ).toBe(true) + }) + + push({ type: EventType.RUN_FINISHED, runId: 'run-b', threadId: 'thread-1', @@ -4020,39 +3953,40 @@ describe('ChatClient', () => { timestamp: Date.now(), finishReason: 'stop', }) - wake.fn?.() - await new Promise((resolve) => setTimeout(resolve, 20)) - // Run A continues streaming - chunks.push({ + await vi.waitFor(() => { + expect(finishedRunIds).toContain('run-b') + }) + + push({ type: EventType.TEXT_MESSAGE_CONTENT, messageId: 'msg-a', model: 'test', timestamp: Date.now(), delta: 'once upon a time', - } as StreamChunk) - wake.fn?.() - await new Promise((resolve) => setTimeout(resolve, 50)) - - // Verify msg-a still has correct content after run-b finished - const messages = client.getMessages() - const msgA = messages.find((m) => m.id === 'msg-a') - const msgB = messages.find((m) => m.id === 'msg-b') + }) - expect(msgA).toBeDefined() - expect(msgB).toBeDefined() - expect(msgA!.parts[0]).toEqual({ - type: 'text', - content: 'Story: once upon a time', + await vi.waitFor(() => { + const msgA = client + .getMessages() + .find((message) => message.id === 'msg-a') + expect(msgA?.parts[0]).toEqual({ + type: 'text', + content: 'Story: once upon a time', + }) }) - expect(msgB!.parts[0]).toEqual({ type: 'text', content: 'Hi!' }) - // No duplicate messages - expect(messages.filter((m) => m.id === 'msg-a')).toHaveLength(1) - expect(messages.filter((m) => m.id === 'msg-b')).toHaveLength(1) + const messages = client.getMessages() + const msgB = messages.find((message) => message.id === 'msg-b') + expect(msgB?.parts[0]).toEqual({ type: 'text', content: 'Hi!' }) + expect(messages.filter((message) => message.id === 'msg-a')).toHaveLength( + 1, + ) + expect(messages.filter((message) => message.id === 'msg-b')).toHaveLength( + 1, + ) - // Finish run A - chunks.push({ + push({ type: EventType.RUN_FINISHED, runId: 'run-a', threadId: 'thread-1', @@ -4060,39 +3994,16 @@ describe('ChatClient', () => { timestamp: Date.now(), finishReason: 'stop', }) - wake.fn?.() - await new Promise((resolve) => setTimeout(resolve, 20)) - expect(client.getSessionGenerating()).toBe(false) + await vi.waitFor(() => { + expect(client.getSessionGenerating()).toBe(false) + }) client.unsubscribe() }) it('reconnect with initialMessages should not duplicate assistant message on content arrival', async () => { - const wake = { fn: null as (() => void) | null } - const chunks: Array = [] - const connection = { - subscribe: async function* (signal?: AbortSignal) { - while (!signal?.aborted) { - if (chunks.length > 0) { - const batch = chunks.splice(0) - for (const chunk of batch) { - yield chunk - } - if (chunks.length > 0) continue - } - await new Promise((resolve) => { - wake.fn = resolve - const onAbort = () => resolve() - signal?.addEventListener('abort', onAbort, { once: true }) - }) - } - }, - send: async () => { - wake.fn?.() - }, - } + const { connection, push } = createPushableSubscribeConnection() - // Simulate reconnect: client created with initialMessages (from SSR/snapshot) const initialMessages: Array = [ { id: 'user-1', @@ -4114,10 +4025,11 @@ describe('ChatClient', () => { }) client.subscribe() - await new Promise((resolve) => setTimeout(resolve, 10)) + await vi.waitFor(() => { + expect(client.getIsSubscribed()).toBe(true) + }) - // Resumed content for in-progress message (no TEXT_MESSAGE_START) - chunks.push( + push( { type: EventType.RUN_STARTED, runId: 'run-1', @@ -4131,7 +4043,7 @@ describe('ChatClient', () => { model: 'test', timestamp: Date.now(), delta: 'time...', - } as StreamChunk, + }, { type: EventType.RUN_FINISHED, runId: 'run-1', @@ -4141,22 +4053,19 @@ describe('ChatClient', () => { finishReason: 'stop', }, ) - wake.fn?.() - await new Promise((resolve) => setTimeout(resolve, 20)) - const messages = client.getMessages() - - // Should still have exactly 2 messages, not 3 - expect(messages).toHaveLength(2) - - // Content should be correctly appended - const asstMsg = messages.find((m) => m.id === 'asst-1') - expect(asstMsg).toBeDefined() - expect(asstMsg!.parts[0]).toEqual({ - type: 'text', - content: 'Once upon a time...', + await vi.waitFor(() => { + const asstMsg = client + .getMessages() + .find((message) => message.id === 'asst-1') + expect(asstMsg?.parts[0]).toEqual({ + type: 'text', + content: 'Once upon a time...', + }) }) + expect(client.getMessages()).toHaveLength(2) + client.unsubscribe() }) }) diff --git a/packages/ai-client/tests/test-utils.ts b/packages/ai-client/tests/test-utils.ts index 1198599e2e..f37948c265 100644 --- a/packages/ai-client/tests/test-utils.ts +++ b/packages/ai-client/tests/test-utils.ts @@ -1,5 +1,8 @@ import { vi } from 'vitest' -import type { ConnectConnectionAdapter } from '../src/connection-adapters' +import type { + ConnectConnectionAdapter, + SubscribeConnectionAdapter, +} from '../src/connection-adapters' import type { ModelMessage, StreamChunk } from '@tanstack/ai/client' import type { ChatClientPersistence, UIMessage } from '../src/types' @@ -131,6 +134,75 @@ 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: + * - rechecks the queue after every yielded batch + * - rechecks again after parking the waiter, so a push in that window still + * wakes + */ +export function createPushableSubscribeConnection(): { + connection: SubscribeConnectionAdapter + push: (...chunks: Array) => void +} { + let wake: (() => void) | null = null + const queue: Array = [] + + const wakeWaiter = () => { + const resolve = wake + wake = null + resolve?.() + } + + const connection: SubscribeConnectionAdapter = { + subscribe: (signal?: AbortSignal) => { + return (async function* () { + while (!signal?.aborted) { + if (queue.length > 0) { + const batch = queue.splice(0) + for (const chunk of batch) { + yield chunk + } + continue + } + await new Promise((resolve) => { + if (signal?.aborted) { + resolve() + return + } + wake = resolve + const onAbort = () => { + if (wake === resolve) { + wake = null + } + resolve() + } + signal?.addEventListener('abort', onAbort, { once: true }) + if (queue.length > 0 || signal?.aborted) { + wakeWaiter() + } + }) + } + })() + }, + send: async () => { + wakeWaiter() + }, + } + + const push = (...chunks: Array) => { + queue.push(...chunks) + wakeWaiter() + } + + return { connection, push } +} + /** * Helper to create simple text content chunks (AG-UI format) */