ai 7.0.82 → 7.0.84

package/src/ui-message-stream/to-ui-message-stream.ts

@@ -7,6 +7,7 @@ import type { UIMessage } from '../ui/ui-messages';
 import { getResponseUIMessageId } from './get-response-ui-message-id';
 import { handleUIMessageStreamFinish } from './handle-ui-message-stream-finish';
 import type { InferUIMessageChunk } from './ui-message-chunks';
+import type { UIMessageStreamOutcome } from './ui-message-stream-outcome';
 import { toUIMessageChunk } from './to-ui-message-chunk';
 
 /**
@@ -37,6 +38,26 @@ export function toUIMessageStream<
 } & UIMessageStreamOptions<UI_MESSAGE>): ReadableStream<
   InferUIMessageChunk<UI_MESSAGE>
 > {
+  let outcome: UIMessageStreamOutcome = { status: 'unknown' };
+  let hasFatalFailure = false;
+
+  const setSourceOutcome = (newOutcome: UIMessageStreamOutcome) => {
+    if (
+      !hasFatalFailure &&
+      outcome.status !== 'completed' &&
+      outcome.status !== 'aborted' &&
+      newOutcome.status !== 'unknown' &&
+      (outcome.status === 'unknown' || newOutcome.status !== 'failed')
+    ) {
+      outcome = newOutcome;
+    }
+  };
+
+  const failOutcome = (error: unknown) => {
+    hasFatalFailure = true;
+    outcome = { status: 'failed', error };
+  };
+
   const responseMessageId =
     generateMessageId != null
       ? getResponseUIMessageId({
@@ -45,37 +66,97 @@ export function toUIMessageStream<
         })
       : undefined;
 
-  const uiMessageChunkStream = stream.pipeThrough(
+  const sourceReader = stream.getReader();
+  let sourceReaderReleased = false;
+  let sourceStreamCancelled = false;
+
+  const releaseSourceReader = () => {
+    if (!sourceReaderReleased) {
+      sourceReader.releaseLock();
+      sourceReaderReleased = true;
+    }
+  };
+
+  const sourceStream = new ReadableStream<TextStreamPart<TOOLS>>({
+    async pull(controller) {
+      try {
+        const { done, value } = await sourceReader.read();
+
+        if (done) {
+          releaseSourceReader();
+          if (!sourceStreamCancelled) {
+            controller.close();
+          }
+        } else {
+          controller.enqueue(value);
+        }
+      } catch (error) {
+        releaseSourceReader();
+        if (!sourceStreamCancelled) {
+          failOutcome(error);
+          controller.error(error);
+        }
+      }
+    },
+
+    async cancel(reason) {
+      sourceStreamCancelled = true;
+      if (sourceReaderReleased) {
+        return;
+      }
+
+      try {
+        await sourceReader.cancel(reason);
+      } finally {
+        releaseSourceReader();
+      }
+    },
+  });
+
+  const uiMessageChunkStream = sourceStream.pipeThrough(
     new TransformStream({
       transform: async (part, controller) => {
-        const messageMetadataValue = messageMetadata?.({ part });
-
-        const uiMessageChunk = toUIMessageChunk(part, {
-          tools,
-          sendReasoning,
-          sendSources,
-          sendStart,
-          sendFinish,
-          onError,
-          messageMetadata: messageMetadataValue,
-          responseMessageId,
-        });
-
-        if (uiMessageChunk != null) {
-          controller.enqueue(uiMessageChunk);
-        }
+        try {
+          const messageMetadataValue = messageMetadata?.({ part });
 
-        // start and finish events already include metadata in the converted
-        // chunk; for other part types emit a separate message-metadata chunk
-        if (
-          messageMetadataValue != null &&
-          part.type !== 'start' &&
-          part.type !== 'finish'
-        ) {
-          controller.enqueue({
-            type: 'message-metadata',
+          const uiMessageChunk = toUIMessageChunk(part, {
+            tools,
+            sendReasoning,
+            sendSources,
+            sendStart,
+            sendFinish,
+            onError,
             messageMetadata: messageMetadataValue,
+            responseMessageId,
           });
+
+          if (uiMessageChunk != null) {
+            controller.enqueue(uiMessageChunk);
+          }
+
+          // start and finish events already include metadata in the converted
+          // chunk; for other part types emit a separate message-metadata chunk
+          if (
+            messageMetadataValue != null &&
+            part.type !== 'start' &&
+            part.type !== 'finish'
+          ) {
+            controller.enqueue({
+              type: 'message-metadata',
+              messageMetadata: messageMetadataValue,
+            });
+          }
+
+          if (part.type === 'finish') {
+            setSourceOutcome({ status: 'completed' });
+          } else if (part.type === 'abort') {
+            setSourceOutcome({ status: 'aborted' });
+          } else if (part.type === 'error') {
+            setSourceOutcome({ status: 'failed', error: part.error });
+          }
+        } catch (error) {
+          failOutcome(error);
+          throw error;
         }
       },
     }),
@@ -87,5 +168,6 @@ export function toUIMessageStream<
     originalMessages,
     onEnd: onEnd ?? onFinish,
     onError,
+    getOutcome: () => outcome,
   });
 }

package/src/ui-message-stream/ui-message-stream-on-end-callback.ts

@@ -1,5 +1,6 @@
 import type { FinishReason } from '../types/language-model';
 import type { UIMessage } from '../ui/ui-messages';
+import type { UIMessageStreamOutcome } from './ui-message-stream-outcome';
 
 export type UIMessageStreamOnEndCallback<UI_MESSAGE extends UIMessage> =
   (event: {
@@ -19,6 +20,12 @@ export type UIMessageStreamOnEndCallback<UI_MESSAGE extends UIMessage> =
      */
     isAborted: boolean;
 
+    /**
+     * The operation-level outcome of the stream. Fatal stream-processing
+     * failures override outcomes declared by the stream owner.
+     */
+    outcome: UIMessageStreamOutcome;
+
     /**
      * The message that was sent to the client as a response
      * (including the original message if it was extended).

package/src/ui-message-stream/ui-message-stream-outcome.ts

@@ -0,0 +1,12 @@
+/**
+ * The operation-level outcome of a UI message stream.
+ *
+ * This is separate from model finish reasons and individual stream chunks.
+ * Fatal stream-processing failures override outcomes declared by the stream
+ * owner.
+ */
+export type UIMessageStreamOutcome =
+  | { status: 'completed' }
+  | { status: 'failed'; error?: unknown }
+  | { status: 'aborted' }
+  | { status: 'unknown' };

package/src/ui-message-stream/ui-message-stream-writer.ts

@@ -1,6 +1,7 @@
 import type { UIMessage } from '../ui';
 import type { ErrorHandler } from '../util/error-handler';
 import type { InferUIMessageChunk } from './ui-message-chunks';
+import type { UIMessageStreamOutcome } from './ui-message-stream-outcome';
 
 export interface UIMessageStreamWriter<
   UI_MESSAGE extends UIMessage = UIMessage,
@@ -22,3 +23,17 @@ export interface UIMessageStreamWriter<
    */
   onError: ErrorHandler | undefined;
 }
+
+export interface UIMessageStreamWriterWithOutcome<
+  UI_MESSAGE extends UIMessage = UIMessage,
+> extends UIMessageStreamWriter<UI_MESSAGE> {
+  /**
+   * Declares the operation-level outcome of the composed stream.
+   *
+   * The first outcome declared through this method is retained. Fatal
+   * execution, merge, error-handling, or downstream processing failures
+   * override declared outcomes. Declaring an outcome does not write a chunk or
+   * close the stream.
+   */
+  setOutcome(outcome: UIMessageStreamOutcome): void;
+}

package/src/util/create-stitchable-stream.ts

@@ -25,9 +25,14 @@ export function createStitchableStream<T>(): {
   }> = [];
   let controller: ReadableStreamDefaultController<T> | null = null;
   let isClosed = false;
+  let isCancelled = false;
   let waitForNewStream = createResolvablePromise<void>();
 
   const terminate = () => {
+    if (isCancelled) {
+      return;
+    }
+
     isClosed = true;
     waitForNewStream.resolve();
 
@@ -40,6 +45,10 @@ export function createStitchableStream<T>(): {
   };
 
   const processPull = async () => {
+    if (isCancelled) {
+      return;
+    }
+
     // Case 1: Outer stream is closed and no more inner streams
     if (isClosed && innerStreams.length === 0) {
       controller?.close();
@@ -59,6 +68,10 @@ export function createStitchableStream<T>(): {
     try {
       const { value, done } = await currentStream.reader.read();
 
+      if (isCancelled) {
+        return;
+      }
+
       if (done) {
         // Case 3: Current inner stream is done
         innerStreams.shift(); // Remove the finished stream
@@ -75,6 +88,10 @@ export function createStitchableStream<T>(): {
         controller?.enqueue(value);
       }
     } catch (error) {
+      if (isCancelled) {
+        return;
+      }
+
       // Case 5: Current inner stream throws an error
       currentStream.onError?.(error);
       controller?.error(error);
@@ -90,12 +107,15 @@ export function createStitchableStream<T>(): {
       },
       pull: processPull,
       async cancel() {
+        isCancelled = true;
+        isClosed = true;
+        waitForNewStream.resolve();
+
         for (const { reader, onCancel } of innerStreams) {
           onCancel?.();
           await reader.cancel();
         }
         innerStreams = [];
-        isClosed = true;
       },
     }),
     addStream: (
@@ -105,6 +125,12 @@ export function createStitchableStream<T>(): {
         onCancel?: () => void;
       },
     ) => {
+      if (isCancelled) {
+        callbacks?.onCancel?.();
+        void innerStream.cancel().catch(() => {});
+        return;
+      }
+
       if (isClosed) {
         throw new Error('Cannot add inner stream: outer stream is closed');
       }
@@ -121,6 +147,10 @@ export function createStitchableStream<T>(): {
      * finish processing and then close the outer stream.
      */
     close: () => {
+      if (isCancelled) {
+        return;
+      }
+
       isClosed = true;
       waitForNewStream.resolve();