[Flight] Add "use ..." boundary after the change instead of before it (#33478)

I noticed that the ThirdPartyComponent in the fixture was showing the
wrong stack and the `"use third-party"` is in the wrong location.

<img width="628" alt="Screenshot 2025-06-06 at 11 22 11 PM"
src="https://github.com/user-attachments/assets/f0013380-d79e-4765-b371-87fd61b3056b"
/>

When creating the initial JSX inside the third party server, we should
make sure that it has no owner. In a real cross-server environment you
get this by default by just executing in different context. But since
the fixture example is inside the same AsyncLocalStorage as the parent
it already has an owner which gets transferred. So we should make sure
that were we create the JSX has no owner to simulate this.

When we then parse a null owner on the receiving side, we replace its
owner/stack with the owner/stack of the call to `createFrom...` to
connect them. This worked fine with only two environments. The bug was
that when we did this and then transferred the result to a third
environment we took the original parsed stack trace. We should instead
parse a new one from the replaced stack in the current environment.

The second bug was that the `"use third-party"` badge ends up in the
wrong place when we do this kind of thing. Because the stack of the
thing entering the new environment is the call to `createFrom...` which
is in the old environment even though the component itself executes in
the new environment. So to see if there's a change we should be
comparing the current environment of the task to the owner's environment
instead of the next environment after the task.

After:

<img width="494" alt="Screenshot 2025-06-07 at 1 13 28 AM"
src="https://github.com/user-attachments/assets/e2e870ba-f125-4526-a853-bd29f164cf09"
/>

Author: sebmarkbage

fixtures/flight/src/App.js

@@ -33,20 +33,26 @@ function Foo({children}) {
   return <div>{children}</div>;
 }
 
+async function delay(text, ms) {
+  return new Promise(resolve => setTimeout(() => resolve(text), ms));
+}
+
 async function Bar({children}) {
-  await new Promise(resolve => setTimeout(() => resolve('deferred text'), 10));
+  await delay('deferred text', 10);
   return <div>{children}</div>;
 }
 
 async function ThirdPartyComponent() {
-  return new Promise(resolve =>
-    setTimeout(() => resolve('hello from a 3rd party'), 30)
-  );
+  return delay('hello from a 3rd party', 30);
 }
 
 // Using Web streams for tee'ing convenience here.
 let cachedThirdPartyReadableWeb;
 
+// We create the Component outside of AsyncLocalStorage so that it has no owner.
+// That way it gets the owner from the call to createFromNodeStream.
+const thirdPartyComponent = <ThirdPartyComponent />;
+
 function fetchThirdParty(noCache) {
   if (cachedThirdPartyReadableWeb && !noCache) {
     const [readableWeb1, readableWeb2] = cachedThirdPartyReadableWeb.tee();
@@ -59,7 +65,7 @@ function fetchThirdParty(noCache) {
   }
 
   const stream = renderToPipeableStream(
-    <ThirdPartyComponent />,
+    thirdPartyComponent,
     {},
     {environmentName: 'third-party'}
   );
@@ -80,8 +86,8 @@ function fetchThirdParty(noCache) {
 }
 
 async function ServerComponent({noCache}) {
-  await new Promise(resolve => setTimeout(() => resolve('deferred text'), 50));
-  return fetchThirdParty(noCache);
+  await delay('deferred text', 50);
+  return await fetchThirdParty(noCache);
 }
 
 export default async function App({prerender, noCache}) {

packages/react-client/src/ReactFlightClient.js

@@ -844,7 +844,7 @@ function createElement(
       // This owner should ideally have already been initialized to avoid getting
       // user stack frames on the stack.
       const ownerTask =
-        owner === null ? null : initializeFakeTask(response, owner, env);
+        owner === null ? null : initializeFakeTask(response, owner);
       if (ownerTask === null) {
         const rootTask = response._debugRootTask;
         if (rootTask != null) {
@@ -2494,7 +2494,6 @@ function getRootTask(
 function initializeFakeTask(
   response: Response,
   debugInfo: ReactComponentInfo | ReactAsyncInfo | ReactIOInfo,
-  childEnvironmentName: string,
 ): null | ConsoleTask {
   if (!supportsCreateTask) {
     return null;
@@ -2504,6 +2503,10 @@ function initializeFakeTask(
     // If it's null, we can't initialize a task.
     return null;
   }
+  const cachedEntry = debugInfo.debugTask;
+  if (cachedEntry !== undefined) {
+    return cachedEntry;
+  }
 
   // Workaround for a bug where Chrome Performance tracking uses the enclosing line/column
   // instead of the callsite. For ReactAsyncInfo/ReactIOInfo, the only thing we're going
@@ -2516,47 +2519,35 @@ function initializeFakeTask(
   const stack = debugInfo.stack;
   const env: string =
     debugInfo.env == null ? response._rootEnvironmentName : debugInfo.env;
-  if (env !== childEnvironmentName) {
+  const ownerEnv: string =
+    debugInfo.owner == null || debugInfo.owner.env == null
+      ? response._rootEnvironmentName
+      : debugInfo.owner.env;
+  const ownerTask =
+    debugInfo.owner == null
+      ? null
+      : initializeFakeTask(response, debugInfo.owner);
+  const taskName =
     // This is the boundary between two environments so we'll annotate the task name.
-    // That is unusual so we don't cache it.
-    const ownerTask =
-      debugInfo.owner == null
-        ? null
-        : initializeFakeTask(response, debugInfo.owner, env);
-    return buildFakeTask(
-      response,
-      ownerTask,
-      stack,
-      '"use ' + childEnvironmentName.toLowerCase() + '"',
-      env,
-      useEnclosingLine,
-    );
-  } else {
-    const cachedEntry = debugInfo.debugTask;
-    if (cachedEntry !== undefined) {
-      return cachedEntry;
-    }
-    const ownerTask =
-      debugInfo.owner == null
-        ? null
-        : initializeFakeTask(response, debugInfo.owner, env);
-    // Some unfortunate pattern matching to refine the type.
-    const taskName =
-      debugInfo.key !== undefined
+    // We assume that the stack frame of the entry into the new environment was done
+    // from the old environment. So we use the owner's environment as the current.
+    env !== ownerEnv
+      ? '"use ' + env.toLowerCase() + '"'
+      : // Some unfortunate pattern matching to refine the type.
+        debugInfo.key !== undefined
         ? getServerComponentTaskName(((debugInfo: any): ReactComponentInfo))
         : debugInfo.name !== undefined
           ? getIOInfoTaskName(((debugInfo: any): ReactIOInfo))
           : getAsyncInfoTaskName(((debugInfo: any): ReactAsyncInfo));
-    // $FlowFixMe[cannot-write]: We consider this part of initialization.
-    return (debugInfo.debugTask = buildFakeTask(
-      response,
-      ownerTask,
-      stack,
-      taskName,
-      env,
-      useEnclosingLine,
-    ));
-  }
+  // $FlowFixMe[cannot-write]: We consider this part of initialization.
+  return (debugInfo.debugTask = buildFakeTask(
+    response,
+    ownerTask,
+    stack,
+    taskName,
+    ownerEnv,
+    useEnclosingLine,
+  ));
 }
 
 function buildFakeTask(
@@ -2658,27 +2649,30 @@ function resolveDebugInfo(
       'resolveDebugInfo should never be called in production mode. This is a bug in React.',
     );
   }
-  // We eagerly initialize the fake task because this resolving happens outside any
-  // render phase so we're not inside a user space stack at this point. If we waited
-  // to initialize it when we need it, we might be inside user code.
-  const env =
-    debugInfo.env === undefined ? response._rootEnvironmentName : debugInfo.env;
   if (debugInfo.stack !== undefined) {
     const componentInfoOrAsyncInfo: ReactComponentInfo | ReactAsyncInfo =
       // $FlowFixMe[incompatible-type]
       debugInfo;
-    initializeFakeTask(response, componentInfoOrAsyncInfo, env);
+    // We eagerly initialize the fake task because this resolving happens outside any
+    // render phase so we're not inside a user space stack at this point. If we waited
+    // to initialize it when we need it, we might be inside user code.
+    initializeFakeTask(response, componentInfoOrAsyncInfo);
   }
-  if (debugInfo.owner === null && response._debugRootOwner != null) {
+  if (debugInfo.owner == null && response._debugRootOwner != null) {
     const componentInfoOrAsyncInfo: ReactComponentInfo | ReactAsyncInfo =
       // $FlowFixMe: By narrowing `owner` to `null`, we narrowed `debugInfo` to `ReactComponentInfo`
       debugInfo;
     // $FlowFixMe[cannot-write]
     componentInfoOrAsyncInfo.owner = response._debugRootOwner;
+    // We clear the parsed stack frames to indicate that it needs to be re-parsed from debugStack.
+    // $FlowFixMe[cannot-write]
+    componentInfoOrAsyncInfo.stack = 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.
     // $FlowFixMe[cannot-write]
     componentInfoOrAsyncInfo.debugStack = response._debugRootStack;
+    // $FlowFixMe[cannot-write]
+    componentInfoOrAsyncInfo.debugTask = response._debugRootTask;
   } else if (debugInfo.stack !== undefined) {
     const componentInfoOrAsyncInfo: ReactComponentInfo | ReactAsyncInfo =
       // $FlowFixMe[incompatible-type]
@@ -2738,7 +2732,7 @@ const replayConsoleWithCallStack = {
         bindToConsole(methodName, args, env),
       );
       if (owner != null) {
-        const task = initializeFakeTask(response, owner, env);
+        const task = initializeFakeTask(response, owner);
         initializeFakeStack(response, owner);
         if (task !== null) {
           task.run(callStack);
@@ -2812,10 +2806,8 @@ function resolveConsoleEntry(
 }
 
 function initializeIOInfo(response: Response, ioInfo: ReactIOInfo): void {
-  const env =
-    ioInfo.env === undefined ? response._rootEnvironmentName : ioInfo.env;
   if (ioInfo.stack !== undefined) {
-    initializeFakeTask(response, ioInfo, env);
+    initializeFakeTask(response, ioInfo);
     initializeFakeStack(response, ioInfo);
   }
   // Adjust the time to the current environment's time space.

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

@@ -320,7 +320,6 @@ describe('ReactFlight', () => {
                 name: 'Greeting',
                 env: 'Server',
                 key: null,
-                owner: null,
                 stack: '    in Object.<anonymous> (at **)',
                 props: {
                   firstName: 'Seb',
@@ -364,7 +363,6 @@ describe('ReactFlight', () => {
                 name: 'Greeting',
                 env: 'Server',
                 key: null,
-                owner: null,
                 stack: '    in Object.<anonymous> (at **)',
                 props: {
                   firstName: 'Seb',
@@ -2812,7 +2810,6 @@ describe('ReactFlight', () => {
                 name: 'ServerComponent',
                 env: 'Server',
                 key: null,
-                owner: null,
                 stack: '    in Object.<anonymous> (at **)',
                 props: {
                   transport: expect.arrayContaining([]),
@@ -2834,7 +2831,6 @@ describe('ReactFlight', () => {
                 name: 'ThirdPartyComponent',
                 env: 'third-party',
                 key: null,
-                owner: null,
                 stack: '    in Object.<anonymous> (at **)',
                 props: {},
               },
@@ -2851,7 +2847,6 @@ describe('ReactFlight', () => {
                 name: 'ThirdPartyLazyComponent',
                 env: 'third-party',
                 key: null,
-                owner: null,
                 stack: '    in myLazy (at **)\n    in lazyInitializer (at **)',
                 props: {},
               },
@@ -2867,7 +2862,6 @@ describe('ReactFlight', () => {
                 name: 'ThirdPartyFragmentComponent',
                 env: 'third-party',
                 key: '3',
-                owner: null,
                 stack: '    in Object.<anonymous> (at **)',
                 props: {},
               },
@@ -2941,7 +2935,6 @@ describe('ReactFlight', () => {
                 name: 'ServerComponent',
                 env: 'Server',
                 key: null,
-                owner: null,
                 stack: '    in Object.<anonymous> (at **)',
                 props: {
                   transport: expect.arrayContaining([]),
@@ -2961,7 +2954,6 @@ describe('ReactFlight', () => {
                 name: 'Keyed',
                 env: 'Server',
                 key: 'keyed',
-                owner: null,
                 stack: '    in ServerComponent (at **)',
                 props: {
                   children: {},
@@ -2980,7 +2972,6 @@ describe('ReactFlight', () => {
                 name: 'ThirdPartyAsyncIterableComponent',
                 env: 'third-party',
                 key: null,
-                owner: null,
                 stack: '    in Object.<anonymous> (at **)',
                 props: {},
               },
@@ -3137,7 +3128,6 @@ describe('ReactFlight', () => {
                 name: 'Component',
                 env: 'A',
                 key: null,
-                owner: null,
                 stack: '    in Object.<anonymous> (at **)',
                 props: {},
               },
@@ -3325,7 +3315,6 @@ describe('ReactFlight', () => {
           name: 'Greeting',
           env: 'Server',
           key: null,
-          owner: null,
           stack: '    in Object.<anonymous> (at **)',
           props: {
             firstName: 'Seb',

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

@@ -708,7 +708,6 @@ describe('ReactFlightDOMBrowser', () => {
             name: 'Server',
             env: 'Server',
             key: null,
-            owner: null,
           }),
         }),
       );
@@ -724,7 +723,6 @@ describe('ReactFlightDOMBrowser', () => {
             name: 'Server',
             env: 'Server',
             key: null,
-            owner: null,
           }),
         }),
       );

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

@@ -1190,7 +1190,6 @@ describe('ReactFlightDOMEdge', () => {
       const greetInfo = expect.objectContaining({
         name: 'Greeting',
         env: 'Server',
-        owner: null,
       });
       expect(lazyWrapper._debugInfo).toEqual([
         {time: 12},