ai 6.0.32 → 6.0.34

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

@@ -0,0 +1,124 @@
+import {
+  generateId as generateIdFunc,
+  getErrorMessage,
+  IdGenerator,
+} from '@ai-sdk/provider-utils';
+import { UIMessage } from '../ui/ui-messages';
+import { handleUIMessageStreamFinish } from './handle-ui-message-stream-finish';
+import { InferUIMessageChunk } from './ui-message-chunks';
+import { UIMessageStreamOnFinishCallback } from './ui-message-stream-on-finish-callback';
+import { UIMessageStreamWriter } from './ui-message-stream-writer';
+
+export function createUIMessageStream<UI_MESSAGE extends UIMessage>({
+  execute,
+  onError = getErrorMessage,
+  originalMessages,
+  onFinish,
+  generateId = generateIdFunc,
+}: {
+  execute: (options: {
+    writer: UIMessageStreamWriter<UI_MESSAGE>;
+  }) => Promise<void> | void;
+  onError?: (error: unknown) => string;
+
+  /**
+   * The original messages. If they are provided, persistence mode is assumed,
+   * and a message ID is provided for the response message.
+   */
+  originalMessages?: UI_MESSAGE[];
+
+  onFinish?: UIMessageStreamOnFinishCallback<UI_MESSAGE>;
+
+  generateId?: IdGenerator;
+}): ReadableStream<InferUIMessageChunk<UI_MESSAGE>> {
+  let controller!: ReadableStreamDefaultController<
+    InferUIMessageChunk<UI_MESSAGE>
+  >;
+
+  const ongoingStreamPromises: Promise<void>[] = [];
+
+  const stream = new ReadableStream({
+    start(controllerArg) {
+      controller = controllerArg;
+    },
+  });
+
+  function safeEnqueue(data: InferUIMessageChunk<UI_MESSAGE>) {
+    try {
+      controller.enqueue(data);
+    } catch (error) {
+      // suppress errors when the stream has been closed
+    }
+  }
+
+  try {
+    const result = execute({
+      writer: {
+        write(part: InferUIMessageChunk<UI_MESSAGE>) {
+          safeEnqueue(part);
+        },
+        merge(streamArg) {
+          ongoingStreamPromises.push(
+            (async () => {
+              const reader = streamArg.getReader();
+              while (true) {
+                const { done, value } = await reader.read();
+                if (done) break;
+                safeEnqueue(value);
+              }
+            })().catch(error => {
+              safeEnqueue({
+                type: 'error',
+                errorText: onError(error),
+              } as InferUIMessageChunk<UI_MESSAGE>);
+            }),
+          );
+        },
+        onError,
+      },
+    });
+
+    if (result) {
+      ongoingStreamPromises.push(
+        result.catch(error => {
+          safeEnqueue({
+            type: 'error',
+            errorText: onError(error),
+          } as InferUIMessageChunk<UI_MESSAGE>);
+        }),
+      );
+    }
+  } catch (error) {
+    safeEnqueue({
+      type: 'error',
+      errorText: onError(error),
+    } as InferUIMessageChunk<UI_MESSAGE>);
+  }
+
+  // Wait until all ongoing streams are done. This approach enables merging
+  // streams even after execute has returned, as long as there is still an
+  // open merged stream. This is important to e.g. forward new streams and
+  // from callbacks.
+  const waitForStreams: Promise<void> = new Promise(async resolve => {
+    while (ongoingStreamPromises.length > 0) {
+      await ongoingStreamPromises.shift();
+    }
+    resolve();
+  });
+
+  waitForStreams.finally(() => {
+    try {
+      controller.close();
+    } catch (error) {
+      // suppress errors when the stream has been closed
+    }
+  });
+
+  return handleUIMessageStreamFinish<UI_MESSAGE>({
+    stream,
+    messageId: generateId(),
+    originalMessages,
+    onFinish,
+    onError,
+  });
+}

package/src/ui-message-stream/get-response-ui-message-id.test.ts

@@ -0,0 +1,55 @@
+import { describe, expect, it } from 'vitest';
+import { getResponseUIMessageId } from './get-response-ui-message-id';
+import { UIMessage } from '../ui/ui-messages';
+
+describe('getResponseUIMessageId', () => {
+  const mockGenerateId = () => 'new-id';
+
+  it('should return undefined when originalMessages is null', () => {
+    const result = getResponseUIMessageId({
+      originalMessages: undefined,
+      responseMessageId: mockGenerateId,
+    });
+    expect(result).toBeUndefined();
+  });
+
+  it('should return the last assistant message id when present', () => {
+    const messages: UIMessage[] = [
+      { id: 'msg-1', role: 'user', parts: [] },
+      { id: 'msg-2', role: 'assistant', parts: [] },
+    ];
+    const result = getResponseUIMessageId({
+      originalMessages: messages,
+      responseMessageId: mockGenerateId,
+    });
+    expect(result).toBe('msg-2');
+  });
+
+  it('should generate new id when last message is not from assistant', () => {
+    const messages: UIMessage[] = [
+      { id: 'msg-1', role: 'assistant', parts: [] },
+      { id: 'msg-2', role: 'user', parts: [] },
+    ];
+    const result = getResponseUIMessageId({
+      originalMessages: messages,
+      responseMessageId: mockGenerateId,
+    });
+    expect(result).toBe('new-id');
+  });
+
+  it('should generate new id when messages array is empty', () => {
+    const result = getResponseUIMessageId({
+      originalMessages: [],
+      responseMessageId: mockGenerateId,
+    });
+    expect(result).toBe('new-id');
+  });
+
+  it('should use the responseMessageId when it is a string', () => {
+    const result = getResponseUIMessageId({
+      originalMessages: [],
+      responseMessageId: 'response-id',
+    });
+    expect(result).toBe('response-id');
+  });
+});

package/src/ui-message-stream/get-response-ui-message-id.ts

@@ -0,0 +1,24 @@
+import { IdGenerator } from '@ai-sdk/provider-utils';
+import { UIMessage } from '../ui/ui-messages';
+
+export function getResponseUIMessageId({
+  originalMessages,
+  responseMessageId,
+}: {
+  originalMessages: UIMessage[] | undefined;
+  responseMessageId: string | IdGenerator;
+}) {
+  // when there are no original messages (i.e. no persistence),
+  // the assistant message id generation is handled on the client side.
+  if (originalMessages == null) {
+    return undefined;
+  }
+
+  const lastMessage = originalMessages[originalMessages.length - 1];
+
+  return lastMessage?.role === 'assistant'
+    ? lastMessage.id
+    : typeof responseMessageId === 'function'
+      ? responseMessageId()
+      : responseMessageId;
+}

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

@@ -0,0 +1,429 @@
+import {
+  convertArrayToReadableStream,
+  convertReadableStreamToArray,
+} from '@ai-sdk/provider-utils/test';
+import { UIMessage } from '../ui/ui-messages';
+import { handleUIMessageStreamFinish } from './handle-ui-message-stream-finish';
+import { UIMessageChunk } from './ui-message-chunks';
+import { beforeEach, describe, expect, it, vi } from 'vitest';
+
+function createUIMessageStream(parts: UIMessageChunk[]) {
+  return convertArrayToReadableStream(parts);
+}
+
+describe('handleUIMessageStreamFinish', () => {
+  const mockErrorHandler = vi.fn();
+
+  beforeEach(() => {
+    mockErrorHandler.mockClear();
+  });
+
+  describe('stream pass-through without onFinish', () => {
+    it('should pass through stream chunks without processing when onFinish is not provided', async () => {
+      const inputChunks: UIMessageChunk[] = [
+        { type: 'start', messageId: 'msg-123' },
+        { 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' },
+      ];
+
+      const stream = createUIMessageStream(inputChunks);
+
+      const resultStream = handleUIMessageStreamFinish<UIMessage>({
+        stream,
+        messageId: 'msg-123',
+        originalMessages: [],
+        onError: mockErrorHandler,
+        // onFinish is not provided
+      });
+
+      const result = await convertReadableStreamToArray(resultStream);
+
+      expect(result).toEqual(inputChunks);
+      expect(mockErrorHandler).not.toHaveBeenCalled();
+    });
+
+    it('should inject messageId when not present in start chunk', async () => {
+      const inputChunks: UIMessageChunk[] = [
+        { type: 'start' }, // no messageId
+        { type: 'text-start', id: 'text-1' },
+        { type: 'text-delta', id: 'text-1', delta: 'Test' },
+        { type: 'text-end', id: 'text-1' },
+        { type: 'finish' },
+      ];
+
+      const stream = createUIMessageStream(inputChunks);
+
+      const resultStream = handleUIMessageStreamFinish<UIMessage>({
+        stream,
+        messageId: 'injected-123',
+        originalMessages: [],
+        onError: mockErrorHandler,
+      });
+
+      const result = await convertReadableStreamToArray(resultStream);
+
+      expect(result[0]).toEqual({ type: 'start', messageId: 'injected-123' });
+      expect(result.slice(1)).toEqual(inputChunks.slice(1));
+    });
+  });
+
+  describe('stream processing with onFinish callback', () => {
+    it('should process stream and call onFinish with correct parameters', async () => {
+      const onFinishCallback = vi.fn();
+      const inputChunks: UIMessageChunk[] = [
+        { type: 'start', messageId: 'msg-456' },
+        { 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' },
+      ];
+
+      const originalMessages: UIMessage[] = [
+        {
+          id: 'user-msg-1',
+          role: 'user',
+          parts: [{ type: 'text', text: 'Hello' }],
+        },
+      ];
+
+      const stream = createUIMessageStream(inputChunks);
+
+      const resultStream = handleUIMessageStreamFinish<UIMessage>({
+        stream,
+        messageId: 'msg-456',
+        originalMessages,
+        onError: mockErrorHandler,
+        onFinish: onFinishCallback,
+      });
+
+      const result = await convertReadableStreamToArray(resultStream);
+
+      expect(result).toEqual(inputChunks);
+      expect(onFinishCallback).toHaveBeenCalledTimes(1);
+
+      const callArgs = onFinishCallback.mock.calls[0][0];
+      expect(callArgs.isContinuation).toBe(false);
+      expect(callArgs.responseMessage.id).toBe('msg-456');
+      expect(callArgs.responseMessage.role).toBe('assistant');
+      expect(callArgs.messages).toHaveLength(2);
+      expect(callArgs.messages[0]).toEqual(originalMessages[0]);
+      expect(callArgs.messages[1]).toEqual(callArgs.responseMessage);
+    });
+
+    it('should handle empty original messages array', async () => {
+      const onFinishCallback = vi.fn();
+      const inputChunks: UIMessageChunk[] = [
+        { type: 'start', messageId: 'msg-789' },
+        { type: 'text-start', id: 'text-1' },
+        { type: 'text-delta', id: 'text-1', delta: 'Response' },
+        { type: 'text-end', id: 'text-1' },
+        { type: 'finish' },
+      ];
+
+      const stream = createUIMessageStream(inputChunks);
+
+      const resultStream = handleUIMessageStreamFinish<UIMessage>({
+        stream,
+        messageId: 'msg-789',
+        originalMessages: [],
+        onError: mockErrorHandler,
+        onFinish: onFinishCallback,
+      });
+
+      await convertReadableStreamToArray(resultStream);
+
+      expect(onFinishCallback).toHaveBeenCalledTimes(1);
+
+      const callArgs = onFinishCallback.mock.calls[0][0];
+      expect(callArgs.isContinuation).toBe(false);
+      expect(callArgs.messages).toHaveLength(1);
+      expect(callArgs.messages[0]).toEqual(callArgs.responseMessage);
+    });
+  });
+
+  describe('stream processing with continuation scenario', () => {
+    it('should handle continuation when last message is assistant', async () => {
+      const onFinishCallback = vi.fn();
+      const inputChunks: UIMessageChunk[] = [
+        { type: 'start', messageId: 'assistant-msg-1' },
+        { type: 'text-start', id: 'text-1' },
+        { type: 'text-delta', id: 'text-1', delta: ' continued' },
+        { type: 'text-end', id: 'text-1' },
+        { type: 'finish' },
+      ];
+
+      const originalMessages: UIMessage[] = [
+        {
+          id: 'user-msg-1',
+          role: 'user',
+          parts: [{ type: 'text', text: 'Continue this' }],
+        },
+        {
+          id: 'assistant-msg-1',
+          role: 'assistant',
+          parts: [{ type: 'text', text: 'This is' }],
+        },
+      ];
+
+      const stream = createUIMessageStream(inputChunks);
+
+      const resultStream = handleUIMessageStreamFinish<UIMessage>({
+        stream,
+        messageId: 'msg-999', // this should be ignored since we're continuing
+        originalMessages,
+        onError: mockErrorHandler,
+        onFinish: onFinishCallback,
+      });
+
+      await convertReadableStreamToArray(resultStream);
+
+      expect(onFinishCallback).toHaveBeenCalledTimes(1);
+
+      const callArgs = onFinishCallback.mock.calls[0][0];
+      expect(callArgs.isContinuation).toBe(true);
+      expect(callArgs.responseMessage.id).toBe('assistant-msg-1'); // uses the existing assistant message id
+      expect(callArgs.messages).toHaveLength(2); // original user message + updated assistant message
+      expect(callArgs.messages[0]).toEqual(originalMessages[0]);
+      expect(callArgs.messages[1]).toEqual(callArgs.responseMessage);
+    });
+
+    it('should not treat user message as continuation', async () => {
+      const onFinishCallback = vi.fn();
+      const inputChunks: UIMessageChunk[] = [
+        { type: 'start', messageId: 'msg-001' },
+        { type: 'text-start', id: 'text-1' },
+        { type: 'text-delta', id: 'text-1', delta: 'New response' },
+        { type: 'text-end', id: 'text-1' },
+        { type: 'finish' },
+      ];
+
+      const originalMessages: UIMessage[] = [
+        {
+          id: 'user-msg-1',
+          role: 'user',
+          parts: [{ type: 'text', text: 'Question' }],
+        },
+        {
+          id: 'user-msg-2',
+          role: 'user',
+          parts: [{ type: 'text', text: 'Another question' }],
+        },
+      ];
+
+      const stream = createUIMessageStream(inputChunks);
+
+      const resultStream = handleUIMessageStreamFinish<UIMessage>({
+        stream,
+        messageId: 'msg-001',
+        originalMessages,
+        onError: mockErrorHandler,
+        onFinish: onFinishCallback,
+      });
+
+      await convertReadableStreamToArray(resultStream);
+
+      expect(onFinishCallback).toHaveBeenCalledTimes(1);
+
+      const callArgs = onFinishCallback.mock.calls[0][0];
+      expect(callArgs.isContinuation).toBe(false);
+      expect(callArgs.responseMessage.id).toBe('msg-001');
+      expect(callArgs.messages).toHaveLength(3); // 2 user messages + 1 new assistant message
+    });
+  });
+
+  describe('abort scenarios', () => {
+    it('should set isAborted to true when abort chunk is encountered', async () => {
+      const onFinishCallback = vi.fn();
+      const inputChunks: UIMessageChunk[] = [
+        { type: 'start', messageId: 'msg-abort-1' },
+        { type: 'text-start', id: 'text-1' },
+        { type: 'text-delta', id: 'text-1', delta: 'Starting text' },
+        { type: 'abort' },
+        { type: 'finish' },
+      ];
+
+      const originalMessages: UIMessage[] = [
+        {
+          id: 'user-msg-1',
+          role: 'user',
+          parts: [{ type: 'text', text: 'Test request' }],
+        },
+      ];
+
+      const stream = createUIMessageStream(inputChunks);
+
+      const resultStream = handleUIMessageStreamFinish<UIMessage>({
+        stream,
+        messageId: 'msg-abort-1',
+        originalMessages,
+        onError: mockErrorHandler,
+        onFinish: onFinishCallback,
+      });
+
+      const result = await convertReadableStreamToArray(resultStream);
+
+      expect(result).toEqual(inputChunks);
+      expect(onFinishCallback).toHaveBeenCalledTimes(1);
+
+      const callArgs = onFinishCallback.mock.calls[0][0];
+      expect(callArgs.isAborted).toBe(true);
+      expect(callArgs.isContinuation).toBe(false);
+      expect(callArgs.responseMessage.id).toBe('msg-abort-1');
+      expect(callArgs.messages).toHaveLength(2);
+    });
+
+    it('should pass through abort reason when provided', async () => {
+      const onFinishCallback = vi.fn();
+      const inputChunks: UIMessageChunk[] = [
+        { type: 'start', messageId: 'msg-abort-reason' },
+        { type: 'text-start', id: 'text-1' },
+        { type: 'text-delta', id: 'text-1', delta: 'Starting text' },
+        { type: 'abort', reason: 'manual abort' },
+        { type: 'finish' },
+      ];
+
+      const stream = createUIMessageStream(inputChunks);
+
+      const resultStream = handleUIMessageStreamFinish<UIMessage>({
+        stream,
+        messageId: 'msg-abort-reason',
+        originalMessages: [],
+        onError: mockErrorHandler,
+        onFinish: onFinishCallback,
+      });
+
+      const result = await convertReadableStreamToArray(resultStream);
+
+      expect(result).toEqual(inputChunks);
+      expect(onFinishCallback).toHaveBeenCalledTimes(1);
+    });
+
+    it('should set isAborted to false when no abort chunk is encountered', async () => {
+      const onFinishCallback = vi.fn();
+      const inputChunks: UIMessageChunk[] = [
+        { type: 'start', messageId: 'msg-normal' },
+        { type: 'text-start', id: 'text-1' },
+        { type: 'text-delta', id: 'text-1', delta: 'Complete text' },
+        { type: 'text-end', id: 'text-1' },
+        { type: 'finish' },
+      ];
+
+      const originalMessages: UIMessage[] = [
+        {
+          id: 'user-msg-1',
+          role: 'user',
+          parts: [{ type: 'text', text: 'Test request' }],
+        },
+      ];
+
+      const stream = createUIMessageStream(inputChunks);
+
+      const resultStream = handleUIMessageStreamFinish<UIMessage>({
+        stream,
+        messageId: 'msg-normal',
+        originalMessages,
+        onError: mockErrorHandler,
+        onFinish: onFinishCallback,
+      });
+
+      await convertReadableStreamToArray(resultStream);
+
+      expect(onFinishCallback).toHaveBeenCalledTimes(1);
+
+      const callArgs = onFinishCallback.mock.calls[0][0];
+      expect(callArgs.isAborted).toBe(false);
+      expect(callArgs.isContinuation).toBe(false);
+      expect(callArgs.responseMessage.id).toBe('msg-normal');
+    });
+
+    it('should handle abort chunk in pass-through mode without onFinish', async () => {
+      const inputChunks: UIMessageChunk[] = [
+        { type: 'start', messageId: 'msg-abort-passthrough' },
+        { type: 'text-start', id: 'text-1' },
+        { type: 'text-delta', id: 'text-1', delta: 'Text before abort' },
+        { type: 'abort' },
+        { type: 'finish' },
+      ];
+
+      const stream = createUIMessageStream(inputChunks);
+
+      const resultStream = handleUIMessageStreamFinish<UIMessage>({
+        stream,
+        messageId: 'msg-abort-passthrough',
+        originalMessages: [],
+        onError: mockErrorHandler,
+        // onFinish is not provided
+      });
+
+      const result = await convertReadableStreamToArray(resultStream);
+
+      expect(result).toEqual(inputChunks);
+      expect(mockErrorHandler).not.toHaveBeenCalled();
+    });
+
+    it('should handle multiple abort chunks correctly', async () => {
+      const onFinishCallback = vi.fn();
+      const inputChunks: UIMessageChunk[] = [
+        { type: 'start', messageId: 'msg-multiple-abort' },
+        { type: 'text-start', id: 'text-1' },
+        { type: 'abort' },
+        { type: 'text-delta', id: 'text-1', delta: 'Some text' },
+        { type: 'abort' },
+        { type: 'finish' },
+      ];
+
+      const stream = createUIMessageStream(inputChunks);
+
+      const resultStream = handleUIMessageStreamFinish<UIMessage>({
+        stream,
+        messageId: 'msg-multiple-abort',
+        originalMessages: [],
+        onError: mockErrorHandler,
+        onFinish: onFinishCallback,
+      });
+
+      const result = await convertReadableStreamToArray(resultStream);
+
+      expect(result).toEqual(inputChunks);
+      expect(onFinishCallback).toHaveBeenCalledTimes(1);
+
+      const callArgs = onFinishCallback.mock.calls[0][0];
+      expect(callArgs.isAborted).toBe(true);
+    });
+
+    it('should call onFinish when reader is cancelled (simulating browser close/navigation)', async () => {
+      const onFinishCallback = vi.fn();
+
+      const inputChunks: UIMessageChunk[] = [
+        { type: 'start', messageId: 'msg-1' },
+        { type: 'text-start', id: 'text-1' },
+        { type: 'text-delta', id: 'text-1', delta: 'Hello' },
+      ];
+
+      const stream = createUIMessageStream(inputChunks);
+
+      const resultStream = handleUIMessageStreamFinish<UIMessage>({
+        stream,
+        messageId: 'msg-1',
+        originalMessages: [],
+        onError: mockErrorHandler,
+        onFinish: onFinishCallback,
+      });
+
+      const reader = resultStream.getReader();
+      await reader.read();
+      await reader.cancel();
+      reader.releaseLock();
+
+      expect(onFinishCallback).toHaveBeenCalledTimes(1);
+
+      const callArgs = onFinishCallback.mock.calls[0][0];
+      expect(callArgs.isAborted).toBe(false);
+      expect(callArgs.responseMessage.id).toBe('msg-1');
+    });
+  });
+});