From 34a11fefb2ddb1c7a47e91e565a0ca4b50e48446 Mon Sep 17 00:00:00 2001
From: Pedro Cattori <pcattori@gmail.com>
Date: Wed, 7 Feb 2024 15:31:38 -0500
Subject: [PATCH] Vite: Enable HMR for .md and .mdx files

Co-authored-by: Hiroshi Ogawa <hi.ogawa.zz@gmail.com>
---
 .changeset/wet-bulldogs-appear.md |  5 +++
 integration/helpers/vite.ts       |  1 +
 integration/vite-hmr-hdr-test.ts  | 56 +++++++++++++++++++++++++++++++
 packages/remix-dev/vite/plugin.ts |  3 +-
 4 files changed, 64 insertions(+), 1 deletion(-)
 create mode 100644 .changeset/wet-bulldogs-appear.md

diff --git a/.changeset/wet-bulldogs-appear.md b/.changeset/wet-bulldogs-appear.md
new file mode 100644
index 00000000000..a9982ad366e
--- /dev/null
+++ b/.changeset/wet-bulldogs-appear.md
@@ -0,0 +1,5 @@
+---
+"@remix-run/dev": patch
+---
+
+Vite: Enable HMR for .md and .mdx files
diff --git a/integration/helpers/vite.ts b/integration/helpers/vite.ts
index ba9b58bcd37..3d6fd1ed531 100644
--- a/integration/helpers/vite.ts
+++ b/integration/helpers/vite.ts
@@ -211,6 +211,7 @@ export const test = base.extend<Fixtures>({
     await use(async (files) => {
       let port = await getPort();
       let cwd = await createProject(await files({ port }));
+      console.log("cwd", cwd);
       stop = await viteDev({ cwd, port });
       return { port, cwd };
     });
diff --git a/integration/vite-hmr-hdr-test.ts b/integration/vite-hmr-hdr-test.ts
index 215561fb528..f4e78f614a8 100644
--- a/integration/vite-hmr-hdr-test.ts
+++ b/integration/vite-hmr-hdr-test.ts
@@ -57,6 +57,62 @@ test("Vite / HMR & HDR / express", async ({ page, browserName, customDev }) => {
   await workflow({ page, browserName, cwd, port });
 });
 
+test("Vite / HMR & HDR / mdx", async ({ page, viteDev }) => {
+  let files: Files = async ({ port }) => ({
+    "vite.config.ts": `
+      import { defineConfig } from "vite";
+      import { unstable_vitePlugin as remix } from "@remix-run/dev";
+      import mdx from "@mdx-js/rollup";
+
+      export default defineConfig({
+        ${await viteConfig.server({ port })}
+        plugins: [
+          mdx(),
+          remix(),
+        ],
+      });
+    `,
+    "app/component.tsx": `
+      import {useState} from "react";
+
+      export const Counter = () => {
+        const [count, setCount] = useState(0);
+        return <button onClick={() => setCount(count => count + 1)}>Count: {count}</button>
+      }
+    `,
+    "app/routes/mdx.mdx": `
+      import { Counter } from "../component";
+
+      # MDX Title (HMR: 0)
+
+      <Counter />
+    `,
+  });
+
+  let { port, cwd } = await viteDev(files);
+  let edit = createEditor(cwd);
+  await page.goto(`http://localhost:${port}/mdx`, {
+    waitUntil: "networkidle",
+  });
+
+  await expect(page.locator("h1")).toHaveText("MDX Title (HMR: 0)");
+  let button = page.locator("button");
+  await expect(button).toHaveText("Count: 0");
+  await button.click();
+  await expect(button).toHaveText("Count: 1");
+
+  // TODO edit
+  await edit("app/routes/mdx.mdx", (contents) =>
+    contents.replace("(HMR: 0)", "(HMR: 1)")
+  );
+  await page.waitForLoadState("networkidle");
+
+  await expect(page.locator("h1")).toHaveText("MDX Title (HMR: 1)");
+  await expect(page.locator("button")).toHaveText("Count: 1");
+
+  expect(page.errors).toEqual([]);
+});
+
 async function workflow({
   page,
   browserName,
diff --git a/packages/remix-dev/vite/plugin.ts b/packages/remix-dev/vite/plugin.ts
index e901947013c..b684b453d24 100644
--- a/packages/remix-dev/vite/plugin.ts
+++ b/packages/remix-dev/vite/plugin.ts
@@ -1491,7 +1491,8 @@ export const remixVitePlugin: RemixVitePlugin = (remixUserConfig = {}) => {
         if (id.includes("/node_modules/")) return;
 
         let [filepath] = id.split("?");
-        if (!/.[tj]sx?$/.test(filepath)) return;
+        let extensionsRE = /\.(jsx?|tsx?|mdx?)$/;
+        if (!extensionsRE.test(filepath)) return;
 
         let devRuntime = "react/jsx-dev-runtime";
         let ssr = options?.ssr === true;