ai 6.0.32 → 6.0.34

package/src/ui-message-stream/handle-ui-message-stream-finish.ts

@@ -0,0 +1,135 @@
+import {
+  createStreamingUIMessageState,
+  processUIMessageStream,
+  StreamingUIMessageState,
+} from '../ui/process-ui-message-stream';
+import { UIMessage } from '../ui/ui-messages';
+import { ErrorHandler } from '../util/error-handler';
+import { InferUIMessageChunk, UIMessageChunk } from './ui-message-chunks';
+import { UIMessageStreamOnFinishCallback } from './ui-message-stream-on-finish-callback';
+
+export function handleUIMessageStreamFinish<UI_MESSAGE extends UIMessage>({
+  messageId,
+  originalMessages = [],
+  onFinish,
+  onError,
+  stream,
+}: {
+  stream: ReadableStream<InferUIMessageChunk<UI_MESSAGE>>;
+
+  /**
+   * The message ID to use for the response message.
+   * If not provided, no id will be set for the response message.
+   */
+  messageId?: string;
+
+  /**
+   * The original messages.
+   */
+  originalMessages?: UI_MESSAGE[];
+
+  onError: ErrorHandler;
+
+  onFinish?: UIMessageStreamOnFinishCallback<UI_MESSAGE>;
+}): ReadableStream<InferUIMessageChunk<UI_MESSAGE>> {
+  // last message is only relevant for assistant messages
+  let lastMessage: UI_MESSAGE | undefined =
+    originalMessages?.[originalMessages.length - 1];
+  if (lastMessage?.role !== 'assistant') {
+    lastMessage = undefined;
+  } else {
+    // appending to the last message, so we need to use the same id
+    messageId = lastMessage.id;
+  }
+
+  let isAborted = false;
+
+  const idInjectedStream = stream.pipeThrough(
+    new TransformStream<
+      InferUIMessageChunk<UI_MESSAGE>,
+      InferUIMessageChunk<UI_MESSAGE>
+    >({
+      transform(chunk, controller) {
+        // when there is no messageId in the start chunk,
+        // but the user checked for persistence,
+        // inject the messageId into the chunk
+        if (chunk.type === 'start') {
+          const startChunk = chunk as UIMessageChunk & { type: 'start' };
+          if (startChunk.messageId == null && messageId != null) {
+            startChunk.messageId = messageId;
+          }
+        }
+
+        if (chunk.type === 'abort') {
+          isAborted = true;
+        }
+
+        controller.enqueue(chunk);
+      },
+    }),
+  );
+
+  if (onFinish == null) {
+    return idInjectedStream;
+  }
+
+  const state = createStreamingUIMessageState<UI_MESSAGE>({
+    lastMessage: lastMessage
+      ? (structuredClone(lastMessage) as UI_MESSAGE)
+      : undefined,
+    messageId: messageId ?? '', // will be overridden by the stream
+  });
+
+  const runUpdateMessageJob = async (
+    job: (options: {
+      state: StreamingUIMessageState<UI_MESSAGE>;
+      write: () => void;
+    }) => Promise<void>,
+  ) => {
+    await job({ state, write: () => {} });
+  };
+
+  let finishCalled = false;
+
+  const callOnFinish = async () => {
+    if (finishCalled || !onFinish) {
+      return;
+    }
+    finishCalled = true;
+
+    const isContinuation = state.message.id === lastMessage?.id;
+    await onFinish({
+      isAborted,
+      isContinuation,
+      responseMessage: state.message as UI_MESSAGE,
+      messages: [
+        ...(isContinuation ? originalMessages.slice(0, -1) : originalMessages),
+        state.message,
+      ] as UI_MESSAGE[],
+      finishReason: state.finishReason,
+    });
+  };
+
+  return processUIMessageStream<UI_MESSAGE>({
+    stream: idInjectedStream,
+    runUpdateMessageJob,
+    onError,
+  }).pipeThrough(
+    new TransformStream<
+      InferUIMessageChunk<UI_MESSAGE>,
+      InferUIMessageChunk<UI_MESSAGE>
+    >({
+      transform(chunk, controller) {
+        controller.enqueue(chunk);
+      },
+      // @ts-expect-error cancel is still new and missing from types https://developer.mozilla.org/en-US/docs/Web/API/TransformStream#browser_compatibility
+      async cancel() {
+        await callOnFinish();
+      },
+
+      async flush() {
+        await callOnFinish();
+      },
+    }),
+  );
+}

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

@@ -0,0 +1,13 @@
+export { createUIMessageStream } from './create-ui-message-stream';
+export { createUIMessageStreamResponse } from './create-ui-message-stream-response';
+export { JsonToSseTransformStream } from './json-to-sse-transform-stream';
+export { pipeUIMessageStreamToResponse } from './pipe-ui-message-stream-to-response';
+export { readUIMessageStream } from './read-ui-message-stream';
+export {
+  uiMessageChunkSchema,
+  type InferUIMessageChunk,
+  type UIMessageChunk,
+} from './ui-message-chunks';
+export { UI_MESSAGE_STREAM_HEADERS } from './ui-message-stream-headers';
+export type { UIMessageStreamOnFinishCallback } from './ui-message-stream-on-finish-callback';
+export type { UIMessageStreamWriter } from './ui-message-stream-writer';

package/src/ui-message-stream/json-to-sse-transform-stream.ts

@@ -0,0 +1,12 @@
+export class JsonToSseTransformStream extends TransformStream<unknown, string> {
+  constructor() {
+    super({
+      transform(part, controller) {
+        controller.enqueue(`data: ${JSON.stringify(part)}\n\n`);
+      },
+      flush(controller) {
+        controller.enqueue('data: [DONE]\n\n');
+      },
+    });
+  }
+}

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

@@ -0,0 +1,90 @@
+import { convertArrayToReadableStream } from '@ai-sdk/provider-utils/test';
+import { createMockServerResponse } from '../test/mock-server-response';
+import { pipeUIMessageStreamToResponse } from './pipe-ui-message-stream-to-response';
+import { describe, it, expect } from 'vitest';
+
+describe('pipeUIMessageStreamToResponse', () => {
+  it('should write to ServerResponse with correct headers and encoded stream', async () => {
+    const mockResponse = createMockServerResponse();
+
+    pipeUIMessageStreamToResponse({
+      response: mockResponse,
+      status: 200,
+      statusText: 'OK',
+      headers: {
+        'Custom-Header': 'test',
+      },
+      stream: convertArrayToReadableStream([
+        { type: 'text-start', id: '1' },
+        { type: 'text-delta', id: '1', delta: 'test-data' },
+        { type: 'text-end', id: '1' },
+      ]),
+    });
+
+    // Wait for the stream to finish writing
+    await mockResponse.waitForEnd();
+
+    // Verify response properties
+    expect(mockResponse.statusCode).toBe(200);
+    expect(mockResponse.statusMessage).toBe('OK');
+
+    // Verify headers
+    expect(mockResponse.headers).toMatchInlineSnapshot(`
+      {
+        "cache-control": "no-cache",
+        "connection": "keep-alive",
+        "content-type": "text/event-stream",
+        "custom-header": "test",
+        "x-accel-buffering": "no",
+        "x-vercel-ai-ui-message-stream": "v1",
+      }
+    `);
+
+    // Verify written data using decoded chunks
+    const decodedChunks = mockResponse.getDecodedChunks();
+    expect(decodedChunks).toMatchInlineSnapshot(`
+      [
+        "data: {"type":"text-start","id":"1"}
+
+      ",
+        "data: {"type":"text-delta","id":"1","delta":"test-data"}
+
+      ",
+        "data: {"type":"text-end","id":"1"}
+
+      ",
+        "data: [DONE]
+
+      ",
+      ]
+    `);
+  });
+
+  it('should handle errors in the stream', async () => {
+    const mockResponse = createMockServerResponse();
+
+    pipeUIMessageStreamToResponse({
+      response: mockResponse,
+      status: 200,
+      stream: convertArrayToReadableStream([
+        { type: 'error', errorText: 'Custom error message' },
+      ]),
+    });
+
+    // Wait for the stream to finish writing
+    await mockResponse.waitForEnd();
+
+    // Verify error handling using decoded chunks
+    const decodedChunks = mockResponse.getDecodedChunks();
+    expect(decodedChunks).toMatchInlineSnapshot(`
+      [
+        "data: {"type":"error","errorText":"Custom error message"}
+
+      ",
+        "data: [DONE]
+
+      ",
+      ]
+    `);
+  });
+});

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

@@ -0,0 +1,40 @@
+import { ServerResponse } from 'node:http';
+import { prepareHeaders } from '../util/prepare-headers';
+import { writeToServerResponse } from '../util/write-to-server-response';
+import { JsonToSseTransformStream } from './json-to-sse-transform-stream';
+import { UI_MESSAGE_STREAM_HEADERS } from './ui-message-stream-headers';
+import { UIMessageChunk } from './ui-message-chunks';
+import { UIMessageStreamResponseInit } from './ui-message-stream-response-init';
+
+export function pipeUIMessageStreamToResponse({
+  response,
+  status,
+  statusText,
+  headers,
+  stream,
+  consumeSseStream,
+}: {
+  response: ServerResponse;
+  stream: ReadableStream<UIMessageChunk>;
+} & UIMessageStreamResponseInit): void {
+  let sseStream = stream.pipeThrough(new JsonToSseTransformStream());
+
+  // when the consumeSseStream is provided, we need to tee the stream
+  // and send the second part to the consumeSseStream function
+  // so that it can be consumed by the client independently
+  if (consumeSseStream) {
+    const [stream1, stream2] = sseStream.tee();
+    sseStream = stream1;
+    consumeSseStream({ stream: stream2 }); // no await (do not block the response)
+  }
+
+  writeToServerResponse({
+    response,
+    status,
+    statusText,
+    headers: Object.fromEntries(
+      prepareHeaders(headers, UI_MESSAGE_STREAM_HEADERS).entries(),
+    ),
+    stream: sseStream.pipeThrough(new TextEncoderStream()),
+  });
+}

package/src/ui-message-stream/read-ui-message-stream.test.ts

@@ -0,0 +1,122 @@
+import {
+  convertArrayToReadableStream,
+  convertAsyncIterableToArray,
+} from '@ai-sdk/provider-utils/test';
+import { UIMessageChunk } from './ui-message-chunks';
+import { readUIMessageStream } from './read-ui-message-stream';
+import { describe, it, expect } from 'vitest';
+
+function createUIMessageStream(parts: UIMessageChunk[]) {
+  return convertArrayToReadableStream(parts);
+}
+
+describe('readUIMessageStream', () => {
+  it('should return a ui message object stream for a basic input stream', async () => {
+    const stream = createUIMessageStream([
+      { type: 'start', messageId: 'msg-123' },
+      { type: 'start-step' },
+      { type: 'text-start', id: 'text-1' },
+      { type: 'text-delta', id: 'text-1', delta: 'Hello, ' },
+      { type: 'text-delta', id: 'text-1', delta: 'world!' },
+      { type: 'text-end', id: 'text-1' },
+      { type: 'finish-step' },
+      { type: 'finish' },
+    ]);
+
+    const uiMessages = readUIMessageStream({ stream });
+
+    expect(await convertAsyncIterableToArray(uiMessages))
+      .toMatchInlineSnapshot(`
+        [
+          {
+            "id": "msg-123",
+            "metadata": undefined,
+            "parts": [],
+            "role": "assistant",
+          },
+          {
+            "id": "msg-123",
+            "metadata": undefined,
+            "parts": [
+              {
+                "type": "step-start",
+              },
+              {
+                "providerMetadata": undefined,
+                "state": "streaming",
+                "text": "",
+                "type": "text",
+              },
+            ],
+            "role": "assistant",
+          },
+          {
+            "id": "msg-123",
+            "metadata": undefined,
+            "parts": [
+              {
+                "type": "step-start",
+              },
+              {
+                "providerMetadata": undefined,
+                "state": "streaming",
+                "text": "Hello, ",
+                "type": "text",
+              },
+            ],
+            "role": "assistant",
+          },
+          {
+            "id": "msg-123",
+            "metadata": undefined,
+            "parts": [
+              {
+                "type": "step-start",
+              },
+              {
+                "providerMetadata": undefined,
+                "state": "streaming",
+                "text": "Hello, world!",
+                "type": "text",
+              },
+            ],
+            "role": "assistant",
+          },
+          {
+            "id": "msg-123",
+            "metadata": undefined,
+            "parts": [
+              {
+                "type": "step-start",
+              },
+              {
+                "providerMetadata": undefined,
+                "state": "done",
+                "text": "Hello, world!",
+                "type": "text",
+              },
+            ],
+            "role": "assistant",
+          },
+        ]
+      `);
+  });
+
+  it('should throw an error when encountering an error UI stream part', async () => {
+    const stream = createUIMessageStream([
+      { type: 'start', messageId: 'msg-123' },
+      { type: 'text-start', id: 'text-1' },
+      { type: 'text-delta', id: 'text-1', delta: 'Hello' },
+      { type: 'error', errorText: 'Test error message' },
+    ]);
+
+    const uiMessages = readUIMessageStream({
+      stream,
+      terminateOnError: true,
+    });
+
+    await expect(convertAsyncIterableToArray(uiMessages)).rejects.toThrow(
+      'Test error message',
+    );
+  });
+});

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

@@ -0,0 +1,87 @@
+import { UIMessage } from '../ui/ui-messages';
+import { UIMessageChunk } from './ui-message-chunks';
+import {
+  createStreamingUIMessageState,
+  processUIMessageStream,
+  StreamingUIMessageState,
+} from '../ui/process-ui-message-stream';
+import {
+  AsyncIterableStream,
+  createAsyncIterableStream,
+} from '../util/async-iterable-stream';
+import { consumeStream } from '../util/consume-stream';
+
+/**
+ * Transforms a stream of `UIMessageChunk`s into an `AsyncIterableStream` of `UIMessage`s.
+ *
+ * @param options.message - The last assistant message to use as a starting point when the conversation is resumed. Otherwise undefined.
+ * @param options.stream - The stream of `UIMessageChunk`s to read.
+ * @param options.terminateOnError - Whether to terminate the stream if an error occurs.
+ * @param options.onError - A function that is called when an error occurs.
+ *
+ * @returns An `AsyncIterableStream` of `UIMessage`s. Each stream part is a different state of the same message
+ * as it is being completed.
+ */
+export function readUIMessageStream<UI_MESSAGE extends UIMessage>({
+  message,
+  stream,
+  onError,
+  terminateOnError = false,
+}: {
+  message?: UI_MESSAGE;
+  stream: ReadableStream<UIMessageChunk>;
+  onError?: (error: unknown) => void;
+  terminateOnError?: boolean;
+}): AsyncIterableStream<UI_MESSAGE> {
+  let controller: ReadableStreamDefaultController<UI_MESSAGE> | undefined;
+  let hasErrored = false;
+
+  const outputStream = new ReadableStream<UI_MESSAGE>({
+    start(controllerParam) {
+      controller = controllerParam;
+    },
+  });
+
+  const state = createStreamingUIMessageState<UI_MESSAGE>({
+    messageId: message?.id ?? '',
+    lastMessage: message,
+  });
+
+  const handleError = (error: unknown) => {
+    onError?.(error);
+
+    if (!hasErrored && terminateOnError) {
+      hasErrored = true;
+      controller?.error(error);
+    }
+  };
+
+  consumeStream({
+    stream: processUIMessageStream({
+      stream,
+      runUpdateMessageJob(
+        job: (options: {
+          state: StreamingUIMessageState<UI_MESSAGE>;
+          write: () => void;
+        }) => Promise<void>,
+      ) {
+        return job({
+          state,
+          write: () => {
+            controller?.enqueue(structuredClone(state.message));
+          },
+        });
+      },
+      onError: handleError,
+    }),
+    onError: handleError,
+  }).finally(() => {
+    // Only close if no error occurred. Calling close() on an errored controller
+    // throws "Invalid state: Controller is already closed" TypeError.
+    if (!hasErrored) {
+      controller?.close();
+    }
+  });
+
+  return createAsyncIterableStream(outputStream);
+}

package/src/ui-message-stream/ui-message-chunks.test-d.ts

@@ -0,0 +1,18 @@
+import { describe, expectTypeOf, it } from 'vitest';
+import { UIMessageChunk, uiMessageChunkSchema } from './ui-message-chunks';
+import { validateTypes } from '@ai-sdk/provider-utils';
+
+describe('UI message chunks type', () => {
+  it('parsed UI message chunk should have UIMessageChunk type', async () => {
+    const chunk = await validateTypes({
+      schema: uiMessageChunkSchema,
+      value: {
+        type: 'text-delta',
+        delta: 'Hello, world!',
+        id: '123',
+      },
+    });
+
+    expectTypeOf(chunk).toEqualTypeOf<UIMessageChunk>();
+  });
+});