fix(ai): preserve tool-call metadata across MESSAGES_SNAPSHOT so addToolResult works (#859)

[IdiotTIQS]

🎯 Changes

Fixes #859 — addToolResult / addToolOutput silently no-op'd after a MESSAGES_SNAPSHOT.

Root cause

handleMessagesSnapshotEvent replaces this.messages with the normalized snapshot. The AG-UI MESSAGES_SNAPSHOT wire shape cannot reconstruct client-side tool-call metadata a server may omit:

• a role: 'tool' message only carries toolCallId + content, and
• an assistant snapshot message may drop the toolCalls the client already observed via TOOL_CALL_* events.

After the snapshot, the matching tool-call part could disappear, so a later addToolResult(toolCallId) couldn't locate the call:

[StreamProcessor] Could not find message with tool call <id>

and silently no-op'd. Removing the snapshot (the reporter's "if I delete it") leaves the streamed assistant message intact, which is why it then worked.

Fix

handleMessagesSnapshotEvent now reconciles the normalized snapshot against the pre-snapshot state in a small post-pass (reconcileSnapshotToolCalls):

• (a) consistency — detached tool-result-only assistant messages (produced by aguiSnapshotMessageToUIMessage for role: 'tool' wire messages) are anchored into the preceding assistant message, matching the streaming fan-out shape assistant: [text, tool-call, tool-result, ...]. This also removes a duplicate tool-result that previously appeared when both the assistant toolCalls and a separate tool message were present.
• (b) bug fix — when a tool-result references a toolCallId whose tool-call part is absent from the snapshot, the tool-call part is carried forward from the pre-snapshot state so addToolResult(toolCallId) can still find it.

The trigger for (b) is deliberately narrow: it only fires when the snapshot explicitly references a toolCallId (via a tool-result) but does not itself supply the corresponding tool-call part. This frames the change as recovering metadata the snapshot payload cannot express, not as making the snapshot non-authoritative.

Tests

Added unit coverage in stream-processor.test.ts (the same location as the prior snapshot fix #698):

• anchors a snapshot tool message into the preceding assistant message
• preserves an in-flight tool-call the snapshot omits so addToolResult still works (the MESSAGES_SNAPSHOT + addToolOutput #859 repro)
• does not duplicate a tool-call the snapshot already supplies

✅ Checklist

• I have followed the steps in the Contributing guide.
• I have tested this code locally.

Local verification

All @tanstack/ai targets are green locally:

• test:lib — 164 tests passing (incl. 3 new)
• test:types, test:eslint (0 errors), build, test:build

I also ran the full repo-wide pnpm run test:pr. The only failures were in packages unrelated to this change and are platform/environment artifacts on my Windows machine, not regressions:

• @tanstack/ai-sandbox › tests/harness-cwd.test.ts — a path-separator assertion (path.join produces \ on Windows; the code returns /). Fails only on Windows; unrelated to this PR.
• @tanstack/ai-angular:build — an intermittent npm peer-dep install crash (Cannot read properties of null (reading 'matches')); re-running it on a clean main checkout succeeds.

Neither touches @tanstack/ai. Linux CI should be unaffected — happy to adjust if CI surfaces anything.

🚀 Release Impact

• This change affects published code, and I have generated a changeset (@tanstack/ai patch).
• This change is docs/CI/dev-only (no release).

---

Note on E2E: there is no existing MESSAGES_SNAPSHOT harness in testing/e2e (snapshots aren't part of the recorded provider flow), and the directly analogous prior snapshot fix (#698) was covered with unit tests in stream-processor.test.ts rather than E2E. I followed that precedent — happy to add an E2E scenario if you'd prefer one.

I raised the two fix directions on the issue first (#859 comment) and implemented the combined approach with the focus on the actual bug fix (b). Let me know if you'd like the scope narrowed/widened.

Summary by CodeRabbit

• Bug Fixes
  • Improved snapshot hydration to correctly reconcile tool-result messages with the originating assistant turn.
  • Preserved client-observed tool-call metadata when snapshots omit it, so follow-up tool result handling continues to work.
  • Prevented duplicate tool-call parts when snapshot data already includes them.
• Tests
  • Extended snapshot hydration tests with assertions for anchoring, metadata preservation, and de-duplication during streaming.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 6fd94515-e88d-4c76-812c-5483596fd3e0

📥 Commits

Reviewing files that changed from the base of the PR and between db6f9b8 and 3064ea1.

📒 Files selected for processing (1)

• packages/ai/src/activities/chat/stream/processor.ts

🚧 Files skipped from review as they are similar to previous changes (1)

• packages/ai/src/activities/chat/stream/processor.ts

---

📝 Walkthrough

Walkthrough

Adds a reconciliation pass in handleMessagesSnapshotEvent that preserves tool-call metadata across MESSAGES_SNAPSHOT hydration, plus tests and a changeset for the behavior.

Changes

Snapshot tool-call reconciliation

Layer / File(s)	Summary
Reconciliation logic in handleMessagesSnapshotEvent packages/ai/src/activities/chat/stream/processor.ts	Adds reconcileSnapshotToolCalls, which indexes prior tool-call parts, anchors detached tool-result-only messages to preceding assistant messages, and carries forward missing tool-call parts without overwriting snapshot-provided metadata; wired into handleMessagesSnapshotEvent.
Tests and changeset for reconciliation behavior packages/ai/tests/stream-processor.test.ts, .changeset/snapshot-preserve-tool-call.md	Adds three MESSAGES_SNAPSHOT test cases for anchoring, preserving in-flight tool-call parts, and de-duplication, and adds a changeset documenting the patch.

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

Possibly related PRs

• TanStack/ai#698: Both PRs modify handleMessagesSnapshotEvent for MESSAGES_SNAPSHOT hydration and related tool-call/tool-result handling.

Suggested reviewers: tombeckenham, AlemTuzlak

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly states the main fix: preserving tool-call metadata across MESSAGES_SNAPSHOT so addToolResult works.
Description check	✅ Passed	The description matches the template and includes changes, checklist, release impact, and testing details.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ 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.}

[IdiotTIQS]

On the E2E coverage requirement in CONTRIBUTING — flagging this explicitly rather than leaving it buried in the PR description, so a maintainer can make the call.

This fix lives entirely in the client-side StreamProcessor (reconciling a MESSAGES_SNAPSHOT against pre-snapshot tool-call state). I looked at testing/e2e and there is no path that emits MESSAGES_SNAPSHOT: the harness drives real provider adapters against aimock, and aimock only replays ordinary LLM streams (text / tool-call / tool-result). MESSAGES_SNAPSHOT is an AG-UI server feature no provider in the matrix produces, so the existing providersFor('feature') model can't exercise this path without new infrastructure (a dedicated snapshot-emitting route + UI page + spec).

Given that, I covered the behaviour with unit tests in stream-processor.test.ts, which is exactly where the directly analogous prior snapshot fix (#698, "normalize AG-UI snapshot messages") was tested — no E2E was added there either.

Options, your call:

1. Keep unit-only (current state), consistent with fix(ai): normalize AG-UI snapshot messages to UIMessage[] in MESSAGES_SNAPSHOT handler #698.
2. I add an @tanstack/ai-client integration test that drives a real ChatClient through a mock connection which injects a MESSAGES_SNAPSHOT chunk mid-stream, then calls addToolResult — proves the end-to-end client behaviour without Playwright.
3. I build out a full Playwright E2E: a new snapshot-emitting route + tools page + spec in testing/e2e.

I'd lean (1) or (2); (3) feels like a lot of new harness for a client-internal fix. Happy to do whichever you prefer.

[coderabbitai]

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

packages/ai/src/activities/chat/stream/processor.ts (1)

875-901: 🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift

Rebuild tool-call routing/state after snapshot reconciliation.

Line 878 clears toolCallToMessage and messageStates, but Lines 899-900 only restore this.messages. After a mid-stream MESSAGES_SNAPSHOT, follow-up TOOL_CALL_ARGS, TOOL_CALL_END, TOOL_CALL_RESULT, approval-requested, and ui-resource events for the reconciled tool call can still get dropped because the routing maps were never rehydrated. Please rebuild the per-message tool state from the reconciled assistant messages before returning here.

Also applies to: 919-988

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@packages/ai/src/activities/chat/stream/processor.ts` around lines 875 - 901,
handleMessagesSnapshotEvent currently restores this.messages after snapshot
reconciliation but leaves the routing/state maps cleared, so follow-up tool
events can no longer resolve the reconciled calls. After
reconcileSnapshotToolCalls returns, rebuild toolCallToMessage and messageStates
from the reconciled assistant messages before emitMessagesChange, using the same
tool-call/tool-result metadata preserved by aguiSnapshotMessageToUIMessage and
reconcileSnapshotToolCalls. Apply the same rehydration logic anywhere the stream
snapshot path resets state so TOOL_CALL_ARGS, TOOL_CALL_END, TOOL_CALL_RESULT,
approval-requested, and ui-resource events can keep routing correctly.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Outside diff comments:
In `@packages/ai/src/activities/chat/stream/processor.ts`:
- Around line 875-901: handleMessagesSnapshotEvent currently restores
this.messages after snapshot reconciliation but leaves the routing/state maps
cleared, so follow-up tool events can no longer resolve the reconciled calls.
After reconcileSnapshotToolCalls returns, rebuild toolCallToMessage and
messageStates from the reconciled assistant messages before emitMessagesChange,
using the same tool-call/tool-result metadata preserved by
aguiSnapshotMessageToUIMessage and reconcileSnapshotToolCalls. Apply the same
rehydration logic anywhere the stream snapshot path resets state so
TOOL_CALL_ARGS, TOOL_CALL_END, TOOL_CALL_RESULT, approval-requested, and
ui-resource events can keep routing correctly.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 64b00b7f-8fc7-449a-b1a7-789f91ed573c

📥 Commits

Reviewing files that changed from the base of the PR and between 1880999 and db6f9b8.

📒 Files selected for processing (3)

• .changeset/snapshot-preserve-tool-call.md
• packages/ai/src/activities/chat/stream/processor.ts
• packages/ai/tests/stream-processor.test.ts

[Copilot]

Pull request overview

Fixes a regression where addToolResult / addToolOutput could no-op after an AG-UI MESSAGES_SNAPSHOT by reconciling snapshot-normalized messages with pre-snapshot tool-call metadata that the snapshot wire format cannot fully represent.

Changes:

• Add a reconciliation post-pass after snapshot normalization to (1) anchor role: 'tool' fan-out messages into the preceding assistant turn and (2) carry forward missing tool-call parts from pre-snapshot state when referenced by a tool-result.
• Add unit tests covering anchoring, metadata preservation for the #859 repro, and avoiding tool-call duplication when the snapshot already includes toolCalls.
• Add a changeset describing the patch-level fix and behavior.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.

File	Description
packages/ai/src/activities/chat/stream/processor.ts	Reconciles snapshot-normalized messages with pre-snapshot tool-call parts so tool-call metadata and tool-result anchoring remain consistent after MESSAGES_SNAPSHOT.
packages/ai/tests/stream-processor.test.ts	Adds regression coverage for anchoring snapshot tool messages and preserving omitted tool-call metadata so addToolResult continues to work.
.changeset/snapshot-preserve-tool-call.md	Publishes release notes for the patch fix addressing addToolResult / addToolOutput no-op after snapshots.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[IdiotTIQS]

Re: the review about rebuilding tool-call routing/state after snapshot reconciliation (toolCallToMessage / messageStates not rehydrated after resetStreamState()).

Good catch, and the mechanism is real — after a MESSAGES_SNAPSHOT, handleToolCallArgsEvent (bails at toolCallToMessage.get(...) then getMessageState(...)), TOOL_CALL_END, TOOL_CALL_RESULT, approval, and ui-resource routing all key off maps that resetStreamState() clears. So follow-up streaming events for a reconciled tool call can be dropped.

Two notes on scope:

1. Pre-existing, not introduced here. On main before this PR, handleMessagesSnapshotEvent already did this.resetStreamState() then this.messages = chunk.messages.map(...) — the routing maps were always cleared on every snapshot and never rebuilt. This PR does not touch those maps.

2. Orthogonal to the bug this PR fixes. MESSAGES_SNAPSHOT + addToolOutput #859 is about addToolResult(toolCallId), which resolves purely against this.messages and never consults toolCallToMessage / messageStates. The fix carries the tool-call part forward into this.messages so addToolResult can find it. Restreaming into a reconciled call is a separate concern.

Rehydrating per-message tool state is a reasonable robustness improvement, but (as the review notes) it is a heavy lift that widens this PR beyond #859. I would prefer to keep this PR minimal and open a dedicated follow-up issue, unless you would rather I fold it in here — happy to do whichever you prefer.