[Fizz] Recover from errors thrown by progressive enhancement form generation (#28611)

This a follow up to #28564. It's alternative to #28609 which takes
#28610 into account.

It used to be possible to return JSX from an action with
`useActionState`.

```js
async function action(errors, payload) {
  "use server";
  try {
    ...
  } catch (x) {
    return <div>Error message</div>;
  }
}
```

```js
const [errors, formAction] = useActionState(action);
return <div>{errors}</div>;
```

Returning JSX from an action is itself not anything problematic. It's
that it also has to return the previous state to the action reducer
again that's the problem. When this happens we accidentally could
serialize an Element back to the server.

I fixed this in #28564 so it's now blocked if you don't have a temporary
reference set.

However, you can't have that for the progressive enhancement case. The
reply is eagerly encoded as part of the SSR render. Typically you
wouldn't have these in the initial state so the common case is that they
show up after the first POST back that yields an error and it's only in
the no-JS case where this happens so it's hard to discover.

As noted in #28609 there's a security implication with allowing elements
to be sent across this kind of payload, so we can't just make it work.

When an error happens during SSR our general policy is to try to recover
on the client instead. After all, SSR is mainly a perf optimization in
React terms and it's not primarily intended for a no JS solution.

This PR takes the approach that if we fail to generate the progressive
enhancement payload. I.e. if the serialization of previous state /
closures throw. Then we fallback to the replaying semantics just client
actions instead which will succeed.

The effect of this is that this pattern mostly just works:

- First render in the typical case doesn't have any JSX in it so it just
renders a progressive enhanced form.
- If JS fails to hydrate or you click early we do a form POST. If that
hits an error and it tries to render it using JSX, then the new page
will render successfully - however this time with a Replaying form
instead.
- If you try to submit the form again it'll have to be using JS.

Meaning if you use JSX as the error return value of form state and you
make a first attempt that fails, then no JS won't work because either
the first or second attempt has to hydrate.

We have ideas for potentially optimizing away serializing unused
arguments like if you don't actually use previous state which would also
solve it but it wouldn't cover all cases such as if it was deeply nested
in complex state.

Another approach that I considered was to poison the prev state if you
passed an element back but let it through to the action but if you try
to render the poisoned value, it wouldn't work. The downside of this is
when to error. Because in the progressive enhancement case it wouldn't
error early but when you actually try to invoke it at which point it
would be too late to fallback to client replaying. It would probably
have to always error even on the client which is unfortunate since this
mostly just works as long as it hydrates.

Author: sebmarkbage

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