ai 7.0.133 → 7.0.135

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/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,25 +601,28 @@ 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
       ) {
-        this.shouldSendAutomatically().then(shouldSend => {
-          if (shouldSend) {
-            // no await to avoid deadlocking
+        const shouldSend = await this.shouldSendAutomatically();
+
+        if (shouldSend) {
+          // no await to avoid deadlocking
+          void this.runAutomaticRequest(() => {
             const messageId =
               messageIndex === -1
                 ? this.lastMessage?.id
                 : messages[messageIndex].id;
 
-            this.makeRequestForToolApproval({
+            return this.makeRequestForToolApproval({
               messageId,
               messageIndex,
               ...options,
             });
-          }
-        });
+          }, shouldSend);
+        }
       }
     });
 
@@ -680,16 +693,20 @@ export abstract class AbstractChat<UI_MESSAGE extends UIMessage> {
         this.status !== 'submitted' &&
         this.sendAutomaticallyWhen
       ) {
-        this.shouldSendAutomatically().then(shouldSend => {
-          if (shouldSend) {
-            // no await to avoid deadlocking
-            this.makeRequest({
-              trigger: 'submit-message',
-              messageId: this.lastMessage?.id,
-              ...options,
-            });
-          }
-        });
+        const shouldSend = await this.shouldSendAutomatically();
+
+        if (shouldSend) {
+          // no await to avoid deadlocking
+          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();
@@ -854,10 +938,8 @@ export abstract class AbstractChat<UI_MESSAGE extends UIMessage> {
       responseMessageIndex === -1
         ? lastMessage
         : this.state.messages[responseMessageIndex];
-    // the continued stream can start with either
-    // 1) input deltas
-    // 2) or the result of an answered tool approval request
-    // Keep the tool part so in case 2, those chunks can find their tool call
+    // The continued stream can start with input deltas or a tool result.
+    // Keep unfinished tool parts so result chunks can find their tool call.
     const resumableResponseMessage =
       trigger === 'resume-stream' &&
       responseMessage?.role === 'assistant' &&
@@ -865,6 +947,7 @@ export abstract class AbstractChat<UI_MESSAGE extends UIMessage> {
         part =>
           isToolUIPart(part) &&
           (part.state === 'input-streaming' ||
+            (part.state === 'input-available' && part.providerExecuted) ||
             part.state === 'approval-responded'),
       )
         ? this.state.snapshot(responseMessage)
@@ -894,6 +977,7 @@ export abstract class AbstractChat<UI_MESSAGE extends UIMessage> {
                 messageId: this.generateId(),
               }),
         abortController,
+        completionPromise,
       } as ActiveResponse<UI_MESSAGE>;
 
       activeResponse = response;
@@ -965,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,
@@ -1060,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