From 7f3e8b6dc86ab38453b0b31f21cad7848e88dc29 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Wed, 8 Apr 2026 13:32:39 +0200 Subject: [PATCH 1/4] fix(ai-client): prevent drainPostStreamActions re-entrancy stealing queued actions When multiple client tools complete in the same round, each addToolResult() queues a checkForContinuation action. The first drain executes one action which calls streamResponse(), whose finally block calls drainPostStreamActions() again (nested). The inner drain steals the remaining actions, permanently stalling the conversation. Add a draining flag to skip nested drain calls. The outer drain processes all actions sequentially, preventing action theft. Also fix shouldAutoSend() to require at least one tool call in the last assistant message. Previously it returned true for text-only responses (areAllToolsComplete() returns true when toolParts.length === 0), causing the second queued checkForContinuation action to incorrectly trigger an extra continuation round and produce duplicate content. Fixes #302 --- .../typescript/ai-client/src/chat-client.ts | 22 +++++-- .../ai-client/tests/chat-client.test.ts | 66 ++++++++++++++++++- 2 files changed, 83 insertions(+), 5 deletions(-) diff --git a/packages/typescript/ai-client/src/chat-client.ts b/packages/typescript/ai-client/src/chat-client.ts index 7272394a..fa453b3c 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 f4cbe68f..823597a6 100644 --- a/packages/typescript/ai-client/tests/chat-client.test.ts +++ b/packages/typescript/ai-client/tests/chat-client.test.ts @@ -8,7 +8,7 @@ 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 +1235,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') From 09585fa816c02e9b6ca89f877c6f8651d7a952a8 Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Wed, 8 Apr 2026 11:38:14 +0000 Subject: [PATCH 2/4] ci: apply automated fixes --- packages/typescript/ai-client/tests/chat-client.test.ts | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/packages/typescript/ai-client/tests/chat-client.test.ts b/packages/typescript/ai-client/tests/chat-client.test.ts index 823597a6..c69754ab 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, ConnectConnectionAdapter } from '../src/connection-adapters' +import type { + ConnectionAdapter, + ConnectConnectionAdapter, +} from '../src/connection-adapters' import type { StreamChunk } from '@tanstack/ai' import type { UIMessage } from '../src/types' From f915ab26405f74305b0df27100c77465028daa0c Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Wed, 8 Apr 2026 13:39:06 +0200 Subject: [PATCH 3/4] changeset: fix drain post-stream re-entrancy --- .changeset/fix-drain-post-stream-reentrance.md | 7 +++++++ 1 file changed, 7 insertions(+) create mode 100644 .changeset/fix-drain-post-stream-reentrance.md diff --git a/.changeset/fix-drain-post-stream-reentrance.md b/.changeset/fix-drain-post-stream-reentrance.md new file mode 100644 index 00000000..91dd8404 --- /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. From 6d27f7022667444165caf64be882f43297bd1b65 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Wed, 8 Apr 2026 13:50:40 +0200 Subject: [PATCH 4/4] fix: resolve type errors in drain re-entrancy test --- packages/typescript/ai-client/tests/chat-client.test.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/typescript/ai-client/tests/chat-client.test.ts b/packages/typescript/ai-client/tests/chat-client.test.ts index c69754ab..30ba9999 100644 --- a/packages/typescript/ai-client/tests/chat-client.test.ts +++ b/packages/typescript/ai-client/tests/chat-client.test.ts @@ -1250,7 +1250,7 @@ describe('ChatClient', () => { let callIndex = 0 const adapter: ConnectConnectionAdapter = { - async *connect(messages, data, abortSignal) { + async *connect(_messages, _data, abortSignal) { callIndex++ const chunks = callIndex === 1 ? round1Chunks : round2Chunks for (const chunk of chunks) {