ai 6.0.296 → 6.0.297

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,
+  });
+}