[Flight] Ensure blocked debug info is handled properly (#34524)

This PR ensures that server components are reliably included in the
DevTools component tree, even if debug info is received delayed, e.g.
when using a debug channel. The fix consists of three parts:

- We must not unset the debug chunk before all debug info entries are
resolved.
- We must ensure that the "RSC Stream" IO debug info entry is pushed
last, after all other entries were resolved.
- We need to transfer the debug info from blocked element chunks onto
the lazy node and the element.

Ideally, we wouldn't even create a lazy node for blocked elements that
are at the root of the JSON payload, because that would basically wrap a
lazy in a lazy. This optimization that ensures that everything around
the blocked element can proceed is only needed for nested elements.
However, we also need it for resolving deduped references in blocked
root elements, unless we adapt that logic, which would be a bigger lift.

When reloading the Flight fixture, the component tree is now displayed
deterministically. Previously, it would sometimes omit synchronous
server components.

<img width="306" height="565" alt="complete"
src="https://github.com/user-attachments/assets/db61aa10-1816-43e6-9903-0e585190cdf1"
/>

---------

Co-authored-by: Sebastian Markbage <[email protected]>

Author: unstubbable

packages/react-client/src/ReactFlightClient.js

@@ -499,10 +499,44 @@ function createErrorChunk<T>(
   return new ReactPromise(ERRORED, null, error);
 }
 
+function moveDebugInfoFromChunkToInnerValue<T>(
+  chunk: InitializedChunk<T>,
+  value: T,
+): void {
+  // Remove the debug info from the initialized chunk, and add it to the inner
+  // value instead. This can be a React element, an array, or an uninitialized
+  // Lazy.
+  const resolvedValue = resolveLazy(value);
+  if (
+    typeof resolvedValue === 'object' &&
+    resolvedValue !== null &&
+    (isArray(resolvedValue) ||
+      typeof resolvedValue[ASYNC_ITERATOR] === 'function' ||
+      resolvedValue.$$typeof === REACT_ELEMENT_TYPE ||
+      resolvedValue.$$typeof === REACT_LAZY_TYPE)
+  ) {
+    const debugInfo = chunk._debugInfo.splice(0);
+    if (isArray(resolvedValue._debugInfo)) {
+      // $FlowFixMe[method-unbinding]
+      resolvedValue._debugInfo.unshift.apply(
+        resolvedValue._debugInfo,
+        debugInfo,
+      );
+    } else {
+      Object.defineProperty((resolvedValue: any), '_debugInfo', {
+        configurable: false,
+        enumerable: false,
+        writable: true,
+        value: debugInfo,
+      });
+    }
+  }
+}
+
 function wakeChunk<T>(
   listeners: Array<InitializationReference | (T => mixed)>,
   value: T,
-  chunk: SomeChunk<T>,
+  chunk: InitializedChunk<T>,
 ): void {
   for (let i = 0; i < listeners.length; i++) {
     const listener = listeners[i];
@@ -512,6 +546,10 @@ function wakeChunk<T>(
       fulfillReference(listener, value, chunk);
     }
   }
+
+  if (__DEV__) {
+    moveDebugInfoFromChunkToInnerValue(chunk, value);
+  }
 }
 
 function rejectChunk(
@@ -649,7 +687,6 @@ function triggerErrorOnChunk<T>(
       }
       try {
         initializeDebugChunk(response, chunk);
-        chunk._debugChunk = null;
         if (initializingHandler !== null) {
           if (initializingHandler.errored) {
             // Ignore error parsing debug info, we'll report the original error instead.
@@ -932,9 +969,9 @@ function initializeModelChunk<T>(chunk: ResolvedModelChunk<T>): void {
   }
 
   if (__DEV__) {
-    // Lazily initialize any debug info and block the initializing chunk on any unresolved entries.
+    // Initialize any debug info and block the initializing chunk on any
+    // unresolved entries.
     initializeDebugChunk(response, chunk);
-    chunk._debugChunk = null;
   }
 
   try {
@@ -946,7 +983,14 @@ function initializeModelChunk<T>(chunk: ResolvedModelChunk<T>): void {
     if (resolveListeners !== null) {
       cyclicChunk.value = null;
       cyclicChunk.reason = null;
-      wakeChunk(resolveListeners, value, cyclicChunk);
+      for (let i = 0; i < resolveListeners.length; i++) {
+        const listener = resolveListeners[i];
+        if (typeof listener === 'function') {
+          listener(value);
+        } else {
+          fulfillReference(listener, value, cyclicChunk);
+        }
+      }
     }
     if (initializingHandler !== null) {
       if (initializingHandler.errored) {
@@ -963,6 +1007,10 @@ function initializeModelChunk<T>(chunk: ResolvedModelChunk<T>): void {
     const initializedChunk: InitializedChunk<T> = (chunk: any);
     initializedChunk.status = INITIALIZED;
     initializedChunk.value = value;
+
+    if (__DEV__) {
+      moveDebugInfoFromChunkToInnerValue(initializedChunk, value);
+    }
   } catch (error) {
     const erroredChunk: ErroredChunk<T> = (chunk: any);
     erroredChunk.status = ERRORED;
@@ -1079,7 +1127,7 @@ function getTaskName(type: mixed): string {
 function initializeElement(
   response: Response,
   element: any,
-  lazyType: null | LazyComponent<
+  lazyNode: null | LazyComponent<
     React$Element<any>,
     SomeChunk<React$Element<any>>,
   >,
@@ -1151,15 +1199,33 @@ function initializeElement(
     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;
+  if (lazyNode !== null) {
+    // In case the JSX runtime has validated the lazy type as a static child, we
+    // need to transfer this information to the element.
+    if (
+      lazyNode._store &&
+      lazyNode._store.validated &&
+      !element._store.validated
+    ) {
+      element._store.validated = lazyNode._store.validated;
+    }
+
+    // If the lazy node is initialized, we move its debug info to the inner
+    // value.
+    if (lazyNode._payload.status === INITIALIZED && lazyNode._debugInfo) {
+      const debugInfo = lazyNode._debugInfo.splice(0);
+      if (element._debugInfo) {
+        // $FlowFixMe[method-unbinding]
+        element._debugInfo.unshift.apply(element._debugInfo, debugInfo);
+      } else {
+        Object.defineProperty(element, '_debugInfo', {
+          configurable: false,
+          enumerable: false,
+          writable: true,
+          value: debugInfo,
+        });
+      }
+    }
   }
 
   // TODO: We should be freezing the element but currently, we might write into
@@ -1279,13 +1345,13 @@ function createElement(
         createBlockedChunk(response);
       handler.value = element;
       handler.chunk = blockedChunk;
-      const lazyType = createLazyChunkWrapper(blockedChunk, validated);
+      const lazyNode = createLazyChunkWrapper(blockedChunk, validated);
       if (__DEV__) {
         // After we have initialized any blocked references, initialize stack etc.
-        const init = initializeElement.bind(null, response, element, lazyType);
+        const init = initializeElement.bind(null, response, element, lazyNode);
         blockedChunk.then(init, init);
       }
-      return lazyType;
+      return lazyNode;
     }
   }
   if (__DEV__) {
@@ -1466,7 +1532,7 @@ function fulfillReference(
     const element: any = handler.value;
     switch (key) {
       case '3':
-        transferReferencedDebugInfo(handler.chunk, fulfilledChunk, mappedValue);
+        transferReferencedDebugInfo(handler.chunk, fulfilledChunk);
         element.props = mappedValue;
         break;
       case '4':
@@ -1482,11 +1548,11 @@ function fulfillReference(
         }
         break;
       default:
-        transferReferencedDebugInfo(handler.chunk, fulfilledChunk, mappedValue);
+        transferReferencedDebugInfo(handler.chunk, fulfilledChunk);
         break;
     }
   } else if (__DEV__ && !reference.isDebug) {
-    transferReferencedDebugInfo(handler.chunk, fulfilledChunk, mappedValue);
+    transferReferencedDebugInfo(handler.chunk, fulfilledChunk);
   }
 
   handler.deps--;
@@ -1808,47 +1874,34 @@ function loadServerReference<A: Iterable<any>, T>(
   return (null: any);
 }
 
+function resolveLazy(value: any): mixed {
+  while (
+    typeof value === 'object' &&
+    value !== null &&
+    value.$$typeof === REACT_LAZY_TYPE
+  ) {
+    const payload: SomeChunk<any> = value._payload;
+    if (payload.status === INITIALIZED) {
+      value = payload.value;
+      continue;
+    }
+    break;
+  }
+
+  return value;
+}
+
 function transferReferencedDebugInfo(
   parentChunk: null | SomeChunk<any>,
   referencedChunk: SomeChunk<any>,
-  referencedValue: mixed,
 ): void {
   if (__DEV__) {
-    const referencedDebugInfo = referencedChunk._debugInfo;
-    // If we have a direct reference to an object that was rendered by a synchronous
-    // server component, it might have some debug info about how it was rendered.
-    // We forward this to the underlying object. This might be a React Element or
-    // an Array fragment.
-    // If this was a string / number return value we lose the debug info. We choose
-    // that tradeoff to allow sync server components to return plain values and not
-    // use them as React Nodes necessarily. We could otherwise wrap them in a Lazy.
-    if (
-      typeof referencedValue === 'object' &&
-      referencedValue !== null &&
-      (isArray(referencedValue) ||
-        typeof referencedValue[ASYNC_ITERATOR] === 'function' ||
-        referencedValue.$$typeof === REACT_ELEMENT_TYPE)
-    ) {
-      // We should maybe use a unique symbol for arrays but this is a React owned array.
-      // $FlowFixMe[prop-missing]: This should be added to elements.
-      const existingDebugInfo: ?ReactDebugInfo =
-        (referencedValue._debugInfo: any);
-      if (existingDebugInfo == null) {
-        Object.defineProperty((referencedValue: any), '_debugInfo', {
-          configurable: false,
-          enumerable: false,
-          writable: true,
-          value: referencedDebugInfo.slice(0), // Clone so that pushing later isn't going into the original
-        });
-      } else {
-        // $FlowFixMe[method-unbinding]
-        existingDebugInfo.push.apply(existingDebugInfo, referencedDebugInfo);
-      }
-    }
-    // We also add the debug info to the initializing chunk since the resolution of that promise is
-    // also blocked by the referenced debug info. By adding it to both we can track it even if the array/element
-    // is extracted, or if the root is rendered as is.
+    // We add the debug info to the initializing chunk since the resolution of
+    // that promise is also blocked by the referenced debug info. By adding it
+    // to both we can track it even if the array/element/lazy is extracted, or
+    // if the root is rendered as is.
     if (parentChunk !== null) {
+      const referencedDebugInfo = referencedChunk._debugInfo;
       const parentDebugInfo = parentChunk._debugInfo;
       for (let i = 0; i < referencedDebugInfo.length; ++i) {
         const debugInfoEntry = referencedDebugInfo[i];
@@ -1999,7 +2052,7 @@ function getOutlinedModel<T>(
         // If we're resolving the "owner" or "stack" slot of an Element array, we don't call
         // transferReferencedDebugInfo because this reference is to a debug chunk.
       } else {
-        transferReferencedDebugInfo(initializingChunk, chunk, chunkValue);
+        transferReferencedDebugInfo(initializingChunk, chunk);
       }
       return chunkValue;
     case PENDING:
@@ -2709,14 +2762,47 @@ function incrementChunkDebugInfo(
   }
 }
 
+function addDebugInfo(chunk: SomeChunk<any>, debugInfo: ReactDebugInfo): void {
+  const value = resolveLazy(chunk.value);
+  if (
+    typeof value === 'object' &&
+    value !== null &&
+    (isArray(value) ||
+      typeof value[ASYNC_ITERATOR] === 'function' ||
+      value.$$typeof === REACT_ELEMENT_TYPE ||
+      value.$$typeof === REACT_LAZY_TYPE)
+  ) {
+    if (isArray(value._debugInfo)) {
+      // $FlowFixMe[method-unbinding]
+      value._debugInfo.push.apply(value._debugInfo, debugInfo);
+    } else {
+      Object.defineProperty((value: any), '_debugInfo', {
+        configurable: false,
+        enumerable: false,
+        writable: true,
+        value: debugInfo,
+      });
+    }
+  } else {
+    // $FlowFixMe[method-unbinding]
+    chunk._debugInfo.push.apply(chunk._debugInfo, debugInfo);
+  }
+}
+
 function resolveChunkDebugInfo(
   streamState: StreamState,
   chunk: SomeChunk<any>,
 ): void {
   if (__DEV__ && enableAsyncDebugInfo) {
-    // Push the currently resolving chunk's debug info representing the stream on the Promise
-    // that was waiting on the stream.
-    chunk._debugInfo.push({awaited: streamState._debugInfo});
+    // Add the currently resolving chunk's debug info representing the stream
+    // to the Promise that was waiting on the stream, or its underlying value.
+    const debugInfo: ReactDebugInfo = [{awaited: streamState._debugInfo}];
+    if (chunk.status === PENDING || chunk.status === BLOCKED) {
+      const boundAddDebugInfo = addDebugInfo.bind(null, chunk, debugInfo);
+      chunk.then(boundAddDebugInfo, boundAddDebugInfo);
+    } else {
+      addDebugInfo(chunk, debugInfo);
+    }
   }
 }
 
@@ -2909,7 +2995,8 @@ function resolveStream<T: ReadableStream | $AsyncIterable<any, any, void>>(
   const resolveListeners = chunk.value;
 
   if (__DEV__) {
-    // Lazily initialize any debug info and block the initializing chunk on any unresolved entries.
+    // Initialize any debug info and block the initializing chunk on any
+    // unresolved entries.
     if (chunk._debugChunk != null) {
       const prevHandler = initializingHandler;
       const prevChunk = initializingChunk;
@@ -2923,7 +3010,6 @@ function resolveStream<T: ReadableStream | $AsyncIterable<any, any, void>>(
       }
       try {
         initializeDebugChunk(response, chunk);
-        chunk._debugChunk = null;
         if (initializingHandler !== null) {
           if (initializingHandler.errored) {
             // Ignore error parsing debug info, we'll report the original error instead.
@@ -2947,7 +3033,7 @@ function resolveStream<T: ReadableStream | $AsyncIterable<any, any, void>>(
   resolvedChunk.value = stream;
   resolvedChunk.reason = controller;
   if (resolveListeners !== null) {
-    wakeChunk(resolveListeners, chunk.value, chunk);
+    wakeChunk(resolveListeners, chunk.value, (chunk: any));
   }
 }