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') 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..bd606ae59 --- /dev/null +++ b/testing/e2e/tests/tools-test/drain-reentrance.spec.ts @@ -0,0 +1,111 @@ +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) + } + }) +})