fix(ai): normalize AG-UI snapshot messages to UIMessage[] in MESSAGES_SNAPSHOT handler#698
fix(ai): normalize AG-UI snapshot messages to UIMessage[] in MESSAGES_SNAPSHOT handler#698amitsaroj wants to merge 6 commits into
Conversation
…_SNAPSHOT handler
ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughReplaces unsafe casting of MESSAGES_SNAPSHOT payloads by importing modelMessageToUIMessage and mapping snapshot messages: reuse messages with parts; convert messages without parts (using existing or generated ids) so downstream code can rely on message.parts. ChangesSafe MESSAGES_SNAPSHOT handling
Sequence Diagram(s) sequenceDiagram
participant StreamProcessor
participant MESSAGES_SNAPSHOT
participant modelMessageToUIMessage
StreamProcessor->>MESSAGES_SNAPSHOT: receive snapshot with messages[]
alt message has parts
StreamProcessor->>StreamProcessor: reuse UIMessage as-is
else message lacks parts
StreamProcessor->>modelMessageToUIMessage: convert AG-UI message (use id or generate)
modelMessageToUIMessage->>StreamProcessor: return UIMessage with parts
end
StreamProcessor->>StreamProcessor: downstream code reads message.parts
🎯 3 (Moderate) | ⏱️ ~20 minutes Suggested reviewers:
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Warning There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure. 🔧 ESLint
ESLint install timed out. The project may have too many dependencies for the sandbox. 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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.
Inline comments:
In `@packages/ai/src/activities/chat/stream/processor.ts`:
- Around line 878-882: The snapshot branch is overwriting incoming message IDs
by always calling modelMessageToUIMessage with generateMessageId(); update the
mapping so when converting non-UI messages you preserve the original msg.id and
only generate a new id if msg.id is missing—i.e., in the chunk.messages mapping
(this.messages assignment) pass msg.id (or msg['id']) as the id argument to
modelMessageToUIMessage when present, otherwise call generateMessageId(); this
keeps existing message identity used by ensureAssistantMessage,
TEXT_MESSAGE_CONTENT, toolCallToMessage, and messageStates.
🪄 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: 1d6f8123-d73d-4983-91d7-fbe5fe45edfa
📒 Files selected for processing (1)
packages/ai/src/activities/chat/stream/processor.ts
|
Thanks for the review @coderabbitai — great catch on the ID preservation issue. The second commit ( This ensures downstream event correlation (e.g. |
|
Tip For best results, initiate chat on the files or code changes.
The fix is well-reasoned. 🐇✨ |
|
Hi @tombeckenham and TanStack maintainers — this is a small but impactful fix normalizing AG-UI |
|
Would you be able to add some tests for this @amitsaroj ? |
|
Hi @tombeckenham — done! Added 6 new tests in
|
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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.
Inline comments:
In `@packages/ai/src/activities/chat/stream/processor.ts`:
- Around line 1713-1715: The global errored-tool check
(this.isToolCallPartErrored(toolCall.id)) can return false positives if an older
message has the same toolCall id; update the check to be scoped to the current
message by passing the current message id (e.g.
this.isToolCallPartErrored(toolCall.id, this.message?.id)) or by adding/using a
scoped helper like isToolCallPartErroredForMessage(toolCall.id, messageId), and
replace the two occurrences (around the tool completion branch and the 1741-1750
block) so they only consider parts belonging to the current message before
returning/skip logic.
🪄 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: 9c7857df-e22c-451a-b9a1-4cc1b9a462b8
📒 Files selected for processing (2)
packages/ai/src/activities/chat/stream/processor.tspackages/ai/tests/stream-processor.test.ts
🚧 Files skipped from review as they are similar to previous changes (1)
- packages/ai/tests/stream-processor.test.ts
There was a problem hiding this comment.
Caution
Inline review comments failed to post. This is likely due to GitHub's internal server error or limits when posting large numbers of comments. If you are seeing this consistently it is likely a permissions issue. Please check "Moderation" -> "Code review limits" under your organization settings.
Actionable comments posted: 1
🤖 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.
Inline comments:
In `@packages/ai/src/activities/chat/stream/processor.ts`:
- Around line 1713-1715: The global errored-tool check
(this.isToolCallPartErrored(toolCall.id)) can return false positives if an older
message has the same toolCall id; update the check to be scoped to the current
message by passing the current message id (e.g.
this.isToolCallPartErrored(toolCall.id, this.message?.id)) or by adding/using a
scoped helper like isToolCallPartErroredForMessage(toolCall.id, messageId), and
replace the two occurrences (around the tool completion branch and the 1741-1750
block) so they only consider parts belonging to the current message before
returning/skip logic.
🪄 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: 9c7857df-e22c-451a-b9a1-4cc1b9a462b8
📒 Files selected for processing (2)
packages/ai/src/activities/chat/stream/processor.tspackages/ai/tests/stream-processor.test.ts
🚧 Files skipped from review as they are similar to previous changes (1)
- packages/ai/tests/stream-processor.test.ts
🛑 Comments failed to post (1)
packages/ai/src/activities/chat/stream/processor.ts (1)
1713-1715:
⚠️ Potential issue | 🟠 Major | ⚡ Quick winScope the errored-tool check to the current message to avoid false positives.
Line 1713 calls a global scan (
isToolCallPartErrored(toolCall.id)). If an older message contains the sametoolCallIdin'error', the current call is incorrectly treated as errored and its completion update is skipped.💡 Suggested fix
- if (this.isToolCallPartErrored(toolCall.id)) { + if (this.isToolCallPartErrored(messageId, toolCall.id)) { return } @@ - private isToolCallPartErrored(toolCallId: string): boolean { - return this.messages.some((msg) => - msg.parts.some( - (part) => - part.type === 'tool-call' && - part.id === toolCallId && - part.state === 'error', - ), - ) + private isToolCallPartErrored( + messageId: string, + toolCallId: string, + ): boolean { + const msg = this.messages.find((m) => m.id === messageId) + if (!msg) return false + return msg.parts.some( + (part) => + part.type === 'tool-call' && + part.id === toolCallId && + part.state === 'error', + ) }Also applies to: 1741-1750
🤖 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 1713 - 1715, The global errored-tool check (this.isToolCallPartErrored(toolCall.id)) can return false positives if an older message has the same toolCall id; update the check to be scoped to the current message by passing the current message id (e.g. this.isToolCallPartErrored(toolCall.id, this.message?.id)) or by adding/using a scoped helper like isToolCallPartErroredForMessage(toolCall.id, messageId), and replace the two occurrences (around the tool completion branch and the 1741-1750 block) so they only consider parts belonging to the current message before returning/skip logic.
2 participants
🎯 Changes
Fixes #688
Root cause
handleMessagesSnapshotEventinprocessor.tswas doing an unsafe type cast:AG-UI
MESSAGES_SNAPSHOTmessages have the shape{ id, role, content }— they do not have thepartsarray thatUIMessagerequires. Theas Array<UIMessage>cast masked this at compile time.When any downstream code called
message.parts.find(...)— for example, theonToolCallStateChangedevtools handler inchat-client.ts— it crashed with:Fix
Normalize each message through
modelMessageToUIMessagewhen it does not already haveparts:This mirrors the pattern already used in
normalizeToUIMessage(used byChatClient.append) and is consistent with how all other message ingestion paths work in this codebase.✅ Checklist
pnpm run test:pr.🚀 Release Impact
🧪 Test Plan
MESSAGES_SNAPSHOTevent (e.g., using an AG-UI agent backend that emits snapshot events).TypeError: Cannot read properties of undefined (reading 'find')is thrown.packages/ai/tests/stream-processor.test.tsforMESSAGES_SNAPSHOT resets transient statecover the core normalization path.Summary by CodeRabbit
Bug Fixes
Tests