Repository: facebook/react
Pull request: [Fiber] Correctly handle replaying when hydrating

When hydrating if something suspends and then resolves in a microtask it is possible that React will resume the render without fully unwinding work in progress. This can cause hydration cursors to be offset and lead to hydration errors. This change adds a restore step when replaying HostComponent to ensure the hydration cursor is in the appropriate position when replaying.

fixes: #35210

Changed files:
filename: packages/react-dom/src/__tests__/ReactDOMServerPartialHydration-test.internal.js
patch: @@ -4063,4 +4063,170 @@ describe('ReactDOMServerPartialHydration', () => {
     expect(span.style.display).toBe('');
     expect(ref.current).toBe(span);
   });
+
+  // Regression for https://github.com/facebook/react/issues/35210 and other issues where lazy elements created in flight
+  // caused hydration issues b/c the replay pathway did not correctly reset the hydration cursor
+  it('Can hydrate even when lazy content resumes immediately inside a HostComponent', async () => {
+    let resolve;
+    const promise = new Promise(r => {
+      resolve = () => r({default: 'value'});
+    });
+
+    const lazyContent = React.lazy(() => {
+      Scheduler.log('Lazy initializer called');
+      return promise;
+    });
+
+    function App() {
+      return <label>{lazyContent}</label>;
+    }
+
+    // Server-rendered HTML
+    const container = document.createElement('div');
+    container.innerHTML = '<label>value</label>';
+
+    const hydrationErrors = [];
+
+    React.startTransition(() => {
+      ReactDOMClient.hydrateRoot(container, <App />, {
+        onRecoverableError(error) {
+          console.log('[DEBUG] hydration error:', error.message);
+          hydrationErrors.push(error.message);
+        },
+      });
+    });
+
+    await waitFor(['Lazy initializer called']);
+    resolve();
+    await waitForAll([]);
+
+    // Without the fix, hydration cursor is wrong and causes mismatch
+    expect(hydrationErrors).toEqual([]);
+    expect(container.innerHTML).toEqual('<label>value</label>');
+  });
+
+  it('Can hydrate even when lazy content resumes immediately inside a HostSingleton', async () => {
+    let resolve;
+    const promise = new Promise(r => {
+      resolve = () => r({default: <div>value</div>});
+    });
+
+    const lazyContent = React.lazy(() => {
+      Scheduler.log('Lazy initializer called');
+      return promise;
+    });
+
+    function App() {
+      return (
+        <html>
+          <body>{lazyContent}</body>
+        </html>
+      );
+    }
+
+    // Server-rendered HTML
+    document.body.innerHTML = '<div>value</div>';
+
+    const hydrationErrors = [];
+
+    React.startTransition(() => {
+      ReactDOMClient.hydrateRoot(document, <App />, {
+        onRecoverableError(error) {
+          console.log('[DEBUG] hydration error:', error.message);
+          hydrationErrors.push(error.message);
+        },
+      });
+    });
+
+    await waitFor(['Lazy initializer called']);
+    resolve();
+    await waitForAll([]);
+
+    expect(hydrationErrors).toEqual([]);
+    expect(document.documentElement.outerHTML).toEqual(
+      '<html><head></head><body><div>value</div></body></html>',
+    );
+  });
+
+  it('Can hydrate even when lazy content resumes immediately inside a Suspense', async () => {
+    let resolve;
+    const promise = new Promise(r => {
+      resolve = () => r({default: 'value'});
+    });
+
+    const lazyContent = React.lazy(() => {
+      Scheduler.log('Lazy initializer called');
+      return promise;
+    });
+
+    function App() {
+      return <Suspense>{lazyContent}</Suspense>;
+    }
+
+    // Server-rendered HTML
+    const container = document.createElement('div');
+    container.innerHTML = '<!--$-->value<!--/$-->';
+
+    const hydrationErrors = [];
+
+    let root;
+    React.startTransition(() => {
+      root = ReactDOMClient.hydrateRoot(container, <App />, {
+        onRecoverableError(error) {
+          console.log('[DEBUG] hydration error:', error.message);
+          hydrationErrors.push(error.message);
+        },
+      });
+    });
+
+    await waitFor(['Lazy initializer called']);
+    resolve();
+    await waitForAll([]);
+
+    expect(hydrationErrors).toEqual([]);
+    expect(container.innerHTML).toEqual('<!--$-->value<!--/$-->');
+    root.unmount();
+    expect(container.innerHTML).toEqual('<!--$--><!--/$-->');
+  });
+
+  it('Can hydrate even when lazy content resumes immediately inside an Activity', async () => {
+    let resolve;
+    const promise = new Promise(r => {
+      resolve = () => r({default: 'value'});
+    });
+
+    const lazyContent = React.lazy(() => {
+      Scheduler.log('Lazy initializer called');
+      return promise;
+    });
+
+    function App() {
+      return <Activity mode="visible">{lazyContent}</Activity>;
+    }
+
+    // Server-rendered HTML
+    const container = document.createElement('div');
+    container.innerHTML = '<!--&-->value<!--/&-->';
+
+    const hydrationErrors = [];
+
+    let root;
+    React.startTransition(() => {
+      root = ReactDOMClient.hydrateRoot(container, <App />, {
+        onRecoverableError(error) {
+          console.log('[DEBUG] hydration error:', error.message);
+          hydrationErrors.push(error.message);
+        },
+      });
+    });
+
+    await waitFor(['Lazy initializer called']);
+    resolve();
+    await waitForAll([]);
+
+    expect(hydrationErrors).toEqual([]);
+    expect(container.innerHTML).toEqual('<!--&-->value<!--/&-->');
+    root.unmount();
+    expect(container.innerHTML).toEqual('<!--&--><!--/&-->');
+  });
 });

filename: packages/react-reconciler/src/ReactFiberHydrationContext.js
patch: @@ -173,6 +173,7 @@ function enterHydrationState(fiber: Fiber): boolean {
   didSuspendOrErrorDEV = false;
   hydrationDiffRootDEV = null;
   rootOrSingletonContext = true;
+
   return true;
 }
 
@@ -834,6 +835,43 @@ function resetHydrationState(): void {
   didSuspendOrErrorDEV = false;
 }
 
+// Restore the hydration cursor when unwinding a HostComponent that already
+// claimed a DOM node. This is a fork of popHydrationState that does all the
+// same validity checks but restores the cursor to this fiber's DOM node
+// instead of advancing past it. It also does NOT clear unhydrated tail nodes
+// or throw on mismatches since we're unwinding, not completing.
+//
+// This is needed when replaySuspendedUnitOfWork calls unwindInterruptedWork
+// before re-running beginWork on the same fiber, or when throwAndUnwindWorkLoop
+// calls unwindWork on ancestor fibers.
+function popHydrationStateOnInterruptedWork(fiber: Fiber): void {
+  if (!supportsHydration) {
+    return;
+  }
+  if (fiber !== hydrationParentFiber) {
+    // We're deeper than the current hydration context, inside an inserted
+    // tree. Don't touch the cursor.
+    return;
+  }
+  if (!isHydrating) {
+    // If we're not currently hydrating but we're in a hydration context, then
+    // we were an insertion and now need to pop up to reenter hydration of our
+    // siblings. Same as popHydrationState.
+    popToNextHostParent(fiber);
+    isHydrating = true;
+    return;
+  }
+
+  // We're in a valid hydration context. Restore the cursor to this fiber's
+  // DOM node so that when beginWork re-runs, it can claim the same node.
+  // Unlike popHydrationState, we do NOT check for unhydrated tail nodes
+  // or advance the cursor - we're restoring, not completing.
+  popToNextHostParent(fiber);
+  if (fiber.tag === HostComponent && fiber.stateNode != null) {
+    nextHydratableInstance = fiber.stateNode;
+  }
+}
+
 export function upgradeHydrationErrorsToRecoverable(): Array<
   CapturedValue<mixed>,
 > | null {
@@ -905,6 +943,7 @@ export {
   reenterHydrationStateFromDehydratedActivityInstance,
   reenterHydrationStateFromDehydratedSuspenseInstance,
   resetHydrationState,
+  popHydrationStateOnInterruptedWork,
   claimHydratableSingleton,
   tryToClaimNextHydratableInstance,
   tryToClaimNextHydratableTextInstance,

filename: packages/react-reconciler/src/ReactFiberWorkLoop.js
patch: @@ -120,7 +120,10 @@ import {
 
 import {createWorkInProgress, resetWorkInProgress} from './ReactFiber';
 import {isRootDehydrated} from './ReactFiberShellHydration';
-import {getIsHydrating} from './ReactFiberHydrationContext';
+import {
+  getIsHydrating,
+  popHydrationStateOnInterruptedWork,
+} from './ReactFiberHydrationContext';
 import {
   NoMode,
   ProfileMode,
@@ -3109,6 +3112,10 @@ function replayBeginWork(unitOfWork: Fiber): null | Fiber {
       // promises that a host component might suspend on are definitely cached
       // because they are controlled by us. So don't bother.
       resetHooksOnUnwind(unitOfWork);
+      // We are about to retry this host component and need to ensure the hydration
+      // state is appropriate for hydrating this unit. Other fiber types hydrate differently
+      // and aren't reliant on the cursor positioning so this function is only for HostComponent
+      popHydrationStateOnInterruptedWork(unitOfWork);
       // Fallthrough to the next branch.
     }
     default: {