From 7f3e8b6dc86ab38453b0b31f21cad7848e88dc29 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Wed, 8 Apr 2026 13:32:39 +0200 Subject: [PATCH 1/6] 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 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..823597a69 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/6] 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 823597a69..c69754ab7 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/6] 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 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. From 6d27f7022667444165caf64be882f43297bd1b65 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Wed, 8 Apr 2026 13:50:40 +0200 Subject: [PATCH 4/6] 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 c69754ab7..30ba99996 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) { From 67c96adf98efd16d398ed67160a0252cc03097ef Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Wed, 15 Apr 2026 11:57:04 +0200 Subject: [PATCH 5/6] test: add e2e regression test for drain re-entrancy stall (#302) Add a Playwright e2e test that verifies parallel client tools complete and the continuation fires with a follow-up text response. Without the drainPostStreamActions() re-entrancy guard, nested drain calls steal queued actions and permanently stall the conversation after both tools complete. The test asserts that the follow-up text "All displayed" arrives, which would time out without the fix. --- .../tests/tools-test/drain-reentrance.spec.ts | 113 ++++++++++++++++++ 1 file changed, 113 insertions(+) create mode 100644 testing/e2e/tests/tools-test/drain-reentrance.spec.ts diff --git a/testing/e2e/tests/tools-test/drain-reentrance.spec.ts b/testing/e2e/tests/tools-test/drain-reentrance.spec.ts new file mode 100644 index 000000000..213e40acd --- /dev/null +++ b/testing/e2e/tests/tools-test/drain-reentrance.spec.ts @@ -0,0 +1,113 @@ +import { test, expect } from '../fixtures' +import { + selectScenario, + runTest, + waitForTestComplete, + getMetadata, + getEventLog, + getToolCalls, +} from './helpers' + +/** + * Drain Re-Entrancy Guard E2E Tests + * + * Regression test for GitHub issue #302 / PR #429. + * + * When multiple client tools complete in the same round, each addToolResult() + * queues a checkForContinuation action. Without a re-entrancy guard on + * drainPostStreamActions(), the first action's streamResponse() → finally → + * drainPostStreamActions() (nested) steals the remaining actions from the + * queue, permanently stalling the conversation. The user sees tool results + * but the model never produces its follow-up text response. + * + * The fix adds a re-entrancy guard to drainPostStreamActions() and a + * shouldAutoSend() check requiring tool-call parts before triggering + * continuation. + * + * This test uses the parallel-client-tools scenario (2 client tools in the + * same turn) and verifies not just that both tools execute, but critically + * that the **continuation fires and the follow-up text response arrives**. + * Without the fix, the test would time out waiting for the follow-up text. + */ + +test.describe('Drain Re-Entrancy Guard (Regression #302)', () => { + test('parallel client tools complete and continuation fires with follow-up text', async ({ + page, + testId, + aimockPort, + }) => { + await selectScenario(page, 'parallel-client-tools', testId, aimockPort) + await runTest(page) + + // Wait for the test to fully complete — this includes the continuation + // round producing the follow-up text. Without the fix, this would + // time out because the continuation never fires. + await waitForTestComplete(page, 20000, 2) + + // Verify both client tools executed + const metadata = await getMetadata(page) + expect(metadata.testComplete).toBe('true') + expect(metadata.isLoading).toBe('false') + + const events = await getEventLog(page) + const toolNames = [...new Set(events.map((e) => e.toolName))] + expect(toolNames).toContain('show_notification') + expect(toolNames).toContain('display_chart') + + // Verify both tools reached execution-complete state + const completionEvents = events.filter( + (e) => e.type === 'execution-complete', + ) + expect(completionEvents.length).toBe(2) + + // CRITICAL ASSERTION: Verify the follow-up text from round 2 was received. + // Without the re-entrancy fix, the conversation stalls after both tools + // complete — the continuation request is never sent, so this text never + // arrives. + const messages = await page.evaluate(() => { + const el = document.getElementById('messages-json-content') + if (!el) return [] + try { + return JSON.parse(el.textContent || '[]') + } catch { + return [] + } + }) + + const assistantMessages = messages.filter( + (m: any) => m.role === 'assistant', + ) + + // There should be at least 2 assistant messages: + // 1. The tool-call round (with both tool calls + results) + // 2. The continuation round (with the follow-up text) + expect(assistantMessages.length).toBeGreaterThanOrEqual(2) + + // The follow-up text from the continuation round should be present + const allTextParts = assistantMessages.flatMap((m: any) => + m.parts + .filter((p: any) => p.type === 'text') + .map((p: any) => p.content), + ) + const allText = allTextParts.join(' ') + expect(allText).toContain('All displayed') + }) + + test.afterEach(async ({ page }, testInfo) => { + if (testInfo.status !== testInfo.expectedStatus) { + await page.screenshot({ + path: `test-results/drain-reentrance-failure-${testInfo.title.replace(/\s+/g, '-')}.png`, + fullPage: true, + }) + + const toolCalls = await getToolCalls(page) + const metadata = await getMetadata(page) + const events = await getEventLog(page) + + console.log('Test failed. Debug info:') + console.log('Metadata:', metadata) + console.log('Tool calls:', toolCalls) + console.log('Events:', events) + } + }) +}) From 29377a335313795ed32063f8c0c86e50ac0e05dc Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Wed, 15 Apr 2026 09:58:06 +0000 Subject: [PATCH 6/6] ci: apply automated fixes --- testing/e2e/tests/tools-test/drain-reentrance.spec.ts | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/testing/e2e/tests/tools-test/drain-reentrance.spec.ts b/testing/e2e/tests/tools-test/drain-reentrance.spec.ts index 213e40acd..bd606ae59 100644 --- a/testing/e2e/tests/tools-test/drain-reentrance.spec.ts +++ b/testing/e2e/tests/tools-test/drain-reentrance.spec.ts @@ -85,9 +85,7 @@ test.describe('Drain Re-Entrancy Guard (Regression #302)', () => { // The follow-up text from the continuation round should be present const allTextParts = assistantMessages.flatMap((m: any) => - m.parts - .filter((p: any) => p.type === 'text') - .map((p: any) => p.content), + m.parts.filter((p: any) => p.type === 'text').map((p: any) => p.content), ) const allText = allTextParts.join(' ') expect(allText).toContain('All displayed')