Skip to content

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

Open
AlemTuzlak wants to merge 7 commits intomainfrom
fix/drain-post-stream-reentrance
Open

fix(ai-client): prevent drainPostStreamActions re-entrancy race condition#429
AlemTuzlak wants to merge 7 commits intomainfrom
fix/drain-post-stream-reentrance

Conversation

@AlemTuzlak
Copy link
Copy Markdown
Contributor

@AlemTuzlak AlemTuzlak commented Apr 8, 2026

Summary

  • Add re-entrancy guard to drainPostStreamActions() preventing nested drain calls from stealing queued actions
  • Add shouldAutoSend() guard requiring at least one tool-call part before triggering auto-continuation (prevents spurious continuation after text-only responses)
  • When multiple client tools complete in the same round, each addToolResult() queues a checkForContinuation action. Without the guard, the first action's streamResponse()finallydrainPostStreamActions() (nested) steals remaining actions, permanently stalling the conversation

Fixes #302

Test plan

  • New test: "should continue after multiple client tools complete in the same round" — verifies two simultaneous tool calls trigger exactly one continuation and final text appears
  • All 199 @tanstack/ai-client tests pass

Summary by CodeRabbit

  • Bug Fixes
    • Prevented re-entrancy that could consume queued actions and stall conversations during simultaneous tool calls and streaming.
    • tightened automatic continuation so it only triggers when a tool-call part is present, avoiding unwanted continuations after text-only responses.
  • Tests
    • Added unit and end-to-end tests covering the multi-round continuation and drain re-entrancy scenario.
  • Chores
    • Added a release changeset documenting the patch.

…ueued 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
@coderabbitai
Copy link
Copy Markdown
Contributor

coderabbitai bot commented Apr 8, 2026

📝 Walkthrough

Walkthrough

This PR prevents nested drainPostStreamActions() invocations from consuming queued actions by adding a draining guard and tightens shouldAutoSend() to require tool-call parts on the last assistant message before auto-continuing, fixing stalled continuations when multiple tools complete in the same round.

Changes

Cohort / File(s) Summary
Changeset metadata
/.changeset/fix-drain-post-stream-reentrance.md
Adds a patch changeset for @tanstack/ai-client documenting the re-entrancy fix for drainPostStreamActions().
Implementation
packages/typescript/ai-client/src/chat-client.ts
Introduces a draining boolean guard to prevent nested/concurrent drainPostStreamActions() execution; updates shouldAutoSend() to require the last assistant message to contain at least one tool-call part before triggering automatic continuation.
Unit tests
packages/typescript/ai-client/tests/chat-client.test.ts
Adds a new test group for drain re-entrancy (imports ConnectConnectionAdapter) that simulates multi-round streams with simultaneous tool calls and asserts continuation text appears from the second round.
E2E tests
testing/e2e/tests/tools-test/drain-reentrance.spec.ts
New Playwright spec exercising the parallel-client-tools scenario to verify the regression: ensures two tool executions, two execution-complete events, loading completes, and the assistant follow-up text ("All displayed") appears. Includes diagnostic capture on failure.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Poem

🐰 I hopped in late to tidy the queue,
A tiny guard now keeps drains from two-step stew.
Tools finish, follow-ups no longer stall,
Rabbity logic mended — hooray for all! ✨

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Title check ✅ Passed The title precisely summarizes the core fix: preventing a re-entrancy race condition in drainPostStreamActions that causes the chat to stall when multiple client tools complete simultaneously.
Description check ✅ Passed The description covers the key changes, root cause, test plan, and issue reference; however, the PR Checklist items are not explicitly marked as completed despite the context indicating they were followed.
Linked Issues check ✅ Passed All code changes directly address issue #302: a re-entrancy guard prevents nested drains from stealing queued actions, and shouldAutoSend() now requires a tool-call part to prevent spurious continuations.
Out of Scope Changes check ✅ Passed All changes are tightly scoped to fixing the re-entrancy issue: guard added to drainPostStreamActions, shouldAutoSend() modified, and tests added to validate the fix.
Docstring Coverage ✅ Passed No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
📝 Generate docstrings
  • Create stacked PR
  • Commit on current branch
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch fix/drain-post-stream-reentrance

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

@github-actions
Copy link
Copy Markdown
Contributor

github-actions bot commented Apr 8, 2026

🚀 Changeset Version Preview

1 package(s) bumped directly, 9 bumped as dependents.

🟩 Patch bumps

Package Version Reason
@tanstack/ai-client 0.7.9 → 0.7.10 Changeset
@tanstack/ai-preact 0.6.14 → 0.6.15 Dependent
@tanstack/ai-react 0.7.10 → 0.7.11 Dependent
@tanstack/ai-solid 0.6.14 → 0.6.15 Dependent
@tanstack/ai-svelte 0.6.14 → 0.6.15 Dependent
@tanstack/ai-vue 0.6.14 → 0.6.15 Dependent
@tanstack/ai-vue-ui 0.1.25 → 0.1.26 Dependent
ts-svelte-chat 0.1.31 → 0.1.32 Dependent
ts-vue-chat 0.1.31 → 0.1.32 Dependent
vanilla-chat 0.0.29 → 0.0.30 Dependent

@nx-cloud
Copy link
Copy Markdown

nx-cloud bot commented Apr 8, 2026

View your CI Pipeline Execution ↗ for commit 29377a3

Command Status Duration Result
nx run-many --targets=build --exclude=examples/** ✅ Succeeded 1m 5s View ↗

☁️ Nx Cloud last updated this comment at 2026-04-15 10:05:28 UTC

@pkg-pr-new
Copy link
Copy Markdown

pkg-pr-new bot commented Apr 8, 2026

Open in StackBlitz

@tanstack/ai

npm i https://pkg.pr.new/@tanstack/ai@429

@tanstack/ai-anthropic

npm i https://pkg.pr.new/@tanstack/ai-anthropic@429

@tanstack/ai-client

npm i https://pkg.pr.new/@tanstack/ai-client@429

@tanstack/ai-code-mode

npm i https://pkg.pr.new/@tanstack/ai-code-mode@429

@tanstack/ai-code-mode-skills

npm i https://pkg.pr.new/@tanstack/ai-code-mode-skills@429

@tanstack/ai-devtools-core

npm i https://pkg.pr.new/@tanstack/ai-devtools-core@429

@tanstack/ai-elevenlabs

npm i https://pkg.pr.new/@tanstack/ai-elevenlabs@429

@tanstack/ai-event-client

npm i https://pkg.pr.new/@tanstack/ai-event-client@429

@tanstack/ai-fal

npm i https://pkg.pr.new/@tanstack/ai-fal@429

@tanstack/ai-gemini

npm i https://pkg.pr.new/@tanstack/ai-gemini@429

@tanstack/ai-grok

npm i https://pkg.pr.new/@tanstack/ai-grok@429

@tanstack/ai-groq

npm i https://pkg.pr.new/@tanstack/ai-groq@429

@tanstack/ai-isolate-cloudflare

npm i https://pkg.pr.new/@tanstack/ai-isolate-cloudflare@429

@tanstack/ai-isolate-node

npm i https://pkg.pr.new/@tanstack/ai-isolate-node@429

@tanstack/ai-isolate-quickjs

npm i https://pkg.pr.new/@tanstack/ai-isolate-quickjs@429

@tanstack/ai-ollama

npm i https://pkg.pr.new/@tanstack/ai-ollama@429

@tanstack/ai-openai

npm i https://pkg.pr.new/@tanstack/ai-openai@429

@tanstack/ai-openrouter

npm i https://pkg.pr.new/@tanstack/ai-openrouter@429

@tanstack/ai-preact

npm i https://pkg.pr.new/@tanstack/ai-preact@429

@tanstack/ai-react

npm i https://pkg.pr.new/@tanstack/ai-react@429

@tanstack/ai-react-ui

npm i https://pkg.pr.new/@tanstack/ai-react-ui@429

@tanstack/ai-solid

npm i https://pkg.pr.new/@tanstack/ai-solid@429

@tanstack/ai-solid-ui

npm i https://pkg.pr.new/@tanstack/ai-solid-ui@429

@tanstack/ai-svelte

npm i https://pkg.pr.new/@tanstack/ai-svelte@429

@tanstack/ai-vue

npm i https://pkg.pr.new/@tanstack/ai-vue@429

@tanstack/ai-vue-ui

npm i https://pkg.pr.new/@tanstack/ai-vue-ui@429

@tanstack/preact-ai-devtools

npm i https://pkg.pr.new/@tanstack/preact-ai-devtools@429

@tanstack/react-ai-devtools

npm i https://pkg.pr.new/@tanstack/react-ai-devtools@429

@tanstack/solid-ai-devtools

npm i https://pkg.pr.new/@tanstack/solid-ai-devtools@429

commit: 29377a3

Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@packages/typescript/ai-client/tests/chat-client.test.ts`:
- Around line 11-14: The import specifiers in the import from
'../src/connection-adapters' are not alphabetized; reorder the named imports so
they are sorted alphabetically (e.g., place ConnectConnectionAdapter before
ConnectionAdapter) to satisfy the linter's sort-imports rule and update the
import line that currently lists ConnectionAdapter and ConnectConnectionAdapter.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 67771ad9-6029-4d21-a7c4-07a93e440236

📥 Commits

Reviewing files that changed from the base of the PR and between e8204b2 and 3d06f8b.

📒 Files selected for processing (3)
  • .changeset/fix-drain-post-stream-reentrance.md
  • packages/typescript/ai-client/src/chat-client.ts
  • packages/typescript/ai-client/tests/chat-client.test.ts

Comment on lines +11 to +14
import type {
ConnectionAdapter,
ConnectConnectionAdapter,
} from '../src/connection-adapters'
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

Fix import member order to satisfy lint.

The import specifiers are not alphabetically sorted, which matches the reported sort-imports error.

🔧 Suggested fix
 import type {
-  ConnectionAdapter,
   ConnectConnectionAdapter,
+  ConnectionAdapter,
 } from '../src/connection-adapters'
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
import type {
ConnectionAdapter,
ConnectConnectionAdapter,
} from '../src/connection-adapters'
import type {
ConnectConnectionAdapter,
ConnectionAdapter,
} from '../src/connection-adapters'
🧰 Tools
🪛 ESLint

[error] 13-13: Member 'ConnectConnectionAdapter' of the import declaration should be sorted alphabetically.

(sort-imports)

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@packages/typescript/ai-client/tests/chat-client.test.ts` around lines 11 -
14, The import specifiers in the import from '../src/connection-adapters' are
not alphabetized; reorder the named imports so they are sorted alphabetically
(e.g., place ConnectConnectionAdapter before ConnectionAdapter) to satisfy the
linter's sort-imports rule and update the import line that currently lists
ConnectionAdapter and ConnectConnectionAdapter.

AlemTuzlak and others added 2 commits April 15, 2026 11:57
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.
Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@testing/e2e/tests/tools-test/drain-reentrance.spec.ts`:
- Line 45: The test waits using waitForTestComplete(page, 20000, 2) but then
immediately asserts the follow-up/continuation message, which causes flakiness;
add an explicit wait for the continuation text before the assertion (e.g., call
an existing helper like waitForText(page, "expected continuation text", timeout)
or implement waitForContinuationText(page, "expected continuation text", 5000))
and replace or add this wait prior to the expect/assert that checks the
follow-up message; update both occurrences around the assertions that follow
waitForTestComplete so the test only asserts after the continuation text is
observed.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: b97d2e60-1822-4b95-a8b0-633fafd2d5c3

📥 Commits

Reviewing files that changed from the base of the PR and between 3d06f8b and 29377a3.

📒 Files selected for processing (1)
  • testing/e2e/tests/tools-test/drain-reentrance.spec.ts

// 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)
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Add an explicit wait for continuation text to reduce E2E flakiness.

Line 45 waits on generic completion state; Line 91 then immediately asserts the follow-up text. If UI/message rendering lags slightly, this can fail intermittently even when continuation did fire.

Suggested patch
@@
     await waitForTestComplete(page, 20000, 2)
+
+    // Ensure continuation text is actually rendered before asserting content.
+    await page.waitForFunction(() => {
+      const el = document.getElementById('messages-json-content')
+      if (!el?.textContent) return false
+      try {
+        const messages = JSON.parse(el.textContent)
+        const text = messages
+          .filter((m: any) => m.role === 'assistant')
+          .flatMap((m: any) =>
+            (m.parts || [])
+              .filter((p: any) => p.type === 'text')
+              .map((p: any) => p.content),
+          )
+          .join(' ')
+        return text.includes('All displayed')
+      } catch {
+        return false
+      }
+    }, { timeout: 5000 })

Also applies to: 67-91

🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@testing/e2e/tests/tools-test/drain-reentrance.spec.ts` at line 45, The test
waits using waitForTestComplete(page, 20000, 2) but then immediately asserts the
follow-up/continuation message, which causes flakiness; add an explicit wait for
the continuation text before the assertion (e.g., call an existing helper like
waitForText(page, "expected continuation text", timeout) or implement
waitForContinuationText(page, "expected continuation text", 5000)) and replace
or add this wait prior to the expect/assert that checks the follow-up message;
update both occurrences around the assertions that follow waitForTestComplete so
the test only asserts after the continuation text is observed.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Client tool continuation stalls when multiple tools complete in the same round

1 participant