From afd9a4e27c30688503c48d78df7885d89d6ab17f Mon Sep 17 00:00:00 2001
From: Jiachi Liu <inbox@huozhi.im>
Date: Wed, 12 Jan 2022 15:13:49 +0100
Subject: [PATCH] ssr app tree

---
 packages/next/client/index.tsx                | 15 +++----------
 packages/next/server/render.tsx               | 21 +++++++++++--------
 .../app/components/container.client.js        |  3 ---
 .../app/components/container.server.js        |  3 +++
 .../test/index.test.js                        | 15 +++++++------
 5 files changed, 27 insertions(+), 30 deletions(-)
 delete mode 100644 test/integration/react-streaming-and-server-components/app/components/container.client.js
 create mode 100644 test/integration/react-streaming-and-server-components/app/components/container.server.js

diff --git a/packages/next/client/index.tsx b/packages/next/client/index.tsx
index f7e521ae8a4d1d3..42cc067bea31577 100644
--- a/packages/next/client/index.tsx
+++ b/packages/next/client/index.tsx
@@ -629,19 +629,10 @@ function AppContainer({
   )
 }
 
-function renderApp(
-  App: AppComponent | React.ComponentType,
-  appProps: AppProps
-) {
+function renderApp(App: AppComponent, appProps: AppProps) {
   if (process.env.__NEXT_RSC) {
-    const { Component, ...props } = appProps
-    const AppServerComponent = App as React.ComponentType
-    const ComponentType = Component as React.ComponentType<typeof props>
-    return (
-      <AppServerComponent>
-        <ComponentType {...props} />
-      </AppServerComponent>
-    )
+    const { Component, err: _, router: __, ...props } = appProps
+    return <Component {...props} />
   } else {
     return <App {...appProps} />
   }
diff --git a/packages/next/server/render.tsx b/packages/next/server/render.tsx
index db599763ec0c2d6..f922f8ca8579fa5 100644
--- a/packages/next/server/render.tsx
+++ b/packages/next/server/render.tsx
@@ -183,18 +183,13 @@ function enhanceComponents(
 }
 
 function renderApp(
-  App: AppType | React.ComponentType,
+  App: AppType,
   Component: React.ComponentType,
   router: ServerRouter,
   props: any
 ) {
   if (process.env.__NEXT_RSC) {
-    let AppServerComponent = App as React.ComponentType
-    return (
-      <AppServerComponent>
-        <Component {...props.pageProps} router={router} />
-      </AppServerComponent>
-    )
+    return <Component {...props.pageProps} router={router} />
   } else {
     return <App {...props} Component={Component} router={router} />
   }
@@ -360,6 +355,7 @@ const useRSCResponse = createRSCHook()
 function createServerComponentRenderer(
   cachePrefix: string,
   transformStream: TransformStream,
+  App: React.ComponentType,
   OriginalComponent: React.ComponentType,
   serverComponentManifest: NonNullable<RenderOpts['serverComponentManifest']>
 ) {
@@ -367,7 +363,9 @@ function createServerComponentRenderer(
   const ServerComponentWrapper = (props: any) => {
     const id = (React as any).useId()
     const reqStream = renderToReadableStream(
-      <OriginalComponent {...props} />,
+      <App>
+        <OriginalComponent {...props} />
+      </App>,
       serverComponentManifest
     )
 
@@ -381,6 +379,7 @@ function createServerComponentRenderer(
     rscCache.delete(id)
     return root
   }
+
   const Component = (props: any) => {
     return (
       <React.Suspense fallback={null}>
@@ -459,6 +458,7 @@ export async function renderToHTML(
     ? createServerComponentRenderer(
         cachePrefix,
         serverComponentsInlinedTransformStream!,
+        App as React.ComponentType,
         OriginalComponent,
         serverComponentManifest
       )
@@ -1095,8 +1095,11 @@ export async function renderToHTML(
   if (isResSent(res) && !isSSG) return null
 
   if (renderServerComponentData) {
+    const AppServerComponent = App as React.ComponentType
     const stream: ReadableStream = renderToReadableStream(
-      <OriginalComponent {...props.pageProps} {...serverComponentProps} />,
+      <AppServerComponent>
+        <OriginalComponent {...props.pageProps} {...serverComponentProps} />
+      </AppServerComponent>,
       serverComponentManifest
     )
     const reader = stream.getReader()
diff --git a/test/integration/react-streaming-and-server-components/app/components/container.client.js b/test/integration/react-streaming-and-server-components/app/components/container.client.js
deleted file mode 100644
index e75ce0e55acbddc..000000000000000
--- a/test/integration/react-streaming-and-server-components/app/components/container.client.js
+++ /dev/null
@@ -1,3 +0,0 @@
-export default function Container({ children }) {
-  return <div className="_app.server">{children}</div>
-}
diff --git a/test/integration/react-streaming-and-server-components/app/components/container.server.js b/test/integration/react-streaming-and-server-components/app/components/container.server.js
new file mode 100644
index 000000000000000..a6746b460b426a1
--- /dev/null
+++ b/test/integration/react-streaming-and-server-components/app/components/container.server.js
@@ -0,0 +1,3 @@
+export default function Container({ children }) {
+  return <div className="container.server">{children}</div>
+}
diff --git a/test/integration/react-streaming-and-server-components/test/index.test.js b/test/integration/react-streaming-and-server-components/test/index.test.js
index c0d9c4a66f96d5d..26607be02f22166 100644
--- a/test/integration/react-streaming-and-server-components/test/index.test.js
+++ b/test/integration/react-streaming-and-server-components/test/index.test.js
@@ -25,6 +25,7 @@ const nativeModuleTestAppDir = join(__dirname, '../unsupported-native-module')
 const distDir = join(__dirname, '../app/.next')
 const documentPage = new File(join(appDir, 'pages/_document.jsx'))
 const appPage = new File(join(appDir, 'pages/_app.js'))
+const appServerPage = new File(join(appDir, 'pages/_app.server.js'))
 const error500Page = new File(join(appDir, 'pages/500.js'))
 
 const documentWithGip = `
@@ -48,7 +49,7 @@ Document.getInitialProps = (ctx) => {
 `
 
 const rscAppPage = `
-import Container from '../components/container.client'
+import Container from '../components/container.server'
 export default function App({children}) {
   return <Container>{children}</Container>
 }
@@ -184,13 +185,15 @@ describe('concurrentFeatures - prod', () => {
 
 const customAppPageSuite = {
   runTests: (context) => {
-    it('should render app page', async () => {
-      const html = await renderViaHTTP(context.appPort, '/')
-      expect(html).toContain('_app.server')
+    it('should render container in app', async () => {
+      const indexHtml = await renderViaHTTP(context.appPort, '/')
+      const indexFlight = await renderViaHTTP(context.appPort, '/?__flight__=1')
+      expect(indexHtml).toContain('container.server')
+      expect(indexFlight).toContain('container.server')
     })
   },
-  before: () => appPage.write(rscAppPage),
-  after: () => appPage.delete(),
+  before: () => appServerPage.write(rscAppPage),
+  after: () => appServerPage.delete(),
 }
 
 runSuite('Custom App', 'dev', customAppPageSuite)