[Flight] Encode React Elements in Replies as Temporary References

packages/react-client/src/ReactFlightClient.js

@@ -35,6 +35,8 @@ import type {
 
 import type {Postpone} from 'react/src/ReactPostpone';
 
+import type {TemporaryReferenceSet} from './ReactFlightTemporaryReferences';
+
 import {
   enableBinaryFlight,
   enablePostpone,
@@ -55,6 +57,8 @@ import {
 
 import {registerServerReference} from './ReactFlightReplyClient';
 
+import {readTemporaryReference} from './ReactFlightTemporaryReferences';
+
 import {
   REACT_LAZY_TYPE,
   REACT_ELEMENT_TYPE,
@@ -224,6 +228,7 @@ export type Response = {
   _rowTag: number, // 0 indicates that we're currently parsing the row ID
   _rowLength: number, // remaining bytes in the row. 0 indicates that we're looking for a newline.
   _buffer: Array<Uint8Array>, // chunks received so far as part of this row
+  _tempRefs: void | TemporaryReferenceSet, // the set temporary references can be resolved from
 };
 
 function readChunk<T>(chunk: SomeChunk<T>): T {
@@ -689,6 +694,18 @@ function parseModelString(
         const metadata = getOutlinedModel(response, id);
         return createServerReferenceProxy(response, metadata);
       }
+      case 'T': {
+        // Temporary Reference
+        const id = parseInt(value.slice(2), 16);
+        const temporaryReferences = response._tempRefs;
+        if (temporaryReferences == null) {
+          throw new Error(
+            'Missing a temporary reference set but the RSC response returned a temporary reference. ' +
+              'Pass a temporaryReference option with the set that was used with the reply.',
+          );
+        }
+        return readTemporaryReference(temporaryReferences, id);
+      }
       case 'Q': {
         // Map
         const id = parseInt(value.slice(2), 16);
@@ -837,6 +854,7 @@ export function createResponse(
   callServer: void | CallServerCallback,
   encodeFormAction: void | EncodeFormActionCallback,
   nonce: void | string,
+  temporaryReferences: void | TemporaryReferenceSet,
 ): Response {
   const chunks: Map<number, SomeChunk<any>> = new Map();
   const response: Response = {
@@ -853,6 +871,7 @@ export function createResponse(
     _rowTag: 0,
     _rowLength: 0,
     _buffer: [],
+    _tempRefs: temporaryReferences,
   };
   // Don't inline this call because it causes closure to outline the call above.
   response._fromJSON = createFromJSONCallback(response);

packages/react-client/src/ReactFlightReplyClient.js

@@ -14,6 +14,9 @@ import type {
   RejectedThenable,
   ReactCustomFormAction,
 } from 'shared/ReactTypes';
+import type {LazyComponent} from 'react/src/ReactLazy';
+import type {TemporaryReferenceSet} from './ReactFlightTemporaryReferences';
+
 import {enableRenderableContext} from 'shared/ReactFeatureFlags';
 
 import {
@@ -30,6 +33,8 @@ import {
   objectName,
 } from 'shared/ReactSerializationErrors';
 
+import {writeTemporaryReference} from './ReactFlightTemporaryReferences';
+
 import isArray from 'shared/isArray';
 import getPrototypeOf from 'shared/getPrototypeOf';
 
@@ -84,9 +89,9 @@ export type ReactServerValue =
 
 type ReactServerObject = {+[key: string]: ReactServerValue};
 
-// function serializeByValueID(id: number): string {
-//   return '$' + id.toString(16);
-// }
+function serializeByValueID(id: number): string {
+  return '$' + id.toString(16);
+}
 
 function serializePromiseID(id: number): string {
   return '$@' + id.toString(16);
@@ -96,6 +101,10 @@ function serializeServerReferenceID(id: number): string {
   return '$F' + id.toString(16);
 }
 
+function serializeTemporaryReferenceID(id: number): string {
+  return '$T' + id.toString(16);
+}
+
 function serializeSymbolReference(name: string): string {
   return '$S' + name;
 }
@@ -158,6 +167,7 @@ function escapeStringValue(value: string): string {
 export function processReply(
   root: ReactServerValue,
   formFieldPrefix: string,
+  temporaryReferences: void | TemporaryReferenceSet,
   resolve: (string | FormData) => void,
   reject: (error: mixed) => void,
 ): void {
@@ -206,6 +216,81 @@ export function processReply(
     }
 
     if (typeof value === 'object') {
+      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) : ''),
+            );
+          }
+          return serializeTemporaryReferenceID(
+            writeTemporaryReference(temporaryReferences, value),
+          );
+        }
+        case REACT_LAZY_TYPE: {
+          // Resolve lazy as if it wasn't here. In the future this will be encoded as a Promise.
+          const lazy: LazyComponent<any, any> = (value: any);
+          const payload = lazy._payload;
+          const init = lazy._init;
+          if (formData === null) {
+            // Upgrade to use FormData to allow us to stream this value.
+            formData = new FormData();
+          }
+          pendingParts++;
+          try {
+            const resolvedModel = init(payload);
+            // We always outline this as a separate part even though we could inline it
+            // because it ensures a more deterministic encoding.
+            const lazyId = nextPartId++;
+            const partJSON = JSON.stringify(resolvedModel, resolveToJSON);
+            // $FlowFixMe[incompatible-type] We know it's not null because we assigned it above.
+            const data: FormData = formData;
+            // eslint-disable-next-line react-internal/safe-string-coercion
+            data.append(formFieldPrefix + lazyId, partJSON);
+            return serializeByValueID(lazyId);
+          } catch (x) {
+            if (
+              typeof x === 'object' &&
+              x !== null &&
+              typeof x.then === 'function'
+            ) {
+              // Suspended
+              pendingParts++;
+              const lazyId = nextPartId++;
+              const thenable: Thenable<any> = (x: any);
+              const retry = function () {
+                // While the first promise resolved, its value isn't necessarily what we'll
+                // resolve into because we might suspend again.
+                try {
+                  const partJSON = JSON.stringify(value, resolveToJSON);
+                  // $FlowFixMe[incompatible-type] We know it's not null because we assigned it above.
+                  const data: FormData = formData;
+                  // eslint-disable-next-line react-internal/safe-string-coercion
+                  data.append(formFieldPrefix + lazyId, partJSON);
+                  pendingParts--;
+                  if (pendingParts === 0) {
+                    resolve(data);
+                  }
+                } catch (reason) {
+                  reject(reason);
+                }
+              };
+              thenable.then(retry, retry);
+              return serializeByValueID(lazyId);
+            } else {
+              // In the future we could consider serializing this as an error
+              // that throws on the server instead.
+              reject(x);
+              return null;
+            }
+          } finally {
+            pendingParts--;
+          }
+        }
+      }
+
       // $FlowFixMe[method-unbinding]
       if (typeof value.then === 'function') {
         // We assume that any object with a .then property is a "Thenable" type,
@@ -219,14 +304,18 @@ export function processReply(
         const thenable: Thenable<any> = (value: any);
         thenable.then(
           partValue => {
-            const partJSON = JSON.stringify(partValue, resolveToJSON);
-            // $FlowFixMe[incompatible-type] We know it's not null because we assigned it above.
-            const data: FormData = formData;
-            // eslint-disable-next-line react-internal/safe-string-coercion
-            data.append(formFieldPrefix + promiseId, partJSON);
-            pendingParts--;
-            if (pendingParts === 0) {
-              resolve(data);
+            try {
+              const partJSON = JSON.stringify(partValue, resolveToJSON);
+              // $FlowFixMe[incompatible-type] We know it's not null because we assigned it above.
+              const data: FormData = formData;
+              // eslint-disable-next-line react-internal/safe-string-coercion
+              data.append(formFieldPrefix + promiseId, partJSON);
+              pendingParts--;
+              if (pendingParts === 0) {
+                resolve(data);
+              }
+            } catch (reason) {
+              reject(reason);
             }
           },
           reason => {
@@ -288,23 +377,19 @@ export function processReply(
         proto !== ObjectPrototype &&
         (proto === null || getPrototypeOf(proto) !== null)
       ) {
-        throw new Error(
-          'Only plain objects, and a few built-ins, can be passed to Server Actions. ' +
-            'Classes or null prototypes are not supported.',
+        if (temporaryReferences === undefined) {
+          throw new Error(
+            'Only plain objects, and a few built-ins, can be passed to Server Actions. ' +
+              'Classes or null prototypes are not supported.',
+          );
+        }
+        // We can serialize class instances as temporary references.
+        return serializeTemporaryReferenceID(
+          writeTemporaryReference(temporaryReferences, value),
         );
       }
       if (__DEV__) {
-        if ((value: any).$$typeof === REACT_ELEMENT_TYPE) {
-          console.error(
-            'React Element cannot be passed to Server Functions from the Client.%s',
-            describeObjectForErrorMessage(parent, key),
-          );
-        } else if ((value: any).$$typeof === REACT_LAZY_TYPE) {
-          console.error(
-            'React Lazy cannot be passed to Server Functions from the Client.%s',
-            describeObjectForErrorMessage(parent, key),
-          );
-        } else if (
+        if (
           (value: any).$$typeof ===
           (enableRenderableContext ? REACT_CONTEXT_TYPE : REACT_PROVIDER_TYPE)
         ) {
@@ -382,9 +467,14 @@ export function processReply(
         formData.set(formFieldPrefix + refId, metaDataJSON);
         return serializeServerReferenceID(refId);
       }
-      throw new Error(
-        'Client Functions cannot be passed directly to Server Functions. ' +
-          'Only Functions passed from the Server can be passed back again.',
+      if (temporaryReferences === undefined) {
+        throw new Error(
+          'Client Functions cannot be passed directly to Server Functions. ' +
+            'Only Functions passed from the Server can be passed back again.',
+        );
+      }
+      return serializeTemporaryReferenceID(
+        writeTemporaryReference(temporaryReferences, value),
       );
     }
 
@@ -443,6 +533,7 @@ function encodeFormData(reference: any): Thenable<FormData> {
   processReply(
     reference,
     '',
+    undefined, // TODO: This means React Elements can't be used as state in progressive enhancement.
     (body: string | FormData) => {
       if (typeof body === 'string') {
         const data = new FormData();

packages/react-client/src/ReactFlightTemporaryReferences.js

@@ -0,0 +1,41 @@
+/**
+ * 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
+ */
+
+interface Reference {}
+
+export opaque type TemporaryReferenceSet = Array<Reference>;
+
+export function createTemporaryReferenceSet(): TemporaryReferenceSet {
+  return [];
+}
+
+export function writeTemporaryReference(
+  set: TemporaryReferenceSet,
+  object: Reference,
+): number {
+  // We always create a new entry regardless if we've already written the same
+  // object. This ensures that we always generate a deterministic encoding of
+  // each slot in the reply for cacheability.
+  const newId = set.length;
+  set.push(object);
+  return newId;
+}
+
+export function readTemporaryReference(
+  set: TemporaryReferenceSet,
+  id: number,
+): Reference {
+  if (id < 0 || id >= set.length) {
+    throw new Error(
+      "The RSC response contained a reference that doesn't exist in the temporary reference set. " +
+        'Always pass the matching set that was used to create the reply when parsing its response.',
+    );
+  }
+  return set[id];
+}

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

@@ -26,11 +26,18 @@ import {
   createServerReference,
 } from 'react-client/src/ReactFlightReplyClient';
 
+import type {TemporaryReferenceSet} from 'react-client/src/ReactFlightTemporaryReferences';
+
+export {createTemporaryReferenceSet} from 'react-client/src/ReactFlightTemporaryReferences';
+
+export type {TemporaryReferenceSet};
+
 type CallServerCallback = <A, T>(string, args: A) => Promise<T>;
 
 export type Options = {
   moduleBaseURL?: string,
   callServer?: CallServerCallback,
+  temporaryReferences?: TemporaryReferenceSet,
 };
 
 function createResponseFromOptions(options: void | Options) {
@@ -40,6 +47,9 @@ function createResponseFromOptions(options: void | Options) {
     options && options.callServer ? options.callServer : undefined,
     undefined, // encodeFormAction
     undefined, // nonce
+    options && options.temporaryReferences
+      ? options.temporaryReferences
+      : undefined,
   );
 }
 
@@ -97,11 +107,20 @@ function createFromFetch<T>(
 
 function encodeReply(
   value: ReactServerValue,
+  options?: {temporaryReferences?: TemporaryReferenceSet},
 ): Promise<
   string | URLSearchParams | FormData,
 > /* We don't use URLSearchParams yet but maybe */ {
   return new Promise((resolve, reject) => {
-    processReply(value, '', resolve, reject);
+    processReply(
+      value,
+      '',
+      options && options.temporaryReferences
+        ? options.temporaryReferences
+        : undefined,
+      resolve,
+      reject,
+    );
   });
 }
 

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

@@ -60,6 +60,7 @@ function createFromNodeStream<T>(
     noServerCall,
     options ? options.encodeFormAction : undefined,
     options && typeof options.nonce === 'string' ? options.nonce : undefined,
+    undefined, // TODO: If encodeReply is supported, this should support temporaryReferences
   );
   stream.on('data', chunk => {
     processBinaryChunk(response, chunk);