[Flight] Allow custom encoding of the form action

packages/react-client/src/ReactFlightClient.js

@@ -28,7 +28,10 @@ import type {
   HintModel,
 } from 'react-server/src/ReactFlightServerConfig';
 
-import type {CallServerCallback} from './ReactFlightReplyClient';
+import type {
+  CallServerCallback,
+  EncodeFormActionCallback,
+} from './ReactFlightReplyClient';
 
 import type {Postpone} from 'react/src/ReactPostpone';
 
@@ -53,7 +56,7 @@ import {
   REACT_POSTPONE_TYPE,
 } from 'shared/ReactSymbols';
 
-export type {CallServerCallback};
+export type {CallServerCallback, EncodeFormActionCallback};
 
 type UninitializedModel = string;
 
@@ -206,6 +209,7 @@ export type Response = {
   _bundlerConfig: SSRModuleMap,
   _moduleLoading: ModuleLoading,
   _callServer: CallServerCallback,
+  _encodeFormAction: void | EncodeFormActionCallback,
   _nonce: ?string,
   _chunks: Map<number, SomeChunk<any>>,
   _fromJSON: (key: string, value: JSONValue) => any,
@@ -592,7 +596,7 @@ function createServerReferenceProxy<A: Iterable<any>, T>(
       },
     );
   };
-  registerServerReference(proxy, metaData);
+  registerServerReference(proxy, metaData, response._encodeFormAction);
   return proxy;
 }
 
@@ -785,13 +789,15 @@ export function createResponse(
   bundlerConfig: SSRModuleMap,
   moduleLoading: ModuleLoading,
   callServer: void | CallServerCallback,
+  encodeFormAction: void | EncodeFormActionCallback,
   nonce: void | string,
 ): Response {
   const chunks: Map<number, SomeChunk<any>> = new Map();
   const response: Response = {
     _bundlerConfig: bundlerConfig,
     _moduleLoading: moduleLoading,
     _callServer: callServer !== undefined ? callServer : missingCall,
+    _encodeFormAction: encodeFormAction,
     _nonce: nonce,
     _chunks: chunks,
     _stringDecoder: createStringDecoder(),

packages/react-client/src/ReactFlightReplyClient.js

@@ -47,6 +47,11 @@ export opaque type ServerReference<T> = T;
 
 export type CallServerCallback = <A, T>(id: any, args: A) => Promise<T>;
 
+export type EncodeFormActionCallback = <A>(
+  id: any,
+  args: Promise<A>,
+) => ReactCustomFormAction;
+
 export type ServerReferenceId = any;
 
 const knownServerReferences: WeakMap<
@@ -454,7 +459,7 @@ function encodeFormData(reference: any): Thenable<FormData> {
   return thenable;
 }
 
-export function encodeFormAction(
+function defaultEncodeFormAction(
   this: any => Promise<any>,
   identifierPrefix: string,
 ): ReactCustomFormAction {
@@ -503,6 +508,25 @@ export function encodeFormAction(
   };
 }
 
+function customEncodeFormAction(
+  proxy: any => Promise<any>,
+  identifierPrefix: string,
+  encodeFormAction: EncodeFormActionCallback,
+): ReactCustomFormAction {
+  const reference = knownServerReferences.get(proxy);
+  if (!reference) {
+    throw new Error(
+      'Tried to encode a Server Action from a different instance than the encoder is from. ' +
+        'This is a bug in React.',
+    );
+  }
+  let boundPromise: Promise<Array<any>> = (reference.bound: any);
+  if (boundPromise === null) {
+    boundPromise = Promise.resolve([]);
+  }
+  return encodeFormAction(reference.id, boundPromise);
+}
+
 function isSignatureEqual(
   this: any => Promise<any>,
   referenceId: ServerReferenceId,
@@ -569,13 +593,27 @@ function isSignatureEqual(
 export function registerServerReference(
   proxy: any,
   reference: {id: ServerReferenceId, bound: null | Thenable<Array<any>>},
+  encodeFormAction: void | EncodeFormActionCallback,
 ) {
   // 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.
+    const $$FORM_ACTION =
+      encodeFormAction === undefined
+        ? defaultEncodeFormAction
+        : function (
+            this: any => Promise<any>,
+            identifierPrefix: string,
+          ): ReactCustomFormAction {
+            return customEncodeFormAction(
+              this,
+              identifierPrefix,
+              encodeFormAction,
+            );
+          };
     Object.defineProperties((proxy: any), {
-      $$FORM_ACTION: {value: encodeFormAction},
+      $$FORM_ACTION: {value: $$FORM_ACTION},
       $$IS_SIGNATURE_EQUAL: {value: isSignatureEqual},
       bind: {value: bind},
     });
@@ -587,7 +625,7 @@ export function registerServerReference(
 const FunctionBind = Function.prototype.bind;
 // $FlowFixMe[method-unbinding]
 const ArraySlice = Array.prototype.slice;
-function bind(this: Function) {
+function bind(this: Function): Function {
   // $FlowFixMe[unsupported-syntax]
   const newFn = FunctionBind.apply(this, arguments);
   const reference = knownServerReferences.get(this);
@@ -601,20 +639,31 @@ function bind(this: Function) {
     } else {
       boundPromise = Promise.resolve(args);
     }
-    registerServerReference(newFn, {id: reference.id, bound: boundPromise});
+    // 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((newFn: any), {
+        $$FORM_ACTION: {value: this.$$FORM_ACTION},
+        $$IS_SIGNATURE_EQUAL: {value: isSignatureEqual},
+        bind: {value: bind},
+      });
+    }
+    knownServerReferences.set(newFn, {id: reference.id, bound: boundPromise});
   }
   return newFn;
 }
 
 export function createServerReference<A: Iterable<any>, T>(
   id: ServerReferenceId,
   callServer: CallServerCallback,
+  encodeFormAction?: EncodeFormActionCallback,
 ): (...A) => Promise<T> {
   const proxy = function (): Promise<T> {
     // $FlowFixMe[method-unbinding]
     const args = Array.prototype.slice.call(arguments);
     return callServer(id, args);
   };
-  registerServerReference(proxy, {id, bound: null});
+  registerServerReference(proxy, {id, bound: null}, encodeFormAction);
   return proxy;
 }

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

@@ -38,6 +38,7 @@ function createResponseFromOptions(options: void | Options) {
     options && options.moduleBaseURL ? options.moduleBaseURL : '',
     null,
     options && options.callServer ? options.callServer : undefined,
+    undefined, // encodeFormAction
     undefined, // nonce
   );
 }

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

@@ -7,7 +7,7 @@
  * @flow
  */
 
-import type {Thenable} from 'shared/ReactTypes.js';
+import type {Thenable, ReactCustomFormAction} from 'shared/ReactTypes.js';
 
 import type {Response} from 'react-client/src/ReactFlightClient';
 
@@ -38,8 +38,14 @@ export function createServerReference<A: Iterable<any>, T>(
   return createServerReferenceImpl(id, noServerCall);
 }
 
+type EncodeFormActionCallback = <A>(
+  id: any,
+  args: Promise<A>,
+) => ReactCustomFormAction;
+
 export type Options = {
   nonce?: string,
+  encodeFormAction?: EncodeFormActionCallback,
 };
 
 function createFromNodeStream<T>(
@@ -52,6 +58,7 @@ function createFromNodeStream<T>(
     moduleRootPath,
     moduleBaseURL,
     noServerCall,
+    options ? options.encodeFormAction : undefined,
     options && typeof options.nonce === 'string' ? options.nonce : undefined,
   );
   stream.on('data', chunk => {

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

@@ -37,6 +37,7 @@ function createResponseFromOptions(options: void | Options) {
     null,
     null,
     options && options.callServer ? options.callServer : undefined,
+    undefined, // encodeFormAction
     undefined, // nonce
   );
 }

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

@@ -7,7 +7,7 @@
  * @flow
  */
 
-import type {Thenable} from 'shared/ReactTypes.js';
+import type {Thenable, ReactCustomFormAction} from 'shared/ReactTypes.js';
 
 import type {Response as FlightResponse} from 'react-client/src/ReactFlightClient';
 
@@ -51,16 +51,23 @@ export function createServerReference<A: Iterable<any>, T>(
   return createServerReferenceImpl(id, noServerCall);
 }
 
+type EncodeFormActionCallback = <A>(
+  id: any,
+  args: Promise<A>,
+) => ReactCustomFormAction;
+
 export type Options = {
   ssrManifest: SSRManifest,
   nonce?: string,
+  encodeFormAction?: EncodeFormActionCallback,
 };
 
 function createResponseFromOptions(options: Options) {
   return createResponse(
     options.ssrManifest.moduleMap,
     options.ssrManifest.moduleLoading,
     noServerCall,
+    options.encodeFormAction,
     typeof options.nonce === 'string' ? options.nonce : undefined,
   );
 }

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

@@ -7,7 +7,7 @@
  * @flow
  */
 
-import type {Thenable} from 'shared/ReactTypes.js';
+import type {Thenable, ReactCustomFormAction} from 'shared/ReactTypes.js';
 
 import type {Response} from 'react-client/src/ReactFlightClient';
 
@@ -40,9 +40,6 @@ function noServerCall() {
       'to pass data to Client Components instead.',
   );
 }
-export type Options = {
-  nonce?: string,
-};
 
 export function createServerReference<A: Iterable<any>, T>(
   id: any,
@@ -51,6 +48,16 @@ export function createServerReference<A: Iterable<any>, T>(
   return createServerReferenceImpl(id, noServerCall);
 }
 
+type EncodeFormActionCallback = <A>(
+  id: any,
+  args: Promise<A>,
+) => ReactCustomFormAction;
+
+export type Options = {
+  nonce?: string,
+  encodeFormAction?: EncodeFormActionCallback,
+};
+
 function createFromNodeStream<T>(
   stream: Readable,
   ssrManifest: SSRManifest,
@@ -60,6 +67,7 @@ function createFromNodeStream<T>(
     ssrManifest.moduleMap,
     ssrManifest.moduleLoading,
     noServerCall,
+    options ? options.encodeFormAction : undefined,
     options && typeof options.nonce === 'string' ? options.nonce : undefined,
   );
   stream.on('data', chunk => {

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

@@ -37,6 +37,7 @@ function createResponseFromOptions(options: void | Options) {
     null,
     null,
     options && options.callServer ? options.callServer : undefined,
+    undefined, // encodeFormAction
     undefined, // nonce
   );
 }

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

@@ -7,7 +7,7 @@
  * @flow
  */
 
-import type {Thenable} from 'shared/ReactTypes.js';
+import type {Thenable, ReactCustomFormAction} from 'shared/ReactTypes.js';
 
 import type {Response as FlightResponse} from 'react-client/src/ReactFlightClient';
 
@@ -51,16 +51,23 @@ export function createServerReference<A: Iterable<any>, T>(
   return createServerReferenceImpl(id, noServerCall);
 }
 
+type EncodeFormActionCallback = <A>(
+  id: any,
+  args: Promise<A>,
+) => ReactCustomFormAction;
+
 export type Options = {
   ssrManifest: SSRManifest,
   nonce?: string,
+  encodeFormAction?: EncodeFormActionCallback,
 };
 
 function createResponseFromOptions(options: Options) {
   return createResponse(
     options.ssrManifest.moduleMap,
     options.ssrManifest.moduleLoading,
     noServerCall,
+    options.encodeFormAction,
     typeof options.nonce === 'string' ? options.nonce : undefined,
   );
 }

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

@@ -7,7 +7,7 @@
  * @flow
  */
 
-import type {Thenable} from 'shared/ReactTypes.js';
+import type {Thenable, ReactCustomFormAction} from 'shared/ReactTypes.js';
 
 import type {Response} from 'react-client/src/ReactFlightClient';
 
@@ -48,8 +48,14 @@ export function createServerReference<A: Iterable<any>, T>(
   return createServerReferenceImpl(id, noServerCall);
 }
 
+type EncodeFormActionCallback = <A>(
+  id: any,
+  args: Promise<A>,
+) => ReactCustomFormAction;
+
 export type Options = {
   nonce?: string,
+  encodeFormAction?: EncodeFormActionCallback,
 };
 
 function createFromNodeStream<T>(
@@ -61,6 +67,7 @@ function createFromNodeStream<T>(
     ssrManifest.moduleMap,
     ssrManifest.moduleLoading,
     noServerCall,
+    options ? options.encodeFormAction : undefined,
     options && typeof options.nonce === 'string' ? options.nonce : undefined,
   );
   stream.on('data', chunk => {