From e2de3fdb080c660846b6cc3b3f2fc90461bbfb26 Mon Sep 17 00:00:00 2001
From: Sebastian Markbage <sebastian@calyptus.eu>
Date: Fri, 14 Apr 2023 22:52:01 -0400
Subject: [PATCH] Add a way to create Server Reference Proxies on the client

This lets the client bundle encode Server References without them first
being passed from an RSC payload. Like if you just import "use server"
from the client.

In the future we could expand this to allow .bind() too.
---
 .../src/ReactFlightReplyClient.js             |  7 +-
 .../src/ReactFlightServerReferenceRegistry.js | 15 +++++
 .../src/ReactFlightDOMClientBrowser.js        | 13 +++-
 .../__tests__/ReactFlightDOMBrowser-test.js   | 64 +++++++++++++++++++
 4 files changed, 96 insertions(+), 3 deletions(-)

diff --git a/packages/react-client/src/ReactFlightReplyClient.js b/packages/react-client/src/ReactFlightReplyClient.js
index 6f9581c36ed83..18fc2834e1da7 100644
--- a/packages/react-client/src/ReactFlightReplyClient.js
+++ b/packages/react-client/src/ReactFlightReplyClient.js
@@ -9,7 +9,10 @@
 
 import type {Thenable} from 'shared/ReactTypes';
 
-import {knownServerReferences} from './ReactFlightServerReferenceRegistry';
+import {
+  knownServerReferences,
+  createServerReference,
+} from './ReactFlightServerReferenceRegistry';
 
 import {
   REACT_ELEMENT_TYPE,
@@ -312,3 +315,5 @@ export function processReply(
     }
   }
 }
+
+export {createServerReference};
diff --git a/packages/react-client/src/ReactFlightServerReferenceRegistry.js b/packages/react-client/src/ReactFlightServerReferenceRegistry.js
index 7436a19915d2e..06ad06e9b3e46 100644
--- a/packages/react-client/src/ReactFlightServerReferenceRegistry.js
+++ b/packages/react-client/src/ReactFlightServerReferenceRegistry.js
@@ -9,9 +9,24 @@
 
 import type {Thenable} from 'shared/ReactTypes';
 
+export type CallServerCallback = <A, T>(id: any, args: A) => Promise<T>;
+
 type ServerReferenceId = any;
 
 export const knownServerReferences: WeakMap<
   Function,
   {id: ServerReferenceId, bound: null | Thenable<Array<any>>},
 > = new WeakMap();
+
+export function createServerReference<A: Iterable<any>, T>(
+  id: ServerReferenceId,
+  callServer: CallServerCallback,
+): (...A) => Promise<T> {
+  const proxy = function (): Promise<T> {
+    // $FlowFixMe[method-unbinding]
+    const args = Array.prototype.slice.call(arguments);
+    return callServer(id, args);
+  };
+  knownServerReferences.set(proxy, {id: id, bound: null});
+  return proxy;
+}
diff --git a/packages/react-server-dom-webpack/src/ReactFlightDOMClientBrowser.js b/packages/react-server-dom-webpack/src/ReactFlightDOMClientBrowser.js
index 537b96187a00d..f847a636e6c89 100644
--- a/packages/react-server-dom-webpack/src/ReactFlightDOMClientBrowser.js
+++ b/packages/react-server-dom-webpack/src/ReactFlightDOMClientBrowser.js
@@ -22,7 +22,10 @@ import {
   close,
 } from 'react-client/src/ReactFlightClientStream';
 
-import {processReply} from 'react-client/src/ReactFlightReplyClient';
+import {
+  processReply,
+  createServerReference,
+} from 'react-client/src/ReactFlightReplyClient';
 
 type CallServerCallback = <A, T>(string, args: A) => Promise<T>;
 
@@ -125,4 +128,10 @@ function encodeReply(
   });
 }
 
-export {createFromXHR, createFromFetch, createFromReadableStream, encodeReply};
+export {
+  createFromXHR,
+  createFromFetch,
+  createFromReadableStream,
+  encodeReply,
+  createServerReference,
+};
diff --git a/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMBrowser-test.js b/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMBrowser-test.js
index 0321c5f75d18d..55c55e19f8779 100644
--- a/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMBrowser-test.js
+++ b/packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMBrowser-test.js
@@ -893,6 +893,70 @@ describe('ReactFlightDOMBrowser', () => {
     expect(result).toBe('Hello Split');
   });
 
+  it('can pass a server function by importing from client back to server', async () => {
+    function greet(transform, text) {
+      return 'Hello ' + transform(text);
+    }
+
+    function upper(text) {
+      return text.toUpperCase();
+    }
+
+    const ServerModuleA = serverExports({
+      greet,
+    });
+    const ServerModuleB = serverExports({
+      upper,
+    });
+
+    let actionProxy;
+
+    // This is a Proxy representing ServerModuleB in the Client bundle.
+    const ServerModuleBImportedOnClient = {
+      upper: ReactServerDOMClient.createServerReference(
+        ServerModuleB.upper.$$id,
+        async function (ref, args) {
+          const body = await ReactServerDOMClient.encodeReply(args);
+          return callServer(ref, body);
+        },
+      ),
+    };
+
+    function Client({action}) {
+      // Client side pass a Server Reference into an action.
+      actionProxy = text => action(ServerModuleBImportedOnClient.upper, text);
+      return 'Click Me';
+    }
+
+    const ClientRef = clientExports(Client);
+
+    const stream = ReactServerDOMServer.renderToReadableStream(
+      <ClientRef action={ServerModuleA.greet} />,
+      webpackMap,
+    );
+
+    const response = ReactServerDOMClient.createFromReadableStream(stream, {
+      async callServer(ref, args) {
+        const body = await ReactServerDOMClient.encodeReply(args);
+        return callServer(ref, body);
+      },
+    });
+
+    function App() {
+      return use(response);
+    }
+
+    const container = document.createElement('div');
+    const root = ReactDOMClient.createRoot(container);
+    await act(() => {
+      root.render(<App />);
+    });
+    expect(container.innerHTML).toBe('Click Me');
+
+    const result = await actionProxy('hi');
+    expect(result).toBe('Hello HI');
+  });
+
   it('can bind arguments to a server reference', async () => {
     let actionProxy;