Add a way to create Server Reference Proxies on the client (#26632)

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. A bundler could already emit these proxies to be called
on the client but the subtle difference is that those proxies couldn't
be passed back into the server by reference. They have to be registered
with React.

We don't currently implement importing `"use server"` from client
components in the reference implementation. It'd need to expand the
Webpack plugin with a loader that rewrites files with the `"use server"`
in the client bundle.

```
"use server";

export async function action() {
   ...
}
```
->
```
import {createServerReference} from "react-server-dom-webpack/client";
import {callServer} from "some-router/call-server";

export const action = createServerReference('1234#action', callServer);
```

The technique I use here is that the compiled output has to call
`createServerReference(id, callServer)` with the `$$id` and proxy
implementation. We then return a proxy function that is registered with
a WeakMap to the particular instance of the Flight Client.

This might be hard to implement because it requires emitting module
imports to a specific stateful runtime module in the compiler. A benefit
is that this ensures that this particular reference is locked to a
specific client if there are multiple - e.g. talking to different
servers.

It's fairly arbitrary whether we use a WeakMap technique (like we do on
the client) vs an `$$id` (like we do on the server). Not sure what's
best overall. The WeakMap is nice because it doesn't leak implementation
details that might be abused to consumers. We should probably pick one
and unify.

Author: sebmarkbage

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};

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;
+}

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,
+};

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

@@ -29,6 +29,13 @@ function noServerCall() {
   );
 }
 
+export function createServerReference<A: Iterable<any>, T>(
+  id: any,
+  callServer: any,
+): (...A) => Promise<T> {
+  return noServerCall;
+}
+
 export type Options = {
   moduleMap?: $NonMaybeType<SSRManifest>,
 };

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

@@ -32,6 +32,13 @@ function noServerCall() {
   );
 }
 
+export function createServerReference<A: Iterable<any>, T>(
+  id: any,
+  callServer: any,
+): (...A) => Promise<T> {
+  return noServerCall;
+}
+
 function createFromNodeStream<T>(
   stream: Readable,
   moduleMap: $NonMaybeType<SSRManifest>,

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;