fix (ui): inject generated response message id (#7198)

## Background

It was not possible to override the generated message id in the ui
message stream.

## Summary

Add `generateMessageId` property to `toUIMessageStream` on stream text
result.

## Related Issues

Fixes #7178

Author: lgrammel

.changeset/giant-ravens-reflect.md

@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+fix (ui): inject generated response message id

packages/ai/src/generate-text/stream-text-result.ts

@@ -1,4 +1,4 @@
-import { ReasoningPart } from '@ai-sdk/provider-utils';
+import { IdGenerator, ReasoningPart } from '@ai-sdk/provider-utils';
 import { ServerResponse } from 'node:http';
 import { InferUIMessageChunk } from '../../src/ui-message-stream/ui-message-chunks';
 import { UIMessageStreamResponseInit } from '../../src/ui-message-stream/ui-message-stream-response-init';
@@ -29,6 +29,14 @@ export type UIMessageStreamOptions<UI_MESSAGE extends UIMessage> = {
    */
   originalMessages?: UI_MESSAGE[];
 
+  /**
+   * Generate a message ID for the response message.
+   *
+   * If not provided, no message ID will be set for the response message (unless
+   * the original messages are provided and the last message is an assistant message).
+   */
+  generateMessageId?: IdGenerator;
+
   onFinish?: (options: {
     /**
      * The updates list of UI messages.

packages/ai/src/generate-text/stream-text.test.ts

@@ -2648,7 +2648,7 @@ describe('streamText', () => {
       `);
     });
 
-    it('should not generate a new message id when onFinish is provided', async () => {
+    it('should not generate a new message id when onFinish is provided and generateMessageId is not provided', async () => {
       const result = streamText({
         model: createTestModel(),
         ...defaultSettings(),
@@ -2702,6 +2702,61 @@ describe('streamText', () => {
         ]
       `);
     });
+
+    it('should generate a new message id when generateMessageId is provided', async () => {
+      const result = streamText({
+        model: createTestModel(),
+        ...defaultSettings(),
+      });
+
+      const uiMessageStream = result.toUIMessageStream({
+        generateMessageId: mockId({ prefix: 'id' }),
+      });
+
+      expect(await convertReadableStreamToArray(uiMessageStream))
+        .toMatchInlineSnapshot(`
+        [
+          {
+            "messageId": "id-0",
+            "messageMetadata": undefined,
+            "type": "start",
+          },
+          {
+            "type": "start-step",
+          },
+          {
+            "id": "1",
+            "type": "text-start",
+          },
+          {
+            "delta": "Hello",
+            "id": "1",
+            "type": "text-delta",
+          },
+          {
+            "delta": ", ",
+            "id": "1",
+            "type": "text-delta",
+          },
+          {
+            "delta": "world!",
+            "id": "1",
+            "type": "text-delta",
+          },
+          {
+            "id": "1",
+            "type": "text-end",
+          },
+          {
+            "type": "finish-step",
+          },
+          {
+            "messageMetadata": undefined,
+            "type": "finish",
+          },
+        ]
+      `);
+    });
   });
 
   describe('result.toUIMessageStreamResponse', () => {

packages/ai/src/generate-text/stream-text.ts

@@ -1573,6 +1573,7 @@ However, the LLM results are expected to be small enough to not cause issues.
 
   toUIMessageStream<UI_MESSAGE extends UIMessage>({
     originalMessages,
+    generateMessageId,
     onFinish,
     messageMetadata,
     sendReasoning = true,
@@ -1806,7 +1807,7 @@ However, the LLM results are expected to be small enough to not cause issues.
 
     return handleUIMessageStreamFinish<UI_MESSAGE>({
       stream: baseStream,
-      messageId: responseMessageId,
+      messageId: responseMessageId ?? generateMessageId?.(),
       originalMessages,
       onFinish,
       onError,
@@ -1817,6 +1818,7 @@ However, the LLM results are expected to be small enough to not cause issues.
     response: ServerResponse,
     {
       originalMessages,
+      generateMessageId,
       onFinish,
       messageMetadata,
       sendReasoning,
@@ -1831,6 +1833,7 @@ However, the LLM results are expected to be small enough to not cause issues.
       response,
       stream: this.toUIMessageStream({
         originalMessages,
+        generateMessageId,
         onFinish,
         messageMetadata,
         sendReasoning,
@@ -1853,6 +1856,7 @@ However, the LLM results are expected to be small enough to not cause issues.
 
   toUIMessageStreamResponse<UI_MESSAGE extends UIMessage>({
     originalMessages,
+    generateMessageId,
     onFinish,
     messageMetadata,
     sendReasoning,
@@ -1866,6 +1870,7 @@ However, the LLM results are expected to be small enough to not cause issues.
     return createUIMessageStreamResponse({
       stream: this.toUIMessageStream({
         originalMessages,
+        generateMessageId,
         onFinish,
         messageMetadata,
         sendReasoning,

packages/ai/src/ui-message-stream/handle-ui-message-stream-finish.ts

@@ -48,10 +48,6 @@ export function handleUIMessageStreamFinish<UI_MESSAGE extends UIMessage>({
     responseMessage: UI_MESSAGE;
   }) => void;
 }): ReadableStream<InferUIMessageChunk<UI_MESSAGE>> {
-  if (onFinish == null) {
-    return stream;
-  }
-
   // last message is only relevant for assistant messages
   let lastMessage: UI_MESSAGE | undefined =
     originalMessages?.[originalMessages.length - 1];
@@ -62,6 +58,31 @@ export function handleUIMessageStreamFinish<UI_MESSAGE extends UIMessage>({
     messageId = lastMessage.id;
   }
 
+  const idInjectedStream = stream.pipeThrough(
+    new TransformStream<
+      InferUIMessageChunk<UI_MESSAGE>,
+      InferUIMessageChunk<UI_MESSAGE>
+    >({
+      transform(chunk, controller) {
+        // when there is no messageId in the start chunk,
+        // but the user checked for persistence,
+        // inject the messageId into the chunk
+        if (chunk.type === 'start') {
+          const startChunk = chunk as UIMessageChunk & { type: 'start' };
+          if (startChunk.messageId == null && messageId != null) {
+            startChunk.messageId = messageId;
+          }
+        }
+
+        controller.enqueue(chunk);
+      },
+    }),
+  );
+
+  if (onFinish == null) {
+    return idInjectedStream;
+  }
+
   const state = createStreamingUIMessageState<UI_MESSAGE>({
     lastMessage: lastMessage
       ? (structuredClone(lastMessage) as UI_MESSAGE)
@@ -79,26 +100,7 @@ export function handleUIMessageStreamFinish<UI_MESSAGE extends UIMessage>({
   };
 
   return processUIMessageStream<UI_MESSAGE>({
-    stream: stream.pipeThrough(
-      new TransformStream<
-        InferUIMessageChunk<UI_MESSAGE>,
-        InferUIMessageChunk<UI_MESSAGE>
-      >({
-        transform(chunk, controller) {
-          // when there is no messageId in the start chunk,
-          // but the user checked for persistence,
-          // inject the messageId into the chunk
-          if (chunk.type === 'start') {
-            const startChunk = chunk as UIMessageChunk & { type: 'start' };
-            if (startChunk.messageId == null && messageId != null) {
-              startChunk.messageId = messageId;
-            }
-          }
-
-          controller.enqueue(chunk);
-        },
-      }),
-    ),
+    stream: idInjectedStream,
     runUpdateMessageJob,
     onError,
   }).pipeThrough(