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