[Flight] Fix encodeReply for JSX with temporary references (#35730)

`encodeReply` throws "React Element cannot be passed to Server Functions
from the Client without a temporary reference set" when a React element
is the root value of a `serializeModel` call (either passed directly or
resolved from a promise), even when a temporary reference set is
provided.

The cause is that `resolveToJSON` hits the `REACT_ELEMENT_TYPE` switch
case before reaching the `existingReference`/`modelRoot` check that
regular objects benefit from. The synthetic JSON root created by
`JSON.stringify` is never tracked in `writtenObjects`, so
`parentReference` is `undefined` and the code falls through to the
throw. This adds a `modelRoot` check in the `REACT_ELEMENT_TYPE` case,
following the same pattern used for promises and plain objects.

The added `JSX as root model` test also uncovered a pre-existing crash
in the Flight Client: when the JSX element round-trips back, it arrives
as a frozen object (client-created elements are frozen in DEV), and
`Object.defineProperty` for `_debugInfo` fails because frozen objects
are non-configurable. The same crash can occur with JSX exported as a
client reference. For now, we're adding `!Object.isFrozen()` guards in
`moveDebugInfoFromChunkToInnerValue` and `addAsyncInfo` to prevent the
crash, which means debug info is silently dropped for frozen elements.
The proper fix would likely be to clone the element so each rendering
context gets its own mutable copy with correct debug info.

closes #34984
closes #35690

packages/react-client/src/ReactFlightClient.js

@@ -552,7 +552,7 @@ function moveDebugInfoFromChunkToInnerValue<T>(
         resolvedValue._debugInfo,
         debugInfo,
       );
-    } else {
+    } else if (!Object.isFrozen(resolvedValue)) {
       Object.defineProperty((resolvedValue: any), '_debugInfo', {
         configurable: false,
         enumerable: false,
@@ -560,6 +560,11 @@ function moveDebugInfoFromChunkToInnerValue<T>(
         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'},