customized flush effects

change cond

inject initial styles to end of head, add test

revert pnpm lock

Author: huozhi

packages/next/build/webpack-config.ts

@@ -974,7 +974,7 @@ export default async function getBaseWebpackConfig(
   }
 
   const rscSharedRegex =
-    /(node_modules\/react\/|\/shared\/lib\/(head-manager-context|router-context)\.js|node_modules\/styled-jsx\/)/
+    /(node_modules\/react\/|\/shared\/lib\/(head-manager-context|router-context|flush-effects)\.js|node_modules\/styled-jsx\/)/
 
   let webpackConfig: webpack.Configuration = {
     parallelism: Number(process.env.NEXT_WEBPACK_PARALLELISM) || undefined,

packages/next/client/components/hooks-client-context.ts

@@ -1,5 +1,5 @@
 import { createContext } from 'react'
-import { NextParsedUrlQuery } from '../../server/request-meta'
+import type { NextParsedUrlQuery } from '../../server/request-meta'
 
 export const SearchParamsContext = createContext<NextParsedUrlQuery>(
   null as any

packages/next/client/components/hooks-client.ts

@@ -12,6 +12,11 @@ import {
   LayoutRouterContext,
 } from '../../shared/lib/app-router-context'
 
+export {
+  FlushEffectsContext,
+  useFlushEffects,
+} from '../../shared/lib/flush-effects'
+
 /**
  * Get the current search params. For example useSearchParams() would return {"foo": "bar"} when ?foo=bar
  */

packages/next/client/components/match-segments.ts

@@ -1,4 +1,4 @@
-import { Segment } from '../../server/app-render'
+import type { Segment } from '../../server/app-render'
 
 export const matchSegment = (
   existingSegment: Segment,

packages/next/server/app-render.tsx

@@ -6,7 +6,6 @@ import React from 'react'
 import { ParsedUrlQuery, stringify as stringifyQuery } from 'querystring'
 import { createFromReadableStream } from 'next/dist/compiled/react-server-dom-webpack'
 import { renderToReadableStream } from 'next/dist/compiled/react-server-dom-webpack/writer.browser.server'
-import { StyleRegistry, createStyleRegistry } from 'styled-jsx'
 import { NextParsedUrlQuery } from './request-meta'
 import RenderResult from './render-result'
 import {
@@ -23,6 +22,7 @@ import { htmlEscapeJsonString } from './htmlescape'
 import { shouldUseReactRoot, stripInternalQueries } from './utils'
 import { NextApiRequestCookies } from './api-utils'
 import { matchSegment } from '../client/components/match-segments'
+import { FlushEffectsContext } from '../client/components/hooks-client'
 
 // this needs to be required lazily so that `next-server` can set
 // the env before we require
@@ -985,18 +985,26 @@ export async function renderToHTMLOrFlight(
     }
   )
 
-  /**
-   * Style registry for styled-jsx
-   */
-  const jsxStyleRegistry = createStyleRegistry()
+  let flushEffectsHandler: (() => React.ReactNode) | null = null
+  function FlushEffects({ children }: { children: JSX.Element }) {
+    // Reset flushEffectsHandler on each render
+    flushEffectsHandler = null
+    const setFlushEffectsHandler = React.useCallback(
+      (handler: () => React.ReactNode) => {
+        if (flushEffectsHandler)
+          throw new Error(
+            'The `useFlushEffects` hook cannot be used more than once.'
+          )
+        flushEffectsHandler = handler
+      },
+      []
+    )
 
-  /**
-   * styled-jsx styles as React Component
-   */
-  const styledJsxFlushEffect = (): React.ReactNode => {
-    const styles = jsxStyleRegistry.styles()
-    jsxStyleRegistry.flush()
-    return <>{styles}</>
+    return (
+      <FlushEffectsContext.Provider value={setFlushEffectsHandler}>
+        {children}
+      </FlushEffectsContext.Provider>
+    )
   }
 
   /**
@@ -1015,11 +1023,18 @@ export async function renderToHTMLOrFlight(
   const generateStaticHTML = supportsDynamicHTML !== true
   const bodyResult = async () => {
     const content = (
-      <StyleRegistry registry={jsxStyleRegistry}>
+      <FlushEffects>
         <ServerComponentsRenderer />
-      </StyleRegistry>
+      </FlushEffects>
     )
 
+    const flushEffectHandler = (): string => {
+      const flushed = ReactDOMServer.renderToString(
+        <>{flushEffectsHandler && flushEffectsHandler()}</>
+      )
+      return flushed
+    }
+
     const renderStream = await renderToInitialStream({
       ReactDOMServer,
       element: content,
@@ -1031,17 +1046,13 @@ export async function renderToHTMLOrFlight(
       },
     })
 
-    const flushEffectHandler = (): string => {
-      const flushed = ReactDOMServer.renderToString(styledJsxFlushEffect())
-      return flushed
-    }
-
     const hasConcurrentFeatures = !!runtime
 
     return await continueFromInitialStream(renderStream, {
       dataStream: serverComponentsInlinedTransformStream?.readable,
       generateStaticHTML: generateStaticHTML || !hasConcurrentFeatures,
       flushEffectHandler,
+      flushEffectsToHead: true,
       initialStylesheets,
     })
   }

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

@@ -137,7 +137,7 @@ export function createFlushEffectStream(
 }
 
 export function createHeadInjectionTransformStream(
-  inject: string
+  inject: () => string
 ): TransformStream<Uint8Array, Uint8Array> {
   let injected = false
   return new TransformStream({
@@ -147,7 +147,7 @@ export function createHeadInjectionTransformStream(
       if (!injected && (index = content.indexOf('</head')) !== -1) {
         injected = true
         const injectedContent =
-          content.slice(0, index) + inject + content.slice(index)
+          content.slice(0, index) + inject() + content.slice(index)
         controller.enqueue(encodeText(injectedContent))
       } else {
         controller.enqueue(chunk)
@@ -175,12 +175,14 @@ export async function continueFromInitialStream(
     dataStream,
     generateStaticHTML,
     flushEffectHandler,
+    flushEffectsToHead,
     initialStylesheets,
   }: {
     suffix?: string
     dataStream?: ReadableStream<Uint8Array>
     generateStaticHTML: boolean
     flushEffectHandler?: () => string
+    flushEffectsToHead: boolean
     initialStylesheets?: string[]
   }
 ): Promise<ReadableStream<Uint8Array>> {
@@ -193,15 +195,22 @@ export async function continueFromInitialStream(
 
   const transforms: Array<TransformStream<Uint8Array, Uint8Array>> = [
     createBufferedTransformStream(),
-    flushEffectHandler ? createFlushEffectStream(flushEffectHandler) : null,
+    flushEffectHandler && !flushEffectsToHead
+      ? createFlushEffectStream(flushEffectHandler)
+      : null,
     suffixUnclosed != null ? createDeferredSuffixStream(suffixUnclosed) : null,
     dataStream ? createInlineDataStream(dataStream) : null,
     suffixUnclosed != null ? createSuffixStream(closeTag) : null,
-    createHeadInjectionTransformStream(
-      (initialStylesheets || [])
+    createHeadInjectionTransformStream(() => {
+      const inlineStyleLinks = (initialStylesheets || [])
         .map((href) => `<link rel="stylesheet" href="/_next/${href}">`)
         .join('')
-    ),
+      // TODO-APP: Inject flush effects to end of head in app layout rendering, to avoid
+      // hydration errors. Remove this once it's ready to be handled by react itself.
+      const flushEffectsContent =
+        flushEffectHandler && flushEffectsToHead ? flushEffectHandler() : ''
+      return inlineStyleLinks + flushEffectsContent
+    }),
   ].filter(nonNullable)
 
   return transforms.reduce(

packages/next/server/render.tsx

@@ -1269,6 +1269,7 @@ export async function renderToHTML(
           dataStream: serverComponentsInlinedTransformStream?.readable,
           generateStaticHTML,
           flushEffectHandler,
+          flushEffectsToHead: false,
         })
       }
 

packages/next/shared/lib/flush-effects.tsx

@@ -0,0 +1,14 @@
+import React, { createContext, useContext } from 'react'
+
+export type FlushEffectsHook = (callbacks: () => React.ReactNode) => void
+
+export const FlushEffectsContext = createContext<FlushEffectsHook | null>(
+  null as any
+)
+
+export function useFlushEffects(callbacks: () => React.ReactNode): void {
+  const flushEffectsImpl = useContext(FlushEffectsContext)
+  // Should have no effects on client where there's no flush effects provider
+  if (!flushEffectsImpl) return
+  return flushEffectsImpl(callbacks)
+}

test/e2e/app-dir/rsc-basic.test.ts

@@ -41,6 +41,8 @@ describe('app dir - react server components', () => {
         'next.config.js': new FileRef(path.join(appDir, 'next.config.js')),
       },
       dependencies: {
+        'styled-jsx': 'latest',
+        'styled-components': '6.0.0-alpha.5',
         react: 'experimental',
         'react-dom': 'experimental',
       },
@@ -320,11 +322,15 @@ describe('app dir - react server components', () => {
     expect(content).toContain('bar.server.js:')
   })
 
-  it.skip('should SSR styled-jsx correctly', async () => {
-    const html = await renderViaHTTP(next.url, '/styled-jsx')
-    const styledJsxClass = getNodeBySelector(html, 'h1').attr('class')
+  it('should render initial styles of css-in-js in SSR correctly', async () => {
+    const html = await renderViaHTTP(next.url, '/css-in-js')
+    const head = getNodeBySelector(html, 'head').html()
 
-    expect(html).toContain(`h1.${styledJsxClass}{color:red}`)
+    // from styled-jsx
+    expect(head).toMatch(/{color:(\s*)purple;?}/)
+
+    // from styled-components
+    expect(head).toMatch(/{color:(\s*)blue;?}/)
   })
 
   it('should support streaming for flight response', async () => {

test/e2e/app-dir/rsc-basic/app/css-in-js/page.server.js

@@ -0,0 +1,11 @@
+import Comp from './styled-jsx.client'
+import StyledComp from './styled-components.client'
+
+export default function Page() {
+  return (
+    <div>
+      <Comp />
+      <StyledComp />
+    </div>
+  )
+}