From b4e299354a5cbf87ab47ec379d65ff0cf42c67b4 Mon Sep 17 00:00:00 2001
From: Marc MacLeod <marbemac@gmail.com>
Date: Thu, 2 Nov 2023 14:46:39 -0500
Subject: [PATCH] fix(vite): fix Fast Refresh load error when using defer
 (#7842)

Signed-off-by: Marc MacLeod <847542+marbemac@users.noreply.github.com>
Co-authored-by: Mark Dalgleish <mark.john.dalgleish@gmail.com>
---
 .changeset/sharp-kids-bake.md     |  5 ++++
 integration/vite-dev-test.ts      | 39 ++++++++++++++++++++++++-------
 packages/remix-dev/vite/plugin.ts | 32 +++++++++++++++++--------
 3 files changed, 58 insertions(+), 18 deletions(-)
 create mode 100644 .changeset/sharp-kids-bake.md

diff --git a/.changeset/sharp-kids-bake.md b/.changeset/sharp-kids-bake.md
new file mode 100644
index 00000000000..bd6455aacdb
--- /dev/null
+++ b/.changeset/sharp-kids-bake.md
@@ -0,0 +1,5 @@
+---
+"@remix-run/dev": patch
+---
+
+Fix React Fast Refresh error on load when using `defer` in Vite dev server
diff --git a/integration/vite-dev-test.ts b/integration/vite-dev-test.ts
index 6e446e19e6c..c9dd3cb2b7b 100644
--- a/integration/vite-dev-test.ts
+++ b/integration/vite-dev-test.ts
@@ -59,20 +59,28 @@ test.describe("Vite dev", () => {
           }
         `,
         "app/routes/_index.tsx": js`
-          import { useState, useEffect } from "react";
+          import { Suspense } from "react";
+          import { defer } from "@remix-run/node";
+          import { Await, useLoaderData } from "@remix-run/react";
+
+          export function loader() {
+            let deferred = new Promise((resolve) => {
+              setTimeout(() => resolve(true), 1000)
+            });
+            return defer({ deferred });
+          }
 
           export default function IndexRoute() {
-            const [mounted, setMounted] = useState(false);
-            useEffect(() => {
-              setMounted(true);
-            }, []);
+            const { deferred } = useLoaderData<typeof loader>();
 
             return (
               <div id="index">
                 <h2 data-title>Index</h2>
                 <input />
-                <p data-mounted>Mounted: {mounted ? "yes" : "no"}</p>
                 <p data-hmr>HMR updated: no</p>
+                <Suspense fallback={<p data-defer>Defer finished: no</p>}>
+                  <Await resolve={deferred}>{() => <p data-defer>Defer finished: yes</p>}</Await>
+                </Suspense>
               </div>
             );
           }
@@ -161,12 +169,19 @@ test.describe("Vite dev", () => {
   });
 
   test("renders matching routes", async ({ page }) => {
+    let pageErrors: unknown[] = [];
+    page.on("pageerror", (error) => pageErrors.push(error));
+
     await page.goto(`http://localhost:${devPort}/`, {
       waitUntil: "networkidle",
     });
+
+    // Ensure no errors on page load
+    expect(pageErrors).toEqual([]);
+
     await expect(page.locator("#index [data-title]")).toHaveText("Index");
-    await expect(page.locator("#index [data-mounted]")).toHaveText(
-      "Mounted: yes"
+    await expect(page.locator("#index [data-defer]")).toHaveText(
+      "Defer finished: yes"
     );
 
     let hmrStatus = page.locator("#index [data-hmr]");
@@ -188,13 +203,21 @@ test.describe("Vite dev", () => {
     await page.waitForLoadState("networkidle");
     await expect(hmrStatus).toHaveText("HMR updated: yes");
     await expect(input).toHaveValue("stateful");
+
+    // Ensure no errors after HMR
+    expect(pageErrors).toEqual([]);
   });
 
   test("handles multiple set-cookie headers", async ({ page }) => {
+    let pageErrors: Error[] = [];
+    page.on("pageerror", (error) => pageErrors.push(error));
+
     await page.goto(`http://localhost:${devPort}/set-cookies`, {
       waitUntil: "networkidle",
     });
 
+    expect(pageErrors).toEqual([]);
+
     // Ensure we redirected
     expect(new URL(page.url()).pathname).toBe("/get-cookies");
 
diff --git a/packages/remix-dev/vite/plugin.ts b/packages/remix-dev/vite/plugin.ts
index b1821387318..56703a8acf6 100644
--- a/packages/remix-dev/vite/plugin.ts
+++ b/packages/remix-dev/vite/plugin.ts
@@ -73,6 +73,7 @@ let serverManifestId = VirtualModule.id("server-manifest");
 let browserManifestId = VirtualModule.id("browser-manifest");
 let remixReactProxyId = VirtualModule.id("remix-react-proxy");
 let hmrRuntimeId = VirtualModule.id("hmr-runtime");
+let injectHmrRuntimeId = VirtualModule.id("inject-hmr-runtime");
 
 const normalizePath = (p: string) => {
   let unixPath = p.replace(/[\\/]+/g, "/").replace(/^([a-zA-Z]+:|\.\/)/, "");
@@ -689,21 +690,32 @@ export const remixVitePlugin: RemixVitePlugin = (options = {}) => {
             'export const LiveReload = process.env.NODE_ENV !== "development" ? () => null : ',
             '() => createElement("script", {',
             ' type: "module",',
-            " suppressHydrationWarning: true,",
-            " dangerouslySetInnerHTML: { __html: `",
-            `   import RefreshRuntime from "${VirtualModule.url(
-              hmrRuntimeId
-            )}"`,
-            "   RefreshRuntime.injectIntoGlobalHook(window)",
-            "   window.$RefreshReg$ = () => {}",
-            "   window.$RefreshSig$ = () => (type) => type",
-            "   window.__vite_plugin_react_preamble_installed__ = true",
-            " `}",
+            " async: true,",
+            ` src: "${VirtualModule.url(injectHmrRuntimeId)}"`,
             "});",
           ].join("\n");
         }
       },
     },
+    {
+      name: "remix-inject-hmr-runtime",
+      enforce: "pre",
+      resolveId(id) {
+        if (id === injectHmrRuntimeId)
+          return VirtualModule.resolve(injectHmrRuntimeId);
+      },
+      async load(id) {
+        if (id !== VirtualModule.resolve(injectHmrRuntimeId)) return;
+
+        return [
+          `import RefreshRuntime from "${hmrRuntimeId}"`,
+          "RefreshRuntime.injectIntoGlobalHook(window)",
+          "window.$RefreshReg$ = () => {}",
+          "window.$RefreshSig$ = () => (type) => type",
+          "window.__vite_plugin_react_preamble_installed__ = true",
+        ].join("\n");
+      },
+    },
     {
       name: "remix-hmr-runtime",
       enforce: "pre",