[Flight] Implement captureStackTrace and owner stacks on the Server

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

@@ -81,6 +81,7 @@ let ErrorBoundary;
 let NoErrorExpected;
 let Scheduler;
 let assertLog;
+let assertConsoleErrorDev;
 
 describe('ReactFlight', () => {
   beforeEach(() => {
@@ -102,6 +103,7 @@ describe('ReactFlight', () => {
     Scheduler = require('scheduler');
     const InternalTestUtils = require('internal-test-utils');
     assertLog = InternalTestUtils.assertLog;
+    assertConsoleErrorDev = InternalTestUtils.assertConsoleErrorDev;
 
     ErrorBoundary = class extends React.Component {
       state = {hasError: false, error: null};
@@ -1441,9 +1443,7 @@ describe('ReactFlight', () => {
         <div>{Array(6).fill(<NoKey />)}</div>,
       );
       ReactNoopFlightClient.read(transport);
-    }).toErrorDev('Each child in a list should have a unique "key" prop.', {
-      withoutStack: gate(flags => flags.enableOwnerStacks),
-    });
+    }).toErrorDev('Each child in a list should have a unique "key" prop.');
   });
 
   it('should warn in DEV a child is missing keys in client component', async () => {
@@ -2728,4 +2728,76 @@ describe('ReactFlight', () => {
 
     expect(ReactNoop).toMatchRenderedOutput(<span>Hello, Seb</span>);
   });
+
+  // @gate __DEV__ && enableOwnerStacks
+  it('can get the component owner stacks during rendering in dev', () => {
+    let stack;
+
+    function Foo() {
+      return ReactServer.createElement(Bar, null);
+    }
+    function Bar() {
+      return ReactServer.createElement(
+        'div',
+        null,
+        ReactServer.createElement(Baz, null),
+      );
+    }
+
+    function Baz() {
+      stack = ReactServer.captureOwnerStack();
+      return ReactServer.createElement('span', null, 'hi');
+    }
+    ReactNoopFlightServer.render(
+      ReactServer.createElement(
+        'div',
+        null,
+        ReactServer.createElement(Foo, null),
+      ),
+    );
+
+    expect(normalizeCodeLocInfo(stack)).toBe(
+      '\n    in Bar (at **)' + '\n    in Foo (at **)',
+    );
+  });
+
+  // @gate (enableOwnerStacks && enableServerComponentLogs) || !__DEV__
+  it('should not include component stacks in replayed logs (unless DevTools add them)', () => {
+    function Foo() {
+      return 'hi';
+    }
+
+    function Bar() {
+      const array = [];
+      // Trigger key warning
+      array.push(ReactServer.createElement(Foo));
+      return ReactServer.createElement('div', null, array);
+    }
+
+    function App() {
+      return ReactServer.createElement(Bar);
+    }
+
+    const transport = ReactNoopFlightServer.render(
+      ReactServer.createElement(App),
+    );
+    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 **)',
+    ]);
+
+    // Replay logs on the client
+    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.',
+      ],
+      // We should not have a stack in the replay because that should be added either by console.createTask
+      // or React DevTools on the client. Neither of which we do here.
+      {withoutStack: true},
+    );
+  });
 });

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

@@ -39,6 +39,15 @@ let ReactServerDOMServer;
 let ReactServerDOMClient;
 let use;
 
+function normalizeCodeLocInfo(str) {
+  return (
+    str &&
+    str.replace(/^ +(?:at|in) ([\S]+)[^\n]*/gm, function (m, name) {
+      return '    in ' + name + (/\d/.test(m) ? ' (at **)' : '');
+    })
+  );
+}
+
 describe('ReactFlightDOMEdge', () => {
   beforeEach(() => {
     jest.resetModules();
@@ -883,4 +892,42 @@ describe('ReactFlightDOMEdge', () => {
       );
     }
   });
+
+  // @gate __DEV__ && enableOwnerStacks
+  it('can get the component owner stacks asynchronously', async () => {
+    let stack;
+
+    function Foo() {
+      return ReactServer.createElement(Bar, null);
+    }
+    function Bar() {
+      return ReactServer.createElement(
+        'div',
+        null,
+        ReactServer.createElement(Baz, null),
+      );
+    }
+
+    const promise = Promise.resolve(0);
+
+    async function Baz() {
+      await promise;
+      stack = ReactServer.captureOwnerStack();
+      return ReactServer.createElement('span', null, 'hi');
+    }
+
+    const stream = ReactServerDOMServer.renderToReadableStream(
+      ReactServer.createElement(
+        'div',
+        null,
+        ReactServer.createElement(Foo, null),
+      ),
+      webpackMap,
+    );
+    await readResult(stream);
+
+    expect(normalizeCodeLocInfo(stack)).toBe(
+      '\n    in Bar (at **)' + '\n    in Foo (at **)',
+    );
+  });
 });

packages/react-server/src/ReactFlightServer.js

@@ -97,6 +97,10 @@ import {DefaultAsyncDispatcher} from './flight/ReactFlightAsyncDispatcher';
 
 import {resolveOwner, setCurrentOwner} from './flight/ReactFlightCurrentOwner';
 
+import {getOwnerStackByComponentInfoInDev} from './flight/ReactFlightComponentStack';
+
+import {isWritingAppendedStack} from 'shared/consoleWithStackDev';
+
 import {
   getIteratorFn,
   REACT_ELEMENT_TYPE,
@@ -263,8 +267,9 @@ function patchConsole(consoleInst: typeof console, methodName: string) {
       'name',
     );
     const wrapperMethod = function (this: typeof console) {
+      let args = arguments;
       const request = resolveRequest();
-      if (methodName === 'assert' && arguments[0]) {
+      if (methodName === 'assert' && args[0]) {
         // assert doesn't emit anything unless first argument is falsy so we can skip it.
       } else if (request !== null) {
         // Extract the stack. Not all console logs print the full stack but they have at
@@ -276,7 +281,22 @@ function patchConsole(consoleInst: typeof console, methodName: string) {
         // refer to previous logs in debug info to associate them with a component.
         const id = request.nextChunkId++;
         const owner: null | ReactComponentInfo = resolveOwner();
-        emitConsoleChunk(request, id, methodName, owner, stack, arguments);
+        if (
+          isWritingAppendedStack &&
+          (methodName === 'error' || methodName === 'warn') &&
+          args.length > 1 &&
+          typeof args[0] === 'string' &&
+          args[0].endsWith('%s')
+        ) {
+          // This looks like we've appended the component stack to the error from our own logs.
+          // We don't want those added to the replayed logs since those have the opportunity to add
+          // their own stacks or use console.createTask on the client as needed.
+          // TODO: Remove this special case once we remove consoleWithStackDev.
+          // $FlowFixMe[method-unbinding]
+          args = Array.prototype.slice.call(args, 0, args.length - 1);
+          args[0] = args[0].slice(0, args[0].length - 2);
+        }
+        emitConsoleChunk(request, id, methodName, owner, stack, args);
       }
       // $FlowFixMe[prop-missing]
       return originalMethod.apply(this, arguments);
@@ -317,6 +337,21 @@ if (
   patchConsole(console, 'warn');
 }
 
+function getCurrentStackInDEV(): string {
+  if (__DEV__) {
+    if (enableOwnerStacks) {
+      const owner: null | ReactComponentInfo = resolveOwner();
+      if (owner === null) {
+        return '';
+      }
+      return getOwnerStackByComponentInfoInDev(owner);
+    }
+    // We don't have Parent Stacks in Flight.
+    return '';
+  }
+  return '';
+}
+
 const ObjectPrototype = Object.prototype;
 
 type JSONValue =
@@ -491,6 +526,12 @@ function RequestInstance(
     );
   }
   ReactSharedInternals.A = DefaultAsyncDispatcher;
+  if (__DEV__) {
+    // Unlike Fizz or Fiber, we don't reset this and just keep it on permanently.
+    // This lets it act more like the AsyncDispatcher so that we can get the
+    // stack asynchronously too.
+    ReactSharedInternals.getCurrentStack = getCurrentStackInDEV;
+  }
 
   const abortSet: Set<Task> = new Set();
   const pingedTasks: Array<Task> = [];

packages/react-server/src/flight/ReactFlightComponentStack.js

@@ -0,0 +1,52 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @flow
+ */
+
+import type {ReactComponentInfo} from 'shared/ReactTypes';
+
+import {describeBuiltInComponentFrame} from 'shared/ReactComponentStackFrame';
+
+import {enableOwnerStacks} from 'shared/ReactFeatureFlags';
+
+export function getOwnerStackByComponentInfoInDev(
+  componentInfo: ReactComponentInfo,
+): string {
+  if (!enableOwnerStacks || !__DEV__) {
+    return '';
+  }
+  try {
+    let info = '';
+
+    // The owner stack of the current component will be where it was created, i.e. inside its owner.
+    // There's no actual name of the currently executing component. Instead, that is available
+    // on the regular stack that's currently executing. However, if there is no owner at all, then
+    // there's no stack frame so we add the name of the root component to the stack to know which
+    // component is currently executing.
+    if (!componentInfo.owner && typeof componentInfo.name === 'string') {
+      return describeBuiltInComponentFrame(componentInfo.name);
+    }
+
+    let owner: void | null | ReactComponentInfo = componentInfo;
+
+    while (owner) {
+      if (typeof owner.stack === 'string') {
+        // Server Component
+        const ownerStack: string = owner.stack;
+        owner = owner.owner;
+        if (owner && ownerStack !== '') {
+          info += '\n' + ownerStack;
+        }
+      } else {
+        break;
+      }
+    }
+    return info;
+  } catch (x) {
+    return '\nError generating stack: ' + x.message + '\n' + x.stack;
+  }
+}

packages/shared/consoleWithStackDev.js

@@ -40,6 +40,8 @@ export function error(format, ...args) {
 // eslint-disable-next-line react-internal/no-production-logging
 const supportsCreateTask = __DEV__ && enableOwnerStacks && !!console.createTask;
 
+export let isWritingAppendedStack = false;
+
 function printWarning(level, format, args, currentStack) {
   // When changing this logic, you might want to also
   // update consoleWithStackDev.www.js as well.
@@ -50,6 +52,7 @@ function printWarning(level, format, args, currentStack) {
       // can be lost while DevTools isn't open but we can't detect this.
       const stack = ReactSharedInternals.getCurrentStack(currentStack);
       if (stack !== '') {
+        isWritingAppendedStack = true;
         format += '%s';
         args = args.concat([stack]);
       }
@@ -60,5 +63,6 @@ function printWarning(level, format, args, currentStack) {
     // breaks IE9: https://github.com/facebook/react/issues/13610
     // eslint-disable-next-line react-internal/no-production-logging
     Function.prototype.apply.call(console[level], console, args);
+    isWritingAppendedStack = false;
   }
 }