@fondation-io/ai 7.0.0-beta.45

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

@@ -0,0 +1,145 @@
+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 { UIMessageStreamOnStepFinishCallback } from './ui-message-stream-on-step-finish-callback';
+import { 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<UI_MESSAGE extends UIMessage>({
+  execute,
+  onError = getErrorMessage,
+  originalMessages,
+  onStepFinish,
+  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[];
+
+  /**
+   * Callback that is called when each step finishes during multi-step agent runs.
+   */
+  onStepFinish?: UIMessageStreamOnStepFinishCallback<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,
+    onStepFinish,
+    onFinish,
+    onError,
+  });
+}

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

@@ -0,0 +1,35 @@
+import { IdGenerator } from '@ai-sdk/provider-utils';
+import { UIMessage } from '../ui/ui-messages';
+
+/**
+ * Determines the message ID to use for a response message.
+ * If the last message is an assistant message, its ID is reused (continuation).
+ * Otherwise, a new ID is generated or the provided ID is used.
+ *
+ * @param options.originalMessages - The original messages. If not provided, returns `undefined`
+ *   since client-side ID generation is used in non-persistence mode.
+ * @param options.responseMessageId - The response message ID or an ID generator function.
+ *
+ * @returns The message ID to use, or `undefined` if no persistence mode.
+ */
+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.ts

@@ -0,0 +1,170 @@
+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';
+import { UIMessageStreamOnStepFinishCallback } from './ui-message-stream-on-step-finish-callback';
+
+export function handleUIMessageStreamFinish<UI_MESSAGE extends UIMessage>({
+  messageId,
+  originalMessages = [],
+  onStepFinish,
+  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;
+
+  /**
+   * Callback that is called when each step finishes during multi-step agent runs.
+   */
+  onStepFinish?: UIMessageStreamOnStepFinishCallback<UI_MESSAGE>;
+
+  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);
+      },
+    }),
+  );
+
+  // Only process the stream if we need to track state for callbacks
+  if (onFinish == null && onStepFinish == 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,
+    });
+  };
+
+  const callOnStepFinish = async () => {
+    if (!onStepFinish) {
+      return;
+    }
+
+    const isContinuation = state.message.id === lastMessage?.id;
+
+    try {
+      await onStepFinish({
+        isContinuation,
+        responseMessage: structuredClone(state.message) as UI_MESSAGE,
+        messages: [
+          ...(isContinuation
+            ? originalMessages.slice(0, -1)
+            : originalMessages),
+          structuredClone(state.message),
+        ] as UI_MESSAGE[],
+      });
+    } catch (error) {
+      onError(error);
+    }
+  };
+
+  return processUIMessageStream<UI_MESSAGE>({
+    stream: idInjectedStream,
+    runUpdateMessageJob,
+    onError,
+  }).pipeThrough(
+    new TransformStream<
+      InferUIMessageChunk<UI_MESSAGE>,
+      InferUIMessageChunk<UI_MESSAGE>
+    >({
+      async transform(chunk, controller) {
+        if (chunk.type === 'finish-step') {
+          await callOnStepFinish();
+        }
+
+        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,14 @@
+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 { UIMessageStreamOnStepFinishCallback } from './ui-message-stream-on-step-finish-callback';
+export type { UIMessageStreamWriter } from './ui-message-stream-writer';

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

@@ -0,0 +1,17 @@
+/**
+ * A TransformStream that converts JSON objects to Server-Sent Events (SSE) format.
+ * Each object is serialized to JSON and wrapped in `data: ...\n\n` format.
+ * When the stream ends, a `data: [DONE]\n\n` message is sent.
+ */
+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.ts

@@ -0,0 +1,51 @@
+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';
+
+/**
+ * Pipes a UI message stream to a Node.js ServerResponse object.
+ * The stream is transformed to Server-Sent Events (SSE) format.
+ *
+ * @param options.response - The Node.js ServerResponse object to write to.
+ * @param options.status - The HTTP status code for the response.
+ * @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.consumeSseStream - Optional callback to consume a copy of the SSE stream independently.
+ */
+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.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);
+}