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

* 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

* ci: apply automated fixes

* changeset: fix drain post-stream re-entrancy

* fix: resolve type errors in drain re-entrancy test

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

* ci: apply automated fixes

---------

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>

Author: AlemTuzlak

.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>()
 
@@ -847,9 +848,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
     }
   }
 
@@ -885,9 +892,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'
 
@@ -1239,6 +1242,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)
+    }
+  })
+})