Merge branch 'fix/more-streaming-fixes' into try/apollo-ssr-stream

* fix/more-streaming-fixes:
  Try injection state isolation
  Also inject before final
  Fix type error
  fix(router): Prevent rerendering authenticated routes on hash change (redwoodjs#9007)

packages/router/src/AuthenticatedRoute.tsx

@@ -12,8 +12,9 @@ interface AuthenticatedRouteProps {
   whileLoadingAuth?: () => React.ReactElement | null
   private?: boolean
 }
-
-export function AuthenticatedRoute(props: AuthenticatedRouteProps) {
+export const AuthenticatedRoute: React.FC<AuthenticatedRouteProps> = (
+  props
+) => {
   const {
     private: isPrivate,
     unauthenticated,

packages/router/src/__tests__/router.test.tsx

@@ -36,7 +36,7 @@ import {
   Route,
   Private,
   Redirect,
-  routes,
+  routes as generatedRoutes,
   Link,
   navigate,
   back,
@@ -45,7 +45,7 @@ import {
 import { useLocation } from '../location'
 import { useParams } from '../params'
 import { Set } from '../Set'
-import { Spec } from '../util'
+import type { Spec, GeneratedRoutesMap } from '../util'
 
 /** running into intermittent test timeout behavior in https://github.com/redwoodjs/redwood/pull/4992
  attempting to work around by bumping the default timeout of 5000 */
@@ -59,9 +59,16 @@ type UnknownAuthContextInterface = AuthContextInterface<
   unknown,
   unknown,
   unknown,
+  unknown,
+  unknown,
+  unknown,
+  unknown,
   unknown
 >
 
+// The types are generated in the user's project
+const routes = generatedRoutes as GeneratedRoutesMap
+
 function createDummyAuthContextValues(
   partial: Partial<UnknownAuthContextInterface>
 ) {

packages/router/src/router.tsx

@@ -27,6 +27,7 @@ import {
   TrailingSlashesTypes,
   validatePath,
 } from './util'
+import type { Wrappers } from './util'
 
 import type { AvailableRoutes } from './index'
 
@@ -203,7 +204,7 @@ const LocationAwareRouter: React.FC<RouterProps> = ({
 }
 
 interface WrappedPageProps {
-  wrappers: ReactNode[]
+  wrappers: Wrappers
   routeLoaderElement: ReactNode
   setProps: Record<any, any>
 }
@@ -220,27 +221,32 @@ interface WrappedPageProps {
  */
 const WrappedPage = memo(
   ({ wrappers, routeLoaderElement, setProps }: WrappedPageProps) => {
+    // @NOTE: don't mutate the wrappers array, it causes full page re-renders
+    // Instead just create a new array with the AuthenticatedRoute wrapper
+    let wrappersWithAuthMaybe = wrappers
     if (setProps.private) {
       if (!setProps.unauthenticated) {
         throw new Error(
           'You must specify an `unauthenticated` route when marking a Route as private'
         )
       }
 
-      wrappers.unshift(AuthenticatedRoute as any)
+      wrappersWithAuthMaybe = [AuthenticatedRoute, ...wrappers]
     }
 
-    if (wrappers.length > 0) {
-      // If wrappers exist e.g. [a,b,c] -> <a><b><c><routeLoader></c></b></a>
-      return wrappers.reduceRight((acc, wrapper) => {
+    if (wrappersWithAuthMaybe.length > 0) {
+      // If wrappers exist e.g. [a,b,c] -> <a><b><c><routeLoader></c></b></a> and returns a single ReactNode
+      // Wrap AuthenticatedRoute this way, because if we mutate the wrappers array itself
+      // it causes full rerenders of the page
+      return wrappersWithAuthMaybe.reduceRight((acc, wrapper) => {
         return React.createElement(
           wrapper as any,
           {
             ...setProps,
           },
           acc ? acc : routeLoaderElement
         )
-      }, undefined) as any
+      }, undefined as ReactNode)
     }
 
     return routeLoaderElement

packages/router/src/util.ts

@@ -452,13 +452,14 @@ export type GeneratedRoutesMap = {
 }
 
 type RoutePath = string
+export type Wrappers = Array<(props: any) => ReactNode>
 interface AnalyzedRoute {
   path: RoutePath
   name: string | null
   whileLoadingPage?: WhileLoadingPage
   page: PageType | null
   redirect: string | null
-  wrappers: ReactNode[]
+  wrappers: Wrappers
   setProps: Record<any, any>
   setId: number
 }
@@ -476,7 +477,7 @@ export function analyzeRoutes(
   interface RecurseParams {
     nodes: ReturnType<typeof Children.toArray>
     whileLoadingPageFromSet?: WhileLoadingPage
-    wrappersFromSet?: ReactNode[]
+    wrappersFromSet?: Wrappers
     // we don't know, or care about, what props users are passing down
     propsFromSet?: Record<string, unknown>
     setId?: number

packages/vite/src/runFeServer.ts

@@ -212,10 +212,10 @@ export async function runFeServer() {
       const pageWithJs = currentRoute.renderMode !== 'html'
       // @NOTE have to add slash so subpaths still pick up the right file
       const bootstrapModules = pageWithJs
-        ? [
+        ? ([
             '/' + indexEntry.file,
             currentRoute.bundle && '/' + currentRoute.bundle,
-          ].filter(Boolean)
+          ].filter(Boolean) as string[])
         : undefined
 
       const isSeoCrawler = checkUaForSeoCrawler(req.get('user-agent'))

packages/vite/src/streamHelpers.ts

@@ -11,6 +11,8 @@ import type { TagDescriptor } from '@redwoodjs/web'
 import {
   ServerHtmlProvider,
   ServerInjectedHtml,
+  createInjector,
+  RenderCallback,
 } from '@redwoodjs/web/dist/components/ServerInject'
 
 interface RenderToStreamArgs {
@@ -50,17 +52,29 @@ export function reactRenderToStream({
     meta: metaTags,
   })
 
+  // This ensures an isolated state for each request
+  const { injectionState, injectToPage } = createInjector()
+  console.log(`👉 \n ~ file: streamHelpers.ts:57 ~ injectToPage:`, injectToPage)
+  console.log(
+    `👉 \n ~ file: streamHelpers.ts:57 ~ injectionState:`,
+    injectionState
+  )
+
+  // 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,
@@ -83,9 +97,11 @@ export function reactRenderToStream({
 function createServerInjectionStream({
   outputStream,
   onFinal,
+  injectionState,
 }: {
   outputStream: Writable
   onFinal: () => void
+  injectionState: Set<RenderCallback>
 }) {
   return new Writable({
     write(chunk, encoding, next) {
@@ -97,7 +113,9 @@ function createServerInjectionStream({
         const [beforeClosingHead, afterClosingHead] = split
 
         const elementsInjectedToHead = renderToString(
-          React.createElement(ServerInjectedHtml)
+          React.createElement(ServerInjectedHtml, {
+            injectionState,
+          })
         )
 
         const outputBuffer = Buffer.from(
@@ -117,6 +135,16 @@ function createServerInjectionStream({
       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/web/src/components/ServerInject.tsx

@@ -1,56 +1,58 @@
-import React, {
-  Fragment,
-  PropsWithChildren,
-  ReactNode,
-  useContext,
-} from 'react'
+import React, { Fragment, ReactNode, useContext } from 'react'
 
 /**
  *
  * Inspired by Next's useServerInsertedHTML, originally designed for CSS-in-JS
  * for now it seems the only way to inject html with streaming is to use a context
  *
+ * We use this for <head> tags, and for apollo cache hydration
+ *
  * Until https://github.com/reactjs/rfcs/pull/219 makes it into react
  *
  */
 
-type RenderCallback = () => ReactNode
-
-const insertCallbacks: Set<RenderCallback> = new Set([])
+export type RenderCallback = () => ReactNode
 
 export const ServerHtmlContext = React.createContext<
   ((things: RenderCallback) => void) | null
 >(null)
 
-const injectToHead = (renderCallback: RenderCallback) => {
-  insertCallbacks.add(renderCallback)
-}
+/**
+ *
+ *  Use this factory, once per request.
+ *  This is to ensure that injectionState is isolated to the request
+ *  and not shared between requests.
+ */
+export const createInjector = () => {
+  const injectionState: Set<RenderCallback> = new Set([])
 
-// @MARK: I don't know why Next don't do this also?
-// My understanding: put the inject function in a context so that
-// ServerInjectedHTML component rerenders, even though the state isn't inside the context
-export const ServerHtmlProvider: React.FC<PropsWithChildren> = ({
-  children,
-}) => {
-  return (
-    <ServerHtmlContext.Provider value={injectToHead}>
-      {children}
-    </ServerHtmlContext.Provider>
-  )
+  const injectToPage = (renderCallback: RenderCallback) => {
+    injectionState.add(renderCallback)
+  }
+
+  return { injectToPage, injectionState }
 }
 
-export const ServerInjectedHtml = () => {
+// @NOTE do not instatiate the provider value here, so that we can ensure
+// context isolation. This is done in streamHelpers currently,
+// using the createInjector factory, once per request
+export const ServerHtmlProvider = ServerHtmlContext.Provider
+
+export const ServerInjectedHtml = ({
+  injectionState,
+}: {
+  injectionState: Set<RenderCallback>
+}) => {
   const serverInsertedHtml = []
-  for (const callback of insertCallbacks) {
+  for (const callback of injectionState) {
     serverInsertedHtml.push(callback())
 
-    // Remove it from the set so its not called again
-    insertCallbacks.delete(callback)
+    // Remove it from the set so its not rendered again
+    injectionState.delete(callback)
   }
 
-  // @MARK: using i as key here might be problematic, no?
   return serverInsertedHtml.map((html, i) => {
-    return <Fragment key={i}>{html}</Fragment>
+    return <Fragment key={`rwjs-server-injected-${i}`}>{html}</Fragment>
   })
 }