From 03614111fa9474c1bd05cb84ffa755b7f08daae6 Mon Sep 17 00:00:00 2001
From: Sebastian Markbage <sebastian@calyptus.eu>
Date: Tue, 20 Feb 2024 12:50:48 -0500
Subject: [PATCH] Don't replay consoles written inside onError/onPostpone.

These aren't conceptually part of the request's render so we exit the
request context for those.
---
 .../src/__tests__/ReactFlight-test.js         | 41 +++++++++++++++++++
 .../react-server/src/ReactFlightServer.js     | 31 ++++++++++++--
 2 files changed, 68 insertions(+), 4 deletions(-)

diff --git a/packages/react-client/src/__tests__/ReactFlight-test.js b/packages/react-client/src/__tests__/ReactFlight-test.js
index 06dee1d54ed22..dd98c7b5b5df5 100644
--- a/packages/react-client/src/__tests__/ReactFlight-test.js
+++ b/packages/react-client/src/__tests__/ReactFlight-test.js
@@ -1978,4 +1978,45 @@ describe('ReactFlight', () => {
       </div>,
     );
   });
+
+  // @gate enableServerComponentLogs
+  it('replays logs, but not onError logs', async () => {
+    function foo() {
+      return 'hello';
+    }
+    function ServerComponent() {
+      console.log('hi', {prop: 123, fn: foo});
+      throw new Error('err');
+    }
+
+    let transport;
+    expect(() => {
+      // Reset the modules so that we get a new overridden console on top of the
+      // one installed by expect. This ensures that we still emit console.error
+      // calls.
+      jest.resetModules();
+      jest.mock('react', () => require('react/react.react-server'));
+      ReactServer = require('react');
+      ReactNoopFlightServer = require('react-noop-renderer/flight-server');
+      transport = ReactNoopFlightServer.render({root: <ServerComponent />});
+    }).toErrorDev('err');
+
+    const log = console.log;
+    try {
+      console.log = jest.fn();
+      // 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.
+      await ReactNoopFlightClient.read(transport);
+
+      expect(console.log).toHaveBeenCalledTimes(1);
+      expect(console.log.mock.calls[0][0]).toBe('hi');
+      expect(console.log.mock.calls[0][1].prop).toBe(123);
+      const loggedFn = console.log.mock.calls[0][1].fn;
+      expect(typeof loggedFn).toBe('function');
+      expect(loggedFn).not.toBe(foo);
+      expect(loggedFn.toString()).toBe(foo.toString());
+    } finally {
+      console.log = log;
+    }
+  });
 });
diff --git a/packages/react-server/src/ReactFlightServer.js b/packages/react-server/src/ReactFlightServer.js
index 547165228c1fc..dee92019324cd 100644
--- a/packages/react-server/src/ReactFlightServer.js
+++ b/packages/react-server/src/ReactFlightServer.js
@@ -1689,13 +1689,36 @@ function renderModelDestructive(
 }
 
 function logPostpone(request: Request, reason: string): void {
-  const onPostpone = request.onPostpone;
-  onPostpone(reason);
+  const prevRequest = currentRequest;
+  currentRequest = null;
+  try {
+    const onPostpone = request.onPostpone;
+    if (supportsRequestStorage) {
+      // Exit the request context while running callbacks.
+      requestStorage.run(undefined, onPostpone, reason);
+    } else {
+      onPostpone(reason);
+    }
+  } finally {
+    currentRequest = prevRequest;
+  }
 }
 
 function logRecoverableError(request: Request, error: mixed): string {
-  const onError = request.onError;
-  const errorDigest = onError(error);
+  const prevRequest = currentRequest;
+  currentRequest = null;
+  let errorDigest;
+  try {
+    const onError = request.onError;
+    if (supportsRequestStorage && false) {
+      // Exit the request context while running callbacks.
+      errorDigest = requestStorage.run(undefined, onError, error);
+    } else {
+      errorDigest = onError(error);
+    }
+  } finally {
+    currentRequest = prevRequest;
+  }
   if (errorDigest != null && typeof errorDigest !== 'string') {
     // eslint-disable-next-line react-internal/prod-error-codes
     throw new Error(