[DO NOT LAND] Allow Previously Seen Client References to be Passed Back to a Reply

packages/react-client/src/ReactFlightClient.js

@@ -55,7 +55,10 @@ import {
   printToConsole,
 } from './ReactFlightClientConfig';
 
-import {registerServerReference} from './ReactFlightReplyClient';
+import {
+  registerServerReference,
+  registerClientReference,
+} from './ReactFlightReplyClient';
 
 import {readTemporaryReference} from './ReactFlightTemporaryReferences';
 
@@ -128,7 +131,7 @@ type ResolvedModelChunk<T> = {
 type ResolvedModuleChunk<T> = {
   status: 'resolved_module',
   value: ClientReference<T>,
-  reason: null,
+  reason: ClientReferenceMetadata,
   _response: Response,
   _debugInfo?: null | ReactDebugInfo,
   then(resolve: (T) => mixed, reject: (mixed) => mixed): void,
@@ -334,9 +337,10 @@ function createResolvedModelChunk<T>(
 function createResolvedModuleChunk<T>(
   response: Response,
   value: ClientReference<T>,
+  metadata: ClientReferenceMetadata,
 ): ResolvedModuleChunk<T> {
   // $FlowFixMe[invalid-constructor] Flow doesn't support functions as constructors
-  return new Chunk(RESOLVED_MODULE, value, null, response);
+  return new Chunk(RESOLVED_MODULE, value, metadata, response);
 }
 
 function createInitializedTextChunk(
@@ -381,6 +385,7 @@ function resolveModelChunk<T>(
 function resolveModuleChunk<T>(
   chunk: SomeChunk<T>,
   value: ClientReference<T>,
+  metadata: ClientReferenceMetadata,
 ): void {
   if (chunk.status !== PENDING && chunk.status !== BLOCKED) {
     // We already resolved. We didn't expect to see this.
@@ -391,6 +396,7 @@ function resolveModuleChunk<T>(
   const resolvedChunk: ResolvedModuleChunk<T> = (chunk: any);
   resolvedChunk.status = RESOLVED_MODULE;
   resolvedChunk.value = value;
+  resolvedChunk.reason = metadata;
   if (resolveListeners !== null) {
     initializeModuleChunk(resolvedChunk);
     wakeChunkIfInitialized(chunk, resolveListeners, rejectListeners);
@@ -450,9 +456,11 @@ function initializeModelChunk<T>(chunk: ResolvedModelChunk<T>): void {
 function initializeModuleChunk<T>(chunk: ResolvedModuleChunk<T>): void {
   try {
     const value: T = requireModule(chunk.value);
+    registerClientReference(value, chunk.reason);
     const initializedChunk: InitializedChunk<T> = (chunk: any);
     initializedChunk.status = INITIALIZED;
     initializedChunk.value = value;
+    initializedChunk.reason = null;
   } catch (error) {
     const erroredChunk: ErroredChunk<T> = (chunk: any);
     erroredChunk.status = ERRORED;
@@ -949,16 +957,28 @@ function resolveModule(
       blockedChunk.status = BLOCKED;
     }
     promise.then(
-      () => resolveModuleChunk(blockedChunk, clientReference),
+      () =>
+        resolveModuleChunk(
+          blockedChunk,
+          clientReference,
+          clientReferenceMetadata,
+        ),
       error => triggerErrorOnChunk(blockedChunk, error),
     );
   } else {
     if (!chunk) {
-      chunks.set(id, createResolvedModuleChunk(response, clientReference));
+      chunks.set(
+        id,
+        createResolvedModuleChunk(
+          response,
+          clientReference,
+          clientReferenceMetadata,
+        ),
+      );
     } else {
       // This can't actually happen because we don't have any forward
       // references to modules.
-      resolveModuleChunk(chunk, clientReference);
+      resolveModuleChunk(chunk, clientReference, clientReferenceMetadata);
     }
   }
 }

packages/react-client/src/ReactFlightReplyClient.js

@@ -16,6 +16,7 @@ import type {
 } from 'shared/ReactTypes';
 import type {LazyComponent} from 'react/src/ReactLazy';
 import type {TemporaryReferenceSet} from './ReactFlightTemporaryReferences';
+import type {ClientReferenceMetadata} from './ReactFlightClientConfig';
 
 import {enableRenderableContext} from 'shared/ReactFeatureFlags';
 
@@ -59,6 +60,9 @@ export type EncodeFormActionCallback = <A>(
   args: Promise<A>,
 ) => ReactCustomFormAction;
 
+const knownClientReferences: WeakMap<Function, ClientReferenceMetadata> =
+  new WeakMap();
+
 export type ServerReferenceId = any;
 
 const knownServerReferences: WeakMap<
@@ -97,6 +101,10 @@ function serializePromiseID(id: number): string {
   return '$@' + id.toString(16);
 }
 
+function serializeClientReferenceID(id: number): string {
+  return '$C' + id.toString(16);
+}
+
 function serializeServerReferenceID(id: number): string {
   return '$F' + id.toString(16);
 }
@@ -219,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),
@@ -454,9 +469,12 @@ export function processReply(
     }
 
     if (typeof value === 'function') {
-      const metaData = knownServerReferences.get(value);
-      if (metaData !== undefined) {
-        const metaDataJSON = JSON.stringify(metaData, resolveToJSON);
+      const serverMetaData = knownServerReferences.get(value);
+      if (serverMetaData !== undefined) {
+        const metaDataJSON: string = JSON.stringify(
+          serverMetaData,
+          resolveToJSON,
+        );
         if (formData === null) {
           // Upgrade to use FormData to allow us to stream this value.
           formData = new FormData();
@@ -468,6 +486,25 @@ export function processReply(
         return serializeServerReferenceID(refId);
       }
       if (temporaryReferences === undefined) {
+        const clientMetadata = knownClientReferences.get(value);
+        if (clientMetadata !== undefined) {
+          // If this function once was loaded by a server response from this same client
+          // we can pass it back given the same metadata that was sent to us. So that the
+          // server can return it to us again.
+          // $FlowFixMe[incompatible-type]: The meta data should be known to never yield undefined.
+          const metaDataJSON: string = JSON.stringify(
+            clientMetadata,
+            resolveToJSON,
+          );
+          if (formData === null) {
+            // Upgrade to use FormData to allow us to stream this value.
+            formData = new FormData();
+          }
+          const refId = nextPartId++;
+          // eslint-disable-next-line react-internal/safe-string-coercion
+          formData.set(formFieldPrefix + refId, metaDataJSON);
+          return serializeClientReferenceID(refId);
+        }
         throw new Error(
           'Client Functions cannot be passed directly to Server Functions. ' +
             'Only Functions passed from the Server can be passed back again.',
@@ -686,6 +723,18 @@ function isSignatureEqual(
   }
 }
 
+export function registerClientReference<T>(
+  reference: T,
+  metadata: ClientReferenceMetadata,
+): void {
+  if (typeof reference === 'function') {
+    // Currently we only track functions, like Components, because objects might be
+    // serializable and if they are the receiving side might expect to be able to
+    // use them. Functions on the other hand would error anywayy.
+    knownClientReferences.set(reference, metadata);
+  }
+}
+
 export function registerServerReference(
   proxy: any,
   reference: {id: ServerReferenceId, bound: null | Thenable<Array<any>>},

packages/react-client/src/forks/ReactFlightClientConfig.custom.js

@@ -29,7 +29,7 @@ export opaque type ModuleLoading = mixed;
 export opaque type SSRModuleMap = mixed;
 export opaque type ServerManifest = mixed;
 export opaque type ServerReferenceId = string;
-export opaque type ClientReferenceMetadata = mixed;
+export type ClientReferenceMetadata = mixed;
 export opaque type ClientReference<T> = mixed; // eslint-disable-line no-unused-vars
 export const resolveClientReference = $$$config.resolveClientReference;
 export const resolveServerReference = $$$config.resolveServerReference;

packages/react-client/src/forks/ReactFlightClientConfig.dom-bun.js

@@ -16,7 +16,7 @@ export opaque type ModuleLoading = mixed;
 export opaque type SSRModuleMap = mixed;
 export opaque type ServerManifest = mixed;
 export opaque type ServerReferenceId = string;
-export opaque type ClientReferenceMetadata = mixed;
+export type ClientReferenceMetadata = mixed;
 export opaque type ClientReference<T> = mixed; // eslint-disable-line no-unused-vars
 export const resolveClientReference: any = null;
 export const resolveServerReference: any = null;

packages/react-client/src/forks/ReactFlightClientConfig.dom-legacy.js

@@ -16,7 +16,7 @@ export opaque type ModuleLoading = mixed;
 export opaque type SSRModuleMap = mixed;
 export opaque type ServerManifest = mixed;
 export opaque type ServerReferenceId = string;
-export opaque type ClientReferenceMetadata = mixed;
+export type ClientReferenceMetadata = mixed;
 export opaque type ClientReference<T> = mixed; // eslint-disable-line no-unused-vars
 export const resolveClientReference: any = null;
 export const resolveServerReference: any = null;

packages/react-server-dom-esm/src/ReactFlightClientConfigBundlerESM.js

@@ -22,7 +22,7 @@ export type ServerReferenceId = string;
 
 import {prepareDestinationForModuleImpl} from 'react-client/src/ReactFlightClientConfig';
 
-export opaque type ClientReferenceMetadata = [
+export type ClientReferenceMetadata = [
   string, // module path
   string, // export name
 ];

packages/react-server-dom-turbopack/src/ReactFlightClientConfigBundlerNode.js

@@ -34,7 +34,7 @@ export type ServerManifest = void;
 
 export type ServerReferenceId = string;
 
-export opaque type ClientReferenceMetadata = ImportMetadata;
+export type ClientReferenceMetadata = ImportMetadata;
 
 // eslint-disable-next-line no-unused-vars
 export opaque type ClientReference<T> = {

packages/react-server-dom-turbopack/src/ReactFlightClientConfigBundlerTurbopack.js

@@ -43,10 +43,10 @@ export type ServerManifest = {
 export type ServerReferenceId = string;
 
 export opaque type ClientReferenceManifestEntry = ImportManifestEntry;
-export opaque type ClientReferenceMetadata = ImportMetadata;
+export type ClientReferenceMetadata = ImportMetadata;
 
 // eslint-disable-next-line no-unused-vars
-export opaque type ClientReference<T> = ClientReferenceMetadata;
+export type ClientReference<T> = ClientReferenceMetadata;
 
 // The reason this function needs to defined here in this file instead of just
 // being exported directly from the TurbopackDestination... file is because the

packages/react-server-dom-webpack/src/ReactFlightClientConfigBundlerNode.js

@@ -34,7 +34,7 @@ export type ServerManifest = void;
 
 export type ServerReferenceId = string;
 
-export opaque type ClientReferenceMetadata = ImportMetadata;
+export type ClientReferenceMetadata = ImportMetadata;
 
 // eslint-disable-next-line no-unused-vars
 export opaque type ClientReference<T> = {

packages/react-server-dom-webpack/src/ReactFlightClientConfigBundlerWebpack.js

@@ -43,7 +43,7 @@ export type ServerManifest = {
 export type ServerReferenceId = string;
 
 export opaque type ClientReferenceManifestEntry = ImportManifestEntry;
-export opaque type ClientReferenceMetadata = ImportMetadata;
+export type ClientReferenceMetadata = ImportMetadata;
 
 // eslint-disable-next-line no-unused-vars
 export opaque type ClientReference<T> = ClientReferenceMetadata;

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

@@ -15,6 +15,8 @@ global.ReadableStream =
 global.TextEncoder = require('util').TextEncoder;
 global.TextDecoder = require('util').TextDecoder;
 
+let clientExports;
+let webpackMap;
 // let serverExports;
 let webpackServerMap;
 let React;
@@ -30,6 +32,8 @@ describe('ReactFlightDOMReply', () => {
       require('react-server-dom-webpack/server.browser'),
     );
     const WebpackMock = require('./utils/WebpackMock');
+    clientExports = WebpackMock.clientExports;
+    webpackMap = WebpackMock.webpackMap;
     // serverExports = WebpackMock.serverExports;
     webpackServerMap = WebpackMock.webpackServerMap;
     React = require('react');
@@ -298,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 />;
@@ -346,4 +332,55 @@ describe('ReactFlightDOMReply', () => {
     // This should've been the same reference that we already saw.
     expect(response.children).toBe(children);
   });
+
+  it('can pass a client reference send by the server back again', async () => {
+    function Component() {
+      return <div />;
+    }
+
+    const ClientComponent = clientExports(Component);
+
+    const stream1 = ReactServerDOMServer.renderToReadableStream(
+      {component: ClientComponent},
+      webpackMap,
+    );
+    const response1 =
+      await ReactServerDOMClient.createFromReadableStream(stream1);
+    expect(response1.component).toBe(Component);
+
+    const body = await ReactServerDOMClient.encodeReply({
+      replied: response1.component,
+    });
+    const serverPayload = await ReactServerDOMServer.decodeReply(body);
+
+    const stream2 = ReactServerDOMServer.renderToReadableStream(serverPayload);
+    const response2 =
+      await ReactServerDOMClient.createFromReadableStream(stream2);
+
+    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);
+  });
 });