feat (ui): add onData callback to Chat (#6984)

## Background

Reacting to data parts as they become available is important for data
parts that are not message related, e.g. events (which can be send as
data parts).

## Summary

Add `onData` callback to `Chat`.

## Verification

Tested manually with `next-openai`` data parts example.

## Future Work

* introduce ephemeral data parts

Author: lgrammel

.changeset/ninety-seahorses-fetch.md

@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+feat (ui): add onData callback to Chat

examples/next-openai/app/use-chat-data-ui-parts/page.tsx

@@ -21,10 +21,13 @@ export default function Chat() {
       transport: new DefaultChatTransport({
         api: '/api/use-chat-data-ui-parts',
       }),
+      onData: dataPart => {
+        console.log('dataPart', JSON.stringify(dataPart, null, 2));
+      },
     });
 
   return (
-    <div className="flex flex-col w-full max-w-md py-24 mx-auto stretch">
+    <div className="flex flex-col py-24 mx-auto w-full max-w-md stretch">
       {messages.map(message => (
         <div key={message.id} className="whitespace-pre-wrap">
           {message.role === 'user' ? 'User: ' : 'AI: '}{' '}
@@ -70,7 +73,7 @@ export default function Chat() {
           {status === 'submitted' && <div>Loading...</div>}
           <button
             type="button"
-            className="px-4 py-2 mt-4 text-blue-500 border border-blue-500 rounded-md"
+            className="px-4 py-2 mt-4 text-blue-500 rounded-md border border-blue-500"
             onClick={stop}
           >
             Stop
@@ -83,7 +86,7 @@ export default function Chat() {
           <div className="text-red-500">An error occurred.</div>
           <button
             type="button"
-            className="px-4 py-2 mt-4 text-blue-500 border border-blue-500 rounded-md"
+            className="px-4 py-2 mt-4 text-blue-500 rounded-md border border-blue-500"
             onClick={() => regenerate()}
           >
             Retry

packages/ai/src/ui/chat.ts

@@ -21,6 +21,7 @@ import {
 } from './should-resubmit-messages';
 import {
   isToolUIPart,
+  type DataUIPart,
   type FileUIPart,
   type InferUIMessageData,
   type InferUIMessageMetadata,
@@ -138,6 +139,13 @@ export interface ChatInit<UI_MESSAGE extends UIMessage> {
    * @param message The message that was streamed.
    */
   onFinish?: (options: { message: UI_MESSAGE }) => void;
+
+  /**
+   * Optional callback function that is called when a data part is received.
+   *
+   * @param data The data part that was received.
+   */
+  onData?: (dataPart: DataUIPart<InferUIMessageData<UI_MESSAGE>>) => void;
 }
 
 export abstract class AbstractChat<UI_MESSAGE extends UIMessage> {
@@ -158,6 +166,7 @@ export abstract class AbstractChat<UI_MESSAGE extends UIMessage> {
   private onError?: ChatInit<UI_MESSAGE>['onError'];
   private onToolCall?: ChatInit<UI_MESSAGE>['onToolCall'];
   private onFinish?: ChatInit<UI_MESSAGE>['onFinish'];
+  private onData?: ChatInit<UI_MESSAGE>['onData'];
 
   private activeResponse: ActiveResponse<UI_MESSAGE> | undefined = undefined;
   private jobExecutor = new SerialJobExecutor();
@@ -173,6 +182,7 @@ export abstract class AbstractChat<UI_MESSAGE extends UIMessage> {
     onError,
     onToolCall,
     onFinish,
+    onData,
   }: Omit<ChatInit<UI_MESSAGE>, 'messages'> & {
     state: ChatState<UI_MESSAGE>;
   }) {
@@ -186,6 +196,7 @@ export abstract class AbstractChat<UI_MESSAGE extends UIMessage> {
     this.onError = onError;
     this.onToolCall = onToolCall;
     this.onFinish = onFinish;
+    this.onData = onData;
   }
 
   /**
@@ -495,6 +506,7 @@ export abstract class AbstractChat<UI_MESSAGE extends UIMessage> {
         stream: processUIMessageStream({
           stream,
           onToolCall: this.onToolCall,
+          onData: this.onData,
           messageMetadataSchema: this.messageMetadataSchema,
           dataPartSchemas: this.dataPartSchemas,
           runUpdateMessageJob,

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

@@ -6,7 +6,7 @@ import {
   processUIMessageStream,
   StreamingUIMessageState,
 } from './process-ui-message-stream';
-import { UIMessage } from './ui-messages';
+import { InferUIMessageData, UIMessage } from './ui-messages';
 
 function createUIMessageStream(parts: UIMessageStreamPart[]) {
   return convertArrayToReadableStream(parts);
@@ -4042,7 +4042,11 @@ describe('processUIMessageStream', () => {
   });
 
   describe('data ui parts (single part)', () => {
+    let dataCalls: InferUIMessageData<UIMessage>[] = [];
+
     beforeEach(async () => {
+      dataCalls = [];
+
       const stream = createUIMessageStream([
         { type: 'start', messageId: 'msg-123' },
         { type: 'start-step' },
@@ -4066,6 +4070,9 @@ describe('processUIMessageStream', () => {
           onError: error => {
             throw error;
           },
+          onData: data => {
+            dataCalls.push(data);
+          },
         }),
       });
     });
@@ -4119,6 +4126,17 @@ describe('processUIMessageStream', () => {
         }
       `);
     });
+
+    it('should call the onData callback with the correct arguments', async () => {
+      expect(dataCalls).toMatchInlineSnapshot(`
+        [
+          {
+            "data": "example-data-can-be-anything",
+            "type": "data-test",
+          },
+        ]
+      `);
+    });
   });
 
   describe('data ui parts (single part with id and replacement update)', () => {

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

@@ -6,6 +6,7 @@ import {
 } from '@ai-sdk/provider-utils';
 import {
   InferUIMessageStreamPart,
+  DataUIMessageStreamPart,
   isDataUIMessageStreamPart,
   UIMessageStreamPart,
 } from '../ui-message-stream/ui-message-stream-parts';
@@ -14,6 +15,7 @@ import { mergeObjects } from '../util/merge-objects';
 import { parsePartialJson } from '../util/parse-partial-json';
 import { UIDataTypesToSchemas } from './chat';
 import {
+  DataUIPart,
   getToolName,
   InferUIMessageData,
   InferUIMessageMetadata,
@@ -69,6 +71,7 @@ export function processUIMessageStream<UI_MESSAGE extends UIMessage>({
   dataPartSchemas,
   runUpdateMessageJob,
   onError,
+  onData,
 }: {
   // input stream is not fully typed yet:
   stream: ReadableStream<UIMessageStreamPart>;
@@ -79,6 +82,7 @@ export function processUIMessageStream<UI_MESSAGE extends UIMessage>({
   onToolCall?: (options: {
     toolCall: ToolCall<string, unknown>;
   }) => void | Promise<unknown> | unknown;
+  onData?: (dataPart: DataUIPart<InferUIMessageData<UI_MESSAGE>>) => void;
   runUpdateMessageJob: (
     job: (options: {
       state: StreamingUIMessageState<UI_MESSAGE>;
@@ -468,25 +472,33 @@ export function processUIMessageStream<UI_MESSAGE extends UIMessage>({
 
             default: {
               if (isDataUIMessageStreamPart(part)) {
+                // TODO validate against dataPartSchemas
+                const dataPart = part as DataUIMessageStreamPart<
+                  InferUIMessageData<UI_MESSAGE>
+                >;
+
                 // TODO improve type safety
                 const existingPart: any =
-                  part.id != null
+                  dataPart.id != null
                     ? state.message.parts.find(
                         (partArg: any) =>
-                          part.type === partArg.type && part.id === partArg.id,
+                          dataPart.type === partArg.type &&
+                          dataPart.id === partArg.id,
                       )
                     : undefined;
 
                 if (existingPart != null) {
-                  // TODO improve type safety
+                  // TODO validate merged data against dataPartSchemas
                   existingPart.data =
-                    isObject(existingPart.data) && isObject(part.data)
-                      ? mergeObjects(existingPart.data, part.data)
-                      : part.data;
+                    isObject(existingPart.data) && isObject(dataPart.data)
+                      ? mergeObjects(existingPart.data, dataPart.data)
+                      : dataPart.data;
                 } else {
-                  // TODO improve type safety
-                  state.message.parts.push(part as any);
+                  state.message.parts.push(dataPart);
                 }
+
+                onData?.(dataPart);
+
                 write();
               }
             }