fix(ai): normalize AG-UI snapshot messages to UIMessage[] in MESSAGES_SNAPSHOT handler

[amitsaroj]

🎯 Changes

Fixes #688

Root cause

handleMessagesSnapshotEvent in processor.ts was doing an unsafe type cast:

// Before (unsafe)
this.messages = [...chunk.messages] as Array<UIMessage>

AG-UI MESSAGES_SNAPSHOT messages have the shape { id, role, content } — they do not have the parts array that UIMessage requires. The as Array<UIMessage> cast masked this at compile time.

When any downstream code called message.parts.find(...) — for example, the onToolCallStateChange devtools handler in chat-client.ts — it crashed with:

TypeError: Cannot read properties of undefined (reading 'find')

Fix

Normalize each message through modelMessageToUIMessage when it does not already have parts:

// After (safe)
this.messages = chunk.messages.map((msg) =>
  'parts' in msg
    ? (msg as UIMessage)
    : modelMessageToUIMessage(msg as any, generateMessageId()),
)

This mirrors the pattern already used in normalizeToUIMessage (used by ChatClient.append) and is consistent with how all other message ingestion paths work in this codebase.

✅ Checklist

• I have followed the steps in the Contributing guide.
• - [ ] I have tested this code locally with pnpm run test:pr.

🚀 Release Impact

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

🧪 Test Plan

1. Send a chat message that triggers a MESSAGES_SNAPSHOT event (e.g., using an AG-UI agent backend that emits snapshot events).
2. 2. Verify no TypeError: Cannot read properties of undefined (reading 'find') is thrown.
3. 3. Existing unit tests in packages/ai/tests/stream-processor.test.ts for MESSAGES_SNAPSHOT resets transient state cover the core normalization path.

Summary by CodeRabbit

• Bug Fixes
  • Normalize streamed chat snapshot messages into a consistent UI message format (converting legacy model messages and assigning IDs where needed) so message content/parts are reliably available and rendering/processing behaves consistently across all message types.

[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: 5d5f33cf-c97e-4b32-982f-3490af88bc7b

📥 Commits

Reviewing files that changed from the base of the PR and between bc8e675 and b091adc.

📒 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

Replaces 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.

Changes

Safe MESSAGES_SNAPSHOT handling

Layer / File(s)	Summary
Safe normalization of MESSAGES_SNAPSHOT payloads packages/ai/src/activities/chat/stream/processor.ts	Adds modelMessageToUIMessage import and updates MESSAGES_SNAPSHOT handler to map snapshot messages: reuse messages that already have parts, convert others via modelMessageToUIMessage (using an existing or generated id) so each resulting UIMessage includes parts.

🎯 3 (Moderate) | ⏱️ ~20 minutes

🐰 I hopped through snapshots, parts awry,
I stitched each message an ID so they’d not cry,
Converted the models, one by one,
Now every part shines in the sun,
A gentle fix — the stream runs spry.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely summarizes the main change: normalizing AG-UI snapshot messages to UIMessage[] in the MESSAGES_SNAPSHOT handler, which is the core fix.
Description check	✅ Passed	The description includes the key sections: root cause analysis with before/after code, the fix explanation, checklist items, release impact (changeset generated), and test plan with specific steps.
Linked Issues check	✅ Passed	The PR directly addresses issue #688 by normalizing MESSAGES_SNAPSHOT messages to include the required 'parts' array, preventing the TypeError when downstream code accesses message.parts.find(...).
Out of Scope Changes check	✅ Passed	All changes are focused on fixing the MESSAGES_SNAPSHOT handler in processor.ts to normalize messages properly, with no out-of-scope modifications detected.
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

• X
• Mastodon
• Reddit
• LinkedIn

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

[coderabbitai]

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

📥 Commits

Reviewing files that changed from the base of the PR and between 769a6c7 and bc8e675.

📒 Files selected for processing (1)

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

[amitsaroj]

Thanks for the review @coderabbitai — great catch on the ID preservation issue.

The second commit (b091adc) addresses this: instead of always generating a fresh generateMessageId(), the normalization now uses (msg as { id?: string }).id ?? generateMessageId() — preserving the AG-UI message's original id when present and only falling back to a generated ID when the message has no id field.

This ensures downstream event correlation (e.g. TEXT_MESSAGE_CONTENT, TOOL_CALL_*, ensureAssistantMessage) continues to work correctly after a MESSAGES_SNAPSHOT reset.

[coderabbitai]

Tip

For best results, initiate chat on the files or code changes.

@amitsaroj: Thanks for the explanation! The approach looks solid — using (msg as { id?: string }).id ?? generateMessageId() is exactly the right pattern here. Preserving the original AG-UI message id ensures that any subsequent events (TEXT_MESSAGE_CONTENT, TOOL_CALL_*, etc.) that reference that ID will correctly correlate with the already-stored message after the snapshot reset, rather than silently breaking due to a mismatched generated ID.

The fix is well-reasoned. 🐇✨