Ensure updating Context.Consumer inside suspended Suspense component

packages/react-reconciler/src/ReactFiberNewContext.new.js

@@ -23,6 +23,7 @@ import {
   ContextProvider,
   ClassComponent,
   DehydratedFragment,
+  SuspenseComponent,
 } from './ReactWorkTags';
 import {
   NoLanes,
@@ -286,6 +287,14 @@ function propagateContextChange_eager<T>(
       // this fiber to indicate that a context has changed.
       scheduleWorkOnParentPath(parentSuspense, renderLanes);
       nextFiber = fiber.sibling;
+    } else if (
+      fiber.tag === SuspenseComponent &&
+      workInProgress.tag === ContextProvider
+    ) {
+      // We don't know if it will have any context consumers in it.
+      // Schedule this fiber as having work on its children.
+      scheduleWorkOnParentPath(fiber.child, renderLanes);
+      nextFiber = fiber.child;
     } else {
       // Traverse down.
       nextFiber = fiber.child;
@@ -365,7 +374,17 @@ function propagateContextChanges<T>(
             if (alternate !== null) {
               alternate.lanes = mergeLanes(alternate.lanes, renderLanes);
             }
-            scheduleWorkOnParentPath(consumer.return, renderLanes);
+
+            if (workInProgress.tag === SuspenseComponent) {
+              // This is intentionally passing this fiber as the parent
+              // because we want to schedule this fiber as having work
+              // on its children. We'll use the childLanes on
+              // this fiber to indicate that a context has changed.
+              const primaryChildFragment = workInProgress.child;
+              scheduleWorkOnParentPath(primaryChildFragment, renderLanes);
+            } else {
+              scheduleWorkOnParentPath(consumer.return, renderLanes);
+            }
 
             if (!forcePropagateEntireTree) {
               // During lazy propagation, when we find a match, we can defer

packages/react-reconciler/src/ReactFiberNewContext.old.js

@@ -23,6 +23,7 @@ import {
   ContextProvider,
   ClassComponent,
   DehydratedFragment,
+  SuspenseComponent,
 } from './ReactWorkTags';
 import {
   NoLanes,
@@ -286,6 +287,14 @@ function propagateContextChange_eager<T>(
       // this fiber to indicate that a context has changed.
       scheduleWorkOnParentPath(parentSuspense, renderLanes);
       nextFiber = fiber.sibling;
+    } else if (
+      fiber.tag === SuspenseComponent &&
+      workInProgress.tag === ContextProvider
+    ) {
+      // We don't know if it will have any context consumers in it.
+      // Schedule this fiber as having work on its children.
+      scheduleWorkOnParentPath(fiber.child, renderLanes);
+      nextFiber = fiber.child;
     } else {
       // Traverse down.
       nextFiber = fiber.child;
@@ -365,7 +374,17 @@ function propagateContextChanges<T>(
             if (alternate !== null) {
               alternate.lanes = mergeLanes(alternate.lanes, renderLanes);
             }
-            scheduleWorkOnParentPath(consumer.return, renderLanes);
+
+            if (workInProgress.tag === SuspenseComponent) {
+              // This is intentionally passing this fiber as the parent
+              // because we want to schedule this fiber as having work
+              // on its children. We'll use the childLanes on
+              // this fiber to indicate that a context has changed.
+              const primaryChildFragment = workInProgress.child;
+              scheduleWorkOnParentPath(primaryChildFragment, renderLanes);
+            } else {
+              scheduleWorkOnParentPath(consumer.return, renderLanes);
+            }
 
             if (!forcePropagateEntireTree) {
               // During lazy propagation, when we find a match, we can defer

packages/react-reconciler/src/__tests__/ReactSuspense-test.internal.js

@@ -1533,4 +1533,63 @@ describe('ReactSuspense', () => {
       expect(root).toMatchRenderedOutput('new value');
     });
   });
+
+  it('updates context consumer within child of suspended suspense component when context updates', () => {
+    const {createContext, useState} = React;
+
+    const ValueContext = createContext(null);
+
+    const promiseThatNeverResolves = new Promise(() => {});
+    function Child() {
+      return (
+        <ValueContext.Consumer>
+          {value => {
+            Scheduler.unstable_yieldValue(`Received context value [${value}]`);
+            if (value === 'default') return <Text text="default" />;
+            throw promiseThatNeverResolves;
+          }}
+        </ValueContext.Consumer>
+      );
+    }
+
+    let setValue;
+    function Wrapper({children}) {
+      const [value, _setValue] = useState('default');
+      setValue = _setValue;
+      return (
+        <ValueContext.Provider value={value}>{children}</ValueContext.Provider>
+      );
+    }
+
+    function App() {
+      return (
+        <Wrapper>
+          <Suspense fallback={<Text text="Loading..." />}>
+            <Child />
+          </Suspense>
+        </Wrapper>
+      );
+    }
+
+    const root = ReactTestRenderer.create(<App />);
+    expect(Scheduler).toHaveYielded([
+      'Received context value [default]',
+      'default',
+    ]);
+    expect(root).toMatchRenderedOutput('default');
+
+    act(() => setValue('new value'));
+    expect(Scheduler).toHaveYielded([
+      'Received context value [new value]',
+      'Loading...',
+    ]);
+    expect(root).toMatchRenderedOutput('Loading...');
+
+    act(() => setValue('default'));
+    expect(Scheduler).toHaveYielded([
+      'Received context value [default]',
+      'default',
+    ]);
+    expect(root).toMatchRenderedOutput('default');
+  });
 });