[pull] canary from vercel:canary

packages/next/errors.json

@@ -1286,5 +1286,6 @@
   "1285": "Cache Components error recovery expected an original resume data cache",
   "1286": "Route \"%s\": Could not validate that a segment in your UI has instant navigation.",
   "1287": "A render that hasn't started yet cannot be abandoned",
-  "1288": "Cannot determine late/early stage before starting the render"
+  "1288": "Cannot determine late/early stage before starting the render",
+  "1289": "Attempted to advance to stage %s but the render is limited to %s"
 }

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

@@ -933,6 +933,7 @@ async function generateStagedDynamicFlightRenderResultWeb(
     // but it can happen e.g. after a revalidation or conditionally for a param that wasn't prerendered.
     // we should change this to track sync IO, log an error and advance to dynamic.
     shouldTrackSyncIO: false,
+    finalStage: null,
   })
 
   // Initialize stale time tracking on the request store.
@@ -1091,6 +1092,7 @@ async function generateStagedDynamicFlightRenderResultNode(
     // but it can happen e.g. after a revalidation or conditionally for a param that wasn't prerendered.
     // we should change this to track sync IO, log an error and advance to dynamic.
     shouldTrackSyncIO: false,
+    finalStage: null,
   })
 
   // Initialize stale time tracking on the request store.
@@ -1278,6 +1280,7 @@ async function stagedRenderWithoutCachesInDevWeb(
     abortSignal: null,
     abandonController: null,
     shouldTrackSyncIO: false, // do not track sync IO (we don't have reliable stages)
+    finalStage: null,
   })
 
   const environmentName = () => {
@@ -1348,6 +1351,7 @@ async function stagedRenderWithoutCachesInDevNode(
     abortSignal: null,
     abandonController: null,
     shouldTrackSyncIO: false, // do not track sync IO (we don't have reliable stages)
+    finalStage: null,
   })
 
   const environmentName = () => {
@@ -1881,6 +1885,7 @@ async function finalRuntimeServerPrerender(
     abortSignal: finalServerController.signal,
     abandonController: null,
     shouldTrackSyncIO: true,
+    finalStage: RenderStage.Runtime,
   })
 
   const varyParamsAccumulator = createResponseVaryParamsAccumulator()
@@ -3708,6 +3713,7 @@ async function renderToStream(
             // but it can happen e.g. after a revalidation or conditionally for a param that wasn't prerendered.
             // we should change this to track sync IO, log an error and advance to dynamic.
             shouldTrackSyncIO: false,
+            finalStage: null,
           })
 
           requestStore.stale = INFINITE_CACHE
@@ -3839,6 +3845,7 @@ async function renderToStream(
             // but it can happen e.g. after a revalidation or conditionally for a param that wasn't prerendered.
             // we should change this to track sync IO, log an error and advance to dynamic.
             shouldTrackSyncIO: false,
+            finalStage: null,
           })
 
           requestStore.stale = INFINITE_CACHE
@@ -4656,6 +4663,7 @@ async function renderWithRestartOnCacheMissInDevWeb(
     abortSignal: initialDataController.signal,
     abandonController: initialAbandonController,
     shouldTrackSyncIO: true,
+    finalStage: null,
   })
 
   // Use a mutable resume data cache for the warmup. After the warmup we'll swap
@@ -4818,6 +4826,7 @@ async function renderWithRestartOnCacheMissInDevWeb(
     abortSignal: null,
     abandonController: null,
     shouldTrackSyncIO: true,
+    finalStage: null,
   })
 
   // We've filled the caches, so now we can render as usual,
@@ -4968,6 +4977,7 @@ async function renderWithRestartOnCacheMissInDevNode(
     abortSignal: initialDataController.signal,
     abandonController: initialAbandonController,
     shouldTrackSyncIO: true,
+    finalStage: null,
   })
 
   // Use a mutable resume data cache for the warmup. After the warmup we'll swap
@@ -5125,6 +5135,7 @@ async function renderWithRestartOnCacheMissInDevNode(
     abortSignal: null,
     abandonController: null,
     shouldTrackSyncIO: true,
+    finalStage: null,
   })
 
   // We've filled the caches, so now we can render as usual,
@@ -6492,6 +6503,7 @@ async function renderWithRestartOnCacheMissInValidation(
     abortSignal: initialDataController.signal,
     abandonController: initialAbandonController,
     shouldTrackSyncIO: true,
+    finalStage: null,
   })
 
   requestStore.resumeDataCache = prerenderResumeDataCache
@@ -6602,6 +6614,7 @@ async function renderWithRestartOnCacheMissInValidation(
     abortSignal: finalDataController.signal,
     abandonController: null,
     shouldTrackSyncIO: true,
+    finalStage: null,
   })
 
   requestStore.resumeDataCache = createRenderResumeDataCache(

packages/next/src/server/app-render/dynamic-rendering.ts

@@ -34,11 +34,12 @@ import React from 'react'
 import { DynamicServerError } from '../../client/components/hooks-server-context'
 import { StaticGenBailoutError } from '../../client/components/static-generation-bailout'
 import {
+  getStagedRenderingController,
   throwForMissingRequestStore,
   workUnitAsyncStorage,
 } from './work-unit-async-storage.external'
 import { workAsyncStorage } from '../app-render/work-async-storage.external'
-import { makeHangingPromise, getRuntimeStage } from '../dynamic-rendering-utils'
+import { makeHangingPromise } from '../dynamic-rendering-utils'
 import {
   METADATA_BOUNDARY_NAME,
   VIEWPORT_BOUNDARY_NAME,
@@ -556,25 +557,20 @@ export function createHangingInputAbortSignal(
       } else {
         // Otherwise we're in the final render and we should already have all
         // our caches filled.
-        // If the prerender uses stages, we have wait until the runtime stage,
-        // at which point all runtime inputs will be resolved.
-        // (otherwise, a runtime prerender might consider `cookies()` hanging
-        //  even though they'd resolve in the next task.)
+        // If the prerender uses stages, we have wait until the final stage.
+        // if an input didn't resolve at that point, then we can assume it never will.
         //
         // We might still be waiting on some microtasks so we
         // wait one tick before giving up. When we give up, we still want to
         // render the content of this cache as deeply as we can so that we can
         // suspend as deeply as possible in the tree or not at all if we don't
         // end up waiting for the input.
-        if (
-          // eslint-disable-next-line no-restricted-syntax -- We are discriminating between two different refined types and don't need an addition exhaustive switch here
-          workUnitStore.type === 'prerender-runtime' &&
-          workUnitStore.stagedRendering
-        ) {
-          const { stagedRendering } = workUnitStore
+
+        const stagedRendering = getStagedRenderingController(workUnitStore)
+        if (stagedRendering && stagedRendering.finalStage !== null) {
           stagedRendering
-            .waitForStage(getRuntimeStage(stagedRendering))
-            .then(() => scheduleOnNextTick(() => controller.abort()))
+            .waitForStage(stagedRendering.finalStage)
+            .then(() => scheduleOnNextTick(() => controller.abort()), noop)
         } else {
           scheduleOnNextTick(() => controller.abort())
         }
@@ -596,6 +592,8 @@ export function createHangingInputAbortSignal(
   }
 }
 
+function noop() {}
+
 export function annotateDynamicAccess(
   expression: string,
   prerenderStore: PrerenderStoreModern | ValidationStoreClient

packages/next/src/server/app-render/staged-rendering.ts

@@ -62,6 +62,7 @@ export class StagedRenderingController {
   private abortSignal: AbortSignal | null
   private abandonController: AbortController | null
   private shouldTrackSyncIO: boolean
+  public readonly finalStage: AdvanceableRenderStage | null
 
   currentStage: RenderStage = RenderStage.Before
 
@@ -81,14 +82,17 @@ export class StagedRenderingController {
     abortSignal,
     abandonController,
     shouldTrackSyncIO,
+    finalStage,
   }: {
     abortSignal: AbortSignal | null
     abandonController: AbortController | null
     shouldTrackSyncIO: boolean
+    finalStage: AdvanceableRenderStage | null
   }) {
     this.abortSignal = abortSignal
     this.abandonController = abandonController
     this.shouldTrackSyncIO = shouldTrackSyncIO
+    this.finalStage = finalStage
 
     if (abortSignal) {
       abortSignal.addEventListener(
@@ -268,6 +272,11 @@ export class StagedRenderingController {
   }
 
   advanceStage(targetStage: AdvanceableRenderStage) {
+    if (this.finalStage !== null && targetStage > this.finalStage) {
+      throw new InvariantError(
+        `Attempted to advance to stage ${RenderStage[targetStage]} but the render is limited to ${RenderStage[this.finalStage]}`
+      )
+    }
     // If we're already at the target stage or beyond, do nothing.
     // (this can happen e.g. if sync IO advanced us to the dynamic stage)
     if (targetStage <= this.currentStage) {
@@ -320,14 +329,17 @@ export class StagedRenderingController {
     stage: AdvanceableRenderStage,
     displayName: string | undefined,
     resolvedValue: T
-  ) {
-    const ioTriggerPromise = this.getStagePromise(stage)
-
-    const promise = makeDevtoolsIOPromiseFromIOTrigger(
-      ioTriggerPromise,
-      displayName,
-      resolvedValue
-    )
+  ): Promise<T> {
+    const stagePromise = this.getStagePromise(stage)
+
+    const promise =
+      process.env.NODE_ENV === 'development'
+        ? makeDevtoolsIOPromiseFromIOTrigger(
+            stagePromise,
+            displayName,
+            resolvedValue
+          )
+        : stagePromise.then(() => resolvedValue)
 
     // Analogously to `makeHangingPromise`, we might reject this promise if the signal is invoked.
     // (e.g. in the case where we don't want want the render to proceed to the dynamic stage and abort it).

packages/next/src/server/dynamic-rendering-utils.ts

@@ -5,10 +5,7 @@ import {
   type StagedRenderingController,
   isEarlyRenderStage,
 } from './app-render/staged-rendering'
-import type {
-  PrerenderStoreModernRuntime,
-  RequestStore,
-} from './app-render/work-unit-async-storage.external'
+import type { RequestStore } from './app-render/work-unit-async-storage.external'
 import { workUnitAsyncStorage } from './app-render/work-unit-async-storage.external'
 import { getServerReact, getClientReact } from './runtime-reacts.external'
 
@@ -130,31 +127,6 @@ export function getRuntimeStage(
     : RenderStage.Runtime
 }
 
-/**
- * Delays until the appropriate runtime stage based on the current stage of
- * the rendering pipeline:
- *
- * - Early stages → wait for EarlyRuntime
- *   (for runtime-prefetchable segments)
- * - Later stages → wait for Runtime
- *   (for segments not using runtime prefetch)
- *
- * This ensures that cookies()/headers()/etc. resolve at the right time for
- * each segment type.
- */
-export function delayUntilRuntimeStage<T>(
-  prerenderStore: PrerenderStoreModernRuntime,
-  result: Promise<T>
-): Promise<T> {
-  const { stagedRendering } = prerenderStore
-  if (!stagedRendering) {
-    return result
-  }
-  return stagedRendering
-    .waitForStage(getRuntimeStage(stagedRendering))
-    .then(() => result)
-}
-
 export function applyOwnerStack(error: Error): Error {
   if (process.env.NODE_ENV !== 'production') {
     let ownerStack: string | undefined | null

packages/next/src/server/request/cookies.ts

@@ -22,9 +22,9 @@ import {
 } from '../app-render/dynamic-rendering'
 import { StaticGenBailoutError } from '../../client/components/static-generation-bailout'
 import {
-  delayUntilRuntimeStage,
   makeDevtoolsIOAwarePromise,
   makeHangingPromise,
+  getRuntimeStage,
 } from '../dynamic-rendering-utils'
 import { createDedupedByCallsiteServerErrorLoggerDev } from '../create-deduped-by-callsite-server-error-logger'
 import { isRequestAPICallableInsideAfter } from './utils'
@@ -104,11 +104,18 @@ export function cookies(): Promise<ReadonlyRequestCookies> {
             workStore,
             workUnitStore
           )
-        case 'prerender-runtime':
-          return delayUntilRuntimeStage(
-            workUnitStore,
-            makeUntrackedCookies(workUnitStore.cookies)
-          )
+        case 'prerender-runtime': {
+          const { stagedRendering } = workUnitStore
+          if (stagedRendering) {
+            return stagedRendering.delayUntilStage(
+              getRuntimeStage(stagedRendering),
+              'cookies',
+              workUnitStore.cookies
+            )
+          } else {
+            return makeUntrackedCookies(workUnitStore.cookies)
+          }
+        }
         case 'private-cache':
           // Private caches are delayed until the runtime stage in use-cache-wrapper,
           // so we don't need an additional delay here.

packages/next/src/server/request/draft-mode.ts

@@ -19,9 +19,8 @@ import { createDedupedByCallsiteServerErrorLoggerDev } from '../create-deduped-b
 import { StaticGenBailoutError } from '../../client/components/static-generation-bailout'
 import { DynamicServerError } from '../../client/components/hooks-server-context'
 import { InvariantError } from '../../shared/lib/invariant-error'
-import { delayUntilRuntimeStage } from '../dynamic-rendering-utils'
 import { ReflectAdapter } from '../web/spec-extension/adapters/reflect'
-import { applyOwnerStack } from '../dynamic-rendering-utils'
+import { applyOwnerStack, getRuntimeStage } from '../dynamic-rendering-utils'
 
 export function draftMode(): Promise<DraftMode> {
   const callingExpression = 'draftMode'
@@ -33,12 +32,19 @@ export function draftMode(): Promise<DraftMode> {
   }
 
   switch (workUnitStore.type) {
-    case 'prerender-runtime':
+    case 'prerender-runtime': {
       // TODO(runtime-ppr): does it make sense to delay this? normally it's always microtasky
-      return delayUntilRuntimeStage(
-        workUnitStore,
-        createOrGetCachedDraftMode(workUnitStore.draftMode, workStore)
-      )
+      const { stagedRendering } = workUnitStore
+      if (stagedRendering) {
+        return stagedRendering.delayUntilStage(
+          getRuntimeStage(stagedRendering),
+          'draftMode',
+          new DraftMode(workUnitStore.draftMode)
+        )
+      } else {
+        return createOrGetCachedDraftMode(workUnitStore.draftMode, workStore)
+      }
+    }
     case 'request':
       return createOrGetCachedDraftMode(workUnitStore.draftMode, workStore)
 

packages/next/src/server/request/headers.ts

@@ -20,9 +20,9 @@ import {
 } from '../app-render/dynamic-rendering'
 import { StaticGenBailoutError } from '../../client/components/static-generation-bailout'
 import {
-  delayUntilRuntimeStage,
   makeDevtoolsIOAwarePromise,
   makeHangingPromise,
+  getRuntimeStage,
 } from '../dynamic-rendering-utils'
 import { createDedupedByCallsiteServerErrorLoggerDev } from '../create-deduped-by-callsite-server-error-logger'
 import { isRequestAPICallableInsideAfter } from './utils'
@@ -131,11 +131,19 @@ export function headers(): Promise<ReadonlyHeaders> {
             workStore,
             workUnitStore
           )
-        case 'prerender-runtime':
-          return delayUntilRuntimeStage(
-            workUnitStore,
-            makeUntrackedHeaders(workUnitStore.headers)
-          )
+        case 'prerender-runtime': {
+          const { stagedRendering } = workUnitStore
+          if (stagedRendering) {
+            // TODO(app-shells): headers should be dynamic instead.
+            return stagedRendering.delayUntilStage(
+              getRuntimeStage(stagedRendering),
+              'headers',
+              workUnitStore.headers
+            )
+          } else {
+            return makeUntrackedHeaders(workUnitStore.headers)
+          }
+        }
         case 'private-cache':
           // Private caches are delayed until the runtime stage in use-cache-wrapper,
           // so we don't need an additional delay here.
@@ -226,7 +234,7 @@ function makeUntrackedHeadersWithDevWarnings(
   const promise = makeDevtoolsIOAwarePromise(
     underlyingHeaders,
     requestStore,
-    RenderStage.Runtime
+    RenderStage.Runtime // TODO(app-shells): headers should be dynamic instead.
   )
 
   const proxiedPromise = instrumentHeadersPromiseWithDevWarnings(promise, route)