feat: load critical CSS on development (#765)

packages/tuono-react-vite-plugin/src/plugin.ts

@@ -1,8 +1,11 @@
 import { normalize } from 'node:path'
 
-import type { Plugin } from 'vite'
+import type { Plugin, ViteDevServer } from 'vite'
 
 import { routeGenerator } from './fs-routing/generator'
+import { getStylesForComponentId, isCssModulesFile } from './styles'
+
+const CRITICAL_CSS_PATH = '/vite-server/tuono_internal__critical_css'
 
 const ROUTES_DIRECTORY_PATH = './src/routes'
 
@@ -30,6 +33,10 @@ export function TuonoReactPlugin(): Plugin {
     }
   }
 
+  // This manifest is used to store the CSS modules contents in dev mode
+  // { [filePath]: cssContent }
+  const cssModulesManifest: Record<string, string> = {}
+
   return {
     name: 'vite-plugin-tuono-react',
     configResolved: async (): Promise<void> => {
@@ -43,5 +50,34 @@ export function TuonoReactPlugin(): Plugin {
         await handleFile(file)
       }
     },
+    transform: (code, id): void => {
+      if (isCssModulesFile(id)) {
+        cssModulesManifest[id] = code
+      }
+    },
+    configureServer: (server: ViteDevServer): void => {
+      // Using middlewares in order to take advantage of async requests out of
+      // the box
+      // eslint-disable-next-line @typescript-eslint/no-misused-promises
+      server.middlewares.use(async (req, res, next): Promise<void> => {
+        const url = new URL(req.url || '', `http://${req.headers.host || ''}`)
+
+        // 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
+        if (url.pathname === CRITICAL_CSS_PATH) {
+          const componentId = url.searchParams.get('componentId')
+          const css = await getStylesForComponentId(
+            server,
+            componentId,
+            cssModulesManifest,
+          )
+
+          res.writeHead(200, { 'Content-Type': 'text/css' })
+          res.end(css)
+          return
+        }
+        next()
+      })
+    },
   }
 }

packages/tuono-react-vite-plugin/src/styles.ts

@@ -0,0 +1,197 @@
+/**
+ * This module is strongly inspired by the remix project.
+ *
+ * source: https://github.com/remix-run/remix/blob/main/packages/remix-dev/vite/styles.ts
+ */
+import path from 'path'
+
+import type { ModuleNode, ViteDevServer } from 'vite'
+
+const isCssFile = (file: string): boolean => cssFileRegExp.test(file)
+
+const cssFileRegExp =
+  /\.(css|less|sass|scss|styl|stylus|pcss|postcss|sss)(?:$|\?)/
+
+const cssModulesRegExp = new RegExp(`\\.module${cssFileRegExp.source}`)
+
+const routesFolder = path.relative(process.cwd(), 'src/routes')
+
+const injectQuery = (url: string, query: string): string =>
+  url.includes('?') ? url.replace('?', `?${query}&`) : `${url}?${query}`
+
+export const isCssModulesFile = (file: string): boolean =>
+  cssModulesRegExp.test(file)
+
+const cssUrlParamsWithoutSideEffects = ['url', 'inline', 'raw', 'inline-css']
+
+const isCssUrlWithoutSideEffects = (url: string): boolean => {
+  const queryString = url.split('?')[1]
+
+  if (!queryString) {
+    return false
+  }
+
+  const params = new URLSearchParams(queryString)
+  for (const paramWithoutSideEffects of cssUrlParamsWithoutSideEffects) {
+    if (
+      // Parameter is blank and not explicitly set, i.e. "?url", not "?url="
+      params.get(paramWithoutSideEffects) === '' &&
+      !url.includes(`?${paramWithoutSideEffects}=`) &&
+      !url.includes(`&${paramWithoutSideEffects}=`)
+    ) {
+      return true
+    }
+  }
+
+  return false
+}
+
+/**
+ * This function transform the componentId into a file path.
+ * File extension is not required for the vite.moduleGraph URL search.
+ */
+function findFileFromComponentId(id: string): string {
+  if (id.endsWith('/')) {
+    return id + 'index'
+  }
+
+  if (id.includes('__root__')) {
+    return id.replaceAll('__root__', '__layout')
+  }
+
+  return id
+}
+
+export const getStylesForComponentId = async (
+  viteDevServer: ViteDevServer,
+  /**
+   * The route name (should match tuono-router specs)
+   */
+  componentId: string | null,
+  /**
+   * All the CSS modules are preloaded and saved in this manifest
+   */
+  cssModulesManifest: Record<string, string>,
+): Promise<string | undefined> => {
+  const relativeFilePath = path.join(
+    routesFolder,
+    findFileFromComponentId(componentId || ''),
+  )
+
+  const fileUrl = path.join(process.cwd(), relativeFilePath)
+
+  const styles: Record<string, string> = {}
+  const deps: Set<ModuleNode> = new Set()
+
+  try {
+    let node: ModuleNode | undefined =
+      await viteDevServer.moduleGraph.getModuleByUrl(fileUrl)
+
+    // If the module is only present in the client module graph, the module
+    // won't have been found on the first request to the server. If so, we
+    // request the module so it's in the module graph, then try again.
+    if (!node) {
+      try {
+        await viteDevServer.transformRequest(fileUrl)
+      } catch (err) {
+        console.error(err)
+      }
+
+      node = await viteDevServer.moduleGraph.getModuleByUrl(fileUrl)
+    }
+
+    if (!node) {
+      console.error(`Could not resolve module for file: ${fileUrl}`)
+      return
+    }
+    await findNodeDependencies(viteDevServer, node, deps)
+  } catch (error) {
+    console.error(error)
+  }
+
+  for (const dep of deps) {
+    if (
+      dep.file &&
+      isCssFile(dep.file) &&
+      !isCssUrlWithoutSideEffects(dep.url) // Ignore styles that resolved as URLs, inline or raw. These shouldn't get injected.
+    ) {
+      try {
+        const css = isCssModulesFile(dep.file)
+          ? cssModulesManifest[dep.file]
+          : ((
+              await viteDevServer.ssrLoadModule(
+                // We need the ?inline query in Vite v6 when loading CSS in SSR
+                // since it does not expose the default export for CSS in a
+                // server environment.
+                injectQuery(dep.url, 'inline'),
+              )
+            ).default as string)
+
+        if (css === undefined) {
+          throw new Error()
+        }
+
+        styles[dep.url] = css
+      } catch {
+        // this can happen with dynamically imported modules
+        console.warn(`Could not load ${dep.file}`)
+      }
+    }
+  }
+
+  return (
+    Object.entries(styles)
+      .map(([fileName, css]) => [
+        `\n/* ${fileName
+          // Escape comment syntax in file paths
+          .replace(/\/\*/g, '/\\*')
+          .replace(/\*\//g, '*\\/')} */`,
+        css,
+      ])
+      .flat()
+      .join('\n') || undefined
+  )
+}
+
+/**
+ * This function is used to find all the dependencies of a module node.
+ * The starting node is always a route.
+ */
+const findNodeDependencies = async (
+  vite: ViteDevServer,
+  node: ModuleNode,
+  deps: Set<ModuleNode>,
+): Promise<void> => {
+  // since `ssrTransformResult.deps` contains URLs instead of `ModuleNode`s, this process is asynchronous.
+  // instead of using `await`, we resolve all branches in parallel.
+  const branches: Array<Promise<void>> = []
+
+  async function addFromNode(innerNode: ModuleNode): Promise<void> {
+    if (!deps.has(innerNode)) {
+      deps.add(innerNode)
+      await findNodeDependencies(vite, innerNode, deps)
+    }
+  }
+
+  async function addFromUrl(url: string): Promise<void> {
+    const innerNode = await vite.moduleGraph.getModuleByUrl(url)
+
+    if (innerNode) {
+      await addFromNode(innerNode)
+    }
+  }
+
+  if (node.ssrTransformResult) {
+    if (node.ssrTransformResult.deps) {
+      node.ssrTransformResult.deps.forEach((url) =>
+        branches.push(addFromUrl(url)),
+      )
+    }
+  } else {
+    node.importedModules.forEach((innerNode: ModuleNode) =>
+      branches.push(addFromNode(innerNode)),
+    )
+  }
+
+  await Promise.all(branches)
+}