ai 7.0.134 → 7.0.136

package/src/embed/embed-many.ts

@@ -20,6 +20,7 @@ import { notify } from '../util/notify';
 import { prepareRetries } from '../util/prepare-retries';
 import { splitArray } from '../util/split-array';
 import type { EmbedEndEvent, EmbedStartEvent } from './embed-events';
+import { validateEmbeddingDimensions } from './validate-embedding-dimensions';
 import type { EmbedManyResult } from './embed-many-result';
 import { VERSION } from '../version';
 
@@ -39,6 +40,8 @@ const originalGenerateCallId = createIdGenerator({
  * @param model - The embedding model to use.
  * @param values - The values that should be embedded.
  *
+ * @param dimensions - Requested output dimensions. Must be a positive integer. Requires provider support.
+ *
  * @param maxRetries - Maximum number of retries. Set to 0 to disable retries. Default: 2.
  * @param abortSignal - An optional abort signal that can be used to cancel the call.
  * @param headers - Additional HTTP headers to be sent with the request. Only applicable for HTTP-based providers.
@@ -59,6 +62,7 @@ const originalGenerateCallId = createIdGenerator({
 export async function embedMany<RUNTIME_CONTEXT extends Context = Context>({
   model: modelArg,
   values,
+  dimensions,
   maxParallelCalls = Infinity,
   maxRetries: maxRetriesArg,
   abortSignal,
@@ -83,6 +87,13 @@ export async function embedMany<RUNTIME_CONTEXT extends Context = Context>({
    */
   values: Array<string>;
 
+  /**
+   * The requested number of dimensions for the output embeddings.
+   * Must be a positive integer. Support and allowed values depend on the model
+   * and provider implementation.
+   */
+  dimensions?: number;
+
   /**
    * Maximum number of retries per embedding model call. Set to 0 to disable retries.
    *
@@ -168,6 +179,8 @@ export async function embedMany<RUNTIME_CONTEXT extends Context = Context>({
     generateCallId?: () => string;
   };
 }): Promise<EmbedManyResult> {
+  validateEmbeddingDimensions(dimensions);
+
   const model = resolveEmbeddingModel(modelArg);
 
   const { maxRetries, retry } = prepareRetries({
@@ -200,6 +213,7 @@ export async function embedMany<RUNTIME_CONTEXT extends Context = Context>({
     provider: model.provider,
     modelId: model.modelId,
     value: values,
+    dimensions,
     maxRetries,
     headers: headersWithUserAgent,
     providerOptions,
@@ -243,12 +257,14 @@ export async function embedMany<RUNTIME_CONTEXT extends Context = Context>({
                   provider: model.provider,
                   modelId: model.modelId,
                   values,
+                  dimensions,
                 },
                 callbacks: [telemetryDispatcher.onEmbedStart],
               });
 
               const modelResponse = await model.doEmbed({
                 values,
+                dimensions,
                 abortSignal,
                 headers: headersWithUserAgent,
                 providerOptions,
@@ -371,12 +387,14 @@ export async function embedMany<RUNTIME_CONTEXT extends Context = Context>({
                     provider: model.provider,
                     modelId: model.modelId,
                     values: chunk,
+                    dimensions,
                   },
                   callbacks: [telemetryDispatcher.onEmbedStart],
                 });
 
                 const modelResponse = await model.doEmbed({
                   values: chunk,
+                  dimensions,
                   abortSignal,
                   headers: headersWithUserAgent,
                   providerOptions: chunkProviderOptions,

package/src/embed/embed.ts

@@ -15,6 +15,7 @@ import { notify } from '../util/notify';
 import { prepareRetries } from '../util/prepare-retries';
 import { VERSION } from '../version';
 import type { EmbedEndEvent, EmbedStartEvent } from './embed-events';
+import { validateEmbeddingDimensions } from './validate-embedding-dimensions';
 import type { EmbedResult } from './embed-result';
 
 const originalGenerateCallId = createIdGenerator({
@@ -28,6 +29,8 @@ const originalGenerateCallId = createIdGenerator({
  * @param model - The embedding model to use.
  * @param value - The value that should be embedded.
  *
+ * @param dimensions - Requested output dimensions. Must be a positive integer. Requires provider support.
+ *
  * @param maxRetries - Maximum number of retries. Set to 0 to disable retries. Default: 2.
  * @param abortSignal - An optional abort signal that can be used to cancel the call.
  * @param headers - Additional HTTP headers to be sent with the request. Only applicable for HTTP-based providers.
@@ -44,6 +47,7 @@ const originalGenerateCallId = createIdGenerator({
 export async function embed<RUNTIME_CONTEXT extends Context = Context>({
   model: modelArg,
   value,
+  dimensions,
   providerOptions,
   maxRetries: maxRetriesArg,
   abortSignal,
@@ -67,6 +71,13 @@ export async function embed<RUNTIME_CONTEXT extends Context = Context>({
    */
   value: string;
 
+  /**
+   * The requested number of dimensions for the output embeddings.
+   * Must be a positive integer. Support and allowed values depend on the model
+   * and provider implementation.
+   */
+  dimensions?: number;
+
   /**
    * Maximum number of retries per embedding model call. Set to 0 to disable retries.
    *
@@ -144,6 +155,8 @@ export async function embed<RUNTIME_CONTEXT extends Context = Context>({
     generateCallId?: () => string;
   };
 }): Promise<EmbedResult> {
+  validateEmbeddingDimensions(dimensions);
+
   const model = resolveEmbeddingModel(modelArg);
 
   const { maxRetries, retry } = prepareRetries({
@@ -176,6 +189,7 @@ export async function embed<RUNTIME_CONTEXT extends Context = Context>({
     provider: model.provider,
     modelId: model.modelId,
     value,
+    dimensions,
     maxRetries,
     headers: headersWithUserAgent,
     providerOptions,
@@ -203,12 +217,14 @@ export async function embed<RUNTIME_CONTEXT extends Context = Context>({
                 provider: model.provider,
                 modelId: model.modelId,
                 values: [value],
+                dimensions,
               },
               callbacks: [telemetryDispatcher.onEmbedStart],
             });
 
             const modelResponse = await model.doEmbed({
               values: [value],
+              dimensions,
               abortSignal,
               headers: headersWithUserAgent,
               providerOptions,

package/src/embed/validate-embedding-dimensions.ts

@@ -0,0 +1,14 @@
+import { InvalidArgumentError } from '../error/invalid-argument-error';
+
+export function validateEmbeddingDimensions(dimensions: number | undefined) {
+  if (
+    dimensions !== undefined &&
+    (!Number.isInteger(dimensions) || dimensions <= 0)
+  ) {
+    throw new InvalidArgumentError({
+      parameter: 'dimensions',
+      value: dimensions,
+      message: 'dimensions must be a positive integer',
+    });
+  }
+}

package/src/generate-text/stream-text.ts

@@ -2275,6 +2275,9 @@ class DefaultStreamTextResult<
           if (chunkTimeoutId != null) {
             clearTimeout(chunkTimeoutId);
           }
+          if (abortSignal?.aborted) {
+            return;
+          }
           chunkTimeoutId = setAbortTimeout({
             abortController: chunkAbortController,
             label: 'Chunk',
@@ -2289,6 +2292,11 @@ class DefaultStreamTextResult<
           }
         }
 
+        function clearModelOutputTimeouts() {
+          clearFirstChunkTimeout();
+          clearChunkTimeout();
+        }
+
         function clearStepTimeout() {
           if (stepTimeoutId != null) {
             clearTimeout(stepTimeoutId);
@@ -2297,8 +2305,7 @@ class DefaultStreamTextResult<
 
         function clearStepTimeouts() {
           clearStepTimeout();
-          clearFirstChunkTimeout();
-          clearChunkTimeout();
+          clearModelOutputTimeouts();
         }
 
         function cleanupStepTimeouts() {
@@ -2509,7 +2516,6 @@ class DefaultStreamTextResult<
           let automaticStreamRetryCount = 0;
           let callbackStreamRetryCount = 0;
           let bufferedAttemptParts: LanguageModelStreamPart<TOOLS>[] = [];
-          const outputChunksHandledBeforeBuffering = new WeakSet<object>();
           const openTextParts = new Set<string>();
           const openReasoningParts = new Set<string>();
           let enqueueStreamRetryAttemptBoundary = false;
@@ -2561,7 +2567,16 @@ class DefaultStreamTextResult<
               };
 
               while (true) {
-                const { done, value } = await languageModelStreamReader.read();
+                let result: ReadableStreamReadResult<
+                  LanguageModelStreamPart<TOOLS>
+                >;
+                try {
+                  result = await languageModelStreamReader.read();
+                } catch (error) {
+                  clearModelOutputTimeouts();
+                  throw error;
+                }
+                const { done, value } = result;
 
                 if (enqueueStreamRetryAttemptBoundary) {
                   controller.enqueue(
@@ -2576,6 +2591,7 @@ class DefaultStreamTextResult<
                 }
 
                 if (done) {
+                  clearModelOutputTimeouts();
                   flushBufferedAttemptParts();
                   controller.close();
                   return;
@@ -2592,22 +2608,24 @@ class DefaultStreamTextResult<
                   value.type === 'tool-error';
 
                 if (value.type === 'model-call-end') {
+                  clearModelOutputTimeouts();
                   flushBufferedAttemptParts();
                   enqueueAttemptPart(value);
                   return;
                 }
 
+                // Observe model output before buffering and tool processing.
+                // Replaying buffered parts must not re-arm finished timers.
+                if (isOutputChunk(value)) {
+                  clearFirstChunkTimeout();
+                  resetChunkTimeout();
+                }
+
                 if (
                   shouldBufferToolParts &&
                   value.type !== 'error' &&
                   (isToolPart || bufferedAttemptParts.length > 0)
                 ) {
-                  if (isOutputChunk(value)) {
-                    clearFirstChunkTimeout();
-                    resetChunkTimeout();
-                    outputChunksHandledBeforeBuffering.add(value);
-                  }
-
                   bufferedAttemptParts.push(value);
                   continue;
                 }
@@ -2617,6 +2635,7 @@ class DefaultStreamTextResult<
                   return;
                 }
 
+                clearModelOutputTimeouts();
                 await notify({
                   event: { chunk: value },
                   callbacks: onChunk,
@@ -2677,10 +2696,12 @@ class DefaultStreamTextResult<
                 response = retryLanguageModelCall.response;
                 languageModelStreamReader =
                   retryLanguageModelCall.stream.getReader();
+                startFirstChunkTimeout();
                 enqueueStreamRetryAttemptBoundary = true;
               }
             },
             cancel(reason) {
+              clearModelOutputTimeouts();
               return languageModelStreamReader.cancel(reason);
             },
           });
@@ -2857,21 +2878,7 @@ class DefaultStreamTextResult<
                   const chunkType = chunk.type;
 
                   if (isOutputChunk(chunk)) {
-                    const timeoutHandledBeforeBuffering =
-                      outputChunksHandledBeforeBuffering.has(chunk);
-
-                    if (
-                      !hasReceivedOutputChunk &&
-                      !timeoutHandledBeforeBuffering
-                    ) {
-                      // Clear before forwarding the first output so a timeout
-                      // cannot race with already-visible generated content.
-                      clearFirstChunkTimeout();
-                    }
                     hasReceivedOutputChunk = true;
-                    if (!timeoutHandledBeforeBuffering) {
-                      resetChunkTimeout();
-                    }
                   }
 
                   switch (chunkType) {

package/src/prompt/request-options.ts

@@ -4,9 +4,9 @@ import type { ToolSet } from '@ai-sdk/provider-utils';
  * Timeout configuration for API calls. Can be specified as:
  * - A number representing milliseconds
  * - An object with `totalMs` property for the total timeout in milliseconds
- * - An object with `stepMs` property for the timeout of each step in milliseconds
+ * - An object with `stepMs` property for the timeout of each complete step, including model generation, retries, and local tool execution, in milliseconds
  * - An object with `firstChunkMs` property for the timeout until the first content chunk of each step (streaming only)
- * - An object with `chunkMs` property for the timeout between content chunks (streaming only)
+ * - An object with `chunkMs` property for the timeout between meaningful model output chunks, starting with the first output and ending when the model response ends (streaming only)
  * - An object with `toolMs` property for the default timeout for all tool executions
  * - An object with `tools` property for per-tool timeout overrides using `{toolName}Ms` keys
  */

package/src/ui/chat-transport.ts

@@ -13,6 +13,22 @@ import type { UIMessage } from './ui-messages';
  * @template UI_MESSAGE - The UI message type extending UIMessage
  */
 export interface ChatTransport<UI_MESSAGE extends UIMessage> {
+  /**
+   * Whether reconnecting returns the complete response from its beginning.
+   * Complete replays must reset the streaming message state before processing
+   * the first start chunk.
+   */
+  readonly resumeStreamIsReplay?: boolean;
+
+  /**
+   * Releases persistent resources owned by the transport.
+   *
+   * The caller that creates a transport owns its lifecycle and should call
+   * this method when no consumers remain. `AbstractChat.dispose()` calls it
+   * for directly managed chat instances.
+   */
+  close?: () => void;
+
   /**
    * Sends messages to the chat API endpoint and returns a streaming response.
    *

package/src/ui/chat.ts

@@ -8,6 +8,7 @@ import { InvalidArgumentError } from '../error/invalid-argument-error';
 import type { FinishReason } from '../types/language-model';
 import type { UIMessageChunk } from '../ui-message-stream/ui-message-chunks';
 import { consumeStream } from '../util/consume-stream';
+import { createResolvablePromise } from '../util/create-resolvable-promise';
 import { SerialJobExecutor } from '../util/serial-job-executor';
 import type { ChatTransport } from './chat-transport';
 import { convertFileListToFileUIParts } from './convert-file-list-to-file-ui-parts';
@@ -135,12 +136,19 @@ export type ChatStatus = 'submitted' | 'streaming' | 'ready' | 'error';
 type ActiveResponse<UI_MESSAGE extends UIMessage> = {
   state: StreamingUIMessageState<UI_MESSAGE>;
   abortController: AbortController;
+  completionPromise: Promise<void>;
 };
 
 type ActiveResumeRequest = {
   abortController: AbortController;
+  completionPromise: Promise<void>;
 };
 
+type MakeRequestOptions = {
+  trigger: 'submit-message' | 'resume-stream' | 'regenerate-message';
+  messageId?: string;
+} & ChatRequestOptions;
+
 export interface ChatState<UI_MESSAGE extends UIMessage> {
   status: ChatStatus;
 
@@ -267,6 +275,8 @@ export abstract class AbstractChat<UI_MESSAGE extends UIMessage> {
     | StreamingUIMessageState<UI_MESSAGE>
     | undefined = undefined;
   private jobExecutor = new SerialJobExecutor();
+  private activeStopCount = 0;
+  private stopGeneration = 0;
 
   constructor({
     generateId = generateIdFunc,
@@ -591,6 +601,7 @@ export abstract class AbstractChat<UI_MESSAGE extends UIMessage> {
 
       // automatically send the message if the sendAutomaticallyWhen function returns true
       if (
+        messageIndex !== -1 &&
         this.status !== 'streaming' &&
         this.status !== 'submitted' &&
         this.sendAutomaticallyWhen
@@ -599,16 +610,18 @@ export abstract class AbstractChat<UI_MESSAGE extends UIMessage> {
 
         if (shouldSend) {
           // no await to avoid deadlocking
-          const messageId =
-            messageIndex === -1
-              ? this.lastMessage?.id
-              : messages[messageIndex].id;
-
-          this.makeRequestForToolApproval({
-            messageId,
-            messageIndex,
-            ...options,
-          });
+          void this.runAutomaticRequest(() => {
+            const messageId =
+              messageIndex === -1
+                ? this.lastMessage?.id
+                : messages[messageIndex].id;
+
+            return this.makeRequestForToolApproval({
+              messageId,
+              messageIndex,
+              ...options,
+            });
+          }, shouldSend);
         }
       }
     });
@@ -684,11 +697,15 @@ export abstract class AbstractChat<UI_MESSAGE extends UIMessage> {
 
         if (shouldSend) {
           // no await to avoid deadlocking
-          this.makeRequest({
-            trigger: 'submit-message',
-            messageId: this.lastMessage?.id,
-            ...options,
-          });
+          void this.runAutomaticRequest(
+            () =>
+              this.makeRequest({
+                trigger: 'submit-message',
+                messageId: this.lastMessage?.id,
+                ...options,
+              }),
+            shouldSend,
+          );
         }
       }
     });
@@ -697,14 +714,42 @@ export abstract class AbstractChat<UI_MESSAGE extends UIMessage> {
   addToolResult = this.addToolOutput;
 
   /**
-   * Abort the current request immediately, keep the generated tokens if any.
+   * Abort the current request, keep the generated tokens if any, and wait for
+   * the request pipeline to finish.
    */
   stop = async () => {
-    for (const controller of this.pendingMessagePreparations) {
-      controller.abort();
+    this.activeStopCount++;
+    this.stopGeneration++;
+
+    try {
+      const activeResumeRequest = this.activeResumeRequest;
+      const activeResponse = this.activeResponse;
+
+      for (const controller of this.pendingMessagePreparations) {
+        controller.abort();
+      }
+      activeResumeRequest?.abortController.abort();
+      activeResponse?.abortController.abort();
+
+      await Promise.all([
+        activeResumeRequest?.completionPromise,
+        activeResponse?.completionPromise,
+      ]);
+
+      // Stream cancellation can complete while a processing job is still
+      // blocked in onToolCall. Drain that job and any message update it queued.
+      await this.jobExecutor.waitForIdle();
+    } finally {
+      this.activeStopCount--;
     }
-    this.activeResumeRequest?.abortController.abort();
-    this.activeResponse?.abortController.abort();
+  };
+
+  /**
+   * Stops active work and releases resources owned by the chat transport.
+   */
+  dispose = async (): Promise<void> => {
+    await this.stop();
+    this.transport.close?.();
   };
 
   private async shouldSendAutomatically(): Promise<boolean> {
@@ -722,6 +767,30 @@ export abstract class AbstractChat<UI_MESSAGE extends UIMessage> {
     return result as boolean;
   }
 
+  private async runAutomaticRequest(
+    request: () => Promise<void>,
+    shouldSend?: boolean,
+  ): Promise<void> {
+    const stopGeneration = this.stopGeneration;
+    const startedWhileStopping = this.activeStopCount > 0;
+
+    if (!(shouldSend ?? (await this.shouldSendAutomatically()))) {
+      return;
+    }
+
+    // A stop invalidates automatic requests that were pending when it started,
+    // including requests queued by callback-driven tool updates.
+    if (
+      startedWhileStopping ||
+      stopGeneration !== this.stopGeneration ||
+      this.activeStopCount > 0
+    ) {
+      return;
+    }
+
+    await request();
+  }
+
   private async makeRequestForToolApproval({
     messageId,
     messageIndex,
@@ -752,23 +821,38 @@ export abstract class AbstractChat<UI_MESSAGE extends UIMessage> {
     }
   }
 
-  private async makeRequest({
+  private async makeRequest(options: MakeRequestOptions) {
+    const completion = createResolvablePromise<void>();
+
+    try {
+      await this.makeRequestImpl({
+        ...options,
+        completionPromise: completion.promise,
+      });
+    } finally {
+      completion.resolve();
+    }
+  }
+
+  private async makeRequestImpl({
     trigger,
     metadata,
     headers,
     body,
     messageId,
-  }: {
-    trigger: 'submit-message' | 'resume-stream' | 'regenerate-message';
-    messageId?: string;
-  } & ChatRequestOptions) {
+    completionPromise,
+  }: MakeRequestOptions & {
+    completionPromise: Promise<void>;
+  }) {
     if (trigger !== 'resume-stream') {
       this.resumableStreamState = undefined;
     }
 
     const abortController = new AbortController();
     const activeResumeRequest =
-      trigger === 'resume-stream' ? { abortController } : undefined;
+      trigger === 'resume-stream'
+        ? { abortController, completionPromise }
+        : undefined;
 
     if (activeResumeRequest) {
       this.activeResumeRequest?.abortController.abort();
@@ -893,6 +977,7 @@ export abstract class AbstractChat<UI_MESSAGE extends UIMessage> {
                 messageId: this.generateId(),
               }),
         abortController,
+        completionPromise,
       } as ActiveResponse<UI_MESSAGE>;
 
       activeResponse = response;
@@ -964,6 +1049,8 @@ export abstract class AbstractChat<UI_MESSAGE extends UIMessage> {
         stream: processUIMessageStream({
           stream,
           resetStateOnMessageIdChange: trigger === 'resume-stream',
+          resetStateOnFirstMessageStart:
+            trigger === 'resume-stream' && this.transport.resumeStreamIsReplay,
           onToolCall: this.onToolCall,
           onData: this.onData,
           messageMetadataSchema: this.messageMetadataSchema,
@@ -1059,14 +1146,16 @@ export abstract class AbstractChat<UI_MESSAGE extends UIMessage> {
     }
 
     // automatically send the message if the sendAutomaticallyWhen function returns true
-    if (!isError && (await this.shouldSendAutomatically())) {
-      await this.makeRequest({
-        trigger: 'submit-message',
-        messageId: this.lastMessage?.id,
-        metadata,
-        headers,
-        body,
-      });
+    if (!isAbort && !isError) {
+      await this.runAutomaticRequest(() =>
+        this.makeRequest({
+          trigger: 'submit-message',
+          messageId: this.lastMessage?.id,
+          metadata,
+          headers,
+          body,
+        }),
+      );
     }
   }
 }

package/src/ui/index.ts

@@ -32,6 +32,19 @@ export {
 export { lastAssistantMessageIsCompleteWithApprovalResponses } from './last-assistant-message-is-complete-with-approval-responses';
 export { lastAssistantMessageIsCompleteWithToolCalls } from './last-assistant-message-is-complete-with-tool-calls';
 export { TextStreamChatTransport } from './text-stream-chat-transport';
+export {
+  safeValidateWebSocketChatTransportRequest,
+  WebSocketChatTransport,
+  type PrepareWebSocketChatTransportReconnectToStreamRequest,
+  type PrepareWebSocketChatTransportSendMessagesRequest,
+  type SafeValidateWebSocketChatTransportRequestResult,
+  type WebSocketChatTransportAbortRequest,
+  type WebSocketChatTransportInitOptions,
+  type WebSocketChatTransportRequest,
+  type WebSocketChatTransportResponse,
+  type WebSocketChatTransportResumeRequest,
+  type WebSocketChatTransportSendRequest,
+} from './websocket-chat-transport';
 export {
   getStaticToolName,
   getToolName,

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

@@ -121,6 +121,7 @@ export function processUIMessageStream<UI_MESSAGE extends UIMessage>({
   onToolCall,
   onData,
   resetStateOnMessageIdChange = false,
+  resetStateOnFirstMessageStart = false,
 }: {
   // input stream is not fully typed yet:
   stream: ReadableStream<UIMessageChunk>;
@@ -139,6 +140,8 @@ export function processUIMessageStream<UI_MESSAGE extends UIMessage>({
   onError: ErrorHandler;
   // During resume, a different message ID identifies a separate response.
   resetStateOnMessageIdChange?: boolean;
+  // A transport can replay a response from the beginning using the same ID.
+  resetStateOnFirstMessageStart?: boolean;
 }): ReadableStream<InferUIMessageChunk<UI_MESSAGE>> {
   return stream.pipeThrough(
     new TransformStream<UIMessageChunk, InferUIMessageChunk<UI_MESSAGE>>({
@@ -957,9 +960,10 @@ export function processUIMessageStream<UI_MESSAGE extends UIMessage>({
 
             case 'start': {
               if (
-                resetStateOnMessageIdChange &&
                 chunk.messageId != null &&
-                chunk.messageId !== state.message.id
+                (resetStateOnFirstMessageStart ||
+                  (resetStateOnMessageIdChange &&
+                    chunk.messageId !== state.message.id))
               ) {
                 // Start the separate response with empty parts and metadata.
                 // Reset the active-part maps too, so earlier chunks stay with