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

Co-authored-by: Lars Grammel <[email protected]>

.changeset/early-hounds-flow.md

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

packages/core/rsc/rsc-shared.mts

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

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

@@ -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/index.ts

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

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

@@ -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;