fix(remix-dev/vite, remix-server-runtime): handle criticalCss in an adapter agnostic way

.changeset/breezy-guests-repair.md

@@ -0,0 +1,6 @@
+---
+"@remix-run/dev": patch
+"@remix-run/server-runtime": patch
+---
+
+Fix flash of unstyled content for non-Express custom servers in Vite dev

packages/remix-dev/vite/node/adapter.ts

@@ -198,16 +198,14 @@ export let createRequestHandler = (
   build: ServerBuild,
   {
     mode = "production",
-    criticalCss,
   }: {
     mode?: string;
-    criticalCss?: string;
   }
 ) => {
   let handler = createBaseRequestHandler(build, mode);
   return async (req: IncomingMessage, res: ServerResponse) => {
     let request = createRequest(req);
-    let response = await handler(request, {}, { __criticalCss: criticalCss });
+    let response = await handler(request, {});
     handleNodeResponse(response, res);
   };
 };

packages/remix-dev/vite/plugin.ts

@@ -5,7 +5,10 @@ import { type BinaryLike, createHash } from "node:crypto";
 import * as path from "node:path";
 import * as fse from "fs-extra";
 import babel from "@babel/core";
-import { type ServerBuild } from "@remix-run/server-runtime";
+import {
+  type ServerBuild,
+  unstable_setDevServerHooks as setDevServerHooks,
+} from "@remix-run/server-runtime";
 import {
   init as initEsModuleLexer,
   parse as esModuleLexer,
@@ -666,22 +669,29 @@ export const remixVitePlugin: RemixVitePlugin = (options = {}) => {
           setTimeout(showUnstableWarning, 50);
         });
 
+        // Give the request handler access to the critical CSS in dev to avoid a
+        // flash of unstyled content since Vite injects CSS file contents via JS
+        setDevServerHooks({
+          getCriticalCss: async (build, url) => {
+            invariant(cachedPluginConfig);
+            return getStylesForUrl(
+              vite,
+              cachedPluginConfig,
+              cssModulesManifest,
+              build,
+              url
+            );
+          },
+        });
+
         // We cache the pluginConfig here to make sure we're only invalidating virtual modules when necessary.
         // This requires a separate cache from `cachedPluginConfig`, which is updated by remix-hmr-updates. If
         // we shared the cache, it would already be refreshed by remix-hmr-updates at this point, and we'd
         // have no way of comparing against the cache to know if the virtual modules need to be invalidated.
         let previousPluginConfig: ResolvedRemixVitePluginConfig | undefined;
 
-        let localsByRequest = new WeakMap<
-          Vite.Connect.IncomingMessage,
-          {
-            build: ServerBuild;
-            criticalCss: string | undefined;
-          }
-        >();
-
         return () => {
-          vite.middlewares.use(async (req, res, next) => {
+          vite.middlewares.use(async (_req, _res, next) => {
             try {
               let pluginConfig = await resolvePluginConfig();
 
@@ -702,36 +712,6 @@ export const remixVitePlugin: RemixVitePlugin = (options = {}) => {
                   }
                 });
               }
-              let { url } = req;
-              let build = await (vite.ssrLoadModule(
-                serverEntryId
-              ) as Promise<ServerBuild>);
-
-              let criticalCss = await getStylesForUrl(
-                vite,
-                pluginConfig,
-                cssModulesManifest,
-                build,
-                url
-              );
-
-              localsByRequest.set(req, {
-                build,
-                criticalCss,
-              });
-
-              // If the middleware is being used in Express, the "res.locals"
-              // object (https://expressjs.com/en/api.html#res.locals) will be
-              // present. If so, we attach the critical CSS as metadata to the
-              // response object so the Remix Express adapter has access to it.
-              if (
-                "locals" in res &&
-                typeof res.locals === "object" &&
-                res.locals !== null
-              ) {
-                (res.locals as Record<string, any>).__remixDevCriticalCss =
-                  criticalCss;
-              }
 
               next();
             } catch (error) {
@@ -744,14 +724,12 @@ export const remixVitePlugin: RemixVitePlugin = (options = {}) => {
           if (!vite.config.server.middlewareMode) {
             vite.middlewares.use(async (req, res, next) => {
               try {
-                let locals = localsByRequest.get(req);
-                invariant(locals, "No Remix locals found for request");
-
-                let { build, criticalCss } = locals;
+                let build = (await vite.ssrLoadModule(
+                  serverEntryId
+                )) as ServerBuild;
 
                 let handle = createRequestHandler(build, {
                   mode: "development",
-                  criticalCss,
                 });
 
                 await handle(req, res);

packages/remix-express/server.ts

@@ -52,14 +52,7 @@ export function createRequestHandler({
       let request = createRemixRequest(req, res);
       let loadContext = await getLoadContext?.(req, res);
 
-      let criticalCss =
-        mode === "production" ? null : res.locals.__remixDevCriticalCss;
-
-      let response = await handleRequest(
-        request,
-        loadContext,
-        criticalCss ? { __criticalCss: criticalCss } : undefined
-      );
+      let response = await handleRequest(request, loadContext);
 
       await sendRemixResponse(res, response);
     } catch (error: unknown) {

packages/remix-server-runtime/dev.ts

@@ -25,3 +25,22 @@ export async function broadcastDevReady(build: ServerBuild, origin?: string) {
 export function logDevReady(build: ServerBuild) {
   console.log(`[REMIX DEV] ${build.assets.version} ready`);
 }
+
+type DevServerHooks = {
+  getCriticalCss: (
+    build: ServerBuild,
+    pathname: string
+  ) => Promise<string | undefined>;
+};
+
+const globalDevServerHooksKey = "__remix_devServerHooks";
+
+export function setDevServerHooks(devServerHooks: DevServerHooks) {
+  // @ts-expect-error
+  globalThis[globalDevServerHooksKey] = devServerHooks;
+}
+
+export function getDevServerHooks(): DevServerHooks | undefined {
+  // @ts-expect-error
+  return globalThis[globalDevServerHooksKey];
+}

packages/remix-server-runtime/index.ts

@@ -15,7 +15,11 @@ export { createCookieSessionStorageFactory } from "./sessions/cookieStorage";
 export { createMemorySessionStorageFactory } from "./sessions/memoryStorage";
 export { createMemoryUploadHandler as unstable_createMemoryUploadHandler } from "./upload/memoryUploadHandler";
 export { MaxPartSizeExceededError } from "./upload/errors";
-export { broadcastDevReady, logDevReady } from "./dev";
+export {
+  broadcastDevReady,
+  logDevReady,
+  setDevServerHooks as unstable_setDevServerHooks,
+} from "./dev";
 
 // Types for the Remix server runtime interface
 export type {

packages/remix-server-runtime/server.ts

@@ -28,16 +28,11 @@ import {
   isResponse,
 } from "./responses";
 import { createServerHandoffString } from "./serverHandoff";
+import { getDevServerHooks } from "./dev";
 
 export type RequestHandler = (
   request: Request,
-  loadContext?: AppLoadContext,
-  args?: {
-    /**
-     * @private This is an internal API intended for use by the Remix Vite plugin in dev mode
-     */
-    __criticalCss?: string;
-  }
+  loadContext?: AppLoadContext
 ) => Promise<Response>;
 
 export type CreateRequestHandlerFunction = (
@@ -80,11 +75,7 @@ export const createRequestHandler: CreateRequestHandlerFunction = (
   let staticHandler: StaticHandler;
   let errorHandler: HandleErrorFunction;
 
-  return async function requestHandler(
-    request,
-    loadContext = {},
-    { __criticalCss: criticalCss } = {}
-  ) {
+  return async function requestHandler(request, loadContext = {}) {
     _build = typeof build === "function" ? await build() : build;
     if (typeof build === "function") {
       let derived = derive(_build, mode);
@@ -144,6 +135,11 @@ export const createRequestHandler: CreateRequestHandlerFunction = (
         handleError
       );
     } else {
+      let criticalCss =
+        mode === ServerMode.Development
+          ? await getDevServerHooks()?.getCriticalCss(_build, url.pathname)
+          : undefined;
+
       response = await handleDocumentRequestRR(
         serverMode,
         _build,