Turn AsyncGenerator into an AsyncIterable if it's an Async Generator ServerComponent

Author: sebmarkbage

packages/react-client/src/__tests__/ReactFlight-test.js

@@ -277,6 +277,29 @@ describe('ReactFlight', () => {
     expect(ReactNoop).toMatchRenderedOutput(<span>ABC</span>);
   });
 
+  it('can render a Generator Server Component as a fragment', async () => {
+    function ItemListClient(props) {
+      return <span>{props.children}</span>;
+    }
+    const ItemList = clientReference(ItemListClient);
+
+    function * Items() {
+      yield 'A';
+      yield 'B';
+      yield 'C';
+    }
+
+    const model =  <ItemList><Items /></ItemList>;
+
+    const transport = ReactNoopFlightServer.render(model);
+
+    await act(async () => {
+      ReactNoop.render(await ReactNoopFlightClient.read(transport));
+    });
+
+    expect(ReactNoop).toMatchRenderedOutput(<span>ABC</span>);
+  });
+
   it('can render undefined', async () => {
     function Undefined() {
       return undefined;
@@ -2133,16 +2156,9 @@ describe('ReactFlight', () => {
     }
     const Stateful = clientReference(StatefulClient);
 
-    function ServerComponent({item, initial}) {
-      // While the ServerComponent itself could be an async generator, single-shot iterables
-      // are not supported as React children since React might need to re-map them based on
-      // state updates. So we create an AsyncIterable instead.
-      return {
-        async *[Symbol.asyncIterator]() {
-          yield <Stateful key="a" initial={'a' + initial} />;
-          yield <Stateful key="b" initial={'b' + initial} />;
-        },
-      };
+    async function* ServerComponent({item, initial}) {
+      yield <Stateful key="a" initial={'a' + initial} />;
+      yield <Stateful key="b" initial={'b' + initial} />;
     }
 
     function ListClient({children}) {
@@ -2154,6 +2170,11 @@ describe('ReactFlight', () => {
         expect(fragment.type).toBe(React.Fragment);
         const fragmentChildren = [];
         const iterator = fragment.props.children[Symbol.asyncIterator]();
+        if (iterator === fragment.props.children) {
+          console.error(
+            'AyncIterators are not valid children of React. It must be a multi-shot AsyncIterable.',
+          );
+        }
         for (let entry; !(entry = React.use(iterator.next())).done; ) {
           fragmentChildren.push(entry.value);
         }
@@ -2298,23 +2319,21 @@ describe('ReactFlight', () => {
     let resolve;
     const iteratorPromise = new Promise(r => (resolve = r));
 
-    function ThirdPartyAsyncIterableComponent({item, initial}) {
-      // While the ServerComponent itself could be an async generator, single-shot iterables
-      // are not supported as React children since React might need to re-map them based on
-      // state updates. So we create an AsyncIterable instead.
-      return {
-        async *[Symbol.asyncIterator]() {
-          yield <span>Who</span>;
-          yield <span>dis?</span>;
-          resolve();
-        },
-      };
+    async function* ThirdPartyAsyncIterableComponent({item, initial}) {
+      yield <span>Who</span>;
+      yield <span>dis?</span>;
+      resolve();
     }
 
     function ListClient({children: fragment}) {
       // TODO: Unwrap AsyncIterables natively in React. For now we do it in this wrapper.
       const resolvedChildren = [];
       const iterator = fragment.props.children[Symbol.asyncIterator]();
+      if (iterator === fragment.props.children) {
+        console.error(
+          'AyncIterators are not valid children of React. It must be a multi-shot AsyncIterable.',
+        );
+      }
       for (let entry; !(entry = React.use(iterator.next())).done; ) {
         resolvedChildren.push(entry.value);
       }

packages/react-server/src/ReactFlightServer.js

@@ -864,20 +864,55 @@ function renderFunctionComponent<Props>(
   } else {
     result = Component(props, secondArg);
   }
-  if (
-    typeof result === 'object' &&
-    result !== null &&
-    typeof result.then === 'function'
-  ) {
-    // When the return value is in children position we can resolve it immediately,
-    // to its value without a wrapper if it's synchronously available.
-    const thenable: Thenable<any> = result;
-    if (thenable.status === 'fulfilled') {
-      return thenable.value;
+  if (typeof result === 'object' && result !== null) {
+    if (typeof result.then === 'function') {
+      // When the return value is in children position we can resolve it immediately,
+      // to its value without a wrapper if it's synchronously available.
+      const thenable: Thenable<any> = result;
+      if (thenable.status === 'fulfilled') {
+        return thenable.value;
+      }
+      // TODO: Once we accept Promises as children on the client, we can just return
+      // the thenable here.
+      result = createLazyWrapperAroundWakeable(result);
+    }
+    if (
+      enableFlightReadableStream &&
+      typeof (result: any)[ASYNC_ITERATOR] === 'function' &&
+      (typeof ReadableStream !== 'function' ||
+        !(result instanceof ReadableStream))
+    ) {
+      // Normally we'd serialize an AsyncIterator as a single-shot which is not compatible
+      // to be rendered as a React Child. However, because we have the function to recreate
+      // an iterable from rendering the element again, we can effectively treat it as multi-
+      // shot. Therefore we treat this as an AsyncIterable, whether it was one or not, by
+      // adding a wrapper so that this component effectively renders down to an AsyncIterable.
+      const iterableChild = result;
+      result = {
+        [ASYNC_ITERATOR]: function () {
+          const iterator = (iterableChild: any)[ASYNC_ITERATOR]();
+          if (__DEV__) {
+            // If this was an AsyncIterator but not an AsyncGeneratorFunction we warn because
+            // it might have been a mistake. Technically you can make this mistake with
+            // AsyncGeneratorFunctions and even single-shot AsyncIterables too but it's extra
+            // tempting to try to return the value from a generator.
+            if (
+              iterator === iterableChild &&
+              // $FlowIgnore[method-unbinding]
+              Object.prototype.toString.call(iterableChild) !==
+                '[object AsyncGenerator]'
+            ) {
+              console.error(
+                'Returning an AsyncIterator from a Server Component is not supported ' +
+                  'since it cannot be looped over more than once. ',
+              );
+            }
+          }
+          return iterator;
+        },
+        _debugInfo: iterableChild._debugInfo,
+      };
     }
-    // TODO: Once we accept Promises as children on the client, we can just return
-    // the thenable here.
-    result = createLazyWrapperAroundWakeable(result);
   }
   // Track this element's key on the Server Component on the keyPath context..
   const prevKeyPath = task.keyPath;
@@ -1911,6 +1946,8 @@ function renderModelDestructive(
 
     const iteratorFn = getIteratorFn(value);
     if (iteratorFn) {
+      // TODO: Distinguish Iterator and Iterable since client-side React should
+      // warn if you try to use an Iterator.
       return renderFragment(request, task, Array.from((value: any)));
     }