fix(ai): reset segment state when a tool-first message's TEXT_MESSAGE_START arrives

[ramigs]

A tool call that opens an assistant turn drops the first text word that follows it. TOOL_CALL_START with a parentMessageId that has no TEXT_MESSAGE_START yet auto-creates the message and leaves stale state behind, so the second text delta wipes out the first. This PR resets that state in the code path that was missing it.

🎯 Changes

StreamProcessor dropped the first TEXT_MESSAGE_CONTENT delta whenever a tool call opened an assistant turn (no leading text before it) and the model replied with two or more deltas after the tool result. "Hello, world." rendered as "world.".

This fix resets the leftover tool-call state at the point TEXT_MESSAGE_START arrives, so no stale state reaches the content handler at all.

Root cause

Issue. StreamProcessor drops the first TEXT_MESSAGE_CONTENT delta of the text that follows a tool call, when that tool call's parentMessageId had no prior TEXT_MESSAGE_START. This is the common AG-UI shape "call a tool, then explain the result," as one assistant turn.

Cause. In handleToolCallStartEvent (packages/ai/src/activities/chat/stream/processor.ts), ensureAssistantMessage() auto-creates the message and sets pendingManualMessageId, then marks hasToolCallsSinceTextStart = true. When the real TEXT_MESSAGE_START for that message arrives, handleTextMessageStartEvent finds pendingManualMessageId still set and takes Case 1 ("manual message"), which clears the pending id and returns. Only Case 2 ("message already exists") resets hasToolCallsSinceTextStart, currentSegmentText, and lastEmittedText. Case 1 skips that reset. The stale hasToolCallsSinceTextStart survives into the content handler, and its "new segment after a tool call" check fires one delta late — on the second delta, not the first — wiping the already-accumulated first delta before the second one is added.

Fix. Case 1 now performs the same reset Case 2 already does, so hasToolCallsSinceTextStart never survives past the message's real TEXT_MESSAGE_START.

Possible alternatives

• PR fix(ai): keep first text delta after a tool call in the same turn #1246 (fix(ai): keep first text delta after a tool call in the same turn), already open, fixes the same root state at a different point: it patches handleTextMessageContentEvent so the "new segment" reset fires on the first post-tool delta (when the prior segment is empty) instead of the second. I found this bug and wrote this fix independently — issue StreamProcessor drops the first TEXT_MESSAGE_CONTENT delta when a tool call precedes any text in the same message #1247 predates my discovery of fix(ai): keep first text delta after a tool call in the same turn #1246 — but fix(ai): keep first text delta after a tool call in the same turn #1246 got there first and is a valid, different-layer fix for the same underlying state bug. A maintainer may prefer one over the other; I am not closing this PR so both are visible for comparison.
• Reset inside handleTextMessageContentEvent only (the approach fix(ai): keep first text delta after a tool call in the same turn #1246 takes): works, but leaves the two TEXT_MESSAGE_START code paths (Case 1 and Case 2) diverging on state they should agree on, which is what let this bug in undetected in the first place (see the existing "tool-first flows" test that only exercised a single delta).

Testing

Commands run:

• pnpm exec vitest run tests/stream-processor.test.ts (packages/ai) — 206/206 pass
• pnpm test:lib (packages/ai) — 1720/1720 pass
• pnpm test:types, pnpm test:build (packages/ai) — clean
• nx affected --targets=test:sherif,test:knip,test:docs,test:kiira,test:oxlint,test:lib,test:types,test:build,build (repo-wide, matches pnpm test:pr's target set) — all pass
• pnpm exec playwright test tool-first-text.spec.ts (testing/e2e) — pass on this branch; also confirmed it fails on unpatched main (rebuilt @tanstack/ai with the source fix reverted, reran, got "world." instead of "Hello, world.", then restored the fix and reran green)
• Did not run the full E2E suite locally (100+ specs across every provider) — the change is isolated to one function plus new, self-contained harness files that no other spec touches, so CI's full run is the right place for that

Manual test (reproduce, then confirm fixed):

1. On main, run the repro in issue StreamProcessor drops the first TEXT_MESSAGE_CONTENT delta when a tool call precedes any text in the same message #1247 (or pnpm exec vitest run tests/stream-processor.test.ts -t "1247" on this branch, pointed at main's processor.ts) — the assembled text comes out as "world.".
2. On this branch, run the same command — the text comes out as "Hello, world.".
3. Or: cd testing/e2e && pnpm exec playwright test tool-first-text.spec.ts — visits /tool-first-text, a harness page that streams the exact wire shape from the issue, and asserts the rendered text.

How this PR makes testing easy: a new unit test (packages/ai/tests/stream-processor.test.ts) reproduces the bug in isolation, and a new E2E harness (api.tool-first-text-wire.ts + tool-first-text.tsx + tool-first-text.spec.ts) reproduces it through the real browser client, following the existing foreign-interrupt wire-harness pattern.

Linked issues

Fixes #1247

Risk / rollback

Low risk: the change only runs when hasToolCallsSinceTextStart is already true in Case 1, mirrors logic Case 2 already ships, and the existing "tool-first flows" test (stream-processor.test.ts:4065) still passes unchanged. Revert is a plain git revert of this PR; no flags or migrations involved.

✅ Checklist

• I have followed the steps in the Contributing guide.
• I have tested code changes locally with pnpm run test:pr (via the equivalent nx affected target set — see Testing).
• I fully understand the code in this pull request, including any code generated with AI assistance.
• Docs: this change is not user-facing — it restores behavior already documented in docs/chat-architecture.md.
• Changeset: I added a changeset (pnpm changeset).

🚀 Release Impact

• This change affects published code, and I have generated a changeset.

Summary by CodeRabbit

• Bug Fixes

  • Fixed streaming responses so the first assistant text is preserved when a tool call occurs before text generation.
  • Ensured multiple text updates are combined correctly into a single message.

• Tests

  • Added regression and end-to-end coverage for tool-first streaming scenarios.
  • Added a test page and simulated streaming endpoint to verify complete assistant responses.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 7e149561-d330-434a-bd45-bca9692cb3b8

📥 Commits

Reviewing files that changed from the base of the PR and between 4e9c5d2 and 75965c7.

📒 Files selected for processing (7)

• .changeset/tool-first-text-drop.md
• packages/ai/src/activities/chat/stream/processor.ts
• packages/ai/tests/stream-processor.test.ts
• testing/e2e/src/routeTree.gen.ts
• testing/e2e/src/routes/api.tool-first-text-wire.ts
• testing/e2e/src/routes/tool-first-text.tsx
• testing/e2e/tests/tool-first-text.spec.ts

Included review availability: Your plan provides up to 8 included reviews per hour; 7 remain after this review.

---

📝 Walkthrough

Walkthrough

The change fixes StreamProcessor handling when tool calls precede TEXT_MESSAGE_START. It preserves all text deltas, adds unit coverage, and adds an end-to-end stream, page, route registration, and Playwright test.

Changes

Tool-first assistant text

Layer / File(s)	Summary
Processor state reset and regression coverage packages/ai/src/activities/chat/stream/processor.ts, packages/ai/tests/stream-processor.test.ts, .changeset/tool-first-text-drop.md	TEXT_MESSAGE_START now clears stale tool-call text state before processing new content. The unit test verifies that multiple deltas combine into the complete text. The changeset records the patch release.
Tool-first SSE route and route registration testing/e2e/src/routes/api.tool-first-text-wire.ts, testing/e2e/src/routeTree.gen.ts	The API route emits tool-call events followed by two text deltas and returns them as SSE. The generated route tree registers the API and page routes.
End-to-end harness and assertion testing/e2e/src/routes/tool-first-text.tsx, testing/e2e/tests/tool-first-text.spec.ts	The harness sends one message, combines assistant text parts, and renders the result. The Playwright test asserts Hello, world..

Estimated code review effort: 3 (Moderate) | ~20 minutes

Merge Risk: ⚪ Minimal · up to 75965

This localized fix restores the first text emitted after a tool call without changing the broader message flow; no actionable merge-blocking risk remains after normal checks and review.

Suggested reviewers: alemtuzlak, tombeckenham

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely identifies the primary fix: resetting segment state when a tool-first message receives TEXT_MESSAGE_START.
Description check	✅ Passed	The description follows the repository template. It explains the change, root cause, testing, linked issue, risk, checklist, and release impact. All required checklist items are addressed.
Linked Issues check	✅ Passed	The changes satisfy issue #1247. They reset the stale state in the pendingManualMessageId path and add unit and end-to-end regression coverage for multiple text deltas after a tool call.
Out of Scope Changes check	✅ Passed	The changes remain within scope. The processor fix, regression tests, changeset, generated route updates, and dedicated end-to-end harness all support the linked issue. No unrelated functional changes…
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 2 functions across 6 files. (1 skipped: 1 …

Full details: Out of Scope Changes check

Explanation

The changes remain within scope. The processor fix, regression tests, changeset, generated route updates, and dedicated end-to-end harness all support the linked issue. No unrelated functional changes are evident.

Full details: Docstring Coverage

Explanation

Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 2 functions across 6 files. (1 skipped: 1 unsupported.)

✨ 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

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}