[Fiber] Propagate context to dehydrated Suspense boundaries before bailout (#36594)

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

@@ -476,6 +476,115 @@ describe('ReactDOMFizzServer', () => {
     );
   });
 
+  it('does not produce a hydration mismatch when context updates in an already-hydrated ancestor before a streamed boundary hydrates', async () => {
+    // Regression test: Updating a context value in an already-hydrated ancestor
+    // (via a state update on a provider with referentially stable children)
+    // before a streamed Suspense child that reads that context finishes
+    // streaming used to cause a hydration mismatch. The streamed segment was
+    // rendered on the server with the initial context value, but the client
+    // hydrated it against the updated context. React should recover without a
+    // mismatch error.
+
+    const NumberContext = React.createContext(0);
+    let setNumberExternal = null;
+
+    function NumberProvider({children}) {
+      const [number, setNumber] = React.useState(0);
+      setNumberExternal = setNumber;
+      return (
+        <NumberContext.Provider value={number}>
+          {children}
+        </NumberContext.Provider>
+      );
+    }
+
+    function DisplayNumber() {
+      const number = React.useContext(NumberContext);
+      readText('display');
+      return <div>Number: {number}</div>;
+    }
+
+    function App() {
+      return (
+        <div>
+          <NumberProvider>
+            <Suspense fallback={<div>Loading...</div>}>
+              <DisplayNumber />
+            </Suspense>
+          </NumberProvider>
+        </div>
+      );
+    }
+
+    // The same context object is shared by the Fizz server and the client
+    // renderer in this single-process test, which trips the "multiple
+    // renderers" dev warning. That doesn't happen in real apps where server and
+    // client are separate processes, so filter it out here.
+    const realConsoleError = console.error;
+    console.error = (...args) => {
+      if (
+        typeof args[0] === 'string' &&
+        args[0].includes('Detected multiple renderers')
+      ) {
+        return;
+      }
+      realConsoleError.apply(console, args);
+    };
+
+    try {
+      await act(() => {
+        const {pipe} = renderToPipeableStream(<App />);
+        pipe(writable);
+      });
+
+      expect(getVisibleChildren(container)).toEqual(
+        <div>
+          <div>Loading...</div>
+        </div>,
+      );
+
+      const errors = [];
+      ReactDOMClient.hydrateRoot(container, <App />, {
+        onRecoverableError(error) {
+          errors.push(normalizeError(error.message));
+        },
+      });
+      await waitForAll([]);
+
+      // Shell hydrated, boundary still pending.
+      expect(getVisibleChildren(container)).toEqual(
+        <div>
+          <div>Loading...</div>
+        </div>,
+      );
+
+      // Update the context value in the already-hydrated ancestor BEFORE the
+      // streamed boundary content arrives and hydrates.
+      await clientAct(() => {
+        React.startTransition(() => {
+          setNumberExternal(1);
+        });
+      });
+
+      // Now the boundary content becomes available.
+      await act(() => {
+        resolveText('display');
+      });
+      await clientAct(async () => {});
+
+      // No hydration mismatch error should have been reported, and the boundary
+      // should reflect the updated context value.
+      expect(errors).toEqual([]);
+      expect(getVisibleChildren(container)).toEqual(
+        <div>
+          <div>Number: {'1'}</div>
+        </div>,
+      );
+    } finally {
+      console.error = realConsoleError;
+    }
+  });
+
   it('#23331: does not warn about hydration mismatches if something suspended in an earlier sibling', async () => {
     const makeApp = () => {
       let resolve;

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -3982,6 +3982,30 @@ function attemptEarlyBailoutIfNoScheduledUpdate(
       const state: SuspenseState | null = workInProgress.memoizedState;
       if (state !== null) {
         if (state.dehydrated !== null) {
+          // Before we bail out on a dehydrated boundary, we need to check
+          // whether a parent provider's context changed. The boundary's
+          // children only exist as server-rendered HTML, so normal propagation
+          // can't find context consumers inside it. If we bail out without
+          // propagating, a context change in an already-hydrated ancestor
+          // (e.g. a `useState` update in a provider with referentially stable
+          // children) will never be recorded on this boundary. When the
+          // streamed content later arrives and hydrates, it would read the
+          // updated context value and mismatch the server HTML. By propagating
+          // now, the boundary's childLanes records the change and
+          // `updateDehydratedSuspenseComponent` can recover instead of
+          // producing a hydration mismatch.
+          const contextChanged = lazilyPropagateParentContextChanges(
+            current,
+            workInProgress,
+            renderLanes,
+          );
+          if (contextChanged) {
+            return updateSuspenseComponent(
+              current,
+              workInProgress,
+              renderLanes,
+            );
+          }
           // We're not going to render the children, so this is just to maintain
           // push/pop symmetry
           pushPrimaryTreeSuspenseHandler(workInProgress);