Create a root task for every Flight response (#29673)

This lets any element created from the server, to bottom out with a
client "owner" which is the creator of the Flight request. This could be
a Server Action being invoked or a router.

This is similar to how a client element bottoms out in the creator of
the root element without an owner. E.g. where the root app element was
created.

Without this, we inherit the task of whatever is currently executing
when we're parsing which can be misleading.

Before:
<img width="507" alt="Screenshot 2024-05-30 at 12 06 57 PM"
src="https://github.com/facebook/react/assets/63648/e234db7e-67f7-404c-958a-5c5500ffdf1f">

After:
<img width="555" alt="Screenshot 2024-05-30 at 4 59 04 PM"
src="https://github.com/facebook/react/assets/63648/8ba6acb4-2ffd-49d4-bd44-08228ad4200e">

The before/after doesn't show much of a difference here but that's just
because our Flight parsing loop is an async, which maybe it shouldn't be
because it can be unnecessarily deep, and it creates a hidden line for
every loop. That's what the `Promise.then` is. If the element is lazily
initialized it's worse because we can end up in an unrelated render task
as the owner - although that's its own problem.

Author: sebmarkbage

packages/react-client/src/ReactFlightClient.js

@@ -254,6 +254,7 @@ 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
+  _debugRootTask?: null | ConsoleTask, // DEV-only
 };
 
 function readChunk<T>(chunk: SomeChunk<T>): T {
@@ -614,6 +615,7 @@ function getTaskName(type: mixed): string {
 }
 
 function createElement(
+  response: Response,
   type: mixed,
   key: mixed,
   props: mixed,
@@ -697,9 +699,15 @@ function createElement(
         const callStack = buildFakeCallStack(stack, createTaskFn);
         // This owner should ideally have already been initialized to avoid getting
         // user stack frames on the stack.
-        const ownerTask = owner === null ? null : initializeFakeTask(owner);
+        const ownerTask =
+          owner === null ? null : initializeFakeTask(response, owner);
         if (ownerTask === null) {
-          task = callStack();
+          const rootTask = response._debugRootTask;
+          if (rootTask != null) {
+            task = rootTask.run(callStack);
+          } else {
+            task = callStack();
+          }
         } else {
           task = ownerTask.run(callStack);
         }
@@ -1106,6 +1114,7 @@ function parseModelTuple(
     // TODO: Consider having React just directly accept these arrays as elements.
     // Or even change the ReactElement type to be an array.
     return createElement(
+      response,
       tuple[1],
       tuple[2],
       tuple[3],
@@ -1149,6 +1158,14 @@ export function createResponse(
     _buffer: [],
     _tempRefs: temporaryReferences,
   };
+  if (supportsCreateTask) {
+    // Any stacks that appear on the server need to be rooted somehow on the client
+    // so we create a root Task for this response which will be the root owner for any
+    // elements created by the server. We use the "use server" string to indicate that
+    // this is where we enter the server from the client.
+    // TODO: Make this string configurable.
+    response._debugRootTask = (console: any).createTask('"use server"');
+  }
   // Don't inline this call because it causes closure to outline the call above.
   response._fromJSON = createFromJSONCallback(response);
   return response;
@@ -1730,6 +1747,7 @@ function buildFakeCallStack<T>(stack: string, innerCall: () => T): () => T {
 }
 
 function initializeFakeTask(
+  response: Response,
   debugInfo: ReactComponentInfo | ReactAsyncInfo,
 ): null | ConsoleTask {
   if (taskCache === null || typeof debugInfo.stack !== 'string') {
@@ -1745,7 +1763,7 @@ function initializeFakeTask(
   const ownerTask =
     componentInfo.owner == null
       ? null
-      : initializeFakeTask(componentInfo.owner);
+      : initializeFakeTask(response, componentInfo.owner);
 
   // eslint-disable-next-line react-internal/no-production-logging
   const createTaskFn = (console: any).createTask.bind(
@@ -1755,7 +1773,12 @@ function initializeFakeTask(
   const callStack = buildFakeCallStack(stack, createTaskFn);
 
   if (ownerTask === null) {
-    return callStack();
+    const rootTask = response._debugRootTask;
+    if (rootTask != null) {
+      return rootTask.run(callStack);
+    } else {
+      return callStack();
+    }
   } else {
     return ownerTask.run(callStack);
   }
@@ -1776,7 +1799,7 @@ function resolveDebugInfo(
   // 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(debugInfo);
+  initializeFakeTask(response, debugInfo);
   const chunk = getChunk(response, id);
   const chunkDebugInfo: ReactDebugInfo =
     chunk._debugInfo || (chunk._debugInfo = []);
@@ -1813,12 +1836,17 @@ function resolveConsoleEntry(
     printToConsole.bind(null, methodName, args, env),
   );
   if (owner != null) {
-    const task = initializeFakeTask(owner);
+    const task = initializeFakeTask(response, owner);
     if (task !== null) {
       task.run(callStack);
       return;
     }
   }
+  const rootTask = response._debugRootTask;
+  if (rootTask != null) {
+    rootTask.run(callStack);
+    return;
+  }
   callStack();
 }