[Flight] Client and Server Reference Creation into Runtime (#27033)

We already did this for Server References on the Client so this brings
us parity with that. This gives us some more flexibility with changing
the runtime implementation without having to affect the loaders.

We can also do more in the runtime such as adding `.bind()` support to
Server References.

I also moved the CommonJS Proxy creation into the runtime helper from
the register so that it can be handled in one place.

This lets us remove the forks from Next.js since the loaders can be
simplified there to just use these helpers.

This PR doesn't change the protocol or shape of the objects. They're
still specific to each bundler but ideally we should probably move this
to shared helpers that can be used by multiple bundler implementations.

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

@@ -38,6 +38,11 @@ import {
 
 import {decodeAction} from 'react-server/src/ReactFlightActionServer';
 
+export {
+  registerServerReference,
+  registerClientReference,
+} from './ReactFlightESMReferences';
+
 function createDrainHandler(destination: Destination, request: Request) {
   return () => startFlowing(request, destination);
 }

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

@@ -178,18 +178,20 @@ function transformServerModule(
         continue;
     }
   }
-
+  if (localNames.size === 0) {
+    return source;
+  }
   let newSrc = source + '\n\n;';
+  newSrc +=
+    'import {registerServerReference} from "react-server-dom-esm/server";\n';
   localNames.forEach(function (exported, local) {
     if (localTypes.get(local) !== 'function') {
       // We first check if the export is a function and if so annotate it.
       newSrc += 'if (typeof ' + local + ' === "function") ';
     }
-    newSrc += 'Object.defineProperties(' + local + ',{';
-    newSrc += '$$typeof: {value: Symbol.for("react.server.reference")},';
-    newSrc += '$$id: {value: ' + JSON.stringify(url + '#' + exported) + '},';
-    newSrc += '$$bound: { value: null }';
-    newSrc += '});\n';
+    newSrc += 'registerServerReference(' + local + ',';
+    newSrc += JSON.stringify(url) + ',';
+    newSrc += JSON.stringify(exported) + ');\n';
   });
   return newSrc;
 }
@@ -313,13 +315,17 @@ async function transformClientModule(
 
   await parseExportNamesInto(body, names, url, loader);
 
+  if (names.length === 0) {
+    return '';
+  }
+
   let newSrc =
-    "const CLIENT_REFERENCE = Symbol.for('react.client.reference');\n";
+    'import {registerClientReference} from "react-server-dom-esm/server";\n';
   for (let i = 0; i < names.length; i++) {
     const name = names[i];
     if (name === 'default') {
       newSrc += 'export default ';
-      newSrc += 'Object.defineProperties(function() {';
+      newSrc += 'registerClientReference(function() {';
       newSrc +=
         'throw new Error(' +
         JSON.stringify(
@@ -331,7 +337,7 @@ async function transformClientModule(
         ');';
     } else {
       newSrc += 'export const ' + name + ' = ';
-      newSrc += 'Object.defineProperties(function() {';
+      newSrc += 'registerClientReference(function() {';
       newSrc +=
         'throw new Error(' +
         JSON.stringify(
@@ -341,10 +347,9 @@ async function transformClientModule(
         ) +
         ');';
     }
-    newSrc += '},{';
-    newSrc += '$$typeof: {value: CLIENT_REFERENCE},';
-    newSrc += '$$id: {value: ' + JSON.stringify(url + '#' + name) + '}';
-    newSrc += '});\n';
+    newSrc += '},';
+    newSrc += JSON.stringify(url) + ',';
+    newSrc += JSON.stringify(name) + ');\n';
   }
   return newSrc;
 }

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

@@ -0,0 +1,74 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @flow
+ */
+
+import type {ReactClientValue} from 'react-server/src/ReactFlightServer';
+
+export type ServerReference<T: Function> = T & {
+  $$typeof: symbol,
+  $$id: string,
+  $$bound: null | Array<ReactClientValue>,
+};
+
+// eslint-disable-next-line no-unused-vars
+export type ClientReference<T> = {
+  $$typeof: symbol,
+  $$id: string,
+};
+
+const CLIENT_REFERENCE_TAG = Symbol.for('react.client.reference');
+const SERVER_REFERENCE_TAG = Symbol.for('react.server.reference');
+
+export function isClientReference(reference: Object): boolean {
+  return reference.$$typeof === CLIENT_REFERENCE_TAG;
+}
+
+export function isServerReference(reference: Object): boolean {
+  return reference.$$typeof === SERVER_REFERENCE_TAG;
+}
+
+export function registerClientReference<T>(
+  proxyImplementation: any,
+  id: string,
+  exportName: string,
+): ClientReference<T> {
+  return Object.defineProperties(proxyImplementation, {
+    $$typeof: {value: CLIENT_REFERENCE_TAG},
+    $$id: {value: id + '#' + exportName},
+  });
+}
+
+// $FlowFixMe[method-unbinding]
+const FunctionBind = Function.prototype.bind;
+// $FlowFixMe[method-unbinding]
+const ArraySlice = Array.prototype.slice;
+function bind(this: ServerReference<any>) {
+  // $FlowFixMe[unsupported-syntax]
+  const newFn = FunctionBind.apply(this, arguments);
+  if (this.$$typeof === SERVER_REFERENCE_TAG) {
+    // $FlowFixMe[method-unbinding]
+    const args = ArraySlice.call(arguments, 1);
+    newFn.$$typeof = SERVER_REFERENCE_TAG;
+    newFn.$$id = this.$$id;
+    newFn.$$bound = this.$$bound ? this.$$bound.concat(args) : args;
+  }
+  return newFn;
+}
+
+export function registerServerReference<T>(
+  reference: ServerReference<T>,
+  id: string,
+  exportName: string,
+): ServerReference<T> {
+  return Object.defineProperties((reference: any), {
+    $$typeof: {value: SERVER_REFERENCE_TAG},
+    $$id: {value: id + '#' + exportName},
+    $$bound: {value: null},
+    bind: {value: bind},
+  });
+}

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

@@ -9,22 +9,17 @@
 
 import type {ReactClientValue} from 'react-server/src/ReactFlightServer';
 
+import type {
+  ClientReference,
+  ServerReference,
+} from './ReactFlightESMReferences';
+
+export type {ClientReference, ServerReference};
+
 export type ClientManifest = string; // base URL on the file system
 
-export type ServerReference<T: Function> = T & {
-  $$typeof: symbol,
-  $$id: string,
-  $$bound: null | Array<ReactClientValue>,
-};
-
 export type ServerReferenceId = string;
 
-// eslint-disable-next-line no-unused-vars
-export type ClientReference<T> = {
-  $$typeof: symbol,
-  $$id: string,
-};
-
 export type ClientReferenceMetadata = [
   string, // module path
   string, // export name
@@ -32,8 +27,7 @@ export type ClientReferenceMetadata = [
 
 export type ClientReferenceKey = string;
 
-const CLIENT_REFERENCE_TAG = Symbol.for('react.client.reference');
-const SERVER_REFERENCE_TAG = Symbol.for('react.server.reference');
+export {isClientReference, isServerReference} from './ReactFlightESMReferences';
 
 export function getClientReferenceKey(
   reference: ClientReference<any>,
@@ -41,14 +35,6 @@ export function getClientReferenceKey(
   return reference.$$id;
 }
 
-export function isClientReference(reference: Object): boolean {
-  return reference.$$typeof === CLIENT_REFERENCE_TAG;
-}
-
-export function isServerReference(reference: Object): boolean {
-  return reference.$$typeof === SERVER_REFERENCE_TAG;
-}
-
 export function resolveClientReferenceMetadata<T>(
   config: ClientManifest,
   clientReference: ClientReference<T>,

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

@@ -27,6 +27,12 @@ import {
 
 import {decodeAction} from 'react-server/src/ReactFlightActionServer';
 
+export {
+  registerServerReference,
+  registerClientReference,
+  createClientModuleProxy,
+} from './ReactFlightWebpackReferences';
+
 type Options = {
   identifierPrefix?: string,
   signal?: AbortSignal,

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

@@ -27,6 +27,12 @@ import {
 
 import {decodeAction} from 'react-server/src/ReactFlightActionServer';
 
+export {
+  registerServerReference,
+  registerClientReference,
+  createClientModuleProxy,
+} from './ReactFlightWebpackReferences';
+
 type Options = {
   identifierPrefix?: string,
   signal?: AbortSignal,

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

@@ -38,6 +38,12 @@ import {
 
 import {decodeAction} from 'react-server/src/ReactFlightActionServer';
 
+export {
+  registerServerReference,
+  registerClientReference,
+  createClientModuleProxy,
+} from './ReactFlightWebpackReferences';
+
 function createDrainHandler(destination: Destination, request: Request) {
   return () => startFlowing(request, destination);
 }

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

@@ -9,25 +9,19 @@
 
 import type {ReactClientValue} from 'react-server/src/ReactFlightServer';
 
+import type {
+  ClientReference,
+  ServerReference,
+} from './ReactFlightWebpackReferences';
+
+export type {ClientReference, ServerReference};
+
 export type ClientManifest = {
   [id: string]: ClientReferenceMetadata,
 };
 
-export type ServerReference<T: Function> = T & {
-  $$typeof: symbol,
-  $$id: string,
-  $$bound: null | Array<ReactClientValue>,
-};
-
 export type ServerReferenceId = string;
 
-// eslint-disable-next-line no-unused-vars
-export type ClientReference<T> = {
-  $$typeof: symbol,
-  $$id: string,
-  $$async: boolean,
-};
-
 export type ClientReferenceMetadata = {
   id: string,
   chunks: Array<string>,
@@ -37,8 +31,10 @@ export type ClientReferenceMetadata = {
 
 export type ClientReferenceKey = string;
 
-const CLIENT_REFERENCE_TAG = Symbol.for('react.client.reference');
-const SERVER_REFERENCE_TAG = Symbol.for('react.server.reference');
+export {
+  isClientReference,
+  isServerReference,
+} from './ReactFlightWebpackReferences';
 
 export function getClientReferenceKey(
   reference: ClientReference<any>,
@@ -46,14 +42,6 @@ export function getClientReferenceKey(
   return reference.$$async ? reference.$$id + '#async' : reference.$$id;
 }
 
-export function isClientReference(reference: Object): boolean {
-  return reference.$$typeof === CLIENT_REFERENCE_TAG;
-}
-
-export function isServerReference(reference: Object): boolean {
-  return reference.$$typeof === SERVER_REFERENCE_TAG;
-}
-
 export function resolveClientReferenceMetadata<T>(
   config: ClientManifest,
   clientReference: ClientReference<T>,

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

@@ -178,18 +178,20 @@ function transformServerModule(
         continue;
     }
   }
-
+  if (localNames.size === 0) {
+    return source;
+  }
   let newSrc = source + '\n\n;';
+  newSrc +=
+    'import {registerServerReference} from "react-server-dom-webpack/server";\n';
   localNames.forEach(function (exported, local) {
     if (localTypes.get(local) !== 'function') {
       // We first check if the export is a function and if so annotate it.
       newSrc += 'if (typeof ' + local + ' === "function") ';
     }
-    newSrc += 'Object.defineProperties(' + local + ',{';
-    newSrc += '$$typeof: {value: Symbol.for("react.server.reference")},';
-    newSrc += '$$id: {value: ' + JSON.stringify(url + '#' + exported) + '},';
-    newSrc += '$$bound: { value: null }';
-    newSrc += '});\n';
+    newSrc += 'registerServerReference(' + local + ',';
+    newSrc += JSON.stringify(url) + ',';
+    newSrc += JSON.stringify(exported) + ');\n';
   });
   return newSrc;
 }
@@ -313,13 +315,17 @@ async function transformClientModule(
 
   await parseExportNamesInto(body, names, url, loader);
 
+  if (names.length === 0) {
+    return '';
+  }
+
   let newSrc =
-    "const CLIENT_REFERENCE = Symbol.for('react.client.reference');\n";
+    'import {registerClientReference} from "react-server-dom-webpack/server";\n';
   for (let i = 0; i < names.length; i++) {
     const name = names[i];
     if (name === 'default') {
       newSrc += 'export default ';
-      newSrc += 'Object.defineProperties(function() {';
+      newSrc += 'registerClientReference(function() {';
       newSrc +=
         'throw new Error(' +
         JSON.stringify(
@@ -331,7 +337,7 @@ async function transformClientModule(
         ');';
     } else {
       newSrc += 'export const ' + name + ' = ';
-      newSrc += 'Object.defineProperties(function() {';
+      newSrc += 'registerClientReference(function() {';
       newSrc +=
         'throw new Error(' +
         JSON.stringify(
@@ -341,10 +347,9 @@ async function transformClientModule(
         ) +
         ');';
     }
-    newSrc += '},{';
-    newSrc += '$$typeof: {value: CLIENT_REFERENCE},';
-    newSrc += '$$id: {value: ' + JSON.stringify(url + '#' + name) + '}';
-    newSrc += '});\n';
+    newSrc += '},';
+    newSrc += JSON.stringify(url) + ',';
+    newSrc += JSON.stringify(name) + ');\n';
   }
   return newSrc;
 }

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

@@ -14,198 +14,9 @@ const url = require('url');
 const Module = require('module');
 
 module.exports = function register() {
-  const CLIENT_REFERENCE = Symbol.for('react.client.reference');
-  const SERVER_REFERENCE = Symbol.for('react.server.reference');
-  const PROMISE_PROTOTYPE = Promise.prototype;
-
-  // Patch bind on the server to ensure that this creates another
-  // bound server reference with the additional arguments.
-  const originalBind = Function.prototype.bind;
-  /*eslint-disable no-extend-native */
-  Function.prototype.bind = (function bind(this: any, self: any) {
-    // $FlowFixMe[unsupported-syntax]
-    const newFn = originalBind.apply(this, arguments);
-    if (this.$$typeof === SERVER_REFERENCE) {
-      // $FlowFixMe[method-unbinding]
-      const args = Array.prototype.slice.call(arguments, 1);
-      newFn.$$typeof = SERVER_REFERENCE;
-      newFn.$$id = this.$$id;
-      newFn.$$bound = this.$$bound ? this.$$bound.concat(args) : args;
-    }
-    return newFn;
-  }: any);
-
-  const deepProxyHandlers = {
-    get: function (target: Function, name: string, receiver: Proxy<Function>) {
-      switch (name) {
-        // These names are read by the Flight runtime if you end up using the exports object.
-        case '$$typeof':
-          // These names are a little too common. We should probably have a way to
-          // have the Flight runtime extract the inner target instead.
-          return target.$$typeof;
-        case '$$id':
-          return target.$$id;
-        case '$$async':
-          return target.$$async;
-        case 'name':
-          return target.name;
-        case 'displayName':
-          return undefined;
-        // We need to special case this because createElement reads it if we pass this
-        // reference.
-        case 'defaultProps':
-          return undefined;
-        // Avoid this attempting to be serialized.
-        case 'toJSON':
-          return undefined;
-        case Symbol.toPrimitive:
-          // $FlowFixMe[prop-missing]
-          return Object.prototype[Symbol.toPrimitive];
-        case 'Provider':
-          throw new Error(
-            `Cannot render a Client Context Provider on the Server. ` +
-              `Instead, you can export a Client Component wrapper ` +
-              `that itself renders a Client Context Provider.`,
-          );
-      }
-      // eslint-disable-next-line react-internal/safe-string-coercion
-      const expression = String(target.name) + '.' + String(name);
-      throw new Error(
-        `Cannot access ${expression} on the server. ` +
-          'You cannot dot into a client module from a server component. ' +
-          'You can only pass the imported name through.',
-      );
-    },
-    set: function () {
-      throw new Error('Cannot assign to a client module from a server module.');
-    },
-  };
-
-  const proxyHandlers = {
-    get: function (
-      target: Function,
-      name: string,
-      receiver: Proxy<Function>,
-    ): $FlowFixMe {
-      switch (name) {
-        // These names are read by the Flight runtime if you end up using the exports object.
-        case '$$typeof':
-          return target.$$typeof;
-        case '$$id':
-          return target.$$id;
-        case '$$async':
-          return target.$$async;
-        case 'name':
-          return target.name;
-        // We need to special case this because createElement reads it if we pass this
-        // reference.
-        case 'defaultProps':
-          return undefined;
-        // Avoid this attempting to be serialized.
-        case 'toJSON':
-          return undefined;
-        case Symbol.toPrimitive:
-          // $FlowFixMe[prop-missing]
-          return Object.prototype[Symbol.toPrimitive];
-        case '__esModule':
-          // Something is conditionally checking which export to use. We'll pretend to be
-          // an ESM compat module but then we'll check again on the client.
-          const moduleId = target.$$id;
-          target.default = Object.defineProperties(
-            (function () {
-              throw new Error(
-                `Attempted to call the default export of ${moduleId} from the server ` +
-                  `but it's on the client. It's not possible to invoke a client function from ` +
-                  `the server, it can only be rendered as a Component or passed to props of a ` +
-                  `Client Component.`,
-              );
-            }: any),
-            {
-              $$typeof: {value: CLIENT_REFERENCE},
-              // This a placeholder value that tells the client to conditionally use the
-              // whole object or just the default export.
-              $$id: {value: target.$$id + '#'},
-              $$async: {value: target.$$async},
-            },
-          );
-          return true;
-        case 'then':
-          if (target.then) {
-            // Use a cached value
-            return target.then;
-          }
-          if (!target.$$async) {
-            // If this module is expected to return a Promise (such as an AsyncModule) then
-            // we should resolve that with a client reference that unwraps the Promise on
-            // the client.
-
-            const clientReference = Object.defineProperties(({}: any), {
-              $$typeof: {value: CLIENT_REFERENCE},
-              $$id: {value: target.$$id},
-              $$async: {value: true},
-            });
-            const proxy = new Proxy(clientReference, proxyHandlers);
-
-            // Treat this as a resolved Promise for React's use()
-            target.status = 'fulfilled';
-            target.value = proxy;
-
-            const then = (target.then = Object.defineProperties(
-              (function then(resolve, reject: any) {
-                // Expose to React.
-                return Promise.resolve(resolve(proxy));
-              }: any),
-              // If this is not used as a Promise but is treated as a reference to a `.then`
-              // export then we should treat it as a reference to that name.
-              {
-                $$typeof: {value: CLIENT_REFERENCE},
-                $$id: {value: target.$$id + '#then'},
-                $$async: {value: false},
-              },
-            ));
-            return then;
-          } else {
-            // Since typeof .then === 'function' is a feature test we'd continue recursing
-            // indefinitely if we return a function. Instead, we return an object reference
-            // if we check further.
-            return undefined;
-          }
-      }
-      let cachedReference = target[name];
-      if (!cachedReference) {
-        const reference = Object.defineProperties(
-          (function () {
-            throw new Error(
-              // eslint-disable-next-line react-internal/safe-string-coercion
-              `Attempted to call ${String(name)}() from the server but ${String(
-                name,
-              )} is on the client. ` +
-                `It's not possible to invoke a client function from the server, it can ` +
-                `only be rendered as a Component or passed to props of a Client Component.`,
-            );
-          }: any),
-          {
-            name: {value: name},
-            $$typeof: {value: CLIENT_REFERENCE},
-            $$id: {value: target.$$id + '#' + name},
-            $$async: {value: target.$$async},
-          },
-        );
-        cachedReference = target[name] = new Proxy(
-          reference,
-          deepProxyHandlers,
-        );
-      }
-      return cachedReference;
-    },
-    getPrototypeOf(target: Function): Object {
-      // Pretend to be a Promise in case anyone asks.
-      return PROMISE_PROTOTYPE;
-    },
-    set: function (): empty {
-      throw new Error('Cannot assign to a client module from a server module.');
-    },
-  };
+  const Server: any = require('react-server-dom-webpack/server');
+  const registerServerReference = Server.registerServerReference;
+  const createClientModuleProxy = Server.createClientModuleProxy;
 
   // $FlowFixMe[prop-missing] found when upgrading Flow
   const originalCompile = Module.prototype._compile;
@@ -264,13 +75,7 @@ module.exports = function register() {
 
     if (useClient) {
       const moduleId: string = (url.pathToFileURL(filename).href: any);
-      const clientReference = Object.defineProperties(({}: any), {
-        $$typeof: {value: CLIENT_REFERENCE},
-        // Represents the whole Module object instead of a particular import.
-        $$id: {value: moduleId},
-        $$async: {value: false},
-      });
-      this.exports = new Proxy(clientReference, proxyHandlers);
+      this.exports = createClientModuleProxy(moduleId);
     }
 
     if (useServer) {
@@ -284,23 +89,19 @@ module.exports = function register() {
       // reference. If there are any functions in the export.
       if (typeof exports === 'function') {
         // The module exports a function directly,
-        Object.defineProperties((exports: any), {
-          $$typeof: {value: SERVER_REFERENCE},
+        registerServerReference(
+          (exports: any),
+          moduleId,
           // Represents the whole Module object instead of a particular import.
-          $$id: {value: moduleId},
-          $$bound: {value: null},
-        });
+          null,
+        );
       } else {
         const keys = Object.keys(exports);
         for (let i = 0; i < keys.length; i++) {
           const key = keys[i];
           const value = exports[keys[i]];
           if (typeof value === 'function') {
-            Object.defineProperties((value: any), {
-              $$typeof: {value: SERVER_REFERENCE},
-              $$id: {value: moduleId + '#' + key},
-              $$bound: {value: null},
-            });
+            registerServerReference((value: any), moduleId, key);
           }
         }
       }

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

@@ -0,0 +1,256 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @flow
+ */
+
+import type {ReactClientValue} from 'react-server/src/ReactFlightServer';
+
+export type ServerReference<T: Function> = T & {
+  $$typeof: symbol,
+  $$id: string,
+  $$bound: null | Array<ReactClientValue>,
+};
+
+// eslint-disable-next-line no-unused-vars
+export type ClientReference<T> = {
+  $$typeof: symbol,
+  $$id: string,
+  $$async: boolean,
+};
+
+const CLIENT_REFERENCE_TAG = Symbol.for('react.client.reference');
+const SERVER_REFERENCE_TAG = Symbol.for('react.server.reference');
+
+export function isClientReference(reference: Object): boolean {
+  return reference.$$typeof === CLIENT_REFERENCE_TAG;
+}
+
+export function isServerReference(reference: Object): boolean {
+  return reference.$$typeof === SERVER_REFERENCE_TAG;
+}
+
+export function registerClientReference<T>(
+  proxyImplementation: any,
+  id: string,
+  exportName: string,
+): ClientReference<T> {
+  return registerClientReferenceImpl(
+    proxyImplementation,
+    id + '#' + exportName,
+    false,
+  );
+}
+
+function registerClientReferenceImpl<T>(
+  proxyImplementation: any,
+  id: string,
+  async: boolean,
+): ClientReference<T> {
+  return Object.defineProperties(proxyImplementation, {
+    $$typeof: {value: CLIENT_REFERENCE_TAG},
+    $$id: {value: id},
+    $$async: {value: async},
+  });
+}
+
+// $FlowFixMe[method-unbinding]
+const FunctionBind = Function.prototype.bind;
+// $FlowFixMe[method-unbinding]
+const ArraySlice = Array.prototype.slice;
+function bind(this: ServerReference<any>) {
+  // $FlowFixMe[unsupported-syntax]
+  const newFn = FunctionBind.apply(this, arguments);
+  if (this.$$typeof === SERVER_REFERENCE_TAG) {
+    const args = ArraySlice.call(arguments, 1);
+    newFn.$$typeof = SERVER_REFERENCE_TAG;
+    newFn.$$id = this.$$id;
+    newFn.$$bound = this.$$bound ? this.$$bound.concat(args) : args;
+  }
+  return newFn;
+}
+
+export function registerServerReference<T>(
+  reference: ServerReference<T>,
+  id: string,
+  exportName: null | string,
+): ServerReference<T> {
+  return Object.defineProperties((reference: any), {
+    $$typeof: {value: SERVER_REFERENCE_TAG},
+    $$id: {value: exportName === null ? id : id + '#' + exportName},
+    $$bound: {value: null},
+    bind: {value: bind},
+  });
+}
+
+const PROMISE_PROTOTYPE = Promise.prototype;
+
+const deepProxyHandlers = {
+  get: function (target: Function, name: string, receiver: Proxy<Function>) {
+    switch (name) {
+      // These names are read by the Flight runtime if you end up using the exports object.
+      case '$$typeof':
+        // These names are a little too common. We should probably have a way to
+        // have the Flight runtime extract the inner target instead.
+        return target.$$typeof;
+      case '$$id':
+        return target.$$id;
+      case '$$async':
+        return target.$$async;
+      case 'name':
+        return target.name;
+      case 'displayName':
+        return undefined;
+      // We need to special case this because createElement reads it if we pass this
+      // reference.
+      case 'defaultProps':
+        return undefined;
+      // Avoid this attempting to be serialized.
+      case 'toJSON':
+        return undefined;
+      case Symbol.toPrimitive:
+        // $FlowFixMe[prop-missing]
+        return Object.prototype[Symbol.toPrimitive];
+      case 'Provider':
+        throw new Error(
+          `Cannot render a Client Context Provider on the Server. ` +
+            `Instead, you can export a Client Component wrapper ` +
+            `that itself renders a Client Context Provider.`,
+        );
+    }
+    // eslint-disable-next-line react-internal/safe-string-coercion
+    const expression = String(target.name) + '.' + String(name);
+    throw new Error(
+      `Cannot access ${expression} on the server. ` +
+        'You cannot dot into a client module from a server component. ' +
+        'You can only pass the imported name through.',
+    );
+  },
+  set: function () {
+    throw new Error('Cannot assign to a client module from a server module.');
+  },
+};
+
+const proxyHandlers = {
+  get: function (
+    target: Function,
+    name: string,
+    receiver: Proxy<Function>,
+  ): $FlowFixMe {
+    switch (name) {
+      // These names are read by the Flight runtime if you end up using the exports object.
+      case '$$typeof':
+        return target.$$typeof;
+      case '$$id':
+        return target.$$id;
+      case '$$async':
+        return target.$$async;
+      case 'name':
+        return target.name;
+      // We need to special case this because createElement reads it if we pass this
+      // reference.
+      case 'defaultProps':
+        return undefined;
+      // Avoid this attempting to be serialized.
+      case 'toJSON':
+        return undefined;
+      case Symbol.toPrimitive:
+        // $FlowFixMe[prop-missing]
+        return Object.prototype[Symbol.toPrimitive];
+      case '__esModule':
+        // Something is conditionally checking which export to use. We'll pretend to be
+        // an ESM compat module but then we'll check again on the client.
+        const moduleId = target.$$id;
+        target.default = registerClientReferenceImpl(
+          (function () {
+            throw new Error(
+              `Attempted to call the default export of ${moduleId} from the server ` +
+                `but it's on the client. It's not possible to invoke a client function from ` +
+                `the server, it can only be rendered as a Component or passed to props of a ` +
+                `Client Component.`,
+            );
+          }: any),
+          target.$$id + '#',
+          target.$$async,
+        );
+        return true;
+      case 'then':
+        if (target.then) {
+          // Use a cached value
+          return target.then;
+        }
+        if (!target.$$async) {
+          // If this module is expected to return a Promise (such as an AsyncModule) then
+          // we should resolve that with a client reference that unwraps the Promise on
+          // the client.
+
+          const clientReference: ClientReference<any> =
+            registerClientReferenceImpl(({}: any), target.$$id, true);
+          const proxy = new Proxy(clientReference, proxyHandlers);
+
+          // Treat this as a resolved Promise for React's use()
+          target.status = 'fulfilled';
+          target.value = proxy;
+
+          const then = (target.then = registerClientReferenceImpl(
+            (function then(resolve, reject: any) {
+              // Expose to React.
+              return Promise.resolve(resolve(proxy));
+            }: any),
+            // If this is not used as a Promise but is treated as a reference to a `.then`
+            // export then we should treat it as a reference to that name.
+            target.$$id + '#then',
+            false,
+          ));
+          return then;
+        } else {
+          // Since typeof .then === 'function' is a feature test we'd continue recursing
+          // indefinitely if we return a function. Instead, we return an object reference
+          // if we check further.
+          return undefined;
+        }
+    }
+    let cachedReference = target[name];
+    if (!cachedReference) {
+      const reference: ClientReference<any> = registerClientReferenceImpl(
+        (function () {
+          throw new Error(
+            // eslint-disable-next-line react-internal/safe-string-coercion
+            `Attempted to call ${String(name)}() from the server but ${String(
+              name,
+            )} is on the client. ` +
+              `It's not possible to invoke a client function from the server, it can ` +
+              `only be rendered as a Component or passed to props of a Client Component.`,
+          );
+        }: any),
+        target.$$id + '#' + name,
+        target.$$async,
+      );
+      Object.defineProperty((reference: any), 'name', {value: name});
+      cachedReference = target[name] = new Proxy(reference, deepProxyHandlers);
+    }
+    return cachedReference;
+  },
+  getPrototypeOf(target: Function): Object {
+    // Pretend to be a Promise in case anyone asks.
+    return PROMISE_PROTOTYPE;
+  },
+  set: function (): empty {
+    throw new Error('Cannot assign to a client module from a server module.');
+  },
+};
+
+export function createClientModuleProxy<T>(
+  moduleId: string,
+): ClientReference<T> {
+  const clientReference: ClientReference<T> = registerClientReferenceImpl(
+    ({}: any),
+    // Represents the whole Module object instead of a particular import.
+    moduleId,
+    false,
+  );
+  return new Proxy(clientReference, proxyHandlers);
+}

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

@@ -36,6 +36,14 @@ let ErrorBoundary;
 describe('ReactFlightDOM', () => {
   beforeEach(() => {
     jest.resetModules();
+
+    // Simulate the condition resolution
+    jest.mock('react-server-dom-webpack/server', () =>
+      require('react-server-dom-webpack/server.node.unbundled'),
+    );
+
+    ReactServerDOMClient = require('react-server-dom-webpack/client');
+
     act = require('internal-test-utils').act;
     const WebpackMock = require('./utils/WebpackMock');
     clientExports = WebpackMock.clientExports;
@@ -49,7 +57,6 @@ describe('ReactFlightDOM', () => {
     use = React.use;
     Suspense = React.Suspense;
     ReactDOMClient = require('react-dom/client');
-    ReactServerDOMClient = require('react-server-dom-webpack/client');
     ReactServerDOMServer = require('react-server-dom-webpack/server.node.unbundled');
 
     ErrorBoundary = class extends React.Component {

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

@@ -32,6 +32,12 @@ let use;
 describe('ReactFlightDOMBrowser', () => {
   beforeEach(() => {
     jest.resetModules();
+
+    // Simulate the condition resolution
+    jest.mock('react-server-dom-webpack/server', () =>
+      require('react-server-dom-webpack/server.browser'),
+    );
+
     act = require('internal-test-utils').act;
     const WebpackMock = require('./utils/WebpackMock');
     clientExports = WebpackMock.clientExports;

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

@@ -31,6 +31,12 @@ let use;
 describe('ReactFlightDOMEdge', () => {
   beforeEach(() => {
     jest.resetModules();
+
+    // Simulate the condition resolution
+    jest.mock('react-server-dom-webpack/server', () =>
+      require('react-server-dom-webpack/server.edge'),
+    );
+
     const WebpackMock = require('./utils/WebpackMock');
     clientExports = WebpackMock.clientExports;
     webpackMap = WebpackMock.webpackMap;

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

@@ -32,6 +32,10 @@ let ReactServerDOMClient;
 describe('ReactFlightDOMForm', () => {
   beforeEach(() => {
     jest.resetModules();
+    // Simulate the condition resolution
+    jest.mock('react-server-dom-webpack/server', () =>
+      require('react-server-dom-webpack/server.edge'),
+    );
     const WebpackMock = require('./utils/WebpackMock');
     serverExports = WebpackMock.serverExports;
     webpackServerMap = WebpackMock.webpackServerMap;

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

@@ -26,6 +26,12 @@ let use;
 describe('ReactFlightDOMNode', () => {
   beforeEach(() => {
     jest.resetModules();
+
+    // Simulate the condition resolution
+    jest.mock('react-server-dom-webpack/server', () =>
+      require('react-server-dom-webpack/server.node'),
+    );
+
     const WebpackMock = require('./utils/WebpackMock');
     clientExports = WebpackMock.clientExports;
     webpackMap = WebpackMock.webpackMap;

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

@@ -23,6 +23,10 @@ let ReactServerDOMClient;
 describe('ReactFlightDOMReply', () => {
   beforeEach(() => {
     jest.resetModules();
+    // Simulate the condition resolution
+    jest.mock('react-server-dom-webpack/server', () =>
+      require('react-server-dom-webpack/server.browser'),
+    );
     const WebpackMock = require('./utils/WebpackMock');
     // serverExports = WebpackMock.serverExports;
     webpackServerMap = WebpackMock.webpackServerMap;

scripts/rollup/bundles.js

@@ -447,7 +447,7 @@ const bundles = [
     global: 'ReactFlightWebpackNodeRegister',
     minifyWithProdErrorCodes: false,
     wrapWithModuleBoundaries: false,
-    externals: ['url', 'module'],
+    externals: ['url', 'module', 'react-server-dom-webpack/server'],
   },
 
   /******* React Server DOM ESM Server *******/

scripts/shared/inlinedHostConfigs.js

@@ -145,6 +145,8 @@ module.exports = [
       'react-server-dom-webpack/server',
       'react-server-dom-webpack/server.node',
       'react-server-dom-webpack/src/ReactFlightDOMServerNode.js', // react-server-dom-webpack/server.node
+      'react-server-dom-webpack/node-register',
+      'react-server-dom-webpack/src/ReactFlightWebpackNodeRegister.js',
       'react-devtools',
       'react-devtools-core',
       'react-devtools-shell',