[Flight] Fix wrong missing key warning when static child is blocked (#34350)

Author: unstubbable

packages/react-client/src/ReactFlightClient.js

@@ -1074,7 +1074,14 @@ function getTaskName(type: mixed): string {
   }
 }
 
-function initializeElement(response: Response, element: any): void {
+function initializeElement(
+  response: Response,
+  element: any,
+  lazyType: null | LazyComponent<
+    React$Element<any>,
+    SomeChunk<React$Element<any>>,
+  >,
+): void {
   if (!__DEV__) {
     return;
   }
@@ -1141,6 +1148,18 @@ function initializeElement(response: Response, element: any): void {
   if (owner !== null) {
     initializeFakeStack(response, owner);
   }
+
+  // In case the JSX runtime has validated the lazy type as a static child, we
+  // need to transfer this information to the element.
+  if (
+    lazyType &&
+    lazyType._store &&
+    lazyType._store.validated &&
+    !element._store.validated
+  ) {
+    element._store.validated = lazyType._store.validated;
+  }
+
   // TODO: We should be freezing the element but currently, we might write into
   // _debugInfo later. We could move it into _store which remains mutable.
   Object.freeze(element.props);
@@ -1230,7 +1249,7 @@ function createElement(
         handler.reason,
       );
       if (__DEV__) {
-        initializeElement(response, element);
+        initializeElement(response, element, null);
         // Conceptually the error happened inside this Element but right before
         // it was rendered. We don't have a client side component to render but
         // we can add some DebugInfo to explain that this was conceptually a
@@ -1258,16 +1277,17 @@ function createElement(
         createBlockedChunk(response);
       handler.value = element;
       handler.chunk = blockedChunk;
+      const lazyType = createLazyChunkWrapper(blockedChunk);
       if (__DEV__) {
-        /// After we have initialized any blocked references, initialize stack etc.
-        const init = initializeElement.bind(null, response, element);
+        // After we have initialized any blocked references, initialize stack etc.
+        const init = initializeElement.bind(null, response, element, lazyType);
         blockedChunk.then(init, init);
       }
-      return createLazyChunkWrapper(blockedChunk);
+      return lazyType;
     }
   }
   if (__DEV__) {
-    initializeElement(response, element);
+    initializeElement(response, element, null);
   }
 
   return element;
@@ -1279,6 +1299,7 @@ function createLazyChunkWrapper<T>(
   const lazyType: LazyComponent<T, SomeChunk<T>> = {
     $$typeof: REACT_LAZY_TYPE,
     _payload: chunk,
+    _store: {validated: 0},
     _init: readChunk,
   };
   if (__DEV__) {

packages/react-server-dom-webpack/src/__tests__/ReactFlightDOMBrowser-test.js

@@ -2846,4 +2846,64 @@ describe('ReactFlightDOMBrowser', () => {
 
     expect(container.innerHTML).toBe('<p>Hi</p>');
   });
+
+  it('should not have missing key warnings when a static child is blocked on debug info', async () => {
+    const ClientComponent = clientExports(function ClientComponent({element}) {
+      return (
+        <div>
+          <span>Hi</span>
+          {element}
+        </div>
+      );
+    });
+
+    let debugReadableStreamController;
+
+    const debugReadableStream = new ReadableStream({
+      start(controller) {
+        debugReadableStreamController = controller;
+      },
+    });
+
+    const stream = await serverAct(() =>
+      ReactServerDOMServer.renderToReadableStream(
+        <ClientComponent element={<span>Sebbie</span>} />,
+        webpackMap,
+        {
+          debugChannel: {
+            writable: new WritableStream({
+              write(chunk) {
+                debugReadableStreamController.enqueue(chunk);
+              },
+              close() {
+                debugReadableStreamController.close();
+              },
+            }),
+          },
+        },
+      ),
+    );
+
+    function ClientRoot({response}) {
+      return use(response);
+    }
+
+    const response = ReactServerDOMClient.createFromReadableStream(stream, {
+      debugChannel: {readable: createDelayedStream(debugReadableStream)},
+    });
+
+    const container = document.createElement('div');
+    const root = ReactDOMClient.createRoot(container);
+
+    await act(() => {
+      root.render(<ClientRoot response={response} />);
+    });
+
+    // Wait for the debug info to be processed.
+    await act(() => {});
+
+    expect(container.innerHTML).toBe(
+      '<div><span>Hi</span><span>Sebbie</span></div>',
+    );
+  });
 });

packages/react/src/ReactLazy.js

@@ -60,6 +60,8 @@ export type LazyComponent<T, P> = {
   _payload: P,
   _init: (payload: P) => T,
   _debugInfo?: null | ReactDebugInfo,
+  // __DEV__
+  _store?: {validated: 0 | 1 | 2, ...}, // 0: not validated, 1: validated, 2: force fail
 };
 
 function lazyInitializer<T>(payload: Payload<T>): T {

packages/react/src/jsx/ReactJSXElement.js

@@ -804,6 +804,14 @@ function validateChildKeys(node) {
       if (node._store) {
         node._store.validated = 1;
       }
+    } else if (isLazyType(node)) {
+      if (node._payload.status === 'fulfilled') {
+        if (isValidElement(node._payload.value) && node._payload.value._store) {
+          node._payload.value._store.validated = 1;
+        }
+      } else if (node._store) {
+        node._store.validated = 1;
+      }
     }
   }
 }
@@ -822,3 +830,11 @@ export function isValidElement(object) {
     object.$$typeof === REACT_ELEMENT_TYPE
   );
 }
+
+export function isLazyType(object) {
+  return (
+    typeof object === 'object' &&
+    object !== null &&
+    object.$$typeof === REACT_LAZY_TYPE
+  );
+}