[react-ssr-prepass] Fix not resolving client components

Change-Id: I127288025c76a12685ead339200396c0a14acb93

Author: victoragnez

package.json

@@ -1,6 +1,6 @@
 {
   "name": "@plasmicapp/react-ssr-prepass",
-  "version": "2.0.5",
+  "version": "2.0.6",
   "description": "A custom partial React SSR renderer for prefetching and suspense",
   "main": "dist/react-ssr-prepass.js",
   "module": "dist/react-ssr-prepass.es.js",

src/__tests__/visitor.test.js

@@ -29,6 +29,8 @@ import {
 
 import { visitElement } from '../visitor'
 
+import { REACT_ELEMENT_TYPE, REACT_LAZY_TYPE } from '../symbols'
+
 const { ReactCurrentDispatcher } = (React: any)
   .__SECRET_INTERNALS_DO_NOT_USE_OR_YOU_WILL_BE_FIRED
 
@@ -159,6 +161,35 @@ describe('visitElement', () => {
     })
   })
 
+  it('resolves lazy components of resolved modules', () => {
+    const ClientComponent = jest.fn().mockReturnValue(null)
+    const queue = []
+    const children = visitElement(
+      {
+        $$typeof: REACT_ELEMENT_TYPE,
+        props: {},
+        key: null,
+        ref: null,
+        type: {
+          $$typeof: REACT_LAZY_TYPE,
+          _payload: {
+            reason: null,
+            status: 'resolved_module',
+            _response: {},
+            _testComponent: ClientComponent
+          },
+          _init: (payload) => payload._testComponent
+        }
+      },
+      queue,
+      () => {}
+    )
+
+    expect(children.length).toBe(1)
+    expect(queue.length).toBe(0)
+    expect(children[0].type).toBe(ClientComponent)
+  })
+
   it('walks over forwardRef components', () => {
     const Test = React.forwardRef(Noop)
     const children = visitElement(<Test />, [], () => {})

src/render/lazyComponent.js

@@ -55,6 +55,20 @@ const resolve = (type: LazyComponent): Promise<void> => {
     }))
 }
 
+const makeFrame = (
+  type: LazyComponent,
+  props: DefaultProps,
+  thenable: Promise<any>
+) => ({
+  kind: 'frame.lazy',
+  contextMap: getCurrentContextMap(),
+  contextStore: getCurrentContextStore(),
+  errorFrame: getCurrentErrorFrame(),
+  thenable,
+  props,
+  type
+})
+
 const render = (
   type: LazyComponent,
   props: DefaultProps,
@@ -63,11 +77,18 @@ const render = (
   // Component has previously been fetched successfully,
   // so create the element with passed props and return it
   const payload = ((type._payload || type: any): LazyComponentPayload)
-  if (payload._status === 1) {
+  if (payload._status === 1 && payload._result) {
     return createElement(payload._result, props)
   }
 
-  return null
+  try {
+    return createElement((type: any)._init((type: any)._payload), props)
+  } catch (err) {
+    if (!!err && typeof err.then === 'function') {
+      queue.push(makeFrame(type, props, err))
+    }
+    return null
+  }
 }
 
 export const mount = (
@@ -77,16 +98,12 @@ export const mount = (
 ): Node => {
   // If the component has not been fetched yet, suspend this component
   const payload = ((type._payload || type: any): LazyComponentPayload)
-  if (payload._status <= 0) {
-    queue.push({
-      kind: 'frame.lazy',
-      contextMap: getCurrentContextMap(),
-      contextStore: getCurrentContextStore(),
-      errorFrame: getCurrentErrorFrame(),
-      thenable: resolve(type),
-      props,
-      type
-    })
+  // For lazy components from resolved modules (e.g. Client Components)
+  // The `payload._status` is null.
+  // Reference code for "normal" lazy components: https://github.com/facebook/react/blob/main/packages/react/src/ReactLazy.js#L53
+  // Reference code for Client Component chunks: https://github.com/facebook/react/blob/main/packages/react-client/src/ReactFlightClient.js#L955
+  if (payload._status != null && payload._status <= 0) {
+    queue.push(makeFrame(type, props, resolve(type)))
 
     return null
   }