fix(ai/rsc): New createStreamableUI implementation (#1825)

Co-authored-by: Lars Grammel <lars.grammel@gmail.com>

Author: shuding

.changeset/early-hounds-flow.md

@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+fix(ai/rsc): Refactor streamable UI internal implementation

packages/core/rsc/__snapshots__/streamable.ui.test.tsx.snap

@@ -8,4 +8,5 @@ export {
   useActions,
   useSyncUIState,
   InternalAIProvider,
+  InternalStreamableUIClient,
 } from './shared-client';

packages/core/rsc/rsc-shared.mts

@@ -0,0 +1,56 @@
+'use client';
+
+import { useState, useEffect } from 'react';
+import { StreamableValue } from '../types';
+import { readStreamableValue } from './streamable';
+
+export function InternalStreamableUIClient<T>({
+  s,
+}: {
+  s: StreamableValue<T>;
+}) {
+  // Set the value to the initial value of the streamable, if it has one.
+  const [value, setValue] = useState<T | undefined>(s.curr);
+
+  // Error state for the streamable. It might be errored initially and we want
+  // to error out as soon as possible.
+  const [error, setError] = useState<Error | undefined>(s.error);
+
+  useEffect(() => {
+    let canceled = false;
+    setError(undefined);
+
+    (async () => {
+      try {
+        // Read the streamable value and update the state with the new value.
+        for await (const v of readStreamableValue(s)) {
+          if (canceled) {
+            break;
+          }
+
+          setValue(v);
+        }
+      } catch (e) {
+        if (canceled) {
+          return;
+        }
+
+        setError(e as Error);
+      }
+    })();
+
+    return () => {
+      // If the component is unmounted, we want to cancel the stream.
+      canceled = true;
+    };
+  }, [s]);
+
+  // This ensures that errors from the streamable UI are thrown during the
+  // render phase, so that they can be caught by error boundary components.
+  // This is necessary for React's declarative model.
+  if (error) {
+    throw error;
+  }
+
+  return value;
+}

packages/core/rsc/shared-client/client-wrapper.tsx

@@ -8,3 +8,4 @@ export {
   useSyncUIState,
   InternalAIProvider,
 } from './context';
+export { InternalStreamableUIClient } from './client-wrapper';

packages/core/rsc/shared-client/index.ts

@@ -1,4 +1,9 @@
-import { startTransition, useLayoutEffect, useState } from 'react';
+import {
+  ReactElement,
+  startTransition,
+  useLayoutEffect,
+  useState,
+} from 'react';
 import { STREAMABLE_VALUE_TYPE } from '../constants';
 import type { StreamableValue } from '../types';
 
@@ -97,6 +102,14 @@ export function readStreamableValue<T = unknown>(
                     (curr as string) = curr + row.diff[1];
                   }
                   break;
+                case 1:
+                  (curr as ReactElement) = (
+                    <>
+                      {curr}
+                      {row.diff[1]}
+                    </>
+                  );
+                  break;
               }
             } else {
               curr = row.curr;

packages/core/rsc/shared-client/streamable.tsx

@@ -1,4 +1,4 @@
-import type { ReactNode } from 'react';
+import { type ReactNode, isValidElement, ReactElement } from 'react';
 import type OpenAI from 'openai';
 import { z } from 'zod';
 import zodToJsonSchema from 'zod-to-json-schema';
@@ -10,12 +10,9 @@ import {
   STREAMABLE_VALUE_TYPE,
   DEV_DEFAULT_STREAMABLE_WARNING_TIME,
 } from './constants';
-import {
-  createResolvablePromise,
-  createSuspensedChunk,
-  consumeStream,
-} from './utils';
+import { createResolvablePromise, consumeStream } from './utils';
 import type { StreamablePatch, StreamableValue } from './types';
+import { InternalStreamableUIClient } from './rsc-shared.mjs';
 
 // 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.
@@ -68,9 +65,9 @@ type StreamableUIWrapper = {
  * On the client side, it can be rendered as a normal React node.
  */
 function createStreamableUI(initialValue?: React.ReactNode) {
-  let currentValue = initialValue;
+  const innerStreamable = createStreamableValue<React.ReactNode>(initialValue);
+
   let closed = false;
-  let { row, resolve, reject } = createSuspensedChunk(initialValue);
 
   function assertStream(method: string) {
     if (closed) {
@@ -94,37 +91,19 @@ function createStreamableUI(initialValue?: React.ReactNode) {
   warnUnclosedStream();
 
   const streamable: StreamableUIWrapper = {
-    value: row,
+    value: <InternalStreamableUIClient s={innerStreamable.value} />,
     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;
-
+      innerStreamable.update(value);
       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;
-
+      innerStreamable.append(value);
       warnUnclosedStream();
 
       return streamable;
@@ -136,7 +115,7 @@ function createStreamableUI(initialValue?: React.ReactNode) {
         clearTimeout(warningTimeout);
       }
       closed = true;
-      reject(error);
+      innerStreamable.error(error);
 
       return streamable;
     },
@@ -148,11 +127,11 @@ function createStreamableUI(initialValue?: React.ReactNode) {
       }
       closed = true;
       if (args.length) {
-        resolve({ value: args[0], done: true });
+        innerStreamable.done(args[0]);
         return streamable;
       }
-      resolve({ value: currentValue, done: true });
 
+      innerStreamable.done();
       return streamable;
     },
   };
@@ -377,31 +356,45 @@ function createStreamableValueImpl<T = any, E = any>(initialValue?: T) {
     append(value: T) {
       assertStream('.append()');
 
-      if (
-        typeof currentValue !== 'string' &&
-        typeof currentValue !== 'undefined'
-      ) {
+      if (typeof value !== 'string' && !isValidElement(value)) {
         throw new Error(
-          `.append(): The current value is not a string. Received: ${typeof currentValue}`,
+          `.append(): The value type can't be appended to the stream. Received: ${typeof value}`,
         );
       }
-      if (typeof value !== 'string') {
+
+      if (typeof currentValue === 'undefined') {
+        currentPatchValue = undefined;
+        currentValue = value;
+      } else if (typeof currentValue === 'string') {
+        if (typeof value === 'string') {
+          currentPatchValue = [0, value];
+          (currentValue as string) = currentValue + value;
+        } else {
+          currentPatchValue = [1, value];
+          (currentValue as unknown as ReactElement) = (
+            <>
+              {currentValue}
+              {value}
+            </>
+          );
+        }
+      } else if (isValidElement(currentValue)) {
+        currentPatchValue = [1, value];
+        (currentValue as ReactElement) = (
+          <>
+            {currentValue}
+            {value}
+          </>
+        );
+      } else {
         throw new Error(
-          `.append(): The value is not a string. Received: ${typeof value}`,
+          `.append(): The current value doesn't support appending data. Type: ${typeof currentValue}`,
         );
       }
 
       const resolvePrevious = resolvable.resolve;
       resolvable = createResolvablePromise();
 
-      if (typeof currentValue === 'string') {
-        currentPatchValue = [0, value];
-        (currentValue as string) = currentValue + value;
-      } else {
-        currentPatchValue = undefined;
-        currentValue = value;
-      }
-
       currentPromise = resolvable.promise;
       resolvePrevious(createWrapped());
 

packages/core/rsc/stream-ui/__snapshots__/stream-ui.ui.test.tsx.snap

@@ -87,7 +87,10 @@ export type MutableAIState<AIState> = {
   done: ((newState: AIState) => void) | (() => void);
 };
 
-export type StreamablePatch = undefined | [0, string]; // Append string.
+export type StreamablePatch =
+  | undefined
+  | [0, string] // Append string
+  | [1, React.ReactElement | string]; // Append element
 
 declare const __internal_curr: unique symbol;
 declare const __internal_error: unique symbol;

packages/core/rsc/stream-ui/stream-ui.ui.test.tsx

@@ -1,5 +1,3 @@
-import React, { Suspense } from 'react';
-
 export function createResolvablePromise<T = any>() {
   let resolve: (value: T) => void, reject: (error: unknown) => void;
   const promise = new Promise<T>((res, rej) => {
@@ -13,56 +11,6 @@ export function createResolvablePromise<T = any>() {
   };
 }
 
-// Use the name `R` for `Row` as it will be shorter in the RSC payload.
-const R = [
-  (async ({
-    c, // current
-    n, // next
-  }: {
-    c: React.ReactNode;
-    n: Promise<any>;
-  }) => {
-    const chunk = await n;
-    if (chunk.done) {
-      return chunk.value;
-    }
-
-    if (chunk.append) {
-      return (
-        <>
-          {c}
-          <Suspense fallback={chunk.value}>
-            <R c={chunk.value} n={chunk.next} />
-          </Suspense>
-        </>
-      );
-    }
-
-    return (
-      <Suspense fallback={chunk.value}>
-        <R c={chunk.value} n={chunk.next} />
-      </Suspense>
-    );
-  }) as unknown as React.FC<{
-    c: React.ReactNode;
-    n: Promise<any>;
-  }>,
-][0];
-
-export function createSuspensedChunk(initialValue: React.ReactNode) {
-  const { promise, resolve, reject } = createResolvablePromise();
-
-  return {
-    row: (
-      <Suspense fallback={initialValue}>
-        <R c={initialValue} n={promise} />
-      </Suspense>
-    ) as React.ReactNode,
-    resolve,
-    reject,
-  };
-}
-
 export const isFunction = (x: unknown): x is Function =>
   typeof x === 'function';
 

packages/core/rsc/streamable.tsx

@@ -0,0 +1,81 @@
+'use server';
+
+import { streamUI } from 'ai/rsc';
+import { z } from 'zod';
+
+import { MockLanguageModelV1 } from '../../../../../core/test/mock-language-model-v1';
+import { convertArrayToReadableStream } from '../../../../../core/test/convert-array-to-readable-stream';
+
+const mockTextModel = new MockLanguageModelV1({
+  doStream: async () => {
+    return {
+      stream: convertArrayToReadableStream([
+        { type: 'text-delta', textDelta: `"Hello, ` },
+        { type: 'text-delta', textDelta: `world` },
+        { type: 'text-delta', textDelta: `!"` },
+      ]),
+      rawCall: { rawPrompt: 'prompt', rawSettings: {} },
+    };
+  },
+});
+
+const mockToolModel = new MockLanguageModelV1({
+  doStream: async () => {
+    return {
+      stream: convertArrayToReadableStream([
+        {
+          type: 'tool-call',
+          toolCallType: 'function',
+          toolCallId: 'call-1',
+          toolName: 'tool1',
+          args: `{ "value": "value" }`,
+        },
+      ]),
+      rawCall: { rawPrompt: 'prompt', rawSettings: {} },
+    };
+  },
+});
+
+function sleep(ms = 0) {
+  return new Promise(resolve => setTimeout(resolve, ms));
+}
+
+export async function action(testCase) {
+  switch (testCase) {
+    case 'text': {
+      const result = await streamUI({
+        model: mockTextModel,
+        prompt: '',
+      });
+      return result.value;
+    }
+    case 'wrapped-text': {
+      const result = await streamUI({
+        model: mockTextModel,
+        prompt: '',
+        text: ({ content }) => <p>AI: {content}</p>,
+      });
+      return result.value;
+    }
+    case 'tool': {
+      const result = await streamUI({
+        model: mockToolModel,
+        prompt: '',
+        tools: {
+          tool1: {
+            description: 'test tool 1',
+            parameters: z.object({
+              value: z.string(),
+            }),
+            generate: async function* ({ value }) {
+              yield 'Loading...';
+              await sleep(10);
+              return <div>tool1: {value}</div>;
+            },
+          },
+        },
+      });
+      return result.value;
+    }
+  }
+}

packages/core/rsc/streamable.ui.test.tsx

@@ -0,0 +1,43 @@
+'use client';
+
+import { useState } from 'react';
+
+export function Client({ action }) {
+  const [log, setLog] = useState('');
+
+  return (
+    <div>
+      <pre id="log" style={{ border: '1px solid #ccc', padding: 5 }}>
+        {log}
+      </pre>
+
+      {/* Test suites */}
+      <div style={{ display: 'inline-flex', flexDirection: 'column', gap: 5 }}>
+        <button
+          id="test-streamui-text"
+          onClick={async () => {
+            setLog(await action('text'));
+          }}
+        >
+          Test streamUI() Text UI
+        </button>
+        <button
+          id="test-streamui-wrapped-text"
+          onClick={async () => {
+            setLog(await action('wrapped-text'));
+          }}
+        >
+          Test streamUI() Wrapped Text UI
+        </button>
+        <button
+          id="test-streamui-tool"
+          onClick={async () => {
+            setLog(await action('tool'));
+          }}
+        >
+          Test streamUI() Tool
+        </button>
+      </div>
+    </div>
+  );
+}

packages/core/rsc/types.ts

@@ -0,0 +1,6 @@
+import { action } from './action';
+import { Client } from './client';
+
+export default function Page() {
+  return <Client action={action} />;
+}

packages/core/rsc/utils.tsx

@@ -7,22 +7,46 @@ function sleep(ms = 0) {
   return new Promise(resolve => setTimeout(resolve, ms));
 }
 
+export async function streamableUIError() {
+  const streamable = createStreamableUI();
+  (async () => {
+    await sleep(10);
+    streamable.update(<p>foo</p>);
+    await sleep(10);
+    streamable.error('This is an error');
+  })();
+  return streamable.value;
+}
+
+export async function streamableUIAppend() {
+  const streamable = createStreamableUI();
+  (async () => {
+    await sleep(10);
+    streamable.update(<p>foo</p>);
+    await sleep(10);
+    streamable.append(<p>bar</p>);
+    await sleep(10);
+    streamable.done();
+  })();
+  return streamable.value;
+}
+
 export async function streamableUI() {
   const streamable = createStreamableUI();
   (async () => {
-    await sleep();
+    await sleep(10);
     streamable.update(
       <ClientInfo>
         <p>I am a paragraph</p>
       </ClientInfo>,
     );
-    await sleep();
+    await sleep(10);
     streamable.update(
       <ClientInfo>
         <button>I am a button</button>
       </ClientInfo>,
     );
-    await sleep();
+    await sleep(10);
     streamable.done();
   })();
   return streamable.value;

packages/core/tests/e2e/next-server/app/rsc/client-utils.js

@@ -0,0 +1,21 @@
+'use client';
+
+import { useEffect, useState } from 'react';
+
+export function ClientInfo({ children }) {
+  const [renders, setRenders] = useState(0);
+
+  // Increment the render count whenever the children change.
+  // This can be used to verify that the component is re-rendered with the
+  // state kept.
+  useEffect(() => {
+    setRenders(r => r + 1);
+  }, [children]);
+
+  return (
+    <div>
+      <p>{renders > 1 ? '(Rerendered) ' : ''}</p>
+      {children}
+    </div>
+  );
+}

packages/core/tests/e2e/next-server/app/stream-ui/action.jsx

@@ -1,8 +1,26 @@
 'use client';
 
-import { useState } from 'react';
+import React, { useState } from 'react';
 import { readStreamableValue } from 'ai/rsc';
 
+class ErrorBoundary extends React.Component {
+  constructor(props) {
+    super(props);
+    this.state = { error: null };
+  }
+
+  static getDerivedStateFromError(error) {
+    return { error };
+  }
+
+  render() {
+    if (this.state.error) {
+      return <div>Caught by Error Boundary: {this.state.error}</div>;
+    }
+    return this.props.children;
+  }
+}
+
 export function Client({ actions }) {
   const [log, setLog] = useState('');
 
@@ -19,9 +37,20 @@ export function Client({ actions }) {
 
   // Test `createStreamableUI` API
   async function testStreamableUI() {
+    setLog(null);
     const value = await actions.streamableUI();
     setLog(value);
   }
+  async function testStreamableUIAppend() {
+    setLog(null);
+    const value = await actions.streamableUIAppend();
+    setLog(value);
+  }
+  async function testStreamableUIError() {
+    setLog(null);
+    const value = await actions.streamableUIError();
+    setLog(<ErrorBoundary>{value}</ErrorBoundary>);
+  }
 
   return (
     <div>
@@ -37,6 +66,12 @@ export function Client({ actions }) {
         <button id="test-streamable-ui" onClick={testStreamableUI}>
           Test Streamable UI
         </button>
+        <button id="test-streamable-ui-append" onClick={testStreamableUIAppend}>
+          Test Streamable UI (Append)
+        </button>
+        <button id="test-streamable-ui-error" onClick={testStreamableUIError}>
+          Test Streamable UI (Error)
+        </button>
       </div>
     </div>
   );

packages/core/tests/e2e/next-server/app/stream-ui/client.js

@@ -1,4 +1,9 @@
-import { streamableUI, streamableValue } from './actions';
+import {
+  streamableUI,
+  streamableValue,
+  streamableUIAppend,
+  streamableUIError,
+} from './actions';
 import { Client } from './client';
 
 export default function Page() {
@@ -7,6 +12,8 @@ export default function Page() {
       actions={{
         streamableUI,
         streamableValue,
+        streamableUIAppend,
+        streamableUIError,
       }}
     />
   );

packages/core/tests/e2e/next-server/app/stream-ui/page.js

@@ -8,6 +8,7 @@
     "next": "canary",
     "react": "rc",
     "react-dom": "rc",
+    "zod": "*",
     "ai": "workspace:*"
   }
 }

packages/core/tests/e2e/next-server/app/rsc/actions.jsx packages/core/tests/e2e/next-server/app/streamable/actions.jsx

@@ -0,0 +1,25 @@
+import { test, expect } from '@playwright/test';
+
+test('streamUI() with text', async ({ page }) => {
+  await page.goto('/stream-ui');
+  await page.click('#test-streamui-text');
+
+  const logs = page.locator('#log');
+  await expect(logs).toHaveText('"Hello, world!"');
+});
+
+test('streamUI() with wrapped text', async ({ page }) => {
+  await page.goto('/stream-ui');
+  await page.click('#test-streamui-wrapped-text');
+
+  const logs = page.locator('#log');
+  await expect(logs).toHaveText('AI: "Hello, world!"');
+});
+
+test('streamUI() with tool call', async ({ page }) => {
+  await page.goto('/stream-ui');
+  await page.click('#test-streamui-tool');
+
+  const logs = page.locator('#log');
+  await expect(logs).toHaveText('tool1: value');
+});

packages/core/tests/e2e/next-server/app/streamable/client-utils.js

@@ -1,7 +1,7 @@
 import { test, expect } from '@playwright/test';
 
-test('createStreamableValue and readStreamableValue', async ({ page }) => {
-  await page.goto('/rsc');
+test('createStreamableValue() and readStreamableValue()', async ({ page }) => {
+  await page.goto('/streamable');
   await page.click('#test-streamable-value');
 
   const logs = page.locator('#log');
@@ -10,10 +10,29 @@ test('createStreamableValue and readStreamableValue', async ({ page }) => {
   );
 });
 
-test('test-streamable-ui', async ({ page }) => {
-  await page.goto('/rsc');
+test('createStreamableUI()', async ({ page }) => {
+  await page.goto('/streamable');
   await page.click('#test-streamable-ui');
 
   const logs = page.locator('#log');
-  await expect(logs).toHaveText('I am a button');
+
+  // It should update the UI but reuse the same component instance and its state
+  // to avoid re-mounting.
+  await expect(logs).toHaveText('(Rerendered) I am a button');
+});
+
+test('createStreamableUI() .append() method', async ({ page }) => {
+  await page.goto('/streamable');
+  await page.click('#test-streamable-ui-append');
+
+  const logs = page.locator('#log');
+  await expect(logs).toHaveText('foobar');
+});
+
+test('createStreamableUI() .error() method', async ({ page }) => {
+  await page.goto('/streamable');
+  await page.click('#test-streamable-ui-error');
+
+  const logs = page.locator('#log');
+  await expect(logs).toHaveText('Caught by Error Boundary: This is an error');
 });