feat (ui): transient data parts (#6987)

## Background

In some cases, the data sent to the client is ephemeral and should not
be added to the message history.

## Summary

Add a flag on data part chunks that prevents them from being added to
the message parts.

## Related Issues

#6975

Author: lgrammel

.changeset/quiet-peas-end.md

@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+feat (ui): transient data parts

packages/ai/src/ui-message-stream/ui-message-stream-parts.ts

@@ -104,6 +104,7 @@ export const uiMessageStreamPartSchema = z.union([
     type: z.string().startsWith('data-'),
     id: z.string().optional(),
     data: z.unknown(),
+    transient: z.boolean().optional(),
   }),
   z.strictObject({
     type: z.literal('start-step'),
@@ -131,6 +132,7 @@ export type DataUIMessageStreamPart<DATA_TYPES extends UIDataTypes> = ValueOf<{
     type: `data-${NAME}`;
     id?: string;
     data: DATA_TYPES[NAME];
+    transient?: boolean;
   };
 }>;
 

packages/ai/src/ui/process-ui-message-stream.test.ts

@@ -4139,6 +4139,86 @@ describe('processUIMessageStream', () => {
     });
   });
 
+  describe('data ui parts (transient part)', () => {
+    let dataCalls: InferUIMessageData<UIMessage>[] = [];
+
+    beforeEach(async () => {
+      dataCalls = [];
+
+      const stream = createUIMessageStream([
+        { type: 'start', messageId: 'msg-123' },
+        { type: 'start-step' },
+        {
+          type: 'data-test',
+          data: 'example-data-can-be-anything',
+          transient: true,
+        },
+        { type: 'finish-step' },
+        { type: 'finish' },
+      ]);
+
+      state = createStreamingUIMessageState({
+        messageId: 'msg-123',
+        lastMessage: undefined,
+      });
+
+      await consumeStream({
+        stream: processUIMessageStream({
+          stream,
+          runUpdateMessageJob,
+          onError: error => {
+            throw error;
+          },
+          onData: data => {
+            dataCalls.push(data);
+          },
+        }),
+      });
+    });
+
+    it('should not call the update function with the transient part', async () => {
+      expect(writeCalls).toMatchInlineSnapshot(`
+        [
+          {
+            "message": {
+              "id": "msg-123",
+              "metadata": undefined,
+              "parts": [],
+              "role": "assistant",
+            },
+          },
+        ]
+      `);
+    });
+
+    it('should not have the transient part in the final message state', async () => {
+      expect(state!.message).toMatchInlineSnapshot(`
+        {
+          "id": "msg-123",
+          "metadata": undefined,
+          "parts": [
+            {
+              "type": "step-start",
+            },
+          ],
+          "role": "assistant",
+        }
+      `);
+    });
+
+    it('should call the onData callback with the transient part', async () => {
+      expect(dataCalls).toMatchInlineSnapshot(`
+        [
+          {
+            "data": "example-data-can-be-anything",
+            "transient": true,
+            "type": "data-test",
+          },
+        ]
+      `);
+    });
+  });
+
   describe('data ui parts (single part with id and replacement update)', () => {
     beforeEach(async () => {
       const stream = createUIMessageStream([

packages/ai/src/ui/process-ui-message-stream.ts

@@ -477,6 +477,12 @@ export function processUIMessageStream<UI_MESSAGE extends UIMessage>({
                   InferUIMessageData<UI_MESSAGE>
                 >;
 
+                // transient parts are not added to the message state
+                if (dataPart.transient) {
+                  onData?.(dataPart);
+                  break;
+                }
+
                 // TODO improve type safety
                 const existingPart: any =
                   dataPart.id != null