[Flight] Add support for Module References in transport protocol

packages/react-client/src/ReactFlightClient.js

@@ -41,8 +41,9 @@ export type JSONValue =
 
 const PENDING = 0;
 const RESOLVED_MODEL = 1;
-const INITIALIZED = 2;
-const ERRORED = 3;
+const RESOLVED_MODULE = 2;
+const INITIALIZED = 3;
+const ERRORED = 4;
 
 type PendingChunk = {
   _status: 0,
@@ -56,21 +57,28 @@ type ResolvedModelChunk = {
   _response: Response,
   then(resolve: () => mixed): void,
 };
-type InitializedChunk<T> = {
+type ResolvedModuleChunk<T> = {
   _status: 2,
+  _value: ModuleReference<T>,
+  _response: Response,
+  then(resolve: () => mixed): void,
+};
+type InitializedChunk<T> = {
+  _status: 3,
   _value: T,
   _response: Response,
   then(resolve: () => mixed): void,
 };
 type ErroredChunk = {
-  _status: 3,
+  _status: 4,
   _value: Error,
   _response: Response,
   then(resolve: () => mixed): void,
 };
 type SomeChunk<T> =
   | PendingChunk
   | ResolvedModelChunk
+  | ResolvedModuleChunk<T>
   | InitializedChunk<T>
   | ErroredChunk;
 
@@ -105,6 +113,8 @@ function readChunk<T>(chunk: SomeChunk<T>): T {
       return chunk._value;
     case RESOLVED_MODEL:
       return initializeModelChunk(chunk);
+    case RESOLVED_MODULE:
+      return initializeModuleChunk(chunk);
     case PENDING:
       // eslint-disable-next-line no-throw-literal
       throw (chunk: Wakeable);
@@ -155,6 +165,13 @@ function createResolvedModelChunk(
   return new Chunk(RESOLVED_MODEL, value, response);
 }
 
+function createResolvedModuleChunk<T>(
+  response: Response,
+  value: ModuleReference<T>,
+): ResolvedModuleChunk<T> {
+  return new Chunk(RESOLVED_MODULE, value, response);
+}
+
 function resolveModelChunk<T>(
   chunk: SomeChunk<T>,
   value: UninitializedModel,
@@ -170,6 +187,21 @@ function resolveModelChunk<T>(
   wakeChunk(listeners);
 }
 
+function resolveModuleChunk<T>(
+  chunk: SomeChunk<T>,
+  value: ModuleReference<T>,
+): void {
+  if (chunk._status !== PENDING) {
+    // We already resolved. We didn't expect to see this.
+    return;
+  }
+  const listeners = chunk._value;
+  const resolvedChunk: ResolvedModuleChunk<T> = (chunk: any);
+  resolvedChunk._status = RESOLVED_MODULE;
+  resolvedChunk._value = value;
+  wakeChunk(listeners);
+}
+
 function initializeModelChunk<T>(chunk: ResolvedModelChunk): T {
   const value: T = parseModel(chunk._response, chunk._value);
   const initializedChunk: InitializedChunk<T> = (chunk: any);
@@ -178,6 +210,14 @@ function initializeModelChunk<T>(chunk: ResolvedModelChunk): T {
   return value;
 }
 
+function initializeModuleChunk<T>(chunk: ResolvedModuleChunk<T>): T {
+  const value: T = requireModule(chunk._value);
+  const initializedChunk: InitializedChunk<T> = (chunk: any);
+  initializedChunk._status = INITIALIZED;
+  initializedChunk._value = value;
+  return value;
+}
+
 // Report that any missing chunks in the model is now going to throw this
 // error upon read. Also notify any pending promises.
 export function reportGlobalError(response: Response, error: Error): void {
@@ -241,7 +281,7 @@ function createElement(type, key, props): React$Element<any> {
 
 type UninitializedBlockPayload<Data> = [
   mixed,
-  ModuleMetaData | SomeChunk<ModuleMetaData>,
+  BlockRenderFunction<any, Data> | SomeChunk<BlockRenderFunction<any, Data>>,
   Data | SomeChunk<Data>,
   Response,
 ];
@@ -250,14 +290,7 @@ function initializeBlock<Props, Data>(
   tuple: UninitializedBlockPayload<Data>,
 ): BlockComponent<Props, Data> {
   // Require module first and then data. The ordering matters.
-  const moduleMetaData: ModuleMetaData = readMaybeChunk(tuple[1]);
-  const moduleReference: ModuleReference<
-    BlockRenderFunction<Props, Data>,
-  > = resolveModuleReference(moduleMetaData);
-  // TODO: Do this earlier, as the chunk is resolved.
-  preloadModule(moduleReference);
-
-  const moduleExport = requireModule(moduleReference);
+  const moduleExport = readMaybeChunk(tuple[1]);
 
   // The ordering here is important because this call might suspend.
   // We don't want that to prevent the module graph for being initialized.
@@ -363,6 +396,28 @@ export function resolveModel(
   }
 }
 
+export function resolveModule(
+  response: Response,
+  id: number,
+  model: UninitializedModel,
+): void {
+  const chunks = response._chunks;
+  const chunk = chunks.get(id);
+  const moduleMetaData: ModuleMetaData = parseModel(response, model);
+  const moduleReference = resolveModuleReference(moduleMetaData);
+
+  // TODO: Add an option to encode modules that are lazy loaded.
+  // For now we preload all modules as early as possible since it's likely
+  // that we'll need them.
+  preloadModule(moduleReference);
+
+  if (!chunk) {
+    chunks.set(id, createResolvedModuleChunk(response, moduleReference));
+  } else {
+    resolveModuleChunk(chunk, moduleReference);
+  }
+}
+
 export function resolveError(
   response: Response,
   id: number,

packages/react-client/src/ReactFlightClientStream.js

@@ -10,6 +10,7 @@
 import type {Response} from './ReactFlightClientHostConfigStream';
 
 import {
+  resolveModule,
   resolveModel,
   resolveError,
   createResponse as createResponseBase,
@@ -39,6 +40,13 @@ function processFullRow(response: Response, row: string): void {
       resolveModel(response, id, json);
       return;
     }
+    case 'M': {
+      const colon = row.indexOf(':', 1);
+      const id = parseInt(row.substring(1, colon), 16);
+      const json = row.substring(colon + 1);
+      resolveModule(response, id, json);
+      return;
+    }
     case 'E': {
       const colon = row.indexOf(':', 1);
       const id = parseInt(row.substring(1, colon), 16);
@@ -48,9 +56,9 @@ function processFullRow(response: Response, row: string): void {
       return;
     }
     default: {
-      // Assume this is the root model.
-      resolveModel(response, 0, row);
-      return;
+      throw new Error(
+        "Error parsing the data. It's probably an error code or network corruption.",
+      );
     }
   }
 }

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

@@ -53,17 +53,29 @@ describe('ReactFlight', () => {
     };
   });
 
+  function moduleReference(value) {
+    return {
+      $$typeof: Symbol.for('react.module.reference'),
+      value: value,
+    };
+  }
+
   function block(render, load) {
     if (load === undefined) {
       return () => {
-        return ReactNoopFlightServerRuntime.serverBlockNoData(render);
+        return ReactNoopFlightServerRuntime.serverBlockNoData(
+          moduleReference(render),
+        );
       };
     }
     return function(...args) {
       const curriedLoad = () => {
         return load(...args);
       };
-      return ReactNoopFlightServerRuntime.serverBlock(render, curriedLoad);
+      return ReactNoopFlightServerRuntime.serverBlock(
+        moduleReference(render),
+        curriedLoad,
+      );
     };
   }
 
@@ -97,6 +109,35 @@ describe('ReactFlight', () => {
     });
   });
 
+  it('can render a client component using a module reference and render there', () => {
+    function UserClient(props) {
+      return (
+        <span>
+          {props.greeting}, {props.name}
+        </span>
+      );
+    }
+    const User = moduleReference(UserClient);
+
+    function Greeting({firstName, lastName}) {
+      return <User greeting="Hello" name={firstName + ' ' + lastName} />;
+    }
+
+    const model = {
+      greeting: <Greeting firstName="Seb" lastName="Smith" />,
+    };
+
+    const transport = ReactNoopFlightServer.render(model);
+
+    act(() => {
+      const rootModel = ReactNoopFlightClient.read(transport);
+      const greeting = rootModel.greeting;
+      ReactNoop.render(greeting);
+    });
+
+    expect(ReactNoop).toMatchRenderedOutput(<span>Hello, Seb Smith</span>);
+  });
+
   if (ReactFeatureFlags.enableBlocksAPI) {
     it('can transfer a Block to the client and render there, without data', () => {
       function User(props, data) {

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

@@ -42,8 +42,14 @@ const ReactNoopFlightServer = ReactFlightServer({
   formatChunk(type: string, props: Object): Uint8Array {
     return Buffer.from(JSON.stringify({type, props}), 'utf8');
   },
-  resolveModuleMetaData(config: void, renderFn: Function) {
-    return saveModule(renderFn);
+  isModuleReference(reference: Object): boolean {
+    return reference.$$typeof === Symbol.for('react.module.reference');
+  },
+  resolveModuleMetaData(
+    config: void,
+    reference: {$$typeof: Symbol, value: any},
+  ) {
+    return saveModule(reference.value);
   },
 });
 

packages/react-server/src/ReactFlightServer.js

@@ -24,8 +24,10 @@ import {
   flushBuffered,
   close,
   processModelChunk,
+  processModuleChunk,
   processErrorChunk,
   resolveModuleMetaData,
+  isModuleReference,
 } from './ReactFlightServerConfig';
 
 import {
@@ -83,6 +85,7 @@ export type Request = {
   nextChunkId: number,
   pendingChunks: number,
   pingedSegments: Array<Segment>,
+  completedModuleChunks: Array<Chunk>,
   completedJSONChunks: Array<Chunk>,
   completedErrorChunks: Array<Chunk>,
   flowing: boolean,
@@ -103,6 +106,7 @@ export function createRequest(
     nextChunkId: 0,
     pendingChunks: 0,
     pingedSegments: pingedSegments,
+    completedModuleChunks: [],
     completedJSONChunks: [],
     completedErrorChunks: [],
     flowing: false,
@@ -151,6 +155,10 @@ function attemptResolveElement(element: React$Element<any>): ReactModel {
   ) {
     return element.props.children;
   } else if (type != null && typeof type === 'object') {
+    if (isModuleReference(type)) {
+      // This is a reference to a client component.
+      return [REACT_ELEMENT_TYPE, type, element.key, element.props];
+    }
     switch (type.$$typeof) {
       case REACT_FORWARD_REF_TYPE: {
         const render = type.render;
@@ -391,19 +399,8 @@ export function resolveModelToJSON(
     switch (key) {
       case '1': {
         // Module reference
-        const moduleReference: ModuleReference<any> = (value: any);
-        try {
-          const moduleMetaData: ModuleMetaData = resolveModuleMetaData(
-            request.bundlerConfig,
-            moduleReference,
-          );
-          return (moduleMetaData: ReactJSONValue);
-        } catch (x) {
-          request.pendingChunks++;
-          const errorId = request.nextChunkId++;
-          emitErrorChunk(request, errorId, x);
-          return serializeIDRef(errorId);
-        }
+        // Encode as a normal value.
+        break;
       }
       case '2': {
         // Load function
@@ -467,7 +464,30 @@ export function resolveModelToJSON(
     }
   }
 
+  if (value === null) {
+    return null;
+  }
+
   if (typeof value === 'object') {
+    if (isModuleReference(value)) {
+      const moduleReference: ModuleReference<any> = (value: any);
+      try {
+        const moduleMetaData: ModuleMetaData = resolveModuleMetaData(
+          request.bundlerConfig,
+          moduleReference,
+        );
+        request.pendingChunks++;
+        const moduleId = request.nextChunkId++;
+        emitModuleChunk(request, moduleId, moduleMetaData);
+        return serializeIDRef(moduleId);
+      } catch (x) {
+        request.pendingChunks++;
+        const errorId = request.nextChunkId++;
+        emitErrorChunk(request, errorId, x);
+        return serializeIDRef(errorId);
+      }
+    }
+
     if (__DEV__) {
       if (value !== null && !isArray(value)) {
         // Verify that this is a simple plain object.
@@ -595,6 +615,15 @@ function emitErrorChunk(request: Request, id: number, error: mixed): void {
   request.completedErrorChunks.push(processedChunk);
 }
 
+function emitModuleChunk(
+  request: Request,
+  id: number,
+  moduleMetaData: ModuleMetaData,
+): void {
+  const processedChunk = processModuleChunk(request, id, moduleMetaData);
+  request.completedModuleChunks.push(processedChunk);
+}
+
 function retrySegment(request: Request, segment: Segment): void {
   const query = segment.query;
   let value;
@@ -654,8 +683,22 @@ function flushCompletedChunks(request: Request): void {
   const destination = request.destination;
   beginWriting(destination);
   try {
-    const jsonChunks = request.completedJSONChunks;
+    // We emit module chunks first in the stream so that
+    // they can be preloaded as early as possible.
+    const moduleChunks = request.completedModuleChunks;
     let i = 0;
+    for (; i < moduleChunks.length; i++) {
+      request.pendingChunks--;
+      const chunk = moduleChunks[i];
+      if (!writeChunk(destination, chunk)) {
+        request.flowing = false;
+        i++;
+        break;
+      }
+    }
+    // Next comes model data.
+    const jsonChunks = request.completedJSONChunks;
+    i = 0;
     for (; i < jsonChunks.length; i++) {
       request.pendingChunks--;
       const chunk = jsonChunks[i];
@@ -666,6 +709,9 @@ function flushCompletedChunks(request: Request): void {
       }
     }
     jsonChunks.splice(0, i);
+    // Finally, errors are sent. The idea is that it's ok to delay
+    // any error messages and prioritize display of other parts of
+    // the page.
     const errorChunks = request.completedErrorChunks;
     i = 0;
     for (; i < errorChunks.length; i++) {

packages/react-server/src/ReactFlightServerBundlerConfigCustom.js

@@ -12,4 +12,5 @@ declare var $$$hostConfig: any;
 export opaque type BundlerConfig = mixed; // eslint-disable-line no-undef
 export opaque type ModuleReference<T> = mixed; // eslint-disable-line no-undef
 export opaque type ModuleMetaData: any = mixed; // eslint-disable-line no-undef
+export const isModuleReference = $$$hostConfig.isModuleReference;
 export const resolveModuleMetaData = $$$hostConfig.resolveModuleMetaData;

packages/react-server/src/ReactFlightServerConfigStream.js

@@ -14,15 +14,15 @@
 FLIGHT PROTOCOL GRAMMAR
 
 Response
-- JSONData RowSequence
-- JSONData
+- RowSequence
 
 RowSequence
 - Row RowSequence
 - Row
 
 Row
 - "J" RowID JSONData
+- "M" RowID JSONModuleData
 - "H" RowID HTMLData
 - "B" RowID BlobData
 - "U" RowID URLData
@@ -95,12 +95,17 @@ export function processModelChunk(
   model: ReactModel,
 ): Chunk {
   const json = stringify(model, request.toJSON);
-  let row;
-  if (id === 0) {
-    row = json + '\n';
-  } else {
-    row = serializeRowHeader('J', id) + json + '\n';
-  }
+  const row = serializeRowHeader('J', id) + json + '\n';
+  return convertStringToBuffer(row);
+}
+
+export function processModuleChunk(
+  request: Request,
+  id: number,
+  moduleMetaData: ReactModel,
+): Chunk {
+  const json = stringify(moduleMetaData);
+  const row = serializeRowHeader('M', id) + json + '\n';
   return convertStringToBuffer(row);
 }
 

packages/react-transport-dom-relay/src/ReactFlightDOMRelayClient.js

@@ -10,6 +10,7 @@
 export {
   createResponse,
   resolveModel,
+  resolveModule,
   resolveError,
   close,
 } from 'react-client/src/ReactFlightClient';

packages/react-transport-dom-relay/src/ReactFlightDOMRelayClientHostConfig.js

@@ -9,6 +9,10 @@
 
 import type {JSONValue, ResponseBase} from 'react-client/src/ReactFlightClient';
 
+import type JSResourceReference from 'JSResourceReference';
+
+export type ModuleReference<T> = JSResourceReference<T>;
+
 import {
   parseModelString,
   parseModelTuple,
@@ -20,10 +24,7 @@ export {
   requireModule,
 } from 'ReactFlightDOMRelayClientIntegration';
 
-export type {
-  ModuleReference,
-  ModuleMetaData,
-} from 'ReactFlightDOMRelayClientIntegration';
+export type {ModuleMetaData} from 'ReactFlightDOMRelayClientIntegration';
 
 export opaque type UninitializedModel = JSONValue;
 

packages/react-transport-dom-relay/src/ReactFlightDOMRelayServerHostConfig.js

@@ -9,28 +9,35 @@
 
 import type {Request, ReactModel} from 'react-server/src/ReactFlightServer';
 
+import JSResourceReference from 'JSResourceReference';
+
+export type ModuleReference<T> = JSResourceReference<T>;
+
 import type {
   Destination,
   BundlerConfig,
-  ModuleReference,
   ModuleMetaData,
 } from 'ReactFlightDOMRelayServerIntegration';
 
 import {resolveModelToJSON} from 'react-server/src/ReactFlightServer';
 
 import {
   emitModel,
+  emitModule,
   emitError,
   resolveModuleMetaData as resolveModuleMetaDataImpl,
 } from 'ReactFlightDOMRelayServerIntegration';
 
 export type {
   Destination,
   BundlerConfig,
-  ModuleReference,
   ModuleMetaData,
 } from 'ReactFlightDOMRelayServerIntegration';
 
+export function isModuleReference(reference: Object): boolean {
+  return reference instanceof JSResourceReference;
+}
+
 export function resolveModuleMetaData<T>(
   config: BundlerConfig,
   resource: ModuleReference<T>,
@@ -52,6 +59,11 @@ export type Chunk =
       id: number,
       json: JSONValue,
     }
+  | {
+      type: 'module',
+      id: number,
+      json: ModuleMetaData,
+    }
   | {
       type: 'error',
       id: number,
@@ -121,6 +133,19 @@ export function processModelChunk(
   };
 }
 
+export function processModuleChunk(
+  request: Request,
+  id: number,
+  moduleMetaData: ModuleMetaData,
+): Chunk {
+  // The moduleMetaData is already a JSON serializable value.
+  return {
+    type: 'module',
+    id: id,
+    json: moduleMetaData,
+  };
+}
+
 export function scheduleWork(callback: () => void) {
   callback();
 }
@@ -132,6 +157,8 @@ export function beginWriting(destination: Destination) {}
 export function writeChunk(destination: Destination, chunk: Chunk): boolean {
   if (chunk.type === 'json') {
     emitModel(destination, chunk.id, chunk.json);
+  } else if (chunk.type === 'module') {
+    emitModule(destination, chunk.id, chunk.json);
   } else {
     emitError(destination, chunk.id, chunk.json.message, chunk.json.stack);
   }

packages/react-transport-dom-relay/src/__mocks__/JSResourceReference.js

@@ -0,0 +1,19 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+'use strict';
+
+class JSResourceReference {
+  constructor(exportedValue) {
+    this._moduleId = exportedValue;
+  }
+  getModuleID() {
+    return this._moduleId;
+  }
+}
+
+module.exports = JSResourceReference;

packages/react-transport-dom-relay/src/__mocks__/ReactFlightDOMRelayClientIntegration.js

@@ -7,13 +7,15 @@
 
 'use strict';
 
+import JSResourceReference from 'JSResourceReference';
+
 const ReactFlightDOMRelayClientIntegration = {
   resolveModuleReference(moduleData) {
-    return moduleData;
+    return new JSResourceReference(moduleData);
   },
   preloadModule(moduleReference) {},
   requireModule(moduleReference) {
-    return moduleReference;
+    return moduleReference._moduleId;
   },
 };
 

packages/react-transport-dom-relay/src/__mocks__/ReactFlightDOMRelayServerIntegration.js

@@ -15,6 +15,13 @@ const ReactFlightDOMRelayServerIntegration = {
       json: json,
     });
   },
+  emitModule(destination, id, json) {
+    destination.push({
+      type: 'module',
+      id: id,
+      json: json,
+    });
+  },
   emitError(destination, id, message, stack) {
     destination.push({
       type: 'error',
@@ -24,7 +31,7 @@ const ReactFlightDOMRelayServerIntegration = {
   },
   close(destination) {},
   resolveModuleMetaData(config, resource) {
-    return resource;
+    return resource._moduleId;
   },
 };
 

packages/react-transport-dom-relay/src/__tests__/ReactFlightDOMRelay-test.internal.js

@@ -10,6 +10,7 @@
 let act;
 let React;
 let ReactDOM;
+let JSResourceReference;
 let ReactDOMFlightRelayServer;
 let ReactDOMFlightRelayServerRuntime;
 let ReactDOMFlightRelayClient;
@@ -24,6 +25,7 @@ describe('ReactFlightDOMRelay', () => {
     ReactDOMFlightRelayServer = require('react-transport-dom-relay/server');
     ReactDOMFlightRelayServerRuntime = require('react-transport-dom-relay/server-runtime');
     ReactDOMFlightRelayClient = require('react-transport-dom-relay');
+    JSResourceReference = require('JSResourceReference');
   });
 
   function readThrough(data) {
@@ -32,6 +34,8 @@ describe('ReactFlightDOMRelay', () => {
       const chunk = data[i];
       if (chunk.type === 'json') {
         ReactDOMFlightRelayClient.resolveModel(response, chunk.id, chunk.json);
+      } else if (chunk.type === 'module') {
+        ReactDOMFlightRelayClient.resolveModule(response, chunk.id, chunk.json);
       } else {
         ReactDOMFlightRelayClient.resolveError(
           response,
@@ -47,14 +51,18 @@ describe('ReactFlightDOMRelay', () => {
   }
 
   function block(render, load) {
+    const reference = new JSResourceReference(render);
     if (load === undefined) {
-      return ReactDOMFlightRelayServerRuntime.serverBlock(render);
+      return ReactDOMFlightRelayServerRuntime.serverBlock(reference);
     }
     return function(...args) {
       const curriedLoad = () => {
         return load(...args);
       };
-      return ReactDOMFlightRelayServerRuntime.serverBlock(render, curriedLoad);
+      return ReactDOMFlightRelayServerRuntime.serverBlock(
+        reference,
+        curriedLoad,
+      );
     };
   }
 
@@ -93,6 +101,39 @@ describe('ReactFlightDOMRelay', () => {
     });
   });
 
+  // @gate experimental
+  it('can render a client component using a module reference and render there', () => {
+    function UserClient(props) {
+      return (
+        <span>
+          {props.greeting}, {props.name}
+        </span>
+      );
+    }
+    const User = new JSResourceReference(UserClient);
+
+    function Greeting({firstName, lastName}) {
+      return <User greeting="Hello" name={firstName + ' ' + lastName} />;
+    }
+
+    const model = {
+      greeting: <Greeting firstName="Seb" lastName="Smith" />,
+    };
+
+    const transport = [];
+    ReactDOMFlightRelayServer.render(model, transport);
+
+    const modelClient = readThrough(transport);
+
+    const container = document.createElement('div');
+    const root = ReactDOM.createRoot(container);
+    act(() => {
+      root.render(modelClient.greeting);
+    });
+
+    expect(container.innerHTML).toEqual('<span>Hello, Seb Smith</span>');
+  });
+
   // @gate experimental
   it('can transfer a Block to the client and render there', () => {
     function load(firstName, lastName) {

packages/react-transport-dom-webpack/src/ReactFlightServerWebpackBundlerConfig.js

@@ -14,17 +14,26 @@ type WebpackMap = {
 export type BundlerConfig = WebpackMap;
 
 // eslint-disable-next-line no-unused-vars
-export type ModuleReference<T> = string;
+export type ModuleReference<T> = {
+  $$typeof: Symbol,
+  name: string,
+};
 
 export type ModuleMetaData = {
   id: string,
   chunks: Array<string>,
   name: string,
 };
 
+const MODULE_TAG = Symbol.for('react.module.reference');
+
+export function isModuleReference(reference: Object): boolean {
+  return reference.$$typeof === MODULE_TAG;
+}
+
 export function resolveModuleMetaData<T>(
   config: BundlerConfig,
-  modulePath: ModuleReference<T>,
+  moduleReference: ModuleReference<T>,
 ): ModuleMetaData {
-  return config[modulePath];
+  return config[moduleReference.name];
 }

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

@@ -83,8 +83,9 @@ describe('ReactFlightDOM', () => {
       const curriedLoad = () => {
         return load(...args);
       };
+      const MODULE_TAG = Symbol.for('react.module.reference');
       return ReactTransportDOMServerRuntime.serverBlock(
-        'path/' + idx,
+        {$$typeof: MODULE_TAG, name: 'path/' + idx},
         curriedLoad,
       );
     };

packages/shared/isValidElementType.js

@@ -27,6 +27,11 @@ import {
 } from 'shared/ReactSymbols';
 import {enableScopeAPI} from './ReactFeatureFlags';
 
+let REACT_MODULE_REFERENCE: number | Symbol = 0;
+if (typeof Symbol === 'function') {
+  REACT_MODULE_REFERENCE = Symbol.for('react.module.reference');
+}
+
 export default function isValidElementType(type: mixed) {
   if (typeof type === 'string' || typeof type === 'function') {
     return true;
@@ -54,6 +59,12 @@ export default function isValidElementType(type: mixed) {
       type.$$typeof === REACT_CONTEXT_TYPE ||
       type.$$typeof === REACT_FORWARD_REF_TYPE ||
       type.$$typeof === REACT_FUNDAMENTAL_TYPE ||
+      // This needs to include all possible module reference object
+      // types supported by any Flight configuration anywhere since
+      // we don't know which Flight build this will end up being used
+      // with.
+      type.$$typeof === REACT_MODULE_REFERENCE ||
+      type.getModuleID !== undefined ||
       type.$$typeof === REACT_BLOCK_TYPE ||
       type[(0: any)] === REACT_SERVER_BLOCK_TYPE
     ) {

scripts/flow/react-relay-hooks.js

@@ -15,6 +15,14 @@ type JSONValue =
   | {+[key: string]: JSONValue}
   | $ReadOnlyArray<JSONValue>;
 
+declare class JSResourceReference<T> {
+  _moduleId: T;
+}
+
+declare module 'JSResourceReference' {
+  declare export default typeof JSResourceReference;
+}
+
 declare module 'ReactFlightDOMRelayServerIntegration' {
   declare export opaque type Destination;
   declare export opaque type BundlerConfig;
@@ -23,6 +31,11 @@ declare module 'ReactFlightDOMRelayServerIntegration' {
     id: number,
     json: JSONValue,
   ): void;
+  declare export function emitModule(
+    destination: Destination,
+    id: number,
+    json: ModuleMetaData,
+  ): void;
   declare export function emitError(
     destination: Destination,
     id: number,
@@ -31,24 +44,22 @@ declare module 'ReactFlightDOMRelayServerIntegration' {
   ): void;
   declare export function close(destination: Destination): void;
 
-  declare export opaque type ModuleReference<T>;
   declare export type ModuleMetaData = JSONValue;
   declare export function resolveModuleMetaData<T>(
     config: BundlerConfig,
-    resourceReference: ModuleReference<T>,
+    resourceReference: JSResourceReference<T>,
   ): ModuleMetaData;
 }
 
 declare module 'ReactFlightDOMRelayClientIntegration' {
-  declare export opaque type ModuleReference<T>;
   declare export opaque type ModuleMetaData;
   declare export function resolveModuleReference<T>(
     moduleData: ModuleMetaData,
-  ): ModuleReference<T>;
+  ): JSResourceReference<T>;
   declare export function preloadModule<T>(
-    moduleReference: ModuleReference<T>,
+    moduleReference: JSResourceReference<T>,
   ): void;
   declare export function requireModule<T>(
-    moduleReference: ModuleReference<T>,
+    moduleReference: JSResourceReference<T>,
   ): T;
 }

scripts/jest/setupHostConfigs.js

@@ -34,6 +34,7 @@ jest.mock('react-server/flight', () => {
     jest.mock(shimServerStreamConfigPath, () => config);
     jest.mock(shimServerFormatConfigPath, () => config);
     jest.mock('react-server/src/ReactFlightServerBundlerConfigCustom', () => ({
+      isModuleReference: config.isModuleReference,
       resolveModuleMetaData: config.resolveModuleMetaData,
     }));
     jest.mock(shimFlightServerConfigPath, () =>

scripts/rollup/bundles.js

@@ -301,14 +301,19 @@ const bundles = [
       'react',
       'react-dom/server',
       'ReactFlightDOMRelayServerIntegration',
+      'JSResourceReference',
     ],
   },
   {
     bundleTypes: [FB_WWW_DEV, FB_WWW_PROD],
     moduleType: RENDERER,
     entry: 'react-transport-dom-relay/server-runtime',
     global: 'ReactFlightDOMRelayServerRuntime',
-    externals: ['react', 'ReactFlightDOMRelayServerIntegration'],
+    externals: [
+      'react',
+      'ReactFlightDOMRelayServerIntegration',
+      'JSResourceReference',
+    ],
   },
 
   /******* React DOM Flight Client Relay *******/
@@ -317,7 +322,11 @@ const bundles = [
     moduleType: RENDERER,
     entry: 'react-transport-dom-relay',
     global: 'ReactFlightDOMRelayClient',
-    externals: ['react', 'ReactFlightDOMRelayClientIntegration'],
+    externals: [
+      'react',
+      'ReactFlightDOMRelayClientIntegration',
+      'JSResourceReference',
+    ],
   },
 
   /******* React ART *******/