[Flight] Track owner/stack where the Flight Client reads as the root (#30933)

This means that the owner of a Component rendered on the remote server
becomes the Component on this server.

Ideally we'd support this for the Client side too. In particular Fiber
but currently ReactComponentInfo's owner is typed as only supporting
other ReactComponentInfo and it's a bigger lift to support that.

packages/react-client/src/ReactFlightClient.js

@@ -86,14 +86,19 @@ import isArray from 'shared/isArray';
 
 import * as React from 'react';
 
+import type {SharedStateServer} from 'react/src/ReactSharedInternalsServer';
+import type {SharedStateClient} from 'react/src/ReactSharedInternalsClient';
+
 // TODO: This is an unfortunate hack. We shouldn't feature detect the internals
 // like this. It's just that for now we support the same build of the Flight
 // client both in the RSC environment, in the SSR environments as well as the
 // browser client. We should probably have a separate RSC build. This is DEV
 // only though.
-const ReactSharedInternals =
+const ReactSharedInteralsServer: void | SharedStateServer = (React: any)
+  .__SERVER_INTERNALS_DO_NOT_USE_OR_WARN_USERS_THEY_CANNOT_UPGRADE;
+const ReactSharedInternals: SharedStateServer | SharedStateClient =
   React.__CLIENT_INTERNALS_DO_NOT_USE_OR_WARN_USERS_THEY_CANNOT_UPGRADE ||
-  React.__SERVER_INTERNALS_DO_NOT_USE_OR_WARN_USERS_THEY_CANNOT_UPGRADE;
+  ReactSharedInteralsServer;
 
 export type {CallServerCallback, EncodeFormActionCallback};
 
@@ -277,6 +282,8 @@ export type Response = {
   _rowLength: number, // remaining bytes in the row. 0 indicates that we're looking for a newline.
   _buffer: Array<Uint8Array>, // chunks received so far as part of this row
   _tempRefs: void | TemporaryReferenceSet, // the set temporary references can be resolved from
+  _debugRootOwner?: null | ReactComponentInfo, // DEV-only
+  _debugRootStack?: null | Error, // DEV-only
   _debugRootTask?: null | ConsoleTask, // DEV-only
   _debugFindSourceMapURL?: void | FindSourceMapURLCallback, // DEV-only
   _replayConsole: boolean, // DEV-only
@@ -672,7 +679,7 @@ function createElement(
       type,
       key,
       props,
-      _owner: owner,
+      _owner: __DEV__ && owner === null ? response._debugRootOwner : owner,
     }: any);
     Object.defineProperty(element, 'ref', {
       enumerable: false,
@@ -699,7 +706,7 @@ function createElement(
       props,
 
       // Record the component responsible for creating this element.
-      _owner: owner,
+      _owner: __DEV__ && owner === null ? response._debugRootOwner : owner,
     }: any);
   }
 
@@ -733,7 +740,11 @@ function createElement(
         env = owner.env;
       }
       let normalizedStackTrace: null | Error = null;
-      if (stack !== null) {
+      if (owner === null && response._debugRootStack != null) {
+        // We override the stack if we override the owner since the stack where the root JSX
+        // was created on the server isn't very useful but where the request was made is.
+        normalizedStackTrace = response._debugRootStack;
+      } else if (stack !== null) {
         // We create a fake stack and then create an Error object inside of it.
         // This means that the stack trace is now normalized into the native format
         // of the browser and the stack frames will have been registered with
@@ -821,8 +832,10 @@ function createElement(
         if (enableOwnerStacks) {
           // $FlowFixMe[cannot-write]
           erroredComponent.debugStack = element._debugStack;
-          // $FlowFixMe[cannot-write]
-          erroredComponent.debugTask = element._debugTask;
+          if (supportsCreateTask) {
+            // $FlowFixMe[cannot-write]
+            erroredComponent.debugTask = element._debugTask;
+          }
         }
         erroredChunk._debugInfo = [erroredComponent];
       }
@@ -998,8 +1011,10 @@ function waitForReference<T>(
         if (enableOwnerStacks) {
           // $FlowFixMe[cannot-write]
           erroredComponent.debugStack = element._debugStack;
-          // $FlowFixMe[cannot-write]
-          erroredComponent.debugTask = element._debugTask;
+          if (supportsCreateTask) {
+            // $FlowFixMe[cannot-write]
+            erroredComponent.debugTask = element._debugTask;
+          }
         }
         const chunkDebugInfo: ReactDebugInfo =
           chunk._debugInfo || (chunk._debugInfo = []);
@@ -1408,6 +1423,25 @@ function ResponseInstance(
   this._buffer = [];
   this._tempRefs = temporaryReferences;
   if (__DEV__) {
+    // TODO: The Flight Client can be used in a Client Environment too and we should really support
+    // getting the owner there as well, but currently the owner of ReactComponentInfo is typed as only
+    // supporting other ReactComponentInfo as owners (and not Fiber or Fizz's ComponentStackNode).
+    // We need to update all the callsites consuming ReactComponentInfo owners to support those.
+    // In the meantime we only check ReactSharedInteralsServer since we know that in an RSC environment
+    // the only owners will be ReactComponentInfo.
+    const rootOwner: null | ReactComponentInfo =
+      ReactSharedInteralsServer === undefined ||
+      ReactSharedInteralsServer.A === null
+        ? null
+        : (ReactSharedInteralsServer.A.getOwner(): any);
+
+    this._debugRootOwner = rootOwner;
+    this._debugRootStack =
+      rootOwner !== null
+        ? // TODO: Consider passing the top frame in so we can avoid internals showing up.
+          new Error('react-stack-top-frame')
+        : null;
+
     const rootEnv = environmentName === undefined ? 'Server' : environmentName;
     if (supportsCreateTask) {
       // Any stacks that appear on the server need to be rooted somehow on the client
@@ -2308,7 +2342,16 @@ function resolveDebugInfo(
   const env =
     debugInfo.env === undefined ? response._rootEnvironmentName : debugInfo.env;
   initializeFakeTask(response, debugInfo, env);
-  initializeFakeStack(response, debugInfo);
+  if (debugInfo.owner === null && response._debugRootOwner != null) {
+    // $FlowFixMe
+    debugInfo.owner = response._debugRootOwner;
+    // We override the stack if we override the owner since the stack where the root JSX
+    // was created on the server isn't very useful but where the request was made is.
+    // $FlowFixMe
+    debugInfo.debugStack = response._debugRootStack;
+  } else {
+    initializeFakeStack(response, debugInfo);
+  }
 
   const chunk = getChunk(response, id);
   const chunkDebugInfo: ReactDebugInfo =
@@ -2344,7 +2387,8 @@ const replayConsoleWithCallStack = {
     // There really shouldn't be anything else on the stack atm.
     const prevStack = ReactSharedInternals.getCurrentStack;
     ReactSharedInternals.getCurrentStack = getCurrentStackInDEV;
-    currentOwnerInDEV = owner;
+    currentOwnerInDEV =
+      owner === null ? (response._debugRootOwner: any) : owner;
 
     try {
       const callStack = buildFakeCallStack(

packages/react-client/src/__tests__/ReactFlight-test.js

@@ -24,6 +24,10 @@ function normalizeCodeLocInfo(str) {
   return (
     str &&
     str.replace(/^ +(?:at|in) ([\S]+)[^\n]*/gm, function (m, name) {
+      const dot = name.lastIndexOf('.');
+      if (dot !== -1) {
+        name = name.slice(dot + 1);
+      }
       return '    in ' + name + (/\d/.test(m) ? ' (at **)' : '');
     })
   );
@@ -3124,6 +3128,69 @@ describe('ReactFlight', () => {
     );
   });
 
+  // @gate __DEV__ && enableOwnerStacks
+  it('can track owner for a flight response created in another render', async () => {
+    jest.resetModules();
+    jest.mock('react', () => ReactServer);
+    // For this to work the Flight Client needs to be the react-server version.
+    const ReactNoopFlightClienOnTheServer = require('react-noop-renderer/flight-client');
+    jest.resetModules();
+    jest.mock('react', () => React);
+
+    let stack;
+
+    function Component() {
+      stack = ReactServer.captureOwnerStack();
+      return ReactServer.createElement('span', null, 'hi');
+    }
+
+    const ClientComponent = clientReference(Component);
+
+    function ThirdPartyComponent() {
+      return ReactServer.createElement(ClientComponent);
+    }
+
+    // This is rendered outside the render to ensure we don't inherit anything accidental
+    // by being in the same environment which would make it seem like it works when it doesn't.
+    const thirdPartyTransport = ReactNoopFlightServer.render(
+      {children: ReactServer.createElement(ThirdPartyComponent)},
+      {
+        environmentName: 'third-party',
+      },
+    );
+
+    async function fetchThirdParty() {
+      return ReactNoopFlightClienOnTheServer.read(thirdPartyTransport);
+    }
+
+    async function FirstPartyComponent() {
+      // This component fetches from a third party
+      const thirdParty = await fetchThirdParty();
+      return thirdParty.children;
+    }
+    function App() {
+      return ReactServer.createElement(FirstPartyComponent);
+    }
+
+    const transport = ReactNoopFlightServer.render(
+      ReactServer.createElement(App),
+    );
+
+    await act(async () => {
+      const root = await ReactNoopFlightClient.read(transport);
+      ReactNoop.render(root);
+    });
+
+    expect(normalizeCodeLocInfo(stack)).toBe(
+      '\n    in ThirdPartyComponent (at **)' +
+        '\n    in createResponse (at **)' + // These two internal frames should
+        '\n    in read (at **)' + // ideally not be included.
+        '\n    in fetchThirdParty (at **)' +
+        '\n    in FirstPartyComponent (at **)' +
+        '\n    in App (at **)',
+    );
+  });
+
   // @gate __DEV__ && enableOwnerStacks
   it('can get the component owner stacks for onError in dev', async () => {
     const thrownError = new Error('hi');

packages/react-server/src/ReactFlightServer.js

@@ -2271,7 +2271,7 @@ function isReactComponentInfo(value: any): boolean {
       typeof value.debugTask.run === 'function') ||
       value.debugStack instanceof Error) &&
     (enableOwnerStacks
-      ? isArray((value: any).stack)
+      ? isArray((value: any).stack) || (value: any).stack === null
       : typeof (value: any).stack === 'undefined') &&
     typeof value.name === 'string' &&
     typeof value.env === 'string' &&