Client implementation of useFormState (#27278)

This implements useFormState in Fiber. (It does not include any
progressive enhancement features; those will be added later.)

useFormState is a hook for tracking state produced by async actions. It
has a signature similar to useReducer, but instead of a reducer, it
accepts an async action function.

```js
async function action(prevState, payload) {
  // ..
}
const [state, dispatch] = useFormState(action, initialState)
```

Calling dispatch runs the async action and updates the state to the
returned value.

Async actions run before React's render cycle, so unlike reducers, they
can contain arbitrary side effects.

packages/react-dom-bindings/src/shared/ReactDOMFormActions.js

@@ -76,7 +76,7 @@ export function useFormStatus(): FormStatus {
 }
 
 export function useFormState<S, P>(
-  action: (S, P) => S,
+  action: (S, P) => Promise<S>,
   initialState: S,
   url?: string,
 ): [S, (P) => void] {

packages/react-dom/src/__tests__/ReactDOMForm-test.js

@@ -35,6 +35,7 @@ describe('ReactDOMForm', () => {
   let ReactDOMClient;
   let Scheduler;
   let assertLog;
+  let waitForThrow;
   let useState;
   let Suspense;
   let startTransition;
@@ -50,6 +51,7 @@ describe('ReactDOMForm', () => {
     Scheduler = require('scheduler');
     act = require('internal-test-utils').act;
     assertLog = require('internal-test-utils').assertLog;
+    waitForThrow = require('internal-test-utils').waitForThrow;
     useState = React.useState;
     Suspense = React.Suspense;
     startTransition = React.startTransition;
@@ -974,15 +976,28 @@ describe('ReactDOMForm', () => {
 
   // @gate enableFormActions
   // @gate enableAsyncActions
-  test('useFormState exists', async () => {
-    // TODO: Not yet implemented. This just tests that the API is wired up.
+  test('useFormState updates state asynchronously and queues multiple actions', async () => {
+    let actionCounter = 0;
+    async function action(state, type) {
+      actionCounter++;
 
-    async function action(state) {
-      return state;
+      Scheduler.log(`Async action started [${actionCounter}]`);
+      await getText(`Wait [${actionCounter}]`);
+
+      switch (type) {
+        case 'increment':
+          return state + 1;
+        case 'decrement':
+          return state - 1;
+        default:
+          return state;
+      }
     }
 
+    let dispatch;
     function App() {
-      const [state] = useFormState(action, 0);
+      const [state, _dispatch] = useFormState(action, 0);
+      dispatch = _dispatch;
       return <Text text={state} />;
     }
 
@@ -990,5 +1005,108 @@ describe('ReactDOMForm', () => {
     await act(() => root.render(<App />));
     assertLog([0]);
     expect(container.textContent).toBe('0');
+
+    await act(() => dispatch('increment'));
+    assertLog(['Async action started [1]']);
+    expect(container.textContent).toBe('0');
+
+    // Dispatch a few more actions. None of these will start until the previous
+    // one finishes.
+    await act(() => dispatch('increment'));
+    await act(() => dispatch('decrement'));
+    await act(() => dispatch('increment'));
+    assertLog([]);
+
+    // Each action starts as soon as the previous one finishes.
+    // NOTE: React does not render in between these actions because they all
+    // update the same queue, which means they get entangled together. This is
+    // intentional behavior.
+    await act(() => resolveText('Wait [1]'));
+    assertLog(['Async action started [2]']);
+    await act(() => resolveText('Wait [2]'));
+    assertLog(['Async action started [3]']);
+    await act(() => resolveText('Wait [3]'));
+    assertLog(['Async action started [4]']);
+    await act(() => resolveText('Wait [4]'));
+
+    // Finally the last action finishes and we can render the result.
+    assertLog([2]);
+    expect(container.textContent).toBe('2');
+  });
+
+  // @gate enableFormActions
+  // @gate enableAsyncActions
+  test('useFormState supports inline actions', async () => {
+    let increment;
+    function App({stepSize}) {
+      const [state, dispatch] = useFormState(async prevState => {
+        return prevState + stepSize;
+      }, 0);
+      increment = dispatch;
+      return <Text text={state} />;
+    }
+
+    // Initial render
+    const root = ReactDOMClient.createRoot(container);
+    await act(() => root.render(<App stepSize={1} />));
+    assertLog([0]);
+
+    // Perform an action. This will increase the state by 1, as defined by the
+    // stepSize prop.
+    await act(() => increment());
+    assertLog([1]);
+
+    // Now increase the stepSize prop to 10. Subsequent steps will increase
+    // by this amount.
+    await act(() => root.render(<App stepSize={10} />));
+    assertLog([1]);
+
+    // Increment again. The state should increase by 10.
+    await act(() => increment());
+    assertLog([11]);
+  });
+
+  // @gate enableFormActions
+  // @gate enableAsyncActions
+  test('useFormState: dispatch throws if called during render', async () => {
+    function App() {
+      const [state, dispatch] = useFormState(async () => {}, 0);
+      dispatch();
+      return <Text text={state} />;
+    }
+
+    const root = ReactDOMClient.createRoot(container);
+    await act(async () => {
+      root.render(<App />);
+      await waitForThrow('Cannot update form state while rendering.');
+    });
+  });
+
+  // @gate enableFormActions
+  // @gate enableAsyncActions
+  test('useFormState: warns if action is not async', async () => {
+    let dispatch;
+    function App() {
+      const [state, _dispatch] = useFormState(() => {}, 0);
+      dispatch = _dispatch;
+      return <Text text={state} />;
+    }
+
+    const root = ReactDOMClient.createRoot(container);
+    await act(async () => {
+      root.render(<App />);
+    });
+    assertLog([0]);
+
+    expect(() => {
+      // This throws because React expects the action to return a promise.
+      expect(() => dispatch()).toThrow('Cannot read properties of undefined');
+    }).toErrorDev(
+      [
+        // In dev we also log a warning.
+        'The action passed to useFormState must be an async function',
+      ],
+      {withoutStack: true},
+    );
   });
 });

packages/react-reconciler/src/ReactFiberAsyncAction.js

@@ -34,97 +34,108 @@ let currentEntangledPendingCount: number = 0;
 let currentEntangledLane: Lane = NoLane;
 
 export function requestAsyncActionContext<S>(
-  actionReturnValue: mixed,
-  finishedState: S,
-): Thenable<S> | S {
-  if (
-    actionReturnValue !== null &&
-    typeof actionReturnValue === 'object' &&
-    typeof actionReturnValue.then === 'function'
-  ) {
-    // This is an async action.
-    //
-    // Return a thenable that resolves once the action scope (i.e. the async
-    // function passed to startTransition) has finished running.
+  actionReturnValue: Thenable<mixed>,
+  // If this is provided, this resulting thenable resolves to this value instead
+  // of the return value of the action. This is a perf trick to avoid composing
+  // an extra async function.
+  overrideReturnValue: S | null,
+): Thenable<S> {
+  // This is an async action.
+  //
+  // Return a thenable that resolves once the action scope (i.e. the async
+  // function passed to startTransition) has finished running.
 
-    const thenable: Thenable<mixed> = (actionReturnValue: any);
-    let entangledListeners;
-    if (currentEntangledListeners === null) {
-      // There's no outer async action scope. Create a new one.
-      entangledListeners = currentEntangledListeners = [];
-      currentEntangledPendingCount = 0;
-      currentEntangledLane = requestTransitionLane();
-    } else {
-      entangledListeners = currentEntangledListeners;
-    }
-
-    currentEntangledPendingCount++;
-    let resultStatus = 'pending';
-    let rejectedReason;
-    thenable.then(
-      () => {
-        resultStatus = 'fulfilled';
-        pingEngtangledActionScope();
-      },
-      error => {
-        resultStatus = 'rejected';
-        rejectedReason = error;
-        pingEngtangledActionScope();
-      },
-    );
-
-    // Create a thenable that represents the result of this action, but doesn't
-    // resolve until the entire entangled scope has finished.
-    //
-    // Expressed using promises:
-    //   const [thisResult] = await Promise.all([thisAction, entangledAction]);
-    //   return thisResult;
-    const resultThenable = createResultThenable<S>(entangledListeners);
-
-    // Attach a listener to fill in the result.
-    entangledListeners.push(() => {
-      switch (resultStatus) {
-        case 'fulfilled': {
-          const fulfilledThenable: FulfilledThenable<S> = (resultThenable: any);
-          fulfilledThenable.status = 'fulfilled';
-          fulfilledThenable.value = finishedState;
-          break;
-        }
-        case 'rejected': {
-          const rejectedThenable: RejectedThenable<S> = (resultThenable: any);
-          rejectedThenable.status = 'rejected';
-          rejectedThenable.reason = rejectedReason;
-          break;
-        }
-        case 'pending':
-        default: {
-          // The listener above should have been called first, so `resultStatus`
-          // should already be set to the correct value.
-          throw new Error(
-            'Thenable should have already resolved. This ' +
-              'is a bug in React.',
-          );
-        }
-      }
-    });
-
-    return resultThenable;
+  const thenable: Thenable<S> = (actionReturnValue: any);
+  let entangledListeners;
+  if (currentEntangledListeners === null) {
+    // There's no outer async action scope. Create a new one.
+    entangledListeners = currentEntangledListeners = [];
+    currentEntangledPendingCount = 0;
+    currentEntangledLane = requestTransitionLane();
   } else {
-    // This is not an async action, but it may be part of an outer async action.
-    if (currentEntangledListeners === null) {
-      return finishedState;
-    } else {
-      // Return a thenable that does not resolve until the entangled actions
-      // have finished.
-      const entangledListeners = currentEntangledListeners;
-      const resultThenable = createResultThenable<S>(entangledListeners);
-      entangledListeners.push(() => {
+    entangledListeners = currentEntangledListeners;
+  }
+
+  currentEntangledPendingCount++;
+
+  // Create a thenable that represents the result of this action, but doesn't
+  // resolve until the entire entangled scope has finished.
+  //
+  // Expressed using promises:
+  //   const [thisResult] = await Promise.all([thisAction, entangledAction]);
+  //   return thisResult;
+  const resultThenable = createResultThenable<S>(entangledListeners);
+
+  let resultStatus = 'pending';
+  let resultValue;
+  let rejectedReason;
+  thenable.then(
+    (value: S) => {
+      resultStatus = 'fulfilled';
+      resultValue = overrideReturnValue !== null ? overrideReturnValue : value;
+      pingEngtangledActionScope();
+    },
+    error => {
+      resultStatus = 'rejected';
+      rejectedReason = error;
+      pingEngtangledActionScope();
+    },
+  );
+
+  // Attach a listener to fill in the result.
+  entangledListeners.push(() => {
+    switch (resultStatus) {
+      case 'fulfilled': {
         const fulfilledThenable: FulfilledThenable<S> = (resultThenable: any);
         fulfilledThenable.status = 'fulfilled';
-        fulfilledThenable.value = finishedState;
-      });
-      return resultThenable;
+        fulfilledThenable.value = resultValue;
+        break;
+      }
+      case 'rejected': {
+        const rejectedThenable: RejectedThenable<S> = (resultThenable: any);
+        rejectedThenable.status = 'rejected';
+        rejectedThenable.reason = rejectedReason;
+        break;
+      }
+      case 'pending':
+      default: {
+        // The listener above should have been called first, so `resultStatus`
+        // should already be set to the correct value.
+        throw new Error(
+          'Thenable should have already resolved. This ' + 'is a bug in React.',
+        );
+      }
     }
+  });
+
+  return resultThenable;
+}
+
+export function requestSyncActionContext<S>(
+  actionReturnValue: mixed,
+  // If this is provided, this resulting thenable resolves to this value instead
+  // of the return value of the action. This is a perf trick to avoid composing
+  // an extra async function.
+  overrideReturnValue: S | null,
+): Thenable<S> | S {
+  const resultValue: S =
+    overrideReturnValue !== null
+      ? overrideReturnValue
+      : (actionReturnValue: any);
+  // This is not an async action, but it may be part of an outer async action.
+  if (currentEntangledListeners === null) {
+    return resultValue;
+  } else {
+    // Return a thenable that does not resolve until the entangled actions
+    // have finished.
+    const entangledListeners = currentEntangledListeners;
+    const resultThenable = createResultThenable<S>(entangledListeners);
+    entangledListeners.push(() => {
+      const fulfilledThenable: FulfilledThenable<S> = (resultThenable: any);
+      fulfilledThenable.status = 'fulfilled';
+      fulfilledThenable.value = resultValue;
+    });
+    return resultThenable;
   }
 }
 

packages/react-reconciler/src/ReactFiberHooks.js

@@ -137,7 +137,10 @@ import {
 } from './ReactFiberThenable';
 import type {ThenableState} from './ReactFiberThenable';
 import type {BatchConfigTransition} from './ReactFiberTracingMarkerComponent';
-import {requestAsyncActionContext} from './ReactFiberAsyncAction';
+import {
+  requestAsyncActionContext,
+  requestSyncActionContext,
+} from './ReactFiberAsyncAction';
 import {HostTransitionContext} from './ReactFiberHostContext';
 import {requestTransitionLane} from './ReactFiberRootScheduler';
 
@@ -1854,35 +1857,304 @@ function rerenderOptimistic<S, A>(
   return [passthrough, dispatch];
 }
 
-function TODO_formStateDispatch() {
-  throw new Error('Not implemented.');
+// useFormState actions run sequentially, because each action receives the
+// previous state as an argument. We store pending actions on a queue.
+type FormStateActionQueue<S, P> = {
+  // This is the most recent state returned from an action. It's updated as
+  // soon as the action finishes running.
+  state: S,
+  // A stable dispatch method, passed to the user.
+  dispatch: Dispatch<P>,
+  // This is the most recent action function that was rendered. It's updated
+  // during the commit phase.
+  action: (S, P) => Promise<S>,
+  // This is a circular linked list of pending action payloads. It incudes the
+  // action that is currently running.
+  pending: FormStateActionQueueNode<P> | null,
+};
+
+type FormStateActionQueueNode<P> = {
+  payload: P,
+  // This is never null because it's part of a circular linked list.
+  next: FormStateActionQueueNode<P>,
+};
+
+function dispatchFormState<S, P>(
+  fiber: Fiber,
+  actionQueue: FormStateActionQueue<S, P>,
+  setState: Dispatch<Thenable<S>>,
+  payload: P,
+): void {
+  if (isRenderPhaseUpdate(fiber)) {
+    throw new Error('Cannot update form state while rendering.');
+  }
+  const last = actionQueue.pending;
+  if (last === null) {
+    // There are no pending actions; this is the first one. We can run
+    // it immediately.
+    const newLast: FormStateActionQueueNode<P> = {
+      payload,
+      next: (null: any), // circular
+    };
+    newLast.next = actionQueue.pending = newLast;
+
+    runFormStateAction(actionQueue, setState, payload);
+  } else {
+    // There's already an action running. Add to the queue.
+    const first = last.next;
+    const newLast: FormStateActionQueueNode<P> = {
+      payload,
+      next: first,
+    };
+    last.next = newLast;
+  }
+}
+
+function runFormStateAction<S, P>(
+  actionQueue: FormStateActionQueue<S, P>,
+  setState: Dispatch<Thenable<S>>,
+  payload: P,
+) {
+  const action = actionQueue.action;
+  const prevState = actionQueue.state;
+
+  // This is a fork of startTransition
+  const prevTransition = ReactCurrentBatchConfig.transition;
+  ReactCurrentBatchConfig.transition = ({}: BatchConfigTransition);
+  const currentTransition = ReactCurrentBatchConfig.transition;
+  if (__DEV__) {
+    ReactCurrentBatchConfig.transition._updatedFibers = new Set();
+  }
+  try {
+    const promise = action(prevState, payload);
+
+    if (__DEV__) {
+      if (
+        promise === null ||
+        typeof promise !== 'object' ||
+        typeof (promise: any).then !== 'function'
+      ) {
+        console.error(
+          'The action passed to useFormState must be an async function.',
+        );
+      }
+    }
+
+    // Attach a listener to read the return state of the action. As soon as this
+    // resolves, we can run the next action in the sequence.
+    promise.then(
+      (nextState: S) => {
+        actionQueue.state = nextState;
+        finishRunningFormStateAction(actionQueue, setState);
+      },
+      () => finishRunningFormStateAction(actionQueue, setState),
+    );
+
+    // Create a thenable that resolves once the current async action scope has
+    // finished. Then stash that thenable in state. We'll unwrap it with the
+    // `use` algorithm during render. This is the same logic used
+    // by startTransition.
+    const entangledThenable: Thenable<S> = requestAsyncActionContext(
+      promise,
+      null,
+    );
+    setState(entangledThenable);
+  } finally {
+    ReactCurrentBatchConfig.transition = prevTransition;
+
+    if (__DEV__) {
+      if (prevTransition === null && currentTransition._updatedFibers) {
+        const updatedFibersCount = currentTransition._updatedFibers.size;
+        currentTransition._updatedFibers.clear();
+        if (updatedFibersCount > 10) {
+          console.warn(
+            'Detected a large number of updates inside startTransition. ' +
+              'If this is due to a subscription please re-write it to use React provided hooks. ' +
+              'Otherwise concurrent mode guarantees are off the table.',
+          );
+        }
+      }
+    }
+  }
+}
+
+function finishRunningFormStateAction<S, P>(
+  actionQueue: FormStateActionQueue<S, P>,
+  setState: Dispatch<Thenable<S>>,
+) {
+  // The action finished running. Pop it from the queue and run the next pending
+  // action, if there are any.
+  const last = actionQueue.pending;
+  if (last !== null) {
+    const first = last.next;
+    if (first === last) {
+      // This was the last action in the queue.
+      actionQueue.pending = null;
+    } else {
+      // Remove the first node from the circular queue.
+      const next = first.next;
+      last.next = next;
+
+      // Run the next action.
+      runFormStateAction(actionQueue, setState, next.payload);
+    }
+  }
+}
+
+function formStateReducer<S>(oldState: S, newState: S): S {
+  return newState;
 }
 
 function mountFormState<S, P>(
-  action: (S, P) => S,
+  action: (S, P) => Promise<S>,
   initialState: S,
   url?: string,
 ): [S, (P) => void] {
-  // TODO: Not yet implemented
-  return [initialState, TODO_formStateDispatch];
+  // State hook. The state is stored in a thenable which is then unwrapped by
+  // the `use` algorithm during render.
+  const stateHook = mountWorkInProgressHook();
+  stateHook.memoizedState = stateHook.baseState = {
+    status: 'fulfilled',
+    value: initialState,
+  };
+  const stateQueue: UpdateQueue<Thenable<S>, Thenable<S>> = {
+    pending: null,
+    lanes: NoLanes,
+    dispatch: null,
+    lastRenderedReducer: formStateReducer,
+    lastRenderedState: (initialState: any),
+  };
+  stateHook.queue = stateQueue;
+  const setState: Dispatch<Thenable<S>> = (dispatchSetState.bind(
+    null,
+    currentlyRenderingFiber,
+    stateQueue,
+  ): any);
+  stateQueue.dispatch = setState;
+
+  // Action queue hook. This is used to queue pending actions. The queue is
+  // shared between all instances of the hook. Similar to a regular state queue,
+  // but different because the actions are run sequentially, and they run in
+  // an event instead of during render.
+  const actionQueueHook = mountWorkInProgressHook();
+  const actionQueue: FormStateActionQueue<S, P> = {
+    state: initialState,
+    dispatch: (null: any), // circular
+    action,
+    pending: null,
+  };
+  actionQueueHook.queue = actionQueue;
+  const dispatch = dispatchFormState.bind(
+    null,
+    currentlyRenderingFiber,
+    actionQueue,
+    setState,
+  );
+  actionQueue.dispatch = dispatch;
+
+  // Stash the action function on the memoized state of the hook. We'll use this
+  // to detect when the action function changes so we can update it in
+  // an effect.
+  actionQueueHook.memoizedState = action;
+
+  return [initialState, dispatch];
 }
 
 function updateFormState<S, P>(
-  action: (S, P) => S,
+  action: (S, P) => Promise<S>,
   initialState: S,
   url?: string,
 ): [S, (P) => void] {
-  // TODO: Not yet implemented
-  return [initialState, TODO_formStateDispatch];
+  const stateHook = updateWorkInProgressHook();
+  const currentStateHook = ((currentHook: any): Hook);
+  return updateFormStateImpl(
+    stateHook,
+    currentStateHook,
+    action,
+    initialState,
+    url,
+  );
+}
+
+function updateFormStateImpl<S, P>(
+  stateHook: Hook,
+  currentStateHook: Hook,
+  action: (S, P) => Promise<S>,
+  initialState: S,
+  url?: string,
+): [S, (P) => void] {
+  const [thenable] = updateReducerImpl<Thenable<S>, Thenable<S>>(
+    stateHook,
+    currentStateHook,
+    formStateReducer,
+  );
+
+  // This will suspend until the action finishes.
+  const state = useThenable(thenable);
+
+  const actionQueueHook = updateWorkInProgressHook();
+  const actionQueue = actionQueueHook.queue;
+  const dispatch = actionQueue.dispatch;
+
+  // Check if a new action was passed. If so, update it in an effect.
+  const prevAction = actionQueueHook.memoizedState;
+  if (action !== prevAction) {
+    currentlyRenderingFiber.flags |= PassiveEffect;
+    pushEffect(
+      HookHasEffect | HookPassive,
+      formStateActionEffect.bind(null, actionQueue, action),
+      createEffectInstance(),
+      null,
+    );
+  }
+
+  return [state, dispatch];
+}
+
+function formStateActionEffect<S, P>(
+  actionQueue: FormStateActionQueue<S, P>,
+  action: (S, P) => Promise<S>,
+): void {
+  actionQueue.action = action;
 }
 
 function rerenderFormState<S, P>(
-  action: (S, P) => S,
+  action: (S, P) => Promise<S>,
   initialState: S,
   url?: string,
 ): [S, (P) => void] {
-  // TODO: Not yet implemented
-  return [initialState, TODO_formStateDispatch];
+  // Unlike useState, useFormState doesn't support render phase updates.
+  // Also unlike useState, we need to replay all pending updates again in case
+  // the passthrough value changed.
+  //
+  // So instead of a forked re-render implementation that knows how to handle
+  // render phase udpates, we can use the same implementation as during a
+  // regular mount or update.
+  const stateHook = updateWorkInProgressHook();
+  const currentStateHook = currentHook;
+
+  if (currentStateHook !== null) {
+    // This is an update. Process the update queue.
+    return updateFormStateImpl(
+      stateHook,
+      currentStateHook,
+      action,
+      initialState,
+      url,
+    );
+  }
+
+  // This is a mount. No updates to process.
+  const state = stateHook.memoizedState;
+
+  const actionQueueHook = updateWorkInProgressHook();
+  const actionQueue = actionQueueHook.queue;
+  const dispatch = actionQueue.dispatch;
+
+  // This may have changed during the rerender.
+  actionQueueHook.memoizedState = action;
+
+  return [state, dispatch];
 }
 
 function pushEffect(
@@ -2459,15 +2731,37 @@ function startTransition<S>(
     if (enableAsyncActions) {
       const returnValue = callback();
 
-      // This is either `finishedState` or a thenable that resolves to
-      // `finishedState`, depending on whether the action scope is an async
-      // function. In the async case, the resulting render will suspend until
-      // the async action scope has finished.
-      const maybeThenable = requestAsyncActionContext(
-        returnValue,
-        finishedState,
-      );
-      dispatchSetState(fiber, queue, maybeThenable);
+      // Check if we're inside an async action scope. If so, we'll entangle
+      // this new action with the existing scope.
+      //
+      // If we're not already inside an async action scope, and this action is
+      // async, then we'll create a new async scope.
+      //
+      // In the async case, the resulting render will suspend until the async
+      // action scope has finished.
+      if (
+        returnValue !== null &&
+        typeof returnValue === 'object' &&
+        typeof returnValue.then === 'function'
+      ) {
+        const thenable = ((returnValue: any): Thenable<mixed>);
+        // This is a thenable that resolves to `finishedState` once the async
+        // action scope has finished.
+        const entangledResult = requestAsyncActionContext(
+          thenable,
+          finishedState,
+        );
+        dispatchSetState(fiber, queue, entangledResult);
+      } else {
+        // This is either `finishedState` or a thenable that resolves to
+        // `finishedState`, depending on whether we're inside an async
+        // action scope.
+        const entangledResult = requestSyncActionContext(
+          returnValue,
+          finishedState,
+        );
+        dispatchSetState(fiber, queue, entangledResult);
+      }
     } else {
       // Async actions are not enabled.
       dispatchSetState(fiber, queue, finishedState);
@@ -3332,7 +3626,7 @@ if (__DEV__) {
       useHostTransitionStatus;
     (HooksDispatcherOnMountInDEV: Dispatcher).useFormState =
       function useFormState<S, P>(
-        action: (S, P) => S,
+        action: (S, P) => Promise<S>,
         initialState: S,
         url?: string,
       ): [S, (P) => void] {
@@ -3502,7 +3796,7 @@ if (__DEV__) {
       useHostTransitionStatus;
     (HooksDispatcherOnMountWithHookTypesInDEV: Dispatcher).useFormState =
       function useFormState<S, P>(
-        action: (S, P) => S,
+        action: (S, P) => Promise<S>,
         initialState: S,
         url?: string,
       ): [S, (P) => void] {
@@ -3674,7 +3968,7 @@ if (__DEV__) {
       useHostTransitionStatus;
     (HooksDispatcherOnUpdateInDEV: Dispatcher).useFormState =
       function useFormState<S, P>(
-        action: (S, P) => S,
+        action: (S, P) => Promise<S>,
         initialState: S,
         url?: string,
       ): [S, (P) => void] {
@@ -3846,7 +4140,7 @@ if (__DEV__) {
       useHostTransitionStatus;
     (HooksDispatcherOnRerenderInDEV: Dispatcher).useFormState =
       function useFormState<S, P>(
-        action: (S, P) => S,
+        action: (S, P) => Promise<S>,
         initialState: S,
         url?: string,
       ): [S, (P) => void] {
@@ -4039,7 +4333,7 @@ if (__DEV__) {
       useHostTransitionStatus;
     (InvalidNestedHooksDispatcherOnMountInDEV: Dispatcher).useFormState =
       function useFormState<S, P>(
-        action: (S, P) => S,
+        action: (S, P) => Promise<S>,
         initialState: S,
         url?: string,
       ): [S, (P) => void] {
@@ -4237,7 +4531,7 @@ if (__DEV__) {
       useHostTransitionStatus;
     (InvalidNestedHooksDispatcherOnUpdateInDEV: Dispatcher).useFormState =
       function useFormState<S, P>(
-        action: (S, P) => S,
+        action: (S, P) => Promise<S>,
         initialState: S,
         url?: string,
       ): [S, (P) => void] {
@@ -4435,7 +4729,7 @@ if (__DEV__) {
       useHostTransitionStatus;
     (InvalidNestedHooksDispatcherOnRerenderInDEV: Dispatcher).useFormState =
       function useFormState<S, P>(
-        action: (S, P) => S,
+        action: (S, P) => Promise<S>,
         initialState: S,
         url?: string,
       ): [S, (P) => void] {

packages/react-reconciler/src/ReactInternalTypes.js

@@ -415,7 +415,7 @@ export type Dispatcher = {
     reducer: ?(S, A) => S,
   ) => [S, (A) => void],
   useFormState?: <S, P>(
-    action: (S, P) => S,
+    action: (S, P) => Promise<S>,
     initialState: S,
     url?: string,
   ) => [S, (P) => void],

packages/react-server/src/ReactFizzHooks.js

@@ -555,7 +555,7 @@ function useOptimistic<S, A>(
 }
 
 function useFormState<S, P>(
-  action: (S, P) => S,
+  action: (S, P) => Promise<S>,
   initialState: S,
   url?: string,
 ): [S, (P) => void] {