[Cache Components] defer cookies(), headers(), and draftMode() to a new task in dev

When resolving these Request Data APIs in dev we now do so in a new Task. This allows the environment tracking to properly move stuff that was previously in the Prerender environment to the Server environment when a Request Data API is used before them. This also has the nice side effect of causing these functions to appear as IO in the new suspended-by features of React Devtools.

In Prod we don't need to do any additional delaying because the all of the features this is designed to interact with are dev only

Author: gnoff

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

@@ -1765,6 +1765,8 @@ async function renderToHTMLOrFlightImpl(
       renderOpts.renderResumeDataCache ?? postponedState?.renderResumeDataCache
 
     const rootParams = getRootParams(loaderTree, ctx.getDynamicParamFromSegment)
+    const devValidatingFallbackParams =
+      getRequestMeta(req, 'devValidatingFallbackParams') || null
     const requestStore = createRequestStoreForRender(
       req,
       res,
@@ -1775,7 +1777,8 @@ async function renderToHTMLOrFlightImpl(
       renderOpts.previewProps,
       isHmrRefresh,
       serverComponentsHmrCache,
-      renderResumeDataCache
+      renderResumeDataCache,
+      devValidatingFallbackParams
     )
 
     if (
@@ -1847,7 +1850,8 @@ async function renderToHTMLOrFlightImpl(
             notFoundLoaderTree,
             formState,
             postponedState,
-            metadata
+            metadata,
+            devValidatingFallbackParams
           )
 
           return new RenderResult(stream, {
@@ -1878,7 +1882,8 @@ async function renderToHTMLOrFlightImpl(
       loaderTree,
       formState,
       postponedState,
-      metadata
+      metadata,
+      devValidatingFallbackParams
     )
 
     // Invalid dynamic usages should only error the request in development.
@@ -2064,7 +2069,8 @@ async function renderToStream(
   tree: LoaderTree,
   formState: any,
   postponedState: PostponedState | null,
-  metadata: AppPageRenderResultMetadata
+  metadata: AppPageRenderResultMetadata,
+  devValidatingFallbackParams: FallbackRouteParams | null
 ): Promise<ReadableStream<Uint8Array>> {
   const { assetPrefix, nonce, pagePath, renderOpts } = ctx
 
@@ -2212,9 +2218,6 @@ async function renderToStream(
         }
       )
 
-      const devValidatingFallbackParams =
-        getRequestMeta(req, 'devValidatingFallbackParams') || null
-
       spawnDynamicValidationInDev(
         resolveValidation,
         tree,

packages/next/src/server/app-render/work-unit-async-storage.external.ts

@@ -18,6 +18,7 @@ import type { Params } from '../request/params'
 import type { ImplicitTags } from '../lib/implicit-tags'
 import type { WorkStore } from './work-async-storage.external'
 import { NEXT_HMR_REFRESH_HASH_COOKIE } from '../../client/components/app-router-headers'
+import { InvariantError } from '../../shared/lib/invariant-error'
 
 export type WorkUnitPhase = 'action' | 'render' | 'after'
 
@@ -67,6 +68,7 @@ export interface RequestStore extends CommonWorkUnitStore {
   // DEV-only
   usedDynamic?: boolean
   prerenderPhase?: boolean
+  devFallbackParams?: FallbackRouteParams | null
 }
 
 /**
@@ -311,6 +313,10 @@ export function throwForMissingRequestStore(callingExpression: string): never {
   )
 }
 
+export function throwInvariantForMissingStore(): never {
+  throw new InvariantError('Expected workUnitAsyncStorage to have a store.')
+}
+
 export function getPrerenderResumeDataCache(
   workUnitStore: WorkUnitStore
 ): PrerenderResumeDataCache | null {

packages/next/src/server/async-storage/request-store.ts

@@ -24,6 +24,7 @@ import type { ServerComponentsHmrCache } from '../response-cache'
 import type { RenderResumeDataCache } from '../resume-data-cache/resume-data-cache'
 import type { Params } from '../request/params'
 import type { ImplicitTags } from '../lib/implicit-tags'
+import type { FallbackRouteParams } from '../request/fallback-params'
 
 function getHeaders(headers: Headers | IncomingHttpHeaders): ReadonlyHeaders {
   const cleaned = HeadersAdapter.from(headers)
@@ -114,7 +115,8 @@ export function createRequestStoreForRender(
   previewProps: WrapperRenderOpts['previewProps'],
   isHmrRefresh: RequestContext['isHmrRefresh'],
   serverComponentsHmrCache: RequestContext['serverComponentsHmrCache'],
-  renderResumeDataCache: RenderResumeDataCache | undefined
+  renderResumeDataCache: RenderResumeDataCache | undefined,
+  devFallbackParams: FallbackRouteParams | null
 ): RequestStore {
   return createRequestStoreImpl(
     // Pages start in render phase by default
@@ -128,7 +130,8 @@ export function createRequestStoreForRender(
     renderResumeDataCache,
     previewProps,
     isHmrRefresh,
-    serverComponentsHmrCache
+    serverComponentsHmrCache,
+    devFallbackParams
   )
 }
 
@@ -151,7 +154,8 @@ export function createRequestStoreForAPI(
     undefined,
     previewProps,
     false,
-    undefined
+    undefined,
+    null
   )
 }
 
@@ -166,7 +170,8 @@ function createRequestStoreImpl(
   renderResumeDataCache: RenderResumeDataCache | undefined,
   previewProps: WrapperRenderOpts['previewProps'],
   isHmrRefresh: RequestContext['isHmrRefresh'],
-  serverComponentsHmrCache: RequestContext['serverComponentsHmrCache']
+  serverComponentsHmrCache: RequestContext['serverComponentsHmrCache'],
+  devFallbackParams: FallbackRouteParams | null | undefined
 ): RequestStore {
   function defaultOnUpdateCookies(cookies: string[]) {
     if (res) {
@@ -258,6 +263,7 @@ function createRequestStoreImpl(
     serverComponentsHmrCache:
       serverComponentsHmrCache ||
       (globalThis as any).__serverComponentsHmrCache,
+    devFallbackParams,
   }
 }
 

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

@@ -72,3 +72,14 @@ export function makeHangingPromise<T>(
 }
 
 function ignoreReject() {}
+
+export function makeDevtoolsIOAwarePromise<T>(underlying: T): Promise<T> {
+  // in React DevTools if we resolve in a setTimeout we will observe
+  // the promise resolution as something that can suspend a boundary or root.
+  return new Promise<T>((resolve) => {
+    // Must use setTimeout to be considered IO React DevTools. setImmediate will not work.
+    setTimeout(() => {
+      resolve(underlying)
+    }, 0)
+  })
+}

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

@@ -1,12 +1,18 @@
 import { workAsyncStorage } from '../app-render/work-async-storage.external'
-import { workUnitAsyncStorage } from '../app-render/work-unit-async-storage.external'
+import {
+  throwForMissingRequestStore,
+  workUnitAsyncStorage,
+} from '../app-render/work-unit-async-storage.external'
 import {
   postponeWithTracking,
   throwToInterruptStaticGeneration,
   trackDynamicDataInDynamicRender,
 } from '../app-render/dynamic-rendering'
 import { StaticGenBailoutError } from '../../client/components/static-generation-bailout'
-import { makeHangingPromise } from '../dynamic-rendering-utils'
+import {
+  makeHangingPromise,
+  makeDevtoolsIOAwarePromise,
+} from '../dynamic-rendering-utils'
 import { isRequestAPICallableInsideAfter } from './utils'
 
 /**
@@ -15,6 +21,7 @@ import { isRequestAPICallableInsideAfter } from './utils'
  * During prerendering it will never resolve and during rendering it resolves immediately.
  */
 export function connection(): Promise<void> {
+  const callingExpression = 'connection'
   const workStore = workAsyncStorage.getStore()
   const workUnitStore = workUnitAsyncStorage.getStore()
 
@@ -94,12 +101,20 @@ export function connection(): Promise<void> {
           )
         case 'request':
           trackDynamicDataInDynamicRender(workUnitStore)
-          break
+          if (process.env.NODE_ENV === 'development') {
+            // Semantically we only need the dev tracking when running in `next dev`
+            // but since you would never use next dev with production NODE_ENV we use this
+            // as a proxy so we can statically exclude this code from production builds.
+            return makeDevtoolsIOAwarePromise(undefined)
+          } else {
+            return Promise.resolve(undefined)
+          }
         default:
           workUnitStore satisfies never
       }
     }
   }
 
-  return Promise.resolve(undefined)
+  // If we end up here, there was no work store or work unit store present.
+  throwForMissingRequestStore(callingExpression)
 }

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

@@ -21,9 +21,11 @@ import {
   trackSynchronousRequestDataAccessInDev,
 } from '../app-render/dynamic-rendering'
 import { StaticGenBailoutError } from '../../client/components/static-generation-bailout'
-import { makeHangingPromise } from '../dynamic-rendering-utils'
+import {
+  makeDevtoolsIOAwarePromise,
+  makeHangingPromise,
+} from '../dynamic-rendering-utils'
 import { createDedupedByCallsiteServerErrorLoggerDev } from '../create-deduped-by-callsite-server-error-logger'
-import { scheduleImmediate } from '../../lib/scheduler'
 import { isRequestAPICallableInsideAfter } from './utils'
 import { InvariantError } from '../../shared/lib/invariant-error'
 import { ReflectAdapter } from '../web/spec-extension/adapters/reflect'
@@ -134,10 +136,10 @@ export function cookies(): Promise<ReadonlyRequestCookies> {
             underlyingCookies = workUnitStore.cookies
           }
 
-          if (
-            process.env.NODE_ENV === 'development' &&
-            !workStore?.isPrefetchRequest
-          ) {
+          if (process.env.NODE_ENV === 'development') {
+            // Semantically we only need the dev tracking when running in `next dev`
+            // but since you would never use next dev with production NODE_ENV we use this
+            // as a proxy so we can statically exclude this code from production builds.
             if (process.env.__NEXT_CACHE_COMPONENTS) {
               return makeUntrackedCookiesWithDevWarnings(
                 underlyingCookies,
@@ -263,9 +265,7 @@ function makeUntrackedExoticCookiesWithDevWarnings(
     return cachedCookies
   }
 
-  const promise = new Promise<ReadonlyRequestCookies>((resolve) =>
-    scheduleImmediate(() => resolve(underlyingCookies))
-  )
+  const promise = makeDevtoolsIOAwarePromise(underlyingCookies)
   CachedCookies.set(underlyingCookies, promise)
 
   Object.defineProperties(promise, {
@@ -416,9 +416,7 @@ function makeUntrackedCookiesWithDevWarnings(
     return cachedCookies
   }
 
-  const promise = new Promise<ReadonlyRequestCookies>((resolve) =>
-    scheduleImmediate(() => resolve(underlyingCookies))
-  )
+  const promise = makeDevtoolsIOAwarePromise(underlyingCookies)
 
   const proxiedPromise = new Proxy(promise, {
     get(target, prop, receiver) {

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

@@ -6,8 +6,8 @@ import {
   workAsyncStorage,
   type WorkStore,
 } from '../app-render/work-async-storage.external'
-import { throwForMissingRequestStore } from '../app-render/work-unit-async-storage.external'
 import {
+  throwForMissingRequestStore,
   workUnitAsyncStorage,
   type PrerenderStoreModern,
 } from '../app-render/work-unit-async-storage.external'
@@ -18,9 +18,11 @@ import {
   trackSynchronousRequestDataAccessInDev,
 } from '../app-render/dynamic-rendering'
 import { StaticGenBailoutError } from '../../client/components/static-generation-bailout'
-import { makeHangingPromise } from '../dynamic-rendering-utils'
+import {
+  makeDevtoolsIOAwarePromise,
+  makeHangingPromise,
+} from '../dynamic-rendering-utils'
 import { createDedupedByCallsiteServerErrorLoggerDev } from '../create-deduped-by-callsite-server-error-logger'
-import { scheduleImmediate } from '../../lib/scheduler'
 import { isRequestAPICallableInsideAfter } from './utils'
 import { InvariantError } from '../../shared/lib/invariant-error'
 import { ReflectAdapter } from '../web/spec-extension/adapters/reflect'
@@ -153,10 +155,10 @@ export function headers(): Promise<ReadonlyHeaders> {
         case 'request':
           trackDynamicDataInDynamicRender(workUnitStore)
 
-          if (
-            process.env.NODE_ENV === 'development' &&
-            !workStore?.isPrefetchRequest
-          ) {
+          if (process.env.NODE_ENV === 'development') {
+            // Semantically we only need the dev tracking when running in `next dev`
+            // but since you would never use next dev with production NODE_ENV we use this
+            // as a proxy so we can statically exclude this code from production builds.
             if (process.env.__NEXT_CACHE_COMPONENTS) {
               return makeUntrackedHeadersWithDevWarnings(
                 workUnitStore.headers,
@@ -263,9 +265,7 @@ function makeUntrackedExoticHeadersWithDevWarnings(
     return cachedHeaders
   }
 
-  const promise = new Promise<ReadonlyHeaders>((resolve) =>
-    scheduleImmediate(() => resolve(underlyingHeaders))
-  )
+  const promise = makeDevtoolsIOAwarePromise(underlyingHeaders)
 
   CachedHeaders.set(underlyingHeaders, promise)
 
@@ -384,9 +384,7 @@ function makeUntrackedHeadersWithDevWarnings(
     return cachedHeaders
   }
 
-  const promise = new Promise<ReadonlyHeaders>((resolve) =>
-    scheduleImmediate(() => resolve(underlyingHeaders))
-  )
+  const promise = makeDevtoolsIOAwarePromise(underlyingHeaders)
 
   const proxiedPromise = new Proxy(promise, {
     get(target, prop, receiver) {