fix(ai-client): prevent drainPostStreamActions re-entrancy race condition

.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.

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<string>()
 
@@ -846,9 +847,15 @@ export class ChatClient {
    * Drain and execute all queued post-stream actions
    */
   private async drainPostStreamActions(): Promise<void> {
-    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()
   }
 

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')

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)
+    }
+  })
+})