[Flight] Add serverModuleMap option for mapping ServerReferences (#31300)

Stacked on #31299.

We already have an option for resolving Client References to other
Client References when consuming an RSC payload on the server.

This lets you resolve Server References on the consuming side when the
environment where you're consuming the RSC payload also has access to
those Server References. Basically they becomes like Client References
for this consumer but for another consumer they wouldn't be.

packages/react-client/src/ReactFlightClient.js

@@ -21,6 +21,7 @@ import type {
   ClientReference,
   ClientReferenceMetadata,
   ServerConsumerModuleMap,
+  ServerManifest,
   StringDecoder,
   ModuleLoading,
 } from './ReactFlightClientConfig';
@@ -51,6 +52,7 @@ import {
 
 import {
   resolveClientReference,
+  resolveServerReference,
   preloadModule,
   requireModule,
   dispatchHint,
@@ -270,6 +272,7 @@ export type FindSourceMapURLCallback = (
 
 export type Response = {
   _bundlerConfig: ServerConsumerModuleMap,
+  _serverReferenceConfig: null | ServerManifest,
   _moduleLoading: ModuleLoading,
   _callServer: CallServerCallback,
   _encodeFormAction: void | EncodeFormActionCallback,
@@ -896,7 +899,7 @@ function waitForReference<T>(
   parentObject: Object,
   key: string,
   response: Response,
-  map: (response: Response, model: any) => T,
+  map: (response: Response, model: any, parentObject: Object, key: string) => T,
   path: Array<string>,
 ): T {
   let handler: InitializationHandler;
@@ -938,7 +941,7 @@ function waitForReference<T>(
       }
       value = value[path[i]];
     }
-    const mappedValue = map(response, value);
+    const mappedValue = map(response, value, parentObject, key);
     parentObject[key] = mappedValue;
 
     // If this is the root object for a model reference, where `handler.value`
@@ -1041,7 +1044,7 @@ function waitForReference<T>(
   return (null: any);
 }
 
-function createServerReferenceProxy<A: Iterable<any>, T>(
+function loadServerReference<A: Iterable<any>, T>(
   response: Response,
   metaData: {
     id: any,
@@ -1050,13 +1053,147 @@ function createServerReferenceProxy<A: Iterable<any>, T>(
     env?: string, // DEV-only
     location?: ReactCallSite, // DEV-only
   },
+  parentObject: Object,
+  key: string,
 ): (...A) => Promise<T> {
-  return createBoundServerReference(
-    metaData,
-    response._callServer,
-    response._encodeFormAction,
-    __DEV__ ? response._debugFindSourceMapURL : undefined,
-  );
+  if (!response._serverReferenceConfig) {
+    // In the normal case, we can't load this Server Reference in the current environment and
+    // we just return a proxy to it.
+    return createBoundServerReference(
+      metaData,
+      response._callServer,
+      response._encodeFormAction,
+      __DEV__ ? response._debugFindSourceMapURL : undefined,
+    );
+  }
+  // If we have a module mapping we can load the real version of this Server Reference.
+  const serverReference: ClientReference<T> =
+    resolveServerReference<$FlowFixMe>(
+      response._serverReferenceConfig,
+      metaData.id,
+    );
+
+  const promise = preloadModule(serverReference);
+  if (!promise) {
+    return (requireModule(serverReference): any);
+  }
+
+  let handler: InitializationHandler;
+  if (initializingHandler) {
+    handler = initializingHandler;
+    handler.deps++;
+  } else {
+    handler = initializingHandler = {
+      parent: null,
+      chunk: null,
+      value: null,
+      deps: 1,
+      errored: false,
+    };
+  }
+
+  function fulfill(): void {
+    const resolvedValue = (requireModule(serverReference): any);
+    parentObject[key] = resolvedValue;
+
+    // If this is the root object for a model reference, where `handler.value`
+    // is a stale `null`, the resolved value can be used directly.
+    if (key === '' && handler.value === null) {
+      handler.value = resolvedValue;
+    }
+
+    // If the parent object is an unparsed React element tuple, we also need to
+    // update the props and owner of the parsed element object (i.e.
+    // handler.value).
+    if (
+      parentObject[0] === REACT_ELEMENT_TYPE &&
+      typeof handler.value === 'object' &&
+      handler.value !== null &&
+      handler.value.$$typeof === REACT_ELEMENT_TYPE
+    ) {
+      const element: any = handler.value;
+      switch (key) {
+        case '3':
+          element.props = resolvedValue;
+          break;
+        case '4':
+          if (__DEV__) {
+            element._owner = resolvedValue;
+          }
+          break;
+      }
+    }
+
+    handler.deps--;
+
+    if (handler.deps === 0) {
+      const chunk = handler.chunk;
+      if (chunk === null || chunk.status !== BLOCKED) {
+        return;
+      }
+      const resolveListeners = chunk.value;
+      const initializedChunk: InitializedChunk<T> = (chunk: any);
+      initializedChunk.status = INITIALIZED;
+      initializedChunk.value = handler.value;
+      if (resolveListeners !== null) {
+        wakeChunk(resolveListeners, handler.value);
+      }
+    }
+  }
+
+  function reject(error: mixed): void {
+    if (handler.errored) {
+      // We've already errored. We could instead build up an AggregateError
+      // but if there are multiple errors we just take the first one like
+      // Promise.all.
+      return;
+    }
+    const blockedValue = handler.value;
+    handler.errored = true;
+    handler.value = error;
+    const chunk = handler.chunk;
+    if (chunk === null || chunk.status !== BLOCKED) {
+      return;
+    }
+
+    if (__DEV__) {
+      if (
+        typeof blockedValue === 'object' &&
+        blockedValue !== null &&
+        blockedValue.$$typeof === REACT_ELEMENT_TYPE
+      ) {
+        const element = blockedValue;
+        // Conceptually the error happened inside this Element but right before
+        // it was rendered. We don't have a client side component to render but
+        // we can add some DebugInfo to explain that this was conceptually a
+        // Server side error that errored inside this element. That way any stack
+        // traces will point to the nearest JSX that errored - e.g. during
+        // serialization.
+        const erroredComponent: ReactComponentInfo = {
+          name: getComponentNameFromType(element.type) || '',
+          owner: element._owner,
+        };
+        if (enableOwnerStacks) {
+          // $FlowFixMe[cannot-write]
+          erroredComponent.debugStack = element._debugStack;
+          if (supportsCreateTask) {
+            // $FlowFixMe[cannot-write]
+            erroredComponent.debugTask = element._debugTask;
+          }
+        }
+        const chunkDebugInfo: ReactDebugInfo =
+          chunk._debugInfo || (chunk._debugInfo = []);
+        chunkDebugInfo.push(erroredComponent);
+      }
+    }
+
+    triggerErrorOnChunk(chunk, error);
+  }
+
+  promise.then(fulfill, reject);
+
+  // Return a place holder value for now.
+  return (null: any);
 }
 
 function getOutlinedModel<T>(
@@ -1064,7 +1201,7 @@ function getOutlinedModel<T>(
   reference: string,
   parentObject: Object,
   key: string,
-  map: (response: Response, model: any) => T,
+  map: (response: Response, model: any, parentObject: Object, key: string) => T,
 ): T {
   const path = reference.split(':');
   const id = parseInt(path[0], 16);
@@ -1099,7 +1236,7 @@ function getOutlinedModel<T>(
         }
         value = value[path[i]];
       }
-      const chunkValue = map(response, value);
+      const chunkValue = map(response, value, parentObject, key);
       if (__DEV__ && chunk._debugInfo) {
         // If we have a direct reference to an object that was rendered by a synchronous
         // server component, it might have some debug info about how it was rendered.
@@ -1244,7 +1381,7 @@ function parseModelString(
           ref,
           parentObject,
           key,
-          createServerReferenceProxy,
+          loadServerReference,
         );
       }
       case 'T': {
@@ -1421,6 +1558,7 @@ function missingCall() {
 function ResponseInstance(
   this: $FlowFixMe,
   bundlerConfig: ServerConsumerModuleMap,
+  serverReferenceConfig: null | ServerManifest,
   moduleLoading: ModuleLoading,
   callServer: void | CallServerCallback,
   encodeFormAction: void | EncodeFormActionCallback,
@@ -1432,6 +1570,7 @@ function ResponseInstance(
 ) {
   const chunks: Map<number, SomeChunk<any>> = new Map();
   this._bundlerConfig = bundlerConfig;
+  this._serverReferenceConfig = serverReferenceConfig;
   this._moduleLoading = moduleLoading;
   this._callServer = callServer !== undefined ? callServer : missingCall;
   this._encodeFormAction = encodeFormAction;
@@ -1486,6 +1625,7 @@ function ResponseInstance(
 
 export function createResponse(
   bundlerConfig: ServerConsumerModuleMap,
+  serverReferenceConfig: null | ServerManifest,
   moduleLoading: ModuleLoading,
   callServer: void | CallServerCallback,
   encodeFormAction: void | EncodeFormActionCallback,
@@ -1498,6 +1638,7 @@ export function createResponse(
   // $FlowFixMe[invalid-constructor]: the shapes are exact here but Flow doesn't like constructors
   return new ResponseInstance(
     bundlerConfig,
+    serverReferenceConfig,
     moduleLoading,
     callServer,
     encodeFormAction,

packages/react-markup/src/ReactMarkupServer.js

@@ -173,6 +173,7 @@ export function experimental_renderToHTML(
       undefined,
     );
     const flightResponse = createFlightResponse(
+      null,
       null,
       null,
       noServerCallOrFormAction,

packages/react-noop-renderer/src/ReactNoopFlightClient.js

@@ -62,6 +62,7 @@ function read<T>(source: Source, options: ReadOptions): Thenable<T> {
   const response = createResponse(
     source,
     null,
+    null,
     undefined,
     undefined,
     undefined,

packages/react-server-dom-esm/src/client/ReactFlightDOMClientBrowser.js

@@ -51,6 +51,7 @@ function createResponseFromOptions(options: void | Options) {
   return createResponse(
     options && options.moduleBaseURL ? options.moduleBaseURL : '',
     null,
+    null,
     options && options.callServer ? options.callServer : undefined,
     undefined, // encodeFormAction
     undefined, // nonce

packages/react-server-dom-esm/src/client/ReactFlightDOMClientNode.js

@@ -62,6 +62,7 @@ function createFromNodeStream<T>(
 ): Thenable<T> {
   const response: Response = createResponse(
     moduleRootPath,
+    null,
     moduleBaseURL,
     noServerCall,
     options ? options.encodeFormAction : undefined,

packages/react-server-dom-turbopack/src/client/ReactFlightDOMClientBrowser.js

@@ -48,6 +48,7 @@ export type Options = {
 
 function createResponseFromOptions(options: void | Options) {
   return createResponse(
+    null,
     null,
     null,
     options && options.callServer ? options.callServer : undefined,

packages/react-server-dom-turbopack/src/client/ReactFlightDOMClientEdge.js

@@ -19,11 +19,13 @@ import type {ReactServerValue} from 'react-client/src/ReactFlightReplyClient';
 import type {
   ServerConsumerModuleMap,
   ModuleLoading,
+  ServerManifest,
 } from 'react-client/src/ReactFlightClientConfig';
 
 type ServerConsumerManifest = {
   moduleMap: ServerConsumerModuleMap,
   moduleLoading: ModuleLoading,
+  serverModuleMap: null | ServerManifest,
 };
 
 import {
@@ -78,6 +80,7 @@ export type Options = {
 function createResponseFromOptions(options: Options) {
   return createResponse(
     options.serverManifest.moduleMap,
+    options.serverManifest.serverModuleMap,
     options.serverManifest.moduleLoading,
     noServerCall,
     options.encodeFormAction,

packages/react-server-dom-turbopack/src/client/ReactFlightDOMClientNode.js

@@ -17,11 +17,13 @@ import type {
 import type {
   ServerConsumerModuleMap,
   ModuleLoading,
+  ServerManifest,
 } from 'react-client/src/ReactFlightClientConfig';
 
 type ServerConsumerManifest = {
   moduleMap: ServerConsumerModuleMap,
   moduleLoading: ModuleLoading,
+  serverModuleMap: null | ServerManifest,
 };
 
 import type {Readable} from 'stream';
@@ -71,6 +73,7 @@ function createFromNodeStream<T>(
 ): Thenable<T> {
   const response: Response = createResponse(
     serverConsumerManifest.moduleMap,
+    serverConsumerManifest.serverModuleMap,
     serverConsumerManifest.moduleLoading,
     noServerCall,
     options ? options.encodeFormAction : undefined,

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

@@ -26,6 +26,7 @@ global.AsyncLocalStorage = require('async_hooks').AsyncLocalStorage;
 let serverExports;
 let clientExports;
 let webpackMap;
+let webpackServerMap;
 let webpackModules;
 let webpackModuleLoading;
 let React;
@@ -63,6 +64,7 @@ describe('ReactFlightDOMEdge', () => {
     serverExports = WebpackMock.serverExports;
     clientExports = WebpackMock.clientExports;
     webpackMap = WebpackMock.webpackMap;
+    webpackServerMap = WebpackMock.webpackServerMap;
     webpackModules = WebpackMock.webpackModules;
     webpackModuleLoading = WebpackMock.moduleLoading;
 
@@ -212,6 +214,72 @@ describe('ReactFlightDOMEdge', () => {
     expect(result).toEqual('<span>Client Component</span>');
   });
 
+  it('should be able to load a server reference on a consuming server if a mapping exists', async () => {
+    function greet() {
+      return 'hi';
+    }
+    const ServerModule = serverExports({
+      greet,
+    });
+
+    const stream = await serverAct(() =>
+      ReactServerDOMServer.renderToReadableStream(
+        {
+          method: ServerModule.greet,
+        },
+        webpackMap,
+      ),
+    );
+    const response = ReactServerDOMClient.createFromReadableStream(stream, {
+      serverConsumerManifest: {
+        moduleMap: webpackMap,
+        serverModuleMap: webpackServerMap,
+        moduleLoading: webpackModuleLoading,
+      },
+    });
+
+    const result = await response;
+
+    expect(result.method).toBe(greet);
+  });
+
+  it('should be able to load a server reference on a consuming server if a mapping exists (async)', async () => {
+    let resolve;
+    const chunkPromise = new Promise(r => (resolve = r));
+
+    function greet() {
+      return 'hi';
+    }
+    const ServerModule = serverExports(
+      {
+        greet,
+      },
+      chunkPromise,
+    );
+
+    const stream = await serverAct(() =>
+      ReactServerDOMServer.renderToReadableStream(
+        {
+          method: ServerModule.greet,
+        },
+        webpackMap,
+      ),
+    );
+    const response = ReactServerDOMClient.createFromReadableStream(stream, {
+      serverConsumerManifest: {
+        moduleMap: webpackMap,
+        serverModuleMap: webpackServerMap,
+        moduleLoading: webpackModuleLoading,
+      },
+    });
+
+    await resolve();
+
+    const result = await response;
+
+    expect(result.method).toBe(greet);
+  });
+
   it('should encode long string in a compact format', async () => {
     const testString = '"\n\t'.repeat(500) + '🙃';
     const testString2 = 'hello'.repeat(400);

packages/react-server-dom-webpack/src/client/ReactFlightDOMClientBrowser.js

@@ -48,6 +48,7 @@ export type Options = {
 
 function createResponseFromOptions(options: void | Options) {
   return createResponse(
+    null,
     null,
     null,
     options && options.callServer ? options.callServer : undefined,

packages/react-server-dom-webpack/src/client/ReactFlightDOMClientEdge.js

@@ -19,11 +19,13 @@ import type {ReactServerValue} from 'react-client/src/ReactFlightReplyClient';
 import type {
   ServerConsumerModuleMap,
   ModuleLoading,
+  ServerManifest,
 } from 'react-client/src/ReactFlightClientConfig';
 
 type ServerConsumerManifest = {
   moduleMap: ServerConsumerModuleMap,
   moduleLoading: ModuleLoading,
+  serverModuleMap: null | ServerManifest,
 };
 
 import {
@@ -78,6 +80,7 @@ export type Options = {
 function createResponseFromOptions(options: Options) {
   return createResponse(
     options.serverConsumerManifest.moduleMap,
+    options.serverConsumerManifest.serverModuleMap,
     options.serverConsumerManifest.moduleLoading,
     noServerCall,
     options.encodeFormAction,

packages/react-server-dom-webpack/src/client/ReactFlightDOMClientNode.js

@@ -17,11 +17,13 @@ import type {
 import type {
   ServerConsumerModuleMap,
   ModuleLoading,
+  ServerManifest,
 } from 'react-client/src/ReactFlightClientConfig';
 
 type ServerConsumerManifest = {
   moduleMap: ServerConsumerModuleMap,
   moduleLoading: ModuleLoading,
+  serverModuleMap: null | ServerManifest,
 };
 
 import type {Readable} from 'stream';
@@ -72,6 +74,7 @@ function createFromNodeStream<T>(
 ): Thenable<T> {
   const response: Response = createResponse(
     serverConsumerManifest.moduleMap,
+    serverConsumerManifest.serverModuleMap,
     serverConsumerManifest.moduleLoading,
     noServerCall,
     options ? options.encodeFormAction : undefined,