feat: load critical CSS on development (#765)

crates/tuono/Cargo.toml

@@ -24,6 +24,7 @@ tracing-subscriber = {version = "0.3.19", features = ["env-filter"]}
 miette = "7.2.0"
 
 colored = "2.1.0"
+once_cell = "1.19.0"
 watchexec = "5.0.0"
 watchexec-signals = "4.0.0"
 watchexec-events = "4.0.0"

crates/tuono/src/commands/dev.rs

@@ -1,3 +1,6 @@
+use once_cell::sync::Lazy;
+use regex::Regex;
+use std::borrow::Cow;
 use std::collections::HashSet;
 use std::fs;
 use std::path::{Path, PathBuf};
@@ -19,6 +22,12 @@ use crate::source_builder::SourceBuilder;
 use console::Term;
 use spinners::{Spinner, Spinners};
 
+fn is_css_module(file_name: Cow<str>) -> bool {
+    static RE: Lazy<Regex> =
+        Lazy::new(|| Regex::new(r"^.*\.module\.(css|scss|sass|less|styl|stylus)$").unwrap());
+    RE.is_match(&file_name)
+}
+
 fn ssr_reload_needed(path: &Path) -> bool {
     let file_name_starts_with_env = path
         .file_name()
@@ -27,7 +36,12 @@ fn ssr_reload_needed(path: &Path) -> bool {
 
     let file_path = path.to_string_lossy();
 
-    file_name_starts_with_env || file_path.ends_with("sx") || file_path.ends_with("mdx")
+    file_name_starts_with_env
+        || file_path.ends_with("sx")
+        || file_path.ends_with("mdx")
+    // When a CSS module is modified
+    // also the class names get modified.
+        || is_css_module(file_path)
 }
 
 #[allow(

crates/tuono_lib/src/vite_reverse_proxy.rs

@@ -1,12 +1,16 @@
 use crate::config::GLOBAL_CONFIG;
 use axum::body::Body;
-use axum::extract::Path;
+use axum::extract::{Path, Query};
+use std::collections::HashMap;
 
 use axum::http::{HeaderName, HeaderValue};
 use axum::response::{IntoResponse, Response};
 use reqwest::Client;
 
-pub async fn vite_reverse_proxy(Path(path): Path<String>) -> impl IntoResponse {
+pub async fn vite_reverse_proxy(
+    Path(path): Path<String>,
+    query: Query<HashMap<String, String>>,
+) -> impl IntoResponse {
     let client = Client::new();
 
     let config = GLOBAL_CONFIG
@@ -19,7 +23,24 @@ pub async fn vite_reverse_proxy(Path(path): Path<String>) -> impl IntoResponse {
         config.server.port + 1
     );
 
-    match client.get(format!("{vite_url}/{path}")).send().await {
+    let query_string = query
+        .0
+        .iter()
+        .map(|(k, v)| format!("{}={}", k, v))
+        .collect::<Vec<_>>()
+        .join("&");
+
+    let query_string = if query_string.is_empty() {
+        String::new()
+    } else {
+        format!("?{}", query_string)
+    };
+
+    match client
+        .get(format!("{vite_url}/{path}{query_string}"))
+        .send()
+        .await
+    {
         Ok(res) => {
             let mut response_builder = Response::builder().status(res.status().as_u16());
 

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

packages/tuono-router/src/components/CriticalCss.tsx

@@ -0,0 +1,33 @@
+import type { JSX } from 'react'
+
+import type { Mode } from '../types'
+
+const VITE_PROXY_PATH = '/vite-server'
+const CRITICAL_CSS_PATH = VITE_PROXY_PATH + '/tuono_internal__critical_css'
+
+interface CriticalCssProps {
+  routeId?: string
+  mode?: Mode
+}
+
+/**
+ * Returns the critical CSS for the given route
+ * This is required in order to avoid FOUC during development
+ * since vite does not support CSS injection without JS waterfall
+ */
+export function CriticalCss({
+  routeId,
+  mode,
+}: CriticalCssProps): JSX.Element | null {
+  if (!routeId || mode !== 'Dev') {
+    return null
+  }
+
+  return (
+    <link
+      href={`${CRITICAL_CSS_PATH}?componentId=${routeId}`}
+      precedence="high"
+      rel="stylesheet"
+    />
+  )
+}

packages/tuono-router/src/components/Matches.tsx

@@ -2,6 +2,8 @@ import type { JSX } from 'react'
 
 import { useRoute } from '../hooks/useRoute'
 
+import type { Mode } from '../types'
+
 import { RouteMatch } from './RouteMatch'
 import { NotFound } from './NotFound'
 import { useRouterContext } from './RouterContext'
@@ -9,16 +11,26 @@ import { useRouterContext } from './RouterContext'
 interface MatchesProps<TServerPayloadData = unknown> {
   // user defined props
   serverInitialData: TServerPayloadData
+  mode?: Mode
 }
 
-export function Matches({ serverInitialData }: MatchesProps): JSX.Element {
+export function Matches({
+  serverInitialData,
+  mode,
+}: MatchesProps): JSX.Element {
   const { location } = useRouterContext()
 
   const route = useRoute(location.pathname)
 
   if (!route) {
-    return <NotFound />
+    return <NotFound mode={mode} />
   }
 
-  return <RouteMatch route={route} serverInitialData={serverInitialData} />
+  return (
+    <RouteMatch
+      route={route}
+      mode={mode}
+      serverInitialData={serverInitialData}
+    />
+  )
 }

packages/tuono-router/src/components/NotFound.tsx

@@ -3,17 +3,25 @@ import type { JSX } from 'react'
 import { useRouterContext } from '../components/RouterContext'
 import { ROOT_ROUTE_ID } from '../route'
 
+import type { Mode } from '../types'
+
 import { RouteMatch } from './RouteMatch'
 import { NotFoundDefaultContent } from './NotFoundDefaultContent'
+import { CriticalCss } from './CriticalCss'
 
-export function NotFound(): JSX.Element | null {
+export function NotFound({ mode }: { mode?: Mode }): JSX.Element | null {
   const { router } = useRouterContext()
 
   const custom404Route = router.routesById['/404']
 
   // Check if exists a custom 404 error page
   if (custom404Route) {
-    return <RouteMatch route={custom404Route} serverInitialData={{}} />
+    return (
+      <>
+        <CriticalCss routeId={custom404Route.id} mode={mode} />
+        <RouteMatch route={custom404Route} mode={mode} serverInitialData={{}} />
+      </>
+    )
   }
 
   const RootLayout = router.routesById[ROOT_ROUTE_ID]?.component
@@ -22,6 +30,7 @@ export function NotFound(): JSX.Element | null {
 
   return (
     <RootLayout data={null} isLoading={false}>
+      <CriticalCss routeId="__root__" mode={mode} />
       <NotFoundDefaultContent />
     </RootLayout>
   )

packages/tuono-router/src/components/RouteMatch.tsx

@@ -1,16 +1,19 @@
 import type { JSX } from 'react'
 import { memo, Suspense, useMemo } from 'react'
 
+import type { Mode } from '../types'
 import type { Route } from '../route'
 
 import { useServerPayloadData } from '../hooks/useServerPayloadData'
 
 import { useRouterContext } from './RouterContext'
+import { CriticalCss } from './CriticalCss'
 
 interface RouteMatchProps<TServerPayloadData = unknown> {
   route: Route
   // User defined server side props
   serverInitialData: TServerPayloadData
+  mode?: Mode
 }
 
 /**
@@ -21,6 +24,7 @@ interface RouteMatchProps<TServerPayloadData = unknown> {
 export const RouteMatch = ({
   route,
   serverInitialData,
+  mode,
 }: RouteMatchProps): JSX.Element => {
   const { data } = useServerPayloadData(route, serverInitialData)
   const { isTransitioning } = useRouterContext()
@@ -35,8 +39,10 @@ export const RouteMatch = ({
       routes={routes}
       data={routeData}
       isLoading={isTransitioning}
+      mode={mode}
     >
       <Suspense>
+        <CriticalCss routeId={route.id} mode={mode} />
         <route.component data={routeData} isLoading={isTransitioning} />
       </Suspense>
     </TraverseRootComponents>
@@ -49,6 +55,7 @@ interface TraverseRootComponentsProps<TData = unknown> {
   isLoading: boolean
   children?: React.ReactNode
   index?: number
+  mode?: Mode
 }
 
 /**
@@ -63,18 +70,22 @@ const TraverseRootComponents = memo(
     data,
     isLoading,
     index = 0,
+    mode,
     children,
   }: TraverseRootComponentsProps): React.JSX.Element => {
     if (routes.length > index) {
-      const Parent = (routes[index] as Route).component
+      const route = routes[index] as Route
+      const Parent = route.component
 
       return (
         <Parent data={data} isLoading={isLoading}>
+          <CriticalCss routeId={route.id} mode={mode} />
           <TraverseRootComponents
             routes={routes}
             data={data}
             isLoading={isLoading}
             index={index + 1}
+            mode={mode}
           >
             {children}
           </TraverseRootComponents>

packages/tuono-router/src/components/RouterProvider.tsx

@@ -1,6 +1,6 @@
 import type { JSX } from 'react'
 
-import type { ServerInitialLocation } from '../types'
+import type { ServerInitialLocation, Mode } from '../types'
 import type { Router } from '../router'
 
 import { RouterContextProvider } from './RouterContext'
@@ -10,19 +10,21 @@ interface RouterProviderProps {
   router: Router
   serverInitialLocation: ServerInitialLocation
   serverInitialData: unknown
+  mode?: Mode
 }
 
 export function RouterProvider({
   router,
   serverInitialLocation,
   serverInitialData,
+  mode,
 }: RouterProviderProps): JSX.Element {
   return (
     <RouterContextProvider
       router={router}
       serverInitialLocation={serverInitialLocation}
     >
-      <Matches serverInitialData={serverInitialData} />
+      <Matches serverInitialData={serverInitialData} mode={mode} />
     </RouterContextProvider>
   )
 }

packages/tuono-router/src/hooks/useRoute.ts

@@ -16,7 +16,9 @@ export function sanitizePathname(pathname: string): string {
   return pathname
 }
 
-/*
+/**
+ * Returns the route that matches the given pathname
+ *
  * This hook is also implemented on server side to match the bundle
  * file to load at the first rendering.
  *

packages/tuono-router/src/types.ts

@@ -1,5 +1,7 @@
 import type { ReactNode, ComponentType } from 'react'
 
+export type Mode = 'Dev' | 'Prod'
+
 export interface Segment {
   type: 'pathname' | 'param' | 'wildcard'
   value: string

packages/tuono/src/shared/DevResources.tsx

@@ -2,8 +2,8 @@ import type { JSX } from 'react'
 
 import type { TuonoConfigServer } from '../config'
 
-const VITE_PROXY_PATH = '/vite-server'
 const DEFAULT_SERVER_CONFIG = { host: 'localhost', origin: null, port: 3000 }
+const VITE_PROXY_PATH = '/vite-server'
 
 interface DevResourcesProps {
   devServerConfig?: TuonoConfigServer

packages/tuono/src/shared/RouterContextProviderWrapper.tsx

@@ -2,10 +2,13 @@ import type { JSX } from 'react'
 import { RouterProvider } from 'tuono-router'
 import type { RouterInstanceType } from 'tuono-router'
 
+import type { Mode } from '../types'
+
 import { useTuonoContextServerPayload } from './TuonoContext'
 
 interface RouterContextProviderWrapperProps {
   router: RouterInstanceType
+  mode?: Mode
 }
 
 /**
@@ -17,6 +20,7 @@ interface RouterContextProviderWrapperProps {
  */
 export function RouterContextProviderWrapper({
   router,
+  mode,
 }: RouterContextProviderWrapperProps): JSX.Element {
   const serverPayload = useTuonoContextServerPayload()
 
@@ -25,6 +29,7 @@ export function RouterContextProviderWrapper({
       router={router}
       serverInitialLocation={serverPayload.location}
       serverInitialData={serverPayload.data}
+      mode={mode}
     />
   )
 }

packages/tuono/src/shared/TuonoEntryPoint.tsx

@@ -19,7 +19,10 @@ export function TuonoEntryPoint({
   return (
     <StrictMode>
       <TuonoContextProvider serverPayload={serverPayload}>
-        <RouterContextProviderWrapper router={router} />
+        <RouterContextProviderWrapper
+          router={router}
+          mode={serverPayload?.mode}
+        />
       </TuonoContextProvider>
     </StrictMode>
   )

packages/tuono/src/ssr/types.ts

@@ -1 +0,0 @@
-export type Mode = 'Dev' | 'Prod'

packages/tuono/src/types.ts

@@ -2,6 +2,8 @@ import type { ReactNode } from 'react'
 
 import type { TuonoConfigServer } from './config'
 
+export type Mode = 'Dev' | 'Prod'
+
 /**
  * Provided by the rust server and used in the ssr env
  * @see tuono-router {@link ServerInitialLocation}