ai 6.0.295 → 6.0.297

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

@@ -60,22 +60,55 @@ export function createStreamingUIMessageState<UI_MESSAGE extends UIMessage>({
   lastMessage: UI_MESSAGE | undefined;
   messageId: string;
 }): StreamingUIMessageState<UI_MESSAGE> {
+  const message =
+    lastMessage?.role === 'assistant'
+      ? lastMessage
+      : ({
+          id: messageId,
+          metadata: undefined,
+          role: 'assistant',
+          parts: [] as UIMessagePart<
+            InferUIMessageData<UI_MESSAGE>,
+            InferUIMessageTools<UI_MESSAGE>
+          >[],
+        } as UI_MESSAGE);
+  const partialToolCalls: StreamingUIMessageState<UI_MESSAGE>['partialToolCalls'] =
+    createIdMap();
+  const lastStepStartIndex = message.parts.findLastIndex(
+    part => part.type === 'step-start',
+  );
+  let staticToolIndex = 0;
+
+  for (const part of message.parts.slice(lastStepStartIndex + 1)) {
+    if (!isToolUIPart(part)) {
+      continue;
+    }
+
+    const index = staticToolIndex;
+    if (isStaticToolUIPart(part)) {
+      staticToolIndex++;
+    }
+
+    if (part.state !== 'input-streaming') {
+      continue;
+    }
+
+    partialToolCalls[part.toolCallId] = {
+      text: part.rawInput ?? '',
+      index,
+      toolName:
+        part.type === 'dynamic-tool' ? part.toolName : getStaticToolName(part),
+      dynamic: part.type === 'dynamic-tool',
+      title: part.title,
+      toolMetadata: part.toolMetadata,
+    };
+  }
+
   return {
-    message:
-      lastMessage?.role === 'assistant'
-        ? lastMessage
-        : ({
-            id: messageId,
-            metadata: undefined,
-            role: 'assistant',
-            parts: [] as UIMessagePart<
-              InferUIMessageData<UI_MESSAGE>,
-              InferUIMessageTools<UI_MESSAGE>
-            >[],
-          } as UI_MESSAGE),
+    message,
     activeTextParts: createIdMap(),
     activeReasoningParts: createIdMap(),
-    partialToolCalls: createIdMap(),
+    partialToolCalls,
   };
 }
 
@@ -167,6 +200,7 @@ export function processUIMessageStream<UI_MESSAGE extends UIMessage>({
               | {
                   state: 'input-streaming';
                   input: unknown;
+                  rawInput?: string;
                   providerExecuted?: boolean;
                   providerMetadata?: ProviderMetadata;
                 }
@@ -282,6 +316,7 @@ export function processUIMessageStream<UI_MESSAGE extends UIMessage>({
               | {
                   state: 'input-streaming';
                   input: unknown;
+                  rawInput?: string;
                   providerMetadata?: ProviderMetadata;
                 }
               | {
@@ -322,7 +357,7 @@ export function processUIMessageStream<UI_MESSAGE extends UIMessage>({
               anyPart.input = anyOptions.input;
               anyPart.output = anyOptions.output;
               anyPart.errorText = anyOptions.errorText;
-              anyPart.rawInput = anyOptions.rawInput ?? anyPart.rawInput;
+              anyPart.rawInput = anyOptions.rawInput;
               anyPart.preliminary = anyOptions.preliminary;
               if (options.title !== undefined) {
                 anyPart.title = options.title;
@@ -615,6 +650,7 @@ export function processUIMessageStream<UI_MESSAGE extends UIMessage>({
                   toolName: partialToolCall.toolName,
                   state: 'input-streaming',
                   input: partialArgs,
+                  rawInput: partialToolCall.text,
                   title: partialToolCall.title,
                   toolMetadata: partialToolCall.toolMetadata,
                 });
@@ -624,6 +660,7 @@ export function processUIMessageStream<UI_MESSAGE extends UIMessage>({
                   toolName: partialToolCall.toolName,
                   state: 'input-streaming',
                   input: partialArgs,
+                  rawInput: partialToolCall.text,
                   title: partialToolCall.title,
                   toolMetadata: partialToolCall.toolMetadata,
                 });
@@ -757,7 +794,8 @@ export function processUIMessageStream<UI_MESSAGE extends UIMessage>({
                     providerExecuted: chunk.providerExecuted,
                     providerMetadata: chunk.providerMetadata,
                     title: toolInvocation.title,
-                    toolMetadata: toolInvocation.toolMetadata,
+                    toolMetadata:
+                      chunk.toolMetadata ?? toolInvocation.toolMetadata,
                   },
                   toolInvocation,
                 );
@@ -773,7 +811,8 @@ export function processUIMessageStream<UI_MESSAGE extends UIMessage>({
                     preliminary: chunk.preliminary,
                     providerMetadata: chunk.providerMetadata,
                     title: toolInvocation.title,
-                    toolMetadata: toolInvocation.toolMetadata,
+                    toolMetadata:
+                      chunk.toolMetadata ?? toolInvocation.toolMetadata,
                   },
                   toolInvocation as ToolUIPart<InferUIMessageTools<UI_MESSAGE>>,
                 );
@@ -797,7 +836,8 @@ export function processUIMessageStream<UI_MESSAGE extends UIMessage>({
                     providerExecuted: chunk.providerExecuted,
                     providerMetadata: chunk.providerMetadata,
                     title: toolInvocation.title,
-                    toolMetadata: toolInvocation.toolMetadata,
+                    toolMetadata:
+                      chunk.toolMetadata ?? toolInvocation.toolMetadata,
                   },
                   toolInvocation,
                 );
@@ -813,7 +853,8 @@ export function processUIMessageStream<UI_MESSAGE extends UIMessage>({
                     providerExecuted: chunk.providerExecuted,
                     providerMetadata: chunk.providerMetadata,
                     title: toolInvocation.title,
-                    toolMetadata: toolInvocation.toolMetadata,
+                    toolMetadata:
+                      chunk.toolMetadata ?? toolInvocation.toolMetadata,
                   },
                   toolInvocation as ToolUIPart<InferUIMessageTools<UI_MESSAGE>>,
                 );

package/src/ui/ui-messages.ts

@@ -238,6 +238,13 @@ export type UIToolInvocation<TOOL extends UITool | Tool> = {
   | {
       state: 'input-streaming';
       input?: DeepPartial<asUITool<TOOL>['input']> | undefined;
+      /**
+       * The raw tool input text received so far.
+       *
+       * This is used to continue input streaming when a message is persisted
+       * and later resumed.
+       */
+      rawInput?: string;
       output?: never;
       errorText?: never;
       callProviderMetadata?: ProviderMetadata;
@@ -361,6 +368,13 @@ export type DynamicToolUIPart = {
   | {
       state: 'input-streaming';
       input?: unknown;
+      /**
+       * The raw tool input text received so far.
+       *
+       * This is used to continue input streaming when a message is persisted
+       * and later resumed.
+       */
+      rawInput?: string;
       output?: never;
       errorText?: never;
       callProviderMetadata?: ProviderMetadata;

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

@@ -139,6 +139,7 @@ const uiMessagesSchema = lazySchema(() => {
                   toolMetadata: toolMetadataSchema.optional(),
                   state: z.literal('input-streaming'),
                   input: z.unknown().optional(),
+                  rawInput: z.string().optional(),
                   providerExecuted: z.boolean().optional(),
                   callProviderMetadata: providerMetadataSchema.optional(),
                   output: z.never().optional(),
@@ -235,6 +236,7 @@ const uiMessagesSchema = lazySchema(() => {
                   providerExecuted: z.boolean().optional(),
                   callProviderMetadata: providerMetadataSchema.optional(),
                   input: z.unknown().optional(),
+                  rawInput: z.string().optional(),
                   output: z.never().optional(),
                   errorText: z.never().optional(),
                   approval: z.never().optional(),

package/src/ui-message-stream/create-sse-stream-with-keep-alive.ts

@@ -0,0 +1,89 @@
+const STREAM_OPEN_COMMENT = ': stream-open\n\n';
+const KEEP_ALIVE_COMMENT = ': keep-alive\n\n';
+
+export function createSseStreamWithKeepAlive({
+  stream,
+  keepAliveMs,
+}: {
+  stream: ReadableStream<string>;
+  keepAliveMs: number | undefined;
+}): ReadableStream<string> {
+  if (keepAliveMs == null) {
+    return stream;
+  }
+
+  if (
+    !Number.isFinite(keepAliveMs) ||
+    keepAliveMs <= 0 ||
+    keepAliveMs > 2_147_483_647
+  ) {
+    throw new Error(
+      'keepAliveMs must be a positive finite timer duration no greater than 2147483647',
+    );
+  }
+
+  const reader = stream.getReader();
+  let keepAliveTimeout: ReturnType<typeof setTimeout> | undefined;
+  let isCancelled = false;
+
+  const clearKeepAliveTimeout = () => {
+    clearTimeout(keepAliveTimeout);
+    keepAliveTimeout = undefined;
+  };
+
+  const scheduleKeepAlive = (
+    controller: ReadableStreamDefaultController<string>,
+  ) => {
+    clearKeepAliveTimeout();
+    keepAliveTimeout = setTimeout(() => {
+      if (isCancelled) {
+        return;
+      }
+
+      if (controller.desiredSize != null && controller.desiredSize > 0) {
+        controller.enqueue(KEEP_ALIVE_COMMENT);
+      }
+
+      scheduleKeepAlive(controller);
+    }, keepAliveMs);
+  };
+
+  return new ReadableStream<string>({
+    start(controller) {
+      controller.enqueue(STREAM_OPEN_COMMENT);
+      scheduleKeepAlive(controller);
+    },
+
+    pull(controller) {
+      return reader.read().then(
+        result => {
+          clearKeepAliveTimeout();
+
+          if (isCancelled) {
+            return;
+          }
+
+          if (result.done) {
+            controller.close();
+          } else {
+            controller.enqueue(result.value);
+            scheduleKeepAlive(controller);
+          }
+        },
+        error => {
+          clearKeepAliveTimeout();
+
+          if (!isCancelled) {
+            throw error;
+          }
+        },
+      );
+    },
+
+    async cancel(reason) {
+      isCancelled = true;
+      clearKeepAliveTimeout();
+      await reader.cancel(reason);
+    },
+  });
+}

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

@@ -1,4 +1,5 @@
 import { prepareHeaders } from '../util/prepare-headers';
+import { createSseStreamWithKeepAlive } from './create-sse-stream-with-keep-alive';
 import { JsonToSseTransformStream } from './json-to-sse-transform-stream';
 import { UI_MESSAGE_STREAM_HEADERS } from './ui-message-stream-headers';
 import type { UIMessageChunk } from './ui-message-chunks';
@@ -12,6 +13,7 @@ import type { UIMessageStreamResponseInit } from './ui-message-stream-response-i
  * @param options.statusText - The HTTP status text for the response.
  * @param options.headers - Additional HTTP headers to include in the response.
  * @param options.stream - The UI message chunk stream to send.
+ * @param options.keepAliveMs - Optional interval for sending SSE keep-alive comments.
  * @param options.consumeSseStream - Optional callback to consume a copy of the SSE stream independently.
  *
  * @returns A `Response` object with the UI message stream as the body.
@@ -21,11 +23,15 @@ export function createUIMessageStreamResponse({
   statusText,
   headers,
   stream,
+  keepAliveMs,
   consumeSseStream,
 }: UIMessageStreamResponseInit & {
   stream: ReadableStream<UIMessageChunk>;
 }): Response {
-  let sseStream = stream.pipeThrough(new JsonToSseTransformStream());
+  let sseStream = createSseStreamWithKeepAlive({
+    stream: stream.pipeThrough(new JsonToSseTransformStream()),
+    keepAliveMs,
+  });
 
   // when the consumeSseStream is provided, we need to tee the stream
   // and send the second part to the consumeSseStream function

package/src/ui-message-stream/pipe-ui-message-stream-to-response.ts

@@ -1,6 +1,7 @@
 import type { ServerResponse } from 'node:http';
 import { prepareHeaders } from '../util/prepare-headers';
 import { writeToServerResponse } from '../util/write-to-server-response';
+import { createSseStreamWithKeepAlive } from './create-sse-stream-with-keep-alive';
 import { JsonToSseTransformStream } from './json-to-sse-transform-stream';
 import { UI_MESSAGE_STREAM_HEADERS } from './ui-message-stream-headers';
 import type { UIMessageChunk } from './ui-message-chunks';
@@ -15,6 +16,7 @@ import type { UIMessageStreamResponseInit } from './ui-message-stream-response-i
  * @param options.statusText - The HTTP status text for the response.
  * @param options.headers - Additional HTTP headers to include in the response.
  * @param options.stream - The UI message chunk stream to send.
+ * @param options.keepAliveMs - Optional interval for sending SSE keep-alive comments.
  * @param options.consumeSseStream - Optional callback to consume a copy of the SSE stream independently.
  * @returns A promise that resolves when the stream has been written.
  */
@@ -24,12 +26,16 @@ export function pipeUIMessageStreamToResponse({
   statusText,
   headers,
   stream,
+  keepAliveMs,
   consumeSseStream,
 }: {
   response: ServerResponse;
   stream: ReadableStream<UIMessageChunk>;
 } & UIMessageStreamResponseInit): Promise<void> {
-  let sseStream = stream.pipeThrough(new JsonToSseTransformStream());
+  let sseStream = createSseStreamWithKeepAlive({
+    stream: stream.pipeThrough(new JsonToSseTransformStream()),
+    keepAliveMs,
+  });
 
   // when the consumeSseStream is provided, we need to tee the stream
   // and send the second part to the consumeSseStream function

package/src/ui-message-stream/ui-message-stream-response-init.ts

@@ -3,6 +3,13 @@
  * Extends the standard `ResponseInit` with additional streaming options.
  */
 export type UIMessageStreamResponseInit = ResponseInit & {
+  /**
+   * Optional interval in milliseconds for sending SSE keep-alive comments.
+   * When set, an opening comment is sent immediately and additional comments
+   * are sent after the stream has been idle for the configured interval.
+   */
+  keepAliveMs?: number;
+
   /**
    * Optional callback to consume a copy of the SSE stream independently.
    * This is useful for logging, debugging, or processing the stream in parallel.

package/src/util/get-own.ts

@@ -0,0 +1,18 @@
+/**
+ * Reads a property by an untrusted key, ignoring inherited prototype members.
+ *
+ * Tool sets, tool contexts, and similar lookup objects are indexed by names
+ * that can come from model output or client-supplied message history. Plain
+ * bracket access (`obj[name]`) resolves names such as `constructor`,
+ * `toString`, or `__proto__` to values on `Object.prototype`, which would slip
+ * past the `== null` / `!value` guards that treat an unknown name as "not
+ * present". This helper returns `undefined` unless `key` is an own property.
+ */
+export function getOwn<T extends object>(
+  obj: T | undefined | null,
+  key: string,
+): T[keyof T] | undefined {
+  return obj != null && Object.hasOwn(obj, key)
+    ? obj[key as keyof T]
+    : undefined;
+}

package/src/util/set-own.ts

@@ -0,0 +1,16 @@
+/**
+ * Writes an enumerable own property without invoking inherited setters such as
+ * `Object.prototype.__proto__`. Keeps the target object's prototype unchanged.
+ */
+export function setOwn<T>(
+  object: Record<string, T>,
+  key: string,
+  value: T,
+): void {
+  Object.defineProperty(object, key, {
+    value,
+    enumerable: true,
+    configurable: true,
+    writable: true,
+  });
+}