data part is not appended to message

[JerryKwan]

Description

When using DataStreamWriter's writeData() to send custom data, the data part is not assigned the message as the operations when invoking writeMessageAnnotation() which annotations are assigned to message as expected.
In the current implementation the data is cached and returned as data for the whole chat, not accessible by message, although Message has a data field, but it's never assigned.
As you can see from the following code snippets.

https://github.com/vercel/ai/blob/main/packages/ui-utils/src/process-chat-response.ts#L222-L234
https://github.com/vercel/ai/blob/main/packages/ui-utils/src/types.ts#L47-L82

It would be better the data part is accessible by message like annotations.

Code example

No response

AI provider

No response

Additional context

No response

[ifreeman6]

I meet the same issue. Any solutions about it ?

[JerryKwan]

@ifreeman6
I think it just need a little hack of the process-chat-response.ts if you maintain your own fork.

[lgrammel]

data is per chat, messageAnnotations are per message. Please use messageAnnotations for your use case if you need data per message.

[JerryKwan]

@lgrammel
Thanks for paying attention to this issue.
As far as I know, messageAnnotations is not per message all of the time.
Consider the following scenario,
When a FinishStepPart was sent, and several messages should be returned right? but how messageAnnotations are processed? messageAnnotations are not cleared when a new message starts, they are accumulated. and the last message has all of the annotations.

[lgrammel]

On the client you need to accumulate consecutive assistant messages manually (unfortunately). It is related to mapping the messages back to the LLM format. When you do this accumulation, you would also get annotations from both msgs. I'll write this up somewhere, it is a somewhat annoying limitation that I plan to overcome with a new approach to AI SDK UI.