From 47311e0ec8a9d958860a5f44ec4fb0a0ba9fc6c7 Mon Sep 17 00:00:00 2001
From: Mark Dalgleish <mark.john.dalgleish@gmail.com>
Date: Tue, 21 Mar 2023 05:46:10 +1100
Subject: [PATCH] fix: handle CSS bundle updates during HMR (#5823)

---
 .changeset/css-bundle-hmr.md                  |   6 +
 integration/hmr-test.ts                       |  24 ++-
 packages/remix-css-bundle/browser.ts          |  10 +-
 packages/remix-dev/compiler/assets.ts         |   6 +-
 packages/remix-dev/compiler/compileBrowser.ts | 178 +++++++++++-------
 .../compiler/plugins/cssBundleUpdatePlugin.ts |  74 ++++++++
 6 files changed, 219 insertions(+), 79 deletions(-)
 create mode 100644 .changeset/css-bundle-hmr.md
 create mode 100644 packages/remix-dev/compiler/plugins/cssBundleUpdatePlugin.ts

diff --git a/.changeset/css-bundle-hmr.md b/.changeset/css-bundle-hmr.md
new file mode 100644
index 00000000000..08fd8503b76
--- /dev/null
+++ b/.changeset/css-bundle-hmr.md
@@ -0,0 +1,6 @@
+---
+"@remix-run/css-bundle": patch
+"@remix-run/dev": patch
+---
+
+Ensure changes to CSS inserted via `@remix-run/css-bundle` are picked up during HMR
diff --git a/integration/hmr-test.ts b/integration/hmr-test.ts
index 13af69b0261..bd08171da42 100644
--- a/integration/hmr-test.ts
+++ b/integration/hmr-test.ts
@@ -13,6 +13,7 @@ let fixture = (options: { port: number; appServerPort: number }) => ({
       port: options.port,
       appServerPort: options.appServerPort,
     },
+    unstable_cssModules: true,
     unstable_tailwind: true,
     v2_routeConvention: true,
     v2_errorBoundary: true,
@@ -26,6 +27,7 @@ let fixture = (options: { port: number; appServerPort: number }) => ({
         "dev:app": `cross-env NODE_ENV=development nodemon --watch build/ ./server.js`,
       },
       dependencies: {
+        "@remix-run/css-bundle": "0.0.0-local-version",
         "@remix-run/node": "0.0.0-local-version",
         "@remix-run/react": "0.0.0-local-version",
         "cross-env": "0.0.0-local-version",
@@ -90,15 +92,23 @@ let fixture = (options: { port: number; appServerPort: number }) => ({
       @tailwind utilities;
     `,
 
+    "app/styles.module.css": css`
+      .test {
+        color: initial;
+      }
+    `,
+
     "app/root.tsx": js`
       import type { LinksFunction } from "@remix-run/node";
       import { Link, Links, LiveReload, Meta, Outlet, Scripts } from "@remix-run/react";
+      import { cssBundleHref } from "@remix-run/css-bundle";
 
       import Counter from "./components/counter";
       import styles from "./tailwind.css";
 
       export const links: LinksFunction = () => [
         { rel: "stylesheet", href: styles },
+        ...cssBundleHref ? [{ rel: "stylesheet", href: cssBundleHref }] : [],
       ];
 
       export default function Root() {
@@ -246,15 +256,26 @@ test("HMR", async ({ page }) => {
     let originalIndex = fs.readFileSync(indexPath, "utf8");
     let counterPath = path.join(projectDir, "app", "components", "counter.tsx");
     let originalCounter = fs.readFileSync(counterPath, "utf8");
+    let cssModulePath = path.join(projectDir, "app", "styles.module.css");
+    let originalCssModule = fs.readFileSync(cssModulePath, "utf8");
 
     // make content and style changed to index route
+    let newCssModule = `
+      .test {
+        background: black;
+        color: white;
+      }
+    `;
+    fs.writeFileSync(cssModulePath, newCssModule);
+
     let newIndex = `
       import { useLoaderData } from "@remix-run/react";
+      import styles from "~/styles.module.css";
       export default function Index() {
         const t = useLoaderData();
         return (
           <main>
-            <h1 className="text-white bg-black">Changed</h1>
+            <h1 className={styles.test}>Changed</h1>
           </main>
         )
       }
@@ -274,6 +295,7 @@ test("HMR", async ({ page }) => {
 
     // undo change
     fs.writeFileSync(indexPath, originalIndex);
+    fs.writeFileSync(cssModulePath, originalCssModule);
     await page.getByText("Index Title").waitFor({ timeout: 2000 });
     expect(await page.getByLabel("Root Input").inputValue()).toBe("asdfasdf");
     await page.waitForSelector(`#root-counter:has-text("inc 1")`);
diff --git a/packages/remix-css-bundle/browser.ts b/packages/remix-css-bundle/browser.ts
index e17c8ce3d10..02c76cfa629 100644
--- a/packages/remix-css-bundle/browser.ts
+++ b/packages/remix-css-bundle/browser.ts
@@ -2,4 +2,12 @@ import type { AssetsManifest } from "@remix-run/dev/assets-manifest";
 
 let assetsManifest: AssetsManifest = (window as any).__remixManifest;
 
-export const cssBundleHref = assetsManifest.cssBundleHref;
+declare const __INJECT_CSS_BUNDLE_HREF__: string | undefined;
+
+// Injected by `cssBundleUpdatePlugin` on rebuilds
+let updatedHref: string | undefined =
+  typeof __INJECT_CSS_BUNDLE_HREF__ === "string"
+    ? __INJECT_CSS_BUNDLE_HREF__
+    : undefined;
+
+export const cssBundleHref = updatedHref || assetsManifest.cssBundleHref;
diff --git a/packages/remix-dev/compiler/assets.ts b/packages/remix-dev/compiler/assets.ts
index f8c427704cc..fad890502b5 100644
--- a/packages/remix-dev/compiler/assets.ts
+++ b/packages/remix-dev/compiler/assets.ts
@@ -42,12 +42,12 @@ export interface AssetsManifest {
 export async function createAssetsManifest({
   config,
   metafile,
-  cssBundlePath,
+  cssBundleHref,
   hmr,
 }: {
   config: RemixConfig;
   metafile: esbuild.Metafile;
-  cssBundlePath?: string;
+  cssBundleHref?: string;
   hmr?: AssetsManifest["hmr"];
 }): Promise<AssetsManifest> {
   function resolveUrl(outputPath: string): string {
@@ -126,8 +126,6 @@ export async function createAssetsManifest({
     JSON.stringify({ entry, routes, hmrRoutes: hmr?.routes })
   ).slice(0, 8);
 
-  let cssBundleHref = cssBundlePath ? resolveUrl(cssBundlePath) : undefined;
-
   return { version, entry, routes, cssBundleHref, hmr };
 }
 
diff --git a/packages/remix-dev/compiler/compileBrowser.ts b/packages/remix-dev/compiler/compileBrowser.ts
index ff5d00a9e0e..8692c7e089a 100644
--- a/packages/remix-dev/compiler/compileBrowser.ts
+++ b/packages/remix-dev/compiler/compileBrowser.ts
@@ -6,7 +6,8 @@ import { NodeModulesPolyfillPlugin } from "@esbuild-plugins/node-modules-polyfil
 import postcss from "postcss";
 import postcssDiscardDuplicates from "postcss-discard-duplicates";
 
-import type { WriteChannel } from "../channel";
+import type { Channel, WriteChannel } from "../channel";
+import { createChannel } from "../channel";
 import type { RemixConfig } from "../config";
 import type { AssetsManifest } from "./assets";
 import { createAssetsManifest } from "./assets";
@@ -20,6 +21,7 @@ import { deprecatedRemixPackagePlugin } from "./plugins/deprecatedRemixPackagePl
 import { emptyModulesPlugin } from "./plugins/emptyModulesPlugin";
 import { mdxPlugin } from "./plugins/mdx";
 import { urlImportsPlugin } from "./plugins/urlImportsPlugin";
+import { cssBundleUpdatePlugin } from "./plugins/cssBundleUpdatePlugin";
 import { cssModulesPlugin } from "./plugins/cssModulesPlugin";
 import { cssSideEffectImportsPlugin } from "./plugins/cssSideEffectImportsPlugin";
 import { vanillaExtractPlugin } from "./plugins/vanillaExtractPlugin";
@@ -78,6 +80,12 @@ const isCssBundlingEnabled = (config: RemixConfig): boolean =>
       config.future.unstable_cssSideEffectImports ||
       config.future.unstable_vanillaExtract
   );
+
+let cssBundleHrefChannel: Channel<string | undefined>;
+
+// This function gives esbuild access to the latest channel value on rebuilds
+let getCssBundleHref = () => cssBundleHrefChannel.read();
+
 const createEsbuildConfig = (
   build: "app" | "css",
   config: RemixConfig,
@@ -131,7 +139,7 @@ const createEsbuildConfig = (
     NodeModulesPolyfillPlugin(),
   ].filter(isNotNull);
 
-  if (mode === "development" && config.future.unstable_dev) {
+  if (build === "app" && mode === "development" && config.future.unstable_dev) {
     // TODO prebundle deps instead of chunking just these ones
     let isolateChunks = [
       require.resolve("react"),
@@ -149,6 +157,10 @@ const createEsbuildConfig = (
     };
 
     plugins.push(hmrPlugin({ remixConfig: config }));
+
+    if (isCssBundlingEnabled(config)) {
+      plugins.push(cssBundleUpdatePlugin({ getCssBundleHref }));
+    }
   }
 
   return {
@@ -234,82 +246,102 @@ export const createBrowserCompiler = (
         return;
       }
 
-      // The types aren't great when combining write: false and incremental: true
-      //  so we need to assert that it's an incremental build
-      cssCompiler = (await (!cssCompiler
-        ? esbuild.build({
-            ...createEsbuildConfig("css", remixConfig, options, onLoader),
-            metafile: true,
-            incremental: true,
-            write: false,
-          })
-        : cssCompiler.rebuild())) as esbuild.BuildIncremental;
-
-      invariant(
-        cssCompiler.metafile,
-        "Expected CSS compiler metafile to be defined. This is likely a bug in Remix. Please open an issue at https://github.com/remix-run/remix/issues/new"
-      );
-
-      let outputFiles = cssCompiler.outputFiles || [];
-
-      let isCssBundleFile = (
-        outputFile: esbuild.OutputFile,
-        extension: ".css" | ".css.map"
-      ): boolean => {
-        return (
-          path.dirname(outputFile.path) === remixConfig.assetsBuildDirectory &&
-          path.basename(outputFile.path).startsWith("css-bundle") &&
-          outputFile.path.endsWith(extension)
+      try {
+        // The types aren't great when combining write: false and incremental: true
+        //  so we need to assert that it's an incremental build
+        cssCompiler = (await (!cssCompiler
+          ? esbuild.build({
+              ...createEsbuildConfig("css", remixConfig, options, onLoader),
+              metafile: true,
+              incremental: true,
+              write: false,
+            })
+          : cssCompiler.rebuild())) as esbuild.BuildIncremental;
+
+        invariant(
+          cssCompiler.metafile,
+          "Expected CSS compiler metafile to be defined. This is likely a bug in Remix. Please open an issue at https://github.com/remix-run/remix/issues/new"
         );
-      };
 
-      let cssBundleFile = outputFiles.find((outputFile) =>
-        isCssBundleFile(outputFile, ".css")
-      );
+        let outputFiles = cssCompiler.outputFiles || [];
+
+        let isCssBundleFile = (
+          outputFile: esbuild.OutputFile,
+          extension: ".css" | ".css.map"
+        ): boolean => {
+          return (
+            path.dirname(outputFile.path) ===
+              remixConfig.assetsBuildDirectory &&
+            path.basename(outputFile.path).startsWith("css-bundle") &&
+            outputFile.path.endsWith(extension)
+          );
+        };
+
+        let cssBundleFile = outputFiles.find((outputFile) =>
+          isCssBundleFile(outputFile, ".css")
+        );
 
-      if (!cssBundleFile) {
-        return;
+        if (!cssBundleFile) {
+          cssBundleHrefChannel.write(undefined);
+          return;
+        }
+
+        let cssBundlePath = cssBundleFile.path;
+
+        let cssBundleHref =
+          remixConfig.publicPath +
+          path.relative(
+            remixConfig.assetsBuildDirectory,
+            path.resolve(cssBundlePath)
+          );
+
+        cssBundleHrefChannel.write(cssBundleHref);
+
+        let { css, map } = await postcss([
+          // We need to discard duplicate rules since "composes"
+          // in CSS Modules can result in duplicate styles
+          postcssDiscardDuplicates(),
+        ]).process(cssBundleFile.text, {
+          from: cssBundlePath,
+          to: cssBundlePath,
+          map: options.sourcemap && {
+            prev: outputFiles.find((outputFile) =>
+              isCssBundleFile(outputFile, ".css.map")
+            )?.text,
+            inline: false,
+            annotation: false,
+            sourcesContent: true,
+          },
+        });
+
+        await fse.ensureDir(path.dirname(cssBundlePath));
+
+        await Promise.all([
+          fse.writeFile(cssBundlePath, css),
+          options.mode !== "production" && map
+            ? fse.writeFile(`${cssBundlePath}.map`, map.toString()) // Write our updated source map rather than esbuild's
+            : null,
+          ...outputFiles
+            .filter((outputFile) => !/\.(css|js|map)$/.test(outputFile.path))
+            .map(async (asset) => {
+              await fse.ensureDir(path.dirname(asset.path));
+              await fse.writeFile(asset.path, asset.contents);
+            }),
+        ]);
+
+        return cssBundleHref;
+      } catch (error) {
+        cssBundleHrefChannel.write(undefined);
+        throw error;
       }
-
-      let cssBundlePath = cssBundleFile.path;
-
-      let { css, map } = await postcss([
-        // We need to discard duplicate rules since "composes"
-        // in CSS Modules can result in duplicate styles
-        postcssDiscardDuplicates(),
-      ]).process(cssBundleFile.text, {
-        from: cssBundlePath,
-        to: cssBundlePath,
-        map: options.sourcemap && {
-          prev: outputFiles.find((outputFile) =>
-            isCssBundleFile(outputFile, ".css.map")
-          )?.text,
-          inline: false,
-          annotation: false,
-          sourcesContent: true,
-        },
-      });
-
-      await fse.ensureDir(path.dirname(cssBundlePath));
-
-      await Promise.all([
-        fse.writeFile(cssBundlePath, css),
-        options.mode !== "production" && map
-          ? fse.writeFile(`${cssBundlePath}.map`, map.toString()) // Write our updated source map rather than esbuild's
-          : null,
-        ...outputFiles
-          .filter((outputFile) => !/\.(css|js|map)$/.test(outputFile.path))
-          .map(async (asset) => {
-            await fse.ensureDir(path.dirname(asset.path));
-            await fse.writeFile(asset.path, asset.contents);
-          }),
-      ]);
-
-      // Return the CSS bundle path so we can use it to generate the manifest
-      return cssBundlePath;
     };
 
-    let [cssBundlePath, metafile] = await Promise.all([
+    // Reset the channel to co-ordinate the CSS and app builds
+    if (isCssBundlingEnabled(remixConfig)) {
+      cssBundleHrefChannel = createChannel();
+    }
+
+    let [cssBundleHref, metafile] = await Promise.all([
       cssBuildTask(),
       appBuildTask(),
     ]);
@@ -336,7 +368,7 @@ export const createBrowserCompiler = (
     let manifest = await createAssetsManifest({
       config: remixConfig,
       metafile: appCompiler.metafile!,
-      cssBundlePath,
+      cssBundleHref,
       hmr,
     });
     await writeAssetsManifest(remixConfig, manifest);
diff --git a/packages/remix-dev/compiler/plugins/cssBundleUpdatePlugin.ts b/packages/remix-dev/compiler/plugins/cssBundleUpdatePlugin.ts
new file mode 100644
index 00000000000..6fc9298f6d9
--- /dev/null
+++ b/packages/remix-dev/compiler/plugins/cssBundleUpdatePlugin.ts
@@ -0,0 +1,74 @@
+import type { Plugin } from "esbuild";
+import { readFile } from "fs-extra";
+
+const pluginName = "css-bundle-update-plugin";
+const namespace = `${pluginName}-ns`;
+
+/**
+ * This plugin updates the source code for the "css-bundle" package on rebuilds
+ * to contain the latest CSS bundle href so CSS changes get picked up for HMR.
+ * Without this plugin, the "css-bundle" package source code never changes on
+ * disk so it never triggers an update.
+ */
+export function cssBundleUpdatePlugin({
+  getCssBundleHref,
+}: {
+  getCssBundleHref: () => Promise<string | undefined>;
+}): Plugin {
+  return {
+    name: pluginName,
+    async setup(build) {
+      let isRebuild = false;
+      build.onEnd(() => {
+        isRebuild = true;
+      });
+
+      let preventInfiniteLoop = {};
+      build.onResolve({ filter: /^@remix-run\/css-bundle$/ }, async (args) => {
+        // Prevent plugin from infinitely trying to resolve itself
+        if (args.pluginData === preventInfiniteLoop) {
+          return null;
+        }
+
+        // We don't wait for the href on the first build and instead rely on the
+        // default runtime manifest lookup. We only need to update this package
+        // to reflect changes during development so the first build is fine.
+        if (!isRebuild) {
+          return null;
+        }
+
+        let resolvedPath = (
+          await build.resolve(args.path, {
+            resolveDir: args.resolveDir,
+            kind: args.kind,
+            pluginData: preventInfiniteLoop,
+          })
+        ).path;
+
+        return {
+          path: resolvedPath,
+          namespace,
+        };
+      });
+
+      build.onLoad({ filter: /.*/, namespace }, async (args) => {
+        let [cssBundleHref, contents] = await Promise.all([
+          getCssBundleHref(),
+          readFile(args.path, "utf8"),
+        ]);
+
+        if (cssBundleHref) {
+          contents = contents.replace(
+            /__INJECT_CSS_BUNDLE_HREF__/g,
+            JSON.stringify(cssBundleHref)
+          );
+        }
+
+        return {
+          loader: "js",
+          contents,
+        };
+      });
+    },
+  };
+}