ai 7.0.79 → 7.0.83

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

@@ -6,9 +6,10 @@ import type { UIMessage } from '../ui/ui-messages';
 import { handleUIMessageStreamFinish } from './handle-ui-message-stream-finish';
 import type { InferUIMessageChunk } from './ui-message-chunks';
 import type { UIMessageStreamOnEndCallback } from './ui-message-stream-on-end-callback';
+import type { UIMessageStreamOutcome } from './ui-message-stream-outcome';
 import type { UIMessageStreamOnStepEndCallback } from './ui-message-stream-on-step-end-callback';
 import type { UIMessageStreamOnStepFinishCallback } from './ui-message-stream-on-step-finish-callback';
-import type { UIMessageStreamWriter } from './ui-message-stream-writer';
+import type { UIMessageStreamWriterWithOutcome } from './ui-message-stream-writer';
 
 /**
  * Creates a UI message stream that can be used to send messages to the client.
@@ -36,7 +37,7 @@ export function createUIMessageStream<UI_MESSAGE extends UIMessage>({
   generateId = generateIdFunc,
 }: {
   execute: (options: {
-    writer: UIMessageStreamWriter<UI_MESSAGE>;
+    writer: UIMessageStreamWriterWithOutcome<UI_MESSAGE>;
   }) => Promise<void> | void;
   onError?: (error: unknown) => string;
 
@@ -72,6 +73,7 @@ export function createUIMessageStream<UI_MESSAGE extends UIMessage>({
   >;
 
   const ongoingStreamPromises: Promise<void>[] = [];
+  let outcome: UIMessageStreamOutcome = { status: 'unknown' };
 
   const stream = new ReadableStream({
     start(controllerArg) {
@@ -87,6 +89,42 @@ export function createUIMessageStream<UI_MESSAGE extends UIMessage>({
     }
   }
 
+  function setOutcome(newOutcome: UIMessageStreamOutcome) {
+    if (outcome.status === 'unknown' && newOutcome.status !== 'unknown') {
+      outcome = newOutcome;
+    }
+  }
+
+  function failOutcome(error: unknown) {
+    outcome = { status: 'failed', error };
+  }
+
+  function safeError(error: unknown) {
+    try {
+      controller.error(error);
+    } catch {
+      // suppress errors when the stream has been closed
+    }
+  }
+
+  function handleError(error: unknown) {
+    failOutcome(error);
+
+    let errorText: string;
+    try {
+      errorText = onError(error);
+    } catch (onErrorError) {
+      failOutcome(onErrorError);
+      safeError(onErrorError);
+      return;
+    }
+
+    safeEnqueue({
+      type: 'error',
+      errorText,
+    } as InferUIMessageChunk<UI_MESSAGE>);
+  }
+
   try {
     const result = execute({
       writer: {
@@ -103,13 +141,11 @@ export function createUIMessageStream<UI_MESSAGE extends UIMessage>({
                 safeEnqueue(value);
               }
             })().catch(error => {
-              safeEnqueue({
-                type: 'error',
-                errorText: onError(error),
-              } as InferUIMessageChunk<UI_MESSAGE>);
+              handleError(error);
             }),
           );
         },
+        setOutcome,
         onError,
       },
     });
@@ -117,30 +153,23 @@ export function createUIMessageStream<UI_MESSAGE extends UIMessage>({
     if (result) {
       ongoingStreamPromises.push(
         result.catch(error => {
-          safeEnqueue({
-            type: 'error',
-            errorText: onError(error),
-          } as InferUIMessageChunk<UI_MESSAGE>);
+          handleError(error);
         }),
       );
     }
   } catch (error) {
-    safeEnqueue({
-      type: 'error',
-      errorText: onError(error),
-    } as InferUIMessageChunk<UI_MESSAGE>);
+    handleError(error);
   }
 
   // Wait until all ongoing streams are done. This approach enables merging
   // streams even after execute has returned, as long as there is still an
   // open merged stream. This is important to e.g. forward new streams and
   // from callbacks.
-  const waitForStreams: Promise<void> = new Promise(async resolve => {
+  const waitForStreams: Promise<void> = (async () => {
     while (ongoingStreamPromises.length > 0) {
       await ongoingStreamPromises.shift();
     }
-    resolve();
-  });
+  })();
 
   waitForStreams.finally(() => {
     try {
@@ -157,5 +186,6 @@ export function createUIMessageStream<UI_MESSAGE extends UIMessage>({
     onStepEnd: onStepEnd ?? onStepFinish,
     onEnd: onEnd ?? onFinish,
     onError,
+    getOutcome: () => outcome,
   });
 }

package/src/ui-message-stream/handle-ui-message-stream-finish.ts

@@ -8,6 +8,7 @@ import type { UIMessage } from '../ui/ui-messages';
 import type { ErrorHandler } from '../util/error-handler';
 import type { InferUIMessageChunk, UIMessageChunk } from './ui-message-chunks';
 import type { UIMessageStreamOnEndCallback } from './ui-message-stream-on-end-callback';
+import type { UIMessageStreamOutcome } from './ui-message-stream-outcome';
 import type { UIMessageStreamOnStepEndCallback } from './ui-message-stream-on-step-end-callback';
 import type { UIMessageStreamOnStepFinishCallback } from './ui-message-stream-on-step-finish-callback';
 
@@ -20,6 +21,7 @@ export function handleUIMessageStreamFinish<UI_MESSAGE extends UIMessage>({
   onFinish,
   onError,
   stream,
+  getOutcome,
 }: {
   stream: ReadableStream<InferUIMessageChunk<UI_MESSAGE>>;
 
@@ -54,6 +56,11 @@ export function handleUIMessageStreamFinish<UI_MESSAGE extends UIMessage>({
    * @deprecated Use `onEnd` instead.
    */
   onFinish?: UIMessageStreamOnEndCallback<UI_MESSAGE>;
+
+  /**
+   * Returns the operation-level outcome declared by the stream owner.
+   */
+  getOutcome?: () => UIMessageStreamOutcome;
 }): ReadableStream<InferUIMessageChunk<UI_MESSAGE>> {
   // last message is only relevant for assistant messages
   let lastMessage: UI_MESSAGE | undefined =
@@ -66,6 +73,13 @@ export function handleUIMessageStreamFinish<UI_MESSAGE extends UIMessage>({
   }
 
   let isAborted = false;
+  let hasProcessingFailure = false;
+  let processingError: unknown;
+
+  const recordProcessingFailure = (error: unknown) => {
+    hasProcessingFailure = true;
+    processingError = error;
+  };
 
   const idInjectedStream = stream.pipeThrough(
     new TransformStream<
@@ -73,21 +87,31 @@ export function handleUIMessageStreamFinish<UI_MESSAGE extends UIMessage>({
       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;
+        try {
+          let outputChunk = chunk;
+
+          // 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) {
+              outputChunk = {
+                ...startChunk,
+                messageId,
+              } as InferUIMessageChunk<UI_MESSAGE>;
+            }
           }
-        }
 
-        if (chunk.type === 'abort') {
-          isAborted = true;
-        }
+          if (chunk.type === 'abort') {
+            isAborted = true;
+          }
 
-        controller.enqueue(chunk);
+          controller.enqueue(outputChunk);
+        } catch (error) {
+          recordProcessingFailure(error);
+          throw error;
+        }
       },
     }),
   );
@@ -113,7 +137,12 @@ export function handleUIMessageStreamFinish<UI_MESSAGE extends UIMessage>({
       write: (options?: UIMessageStreamWriteOptions) => void;
     }) => Promise<void>,
   ) => {
-    await job({ state, write: () => {} });
+    try {
+      await job({ state, write: () => {} });
+    } catch (error) {
+      recordProcessingFailure(error);
+      throw error;
+    }
   };
 
   let finishCalled = false;
@@ -125,9 +154,17 @@ export function handleUIMessageStreamFinish<UI_MESSAGE extends UIMessage>({
     finishCalled = true;
 
     const isContinuation = state.message.id === lastMessage?.id;
+    const declaredOutcome = getOutcome?.() ?? { status: 'unknown' };
+    const outcome: UIMessageStreamOutcome = hasProcessingFailure
+      ? { status: 'failed', error: processingError }
+      : declaredOutcome.status === 'unknown' && isAborted
+        ? { status: 'aborted' }
+        : declaredOutcome;
+
     await resolvedOnEnd({
-      isAborted,
+      isAborted: isAborted || outcome.status === 'aborted',
       isContinuation,
+      outcome,
       responseMessage: state.message as UI_MESSAGE,
       messages: [
         ...(isContinuation ? originalMessages.slice(0, -1) : originalMessages),

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

@@ -16,6 +16,10 @@ export {
 export { UI_MESSAGE_STREAM_HEADERS } from './ui-message-stream-headers';
 export type { UIMessageStreamOnEndCallback } from './ui-message-stream-on-end-callback';
 export type { UIMessageStreamOnFinishCallback } from './ui-message-stream-on-finish-callback';
+export type { UIMessageStreamOutcome } from './ui-message-stream-outcome';
 export type { UIMessageStreamOnStepEndCallback } from './ui-message-stream-on-step-end-callback';
 export type { UIMessageStreamOnStepFinishCallback } from './ui-message-stream-on-step-finish-callback';
-export type { UIMessageStreamWriter } from './ui-message-stream-writer';
+export type {
+  UIMessageStreamWriter,
+  UIMessageStreamWriterWithOutcome,
+} from './ui-message-stream-writer';

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

@@ -254,6 +254,7 @@ export function toUIMessageChunk<
         type: 'tool-approval-request',
         approvalId: part.approvalId,
         toolCallId: part.toolCall.toolCallId,
+        ...(part.reason != null ? { reason: part.reason } : {}),
         ...(part.isAutomatic != null ? { isAutomatic: part.isAutomatic } : {}),
         ...(part.signature != null ? { signature: part.signature } : {}),
       };

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-chunks.ts

@@ -85,6 +85,7 @@ export const uiMessageChunkSchema = lazySchema(() =>
         type: z.literal('tool-approval-request'),
         approvalId: z.string(),
         toolCallId: z.string(),
+        reason: z.string().optional(),
         isAutomatic: z.boolean().optional(),
         signature: z.string().optional(),
       }),
@@ -298,6 +299,7 @@ export type UIMessageChunk<
       type: 'tool-approval-request';
       approvalId: string;
       toolCallId: string;
+      reason?: string;
       isAutomatic?: boolean;
       signature?: string;
     }

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();