import { generateId as generateIdFunc, getErrorMessage, type IdGenerator, } from '@ai-sdk/provider-utils'; import type { UIMessage } from '../ui/ui-messages'; import { handleUIMessageStreamFinish } from './handle-ui-message-stream-finish'; import type { InferUIMessageChunk } from './ui-message-chunks'; import type { UIMessageStreamOnFinishCallback } from './ui-message-stream-on-finish-callback'; import type { UIMessageStreamOnStepFinishCallback } from './ui-message-stream-on-step-finish-callback'; import type { UIMessageStreamWriter } from './ui-message-stream-writer'; /** * Creates a UI message stream that can be used to send messages to the client. * * @param options.execute - A function that is called with a writer to write UI message chunks to the stream. * @param options.onError - A function that extracts an error message from an error. Defaults to `getErrorMessage`. * @param options.originalMessages - The original messages. If provided, persistence mode is assumed * and a message ID is provided for the response message. * @param options.onStepFinish - A callback that is called when each step finishes. Useful for persisting intermediate messages. * @param options.onFinish - A callback that is called when the stream finishes. * @param options.generateId - A function that generates a unique ID. Defaults to the built-in ID generator. * * @returns A `ReadableStream` of UI message chunks. */ export function createUIMessageStream({ execute, onError = getErrorMessage, originalMessages, onStepFinish, onFinish, generateId = generateIdFunc, }: { execute: (options: { writer: UIMessageStreamWriter; }) => Promise | 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[]; /** * Callback that is called when each step finishes during multi-step agent runs. */ onStepFinish?: UIMessageStreamOnStepFinishCallback; onFinish?: UIMessageStreamOnFinishCallback; generateId?: IdGenerator; }): ReadableStream> { let controller!: ReadableStreamDefaultController< InferUIMessageChunk >; const ongoingStreamPromises: Promise[] = []; const stream = new ReadableStream({ start(controllerArg) { controller = controllerArg; }, }); function safeEnqueue(data: InferUIMessageChunk) { try { controller.enqueue(data); } catch { // suppress errors when the stream has been closed } } try { const result = execute({ writer: { write(part: InferUIMessageChunk) { 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); }), ); }, onError, }, }); if (result) { ongoingStreamPromises.push( result.catch(error => { safeEnqueue({ type: 'error', errorText: onError(error), } as InferUIMessageChunk); }), ); } } catch (error) { safeEnqueue({ type: 'error', errorText: onError(error), } as InferUIMessageChunk); } // 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 = new Promise(async resolve => { while (ongoingStreamPromises.length > 0) { await ongoingStreamPromises.shift(); } resolve(); }); waitForStreams.finally(() => { try { controller.close(); } catch { // suppress errors when the stream has been closed } }); return handleUIMessageStreamFinish({ stream, messageId: generateId(), originalMessages, onStepFinish, onFinish, onError, }); }