diff --git a/.changeset/fix-drain-post-stream-reentrance.md b/.changeset/fix-drain-post-stream-reentrance.md new file mode 100644 index 000000000..91dd84041 --- /dev/null +++ b/.changeset/fix-drain-post-stream-reentrance.md @@ -0,0 +1,7 @@ +--- +'@tanstack/ai-client': patch +--- + +fix(ai-client): prevent drainPostStreamActions re-entrancy stealing queued actions + +When multiple client tools complete in the same round, nested `drainPostStreamActions()` calls from `streamResponse()`'s `finally` block could steal queued actions, permanently stalling the conversation. Added a re-entrancy guard and a `shouldAutoSend()` check requiring tool-call parts before triggering continuation. diff --git a/packages/typescript/ai-client/src/chat-client.ts b/packages/typescript/ai-client/src/chat-client.ts index 7272394a3..fa453b3ca 100644 --- a/packages/typescript/ai-client/src/chat-client.ts +++ b/packages/typescript/ai-client/src/chat-client.ts @@ -54,6 +54,7 @@ export class ChatClient { // Tracks whether a queued checkForContinuation was skipped because // continuationPending was true (chained approval scenario) private continuationSkipped = false + private draining = false private sessionGenerating = false private activeRunIds = new Set() @@ -846,9 +847,15 @@ export class ChatClient { * Drain and execute all queued post-stream actions */ private async drainPostStreamActions(): Promise { - while (this.postStreamActions.length > 0) { - const action = this.postStreamActions.shift()! - await action() + if (this.draining) return + this.draining = true + try { + while (this.postStreamActions.length > 0) { + const action = this.postStreamActions.shift()! + await action() + } + } finally { + this.draining = false } } @@ -884,9 +891,16 @@ export class ChatClient { } /** - * Check if all tool calls are complete and we should auto-send + * Check if all tool calls are complete and we should auto-send. + * Requires that there is at least one tool call in the last assistant message; + * a text-only response has nothing to auto-send. */ private shouldAutoSend(): boolean { + const messages = this.processor.getMessages() + const lastAssistant = messages.findLast((m) => m.role === 'assistant') + if (!lastAssistant) return false + const hasToolCalls = lastAssistant.parts.some((p) => p.type === 'tool-call') + if (!hasToolCalls) return false return this.processor.areAllToolsComplete() } diff --git a/packages/typescript/ai-client/tests/chat-client.test.ts b/packages/typescript/ai-client/tests/chat-client.test.ts index f4cbe68f5..30ba99996 100644 --- a/packages/typescript/ai-client/tests/chat-client.test.ts +++ b/packages/typescript/ai-client/tests/chat-client.test.ts @@ -8,7 +8,10 @@ import { createApprovalToolCallChunks, createCustomEventChunks, } from './test-utils' -import type { ConnectionAdapter } from '../src/connection-adapters' +import type { + ConnectionAdapter, + ConnectConnectionAdapter, +} from '../src/connection-adapters' import type { StreamChunk } from '@tanstack/ai' import type { UIMessage } from '../src/types' @@ -1235,6 +1238,70 @@ describe('ChatClient', () => { }) }) + describe('drain re-entrancy guard (fix #302)', () => { + it('should continue after multiple client tools complete in the same round', async () => { + // Round 1: two simultaneous tool calls (triggers the re-entrancy bug) + const round1Chunks = createToolCallChunks([ + { id: 'tc-1', name: 'tool_one', arguments: '{}' }, + { id: 'tc-2', name: 'tool_two', arguments: '{}' }, + ]) + // Round 2: final text response + const round2Chunks = createTextChunks('Done!', 'msg-2') + + let callIndex = 0 + const adapter: ConnectConnectionAdapter = { + async *connect(_messages, _data, abortSignal) { + callIndex++ + const chunks = callIndex === 1 ? round1Chunks : round2Chunks + for (const chunk of chunks) { + if (abortSignal?.aborted) return + yield chunk + } + }, + } + + // Both tools execute immediately (synchronously resolve) + const client = new ChatClient({ + connection: adapter, + tools: [ + { + __toolSide: 'client' as const, + name: 'tool_one', + description: 'Tool one', + execute: async () => ({ result: 'one' }), + }, + { + __toolSide: 'client' as const, + name: 'tool_two', + description: 'Tool two', + execute: async () => ({ result: 'two' }), + }, + ], + }) + + // Send initial message — triggers round 1 (two tool calls, both auto-executed) + await client.sendMessage('Run both tools') + + // Wait for loading to stop and the continuation (round 2) to complete + await vi.waitFor( + () => { + expect(client.getIsLoading()).toBe(false) + // Ensure round 2 actually fired + expect(callIndex).toBeGreaterThanOrEqual(2) + }, + { timeout: 2000 }, + ) + + // The final response "Done!" should appear in messages + const messages = client.getMessages() + const lastAssistant = [...messages] + .reverse() + .find((m) => m.role === 'assistant') + const textPart = lastAssistant?.parts.find((p) => p.type === 'text') + expect(textPart?.content).toBe('Done!') + }) + }) + describe('error handling', () => { it('should set error state on connection failure', async () => { const error = new Error('Network error')