ai 6.0.303 → 6.0.304

package/src/ui/chat-transport.ts

@@ -13,6 +13,13 @@ 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;
+
   /**
    * Sends messages to the chat API endpoint and returns a streaming response.
    *

package/src/ui/chat.ts

@@ -518,6 +518,33 @@ export abstract class AbstractChat<UI_MESSAGE extends UIMessage> {
     await this.makeRequest({ trigger: 'resume-stream', ...options });
   };
 
+  /**
+   * Resume the stream if the latest response ended with a resumable network
+   * disconnect. If a response is still active, wait for it to settle first so
+   * disconnects reported shortly after the caller's trigger are handled.
+   *
+   * @internal
+   */
+  '~resumeStreamIfDisconnected' = async ({
+    shouldResume = () => true,
+  }: {
+    shouldResume?: () => boolean;
+  } = {}): Promise<void> => {
+    const activeResponse = this.activeResponse;
+
+    if (this.resumableStreamState == null && activeResponse != null) {
+      await activeResponse.completionPromise;
+
+      if (this.resumableStreamState !== activeResponse.state) {
+        return;
+      }
+    }
+
+    if (this.resumableStreamState != null && shouldResume()) {
+      await this.resumeStream();
+    }
+  };
+
   /**
    * Clear the error state and set the status to ready if the chat is in an error state.
    */
@@ -611,7 +638,6 @@ export abstract class AbstractChat<UI_MESSAGE extends UIMessage> {
   }) =>
     this.jobExecutor.run(async () => {
       const messages = this.state.messages;
-      const lastMessage = messages[messages.length - 1];
 
       const updatePart = (
         part: UIMessagePart<UIDataTypes, UITools>,
@@ -620,11 +646,21 @@ export abstract class AbstractChat<UI_MESSAGE extends UIMessage> {
           ? ({ ...part, state, output, errorText } as typeof part)
           : part;
 
-      // update the message to trigger an immediate UI update
-      this.state.replaceMessage(messages.length - 1, {
-        ...lastMessage,
-        parts: lastMessage.parts.map(updatePart),
-      });
+      const messageIndex = messages.findIndex(message =>
+        message.parts.some(
+          part => isToolUIPart(part) && part.toolCallId === toolCallId,
+        ),
+      );
+
+      if (messageIndex !== -1) {
+        const message = messages[messageIndex];
+
+        // update the message to trigger an immediate UI update
+        this.state.replaceMessage(messageIndex, {
+          ...message,
+          parts: message.parts.map(updatePart),
+        });
+      }
 
       // update the active response if it exists
       if (this.activeResponse) {
@@ -634,6 +670,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
@@ -646,7 +683,7 @@ export abstract class AbstractChat<UI_MESSAGE extends UIMessage> {
             () =>
               this.makeRequest({
                 trigger: 'submit-message',
-                messageId: this.lastMessage?.id,
+                messageId: messages[messageIndex].id,
                 ...options,
               }),
             shouldSend,
@@ -872,6 +909,7 @@ export abstract class AbstractChat<UI_MESSAGE extends UIMessage> {
       responseMessageIndex === -1
         ? lastMessage
         : this.state.messages[responseMessageIndex];
+    const originalResponseMessageId = responseMessage?.id;
     const resumableResponseMessage =
       trigger === 'resume-stream' &&
       responseMessage?.role === 'assistant' &&
@@ -956,16 +994,28 @@ export abstract class AbstractChat<UI_MESSAGE extends UIMessage> {
                 this.setStatus({ status: 'streaming' });
               }
 
-              if (usesEarlierAssistantMessage) {
+              const existingMessageIndex = this.state.messages.findLastIndex(
+                message => message.id === response.state.message.id,
+              );
+
+              if (existingMessageIndex !== -1) {
                 this.state.replaceMessage(
-                  responseMessageIndex,
+                  existingMessageIndex,
                   response.state.message,
                 );
-              } else if (response.state.message.id === this.lastMessage?.id) {
-                this.state.replaceMessage(
-                  this.state.messages.length - 1,
-                  response.state.message,
+              } else if (usesEarlierAssistantMessage) {
+                const originalMessageIndex = this.state.messages.findLastIndex(
+                  message => message.id === originalResponseMessageId,
                 );
+
+                if (originalMessageIndex !== -1) {
+                  this.state.replaceMessage(
+                    originalMessageIndex,
+                    response.state.message,
+                  );
+                } else {
+                  this.state.pushMessage(response.state.message);
+                }
               } else {
                 this.state.pushMessage(response.state.message);
               }
@@ -976,6 +1026,9 @@ export abstract class AbstractChat<UI_MESSAGE extends UIMessage> {
       await consumeStream({
         stream: processUIMessageStream({
           stream,
+          resetStateOnMessageIdChange: trigger === 'resume-stream',
+          resetStateOnFirstMessageStart:
+            trigger === 'resume-stream' && this.transport.resumeStreamIsReplay,
           onToolCall: this.onToolCall,
           onData: this.onData,
           messageMetadataSchema: this.messageMetadataSchema,

package/src/ui/default-chat-transport.ts

@@ -12,8 +12,21 @@ import type { UIMessage } from './ui-messages';
 export class DefaultChatTransport<
   UI_MESSAGE extends UIMessage,
 > extends HttpChatTransport<UI_MESSAGE> {
-  constructor(options: HttpChatTransportInitOptions<UI_MESSAGE> = {}) {
+  readonly resumeStreamIsReplay: boolean;
+
+  constructor({
+    resumeStreamIsReplay = true,
+    ...options
+  }: HttpChatTransportInitOptions<UI_MESSAGE> & {
+    /**
+     * Whether resumed streams replay the response from the beginning.
+     * Defaults to true. Set to false for endpoints that only continue the
+     * interrupted response, even when they emit a start chunk.
+     */
+    resumeStreamIsReplay?: boolean;
+  } = {}) {
     super(options);
+    this.resumeStreamIsReplay = resumeStreamIsReplay;
   }
 
   protected processResponseStream(

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,7 +140,11 @@ 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>> {
+  let isFirstMessageStart = true;
+
   return stream.pipeThrough(
     new TransformStream<UIMessageChunk, InferUIMessageChunk<UI_MESSAGE>>({
       async transform(chunk, controller) {
@@ -881,19 +886,23 @@ export function processUIMessageStream<UI_MESSAGE extends UIMessage>({
             }
 
             case 'start': {
-              if (
-                resetStateOnMessageIdChange &&
-                chunk.messageId != null &&
-                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
-                // the previous response. Keep the state object for its callers.
+              const shouldReset =
+                (resetStateOnFirstMessageStart && isFirstMessageStart) ||
+                (resetStateOnMessageIdChange &&
+                  chunk.messageId != null &&
+                  chunk.messageId !== state.message.id);
+              isFirstMessageStart = false;
+
+              if (shouldReset) {
+                // Rebuild a replay or separate response from empty state. Only
+                // reset a replay once: merged streams can emit more starts.
+                // Keep the state object for its callers and preserve the ID
+                // when the replay does not include one.
                 Object.assign(
                   state,
                   createStreamingUIMessageState({
                     lastMessage: undefined,
-                    messageId: chunk.messageId,
+                    messageId: chunk.messageId ?? state.message.id,
                   }),
                   { finishReason: undefined },
                 );
@@ -905,7 +914,11 @@ export function processUIMessageStream<UI_MESSAGE extends UIMessage>({
 
               await updateMessageMetadata(chunk.messageMetadata);
 
-              if (chunk.messageId != null || chunk.messageMetadata != null) {
+              if (
+                shouldReset ||
+                chunk.messageId != null ||
+                chunk.messageMetadata != null
+              ) {
                 write({ updateStatus: false });
               }
               break;