Fix: useTransition after use gets stuck in pending state

packages/react-reconciler/src/ReactFiberHooks.js

@@ -1083,20 +1083,49 @@ function useThenable<T>(thenable: Thenable<T>): T {
     thenableState = createThenableState();
   }
   const result = trackUsedThenable(thenableState, thenable, index);
-  if (
-    currentlyRenderingFiber.alternate === null &&
-    (workInProgressHook === null
-      ? currentlyRenderingFiber.memoizedState === null
-      : workInProgressHook.next === null)
-  ) {
-    // Initial render, and either this is the first time the component is
-    // called, or there were no Hooks called after this use() the previous
-    // time (perhaps because it threw). Subsequent Hook calls should use the
-    // mount dispatcher.
+
+  // When something suspends with `use`, we replay the component with the
+  // "re-render" dispatcher instead of the "mount" or "update" dispatcher.
+  //
+  // But if there are additional hooks that occur after the `use` invocation
+  // that suspended, they wouldn't have been processed during the previous
+  // attempt. So after we invoke `use` again, we may need to switch from the
+  // "re-render" dispatcher back to the "mount" or "update" dispatcher. That's
+  // what the following logic accounts for.
+  //
+  // TODO: Theoretically this logic only needs to go into the rerender
+  // dispatcher. Could optimize, but probably not be worth it.
+
+  // This is the same logic as in updateWorkInProgressHook.
+  const workInProgressFiber = currentlyRenderingFiber;
+  const nextWorkInProgressHook =
+    workInProgressHook === null
+      ? // We're at the beginning of the list, so read from the first hook from
+        // the fiber.
+        workInProgressFiber.memoizedState
+      : workInProgressHook.next;
+
+  if (nextWorkInProgressHook !== null) {
+    // There are still hooks remaining from the previous attempt.
+  } else {
+    // There are no remaining hooks from the previous attempt. We're no longer
+    // in "re-render" mode. Switch to the normal mount or update dispatcher.
+    //
+    // This is the same as the logic in renderWithHooks, except we don't bother
+    // to track the hook types debug information in this case (sufficient to
+    // only do that when nothing suspends).
+    const currentFiber = workInProgressFiber.alternate;
     if (__DEV__) {
-      ReactSharedInternals.H = HooksDispatcherOnMountInDEV;
+      if (currentFiber !== null && currentFiber.memoizedState !== null) {
+        ReactSharedInternals.H = HooksDispatcherOnUpdateInDEV;
+      } else {
+        ReactSharedInternals.H = HooksDispatcherOnMountInDEV;
+      }
     } else {
-      ReactSharedInternals.H = HooksDispatcherOnMount;
+      ReactSharedInternals.H =
+        currentFiber === null || currentFiber.memoizedState === null
+          ? HooksDispatcherOnMount
+          : HooksDispatcherOnUpdate;
     }
   }
   return result;

packages/react-reconciler/src/__tests__/ReactUse-test.js

@@ -16,6 +16,7 @@ let act;
 let use;
 let useDebugValue;
 let useState;
+let useTransition;
 let useMemo;
 let useEffect;
 let Suspense;
@@ -38,6 +39,7 @@ describe('ReactUse', () => {
     use = React.use;
     useDebugValue = React.useDebugValue;
     useState = React.useState;
+    useTransition = React.useTransition;
     useMemo = React.useMemo;
     useEffect = React.useEffect;
     Suspense = React.Suspense;
@@ -1915,4 +1917,80 @@ describe('ReactUse', () => {
     assertLog(['Hi', 'World']);
     expect(root).toMatchRenderedOutput(<div>Hi World</div>);
   });
+
+  it(
+    'regression: does not get stuck in pending state after `use` suspends ' +
+      '(when `use` comes before all hooks)',
+    async () => {
+      // This is a regression test. The root cause was an issue where we failed to
+      // switch from the "re-render" dispatcher back to the "update" dispatcher
+      // after a `use` suspends and triggers a replay.
+      let update;
+      function App({promise}) {
+        const value = use(promise);
+
+        const [isPending, startLocalTransition] = useTransition();
+        update = () => {
+          startLocalTransition(() => {
+            root.render(<App promise={getAsyncText('Updated')} />);
+          });
+        };
+
+        return <Text text={value + (isPending ? ' (pending...)' : '')} />;
+      }
+
+      const root = ReactNoop.createRoot();
+      await act(() => {
+        root.render(<App promise={Promise.resolve('Initial')} />);
+      });
+      assertLog(['Initial']);
+      expect(root).toMatchRenderedOutput('Initial');
+
+      await act(() => update());
+      assertLog(['Async text requested [Updated]', 'Initial (pending...)']);
+
+      await act(() => resolveTextRequests('Updated'));
+      assertLog(['Updated']);
+      expect(root).toMatchRenderedOutput('Updated');
+    },
+  );
+
+  it(
+    'regression: does not get stuck in pending state after `use` suspends ' +
+      '(when `use` in in the middle of hook list)',
+    async () => {
+      // Same as previous test but `use` comes in between two hooks.
+      let update;
+      function App({promise}) {
+        // This hook is only here to test that `use` resumes correctly after
+        // suspended even if it comes in between other hooks.
+        useState(false);
+
+        const value = use(promise);
+
+        const [isPending, startLocalTransition] = useTransition();
+        update = () => {
+          startLocalTransition(() => {
+            root.render(<App promise={getAsyncText('Updated')} />);
+          });
+        };
+
+        return <Text text={value + (isPending ? ' (pending...)' : '')} />;
+      }
+
+      const root = ReactNoop.createRoot();
+      await act(() => {
+        root.render(<App promise={Promise.resolve('Initial')} />);
+      });
+      assertLog(['Initial']);
+      expect(root).toMatchRenderedOutput('Initial');
+
+      await act(() => update());
+      assertLog(['Async text requested [Updated]', 'Initial (pending...)']);
+
+      await act(() => resolveTextRequests('Updated'));
+      assertLog(['Updated']);
+      expect(root).toMatchRenderedOutput('Updated');
+    },
+  );
 });