fix: wait for a chunk's css on every dynamic import

.changeset/preload-css-every-import.md

@@ -0,0 +1,5 @@
+---
+"@marko/vite": patch
+---
+
+Make every dynamic import of a chunk wait for its stylesheets, so a lazily loaded template imported again while its CSS is still loading no longer renders unstyled.

src/__tests__/preload-css.test.ts

@@ -0,0 +1,146 @@
+import assert from "assert";
+import fs from "fs";
+import os from "os";
+import path from "path";
+import type * as viteNamespace from "vite";
+import vm from "vm";
+
+import preloadCssPlugin, { patchPreloadHelper } from "../preload-css";
+
+const preloadHelperId = "\0vite/preload-helper.js";
+
+// Runs Vite's preload helper against a stub DOM where each stylesheet link
+// loads (or fails) only when the test says so.
+function loadHelper(code: string) {
+  const links: FakeLink[] = [];
+  class FakeLink extends EventTarget {
+    rel = "";
+    as = "";
+    href = "";
+    crossOrigin = "";
+    setAttribute() {}
+  }
+  const context = vm.createContext({
+    __VITE_IS_MODERN__: true,
+    __importMeta: {
+      url: "http://localhost/assets/main.js",
+      resolve: (specifier: string) =>
+        new URL(specifier, "http://localhost/").href,
+    },
+    document: {
+      getElementsByTagName: () => links,
+      querySelector: () => null,
+      createElement: () => new FakeLink(),
+      head: { appendChild: (link: FakeLink) => links.push(link) },
+    },
+    window: new EventTarget(),
+    Event,
+    EventTarget,
+    Promise,
+    URL,
+    Error,
+  });
+  vm.runInContext(
+    code
+      .replace("export const __vitePreload", "globalThis.__vitePreload")
+      .replaceAll("import.meta", "__importMeta"),
+    context,
+  );
+  const preload = context.__vitePreload as (
+    load: () => Promise<unknown>,
+    deps: string[],
+  ) => Promise<unknown>;
+  const importChunk = () => {
+    const state = { settled: false as false | "resolved" | "rejected" };
+    preload(() => Promise.resolve("chunk"), ["assets/lazy.css"]).then(
+      () => (state.settled = "resolved"),
+      () => (state.settled = "rejected"),
+    );
+    return state;
+  };
+  return { links, importChunk };
+}
+
+const flush = () => new Promise((resolve) => setImmediate(resolve));
+
+describe("preload css", () => {
+  let vite: typeof viteNamespace;
+  let dir: string;
+  let helper: string;
+
+  before(async () => {
+    vite = await import("vite");
+    dir = fs.mkdtempSync(path.join(os.tmpdir(), "marko-vite-preload-"));
+    fs.writeFileSync(
+      path.join(dir, "index.html"),
+      `<script type="module" src="./main.js"></script>`,
+    );
+    fs.writeFileSync(path.join(dir, "main.js"), `import("./lazy.js");`);
+    fs.writeFileSync(path.join(dir, "lazy.js"), `import "./lazy.css";`);
+    fs.writeFileSync(path.join(dir, "lazy.css"), `.lazy { color: teal }`);
+    await vite.build({
+      root: dir,
+      logLevel: "silent",
+      build: { write: false },
+      plugins: [
+        {
+          name: "capture-preload-helper",
+          transform: {
+            filter: { id: /^\0vite\/preload-helper\.js$/ },
+            handler(code, id) {
+              if (id === preloadHelperId) helper = code;
+            },
+          },
+        },
+      ],
+    });
+  });
+
+  after(() => fs.rmSync(dir, { recursive: true, force: true }));
+
+  it("patches the helper in a build", async () => {
+    const output = (await vite.build({
+      root: dir,
+      logLevel: "silent",
+      build: { write: false, minify: false },
+      plugins: [preloadCssPlugin()],
+    })) as viteNamespace.Rolldown.RolldownOutput;
+    assert.ok(
+      output.output.some(
+        (chunk) =>
+          chunk.type === "chunk" && chunk.code.includes("__markoCssLoad"),
+      ),
+    );
+  });
+
+  it("makes every import wait for a stylesheet still loading", async () => {
+    const { links, importChunk } = loadHelper(patchPreloadHelper(helper)!);
+    const first = importChunk();
+    const second = importChunk();
+    await flush();
+    assert.deepEqual([first.settled, second.settled], [false, false]);
+
+    links[0].dispatchEvent(new Event("load"));
+    await flush();
+    assert.deepEqual([first.settled, second.settled], ["resolved", "resolved"]);
+    const later = importChunk();
+    await flush();
+    assert.equal(later.settled, "resolved");
+  });
+
+  it("fails every import waiting on a stylesheet that fails", async () => {
+    const { links, importChunk } = loadHelper(patchPreloadHelper(helper)!);
+    const first = importChunk();
+    const second = importChunk();
+    links[0].dispatchEvent(new Event("error"));
+    await flush();
+    assert.deepEqual([first.settled, second.settled], ["rejected", "rejected"]);
+  });
+
+  it("leaves a helper it does not recognize alone", () => {
+    assert.equal(
+      patchPreloadHelper("export const __vitePreload = 1"),
+      undefined,
+    );
+  });
+});

src/index.ts

@@ -29,6 +29,7 @@ import {
   type LinkAssetsDocManifest,
 } from "./manifest-generator";
 import { normalizePath, POSIX_SEP, WINDOWS_SEP } from "./normalize-path";
+import preloadCssPlugin from "./preload-css";
 import { cleanUrl, hasOpaqueQuery } from "./query";
 import { ReadOncePersistedStore } from "./read-once-persisted-store";
 import relativeAssetsTransform from "./relative-assets-transform";
@@ -282,6 +283,7 @@ export default function markoPlugin(opts: Options = {}): vite.Plugin[] {
     isMarkoFile(file) || styleImportReg.test(file) || viteAssetsInclude(file);
 
   return [
+    preloadCssPlugin(),
     {
       name: "marko-vite:pre",
       enforce: "pre", // Must be pre to allow us to resolve assets before vite.

src/preload-css.ts

@@ -0,0 +1,55 @@
+import type * as vite from "vite";
+
+const preloadHelperReg = /^\0vite\/preload-helper\.js$/;
+const seenCheckReg =
+  /if\s*\(dep in seen\)\s*return;?\s*seen\[dep\]\s*=\s*true;?/;
+const appendLinkReg = /document\.head\.appendChild\(link\);?/;
+
+/**
+ * Vite's preload helper only makes the dynamic import that inserts a
+ * stylesheet wait for it to load. Another import needing that stylesheet while
+ * it loads (eg a lazy template imported again) resolves early and renders
+ * unstyled, so the pending load is shared with every caller until it settles.
+ */
+export default function preloadCssPlugin(): vite.Plugin {
+  return {
+    name: "marko-vite:preload-css",
+    apply: "build",
+    transform: {
+      filter: { id: preloadHelperReg },
+      handler(code) {
+        return patchPreloadHelper(code);
+      },
+    },
+  };
+}
+
+/**
+ * Returns the helper with shared stylesheet loads, or `undefined` when its
+ * code no longer has the shape this patches (eg Vite fixed it upstream).
+ */
+export function patchPreloadHelper(code: string): string | undefined {
+  if (!seenCheckReg.test(code) || !appendLinkReg.test(code)) return;
+  return (
+    code
+      .replace(
+        seenCheckReg,
+        "if (dep in seen) return seen[dep];\nseen[dep] = void 0;",
+      )
+      .replace(
+        appendLinkReg,
+        "$&\nif (isCss) seen[dep] = __markoCssLoad(link, dep, seen);",
+      ) +
+    `
+function __markoCssLoad(link, dep, seen) {
+  const loaded = new Promise((res, rej) => {
+    link.addEventListener("load", res);
+    link.addEventListener("error", () => rej(new Error(\`Unable to preload CSS for \${dep}\`)));
+  });
+  const settle = () => { seen[dep] = void 0; };
+  loaded.then(settle, settle);
+  return loaded;
+}
+`
+  );
+}