[Flight] Enable owner stacks on the client when replaying logs

packages/react-client/src/ReactFlightClient.js

@@ -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 {
@@ -2282,6 +2295,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,
@@ -2310,34 +2339,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(

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

@@ -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 () => {

packages/react-server/src/ReactFlightServer.js

@@ -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;
   }