Fix injecting CSS link tags before page content (#38559)

* inject the link tags before page content in layout router

* fix flight handling during client navigation

* fix buffering

Co-authored-by: kodiakhq[bot] <49736102+kodiakhq[bot]@users.noreply.github.com>

packages/next/client/app-index.tsx

@@ -32,10 +32,13 @@ self.__next_require__ = __webpack_require__
 // eslint-disable-next-line no-undef
 ;(self as any).__next_chunk_load__ = (chunk: string) => {
   if (chunk.endsWith('.css')) {
-    const link = document.createElement('link')
-    link.rel = 'stylesheet'
-    link.href = '/_next/' + chunk
-    document.head.appendChild(link)
+    const existingTag = document.querySelector(`link[href="${chunk}"]`)
+    if (!existingTag) {
+      const link = document.createElement('link')
+      link.rel = 'stylesheet'
+      link.href = '/_next/' + chunk
+      document.head.appendChild(link)
+    }
     return Promise.resolve()
   }
 
@@ -169,25 +172,30 @@ function useInitialServerResponse(cacheKey: string) {
   }
 
   const loadCssFromStreamData = (data: string) => {
-    const seg = data.split(':')
-    if (seg[0] === 'CSS') {
-      loadCss(seg.slice(1).join(':'))
+    if (data.startsWith('CSS:')) {
+      loadCss(data.slice(4).trim())
     }
   }
 
+  // TODO-APP: Refine the buffering code here to make it more correct.
   let buffer = ''
   const loadCssFromFlight = new TransformStream({
     transform(chunk, controller) {
       const data = new TextDecoder().decode(chunk)
       buffer += data
       let index
       while ((index = buffer.indexOf('\n')) !== -1) {
-        const line = buffer.slice(0, index)
+        const line = buffer.slice(0, index + 1)
         buffer = buffer.slice(index + 1)
-        loadCssFromStreamData(line)
+        if (line.startsWith('CSS:')) {
+          loadCssFromStreamData(line)
+        } else {
+          controller.enqueue(new TextEncoder().encode(line))
+        }
       }
-      if (!data.startsWith('CSS:')) {
-        controller.enqueue(chunk)
+      if (buffer && !buffer.startsWith('CSS:')) {
+        controller.enqueue(new TextEncoder().encode(buffer))
+        buffer = ''
       }
     },
     flush() {

packages/next/client/components/app-router.client.tsx

@@ -1,5 +1,5 @@
 import React, { useEffect } from 'react'
-import { createFromFetch } from 'next/dist/compiled/react-server-dom-webpack'
+import { createFromReadableStream } from 'next/dist/compiled/react-server-dom-webpack'
 import {
   AppRouterContext,
   AppTreeContext,
@@ -16,24 +16,75 @@ import {
   // LayoutSegmentsContext,
 } from './hooks-client-context'
 
-function fetchFlight(
-  url: URL,
-  flightRouterStateData: string
-): Promise<Response> {
+async function loadCss(cssChunkInfoJson: string) {
+  const data = JSON.parse(cssChunkInfoJson)
+  await Promise.all(
+    data.chunks.map((chunkId: string) => {
+      // load css related chunks
+      return (self as any).__next_chunk_load__(chunkId)
+    })
+  )
+  // In development mode, import css in dev when it's wrapped by style loader.
+  // In production mode, css are standalone chunk that doesn't need to be imported.
+  if (data.id) {
+    ;(self as any).__next_require__(data.id)
+  }
+}
+
+const loadCssFromStreamData = (data: string) => {
+  if (data.startsWith('CSS:')) {
+    loadCss(data.slice(4).trim())
+  }
+}
+
+function fetchFlight(url: URL, flightRouterStateData: string): ReadableStream {
   const flightUrl = new URL(url)
   const searchParams = flightUrl.searchParams
   searchParams.append('__flight__', '1')
   searchParams.append('__flight_router_state_tree__', flightRouterStateData)
 
-  return fetch(flightUrl.toString())
+  const { readable, writable } = new TransformStream()
+
+  // TODO-APP: Refine the buffering code here to make it more correct.
+  let buffer = ''
+  const loadCssFromFlight = new TransformStream({
+    transform(chunk, controller) {
+      const data = new TextDecoder().decode(chunk)
+      buffer += data
+      let index
+      while ((index = buffer.indexOf('\n')) !== -1) {
+        const line = buffer.slice(0, index + 1)
+        buffer = buffer.slice(index + 1)
+
+        if (line.startsWith('CSS:')) {
+          loadCssFromStreamData(line)
+        } else {
+          controller.enqueue(new TextEncoder().encode(line))
+        }
+      }
+      if (buffer && !buffer.startsWith('CSS:')) {
+        controller.enqueue(new TextEncoder().encode(buffer))
+        buffer = ''
+      }
+    },
+    flush() {
+      loadCssFromStreamData(buffer)
+    },
+  })
+
+  fetch(flightUrl.toString()).then((res) => {
+    res.body?.pipeThrough(loadCssFromFlight).pipeTo(writable)
+  })
+
+  return readable
 }
 
 export function fetchServerResponse(
   url: URL,
   flightRouterState: FlightRouterState
 ): { readRoot: () => FlightData } {
   const flightRouterStateData = JSON.stringify(flightRouterState)
-  return createFromFetch(fetchFlight(url, flightRouterStateData))
+  return createFromReadableStream(fetchFlight(url, flightRouterStateData))
 }
 
 function ErrorOverlay({
@@ -56,11 +107,13 @@ let initialParallelRoutes: CacheNode['parallelRoutes'] =
 export default function AppRouter({
   initialTree,
   initialCanonicalUrl,
+  initialStylesheets,
   children,
   hotReloader,
 }: {
   initialTree: FlightRouterState
   initialCanonicalUrl: string
+  initialStylesheets: string[]
   children: React.ReactNode
   hotReloader?: React.ReactNode
 }) {
@@ -259,6 +312,7 @@ export default function AppRouter({
                 // Root node always has `url`
                 // Provided in AppTreeContext to ensure it can be overwritten in layout-router
                 url: canonicalUrl,
+                stylesheets: initialStylesheets,
               }}
             >
               <ErrorOverlay>{cache.subTreeData}</ErrorOverlay>

packages/next/client/components/layout-router.client.tsx

@@ -235,7 +235,7 @@ export default function OuterLayoutRouter({
   childProp: ChildProp
   loading: React.ReactNode | undefined
 }) {
-  const { childNodes, tree, url } = useContext(AppTreeContext)
+  const { childNodes, tree, url, stylesheets } = useContext(AppTreeContext)
 
   let childNodesForParallelRouter = childNodes.get(parallelRouterKey)
   if (!childNodesForParallelRouter) {
@@ -256,6 +256,11 @@ export default function OuterLayoutRouter({
 
   return (
     <>
+      {stylesheets
+        ? stylesheets.map((href) => (
+            <link rel="stylesheet" href={`/_next/${href}`} key={href} />
+          ))
+        : null}
       {preservedSegments.map((preservedSegment) => {
         return (
           <LoadingBoundary loading={loading} key={preservedSegment}>

packages/next/server/app-render.tsx

@@ -209,7 +209,7 @@ function createServerComponentRenderer(
   }
 
   const writable = transformStream.writable
-  const ServerComponentWrapper = () => {
+  return function ServerComponentWrapper() {
     const reqStream = createRSCStream()
     const response = useFlightResponse(
       writable,
@@ -218,11 +218,8 @@ function createServerComponentRenderer(
       serverComponentManifest,
       cssFlightData
     )
-    const root = response.readRoot()
-    return root
+    return response.readRoot()
   }
-
-  return ServerComponentWrapper
 }
 
 type DynamicParamTypes = 'catchall' | 'optional-catchall' | 'dynamic'
@@ -327,6 +324,26 @@ function getSegmentParam(segment: string): {
   return null
 }
 
+function getCSSInlinedLinkTags(
+  ComponentMod: any,
+  serverComponentManifest: any
+) {
+  const importedServerCSSFiles: string[] =
+    ComponentMod.__client__?.__next_rsc_css__ || []
+
+  return Array.from(
+    new Set(
+      importedServerCSSFiles
+        .map((css) =>
+          css.endsWith('.css')
+            ? serverComponentManifest[css].default.chunks
+            : []
+        )
+        .flat()
+    )
+  )
+}
+
 function getCssFlightData(ComponentMod: any, serverComponentManifest: any) {
   const importedServerCSSFiles: string[] =
     ComponentMod.__client__?.__next_rsc_css__ || []
@@ -335,7 +352,7 @@ function getCssFlightData(ComponentMod: any, serverComponentManifest: any) {
     (css) => serverComponentManifest[css].default
   )
   if (process.env.NODE_ENV === 'development') {
-    return cssFiles.map((css) => `CSS:${JSON.stringify(css)}`).join('\n')
+    return cssFiles.map((css) => `CSS:${JSON.stringify(css)}`).join('\n') + '\n'
   }
 
   // Multiple css chunks could be merged into one by mini-css-extract-plugin,
@@ -345,10 +362,7 @@ function getCssFlightData(ComponentMod: any, serverComponentManifest: any) {
     return res
   }, new Set())
 
-  const cssFlight = Array.from(cssSet)
-    .map((css) => `CSS:${JSON.stringify({ chunks: [css] })}`)
-    .join('\n')
-  return cssFlight
+  return `CSS:${JSON.stringify({ chunks: [...cssSet] })}\n`
 }
 
 export async function renderToHTML(
@@ -793,6 +807,11 @@ export async function renderToHTML(
   // /blog/[slug] /blog/hello-world -> ['children', 'blog', 'children', ['slug', 'hello-world']]
   const initialTree = createFlightRouterStateFromLoaderTree(tree)
 
+  const initialStylesheets = getCSSInlinedLinkTags(
+    ComponentMod,
+    serverComponentManifest
+  )
+
   const { Component: ComponentTree } = createComponentTree({
     createSegmentPath: (child) => child,
     tree,
@@ -818,6 +837,7 @@ export async function renderToHTML(
             hotReloader={HotReloader && <HotReloader assetPrefix="" />}
             initialCanonicalUrl={initialCanonicalUrl}
             initialTree={initialTree}
+            initialStylesheets={initialStylesheets}
           >
             <ComponentTree />
           </AppRouter>
@@ -896,7 +916,6 @@ export async function renderToHTML(
     }
 
     return await continueFromInitialStream(renderStream, {
-      suffix: '',
       dataStream: serverComponentsInlinedTransformStream?.readable,
       generateStaticHTML: generateStaticHTML || !hasConcurrentFeatures,
       flushEffectHandler,

packages/next/server/node-web-streams-helper.ts

@@ -172,7 +172,7 @@ export async function continueFromInitialStream(
   const transforms: Array<TransformStream<Uint8Array, Uint8Array>> = [
     createBufferedTransformStream(),
     flushEffectHandler ? createFlushEffectStream(flushEffectHandler) : null,
-    suffixUnclosed != null ? createPrefixStream(suffixUnclosed) : null,
+    suffixUnclosed != null ? createDeferredPrefixStream(suffixUnclosed) : null,
     dataStream ? createInlineDataStream(dataStream) : null,
     suffixUnclosed != null ? createSuffixStream(closeTag) : null,
   ].filter(nonNullable)
@@ -197,6 +197,21 @@ export function createSuffixStream(
 
 export function createPrefixStream(
   prefix: string
+): TransformStream<Uint8Array, Uint8Array> {
+  let prefixEnqueued = false
+  return new TransformStream({
+    transform(chunk, controller) {
+      if (!prefixEnqueued) {
+        controller.enqueue(encodeText(prefix))
+        prefixEnqueued = true
+      }
+      controller.enqueue(chunk)
+    },
+  })
+}
+
+export function createDeferredPrefixStream(
+  prefix: string
 ): TransformStream<Uint8Array, Uint8Array> {
   let prefixFlushed = false
   let prefixPrefixFlushFinished: Promise<void> | null = null

packages/next/shared/lib/app-router-context.ts

@@ -28,6 +28,7 @@ export const AppTreeContext = React.createContext<{
   childNodes: CacheNode['parallelRoutes']
   tree: FlightRouterState
   url: string
+  stylesheets?: string[]
 }>(null as any)
 export const FullAppTreeContext = React.createContext<{
   tree: FlightRouterState