Skip to content

fix(ai-client): capture abort signal before await to prevent race condition#377

Open
FranDias wants to merge 5 commits intoTanStack:mainfrom
FranDias:fix/chat-client-abort-signal-race
Open

fix(ai-client): capture abort signal before await to prevent race condition#377
FranDias wants to merge 5 commits intoTanStack:mainfrom
FranDias:fix/chat-client-abort-signal-race

Conversation

@FranDias
Copy link
Copy Markdown

@FranDias FranDias commented Mar 14, 2026

Summary

Fixes a race condition in ChatClient.streamResponse() where this.abortController.signal could reference a stale or null controller by the time it is passed to this.connection.connect().

  • Root cause: Between this.abortController = new AbortController() (line 432) and the this.connection.connect(..., this.abortController.signal) call (line 459), there is an await this.callbacksRef.current.onResponse(). During that await, a concurrent stop() call nulls out this.abortController, or a rapid second sendMessage() could replace it with a new controller.
  • Fix: Capture const signal = this.abortController.signal immediately after creation, then pass the local signal variable to connect(), ensuring it always receives the signal from the correct AbortController regardless of concurrent mutations.

How to reproduce

  1. Create a ChatClient and call sendMessage("hello")
  2. Before the stream connects (e.g., during the onResponse callback), call stop() or fire another sendMessage()
  3. this.abortController is now null or points to a different controller
  4. connect() receives a wrong/null signal, leading to either a runtime error or the inability to abort the correct stream

Changes

  • packages/typescript/ai-client/src/chat-client.ts: Capture signal locally right after AbortController creation; pass it to connect() instead of re-reading this.abortController.signal

Test plan

  • All 165 existing tests pass (8 test files in @tanstack/ai-client)
  • Abort-specific tests in chat-client-abort.test.ts all pass (6 tests)
  • Verified the diff is minimal and correct

Summary by CodeRabbit

  • Bug Fixes

    • Ensured the correct cancellation signal is captured and used during streaming so stop/send races no longer produce stale or null signals, preventing interrupted or unrecoverable streams.
  • Tests

    • Added tests validating cancellation behavior when stop is called during a response and across sequential messages, asserting distinct, valid cancellation signals are passed.

@changeset-bot
Copy link
Copy Markdown

changeset-bot bot commented Mar 14, 2026

⚠️ No Changeset found

Latest commit: d70c39f

Merging this PR will not cause a version bump for any packages. If these changes should not result in a new version, you're good to go. If these changes should result in a version bump, you need to add a changeset.

This PR includes no changesets

When changesets are added to this PR, you'll see the packages that this PR includes changesets for and the associated semver types

Click here to learn what changesets are, and how to add one.

Click here if you're a maintainer who wants to add a changeset to this PR

@coderabbitai
Copy link
Copy Markdown
Contributor

coderabbitai bot commented Mar 14, 2026

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

  • @coderabbitai resume to resume automatic reviews.
  • @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

  • ▶️ Resume reviews
  • 🔍 Trigger review

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: ba861c93-c009-44df-951f-c972fa8fd3a1

📥 Commits

Reviewing files that changed from the base of the PR and between d70c39f and e6586fb.

📒 Files selected for processing (1)
  • packages/typescript/ai-client/src/chat-client.ts
🚧 Files skipped from review as they are similar to previous changes (1)
  • packages/typescript/ai-client/src/chat-client.ts

📝 Walkthrough

Walkthrough

Capture the newly created AbortController's signal into a local variable and use that captured signal when invoking the connection adapter and processing streams to avoid races if stop() or sendMessage() reassigns the controller. Adds tests validating captured-signal behavior during onResponse.

Changes

Cohort / File(s) Summary
Signal capture change
packages/typescript/ai-client/src/chat-client.ts
Capture the newly created AbortController's signal in a local signal constant and pass that captured signal into connection.send(...) and downstream stream handling instead of referencing this.abortController.signal.
Tests for abort behavior
packages/typescript/ai-client/tests/chat-client-abort.test.ts
Add two tests: one ensures the captured AbortSignal is passed to the connection when stop() is invoked during onResponse; the other ensures two distinct AbortSignal instances are passed across consecutive connection calls when a second message is queued during onResponse.
Manifest
package.json
Minor single-line manifest edit reported.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Poem

🐰 I caught a signal snug and tight,
Held it steady through the night.
When stop hops in or messages stack,
Each little signal keeps its track.
Hooray — no races, all intact!

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Description check ⚠️ Warning The description provides thorough context including root cause, fix details, reproduction steps, and test plan, but does not follow the repository's required template structure with 🎯 Changes, ✅ Checklist, and 🚀 Release Impact sections. Restructure the description to follow the required template format with the specified section headings and checklist items, while preserving the detailed technical content.
✅ Passed checks (2 passed)
Check name Status Explanation
Title check ✅ Passed The title accurately summarizes the main change: capturing the abort signal before an await to prevent a race condition, which is the core fix in this PR.
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 unit tests (beta)
  • Create PR with unit tests

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.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
@FranDias
Copy link
Copy Markdown
Author

Unit tests for the signal capture fix

Added two tests to chat-client-abort.test.ts that exercise the race condition described in this PR:

1. stop() called during onResponse callback

Calls client.stop() inside the onResponse callback, which runs between AbortController creation and connect(). This sets this.abortController = null.

  • Without the fix: Failsthis.abortController.signal dereferences null, connect() never runs, and the signal is undefined.
  • With the fix: Passes — the locally captured signal is a valid AbortSignal regardless of what happens to this.abortController.

2. sendMessage() called during onResponse callback

Calls client.append() inside onResponse, which queues a second stream that would create a new AbortController.

  • Without the fix: Passes — the queued second stream runs after the first completes (via drainPostStreamActions), so the controllers don't actually interleave in this code path.
  • With the fix: Passes — both calls receive distinct, valid AbortSignal instances.

Test 1 is the key regression test — it directly reproduces the crash from the original bug report.

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-abort.test.ts`:
- Around line 372-398: The test currently fires a nested client.append(...)
inside the onResponse callback without awaiting it and then uses a fixed 50ms
setTimeout to wait, which is non-deterministic and can hide rejections; change
the onResponse callback to capture the returned promise (e.g., assign the result
of client.append(...) to a variable like secondAppendPromise) and then await
that promise after the initial await client.append({...}) instead of using
setTimeout, ensuring any rejection from the second append is surfaced to the
test; reference the onResponse callback and client.append calls to locate and
update the logic accordingly.

ℹ️ Review info
⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: baa8584d-a01d-49a3-89fb-9c4840ead742

📥 Commits

Reviewing files that changed from the base of the PR and between 66cbdca and b5b776e.

📒 Files selected for processing (1)
  • packages/typescript/ai-client/tests/chat-client-abort.test.ts

FranDias and others added 3 commits March 13, 2026 23:37
Capture the nested append() promise and await it directly instead of
relying on a fixed 50ms setTimeout.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
@FranDias
Copy link
Copy Markdown
Author

Just a note that this is actively causing a regression for us after upgrading to vite 8 + @tanstack/ai-client@0.7.2. The timing change in vite 8's Rolldown bundler makes this race condition trigger consistently on the first message sent in a new chat thread — the stream is silently aborted.

We're carrying a local patch (patchedDependencies in bun) with the same fix to unblock, but would love to get this merged upstream.

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.

1 participant