Repository: facebook/react
Pull request: [Flight] Enable owner stacks on the client when replaying logs

There's a special case that happens when we replay logs on the client because this doesn't happen within the context of any particular rendered component. So we need to reimplement things that would normally be handled by a full client like Fiber.

The implementation of `getOwnerStackByComponentInfoInDev` is the simplest version since it doesn't have any client components in it so I move it to `shared/`. It's only used by Flight but both `react-server/` and `react-client/` packages. The ReactComponentInfo type is also more generic than just Flight anyway.

In a follow up I still need to implement this in React DevTools when native tasks are not available so that it appends it to the console.

Changed files:
filename: packages/react-client/src/ReactFlightClient.js
patch: @@ -73,8 +73,21 @@ import {
 
 import getComponentNameFromType from 'shared/getComponentNameFromType';
 
+import {getOwnerStackByComponentInfoInDev} from 'shared/ReactComponentInfoStack';
+
 import isArray from 'shared/isArray';
 
+import * as React from 'react';
+
+// 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 =
+  React.__CLIENT_INTERNALS_DO_NOT_USE_OR_WARN_USERS_THEY_CANNOT_UPGRADE ||
+  React.__SERVER_INTERNALS_DO_NOT_USE_OR_WARN_USERS_THEY_CANNOT_UPGRADE;
+
 export type {CallServerCallback, EncodeFormActionCallback};
 
 interface FlightStreamController {
@@ -2296,6 +2309,22 @@ function resolveDebugInfo(
   chunkDebugInfo.push(debugInfo);
 }
 
+let currentOwnerInDEV: null | ReactComponentInfo = null;
+function getCurrentStackInDEV(): string {
+  if (__DEV__) {
+    if (enableOwnerStacks) {
+      const owner: null | ReactComponentInfo = currentOwnerInDEV;
+      if (owner === null) {
+        return '';
+      }
+      return getOwnerStackByComponentInfoInDev(owner);
+    }
+    // We don't have Parent Stacks in Flight.
+    return '';
+  }
+  return '';
+}
+
 function resolveConsoleEntry(
   response: Response,
   value: UninitializedModel,
@@ -2324,34 +2353,44 @@ function resolveConsoleEntry(
   const owner = payload[2];
   const env = payload[3];
   const args = payload.slice(4);
-  if (!enableOwnerStacks) {
-    // Printing with stack isn't really limited to owner stacks but
-    // we gate it behind the same flag for now while iterating.
-    bindToConsole(methodName, args, env)();
-    return;
-  }
-  const callStack = buildFakeCallStack(
-    response,
-    stackTrace,
-    env,
-    bindToConsole(methodName, args, env),
-  );
-  if (owner != null) {
-    const task = initializeFakeTask(response, owner, env);
-    initializeFakeStack(response, owner);
-    if (task !== null) {
-      task.run(callStack);
+
+  // There really shouldn't be anything else on the stack atm.
+  const prevStack = ReactSharedInternals.getCurrentStack;
+  ReactSharedInternals.getCurrentStack = getCurrentStackInDEV;
+  currentOwnerInDEV = owner;
+
+  try {
+    if (!enableOwnerStacks) {
+      // Printing with stack isn't really limited to owner stacks but
+      // we gate it behind the same flag for now while iterating.
+      bindToConsole(methodName, args, env)();
       return;
     }
-    // TODO: Set the current owner so that captureOwnerStack() adds the component
-    // stack during the replay - if needed.
-  }
-  const rootTask = getRootTask(response, env);
-  if (rootTask != null) {
-    rootTask.run(callStack);
-    return;
+    const callStack = buildFakeCallStack(
+      response,
+      stackTrace,
+      env,
+      bindToConsole(methodName, args, env),
+    );
+    if (owner != null) {
+      const task = initializeFakeTask(response, owner, env);
+      initializeFakeStack(response, owner);
+      if (task !== null) {
+        task.run(callStack);
+        return;
+      }
+      // TODO: Set the current owner so that captureOwnerStack() adds the component
+      // stack during the replay - if needed.
+    }
+    const rootTask = getRootTask(response, env);
+    if (rootTask != null) {
+      rootTask.run(callStack);
+      return;
+    }
+    callStack();
+  } finally {
+    ReactSharedInternals.getCurrentStack = prevStack;
   }
-  callStack();
 }
 
 function mergeBuffer(

filename: packages/react-client/src/__tests__/ReactFlight-test.js
patch: @@ -2918,7 +2918,7 @@ describe('ReactFlight', () => {
     expect(ReactNoop).toMatchRenderedOutput(<div>hi</div>);
   });
 
-  // @gate enableServerComponentLogs && __DEV__
+  // @gate enableServerComponentLogs && __DEV__ && enableOwnerStacks
   it('replays logs, but not onError logs', async () => {
     function foo() {
       return 'hello';
@@ -2928,12 +2928,21 @@ describe('ReactFlight', () => {
       throw new Error('err');
     }
 
+    function App() {
+      return ReactServer.createElement(ServerComponent);
+    }
+
+    let ownerStacks = [];
+
     // These tests are specifically testing console.log.
     // Assign to `mockConsoleLog` so we can still inspect it when `console.log`
     // is overridden by the test modules. The original function will be restored
     // after this test finishes by `jest.restoreAllMocks()`.
     const mockConsoleLog = spyOnDevAndProd(console, 'log').mockImplementation(
-      () => {},
+      () => {
+        // Uses server React.
+        ownerStacks.push(normalizeCodeLocInfo(ReactServer.captureOwnerStack()));
+      },
     );
 
     let transport;
@@ -2946,14 +2955,20 @@ describe('ReactFlight', () => {
       ReactServer = require('react');
       ReactNoopFlightServer = require('react-noop-renderer/flight-server');
       transport = ReactNoopFlightServer.render({
-        root: ReactServer.createElement(ServerComponent),
+        root: ReactServer.createElement(App),
       });
     }).toErrorDev('err');
 
     expect(mockConsoleLog).toHaveBeenCalledTimes(1);
     expect(mockConsoleLog.mock.calls[0][0]).toBe('hi');
     expect(mockConsoleLog.mock.calls[0][1].prop).toBe(123);
+    expect(ownerStacks).toEqual(['\n    in App (at **)']);
     mockConsoleLog.mockClear();
+    mockConsoleLog.mockImplementation(() => {
+      // Switching to client React.
+      ownerStacks.push(normalizeCodeLocInfo(React.captureOwnerStack()));
+    });
+    ownerStacks = [];
 
     // The error should not actually get logged because we're not awaiting the root
     // so it's not thrown but the server log also shouldn't be replayed.
@@ -2973,6 +2988,8 @@ describe('ReactFlight', () => {
     expect(typeof loggedFn2).toBe('function');
     expect(loggedFn2).not.toBe(foo);
     expect(loggedFn2.toString()).toBe(foo.toString());
+
+    expect(ownerStacks).toEqual(['\n    in App (at **)']);
   });
 
   it('uses the server component debug info as the element owner in DEV', async () => {
@@ -3159,18 +3176,18 @@ describe('ReactFlight', () => {
     jest.resetModules();
     jest.mock('react', () => React);
     ReactNoopFlightClient.read(transport);
-    assertConsoleErrorDev(
-      [
-        'Each child in a list should have a unique "key" prop.' +
-          ' See https://react.dev/link/warning-keys for more information.',
-        'Error objects cannot be rendered as text children. Try formatting it using toString().\n' +
-          '  <div>Womp womp: {Error}</div>\n' +
-          '                  ^^^^^^^',
-      ],
-      // We should have a stack in the replay but we don't yet set the owner from the Flight replaying
-      // so our simulated polyfill doesn't end up getting any component stacks yet.
-      {withoutStack: true},
-    );
+    assertConsoleErrorDev([
+      'Each child in a list should have a unique "key" prop.' +
+        ' See https://react.dev/link/warning-keys for more information.\n' +
+        '    in Bar (at **)\n' +
+        '    in App (at **)',
+      'Error objects cannot be rendered as text children. Try formatting it using toString().\n' +
+        '  <div>Womp womp: {Error}</div>\n' +
+        '                  ^^^^^^^\n' +
+        '    in Foo (at **)\n' +
+        '    in Bar (at **)\n' +
+        '    in App (at **)',
+    ]);
   });
 
   it('can filter out stack frames of a serialized error in dev', async () => {

filename: packages/react-server/src/ReactFlightServer.js
patch: @@ -99,7 +99,7 @@ import {DefaultAsyncDispatcher} from './flight/ReactFlightAsyncDispatcher';
 
 import {resolveOwner, setCurrentOwner} from './flight/ReactFlightCurrentOwner';
 
-import {getOwnerStackByComponentInfoInDev} from './flight/ReactFlightComponentStack';
+import {getOwnerStackByComponentInfoInDev} from 'shared/ReactComponentInfoStack';
 
 import {
   callComponentInDEV,
@@ -968,6 +968,7 @@ function callWithDebugContextInDEV<A, T>(
   // a fake owner during this callback so we can get the stack trace from it.
   // This also gets sent to the client as the owner for the replaying log.
   const componentDebugInfo: ReactComponentInfo = {
+    name: '',
     env: task.environmentName,
     owner: task.debugOwner,
   };
@@ -2063,6 +2064,23 @@ function escapeStringValue(value: string): string {
   }
 }
 
+function isReactComponentInfo(value: any): boolean {
+  // TODO: We don't currently have a brand check on ReactComponentInfo. Reconsider.
+  return (
+    ((typeof value.debugTask === 'object' &&
+      value.debugTask !== null &&
+      // $FlowFixMe[method-unbinding]
+      typeof value.debugTask.run === 'function') ||
+      value.debugStack instanceof Error) &&
+    (enableOwnerStacks
+      ? isArray((value: any).stack)
+      : typeof (value: any).stack === 'undefined') &&
+    typeof value.name === 'string' &&
+    typeof value.env === 'string' &&
+    value.owner !== undefined
+  );
+}
+
 let modelRoot: null | ReactClientValue = false;
 
 function renderModel(
@@ -2574,28 +2592,15 @@ function renderModelDestructive(
       );
     }
     if (__DEV__) {
-      if (
-        // TODO: We don't currently have a brand check on ReactComponentInfo. Reconsider.
-        ((typeof value.debugTask === 'object' &&
-          value.debugTask !== null &&
-          // $FlowFixMe[method-unbinding]
-          typeof value.debugTask.run === 'function') ||
-          value.debugStack instanceof Error) &&
-        (enableOwnerStacks
-          ? isArray((value: any).stack)
-          : typeof (value: any).stack === 'undefined') &&
-        typeof value.name === 'string' &&
-        typeof value.env === 'string' &&
-        value.owner !== undefined
-      ) {
+      if (isReactComponentInfo(value)) {
         // This looks like a ReactComponentInfo. We can't serialize the ConsoleTask object so we
         // need to omit it before serializing.
         const componentDebugInfo: Omit<
           ReactComponentInfo,
           'debugTask' | 'debugStack',
         > = {
-          name: value.name,
-          env: value.env,
+          name: (value: any).name,
+          env: (value: any).env,
           owner: (value: any).owner,
         };
         if (enableOwnerStacks) {
@@ -3259,6 +3264,24 @@ function renderConsoleValue(
       return Array.from((value: any));
     }
 
+    if (isReactComponentInfo(value)) {
+      // This looks like a ReactComponentInfo. We can't serialize the ConsoleTask object so we
+      // need to omit it before serializing.
+      const componentDebugInfo: Omit<
+        ReactComponentInfo,
+        'debugTask' | 'debugStack',
+      > = {
+        name: (value: any).name,
+        env: (value: any).env,
+        owner: (value: any).owner,
+      };
+      if (enableOwnerStacks) {
+        // $FlowFixMe[cannot-write]
+        componentDebugInfo.stack = (value: any).stack;
+      }
+      return componentDebugInfo;
+    }
+
     // $FlowFixMe[incompatible-return]
     return value;
   }