[Fast Refresh] Support injecting runtime after renderer executes

packages/react-devtools-shared/src/hook.js

@@ -277,6 +277,8 @@ export function installHook(target: any): DevToolsHook | null {
   const hook: DevToolsHook = {
     rendererInterfaces,
     listeners,
+
+    // Fast Refresh for web relies on this.
     renderers,
 
     emit,

packages/react-refresh/src/ReactFreshRuntime.js

@@ -435,6 +435,7 @@ export function injectIntoGlobalHook(globalObject: any): void {
       // Otherwise, the renderer will think that there is no global hook, and won't do the injection.
       let nextID = 0;
       globalObject.__REACT_DEVTOOLS_GLOBAL_HOOK__ = hook = {
+        renderers: new Map(),
         supportsFiber: true,
         inject(injected) {
           return nextID++;
@@ -468,6 +469,19 @@ export function injectIntoGlobalHook(globalObject: any): void {
       return id;
     };
 
+    // Do the same for any already injected roots.
+    // This is useful if ReactDOM has already been initialized.
+    // https://github.com/facebook/react/issues/17626
+    hook.renderers.forEach((injected, id) => {
+      if (
+        typeof injected.scheduleRefresh === 'function' &&
+        typeof injected.setRefreshHandler === 'function'
+      ) {
+        // This version supports React Refresh.
+        helpersByRendererID.set(id, ((injected: any): RendererHelpers));
+      }
+    });
+
     // We also want to track currently mounted roots.
     const oldOnCommitFiberRoot = hook.onCommitFiberRoot;
     const oldOnScheduleFiberRoot = hook.onScheduleFiberRoot || (() => {});

packages/react-refresh/src/__tests__/ReactFresh-test.js

@@ -42,6 +42,7 @@ describe('ReactFresh', () => {
 
   afterEach(() => {
     if (__DEV__) {
+      delete global.__REACT_DEVTOOLS_GLOBAL_HOOK__;
       document.body.removeChild(container);
     }
   });
@@ -3707,4 +3708,80 @@ describe('ReactFresh', () => {
       // For example, we can use this to print a log of what was updated.
     }
   });
+
+  // This simulates the scenario in https://github.com/facebook/react/issues/17626.
+  it('can inject the runtime after the renderer executes', () => {
+    if (__DEV__) {
+      // This is a minimal shim for the global hook installed by DevTools.
+      // The real one is in packages/react-devtools-shared/src/hook.js.
+      let idCounter = 0;
+      let renderers = new Map();
+      global.__REACT_DEVTOOLS_GLOBAL_HOOK__ = {
+        renderers,
+        supportsFiber: true,
+        inject(renderer) {
+          const id = ++idCounter;
+          renderers.set(id, renderer);
+          return id;
+        },
+        onCommitFiberRoot() {},
+        onCommitFiberUnmount() {},
+      };
+
+      // Load these first, as if they're coming from a CDN.
+      jest.resetModules();
+      React = require('react');
+      ReactDOM = require('react-dom');
+      Scheduler = require('scheduler');
+      act = require('react-dom/test-utils').act;
+
+      // Important! Inject into the global hook *after* ReactDOM runs:
+      ReactFreshRuntime = require('react-refresh/runtime');
+      ReactFreshRuntime.injectIntoGlobalHook(global);
+
+      // We're verifying that we're able to track roots mounted after this point.
+      // The rest of this test is taken from the simplest first test case.
+
+      render(() => {
+        function Hello() {
+          const [val, setVal] = React.useState(0);
+          return (
+            <p style={{color: 'blue'}} onClick={() => setVal(val + 1)}>
+              {val}
+            </p>
+          );
+        }
+        $RefreshReg$(Hello, 'Hello');
+        return Hello;
+      });
+
+      // Bump the state before patching.
+      const el = container.firstChild;
+      expect(el.textContent).toBe('0');
+      expect(el.style.color).toBe('blue');
+      act(() => {
+        el.dispatchEvent(new MouseEvent('click', {bubbles: true}));
+      });
+      expect(el.textContent).toBe('1');
+
+      // Perform a hot update.
+      patch(() => {
+        function Hello() {
+          const [val, setVal] = React.useState(0);
+          return (
+            <p style={{color: 'red'}} onClick={() => setVal(val + 1)}>
+              {val}
+            </p>
+          );
+        }
+        $RefreshReg$(Hello, 'Hello');
+        return Hello;
+      });
+
+      // Assert the state was preserved but color changed.
+      expect(container.firstChild).toBe(el);
+      expect(el.textContent).toBe('1');
+      expect(el.style.color).toBe('red');
+    }
+  });
 });