[Flight] Let environmentName vary over time by making it a function of string

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

@@ -2565,6 +2565,50 @@ describe('ReactFlight', () => {
     );
   });
 
+  it('can change the environment name inside a component', async () => {
+    let env = 'A';
+    function Component(props) {
+      env = 'B';
+      return <div>hi</div>;
+    }
+
+    const transport = ReactNoopFlightServer.render(
+      {
+        greeting: <Component />,
+      },
+      {
+        environmentName() {
+          return env;
+        },
+      },
+    );
+
+    await act(async () => {
+      const rootModel = await ReactNoopFlightClient.read(transport);
+      const greeting = rootModel.greeting;
+      expect(getDebugInfo(greeting)).toEqual(
+        __DEV__
+          ? [
+              {
+                name: 'Component',
+                env: 'A',
+                owner: null,
+                stack: gate(flag => flag.enableOwnerStacks)
+                  ? '    in Object.<anonymous> (at **)'
+                  : undefined,
+              },
+              {
+                env: 'B',
+              },
+            ]
+          : undefined,
+      );
+      ReactNoop.render(greeting);
+    });
+
+    expect(ReactNoop).toMatchRenderedOutput(<div>hi</div>);
+  });
+
   // @gate enableServerComponentLogs && __DEV__
   it('replays logs, but not onError logs', async () => {
     function foo() {

packages/react-noop-renderer/src/ReactNoopFlightServer.js

@@ -67,7 +67,7 @@ const ReactNoopFlightServer = ReactFlightServer({
 });
 
 type Options = {
-  environmentName?: string,
+  environmentName?: string | (() => string),
   identifierPrefix?: string,
   onError?: (error: mixed) => void,
   onPostpone?: (reason: string) => void,

packages/react-server-dom-esm/src/ReactFlightDOMServerNode.js

@@ -66,7 +66,7 @@ function createCancelHandler(request: Request, reason: string) {
 }
 
 type Options = {
-  environmentName?: string,
+  environmentName?: string | (() => string),
   onError?: (error: mixed) => void,
   onPostpone?: (reason: string) => void,
   identifierPrefix?: string,

packages/react-server-dom-turbopack/src/ReactFlightDOMServerBrowser.js

@@ -44,7 +44,7 @@ export {createTemporaryReferenceSet} from 'react-server/src/ReactFlightServerTem
 export type {TemporaryReferenceSet};
 
 type Options = {
-  environmentName?: string,
+  environmentName?: string | (() => string),
   identifierPrefix?: string,
   signal?: AbortSignal,
   temporaryReferences?: TemporaryReferenceSet,

packages/react-server-dom-turbopack/src/ReactFlightDOMServerEdge.js

@@ -44,7 +44,7 @@ export {createTemporaryReferenceSet} from 'react-server/src/ReactFlightServerTem
 export type {TemporaryReferenceSet};
 
 type Options = {
-  environmentName?: string,
+  environmentName?: string | (() => string),
   identifierPrefix?: string,
   signal?: AbortSignal,
   temporaryReferences?: TemporaryReferenceSet,

packages/react-server-dom-turbopack/src/ReactFlightDOMServerNode.js

@@ -67,7 +67,7 @@ function createCancelHandler(request: Request, reason: string) {
 }
 
 type Options = {
-  environmentName?: string,
+  environmentName?: string | (() => string),
   onError?: (error: mixed) => void,
   onPostpone?: (reason: string) => void,
   identifierPrefix?: string,

packages/react-server-dom-webpack/src/ReactFlightDOMServerBrowser.js

@@ -44,7 +44,7 @@ export {createTemporaryReferenceSet} from 'react-server/src/ReactFlightServerTem
 export type {TemporaryReferenceSet};
 
 type Options = {
-  environmentName?: string,
+  environmentName?: string | (() => string),
   identifierPrefix?: string,
   signal?: AbortSignal,
   temporaryReferences?: TemporaryReferenceSet,

packages/react-server-dom-webpack/src/ReactFlightDOMServerEdge.js

@@ -44,7 +44,7 @@ export {createTemporaryReferenceSet} from 'react-server/src/ReactFlightServerTem
 export type {TemporaryReferenceSet};
 
 type Options = {
-  environmentName?: string,
+  environmentName?: string | (() => string),
   identifierPrefix?: string,
   signal?: AbortSignal,
   temporaryReferences?: TemporaryReferenceSet,

packages/react-server-dom-webpack/src/ReactFlightDOMServerNode.js

@@ -67,7 +67,7 @@ function createCancelHandler(request: Request, reason: string) {
 }
 
 type Options = {
-  environmentName?: string,
+  environmentName?: string | (() => string),
   onError?: (error: mixed) => void,
   onPostpone?: (reason: string) => void,
   identifierPrefix?: string,

packages/react-server/src/ReactFlightServer.js

@@ -393,6 +393,7 @@ type Task = {
   keyPath: null | string, // parent server component keys
   implicitSlot: boolean, // true if the root server component of this sequence had a null key
   thenableState: ThenableState | null,
+  environmentName: string, // DEV-only. Used to track if the environment for this task changed.
 };
 
 interface Reference {}
@@ -425,7 +426,7 @@ export type Request = {
   onError: (error: mixed) => ?string,
   onPostpone: (reason: string) => void,
   // DEV-only
-  environmentName: string,
+  environmentName: () => string,
   didWarnForKey: null | WeakSet<ReactComponentInfo>,
 };
 
@@ -481,7 +482,7 @@ function RequestInstance(
   onError: void | ((error: mixed) => ?string),
   identifierPrefix?: string,
   onPostpone: void | ((reason: string) => void),
-  environmentName: void | string,
+  environmentName: void | string | (() => string),
   temporaryReferences: void | TemporaryReferenceSet,
 ) {
   if (
@@ -531,7 +532,11 @@ function RequestInstance(
 
   if (__DEV__) {
     this.environmentName =
-      environmentName === undefined ? 'Server' : environmentName;
+      environmentName === undefined
+        ? () => 'Server'
+        : typeof environmentName !== 'function'
+        ? () => environmentName
+        : environmentName;
     this.didWarnForKey = null;
   }
   const rootTask = createTask(this, model, null, false, abortSet);
@@ -544,7 +549,7 @@ export function createRequest(
   onError: void | ((error: mixed) => ?string),
   identifierPrefix?: string,
   onPostpone: void | ((reason: string) => void),
-  environmentName: void | string,
+  environmentName: void | string | (() => string),
   temporaryReferences: void | TemporaryReferenceSet,
 ): Request {
   // $FlowFixMe[invalid-constructor]: the shapes are exact here but Flow doesn't like constructors
@@ -1049,14 +1054,14 @@ function renderFunctionComponent<Props>(
       componentDebugInfo = (prevThenableState: any)._componentDebugInfo;
     } else {
       // This is a new component in the same task so we can emit more debug info.
+      const componentDebugID = debugID;
       const componentName =
         (Component: any).displayName || Component.name || '';
+      const componentEnv = request.environmentName();
       request.pendingChunks++;
-
-      const componentDebugID = debugID;
       componentDebugInfo = ({
         name: componentName,
-        env: request.environmentName,
+        env: componentEnv,
         owner: owner,
       }: ReactComponentInfo);
       if (enableOwnerStacks) {
@@ -1069,6 +1074,9 @@ function renderFunctionComponent<Props>(
       outlineModel(request, componentDebugInfo);
       emitDebugChunk(request, componentDebugID, componentDebugInfo);
 
+      // We've emitted the latest environment for this task so we track that.
+      task.environmentName = componentEnv;
+
       if (enableOwnerStacks) {
         warnForMissingKey(request, key, validated, componentDebugInfo);
       }
@@ -1644,7 +1652,7 @@ function createTask(
       request.writtenObjects.set(model, serializeByValueID(id));
     }
   }
-  const task: Task = {
+  const task: Task = (({
     id,
     status: PENDING,
     model,
@@ -1697,7 +1705,10 @@ function createTask(
       return renderModel(request, task, parent, parentPropertyName, value);
     },
     thenableState: null,
-  };
+  }: Omit<Task, 'environmentName'>): any);
+  if (__DEV__) {
+    task.environmentName = request.environmentName();
+  }
   abortSet.add(task);
   return task;
 }
@@ -3252,7 +3263,7 @@ function emitConsoleChunk(
   }
 
   // TODO: Don't double badge if this log came from another Flight Client.
-  const env = request.environmentName;
+  const env = request.environmentName();
   const payload = [methodName, stackTrace, owner, env];
   // $FlowFixMe[method-unbinding]
   payload.push.apply(payload, args);
@@ -3420,6 +3431,15 @@ function retryTask(request: Request, task: Task): void {
       // any future references.
       request.writtenObjects.set(resolvedModel, serializeByValueID(task.id));
 
+      if (__DEV__) {
+        const currentEnv = request.environmentName();
+        if (currentEnv !== task.environmentName) {
+          // The environment changed since we last emitted any debug information for this
+          // task. We emit an entry that just includes the environment name change.
+          emitDebugChunk(request, task.id, {env: currentEnv});
+        }
+      }
+
       // Object might contain unresolved values like additional elements.
       // This is simulating what the JSON loop would do if this was part of it.
       emitChunk(request, task, resolvedModel);
@@ -3428,6 +3448,16 @@ function retryTask(request: Request, task: Task): void {
       // We don't need to escape it again so it's not passed the toJSON replacer.
       // $FlowFixMe[incompatible-type] stringify can return null for undefined but we never do
       const json: string = stringify(resolvedModel);
+
+      if (__DEV__) {
+        const currentEnv = request.environmentName();
+        if (currentEnv !== task.environmentName) {
+          // The environment changed since we last emitted any debug information for this
+          // task. We emit an entry that just includes the environment name change.
+          emitDebugChunk(request, task.id, {env: currentEnv});
+        }
+      }
+
       emitModelChunk(request, task.id, json);
     }