@ai-sdk/rsc 2.0.44 → 2.0.46

package/src/stream-ui/stream-ui.ui.test.tsx

@@ -0,0 +1,321 @@
+import { LanguageModelV3Usage } from '@ai-sdk/provider';
+import { delay } from '@ai-sdk/provider-utils';
+import { convertArrayToReadableStream } from '@ai-sdk/provider-utils/test';
+import { asLanguageModelUsage } from 'ai/internal';
+import { MockLanguageModelV3 } from 'ai/test';
+import { beforeEach, describe, expect, it } from 'vitest';
+import { z } from 'zod/v4';
+import { streamUI } from './stream-ui';
+
+async function recursiveResolve(val: any): Promise<any> {
+  if (val && typeof val === 'object' && typeof val.then === 'function') {
+    return await recursiveResolve(await val);
+  }
+
+  if (Array.isArray(val)) {
+    return await Promise.all(val.map(recursiveResolve));
+  }
+
+  if (val && typeof val === 'object') {
+    const result: any = {};
+    for (const key in val) {
+      result[key] = await recursiveResolve(val[key]);
+    }
+    return result;
+  }
+
+  return val;
+}
+
+async function simulateFlightServerRender(node: React.ReactNode) {
+  async function traverse(node: any): Promise<any> {
+    if (!node) return {};
+
+    // Let's only do one level of promise resolution here. As it's only for testing purposes.
+    const props = await recursiveResolve({ ...node.props });
+
+    const { type } = node;
+    const { children, ...otherProps } = props;
+    const typeName = typeof type === 'function' ? type.name : String(type);
+
+    return {
+      type: typeName,
+      props: otherProps,
+      children:
+        typeof children === 'string'
+          ? children
+          : Array.isArray(children)
+            ? children.map(traverse)
+            : await traverse(children),
+    };
+  }
+
+  return traverse(node);
+}
+
+const testUsage: LanguageModelV3Usage = {
+  inputTokens: {
+    total: 3,
+    noCache: 3,
+    cacheRead: 0,
+    cacheWrite: 0,
+  },
+  outputTokens: {
+    total: 10,
+    text: 10,
+    reasoning: 0,
+  },
+};
+
+const mockTextModel = new MockLanguageModelV3({
+  doStream: async () => {
+    return {
+      stream: convertArrayToReadableStream([
+        { type: 'text-start', id: '0' },
+        { type: 'text-delta', id: '0', delta: '{ ' },
+        { type: 'text-delta', id: '0', delta: '"content": ' },
+        { type: 'text-delta', id: '0', delta: `"Hello, ` },
+        { type: 'text-delta', id: '0', delta: `world` },
+        { type: 'text-delta', id: '0', delta: `!"` },
+        { type: 'text-delta', id: '0', delta: ' }' },
+        { type: 'text-end', id: '0' },
+        {
+          type: 'finish',
+          finishReason: { unified: 'stop', raw: 'stop' },
+          usage: testUsage,
+        },
+      ]),
+    };
+  },
+});
+
+const mockToolModel = new MockLanguageModelV3({
+  doStream: async () => {
+    return {
+      stream: convertArrayToReadableStream([
+        {
+          type: 'tool-call',
+          toolCallType: 'function',
+          toolCallId: 'call-1',
+          toolName: 'tool1',
+          input: `{ "value": "value" }`,
+        },
+        {
+          type: 'finish',
+          finishReason: { unified: 'stop', raw: 'stop' },
+          usage: testUsage,
+        },
+      ]),
+    };
+  },
+});
+
+describe('result.value', () => {
+  it('should render text', async () => {
+    const result = await streamUI({
+      model: mockTextModel,
+      prompt: '',
+    });
+
+    const rendered = await simulateFlightServerRender(result.value);
+    expect(rendered).toMatchSnapshot();
+  });
+
+  it('should render text function returned ui', async () => {
+    const result = await streamUI({
+      model: mockTextModel,
+      prompt: '',
+      text: ({ content }) => <h1>{content}</h1>,
+    });
+
+    const rendered = await simulateFlightServerRender(result.value);
+    expect(rendered).toMatchSnapshot();
+  });
+
+  it('should render tool call results', async () => {
+    const result = await streamUI({
+      model: mockToolModel,
+      prompt: '',
+      tools: {
+        tool1: {
+          description: 'test tool 1',
+          inputSchema: z.object({
+            value: z.string(),
+          }),
+          generate: async ({ value }) => {
+            await delay(100);
+            return <div>tool1: {value}</div>;
+          },
+        },
+      },
+    });
+
+    const rendered = await simulateFlightServerRender(result.value);
+    expect(rendered).toMatchSnapshot();
+  });
+
+  it('should render tool call results with generator render function', async () => {
+    const result = await streamUI({
+      model: mockToolModel,
+      prompt: '',
+      tools: {
+        tool1: {
+          description: 'test tool 1',
+          inputSchema: z.object({
+            value: z.string(),
+          }),
+          generate: async function* ({ value }) {
+            yield <div>Loading...</div>;
+            await delay(100);
+            return <div>tool: {value}</div>;
+          },
+        },
+      },
+    });
+
+    const rendered = await simulateFlightServerRender(result.value);
+    expect(rendered).toMatchSnapshot();
+  });
+
+  it('should show better error messages if legacy options are passed', async () => {
+    try {
+      await streamUI({
+        model: mockToolModel,
+        prompt: '',
+        tools: {
+          tool1: {
+            description: 'test tool 1',
+            inputSchema: z.object({
+              value: z.string(),
+            }),
+            // @ts-expect-error
+            render: async function* () {},
+          },
+        },
+      });
+    } catch (e) {
+      expect(e).toMatchSnapshot();
+    }
+  });
+});
+
+describe('rsc - streamUI() onFinish callback', () => {
+  let result: Parameters<
+    Required<Parameters<typeof streamUI>[0]>['onFinish']
+  >[0];
+
+  beforeEach(async () => {
+    const ui = await streamUI({
+      model: mockToolModel,
+      prompt: '',
+      tools: {
+        tool1: {
+          description: 'test tool 1',
+          inputSchema: z.object({
+            value: z.string(),
+          }),
+          generate: async ({ value }) => {
+            await delay(100);
+            return <div>tool1: {value}</div>;
+          },
+        },
+      },
+      onFinish: event => {
+        result = event;
+      },
+    });
+
+    // consume stream
+    await simulateFlightServerRender(ui.value);
+  });
+
+  it('should contain token usage', () => {
+    expect(result.usage).toStrictEqual(asLanguageModelUsage(testUsage));
+  });
+
+  it('should contain finish reason', async () => {
+    expect(result.finishReason).toBe('stop');
+  });
+
+  it('should contain final React node', async () => {
+    expect(result.value).toMatchSnapshot();
+  });
+});
+
+describe('options.headers', () => {
+  it('should pass headers to model', async () => {
+    const result = await streamUI({
+      model: new MockLanguageModelV3({
+        doStream: async ({ headers }) => {
+          expect(headers).toStrictEqual({
+            'custom-request-header': 'request-header-value',
+          });
+
+          return {
+            stream: convertArrayToReadableStream([
+              { type: 'text-start', id: '0' },
+              {
+                type: 'text-delta',
+                id: '0',
+                delta: '{ "content": "headers test" }',
+              },
+              { type: 'text-end', id: '0' },
+              {
+                type: 'finish',
+                finishReason: {
+                  unified: 'stop',
+                  raw: 'stop',
+                },
+                usage: testUsage,
+              },
+            ]),
+          };
+        },
+      }),
+      prompt: '',
+      headers: { 'custom-request-header': 'request-header-value' },
+    });
+
+    expect(await simulateFlightServerRender(result.value)).toMatchSnapshot();
+  });
+});
+
+describe('options.providerMetadata', () => {
+  it('should pass provider metadata to model', async () => {
+    const result = await streamUI({
+      model: new MockLanguageModelV3({
+        doStream: async ({ providerOptions }) => {
+          expect(providerOptions).toStrictEqual({
+            aProvider: { someKey: 'someValue' },
+          });
+
+          return {
+            stream: convertArrayToReadableStream([
+              { type: 'text-start', id: '0' },
+              {
+                type: 'text-delta',
+                id: '0',
+                delta: '{ "content": "provider metadata test" }',
+              },
+              { type: 'text-end', id: '0' },
+              {
+                type: 'finish',
+                usage: testUsage,
+                finishReason: {
+                  unified: 'stop',
+                  raw: 'stop',
+                },
+              },
+            ]),
+          };
+        },
+      }),
+      prompt: '',
+      providerOptions: {
+        aProvider: { someKey: 'someValue' },
+      },
+    });
+
+    expect(await simulateFlightServerRender(result.value)).toMatchSnapshot();
+  });
+});

package/src/streamable-ui/create-streamable-ui.tsx

@@ -0,0 +1,148 @@
+import { HANGING_STREAM_WARNING_TIME_MS } from '../util/constants';
+import { createResolvablePromise } from '../util/create-resolvable-promise';
+import { createSuspendedChunk } from './create-suspended-chunk';
+
+// It's necessary to define the type manually here, otherwise TypeScript compiler
+// will not be able to infer the correct return type as it's circular.
+type StreamableUIWrapper = {
+  /**
+   * The value of the streamable UI. This can be returned from a Server Action and received by the client.
+   */
+  readonly value: React.ReactNode;
+
+  /**
+   * This method updates the current UI node. It takes a new UI node and replaces the old one.
+   */
+  update(value: React.ReactNode): StreamableUIWrapper;
+
+  /**
+   * This method is used to append a new UI node to the end of the old one.
+   * Once appended a new UI node, the previous UI node cannot be updated anymore.
+   *
+   * @example
+   * ```jsx
+   * const ui = createStreamableUI(<div>hello</div>)
+   * ui.append(<div>world</div>)
+   *
+   * // The UI node will be:
+   * // <>
+   * //   <div>hello</div>
+   * //   <div>world</div>
+   * // </>
+   * ```
+   */
+  append(value: React.ReactNode): StreamableUIWrapper;
+
+  /**
+   * This method is used to signal that there is an error in the UI stream.
+   * It will be thrown on the client side and caught by the nearest error boundary component.
+   */
+  error(error: any): StreamableUIWrapper;
+
+  /**
+   * This method marks the UI node as finalized. You can either call it without any parameters or with a new UI node as the final state.
+   * Once called, the UI node cannot be updated or appended anymore.
+   *
+   * This method is always **required** to be called, otherwise the response will be stuck in a loading state.
+   */
+  done(...args: [React.ReactNode] | []): StreamableUIWrapper;
+};
+
+/**
+ * Create a piece of changeable UI that can be streamed to the client.
+ * On the client side, it can be rendered as a normal React node.
+ */
+function createStreamableUI(initialValue?: React.ReactNode) {
+  let currentValue = initialValue;
+  let closed = false;
+  let { row, resolve, reject } = createSuspendedChunk(initialValue);
+
+  function assertStream(method: string) {
+    if (closed) {
+      throw new Error(method + ': UI stream is already closed.');
+    }
+  }
+
+  let warningTimeout: NodeJS.Timeout | undefined;
+  function warnUnclosedStream() {
+    if (process.env.NODE_ENV === 'development') {
+      if (warningTimeout) {
+        clearTimeout(warningTimeout);
+      }
+      warningTimeout = setTimeout(() => {
+        console.warn(
+          'The streamable UI has been slow to update. This may be a bug or a performance issue or you forgot to call `.done()`.',
+        );
+      }, HANGING_STREAM_WARNING_TIME_MS);
+    }
+  }
+  warnUnclosedStream();
+
+  const streamable: StreamableUIWrapper = {
+    value: row,
+    update(value: React.ReactNode) {
+      assertStream('.update()');
+
+      // There is no need to update the value if it's referentially equal.
+      if (value === currentValue) {
+        warnUnclosedStream();
+        return streamable;
+      }
+
+      const resolvable = createResolvablePromise();
+      currentValue = value;
+
+      resolve({ value: currentValue, done: false, next: resolvable.promise });
+      resolve = resolvable.resolve;
+      reject = resolvable.reject;
+
+      warnUnclosedStream();
+
+      return streamable;
+    },
+    append(value: React.ReactNode) {
+      assertStream('.append()');
+
+      const resolvable = createResolvablePromise();
+      currentValue = value;
+
+      resolve({ value, done: false, append: true, next: resolvable.promise });
+      resolve = resolvable.resolve;
+      reject = resolvable.reject;
+
+      warnUnclosedStream();
+
+      return streamable;
+    },
+    error(error: any) {
+      assertStream('.error()');
+
+      if (warningTimeout) {
+        clearTimeout(warningTimeout);
+      }
+      closed = true;
+      reject(error);
+
+      return streamable;
+    },
+    done(...args: [] | [React.ReactNode]) {
+      assertStream('.done()');
+
+      if (warningTimeout) {
+        clearTimeout(warningTimeout);
+      }
+      closed = true;
+      if (args.length) {
+        resolve({ value: args[0], done: true });
+        return streamable;
+      }
+      resolve({ value: currentValue, done: true });
+
+      return streamable;
+    },
+  };
+
+  return streamable;
+}
+
+export { createStreamableUI };