fix(ai-client): prevent drainPostStreamActions re-entrancy race condition#429
fix(ai-client): prevent drainPostStreamActions re-entrancy race condition#429AlemTuzlak wants to merge 7 commits intomainfrom
Conversation
…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
📝 WalkthroughWalkthroughThis PR prevents nested Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
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. Comment |
🚀 Changeset Version Preview1 package(s) bumped directly, 9 bumped as dependents. 🟩 Patch bumps
|
|
View your CI Pipeline Execution ↗ for commit 29377a3
☁️ Nx Cloud last updated this comment at |
@tanstack/ai
@tanstack/ai-anthropic
@tanstack/ai-client
@tanstack/ai-code-mode
@tanstack/ai-code-mode-skills
@tanstack/ai-devtools-core
@tanstack/ai-elevenlabs
@tanstack/ai-event-client
@tanstack/ai-fal
@tanstack/ai-gemini
@tanstack/ai-grok
@tanstack/ai-groq
@tanstack/ai-isolate-cloudflare
@tanstack/ai-isolate-node
@tanstack/ai-isolate-quickjs
@tanstack/ai-ollama
@tanstack/ai-openai
@tanstack/ai-openrouter
@tanstack/ai-preact
@tanstack/ai-react
@tanstack/ai-react-ui
@tanstack/ai-solid
@tanstack/ai-solid-ui
@tanstack/ai-svelte
@tanstack/ai-vue
@tanstack/ai-vue-ui
@tanstack/preact-ai-devtools
@tanstack/react-ai-devtools
@tanstack/solid-ai-devtools
commit: |
There was a problem hiding this comment.
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
📒 Files selected for processing (3)
.changeset/fix-drain-post-stream-reentrance.mdpackages/typescript/ai-client/src/chat-client.tspackages/typescript/ai-client/tests/chat-client.test.ts
| import type { | ||
| ConnectionAdapter, | ||
| ConnectConnectionAdapter, | ||
| } from '../src/connection-adapters' |
There was a problem hiding this comment.
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.
| 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.
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.
There was a problem hiding this comment.
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
📒 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) |
There was a problem hiding this comment.
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.
Summary
drainPostStreamActions()preventing nested drain calls from stealing queued actionsshouldAutoSend()guard requiring at least onetool-callpart before triggering auto-continuation (prevents spurious continuation after text-only responses)addToolResult()queues acheckForContinuationaction. Without the guard, the first action'sstreamResponse()→finally→drainPostStreamActions()(nested) steals remaining actions, permanently stalling the conversationFixes #302
Test plan
@tanstack/ai-clienttests passSummary by CodeRabbit