feat(streaming-ssr): Fix build and server html injection

packages/vite/ambient.d.ts

@@ -1,5 +1,6 @@
 /* eslint-disable no-var */
 /// <reference types="react/canary" />
+import type { HelmetServerState } from 'react-helmet-async'
 
 declare global {
   var RWJS_ENV: {
@@ -11,6 +12,8 @@ declare global {
   }
 
   var __REDWOOD__PRERENDER_PAGES: any
+
+  var __REDWOOD__HELMET_CONTEXT: { helmet?: HelmetServerState }
 }
 
 export {}

packages/vite/bins/rw-vite-build.mjs

@@ -41,13 +41,14 @@ const buildWebSide = async (webDir) => {
     throw new Error('Could not locate your web/vite.config.{js,ts} file')
   }
 
-  // @NOTE: necessary for keeping the cwd correct for postcss/tailwind
-  process.chdir(webDir)
   process.env.NODE_ENV = 'production'
 
   if (getConfig().experimental?.streamingSsr?.enabled) {
-    await buildFeServer({ verbose })
+    // Webdir checks handled in the rwjs/vite package in new build system
+    await buildFeServer({ verbose, webDir })
   } else {
+    // Ensure cwd to be web: required for postcss/tailwind to work correctly
+    process.chdir(webDir)
     // Right now, the buildWeb function looks up the config file from project-config
     // In the future, if we have multiple web spaces we could pass in the cwd here
     buildWeb({ verbose })

packages/vite/src/buildFeServer.ts

@@ -13,12 +13,16 @@ import { getAppRouteHook, getConfig, getPaths } from '@redwoodjs/project-config'
 
 import { buildRscFeServer } from './buildRscFeServer'
 import { RWRouteManifest } from './types'
+import { ensureProcessDirWeb } from './utils'
 
 export interface BuildOptions {
   verbose?: boolean
+  webDir?: string
 }
 
-export const buildFeServer = async ({ verbose }: BuildOptions) => {
+export const buildFeServer = async ({ verbose, webDir }: BuildOptions = {}) => {
+  ensureProcessDirWeb(webDir)
+
   const rwPaths = getPaths()
   const rwConfig = getConfig()
   const viteConfigPath = rwPaths.web.viteConfig
@@ -143,7 +147,7 @@ export const buildFeServer = async ({ verbose }: BuildOptions) => {
     acc[route.path] = {
       name: route.name,
       bundle: route.relativeFilePath
-        ? clientBuildManifest[route.relativeFilePath].file
+        ? clientBuildManifest[route.relativeFilePath]?.file
         : null,
       matchRegexString: route.matchRegexString,
       // @NOTE this is the path definition, not the actual path

packages/vite/src/devFeServer.ts

@@ -1,17 +1,16 @@
 // TODO (STREAMING) Merge with runFeServer so we only have one file
-import path from 'path'
 
 import express from 'express'
-import { renderToPipeableStream } from 'react-dom/server'
 import { createServer as createViteServer } from 'vite'
 
 import { getProjectRoutes } from '@redwoodjs/internal/dist/routes'
 import { getAppRouteHook, getConfig, getPaths } from '@redwoodjs/project-config'
 import { matchPath } from '@redwoodjs/router'
 import type { TagDescriptor } from '@redwoodjs/web'
 
+import { reactRenderToStream } from './streamHelpers'
 import { loadAndRunRouteHooks } from './triggerRouteHooks'
-import { stripQueryStringAndHashFromPath } from './utils'
+import { ensureProcessDirWeb, stripQueryStringAndHashFromPath } from './utils'
 
 // These values are defined in the vite.config.ts
 globalThis.RWJS_ENV = {}
@@ -20,19 +19,7 @@ globalThis.RWJS_ENV = {}
 globalThis.__REDWOOD__PRERENDER_PAGES = {}
 
 async function createServer() {
-  // Check CWD: make sure its the web/ directory
-  // Without this Postcss can misbehave, and its hard to trace why.
-  if (process.cwd() !== getPaths().web.base) {
-    console.error('⚠️  Warning: CWD is ', process.cwd())
-    console.warn('~'.repeat(50))
-    console.warn(
-      'The FE dev server cwd must be web/. Please use `yarn rw dev` or start the server from the web/ directory.'
-    )
-    console.log(`Changing cwd to ${getPaths().web.base}....`)
-    console.log()
-
-    process.chdir(getPaths().web.base)
-  }
+  ensureProcessDirWeb()
 
   const app = express()
   const rwPaths = getPaths()
@@ -62,6 +49,9 @@ async function createServer() {
 
   app.use('*', async (req, res, next) => {
     const currentPathName = stripQueryStringAndHashFromPath(req.originalUrl)
+    globalThis.__REDWOOD__HELMET_CONTEXT = {}
+
+    res.setHeader('content-type', 'text/html; charset=utf-8')
 
     try {
       const routes = getProjectRoutes()
@@ -120,46 +110,15 @@ async function createServer() {
       //    required, and provides efficient invalidation similar to HMR.
       const { ServerEntry } = await vite.ssrLoadModule(rwPaths.web.entryServer)
 
-      // TODO (STREAMING) CSS is handled by Vite in dev mode, we don't need to
-      // worry about it in dev but..... it causes a flash of unstyled content.
-      // For now I'm just injecting index css here
-      // We believe we saw a fix for this somewhere in the Waku sources. Maybe
-      // it was called something like "Capture Css". And it's also mentioned
-      // in the Vite issues on GitHub
-      const FIXME_HardcodedIndexCss = ['index.css']
-
-      const assetMap = JSON.stringify({
-        css: FIXME_HardcodedIndexCss,
-        meta: metaTags,
-      })
-
-      const bootstrapModules = [
-        path.join(__dirname, '../inject', 'reactRefresh.js'),
-      ]
-
       const pageWithJs = currentRoute?.renderMode !== 'html'
 
-      if (pageWithJs) {
-        bootstrapModules.push(rwPaths.web.entryClient)
-      }
-
-      const { pipe } = renderToPipeableStream(
-        ServerEntry({
-          url: currentPathName,
-          css: FIXME_HardcodedIndexCss,
-          meta: metaTags,
-        }),
-        {
-          bootstrapScriptContent: pageWithJs
-            ? `window.__assetMap = function() { return ${assetMap} }`
-            : undefined,
-          bootstrapModules,
-          onShellReady() {
-            res.setHeader('content-type', 'text/html; charset=utf-8')
-            pipe(res)
-          },
-        }
-      )
+      reactRenderToStream({
+        ServerEntry,
+        currentPathName,
+        metaTags,
+        includeJs: pageWithJs,
+        res,
+      })
     } catch (e) {
       // TODO (STREAMING) Is this what we want to do?
       // send back a SPA page

packages/vite/src/runFeServer.ts

@@ -211,15 +211,16 @@ export async function runFeServer() {
 
       const pageWithJs = currentRoute.renderMode !== 'html'
       // @NOTE have to add slash so subpaths still pick up the right file
-      // Vite is currently producing modules not scripts: https://vitejs.dev/config/build-options.html#build-target
       const bootstrapModules = pageWithJs
-        ? ['/' + indexEntry.file, '/' + currentRoute.bundle]
+        ? ([
+            '/' + indexEntry.file,
+            currentRoute.bundle && '/' + currentRoute.bundle,
+          ].filter(Boolean) as string[])
         : undefined
 
       const isSeoCrawler = checkUaForSeoCrawler(req.get('user-agent'))
 
       const { pipe, abort } = renderToPipeableStream(
-        // we should use the same shape as Remix or Next for the meta object
         ServerEntry({
           url: currentPathName,
           css: indexEntry.css,

packages/vite/src/streamHelpers.ts

@@ -0,0 +1,146 @@
+import { Writable } from 'node:stream'
+import path from 'path'
+
+import React from 'react'
+
+import { renderToPipeableStream, renderToString } from 'react-dom/server'
+
+import { getPaths } from '@redwoodjs/project-config'
+import type { TagDescriptor } from '@redwoodjs/web'
+// @TODO (ESM), use exports field. Cannot import from web because of index exports
+import {
+  ServerHtmlProvider,
+  ServerInjectedHtml,
+  createInjector,
+  RenderCallback,
+} from '@redwoodjs/web/dist/components/ServerInject'
+
+interface RenderToStreamArgs {
+  ServerEntry: any
+  currentPathName: string
+  metaTags: TagDescriptor[]
+  includeJs: boolean
+  res: Writable
+}
+
+export function reactRenderToStream({
+  ServerEntry,
+  currentPathName,
+  metaTags,
+  includeJs,
+  res,
+}: RenderToStreamArgs) {
+  const rwPaths = getPaths()
+
+  const bootstrapModules = [
+    path.join(__dirname, '../inject', 'reactRefresh.js'),
+  ]
+
+  if (includeJs) {
+    // type casting: guaranteed to have entryClient by this stage, because checks run earlier
+    bootstrapModules.push(rwPaths.web.entryClient as string)
+  }
+
+  // TODO (STREAMING) CSS is handled by Vite in dev mode, we don't need to
+  // worry about it in dev but..... it causes a flash of unstyled content.
+  // For now I'm just injecting index css here
+  // Looks at collectStyles in packages/vite/src/fully-react/find-styles.ts
+  const FIXME_HardcodedIndexCss = ['index.css']
+
+  const assetMap = JSON.stringify({
+    css: FIXME_HardcodedIndexCss,
+    meta: metaTags,
+  })
+
+  // This ensures an isolated state for each request
+  const { injectionState, injectToPage } = createInjector()
+
+  // This is effectively a transformer stream
+  const intermediateStream = createServerInjectionStream({
+    outputStream: res,
+    onFinal: () => {
+      res.end()
+    },
+    injectionState,
+  })
+
+  const { pipe } = renderToPipeableStream(
+    React.createElement(
+      ServerHtmlProvider,
+      {
+        value: injectToPage,
+      },
+      ServerEntry({
+        url: currentPathName,
+        css: FIXME_HardcodedIndexCss,
+        meta: metaTags,
+      })
+    ),
+    {
+      bootstrapScriptContent: includeJs
+        ? `window.__assetMap = function() { return ${assetMap} }`
+        : undefined,
+      bootstrapModules,
+      onShellReady() {
+        // Pass the react "input" stream to the injection stream
+        // This intermediate stream will interweave the injected html into the react stream's <head>
+        pipe(intermediateStream)
+      },
+    }
+  )
+}
+function createServerInjectionStream({
+  outputStream,
+  onFinal,
+  injectionState,
+}: {
+  outputStream: Writable
+  onFinal: () => void
+  injectionState: Set<RenderCallback>
+}) {
+  return new Writable({
+    write(chunk, encoding, next) {
+      const chunkAsString = chunk.toString()
+      const split = chunkAsString.split('</head>')
+
+      // If the closing tag exists
+      if (split.length > 1) {
+        const [beforeClosingHead, afterClosingHead] = split
+
+        const elementsInjectedToHead = renderToString(
+          React.createElement(ServerInjectedHtml, {
+            injectionState,
+          })
+        )
+
+        const outputBuffer = Buffer.from(
+          [
+            beforeClosingHead,
+            elementsInjectedToHead,
+            '</head>',
+            afterClosingHead,
+          ].join('')
+        )
+
+        outputStream.write(outputBuffer, encoding)
+      } else {
+        outputStream.write(chunk, encoding)
+      }
+
+      next()
+    },
+    final() {
+      // Before finishing, make sure we flush anything else that has been added to the queue
+      // Because of the implementation in ServerRenderHtml, its safe to call this multiple times (I think!)
+      // This is really for the data fetching usecase, where the promise is resolved after <head> is closed
+      const elementsAtTheEnd = renderToString(
+        React.createElement(ServerInjectedHtml, {
+          injectionState,
+        })
+      )
+
+      outputStream.write(elementsAtTheEnd)
+      onFinal()
+    },
+  })
+}

packages/vite/src/triggerRouteHooks.ts

@@ -80,39 +80,38 @@ export const loadAndRunRouteHooks = async ({
   }
 
   let currentRouteHooks: RouteHooks
-
+  let rhOutput: RouteHookOutput = defaultRouteHookOutput
   // Pull out the first path
   // Remember shift will mutate the array
   const routeHookPath = paths.shift()
 
-  if (!routeHookPath) {
-    return defaultRouteHookOutput
-  } else {
-    try {
+  try {
+    // Sometimes the appRouteHook is null, so we can skip it
+    if (routeHookPath) {
       currentRouteHooks = await loadModule(routeHookPath)
 
       // Step 2, run the route hooks
-      const rhOutput = await triggerRouteHooks({
+      rhOutput = await triggerRouteHooks({
         routeHooks: currentRouteHooks,
         req: reqMeta.req,
         parsedParams: reqMeta.parsedParams,
         previousOutput,
       })
+    }
 
-      if (paths.length > 0) {
-        // Step 3, recursively call this function
-        return loadAndRunRouteHooks({
-          paths,
-          reqMeta,
-          previousOutput: rhOutput,
-          viteDevServer,
-        })
-      } else {
-        return rhOutput
-      }
-    } catch (e) {
-      console.error(`Error loading route hooks in ${routeHookPath}}`)
-      throw new Error(e as any)
+    if (paths.length > 0) {
+      // Step 3, recursively call this function
+      return loadAndRunRouteHooks({
+        paths,
+        reqMeta,
+        previousOutput: rhOutput,
+        viteDevServer,
+      })
+    } else {
+      return rhOutput
     }
+  } catch (e) {
+    console.error(`Error loading route hooks in ${routeHookPath}}`)
+    throw new Error(e as any)
   }
 }