From 1d26103cd3ba49e74b9780ab049657e6ed88862b Mon Sep 17 00:00:00 2001
From: Sebastian Markbage <sebastian@calyptus.eu>
Date: Wed, 20 Mar 2024 15:34:57 -0400
Subject: [PATCH] Reenable JSX elements in Replies

---
 .../src/ReactFlightReplyClient.js             | 17 +++++---
 .../src/__tests__/ReactFlightDOMReply-test.js | 42 +++++++++++--------
 2 files changed, 36 insertions(+), 23 deletions(-)

diff --git a/packages/react-client/src/ReactFlightReplyClient.js b/packages/react-client/src/ReactFlightReplyClient.js
index 944deb0727546..4427e144e1b95 100644
--- a/packages/react-client/src/ReactFlightReplyClient.js
+++ b/packages/react-client/src/ReactFlightReplyClient.js
@@ -227,11 +227,18 @@ export function processReply(
       switch ((value: any).$$typeof) {
         case REACT_ELEMENT_TYPE: {
           if (temporaryReferences === undefined) {
-            throw new Error(
-              'React Element cannot be passed to Server Functions from the Client without a ' +
-                'temporary reference set. Pass a TemporaryReferenceSet to the options.' +
-                (__DEV__ ? describeObjectForErrorMessage(parent, key) : ''),
-            );
+            const element: React$Element<any> = (value: any);
+            // Serialize as a plain object with a symbol property
+            // TODO: Consider if we should use a special encoding for this or restore a proper
+            // element object on the server. E.g. we probably need the _store stuff in case it
+            // is passed as a child. For now we assume it'll just be passed back to Flight.
+            return {
+              $$typeof: REACT_ELEMENT_TYPE,
+              type: element.type,
+              key: element.key,
+              ref: element.ref,
+              props: element.props,
+            };
           }
           return serializeTemporaryReferenceID(
             writeTemporaryReference(temporaryReferences, value),
diff --git a/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMReply-test.js b/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMReply-test.js
index 54d994444be1b..fef4f7e740bb8 100644
--- a/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMReply-test.js
+++ b/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMReply-test.js
@@ -302,24 +302,6 @@ describe('ReactFlightDOMReply', () => {
     expect(await result2.lazy.value).toBe('Hello');
   });
 
-  it('errors when called with JSX by default', async () => {
-    let error;
-    try {
-      await ReactServerDOMClient.encodeReply(<div />);
-    } catch (x) {
-      error = x;
-    }
-    expect(error).toEqual(
-      expect.objectContaining({
-        message: __DEV__
-          ? expect.stringContaining(
-              'React Element cannot be passed to Server Functions from the Client without a temporary reference set.',
-            )
-          : expect.stringContaining(''),
-      }),
-    );
-  });
-
   it('can pass JSX through a round trip using temporary references', async () => {
     function Component() {
       return <div />;
@@ -377,4 +359,28 @@ describe('ReactFlightDOMReply', () => {
 
     expect(response2.replied).toBe(Component);
   });
+
+  it('can pass a client JSX sent by the server back again', async () => {
+    function Component() {
+      return <div />;
+    }
+
+    const ClientComponent = clientExports(Component);
+
+    const stream1 = ReactServerDOMServer.renderToReadableStream(
+      <ClientComponent />,
+      webpackMap,
+    );
+    const response1 =
+      await ReactServerDOMClient.createFromReadableStream(stream1);
+
+    const body = await ReactServerDOMClient.encodeReply(response1);
+    const serverPayload = await ReactServerDOMServer.decodeReply(body);
+
+    const stream2 = ReactServerDOMServer.renderToReadableStream(serverPayload);
+    const response2 =
+      await ReactServerDOMClient.createFromReadableStream(stream2);
+
+    expect(response2.type).toBe(Component);
+  });
 });