Fix critical CSS with custom Vite.DevEnvironment

packages/react-router-dev/vite/plugin.ts

@@ -36,7 +36,7 @@ import type { Cache } from "./cache";
 import { generate, parse } from "./babel";
 import type { NodeRequestHandler } from "./node-adapter";
 import { fromNodeRequest, toNodeRequest } from "./node-adapter";
-import { getStylesForUrl, isCssModulesFile } from "./styles";
+import { getStylesForPathname, isCssModulesFile } from "./styles";
 import * as VirtualModule from "./virtual-module";
 import { resolveFileUrl } from "./resolve-file-url";
 import { combineURLs } from "./combine-urls";
@@ -678,7 +678,22 @@ export const reactRouterVitePlugin: ReactRouterVitePlugin = () => {
         }`;
           })
           .join(",\n  ")}
-      };`;
+      };
+      ${
+        ctx.reactRouterConfig.future.unstable_viteEnvironmentApi &&
+        viteCommand === "serve"
+          ? `
+              export const getCriticalCss = ({ pathname }) => {
+                return {
+                  rel: "stylesheet",
+                  href: "${
+                    viteUserConfig.base ?? "/"
+                  }@react-router/critical.css?pathname=" + pathname,
+                };
+              }
+            `
+          : ""
+      }`;
   };
 
   let loadViteManifest = async (directory: string) => {
@@ -1341,15 +1356,14 @@ export const reactRouterVitePlugin: ReactRouterVitePlugin = () => {
         setDevServerHooks({
           // 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
-          getCriticalCss: async (build, url) => {
-            return getStylesForUrl({
+          getCriticalCss: async (pathname) => {
+            return getStylesForPathname({
               rootDirectory: ctx.rootDirectory,
               entryClientFilePath: ctx.entryClientFilePath,
               reactRouterConfig: ctx.reactRouterConfig,
               viteDevServer,
               loadCssContents,
-              build,
-              url,
+              pathname,
             });
           },
           // If an error is caught within the request handler, let Vite fix the
@@ -1397,6 +1411,30 @@ export const reactRouterVitePlugin: ReactRouterVitePlugin = () => {
           }
         );
 
+        if (ctx.reactRouterConfig.future.unstable_viteEnvironmentApi) {
+          viteDevServer.middlewares.use(async (req, res, next) => {
+            let [reqPathname, reqSearch] = (req.url ?? "").split("?");
+            if (reqPathname === "/@react-router/critical.css") {
+              let pathname = new URLSearchParams(reqSearch).get("pathname");
+              if (!pathname) {
+                return next("No pathname provided");
+              }
+              let css = await getStylesForPathname({
+                rootDirectory: ctx.rootDirectory,
+                entryClientFilePath: ctx.entryClientFilePath,
+                reactRouterConfig: ctx.reactRouterConfig,
+                viteDevServer,
+                loadCssContents,
+                pathname,
+              });
+              res.setHeader("Content-Type", "text/css");
+              res.end(css);
+            } else {
+              next();
+            }
+          });
+        }
+
         return () => {
           // Let user servers handle SSR requests in middleware mode,
           // otherwise the Vite plugin will handle the request

packages/react-router-dev/vite/styles.ts

@@ -1,15 +1,12 @@
 import * as path from "node:path";
-import type { ServerBuild } from "react-router";
 import { matchRoutes } from "react-router";
 import type { ModuleNode, ViteDevServer } from "vite";
 
 import type { ResolvedReactRouterConfig } from "../config/config";
+import type { RouteManifest, RouteManifestEntry } from "../config/routes";
 import type { LoadCssContents } from "./plugin";
 import { resolveFileUrl } from "./resolve-file-url";
 
-type ServerRouteManifest = ServerBuild["routes"];
-type ServerRoute = ServerRouteManifest[string];
-
 // Style collection logic adapted from solid-start: https://github.com/solidjs/solid-start
 
 // Vite doesn't expose these so we just copy the list for now
@@ -163,8 +160,8 @@ const findDeps = async (
   await Promise.all(branches);
 };
 
-const groupRoutesByParentId = (manifest: ServerRouteManifest) => {
-  let routes: Record<string, NonNullable<ServerRoute>[]> = {};
+const groupRoutesByParentId = (manifest: RouteManifest) => {
+  let routes: Record<string, Array<RouteManifestEntry>> = {};
 
   Object.values(manifest).forEach((route) => {
     if (route) {
@@ -179,45 +176,62 @@ const groupRoutesByParentId = (manifest: ServerRouteManifest) => {
   return routes;
 };
 
-// Create a map of routes by parentId to use recursively instead of
-// repeatedly filtering the manifest.
-const createRoutes = (
-  manifest: ServerRouteManifest,
+type RouteManifestEntryWithChildren = Omit<RouteManifestEntry, "index"> &
+  (
+    | { index?: false | undefined; children: RouteManifestEntryWithChildren[] }
+    | { index: true; children?: never }
+  );
+
+const createRoutesWithChildren = (
+  manifest: RouteManifest,
   parentId: string = "",
   routesByParentId = groupRoutesByParentId(manifest)
-): NonNullable<ServerRoute>[] => {
+): RouteManifestEntryWithChildren[] => {
   return (routesByParentId[parentId] || []).map((route) => ({
     ...route,
-    children: createRoutes(manifest, route.id, routesByParentId),
+    ...(route.index
+      ? {
+          index: true,
+        }
+      : {
+          index: false,
+          children: createRoutesWithChildren(
+            manifest,
+            route.id,
+            routesByParentId
+          ),
+        }),
   }));
 };
 
-export const getStylesForUrl = async ({
+export const getStylesForPathname = async ({
   viteDevServer,
   rootDirectory,
   reactRouterConfig,
   entryClientFilePath,
   loadCssContents,
-  build,
-  url,
+  pathname,
 }: {
   viteDevServer: ViteDevServer;
   rootDirectory: string;
-  reactRouterConfig: Pick<ResolvedReactRouterConfig, "appDirectory" | "routes">;
+  reactRouterConfig: Pick<
+    ResolvedReactRouterConfig,
+    "appDirectory" | "routes" | "basename"
+  >;
   entryClientFilePath: string;
   loadCssContents: LoadCssContents;
-  build: ServerBuild;
-  url: string | undefined;
+  pathname: string | undefined;
 }): Promise<string | undefined> => {
-  if (url === undefined || url.includes("?_data=")) {
+  if (pathname === undefined || pathname.includes("?_data=")) {
     return undefined;
   }
 
-  let routes = createRoutes(build.routes);
+  let routesWithChildren = createRoutesWithChildren(reactRouterConfig.routes);
   let appPath = path.relative(process.cwd(), reactRouterConfig.appDirectory);
   let documentRouteFiles =
-    matchRoutes(routes, url, build.basename)?.map((match) =>
-      path.resolve(appPath, reactRouterConfig.routes[match.route.id].file)
+    matchRoutes(routesWithChildren, pathname, reactRouterConfig.basename)?.map(
+      (match) =>
+        path.resolve(appPath, reactRouterConfig.routes[match.route.id].file)
     ) ?? [];
 
   let styles = await getStylesForFiles({

packages/react-router/lib/dom/global.ts

@@ -1,11 +1,11 @@
 import type { HydrationState, Router as DataRouter } from "../router/router";
-import type { AssetsManifest, FutureConfig } from "./ssr/entry";
+import type { AssetsManifest, CriticalCss, FutureConfig } from "./ssr/entry";
 import type { RouteModules } from "./ssr/routeModules";
 
 export type WindowReactRouterContext = {
   basename?: string;
   state: HydrationState;
-  criticalCss?: string;
+  criticalCss?: CriticalCss;
   future: FutureConfig;
   ssr: boolean;
   isSpaMode: boolean;

packages/react-router/lib/dom/ssr/components.tsx

@@ -241,9 +241,12 @@ export function Links() {
 
   return (
     <>
-      {criticalCss ? (
+      {typeof criticalCss === "string" ? (
         <style dangerouslySetInnerHTML={{ __html: criticalCss }} />
       ) : null}
+      {typeof criticalCss === "object" ? (
+        <link rel="stylesheet" href={criticalCss.href} />
+      ) : null}
       {keyedLinks.map(({ key, link }) =>
         isPageLinkDescriptor(link) ? (
           <PrefetchPageLinks key={key} {...link} />

packages/react-router/lib/dom/ssr/entry.ts

@@ -13,7 +13,7 @@ type SerializedError = {
 export interface FrameworkContextObject {
   manifest: AssetsManifest;
   routeModules: RouteModules;
-  criticalCss?: string;
+  criticalCss?: CriticalCss;
   serverHandoffString?: string;
   future: FutureConfig;
   ssr: boolean;
@@ -43,6 +43,8 @@ export interface EntryContext extends FrameworkContextObject {
 
 export interface FutureConfig {}
 
+export type CriticalCss = string | { rel: "stylesheet"; href: string };
+
 export interface AssetsManifest {
   entry: {
     imports: string[];

packages/react-router/lib/server-runtime/build.ts

@@ -1,12 +1,15 @@
 import type { ActionFunctionArgs, LoaderFunctionArgs } from "../router/utils";
 import type {
   AssetsManifest,
+  CriticalCss,
   EntryContext,
   FutureConfig,
 } from "../dom/ssr/entry";
 import type { ServerRouteManifest } from "./routes";
 import type { AppLoadContext } from "./data";
 
+type OptionalCriticalCss = CriticalCss | undefined;
+
 /**
  * The output of the compiler for the server build.
  */
@@ -21,6 +24,9 @@ export interface ServerBuild {
   assetsBuildDirectory: string;
   future: FutureConfig;
   ssr: boolean;
+  getCriticalCss?: (args: {
+    pathname: string;
+  }) => OptionalCriticalCss | Promise<OptionalCriticalCss>;
   /**
    * @deprecated This is now done via a custom header during prerendering
    */

packages/react-router/lib/server-runtime/dev.ts

@@ -1,10 +1,5 @@
-import type { ServerBuild } from "./build";
-
 type DevServerHooks = {
-  getCriticalCss?: (
-    build: ServerBuild,
-    pathname: string
-  ) => Promise<string | undefined>;
+  getCriticalCss?: (pathname: string) => Promise<string | undefined>;
   processRequestError?: (error: unknown) => void;
 };
 

packages/react-router/lib/server-runtime/server.ts

@@ -13,7 +13,7 @@ import {
 } from "../router/router";
 import type { AppLoadContext } from "./data";
 import type { HandleErrorFunction, ServerBuild } from "./build";
-import type { EntryContext } from "../dom/ssr/entry";
+import type { CriticalCss, EntryContext } from "../dom/ssr/entry";
 import { createEntryRouteModules } from "./entry";
 import { sanitizeErrors, serializeError, serializeErrors } from "./errors";
 import { ServerMode, isServerMode } from "./mode";
@@ -263,10 +263,17 @@ export const createRequestHandler: CreateRequestHandlerFunction = (
         handleError
       );
     } else {
-      let criticalCss =
-        mode === ServerMode.Development
-          ? await getDevServerHooks()?.getCriticalCss?.(_build, url.pathname)
-          : undefined;
+      let { pathname } = url;
+
+      let criticalCss: CriticalCss | undefined = undefined;
+      if (_build.getCriticalCss) {
+        criticalCss = await _build.getCriticalCss({ pathname });
+      } else if (
+        mode === ServerMode.Development &&
+        getDevServerHooks()?.getCriticalCss
+      ) {
+        criticalCss = await getDevServerHooks()?.getCriticalCss?.(pathname);
+      }
 
       response = await handleDocumentRequest(
         serverMode,
@@ -389,7 +396,7 @@ async function handleDocumentRequest(
   request: Request,
   loadContext: AppLoadContext,
   handleError: (err: unknown) => void,
-  criticalCss?: string
+  criticalCss?: CriticalCss
 ) {
   let isSpaMode = request.headers.has("X-React-Router-SPA-Mode");
   let context: Awaited<ReturnType<typeof staticHandler.query>>;

packages/react-router/lib/server-runtime/serverHandoff.ts

@@ -1,5 +1,5 @@
 import type { HydrationState } from "../router/router";
-import type { FutureConfig } from "../dom/ssr/entry";
+import type { CriticalCss, FutureConfig } from "../dom/ssr/entry";
 import { escapeHtml } from "./markup";
 
 type ValidateShape<T, Shape> =
@@ -18,7 +18,7 @@ export function createServerHandoffString<T>(serverHandoff: {
   // Don't allow StaticHandlerContext to be passed in verbatim, since then
   // we'd end up including duplicate info
   state?: ValidateShape<T, HydrationState>;
-  criticalCss?: string;
+  criticalCss?: CriticalCss;
   basename: string | undefined;
   future: FutureConfig;
   ssr: boolean;