Fix context propagation through suspended Suspense boundaries (facebook#35839)

When a Suspense boundary suspends during initial mount, the primary
children's fibers are discarded because there is no current tree to
preserve them. If the suspended promise never resolves, the only way to
retry is something external like a context change. However, lazy context
propagation could not find the consumer fibers — they no longer exist in
the tree — so the Suspense boundary was never marked for retry and
remained stuck in fallback state indefinitely.

The fix teaches context propagation to conservatively mark suspended
Suspense boundaries for retry when a parent context changes, even when
the consumer fibers can't be found. This matches the existing
conservative approach used for dehydrated (SSR) Suspense boundaries.

Author: acdlite

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -3991,9 +3991,23 @@ function attemptEarlyBailoutIfNoScheduledUpdate(
         // whether to retry the primary children, or to skip over it and
         // go straight to the fallback. Check the priority of the primary
         // child fragment.
+        //
+        // Propagate context changes first. If a parent context changed
+        // and the primary children's consumer fibers were discarded
+        // during initial mount suspension, normal propagation can't find
+        // them. In that case we conservatively retry the boundary — the
+        // re-mounted children will read the updated context value.
+        const contextChanged = lazilyPropagateParentContextChanges(
+          current,
+          workInProgress,
+          renderLanes,
+        );
         const primaryChildFragment: Fiber = (workInProgress.child: any);
         const primaryChildLanes = primaryChildFragment.childLanes;
-        if (includesSomeLane(renderLanes, primaryChildLanes)) {
+        if (
+          contextChanged ||
+          includesSomeLane(renderLanes, primaryChildLanes)
+        ) {
           // The primary children have pending work. Use the normal path
           // to attempt to render the primary children again.
           return updateSuspenseComponent(current, workInProgress, renderLanes);

packages/react-reconciler/src/ReactFiberNewContext.js

@@ -20,7 +20,11 @@ import type {Hook} from './ReactFiberHooks';
 
 import {isPrimaryRenderer, HostTransitionContext} from './ReactFiberConfig';
 import {createCursor, push, pop} from './ReactFiberStack';
-import {ContextProvider, DehydratedFragment} from './ReactWorkTags';
+import {
+  ContextProvider,
+  DehydratedFragment,
+  SuspenseComponent,
+} from './ReactWorkTags';
 import {NoLanes, isSubsetOfLanes, mergeLanes} from './ReactFiberLane';
 import {
   NoFlags,
@@ -295,6 +299,37 @@ function propagateContextChanges<T>(
         workInProgress,
       );
       nextFiber = null;
+    } else if (
+      fiber.tag === SuspenseComponent &&
+      fiber.memoizedState !== null &&
+      fiber.memoizedState.dehydrated === null
+    ) {
+      // This is a client-rendered Suspense boundary that is currently
+      // showing its fallback. The primary children may include context
+      // consumers, but their fibers may not exist in the tree — during
+      // initial mount, if the primary children suspended, their fibers
+      // were discarded since there was no current tree to preserve them.
+      // We can't walk into the primary tree to find consumers, so
+      // conservatively mark the Suspense boundary itself for retry.
+      // When it re-renders, it will re-mount the primary children,
+      // which will read the updated context value.
+      fiber.lanes = mergeLanes(fiber.lanes, renderLanes);
+      const alternate = fiber.alternate;
+      if (alternate !== null) {
+        alternate.lanes = mergeLanes(alternate.lanes, renderLanes);
+      }
+      scheduleContextWorkOnParentPath(
+        fiber.return,
+        renderLanes,
+        workInProgress,
+      );
+      if (!forcePropagateEntireTree) {
+        // During lazy propagation, we can defer propagating changes to
+        // the children, same as the consumer match above.
+        nextFiber = null;
+      } else {
+        nextFiber = fiber.child;
+      }
     } else {
       // Traverse down.
       nextFiber = fiber.child;
@@ -331,9 +366,9 @@ export function lazilyPropagateParentContextChanges(
   current: Fiber,
   workInProgress: Fiber,
   renderLanes: Lanes,
-) {
+): boolean {
   const forcePropagateEntireTree = false;
-  propagateParentContextChanges(
+  return propagateParentContextChanges(
     current,
     workInProgress,
     renderLanes,
@@ -364,7 +399,7 @@ function propagateParentContextChanges(
   workInProgress: Fiber,
   renderLanes: Lanes,
   forcePropagateEntireTree: boolean,
-) {
+): boolean {
   // Collect all the parent providers that changed. Since this is usually small
   // number, we use an Array instead of Set.
   let contexts = null;
@@ -460,6 +495,7 @@ function propagateParentContextChanges(
   // then we could remove both `DidPropagateContext` and `NeedsPropagation`.
   // Consider this as part of the next refactor to the fiber tree structure.
   workInProgress.flags |= DidPropagateContext;
+  return contexts !== null;
 }
 
 export function checkIfContextChanged(

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

@@ -2,6 +2,7 @@ let React;
 let ReactNoop;
 let Scheduler;
 let act;
+let use;
 let useState;
 let useContext;
 let Suspense;
@@ -19,6 +20,7 @@ describe('ReactLazyContextPropagation', () => {
     ReactNoop = require('react-noop-renderer');
     Scheduler = require('scheduler');
     act = require('internal-test-utils').act;
+    use = React.use;
     useState = React.useState;
     useContext = React.useContext;
     Suspense = React.Suspense;
@@ -937,4 +939,102 @@ describe('ReactLazyContextPropagation', () => {
     assertLog(['B', 'B']);
     expect(root).toMatchRenderedOutput('BB');
   });
+
+  it('regression: context change triggers retry of suspended Suspense boundary on initial mount', async () => {
+    // Regression test for a bug where a context change above a suspended
+    // Suspense boundary would fail to trigger a retry. When a Suspense
+    // boundary suspends during initial mount, the primary children's fibers
+    // are discarded because there is no current tree to preserve them. If
+    // the suspended promise never resolves, the only way to retry is
+    // something external — like a context change. Context propagation must
+    // mark suspended Suspense boundaries for retry even though the consumer
+    // fibers no longer exist in the tree.
+    //
+    // The Provider component owns the state update. The children are
+    // passed in from above, so they are not re-created when the Provider
+    // re-renders — this means the Suspense boundary bails out, exercising
+    // the lazy context propagation path where the bug manifests.
+    const Context = React.createContext(null);
+    const neverResolvingPromise = new Promise(() => {});
+    const resolvedThenable = {status: 'fulfilled', value: 'Result', then() {}};
+
+    function Consumer() {
+      return <Text text={use(use(Context))} />;
+    }
+
+    let setPromise;
+    function Provider({children}) {
+      const [promise, _setPromise] = useState(neverResolvingPromise);
+      setPromise = _setPromise;
+      return <Context.Provider value={promise}>{children}</Context.Provider>;
+    }
+
+    const root = ReactNoop.createRoot();
+    await act(() => {
+      root.render(
+        <Provider>
+          <Suspense fallback={<Text text="Loading" />}>
+            <Consumer />
+          </Suspense>
+        </Provider>,
+      );
+    });
+    assertLog(['Loading']);
+    expect(root).toMatchRenderedOutput('Loading');
+
+    await act(() => {
+      setPromise(resolvedThenable);
+    });
+    assertLog(['Result']);
+    expect(root).toMatchRenderedOutput('Result');
+  });
+
+  it('regression: context change triggers retry of suspended Suspense boundary on initial mount (nested)', async () => {
+    // Same as above, but with an additional indirection component between
+    // the provider and the Suspense boundary. This exercises the
+    // propagateContextChanges walker path rather than the
+    // propagateParentContextChanges path.
+    const Context = React.createContext(null);
+    const neverResolvingPromise = new Promise(() => {});
+    const resolvedThenable = {status: 'fulfilled', value: 'Result', then() {}};
+
+    function Consumer() {
+      return <Text text={use(use(Context))} />;
+    }
+
+    function Indirection({children}) {
+      Scheduler.log('Indirection');
+      return children;
+    }
+
+    let setPromise;
+    function Provider({children}) {
+      const [promise, _setPromise] = useState(neverResolvingPromise);
+      setPromise = _setPromise;
+      return <Context.Provider value={promise}>{children}</Context.Provider>;
+    }
+
+    const root = ReactNoop.createRoot();
+    await act(() => {
+      root.render(
+        <Provider>
+          <Indirection>
+            <Suspense fallback={<Text text="Loading" />}>
+              <Consumer />
+            </Suspense>
+          </Indirection>
+        </Provider>,
+      );
+    });
+    assertLog(['Indirection', 'Loading']);
+    expect(root).toMatchRenderedOutput('Loading');
+
+    // Indirection should not re-render — only the Suspense boundary
+    // should be retried.
+    await act(() => {
+      setPromise(resolvedThenable);
+    });
+    assertLog(['Result']);
+    expect(root).toMatchRenderedOutput('Result');
+  });
 });