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()
+      })
+    },
   }
 }