Override .bind on Server References on the Client (#27282)

That way when you bind arguments to a Server Reference, it's still a
server reference and works with progressive enhancement.

This already works on the Server (RSC) layer.

packages/react-client/src/ReactFlightClient.js

@@ -33,13 +33,9 @@ import {
   readPartialStringChunk,
   readFinalStringChunk,
   createStringDecoder,
-  usedWithSSR,
 } from './ReactFlightClientConfig';
 
-import {
-  encodeFormAction,
-  knownServerReferences,
-} from './ReactFlightReplyClient';
+import {registerServerReference} from './ReactFlightReplyClient';
 
 import {
   REACT_LAZY_TYPE,
@@ -545,12 +541,7 @@ function createServerReferenceProxy<A: Iterable<any>, T>(
       return callServer(metaData.id, bound.concat(args));
     });
   };
-  // Expose encoder for use by SSR.
-  if (usedWithSSR) {
-    // Only expose this in builds that would actually use it. Not needed on the client.
-    (proxy: any).$$FORM_ACTION = encodeFormAction;
-  }
-  knownServerReferences.set(proxy, metaData);
+  registerServerReference(proxy, metaData);
   return proxy;
 }
 

packages/react-client/src/ReactFlightReplyClient.js

@@ -44,7 +44,7 @@ export type CallServerCallback = <A, T>(id: any, args: A) => Promise<T>;
 
 export type ServerReferenceId = any;
 
-export const knownServerReferences: WeakMap<
+const knownServerReferences: WeakMap<
   Function,
   {id: ServerReferenceId, bound: null | Thenable<Array<any>>},
 > = new WeakMap();
@@ -488,6 +488,45 @@ export function encodeFormAction(
   };
 }
 
+export function registerServerReference(
+  proxy: any,
+  reference: {id: ServerReferenceId, bound: null | Thenable<Array<any>>},
+) {
+  // Expose encoder for use by SSR, as well as a special bind that can be used to
+  // keep server capabilities.
+  if (usedWithSSR) {
+    // Only expose this in builds that would actually use it. Not needed on the client.
+    Object.defineProperties((proxy: any), {
+      $$FORM_ACTION: {value: encodeFormAction},
+      bind: {value: bind},
+    });
+  }
+  knownServerReferences.set(proxy, reference);
+}
+
+// $FlowFixMe[method-unbinding]
+const FunctionBind = Function.prototype.bind;
+// $FlowFixMe[method-unbinding]
+const ArraySlice = Array.prototype.slice;
+function bind(this: Function) {
+  // $FlowFixMe[unsupported-syntax]
+  const newFn = FunctionBind.apply(this, arguments);
+  const reference = knownServerReferences.get(this);
+  if (reference) {
+    const args = ArraySlice.call(arguments, 1);
+    let boundPromise = null;
+    if (reference.bound !== null) {
+      boundPromise = Promise.resolve((reference.bound: any)).then(boundArgs =>
+        boundArgs.concat(args),
+      );
+    } else {
+      boundPromise = Promise.resolve(args);
+    }
+    registerServerReference(newFn, {id: reference.id, bound: boundPromise});
+  }
+  return newFn;
+}
+
 export function createServerReference<A: Iterable<any>, T>(
   id: ServerReferenceId,
   callServer: CallServerCallback,
@@ -497,11 +536,6 @@ export function createServerReference<A: Iterable<any>, T>(
     const args = Array.prototype.slice.call(arguments);
     return callServer(id, args);
   };
-  // Expose encoder for use by SSR.
-  if (usedWithSSR) {
-    // Only expose this in builds that would actually use it. Not needed on the client.
-    (proxy: any).$$FORM_ACTION = encodeFormAction;
-  }
-  knownServerReferences.set(proxy, {id: id, bound: null});
+  registerServerReference(proxy, {id, bound: null});
   return proxy;
 }

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

@@ -22,7 +22,9 @@ global.TextDecoder = require('util').TextDecoder;
 global.setTimeout = cb => cb();
 
 let container;
+let clientExports;
 let serverExports;
+let webpackMap;
 let webpackServerMap;
 let React;
 let ReactDOMServer;
@@ -37,7 +39,9 @@ describe('ReactFlightDOMForm', () => {
       require('react-server-dom-webpack/server.edge'),
     );
     const WebpackMock = require('./utils/WebpackMock');
+    clientExports = WebpackMock.clientExports;
     serverExports = WebpackMock.serverExports;
+    webpackMap = WebpackMock.webpackMap;
     webpackServerMap = WebpackMock.webpackServerMap;
     React = require('react');
     ReactServerDOMServer = require('react-server-dom-webpack/server.edge');
@@ -236,4 +240,72 @@ describe('ReactFlightDOMForm', () => {
     expect(result).toBe('helloc');
     expect(foo).toBe('barc');
   });
+
+  // @gate enableFormActions
+  it('can bind an imported server action on the client without hydrating it', async () => {
+    let foo = null;
+
+    const ServerModule = serverExports(function action(bound, formData) {
+      foo = formData.get('foo') + bound.complex;
+      return 'hello';
+    });
+    const serverAction = ReactServerDOMClient.createServerReference(
+      ServerModule.$$id,
+    );
+    function Client() {
+      return (
+        <form action={serverAction.bind(null, {complex: 'object'})}>
+          <input type="text" name="foo" defaultValue="bar" />
+        </form>
+      );
+    }
+
+    const ssrStream = await ReactDOMServer.renderToReadableStream(<Client />);
+    await readIntoContainer(ssrStream);
+
+    const form = container.firstChild;
+
+    expect(foo).toBe(null);
+
+    const result = await submit(form);
+
+    expect(result).toBe('hello');
+    expect(foo).toBe('barobject');
+  });
+
+  // @gate enableFormActions
+  it('can bind a server action on the client without hydrating it', async () => {
+    let foo = null;
+
+    const serverAction = serverExports(function action(bound, formData) {
+      foo = formData.get('foo') + bound.complex;
+      return 'hello';
+    });
+
+    function Client({action}) {
+      return (
+        <form action={action.bind(null, {complex: 'object'})}>
+          <input type="text" name="foo" defaultValue="bar" />
+        </form>
+      );
+    }
+    const ClientRef = await clientExports(Client);
+
+    const rscStream = ReactServerDOMServer.renderToReadableStream(
+      <ClientRef action={serverAction} />,
+      webpackMap,
+    );
+    const response = ReactServerDOMClient.createFromReadableStream(rscStream);
+    const ssrStream = await ReactDOMServer.renderToReadableStream(response);
+    await readIntoContainer(ssrStream);
+
+    const form = container.firstChild;
+
+    expect(foo).toBe(null);
+
+    const result = await submit(form);
+
+    expect(result).toBe('hello');
+    expect(foo).toBe('barobject');
+  });
 });