[Fizz] Recover from errors thrown by progressive enhancement form generation

packages/react-dom-bindings/src/server/ReactFizzConfigDOM.js

@@ -1045,14 +1045,43 @@ function pushAdditionalFormField(
 
 function pushAdditionalFormFields(
   target: Array<Chunk | PrecomputedChunk>,
-  formData: null | FormData,
+  formData: void | null | FormData,
 ) {
-  if (formData !== null) {
+  if (formData != null) {
     // $FlowFixMe[prop-missing]: FormData has forEach.
     formData.forEach(pushAdditionalFormField, target);
   }
 }
 
+function getCustomFormFields(
+  resumableState: ResumableState,
+  formAction: any,
+): null | ReactCustomFormAction {
+  const customAction = formAction.$$FORM_ACTION;
+  if (typeof customAction === 'function') {
+    const prefix = makeFormFieldPrefix(resumableState);
+    try {
+      return formAction.$$FORM_ACTION(prefix);
+    } catch (x) {
+      if (typeof x === 'object' && x !== null && typeof x.then === 'function') {
+        // Rethrow suspense.
+        throw x;
+      }
+      // If we fail to encode the form action for progressive enhancement for some reason,
+      // fallback to trying replaying on the client instead of failing the page. It might
+      // work there.
+      if (__DEV__) {
+        // TODO: Should this be some kind of recoverable error?
+        console.error(
+          'Failed to serialize an action for progressive enhancement:\n%s',
+          x,
+        );
+      }
+    }
+  }
+  return null;
+}
+
 function pushFormActionAttribute(
   target: Array<Chunk | PrecomputedChunk>,
   resumableState: ResumableState,
@@ -1062,7 +1091,7 @@ function pushFormActionAttribute(
   formMethod: any,
   formTarget: any,
   name: any,
-): null | FormData {
+): void | null | FormData {
   let formData = null;
   if (enableFormActions && typeof formAction === 'function') {
     // Function form actions cannot control the form properties
@@ -1092,12 +1121,10 @@ function pushFormActionAttribute(
         );
       }
     }
-    const customAction: ReactCustomFormAction = formAction.$$FORM_ACTION;
-    if (typeof customAction === 'function') {
+    const customFields = getCustomFormFields(resumableState, formAction);
+    if (customFields !== null) {
       // This action has a custom progressive enhancement form that can submit the form
       // back to the server if it's invoked before hydration. Such as a Server Action.
-      const prefix = makeFormFieldPrefix(resumableState);
-      const customFields = formAction.$$FORM_ACTION(prefix);
       name = customFields.name;
       formAction = customFields.action || '';
       formEncType = customFields.encType;
@@ -1882,12 +1909,10 @@ function pushStartForm(
         );
       }
     }
-    const customAction: ReactCustomFormAction = formAction.$$FORM_ACTION;
-    if (typeof customAction === 'function') {
+    const customFields = getCustomFormFields(resumableState, formAction);
+    if (customFields !== null) {
       // This action has a custom progressive enhancement form that can submit the form
       // back to the server if it's invoked before hydration. Such as a Server Action.
-      const prefix = makeFormFieldPrefix(resumableState);
-      const customFields = formAction.$$FORM_ACTION(prefix);
       formAction = customFields.action || '';
       formEncType = customFields.encType;
       formMethod = customFields.method;

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

@@ -897,4 +897,77 @@ describe('ReactFlightDOMForm', () => {
 
     expect(form.action).toBe('http://localhost/permalink');
   });
+
+  // @gate enableFormActions
+  // @gate enableAsyncActions
+  it('useFormState can return JSX state during MPA form submission', async () => {
+    const serverAction = serverExports(
+      async function action(prevState, formData) {
+        return <div>error message</div>;
+      },
+    );
+
+    function Form({action}) {
+      const [errorMsg, dispatch] = useFormState(action, null);
+      return <form action={dispatch}>{errorMsg}</form>;
+    }
+
+    const FormRef = await clientExports(Form);
+
+    const rscStream = ReactServerDOMServer.renderToReadableStream(
+      <FormRef action={serverAction} />,
+      webpackMap,
+    );
+    const response = ReactServerDOMClient.createFromReadableStream(rscStream, {
+      ssrManifest: {
+        moduleMap: null,
+        moduleLoading: null,
+      },
+    });
+    const ssrStream = await ReactDOMServer.renderToReadableStream(response);
+    await readIntoContainer(ssrStream);
+
+    const form1 = container.getElementsByTagName('form')[0];
+    expect(form1.textContent).toBe('');
+
+    async function submitTheForm() {
+      const form = container.getElementsByTagName('form')[0];
+      const {formState} = await submit(form);
+
+      // Simulate an MPA form submission by resetting the container and
+      // rendering again.
+      container.innerHTML = '';
+
+      const postbackRscStream = ReactServerDOMServer.renderToReadableStream(
+        <FormRef action={serverAction} />,
+        webpackMap,
+      );
+      const postbackResponse = ReactServerDOMClient.createFromReadableStream(
+        postbackRscStream,
+        {
+          ssrManifest: {
+            moduleMap: null,
+            moduleLoading: null,
+          },
+        },
+      );
+      const postbackSsrStream = await ReactDOMServer.renderToReadableStream(
+        postbackResponse,
+        {formState: formState},
+      );
+      await readIntoContainer(postbackSsrStream);
+    }
+
+    await expect(submitTheForm).toErrorDev(
+      'Warning: Failed to serialize an action for progressive enhancement:\n' +
+        'Error: React Element cannot be passed to Server Functions from the Client without a temporary reference set. Pass a TemporaryReferenceSet to the options.\n' +
+        '  [<div/>]\n' +
+        '   ^^^^^^',
+    );
+
+    // The error message was returned as JSX.
+    const form2 = container.getElementsByTagName('form')[0];
+    expect(form2.textContent).toBe('error message');
+    expect(form2.firstChild.tagName).toBe('DIV');
+  });
 });