refactor: remove delayUntilRuntimeStage (#94350)

In #93801, we'll be removing `getRuntimeStage` because we want to
differentiate session data and link data, which resolve in different
stages. `delayUntilRuntimeStage` isn't all that useful anyway (we can do
`stagedRendering.delayUntilStage` instead) and gets in the way of that
refactor. This should have no behavior changes otherwise.

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

@@ -320,14 +320,17 @@ export class StagedRenderingController {
     stage: AdvanceableRenderStage,
     displayName: string | undefined,
     resolvedValue: T
-  ) {
-    const ioTriggerPromise = this.getStagePromise(stage)
+  ): Promise<T> {
+    const stagePromise = this.getStagePromise(stage)
 
-    const promise = makeDevtoolsIOPromiseFromIOTrigger(
-      ioTriggerPromise,
-      displayName,
-      resolvedValue
-    )
+    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)