[Flight] Expose registerServerReference from the client builds (#32534)

This is used to register Server References that exist in the current
environment but also exists in the server it might call into. Such as a
remote server.

If the value comes from the remote server in the first place then this
is called automatically to ensure that you can pass a reference back to
where it came from - even if the `serverModuleMap` option is used. This
was already the case when `serverModuleMap` wasn't passed. This is how
you can pass server references back to the server. However, when we
added `serverModuleMap` that pass was skipped because we were getting
real functions instead of proxies.

For functions that wasn't yet passed from the remote server to the
current server, we can register them eagerly just like we do for
`import('/server').registerServerReference()`. You can now also do this
with `import('/client').registerServerReference()`. We could make them
shared so you only have to do this once but it might not be possible to
pass to the remote server and the remote server might not even be the
same RSC renderer. Therefore I split them. It's up to the compiler
whether it should do that or not. It has to know that any function you
might call might be able to receive it. This is currently global to a
specific RSC renderer.

packages/react-client/src/ReactFlightClient.js

@@ -64,7 +64,10 @@ import {
   rendererPackageName,
 } from './ReactFlightClientConfig';
 
-import {createBoundServerReference} from './ReactFlightReplyClient';
+import {
+  createBoundServerReference,
+  registerBoundServerReference,
+} from './ReactFlightReplyClient';
 
 import {readTemporaryReference} from './ReactFlightTemporaryReferences';
 
@@ -1096,7 +1099,14 @@ function loadServerReference<A: Iterable<any>, T>(
   let promise: null | Thenable<any> = preloadModule(serverReference);
   if (!promise) {
     if (!metaData.bound) {
-      return (requireModule(serverReference): any);
+      const resolvedValue = (requireModule(serverReference): any);
+      registerBoundServerReference(
+        resolvedValue,
+        metaData.id,
+        metaData.bound,
+        response._encodeFormAction,
+      );
+      return resolvedValue;
     } else {
       promise = Promise.resolve(metaData.bound);
     }
@@ -1128,6 +1138,13 @@ function loadServerReference<A: Iterable<any>, T>(
       resolvedValue = resolvedValue.bind.apply(resolvedValue, boundArgs);
     }
 
+    registerBoundServerReference(
+      resolvedValue,
+      metaData.id,
+      metaData.bound,
+      response._encodeFormAction,
+    );
+
     parentObject[key] = resolvedValue;
 
     // If this is the root object for a model reference, where `handler.value`

packages/react-client/src/ReactFlightReplyClient.js

@@ -1125,11 +1125,12 @@ function createFakeServerFunction<A: Iterable<any>, T>(
   }
 }
 
-function registerServerReference(
-  proxy: any,
-  reference: {id: ServerReferenceId, bound: null | Thenable<Array<any>>},
+export function registerBoundServerReference<T: Function>(
+  reference: T,
+  id: ServerReferenceId,
+  bound: null | Thenable<Array<any>>,
   encodeFormAction: void | EncodeFormActionCallback,
-) {
+): void {
   // Expose encoder for use by SSR, as well as a special bind that can be used to
   // keep server capabilities.
   if (usedWithSSR) {
@@ -1147,13 +1148,22 @@ function registerServerReference(
               encodeFormAction,
             );
           };
-    Object.defineProperties((proxy: any), {
+    Object.defineProperties((reference: any), {
       $$FORM_ACTION: {value: $$FORM_ACTION},
       $$IS_SIGNATURE_EQUAL: {value: isSignatureEqual},
       bind: {value: bind},
     });
   }
-  knownServerReferences.set(proxy, reference);
+  knownServerReferences.set(reference, {id, bound});
+}
+
+export function registerServerReference<T: Function>(
+  reference: T,
+  id: ServerReferenceId,
+  encodeFormAction?: EncodeFormActionCallback,
+): ServerReference<T> {
+  registerBoundServerReference(reference, id, null, encodeFormAction);
+  return reference;
 }
 
 // $FlowFixMe[method-unbinding]
@@ -1258,7 +1268,7 @@ export function createBoundServerReference<A: Iterable<any>, T>(
       );
     }
   }
-  registerServerReference(action, {id, bound}, encodeFormAction);
+  registerBoundServerReference(action, id, bound, encodeFormAction);
   return action;
 }
 
@@ -1358,6 +1368,6 @@ export function createServerReference<A: Iterable<any>, T>(
       );
     }
   }
-  registerServerReference(action, {id, bound: null}, encodeFormAction);
+  registerBoundServerReference(action, id, null, encodeFormAction);
   return action;
 }

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

@@ -25,9 +25,11 @@ import {
   injectIntoDevTools,
 } from 'react-client/src/ReactFlightClient';
 
-import {
-  processReply,
+import {processReply} from 'react-client/src/ReactFlightReplyClient';
+
+export {
   createServerReference,
+  registerServerReference,
 } from 'react-client/src/ReactFlightReplyClient';
 
 import type {TemporaryReferenceSet} from 'react-client/src/ReactFlightTemporaryReferences';
@@ -151,12 +153,7 @@ function encodeReply(
   });
 }
 
-export {
-  createFromFetch,
-  createFromReadableStream,
-  encodeReply,
-  createServerReference,
-};
+export {createFromFetch, createFromReadableStream, encodeReply};
 
 if (__DEV__) {
   injectIntoDevTools();

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

@@ -26,6 +26,8 @@ import {
 
 import {createServerReference as createServerReferenceImpl} from 'react-client/src/ReactFlightReplyClient';
 
+export {registerServerReference} from 'react-client/src/ReactFlightReplyClient';
+
 function noServerCall() {
   throw new Error(
     'Server Functions cannot be called during initial render. ' +

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

@@ -26,6 +26,8 @@ import {
   createServerReference as createServerReferenceImpl,
 } from 'react-client/src/ReactFlightReplyClient';
 
+export {registerServerReference} from 'react-client/src/ReactFlightReplyClient';
+
 import type {TemporaryReferenceSet} from 'react-client/src/ReactFlightTemporaryReferences';
 
 export {createTemporaryReferenceSet} from 'react-client/src/ReactFlightTemporaryReferences';

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

@@ -25,6 +25,8 @@ import {
   createServerReference as createServerReferenceImpl,
 } from 'react-client/src/ReactFlightReplyClient';
 
+export {registerServerReference} from 'react-client/src/ReactFlightReplyClient';
+
 import type {TemporaryReferenceSet} from 'react-client/src/ReactFlightTemporaryReferences';
 
 export {createTemporaryReferenceSet} from 'react-client/src/ReactFlightTemporaryReferences';

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

@@ -21,6 +21,8 @@ import {
 
 import {createServerReference as createServerReferenceImpl} from 'react-client/src/ReactFlightReplyClient';
 
+export {registerServerReference} from 'react-client/src/ReactFlightReplyClient';
+
 function findSourceMapURL(filename: string, environmentName: string) {
   const devServer = parcelRequire.meta.devServer;
   if (devServer != null) {

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

@@ -25,9 +25,11 @@ import {
   injectIntoDevTools,
 } from 'react-client/src/ReactFlightClient';
 
-import {
-  processReply,
+import {processReply} from 'react-client/src/ReactFlightReplyClient';
+
+export {
   createServerReference,
+  registerServerReference,
 } from 'react-client/src/ReactFlightReplyClient';
 
 import type {TemporaryReferenceSet} from 'react-client/src/ReactFlightTemporaryReferences';
@@ -150,12 +152,7 @@ function encodeReply(
   });
 }
 
-export {
-  createFromFetch,
-  createFromReadableStream,
-  encodeReply,
-  createServerReference,
-};
+export {createFromFetch, createFromReadableStream, encodeReply};
 
 if (__DEV__) {
   injectIntoDevTools();

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

@@ -41,6 +41,8 @@ import {
   createServerReference as createServerReferenceImpl,
 } from 'react-client/src/ReactFlightReplyClient';
 
+export {registerServerReference} from 'react-client/src/ReactFlightReplyClient';
+
 import type {TemporaryReferenceSet} from 'react-client/src/ReactFlightTemporaryReferences';
 
 export {createTemporaryReferenceSet} from 'react-client/src/ReactFlightTemporaryReferences';

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

@@ -38,6 +38,8 @@ import {
 
 import {createServerReference as createServerReferenceImpl} from 'react-client/src/ReactFlightReplyClient';
 
+export {registerServerReference} from 'react-client/src/ReactFlightReplyClient';
+
 function noServerCall() {
   throw new Error(
     'Server Functions cannot be called during initial render. ' +

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

@@ -301,6 +301,48 @@ describe('ReactFlightDOMEdge', () => {
     expect(result.boundMethod()).toBe('hi, there');
   });
 
+  it('should load a server reference on a consuming server and pass it back', async () => {
+    function greet(name) {
+      return 'hi, ' + name;
+    }
+    const ServerModule = serverExports({
+      greet,
+    });
+
+    const stream = await serverAct(() =>
+      ReactServerDOMServer.renderToReadableStream(
+        {
+          method: ServerModule.greet,
+          boundMethod: ServerModule.greet.bind(null, 'there'),
+        },
+        webpackMap,
+      ),
+    );
+    const response = ReactServerDOMClient.createFromReadableStream(stream, {
+      serverConsumerManifest: {
+        moduleMap: webpackMap,
+        serverModuleMap: webpackServerMap,
+        moduleLoading: webpackModuleLoading,
+      },
+    });
+
+    const result = await response;
+
+    expect(result.method).toBe(greet);
+    expect(result.boundMethod()).toBe('hi, there');
+
+    const body = await ReactServerDOMClient.encodeReply({
+      method: result.method,
+      boundMethod: result.boundMethod,
+    });
+    const replyResult = await ReactServerDOMServer.decodeReply(
+      body,
+      webpackServerMap,
+    );
+    expect(replyResult.method).toBe(greet);
+    expect(replyResult.boundMethod()).toBe('hi, there');
+  });
+
   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/__tests__/ReactFlightDOMReplyEdge-test.js

@@ -22,7 +22,7 @@ if (typeof File === 'undefined' || typeof FormData === 'undefined') {
   global.FormData = require('undici').FormData;
 }
 
-// let serverExports;
+let serverExports;
 let webpackServerMap;
 let ReactServerDOMServer;
 let ReactServerDOMClient;
@@ -36,7 +36,7 @@ describe('ReactFlightDOMReplyEdge', () => {
       require('react-server-dom-webpack/server.edge'),
     );
     const WebpackMock = require('./utils/WebpackMock');
-    // serverExports = WebpackMock.serverExports;
+    serverExports = WebpackMock.serverExports;
     webpackServerMap = WebpackMock.webpackServerMap;
     ReactServerDOMServer = require('react-server-dom-webpack/server.edge');
     jest.resetModules();
@@ -308,4 +308,29 @@ describe('ReactFlightDOMReplyEdge', () => {
     expect(await decoded.a).toBe('hello');
     expect(Array.from(await decoded.b)).toEqual(Array.from(buffer));
   });
+
+  it('can pass a registered server reference', async () => {
+    function greet(name) {
+      return 'hi, ' + name;
+    }
+    const ServerModule = serverExports({
+      greet,
+    });
+
+    ReactServerDOMClient.registerServerReference(
+      ServerModule.greet,
+      ServerModule.greet.$$id,
+    );
+
+    const body = await ReactServerDOMClient.encodeReply({
+      method: ServerModule.greet,
+      boundMethod: ServerModule.greet.bind(null, 'there'),
+    });
+    const replyResult = await ReactServerDOMServer.decodeReply(
+      body,
+      webpackServerMap,
+    );
+    expect(replyResult.method).toBe(greet);
+    expect(replyResult.boundMethod()).toBe('hi, there');
+  });
 });

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

@@ -25,9 +25,11 @@ import {
   injectIntoDevTools,
 } from 'react-client/src/ReactFlightClient';
 
-import {
-  processReply,
+import {processReply} from 'react-client/src/ReactFlightReplyClient';
+
+export {
   createServerReference,
+  registerServerReference,
 } from 'react-client/src/ReactFlightReplyClient';
 
 import type {TemporaryReferenceSet} from 'react-client/src/ReactFlightTemporaryReferences';
@@ -150,12 +152,7 @@ function encodeReply(
   });
 }
 
-export {
-  createFromFetch,
-  createFromReadableStream,
-  encodeReply,
-  createServerReference,
-};
+export {createFromFetch, createFromReadableStream, encodeReply};
 
 if (__DEV__) {
   injectIntoDevTools();

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

@@ -41,6 +41,8 @@ import {
   createServerReference as createServerReferenceImpl,
 } from 'react-client/src/ReactFlightReplyClient';
 
+export {registerServerReference} from 'react-client/src/ReactFlightReplyClient';
+
 import type {TemporaryReferenceSet} from 'react-client/src/ReactFlightTemporaryReferences';
 
 export {createTemporaryReferenceSet} from 'react-client/src/ReactFlightTemporaryReferences';

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

@@ -39,6 +39,8 @@ import {
 
 import {createServerReference as createServerReferenceImpl} from 'react-client/src/ReactFlightReplyClient';
 
+export {registerServerReference} from 'react-client/src/ReactFlightReplyClient';
+
 function noServerCall() {
   throw new Error(
     'Server Functions cannot be called during initial render. ' +

packages/react-server/src/ReactFlightReplyServer.js

@@ -936,8 +936,10 @@ function parseModelString(
         // Server Reference
         const ref = value.slice(2);
         // TODO: Just encode this in the reference inline instead of as a model.
-        const metaData: {id: ServerReferenceId, bound: Thenable<Array<any>>} =
-          getOutlinedModel(response, ref, obj, key, createModel);
+        const metaData: {
+          id: ServerReferenceId,
+          bound: null | Thenable<Array<any>>,
+        } = getOutlinedModel(response, ref, obj, key, createModel);
         return loadServerReference(
           response,
           metaData.id,