[Cache Components] correctly label IO promises in devtools (#84928)

Updates to staged rendering (cacheComponents dev) to support "suspended
by" in Suspense Devtools

- promises for `cookies()` and other user-callable APIs are now created
before we start the render. Each call to `cookies()` will return the
same promise (although currently it's wrapped in a fresh proxy). The
promise is created via `new Promise`, triggered by a timeout (for the
relevant stage) and has `displayName`. This marks it as an IO operation.
By re-using the same promise, we make sure that all callsites that await
it are considered to be suspended by the same IO operation.
- promises for `params` and `searchParams` use a similar trick. All
instances of `params` and will be derived from one shared IO promise
`sharedParamsParent`. When determining "suspended by", React will walk
up the promise chain, find this parent, and use its `displayName`, which
we set to `"params"`. `searchParams` work analogously

---------

Co-authored-by: Josh Story <story@hey.com>

Author: lubieowoce

packages/next/errors.json

@@ -887,5 +887,6 @@
   "886": "`cacheTag()` is only available with the `cacheComponents` config.",
   "887": "`cacheLife()` is only available with the `cacheComponents` config.",
   "888": "Unknown \\`cacheLife()\\` profile \"%s\" is not configured in next.config.js\\nmodule.exports = {\n  cacheLife: {\n    \"%s\": ...\\n  }\n}",
-  "889": "Unknown \\`cacheLife()\\` profile \"%s\" is not configured in next.config.js\\nmodule.exports = {\n  cacheLife: {\n      \"%s\": ...\\n    }\n}"
+  "889": "Unknown \\`cacheLife()\\` profile \"%s\" is not configured in next.config.js\\nmodule.exports = {\n  cacheLife: {\n      \"%s\": ...\\n    }\n}",
+  "890": "Received an underlying cookies object that does not match either `cookies` or `mutableCookies`"
 }

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

@@ -2821,6 +2821,12 @@ async function renderWithRestartOnCacheMissInDev(
   // so not having a resume data cache won't break any expectations in case we don't need to restart.
   requestStore.renderResumeDataCache = null
   requestStore.stagedRendering = initialStageController
+  requestStore.asyncApiPromises = createAsyncApiPromisesInDev(
+    initialStageController,
+    requestStore.cookies,
+    requestStore.mutableCookies,
+    requestStore.headers
+  )
   requestStore.cacheSignal = cacheSignal
 
   let debugChannel = setReactDebugChannel && createDebugChannel()
@@ -2921,6 +2927,12 @@ async function renderWithRestartOnCacheMissInDev(
   )
   requestStore.stagedRendering = finalStageController
   requestStore.cacheSignal = null
+  requestStore.asyncApiPromises = createAsyncApiPromisesInDev(
+    finalStageController,
+    requestStore.cookies,
+    requestStore.mutableCookies,
+    requestStore.headers
+  )
 
   // The initial render already wrote to its debug channel.
   // We're not using it, so we need to create a new one.
@@ -2962,6 +2974,48 @@ async function renderWithRestartOnCacheMissInDev(
   }
 }
 
+function createAsyncApiPromisesInDev(
+  stagedRendering: StagedRenderingController,
+  cookies: RequestStore['cookies'],
+  mutableCookies: RequestStore['mutableCookies'],
+  headers: RequestStore['headers']
+): NonNullable<RequestStore['asyncApiPromises']> {
+  return {
+    // Runtime APIs
+    cookies: stagedRendering.delayUntilStage(
+      RenderStage.Runtime,
+      'cookies',
+      cookies
+    ),
+    mutableCookies: stagedRendering.delayUntilStage(
+      RenderStage.Runtime,
+      'cookies',
+      mutableCookies as RequestStore['cookies']
+    ),
+    headers: stagedRendering.delayUntilStage(
+      RenderStage.Runtime,
+      'headers',
+      headers
+    ),
+    // These are not used directly, but we chain other `params`/`searchParams` promises off of them.
+    sharedParamsParent: stagedRendering.delayUntilStage(
+      RenderStage.Runtime,
+      'params',
+      '<internal params>'
+    ),
+    sharedSearchParamsParent: stagedRendering.delayUntilStage(
+      RenderStage.Runtime,
+      'searchParams',
+      '<internal searchParams>'
+    ),
+    connection: stagedRendering.delayUntilStage(
+      RenderStage.Dynamic,
+      'connection',
+      undefined
+    ),
+  }
+}
+
 type DebugChannelPair = {
   serverSide: DebugChannelServer
   clientSide: DebugChannelClient

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

@@ -52,30 +52,37 @@ export class StagedRenderingController {
     }
   }
 
-  delayUntilStage<T>(stage: NonStaticRenderStage, resolvedValue: T) {
-    let stagePromise: Promise<void>
+  private getStagePromise(stage: NonStaticRenderStage): Promise<void> {
     switch (stage) {
       case RenderStage.Runtime: {
-        stagePromise = this.runtimeStagePromise.promise
-        break
+        return this.runtimeStagePromise.promise
       }
       case RenderStage.Dynamic: {
-        stagePromise = this.dynamicStagePromise.promise
-        break
+        return this.dynamicStagePromise.promise
       }
       default: {
         stage satisfies never
         throw new InvariantError(`Invalid render stage: ${stage}`)
       }
     }
+  }
+
+  waitForStage(stage: NonStaticRenderStage) {
+    return this.getStagePromise(stage)
+  }
+
+  delayUntilStage<T>(
+    stage: NonStaticRenderStage,
+    displayName: string | undefined,
+    resolvedValue: T
+  ) {
+    const ioTriggerPromise = this.getStagePromise(stage)
 
-    // FIXME: this seems to be the only form that leads to correct API names
-    // being displayed in React Devtools (in the "suspended by" section).
-    // If we use `promise.then(() => resolvedValue)`, the names are lost.
-    // It's a bit strange that only one of those works right.
-    const promise = new Promise<T>((resolve, reject) => {
-      stagePromise.then(resolve.bind(null, resolvedValue), reject)
-    })
+    const promise = makeDevtoolsIOPromiseFromIOTrigger(
+      ioTriggerPromise,
+      displayName,
+      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).
@@ -88,3 +95,26 @@ export class StagedRenderingController {
 }
 
 function ignoreReject() {}
+
+// TODO(restart-on-cache-miss): the layering of `delayUntilStage`,
+// `makeDevtoolsIOPromiseFromIOTrigger` and and `makeDevtoolsIOAwarePromise`
+// is confusing, we should clean it up.
+function makeDevtoolsIOPromiseFromIOTrigger<T>(
+  ioTrigger: Promise<any>,
+  displayName: string | undefined,
+  resolvedValue: T
+): Promise<T> {
+  // If we create a `new Promise` and give it a displayName
+  // (with no userspace code above us in the stack)
+  // React Devtools will use it as the IO cause when determining "suspended by".
+  // In particular, it should shadow any inner IO that resolved/rejected the promise
+  // (in case of staged rendering, this will be the `setTimeout` that triggers the relevant stage)
+  const promise = new Promise<T>((resolve, reject) => {
+    ioTrigger.then(resolve.bind(null, resolvedValue), reject)
+  })
+  if (displayName !== undefined) {
+    // @ts-expect-error
+    promise.displayName = displayName
+  }
+  return promise
+}

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

@@ -70,10 +70,22 @@ export interface RequestStore extends CommonWorkUnitStore {
   usedDynamic?: boolean
   devFallbackParams?: OpaqueFallbackRouteParams | null
   stagedRendering?: StagedRenderingController | null
+  asyncApiPromises?: DevAsyncApiPromises
   cacheSignal?: CacheSignal | null
   prerenderResumeDataCache?: PrerenderResumeDataCache | null
 }
 
+type DevAsyncApiPromises = {
+  cookies: Promise<ReadonlyRequestCookies>
+  mutableCookies: Promise<ReadonlyRequestCookies>
+  headers: Promise<ReadonlyHeaders>
+
+  sharedParamsParent: Promise<string>
+  sharedSearchParamsParent: Promise<string>
+
+  connection: Promise<undefined>
+}
+
 /**
  * The Prerender store is for tracking information related to prerenders.
  *

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

@@ -83,7 +83,11 @@ export function makeDevtoolsIOAwarePromise<T>(
 ): Promise<T> {
   if (requestStore.stagedRendering) {
     // We resolve each stage in a timeout, so React DevTools will pick this up as IO.
-    return requestStore.stagedRendering.delayUntilStage(stage, underlying)
+    return requestStore.stagedRendering.delayUntilStage(
+      stage,
+      undefined,
+      underlying
+    )
   }
   // in React DevTools if we resolve in a setTimeout we will observe
   // the promise resolution as something that can suspend a boundary or root.

packages/next/src/server/lib/patch-fetch.ts

@@ -562,9 +562,8 @@ export function createPatchedFetcher(
                   cacheSignal.endRead()
                   cacheSignal = null
                 }
-                await workUnitStore.stagedRendering.delayUntilStage(
-                  RenderStage.Dynamic,
-                  undefined
+                await workUnitStore.stagedRendering.waitForStage(
+                  RenderStage.Dynamic
                 )
               }
               break
@@ -689,9 +688,8 @@ export function createPatchedFetcher(
                       cacheSignal.endRead()
                       cacheSignal = null
                     }
-                    await workUnitStore.stagedRendering.delayUntilStage(
-                      RenderStage.Dynamic,
-                      undefined
+                    await workUnitStore.stagedRendering.waitForStage(
+                      RenderStage.Dynamic
                     )
                   }
                   break
@@ -962,9 +960,8 @@ export function createPatchedFetcher(
                     process.env.NODE_ENV === 'development' &&
                     workUnitStore.stagedRendering
                   ) {
-                    await workUnitStore.stagedRendering.delayUntilStage(
-                      RenderStage.Dynamic,
-                      undefined
+                    await workUnitStore.stagedRendering.waitForStage(
+                      RenderStage.Dynamic
                     )
                   }
                   break
@@ -1091,9 +1088,8 @@ export function createPatchedFetcher(
                       cacheSignal.endRead()
                       cacheSignal = null
                     }
-                    await workUnitStore.stagedRendering.delayUntilStage(
-                      RenderStage.Dynamic,
-                      undefined
+                    await workUnitStore.stagedRendering.waitForStage(
+                      RenderStage.Dynamic
                     )
                   }
                   break
@@ -1138,9 +1134,8 @@ export function createPatchedFetcher(
                       process.env.NODE_ENV === 'development' &&
                       workUnitStore.stagedRendering
                     ) {
-                      await workUnitStore.stagedRendering.delayUntilStage(
-                        RenderStage.Dynamic,
-                        undefined
+                      await workUnitStore.stagedRendering.waitForStage(
+                        RenderStage.Dynamic
                       )
                     }
                     break

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

@@ -106,6 +106,9 @@ export function connection(): Promise<void> {
             // 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 (workUnitStore.asyncApiPromises) {
+              return workUnitStore.asyncApiPromises.connection
+            }
             return makeDevtoolsIOAwarePromise(
               undefined,
               workUnitStore,