ai 6.0.270 → 6.0.272

package/src/ui/validate-ui-messages.ts

@@ -5,6 +5,7 @@ import {
 } from '@ai-sdk/provider';
 import {
   lazySchema,
+  safeValidateTypes,
   StandardSchemaV1,
   validateTypes,
   zodSchema,
@@ -17,6 +18,7 @@ import { jsonValueSchema } from '../types/json-value';
 import { providerMetadataSchema } from '../types/provider-metadata';
 import type {
   DataUIPart,
+  DynamicToolUIPart,
   InferUIMessageData,
   InferUIMessageTools,
   ToolUIPart,
@@ -28,6 +30,25 @@ const toolMetadataSchema: z.ZodType<JSONObject> = z.record(
   jsonValueSchema.optional(),
 );
 
+function isEmptyObject(value: unknown): value is Record<string, never> {
+  return (
+    value != null &&
+    typeof value === 'object' &&
+    !Array.isArray(value) &&
+    Object.keys(value).length === 0
+  );
+}
+
+function asDynamicToolPart(toolPart: ToolUIPart): DynamicToolUIPart {
+  const { type, ...part } = toolPart;
+
+  return {
+    ...part,
+    type: 'dynamic-tool',
+    toolName: type.slice(5),
+  } as DynamicToolUIPart;
+}
+
 const uiMessagesSchema = lazySchema(() =>
   zodSchema(
     z
@@ -354,17 +375,7 @@ export type SafeValidateUIMessagesResult<UI_MESSAGE extends UIMessage> =
       error: Error;
     };
 
-/**
- * Validates a list of UI messages like `validateUIMessages`,
- * but instead of throwing it returns `{ success: true, data }`
- * or `{ success: false, error }`.
- */
-export async function safeValidateUIMessages<UI_MESSAGE extends UIMessage>({
-  messages,
-  metadataSchema,
-  dataSchemas,
-  tools,
-}: {
+type ValidateUIMessagesOptions<UI_MESSAGE extends UIMessage> = {
   messages: unknown;
   metadataSchema?: FlexibleSchema<UIMessage['metadata']>;
   dataSchemas?: {
@@ -378,7 +389,21 @@ export async function safeValidateUIMessages<UI_MESSAGE extends UIMessage>({
       InferUIMessageTools<UI_MESSAGE>[NAME]['output']
     >;
   };
-}): Promise<SafeValidateUIMessagesResult<UI_MESSAGE>> {
+};
+
+async function safeValidateUIMessagesInternal<UI_MESSAGE extends UIMessage>(
+  {
+    messages,
+    metadataSchema,
+    dataSchemas,
+    tools,
+  }: ValidateUIMessagesOptions<UI_MESSAGE>,
+  {
+    convertMissingTerminalToolsToDynamic,
+  }: {
+    convertMissingTerminalToolsToDynamic: boolean;
+  },
+): Promise<SafeValidateUIMessagesResult<UI_MESSAGE>> {
   try {
     if (messages == null) {
       return {
@@ -409,7 +434,10 @@ export async function safeValidateUIMessages<UI_MESSAGE extends UIMessage>({
       }
     }
 
-    if (dataSchemas || tools) {
+    const shouldValidateToolParts =
+      tools != null || convertMissingTerminalToolsToDynamic;
+
+    if (dataSchemas || shouldValidateToolParts) {
       for (const [msgIdx, message] of validatedMessages.entries()) {
         for (const [partIdx, part] of message.parts.entries()) {
           // Data part validation
@@ -445,19 +473,26 @@ export async function safeValidateUIMessages<UI_MESSAGE extends UIMessage>({
           }
 
           // Tool part validation
-          if (tools && part.type.startsWith('tool-')) {
+          if (shouldValidateToolParts && part.type.startsWith('tool-')) {
             const toolPart = part as ToolUIPart<
               InferUIMessageTools<UI_MESSAGE>
             >;
             const toolName = toolPart.type.slice(5);
-            const tool = tools[toolName];
+            const tool = tools?.[toolName];
+            const isTerminal =
+              toolPart.state === 'output-available' ||
+              toolPart.state === 'output-error' ||
+              toolPart.state === 'output-denied';
 
-            if (
-              !tool &&
-              (toolPart.state === 'output-available' ||
-                toolPart.state === 'output-error' ||
-                toolPart.state === 'output-denied')
-            ) {
+            if (!tool && isTerminal) {
+              if (tools != null || convertMissingTerminalToolsToDynamic) {
+                // Persisted terminal history can reference tools that are no
+                // longer registered. Normalize those parts so callers do not
+                // receive unvalidated values under current static tool types.
+                message.parts[partIdx] = asDynamicToolPart(
+                  toolPart,
+                ) as (typeof message.parts)[number];
+              }
               continue;
             }
 
@@ -477,19 +512,53 @@ export async function safeValidateUIMessages<UI_MESSAGE extends UIMessage>({
               };
             }
 
+            const inputValidationContext = {
+              field: `messages[${msgIdx}].parts[${partIdx}].input`,
+              entityName: toolName,
+              entityId: toolPart.toolCallId,
+            };
+            let convertToDynamic = false;
+
             // Tool input validation
-            // Note: input is intentionally not re-validated for terminal states.
-            // Terminal tool calls can keep invalid or incomplete input, and
-            // re-validating it on replay would crash follow-up messages.
-            if (toolPart.state === 'input-available') {
+            if (toolPart.state === 'output-error') {
+              // Failed calls can retain invalid input. Keep them loadable, but
+              // expose incompatible input as unknown instead of the current
+              // static tool input type.
+              if (toolPart.input !== undefined) {
+                const result = await safeValidateTypes({
+                  value: toolPart.input,
+                  schema: tool.inputSchema,
+                  context: inputValidationContext,
+                });
+                convertToDynamic = !result.success;
+              }
+            } else if (toolPart.state === 'output-available') {
+              const result = await safeValidateTypes({
+                value: toolPart.input,
+                schema: tool.inputSchema,
+                context: inputValidationContext,
+              });
+
+              if (!result.success) {
+                // Empty terminal input can represent aborted or incomplete
+                // history whose input was never streamed. Preserve it without
+                // claiming that it matches the current static input type.
+                if (isEmptyObject(toolPart.input)) {
+                  convertToDynamic = true;
+                } else {
+                  throw result.error;
+                }
+              }
+            } else if (
+              toolPart.state === 'input-available' ||
+              toolPart.state === 'approval-requested' ||
+              toolPart.state === 'approval-responded' ||
+              toolPart.state === 'output-denied'
+            ) {
               await validateTypes({
                 value: toolPart.input,
                 schema: tool.inputSchema,
-                context: {
-                  field: `messages[${msgIdx}].parts[${partIdx}].input`,
-                  entityName: toolName,
-                  entityId: toolPart.toolCallId,
-                },
+                context: inputValidationContext,
               });
             }
 
@@ -505,6 +574,12 @@ export async function safeValidateUIMessages<UI_MESSAGE extends UIMessage>({
                 },
               });
             }
+
+            if (convertToDynamic) {
+              message.parts[partIdx] = asDynamicToolPart(
+                toolPart,
+              ) as (typeof message.parts)[number];
+            }
           }
         }
       }
@@ -524,6 +599,19 @@ export async function safeValidateUIMessages<UI_MESSAGE extends UIMessage>({
   }
 }
 
+/**
+ * Validates a list of UI messages like `validateUIMessages`,
+ * but instead of throwing it returns `{ success: true, data }`
+ * or `{ success: false, error }`.
+ */
+export async function safeValidateUIMessages<UI_MESSAGE extends UIMessage>(
+  options: ValidateUIMessagesOptions<UI_MESSAGE>,
+): Promise<SafeValidateUIMessagesResult<UI_MESSAGE>> {
+  return safeValidateUIMessagesInternal(options, {
+    convertMissingTerminalToolsToDynamic: false,
+  });
+}
+
 /**
  * Validates a list of UI messages.
  *
@@ -531,31 +619,24 @@ export async function safeValidateUIMessages<UI_MESSAGE extends UIMessage>({
  * the corresponding schemas are provided. Otherwise, they are assumed to be
  * valid.
  */
-export async function validateUIMessages<UI_MESSAGE extends UIMessage>({
-  messages,
-  metadataSchema,
-  dataSchemas,
-  tools,
-}: {
-  messages: unknown;
-  metadataSchema?: FlexibleSchema<UIMessage['metadata']>;
-  dataSchemas?: {
-    [NAME in keyof InferUIMessageData<UI_MESSAGE> & string]?: FlexibleSchema<
-      InferUIMessageData<UI_MESSAGE>[NAME]
-    >;
-  };
-  tools?: {
-    [NAME in keyof InferUIMessageTools<UI_MESSAGE> & string]?: Tool<
-      InferUIMessageTools<UI_MESSAGE>[NAME]['input'],
-      InferUIMessageTools<UI_MESSAGE>[NAME]['output']
-    >;
-  };
-}): Promise<Array<UI_MESSAGE>> {
-  const response = await safeValidateUIMessages({
-    messages,
-    metadataSchema,
-    dataSchemas,
-    tools,
+export async function validateUIMessages<UI_MESSAGE extends UIMessage>(
+  options: ValidateUIMessagesOptions<UI_MESSAGE>,
+): Promise<Array<UI_MESSAGE>> {
+  const response = await safeValidateUIMessages(options);
+
+  if (!response.success) throw response.error;
+
+  return response.data;
+}
+
+export async function validateUIMessagesForAgent<UI_MESSAGE extends UIMessage>(
+  options: ValidateUIMessagesOptions<UI_MESSAGE>,
+): Promise<Array<UI_MESSAGE>> {
+  const response = await safeValidateUIMessagesInternal(options, {
+    // Agent tool sets can include ephemeral tools (for example, tools from a
+    // disconnected MCP server), so terminal history is converted to dynamic
+    // tool parts when those tools are no longer registered.
+    convertMissingTerminalToolsToDynamic: true,
   });
 
   if (!response.success) throw response.error;

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

@@ -6,8 +6,9 @@ import type { UIMessage } from '../ui/ui-messages';
 import { handleUIMessageStreamFinish } from './handle-ui-message-stream-finish';
 import type { InferUIMessageChunk } from './ui-message-chunks';
 import type { UIMessageStreamOnFinishCallback } from './ui-message-stream-on-finish-callback';
+import type { UIMessageStreamOutcome } from './ui-message-stream-outcome';
 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.
@@ -31,7 +32,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;
 
@@ -55,6 +56,7 @@ export function createUIMessageStream<UI_MESSAGE extends UIMessage>({
   >;
 
   const ongoingStreamPromises: Promise<void>[] = [];
+  let outcome: UIMessageStreamOutcome = { status: 'unknown' };
 
   const stream = new ReadableStream({
     start(controllerArg) {
@@ -70,6 +72,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: {
@@ -86,13 +124,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,
       },
     });
@@ -100,30 +136,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 {
@@ -140,5 +169,6 @@ export function createUIMessageStream<UI_MESSAGE extends UIMessage>({
     onStepFinish,
     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 { UIMessageStreamOnFinishCallback } from './ui-message-stream-on-finish-callback';
+import type { UIMessageStreamOutcome } from './ui-message-stream-outcome';
 import type { UIMessageStreamOnStepFinishCallback } from './ui-message-stream-on-step-finish-callback';
 
 export function handleUIMessageStreamFinish<UI_MESSAGE extends UIMessage>({
@@ -17,6 +18,7 @@ export function handleUIMessageStreamFinish<UI_MESSAGE extends UIMessage>({
   onFinish,
   onError,
   stream,
+  getOutcome,
 }: {
   stream: ReadableStream<InferUIMessageChunk<UI_MESSAGE>>;
 
@@ -39,6 +41,11 @@ export function handleUIMessageStreamFinish<UI_MESSAGE extends UIMessage>({
   onStepFinish?: UIMessageStreamOnStepFinishCallback<UI_MESSAGE>;
 
   onFinish?: UIMessageStreamOnFinishCallback<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 =
@@ -51,6 +58,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<
@@ -58,21 +72,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;
+        }
       },
     }),
   );
@@ -95,7 +119,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;
@@ -107,9 +136,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 onFinish({
-      isAborted,
+      isAborted: isAborted || outcome.status === 'aborted',
       isContinuation,
+      outcome,
       responseMessage: state.message as UI_MESSAGE,
       messages: [
         ...(isContinuation ? originalMessages.slice(0, -1) : originalMessages),
@@ -138,7 +175,13 @@ export function handleUIMessageStreamFinish<UI_MESSAGE extends UIMessage>({
         ] as UI_MESSAGE[],
       });
     } catch (error) {
-      onError(error);
+      try {
+        onError(error);
+      } catch (onErrorError) {
+        recordProcessingFailure(onErrorError);
+        await callOnFinish();
+        throw onErrorError;
+      }
     }
   };
 

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

@@ -10,5 +10,9 @@ export {
 } from './ui-message-chunks';
 export { UI_MESSAGE_STREAM_HEADERS } from './ui-message-stream-headers';
 export type { UIMessageStreamOnFinishCallback } from './ui-message-stream-on-finish-callback';
+export type { UIMessageStreamOutcome } from './ui-message-stream-outcome';
 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/ui-message-stream-on-finish-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 UIMessageStreamOnFinishCallback<UI_MESSAGE extends UIMessage> =
   (event: {
@@ -19,6 +20,12 @@ export type UIMessageStreamOnFinishCallback<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;
+}