[Flight] Fix encodeReply for JSX with temporary references

packages/react-client/src/ReactFlightClient.js

@@ -552,14 +552,19 @@ function moveDebugInfoFromChunkToInnerValue<T>(
         resolvedValue._debugInfo,
         debugInfo,
       );
-    } else {
+    } else if (!Object.isFrozen(resolvedValue)) {
       Object.defineProperty((resolvedValue: any), '_debugInfo', {
         configurable: false,
         enumerable: false,
         writable: true,
         value: debugInfo,
       });
     }
+    // TODO: If the resolved value is a frozen element (e.g. a client-created
+    // element from a temporary reference, or a JSX element exported as a client
+    // reference), server debug info is currently dropped because the element
+    // can't be mutated. We should probably clone the element so each rendering
+    // context gets its own mutable copy with the correct debug info.
   }
 }
 
@@ -2900,7 +2905,9 @@ function addAsyncInfo(chunk: SomeChunk<any>, asyncInfo: ReactAsyncInfo): void {
     if (isArray(value._debugInfo)) {
       // $FlowFixMe[method-unbinding]
       value._debugInfo.push(asyncInfo);
-    } else {
+    } else if (!Object.isFrozen(value)) {
+      // TODO: Debug info is dropped for frozen elements. See the TODO in
+      // moveDebugInfoFromChunkToInnerValue.
       Object.defineProperty((value: any), '_debugInfo', {
         configurable: false,
         enumerable: false,

packages/react-client/src/ReactFlightReplyClient.js

@@ -429,6 +429,14 @@ export function processReply(
               return serializeTemporaryReferenceMarker();
             }
           }
+          // This element is the root of a serializeModel call (e.g. JSX
+          // passed directly to encodeReply, or a promise that resolved to
+          // JSX). It was already registered as a temporary reference by
+          // serializeModel so we just need to emit the marker.
+          if (temporaryReferences !== undefined && modelRoot === value) {
+            modelRoot = null;
+            return serializeTemporaryReferenceMarker();
+          }
           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.' +

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

@@ -3941,4 +3941,61 @@ describe('ReactFlight', () => {
     const model = await ReactNoopFlightClient.read(transport);
     expect(model.element.key).toBe(React.optimisticKey);
   });
+
+  it('can use a JSX element exported as a client reference in multiple server components', async () => {
+    const ClientReference = clientReference(React.createElement('span'));
+
+    function Foo() {
+      return ClientReference;
+    }
+
+    function Bar() {
+      return ClientReference;
+    }
+
+    function App() {
+      return ReactServer.createElement(
+        'div',
+        null,
+        ReactServer.createElement(Foo),
+        ReactServer.createElement(Bar),
+      );
+    }
+
+    const transport = ReactNoopFlightServer.render(
+      ReactServer.createElement(App),
+    );
+
+    await act(async () => {
+      const result = await ReactNoopFlightClient.read(transport);
+      ReactNoop.render(result);
+
+      if (__DEV__) {
+        // TODO: Debug info is dropped for frozen elements (client-created JSX
+        // exported as a client reference in this case). Ideally we'd clone the
+        // element so that each context gets its own mutable copy with correct
+        // debug info. When fixed, foo should have Foo's debug info and bar should
+        // have Bar's debug info.
+        const [foo, bar] = result.props.children;
+        expect(getDebugInfo(foo)).toBe(null);
+        expect(getDebugInfo(bar)).toBe(null);
+      }
+    });
+
+    // TODO: With cloning, each context would get its own element copy, so this
+    // key warning should go away.
+    assertConsoleErrorDev([
+      'Each child in a list should have a unique "key" prop.\n\n' +
+        'Check the top-level render call using <div>. ' +
+        'See https://react.dev/link/warning-keys for more information.\n' +
+        '    in span (at **)',
+    ]);
+
+    expect(ReactNoop).toMatchRenderedOutput(
+      <div>
+        <span />
+        <span />
+      </div>,
+    );
+  });
 });

packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMReply-test.js

@@ -394,6 +394,74 @@ describe('ReactFlightDOMReply', () => {
     expect(response.children).toBe(children);
   });
 
+  it('can pass JSX as root model through a round trip using temporary references', async () => {
+    const jsx = <div />;
+
+    const temporaryReferences =
+      ReactServerDOMClient.createTemporaryReferenceSet();
+    const body = await ReactServerDOMClient.encodeReply(jsx, {
+      temporaryReferences,
+    });
+
+    const temporaryReferencesServer =
+      ReactServerDOMServer.createTemporaryReferenceSet();
+    const serverPayload = await ReactServerDOMServer.decodeReply(
+      body,
+      webpackServerMap,
+      {temporaryReferences: temporaryReferencesServer},
+    );
+    const stream = await serverAct(() =>
+      ReactServerDOMServer.renderToReadableStream(serverPayload, null, {
+        temporaryReferences: temporaryReferencesServer,
+      }),
+    );
+    const response = await ReactServerDOMClient.createFromReadableStream(
+      stream,
+      {
+        temporaryReferences,
+      },
+    );
+
+    // This should be the same reference that we already saw.
+    await expect(response).toBe(jsx);
+  });
+
+  it('can pass a promise that resolves to JSX through a round trip using temporary references', async () => {
+    const jsx = <div />;
+    const promise = Promise.resolve(jsx);
+
+    const temporaryReferences =
+      ReactServerDOMClient.createTemporaryReferenceSet();
+    const body = await ReactServerDOMClient.encodeReply(
+      {promise},
+      {
+        temporaryReferences,
+      },
+    );
+
+    const temporaryReferencesServer =
+      ReactServerDOMServer.createTemporaryReferenceSet();
+    const serverPayload = await ReactServerDOMServer.decodeReply(
+      body,
+      webpackServerMap,
+      {temporaryReferences: temporaryReferencesServer},
+    );
+    const stream = await serverAct(() =>
+      ReactServerDOMServer.renderToReadableStream(serverPayload, null, {
+        temporaryReferences: temporaryReferencesServer,
+      }),
+    );
+    const response = await ReactServerDOMClient.createFromReadableStream(
+      stream,
+      {
+        temporaryReferences,
+      },
+    );
+
+    // This should resolve to the same reference that we already saw.
+    await expect(response.promise).resolves.toBe(jsx);
+  });
+
   it('can return the same object using temporary references', async () => {
     const obj = {
       this: {is: 'a large object'},