fix: React.use inside React.lazy-ed component on SSR (facebook#33941)

Author: hi-ogawa

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

@@ -6340,6 +6340,63 @@ describe('ReactDOMFizzServer', () => {
     expect(getVisibleChildren(container)).toEqual('Hi');
   });
 
+  it('should correctly handle different promises in React.use() across lazy components', async () => {
+    let promise1;
+    let promise2;
+    let promiseLazy;
+
+    function Component1() {
+      promise1 ??= new Promise(r => setTimeout(() => r('value1'), 50));
+      const data = React.use(promise1);
+      return (
+        <div>
+          {data}
+          <Component2Lazy />
+        </div>
+      );
+    }
+
+    function Component2() {
+      promise2 ??= new Promise(r => setTimeout(() => r('value2'), 50));
+      const data = React.use(promise2);
+      return <div>{data}</div>;
+    }
+
+    const Component2Lazy = React.lazy(async () => {
+      promiseLazy ??= new Promise(r => setTimeout(r, 50));
+      await promiseLazy;
+      return {default: Component2};
+    });
+
+    function App() {
+      return <Component1 />;
+    }
+
+    await act(async () => {
+      const {pipe} = renderToPipeableStream(<App />);
+      pipe(writable);
+    });
+
+    // Wait for promise to resolve
+    await act(async () => {
+      await promise1;
+    });
+    await act(async () => {
+      await promiseLazy;
+    });
+    await act(async () => {
+      await promise2;
+    });
+
+    // Verify both components received the correct values
+    expect(getVisibleChildren(container)).toEqual(
+      <div>
+        value1
+        <div>value2</div>
+      </div>,
+    );
+  });
+
   it('useActionState hydrates without a mismatch', async () => {
     // This is testing an implementation detail: useActionState emits comment
     // nodes into the SSR stream, so this checks that they are handled correctly

packages/react-server/src/ReactFizzServer.js

@@ -4153,7 +4153,10 @@ function renderNode(
         // $FlowFixMe[method-unbinding]
         if (typeof x.then === 'function') {
           const wakeable: Wakeable = (x: any);
-          const thenableState = getThenableStateAfterSuspending();
+          const thenableState =
+            thrownValue === SuspenseException
+              ? getThenableStateAfterSuspending()
+              : null;
           const newTask = spawnNewSuspendedReplayTask(
             request,
             // $FlowFixMe: Refined.
@@ -4186,7 +4189,10 @@ function renderNode(
           // performance but it can lead to stack overflows in extremely deep trees.
           // We do have the ability to create a trampoile if this happens which makes
           // this kind of zero-cost.
-          const thenableState = getThenableStateAfterSuspending();
+          const thenableState =
+            thrownValue === SuspenseException
+              ? getThenableStateAfterSuspending()
+              : null;
           const newTask = spawnNewSuspendedReplayTask(
             request,
             // $FlowFixMe: Refined.
@@ -4246,7 +4252,10 @@ function renderNode(
         // $FlowFixMe[method-unbinding]
         if (typeof x.then === 'function') {
           const wakeable: Wakeable = (x: any);
-          const thenableState = getThenableStateAfterSuspending();
+          const thenableState =
+            thrownValue === SuspenseException
+              ? getThenableStateAfterSuspending()
+              : null;
           const newTask = spawnNewSuspendedRenderTask(
             request,
             // $FlowFixMe: Refined.
@@ -4317,7 +4326,10 @@ function renderNode(
           // performance but it can lead to stack overflows in extremely deep trees.
           // We do have the ability to create a trampoile if this happens which makes
           // this kind of zero-cost.
-          const thenableState = getThenableStateAfterSuspending();
+          const thenableState =
+            thrownValue === SuspenseException
+              ? getThenableStateAfterSuspending()
+              : null;
           const newTask = spawnNewSuspendedRenderTask(
             request,
             // $FlowFixMe: Refined.
@@ -5233,7 +5245,10 @@ function retryRenderTask(
       if (typeof x.then === 'function') {
         // Something suspended again, let's pick it back up later.
         segment.status = PENDING;
-        task.thenableState = getThenableStateAfterSuspending();
+        task.thenableState =
+          thrownValue === SuspenseException
+            ? getThenableStateAfterSuspending()
+            : null;
         const ping = task.ping;
         // We've asserted that x is a thenable above
         (x: any).then(ping, ping);
@@ -5338,7 +5353,10 @@ function retryReplayTask(request: Request, task: ReplayTask): void {
         // Something suspended again, let's pick it back up later.
         const ping = task.ping;
         x.then(ping, ping);
-        task.thenableState = getThenableStateAfterSuspending();
+        task.thenableState =
+          thrownValue === SuspenseException
+            ? getThenableStateAfterSuspending()
+            : null;
         return;
       }
     }