Expose registerServerReference from the client builds

sebmarkbage · sebmarkbage · commit 7ece2218b0f5 · 2025-03-05T19:46:08.000-05:00
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 .
diff --git a/packages/react-client/src/ReactFlightClient.js b/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`
diff --git a/packages/react-client/src/ReactFlightReplyClient.js b/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;
 }
diff --git a/packages/react-server-dom-esm/src/client/ReactFlightDOMClientBrowser.js b/packages/react-server-dom-esm/src/client/ReactFlightDOMClientBrowser.js
@@ -25,11 +25,15 @@ import {
   injectIntoDevTools,
 } from 'react-client/src/ReactFlightClient';
 
-import {
-  processReply,
+import {processReply} from 'react-client/src/ReactFlightReplyClient';
+
+export {
   createServerReference,
+  registerServerReference,
 } 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';
@@ -151,12 +155,7 @@ function encodeReply(
   });
 }
 
-export {
-  createFromFetch,
-  createFromReadableStream,
-  encodeReply,
-  createServerReference,
-};
+export {createFromFetch, createFromReadableStream, encodeReply};
 
 if (__DEV__) {
   injectIntoDevTools();
diff --git a/packages/react-server-dom-esm/src/client/ReactFlightDOMClientNode.js b/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. ' +
diff --git a/packages/react-server-dom-parcel/src/client/ReactFlightDOMClientBrowser.js b/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';
diff --git a/packages/react-server-dom-parcel/src/client/ReactFlightDOMClientEdge.js b/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';
diff --git a/packages/react-server-dom-parcel/src/client/ReactFlightDOMClientNode.js b/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) {
diff --git a/packages/react-server-dom-turbopack/src/client/ReactFlightDOMClientBrowser.js b/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();
diff --git a/packages/react-server-dom-turbopack/src/client/ReactFlightDOMClientEdge.js b/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';
diff --git a/packages/react-server-dom-turbopack/src/client/ReactFlightDOMClientNode.js b/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. ' +
diff --git a/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMEdge-test.js b/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);
diff --git a/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMReplyEdge-test.js b/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');
+  });
 });
diff --git a/packages/react-server-dom-webpack/src/client/ReactFlightDOMClientBrowser.js b/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();
diff --git a/packages/react-server-dom-webpack/src/client/ReactFlightDOMClientEdge.js b/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';
diff --git a/packages/react-server-dom-webpack/src/client/ReactFlightDOMClientNode.js b/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. ' +
diff --git a/packages/react-server/src/ReactFlightReplyServer.js b/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,

Original file line number	Diff line number	Diff line change
`@@ -1125,11 +1125,12 @@ function createFakeServerFunction<A: Iterable<any>, T>(`
`1125`	`1125`	`}`
`1126`	`1126`	`}`
`1127`	`1127`
`1128`		`-function registerServerReference(`
`1129`		`- proxy: any,`
`1130`		`- reference: {id: ServerReferenceId, bound: null \| Thenable<Array<any>>},`
	`1128`	`+export function registerBoundServerReference<T: Function>(`
	`1129`	`+ reference: T,`
	`1130`	`+ id: ServerReferenceId,`
	`1131`	`+ bound: null \| Thenable<Array<any>>,`
`1131`	`1132`	`encodeFormAction: void \| EncodeFormActionCallback,`
`1132`		`-) {`
	`1133`	`+): void {`
`1133`	`1134`	`// Expose encoder for use by SSR, as well as a special bind that can be used to`
`1134`	`1135`	`// keep server capabilities.`
`1135`	`1136`	`if (usedWithSSR) {`
`@@ -1147,13 +1148,22 @@ function registerServerReference(`
`1147`	`1148`	`encodeFormAction,`
`1148`	`1149`	`);`
`1149`	`1150`	`};`
`1150`		`- Object.defineProperties((proxy: any), {`
	`1151`	`+ Object.defineProperties((reference: any), {`
`1151`	`1152`	`$$FORM_ACTION: {value: $$FORM_ACTION},`
`1152`	`1153`	`$$IS_SIGNATURE_EQUAL: {value: isSignatureEqual},`
`1153`	`1154`	`bind: {value: bind},`
`1154`	`1155`	`});`
`1155`	`1156`	`}`
`1156`		`- knownServerReferences.set(proxy, reference);`
	`1157`	`+ knownServerReferences.set(reference, {id, bound});`
	`1158`	`+}`
	`1159`	`+`
	`1160`	`+export function registerServerReference<T: Function>(`
	`1161`	`+ reference: T,`
	`1162`	`+ id: ServerReferenceId,`
	`1163`	`+ encodeFormAction?: EncodeFormActionCallback,`
	`1164`	`+): ServerReference<T> {`
	`1165`	`+ registerBoundServerReference(reference, id, null, encodeFormAction);`
	`1166`	`+ return reference;`
`1157`	`1167`	`}`
`1158`	`1168`
`1159`	`1169`	`// $FlowFixMe[method-unbinding]`
`@@ -1258,7 +1268,7 @@ export function createBoundServerReference<A: Iterable<any>, T>(`
`1258`	`1268`	`);`
`1259`	`1269`	`}`
`1260`	`1270`	`}`
`1261`		`- registerServerReference(action, {id, bound}, encodeFormAction);`
	`1271`	`+ registerBoundServerReference(action, id, bound, encodeFormAction);`
`1262`	`1272`	`return action;`
`1263`	`1273`	`}`
`1264`	`1274`
`@@ -1358,6 +1368,6 @@ export function createServerReference<A: Iterable<any>, T>(`
`1358`	`1368`	`);`
`1359`	`1369`	`}`
`1360`	`1370`	`}`
`1361`		`- registerServerReference(action, {id, bound: null}, encodeFormAction);`
	`1371`	`+ registerBoundServerReference(action, id, null, encodeFormAction);`
`1362`	`1372`	`return action;`
`1363`	`1373`	`}`