message.annotations missing annotations sent after finish_roundtrip since @ai-sdk/[email protected]

[aloisklink]

Description

Since @ai-sdk/[email protected] (commit aa2dc58), calling data.appendMessageAnnotation to send a message annotation chunk after the finish_roundtrip chunk prevents the annotation from being stored in message.annotations on the client.

For example, something like the following no longer works:

  import { StreamData, convertToCoreMessages, streamText } from 'ai';

  const data = new StreamData();
  // Call the language model
  const result = await streamText({
    model: openai('gpt-4o'),
    messages: convertToCoreMessages(messages),
    async onFinish() {
      await setTimeout(10); // delay, so annotation is sent after "finishReason":"stop"
      data.appendMessageAnnotation("This message annotation cannot be read in @ai-sdk/[email protected]");
      await data.close();
    },
  });

  // Respond with the stream
  return result.toDataStreamResponse({data});

Code example

Try applying this patch on 91bea48, then running the example in examples/sveltekit-openai:

From fd04e353228e684ccb4e465df8809078619f8f21 Mon Sep 17 00:00:00 2001
From: Alois Klink <[email protected]
Date: Fri, 4 Oct 2024 04:53:42 +0900
Subject: [PATCH] Example showing `data.appendMessageAnnotation` not working if
 "finishReason":"stop" has already been sent

---
 examples/sveltekit-openai/src/routes/+page.svelte     |  2 +-
 .../sveltekit-openai/src/routes/api/chat/+server.ts   | 11 +++++++++--
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/examples/sveltekit-openai/src/routes/+page.svelte b/examples/sveltekit-openai/src/routes/+page.svelte
index 9a0e4b31..d83d85ce 100644
--- a/examples/sveltekit-openai/src/routes/+page.svelte
+++ b/examples/sveltekit-openai/src/routes/+page.svelte
@@ -27,7 +27,7 @@
   <h1>useChat</h1>
   <ul>
     {#each $messages as message}
-      <li>{message.role}: {message.content}</li>
+      <li>{message.role}: {message.content}. Annotations was: {JSON.stringify(message.annotations)}</li>
     {/each}
   </ul>

diff --git a/examples/sveltekit-openai/src/routes/api/chat/+server.ts b/examples/sveltekit-openai/src/routes/api/chat/+server.ts
index c6ec142f..227e8c49 100644
--- a/examples/sveltekit-openai/src/routes/api/chat/+server.ts
+++ b/examples/sveltekit-openai/src/routes/api/chat/+server.ts
@@ -1,6 +1,7 @@
 import { createOpenAI } from '@ai-sdk/openai';
-import { convertToCoreMessages, streamText } from 'ai';
+import { StreamData, convertToCoreMessages, streamText } from 'ai';
 import type { RequestHandler } from './$types';
+import { setTimeout } from 'node:timers/promises';

 import { env } from '$env/dynamic/private';

@@ -17,6 +18,9 @@ export const POST = (async ({ request }) => {
   // Extract the `prompt` from the body of the request
   const { messages } = await request.json();

+  const data = new StreamData();
+
+  data.appendMessageAnnotation("Annotation sent before message!");
   // Call the language model
   const result = await streamText({
     model: openai('gpt-4-turbo'),
@@ -24,9 +28,12 @@ export const POST = (async ({ request }) => {
     async onFinish({ text, toolCalls, toolResults, usage, finishReason }) {
       // implement your own logic here, e.g. for storing messages
       // or recording token usage
+      await setTimeout(10); // delay, so annotation is sent after "finishReason":"stop"
+      data.appendMessageAnnotation("Annotation sent after message!");
+      await data.close();
     },
   });

   // Respond with the stream
-  return result.toDataStreamResponse();
+  return result.toDataStreamResponse({data});
 }) satisfies RequestHandler;
--
2.43.0

annotation sent after "finishReason":"stop"

annotation sent before "finishReason":"stop"

Additional context

Maybe something like the following might fix the issue?

diff --git a/packages/ui-utils/src/process-data-protocol-response.ts b/packages/ui-utils/src/process-data-protocol-response.ts
index 2805b9c6..7f66f038 100644
--- a/packages/ui-utils/src/process-data-protocol-response.ts
+++ b/packages/ui-utils/src/process-data-protocol-response.ts
@@ -345,6 +345,13 @@ export async function processDataProtocolResponse({
     update([...previousMessages, ...merged], [...data]); // make a copy of the data array
   }

+  // Append all annotations to all messages
+  if (message_annotations?.length) {
+    update(previousMessages.map(message => ({
+      ...assignAnnotationsToMessage(message, message_annotations),
+    })), [...data]);
+  }
+
   onFinish?.({ message: prefixMap.text, finishReason, usage });

   return {

The message annotations feature isn't really well documented, so it's hard to know what it's supposed to do in a multi-step/multi-message stream.