From f20732a8e1bdd83802693b182eb1eba26627fb45 Mon Sep 17 00:00:00 2001
From: Sebastian Markbage <sebastian@calyptus.eu>
Date: Tue, 5 Mar 2024 19:52:33 -0500
Subject: [PATCH] Move tail hydration mismatch back to hydration context

Fixes an issue where no recoverable error is thrown in this case.
---
 ...DOMServerPartialHydration-test.internal.js |  1 +
 .../src/ReactFiberCompleteWork.js             | 19 ++++---------------
 .../src/ReactFiberHydrationContext.js         |  6 ------
 3 files changed, 5 insertions(+), 21 deletions(-)

diff --git a/packages/react-dom/src/__tests__/ReactDOMServerPartialHydration-test.internal.js b/packages/react-dom/src/__tests__/ReactDOMServerPartialHydration-test.internal.js
index 5b10d81f2220d..d9cfd7e5eee1a 100644
--- a/packages/react-dom/src/__tests__/ReactDOMServerPartialHydration-test.internal.js
+++ b/packages/react-dom/src/__tests__/ReactDOMServerPartialHydration-test.internal.js
@@ -538,6 +538,7 @@ describe('ReactDOMServerPartialHydration', () => {
     assertLog([
       'Server rendered',
       'Client rendered',
+      'Hydration failed because the initial UI does not match what was rendered on the server.',
       'There was an error while hydrating this Suspense boundary. ' +
         'Switched to client rendering.',
     ]);
diff --git a/packages/react-reconciler/src/ReactFiberCompleteWork.js b/packages/react-reconciler/src/ReactFiberCompleteWork.js
index a2daff05de72e..e15192e2854a2 100644
--- a/packages/react-reconciler/src/ReactFiberCompleteWork.js
+++ b/packages/react-reconciler/src/ReactFiberCompleteWork.js
@@ -141,11 +141,9 @@ import {
   prepareToHydrateHostInstance,
   prepareToHydrateHostTextInstance,
   prepareToHydrateHostSuspenseInstance,
-  warnIfUnhydratedTailNodes,
   popHydrationState,
   resetHydrationState,
   getIsHydrating,
-  hasUnhydratedTailNodes,
   upgradeHydrationErrorsToRecoverable,
 } from './ReactFiberHydrationContext';
 import {
@@ -866,18 +864,6 @@ function completeDehydratedSuspenseBoundary(
   workInProgress: Fiber,
   nextState: SuspenseState | null,
 ): boolean {
-  if (
-    hasUnhydratedTailNodes() &&
-    (workInProgress.mode & ConcurrentMode) !== NoMode &&
-    (workInProgress.flags & DidCapture) === NoFlags
-  ) {
-    warnIfUnhydratedTailNodes(workInProgress);
-    resetHydrationState();
-    workInProgress.flags |= ForceClientRender | DidCapture;
-
-    return false;
-  }
-
   const wasHydrated = popHydrationState(workInProgress);
 
   if (nextState !== null && nextState.dehydrated !== null) {
@@ -1337,7 +1323,6 @@ function completeWork(
       return null;
     }
     case SuspenseComponent: {
-      popSuspenseHandler(workInProgress);
       const nextState: null | SuspenseState = workInProgress.memoizedState;
 
       // Special path for dehydrated boundaries. We may eventually move this
@@ -1358,10 +1343,12 @@ function completeWork(
           );
         if (!fallthroughToNormalSuspensePath) {
           if (workInProgress.flags & ForceClientRender) {
+            popSuspenseHandler(workInProgress);
             // Special case. There were remaining unhydrated nodes. We treat
             // this as a mismatch. Revert to client rendering.
             return workInProgress;
           } else {
+            popSuspenseHandler(workInProgress);
             // Did not finish hydrating, either because this is the initial
             // render or because something suspended.
             return null;
@@ -1371,6 +1358,8 @@ function completeWork(
         // Continue with the normal Suspense path.
       }
 
+      popSuspenseHandler(workInProgress);
+
       if ((workInProgress.flags & DidCapture) !== NoFlags) {
         // Something suspended. Re-render with the fallback children.
         workInProgress.lanes = renderLanes;
diff --git a/packages/react-reconciler/src/ReactFiberHydrationContext.js b/packages/react-reconciler/src/ReactFiberHydrationContext.js
index 61be1aa39b940..4d59710c758e2 100644
--- a/packages/react-reconciler/src/ReactFiberHydrationContext.js
+++ b/packages/react-reconciler/src/ReactFiberHydrationContext.js
@@ -893,10 +893,6 @@ function popHydrationState(fiber: Fiber): boolean {
   return true;
 }
 
-function hasUnhydratedTailNodes(): boolean {
-  return isHydrating && nextHydratableInstance !== null;
-}
-
 function warnIfUnhydratedTailNodes(fiber: Fiber) {
   let nextInstance = nextHydratableInstance;
   while (nextInstance) {
@@ -952,6 +948,4 @@ export {
   prepareToHydrateHostTextInstance,
   prepareToHydrateHostSuspenseInstance,
   popHydrationState,
-  hasUnhydratedTailNodes,
-  warnIfUnhydratedTailNodes,
 };