onHeaders updates for App Pages (#58410)

This serves to address a few problems as it relates to headers handling within the App Page render pipeline.

1. During static generation, we should not write to the response directly as the revalidation for static responses is performed out of band from the actual request/response as a stale version can be served early. This is partially addressed by modifying the `onHeaders` to only write when there it is not during static generation, but long term I'd hope to see that we can instead pass a immutable request and `null` for the response to help indicate to the render pipeline that it should instead persist the data into the `RenderResult` objects. Alternatively, the response could use a mocked version that would then supersede the `RenderResult` and be used instead.

2. The types for rendering App Pages were merged incorrectly, so this additionally restructures the way that we initialize those renderers such that only the options that are applicable to each renderer is passed down.

Author: wyattjoh

packages/next/src/export/routes/app-page.ts

@@ -184,7 +184,7 @@ export async function exportAppPage(
       )
     }
 
-    const headers = { ...metadata.extraHeaders }
+    const headers = { ...metadata.headers }
 
     if (fetchTags) {
       headers[NEXT_CACHE_TAGS_HEADER] = fetchTags

packages/next/src/lib/recursive-delete.ts

@@ -2,9 +2,7 @@ import type { Dirent } from 'fs'
 import { promises } from 'fs'
 import { join, isAbsolute, dirname } from 'path'
 import isError from './is-error'
-
-const sleep = (timeout: number) =>
-  new Promise((resolve) => setTimeout(resolve, timeout))
+import { wait } from './wait'
 
 const unlinkPath = async (p: string, isDir = false, t = 1): Promise<void> => {
   try {
@@ -22,7 +20,7 @@ const unlinkPath = async (p: string, isDir = false, t = 1): Promise<void> => {
         code === 'EMFILE') &&
       t < 3
     ) {
-      await sleep(t * 100)
+      await wait(t * 100)
       return unlinkPath(p, isDir, t++)
     }
 

packages/next/src/server/app-render/action-handler.ts

@@ -6,6 +6,9 @@ import type {
 } from 'http'
 import type { WebNextRequest } from '../base-http/web'
 import type { SizeLimit } from '../../../types'
+import type { RequestStore } from '../../client/components/request-async-storage.external'
+import type { AppRenderContext, GenerateFlight } from './app-render'
+import type { AppPageModule } from '../../server/future/route-modules/app-page/module'
 
 import {
   ACTION,
@@ -20,7 +23,6 @@ import {
 import RenderResult from '../render-result'
 import type { StaticGenerationStore } from '../../client/components/static-generation-async-storage.external'
 import { FlightRenderResult } from './flight-render-result'
-import type { ActionAsyncStorage } from '../../client/components/action-async-storage.external'
 import {
   filterReqHeaders,
   actionsForbiddenHeaders,
@@ -30,12 +32,10 @@ import {
   getModifiedCookieValues,
 } from '../web/spec-extension/adapters/request-cookies'
 
-import type { RequestStore } from '../../client/components/request-async-storage.external'
 import {
   NEXT_CACHE_REVALIDATED_TAGS_HEADER,
   NEXT_CACHE_REVALIDATE_TAG_TOKEN_HEADER,
 } from '../../lib/constants'
-import type { AppRenderContext, GenerateFlight } from './app-render'
 
 function formDataFromSearchQueryString(query: string) {
   const searchParams = new URLSearchParams(query)
@@ -244,7 +244,7 @@ export async function handleAction({
 }: {
   req: IncomingMessage
   res: ServerResponse
-  ComponentMod: any
+  ComponentMod: AppPageModule
   serverModuleMap: {
     [id: string]: {
       id: string
@@ -288,6 +288,12 @@ export async function handleAction({
     return
   }
 
+  if (staticGenerationStore.isStaticGeneration) {
+    throw new Error(
+      "Invariant: server actions can't be handled during static rendering"
+    )
+  }
+
   const originDomain =
     typeof req.headers['origin'] === 'string'
       ? new URL(req.headers['origin']).host
@@ -373,9 +379,7 @@ export async function handleAction({
   )
   let bound = []
 
-  const { actionAsyncStorage } = ComponentMod as {
-    actionAsyncStorage: ActionAsyncStorage
-  }
+  const { actionAsyncStorage } = ComponentMod
 
   let actionResult: RenderResult | undefined
   let formState: any | undefined

packages/next/src/server/app-render/app-render.tsx

@@ -73,6 +73,7 @@ import { getAssetQueryString } from './get-asset-query-string'
 import { setReferenceManifestsSingleton } from './action-encryption-utils'
 import { createStaticRenderer } from './static/static-renderer'
 import { MissingPostponeDataError } from './is-missing-postpone-error'
+import { DetachedPromise } from '../../lib/detached-promise'
 
 export type GetDynamicParamFromSegment = (
   // [slug] / [[slug]] / [...slug]
@@ -654,7 +655,13 @@ async function renderToHTMLOrFlightImpl(
     createServerInsertedHTML()
 
   getTracer().getRootSpanAttributes()?.set('next.route', pagePath)
-  const bodyResult = getTracer().wrap(
+
+  // Create a promise that will help us signal when the headers have been
+  // written to the metadata for static generation as they aren't written to the
+  // response directly.
+  const onHeadersFinished = new DetachedPromise<void>()
+
+  const renderToStream = getTracer().wrap(
     AppRenderSpan.getBodyResult,
     {
       spanName: `render route (app) ${pagePath}`,
@@ -703,21 +710,13 @@ async function renderToHTMLOrFlightImpl(
         nonce
       )
 
-      const renderer = createStaticRenderer({
-        ppr: renderOpts.experimental.ppr,
-        isStaticGeneration: staticGenerationStore.isStaticGeneration,
-        postponed: renderOpts.postponed
-          ? JSON.parse(renderOpts.postponed)
-          : null,
-      })
-
       const ServerComponentsRenderer = createServerComponentsRenderer(tree, {
         ctx,
         preinitScripts,
         options: serverComponentsRenderOpts,
       })
 
-      const content = (
+      const children = (
         <HeadManagerContext.Provider
           value={{
             appDir: true,
@@ -736,34 +735,51 @@ async function renderToHTMLOrFlightImpl(
         hasPostponed,
       })
 
-      function onHeaders(headers: Headers): void {
-        // Copy headers created by React into the response object.
-        headers.forEach((value: string, key: string) => {
-          res.appendHeader(key, value)
-          if (!extraRenderResultMeta.extraHeaders) {
-            extraRenderResultMeta.extraHeaders = {}
-          }
-          extraRenderResultMeta.extraHeaders[key] = value
-        })
-      }
-
-      try {
-        const renderStream = await renderer.render(content, {
+      const renderer = createStaticRenderer({
+        ppr: renderOpts.experimental.ppr,
+        isStaticGeneration,
+        // If provided, the postpone state should be parsed as JSON so it can be
+        // provided to React.
+        postponed: renderOpts.postponed
+          ? JSON.parse(renderOpts.postponed)
+          : null,
+        streamOptions: {
           onError: htmlRendererErrorHandler,
-          onHeaders: onHeaders,
+          onHeaders: (headers: Headers) => {
+            // If this is during static generation, we shouldn't write to the
+            // headers object directly, instead we should add to the render
+            // result.
+            if (isStaticGeneration) {
+              headers.forEach((value, key) => {
+                extraRenderResultMeta.headers ??= {}
+                extraRenderResultMeta.headers[key] = value
+              })
+
+              // Resolve the promise to continue the stream.
+              onHeadersFinished.resolve()
+            } else {
+              headers.forEach((value, key) => {
+                res.appendHeader(key, value)
+              })
+            }
+          },
           maxHeadersLength: 600,
           nonce,
           bootstrapScripts: [bootstrapScript],
           formState,
-        })
+        },
+      })
 
-        const { stream, postponed } = renderStream
+      try {
+        let { stream, postponed } = await renderer.render(children)
 
+        // If the stream was postponed, we need to add the result to the
+        // metadata so that it can be resumed later.
         if (postponed) {
           extraRenderResultMeta.postponed = JSON.stringify(postponed)
 
-          // If this render generated a postponed state, we don't want to add
-          // any other data to the response.
+          // We don't need to "continue" this stream now as it's continued when
+          // we resume the stream.
           return stream
         }
 
@@ -786,10 +802,10 @@ async function renderToHTMLOrFlightImpl(
         }
 
         if (renderOpts.postponed) {
-          return continuePostponedFizzStream(stream, options)
+          return await continuePostponedFizzStream(stream, options)
         }
 
-        return continueFizzStream(stream, options)
+        return await continueFizzStream(stream, options)
       } catch (err: any) {
         if (
           err.code === 'NEXT_STATIC_GEN_BAILOUT' ||
@@ -967,8 +983,8 @@ async function renderToHTMLOrFlightImpl(
     ComponentMod,
     serverModuleMap,
     generateFlight,
-    staticGenerationStore: staticGenerationStore,
-    requestStore: requestStore,
+    staticGenerationStore,
+    requestStore,
     serverActions,
     ctx,
   })
@@ -978,7 +994,7 @@ async function renderToHTMLOrFlightImpl(
     if (actionRequestResult.type === 'not-found') {
       const notFoundLoaderTree = createNotFoundLoaderTree(loaderTree)
       return new RenderResult(
-        await bodyResult({
+        await renderToStream({
           asNotFound: true,
           tree: notFoundLoaderTree,
           formState,
@@ -996,15 +1012,20 @@ async function renderToHTMLOrFlightImpl(
   }
 
   const renderResult = new RenderResult(
-    await bodyResult({
+    await renderToStream({
       asNotFound: isNotFoundPath,
       tree: loaderTree,
       formState,
     }),
     {
       ...extraRenderResultMeta,
+      // Wait for and collect the flight payload data if we don't have it
+      // already.
       pageData: await stringifiedFlightPayloadPromise,
-      waitUntil: Promise.all(staticGenerationStore.pendingRevalidates || []),
+      // If we have pending revalidates, wait until they are all resolved.
+      waitUntil: staticGenerationStore.pendingRevalidates
+        ? Promise.all(staticGenerationStore.pendingRevalidates)
+        : undefined,
     }
   )
 
@@ -1015,8 +1036,30 @@ async function renderToHTMLOrFlightImpl(
   })
 
   if (staticGenerationStore.isStaticGeneration) {
+    // Collect the entire render result to a string (by streaming it to a
+    // string).
     const htmlResult = await renderResult.toUnchunkedString(true)
 
+    // Timeout after 1.5 seconds for the headers to write. If it takes
+    // longer than this it's more likely that the stream has stalled and
+    // there is a React bug. The headers will then be updated in the render
+    // result below when the metadata is re-added to the new render result.
+    const onTimeout = new DetachedPromise<never>()
+    const timeout = setTimeout(() => {
+      onTimeout.reject(
+        new Error(
+          'Timeout waiting for headers to be emitted, this is a bug in Next.js'
+        )
+      )
+    }, 1500)
+
+    // Race against the timeout and the headers being written.
+    await Promise.race([onHeadersFinished.promise, onTimeout.promise])
+
+    // It got here, which means it did not reject, so clear the timeout to avoid
+    // it from rejecting again (which is a no-op anyways).
+    clearTimeout(timeout)
+
     if (
       // if PPR is enabled
       renderOpts.experimental.ppr &&