add phase tracking to workUnitStore

packages/next/src/server/after/after-context.test.ts

@@ -2,16 +2,19 @@ import { DetachedPromise } from '../../lib/detached-promise'
 import { AsyncLocalStorage } from 'async_hooks'
 
 import type { WorkStore } from '../app-render/work-async-storage.external'
+import type { WorkUnitStore } from '../app-render/work-unit-async-storage.external'
 import type { AfterContext } from './after-context'
 
 describe('AfterContext', () => {
   // 'async-local-storage.ts' needs `AsyncLocalStorage` on `globalThis` at import time,
   // so we have to do some contortions here to set it up before running anything else
   type WASMod = typeof import('../app-render/work-async-storage.external')
+  type WSMod = typeof import('../app-render/work-unit-async-storage.external')
   type AfterMod = typeof import('./after')
   type AfterContextMod = typeof import('./after-context')
 
   let workAsyncStorage: WASMod['workAsyncStorage']
+  let workUnitAsyncStorage: WSMod['workUnitAsyncStorage']
   let AfterContext: AfterContextMod['AfterContext']
   let after: AfterMod['unstable_after']
 
@@ -22,6 +25,9 @@ describe('AfterContext', () => {
     const WASMod = await import('../app-render/work-async-storage.external')
     workAsyncStorage = WASMod.workAsyncStorage
 
+    const WSMod = await import('../app-render/work-unit-async-storage.external')
+    workUnitAsyncStorage = WSMod.workUnitAsyncStorage
+
     const AfterContextMod = await import('./after-context')
     AfterContext = AfterContextMod.AfterContext
 
@@ -32,7 +38,9 @@ describe('AfterContext', () => {
   const createRun =
     (_afterContext: AfterContext, workStore: WorkStore) =>
     <T>(cb: () => T): T => {
-      return workAsyncStorage.run(workStore, cb)
+      return workAsyncStorage.run(workStore, () =>
+        workUnitAsyncStorage.run(createMockWorkUnitStore(), cb)
+      )
     }
 
   it('runs after() callbacks from a run() callback that resolves', async () => {
@@ -362,11 +370,13 @@ describe('AfterContext', () => {
     const promise3 = new DetachedPromise<string>()
     const afterCallback3 = jest.fn(() => promise3.promise)
 
-    workAsyncStorage.run(workStore, () => {
-      after(afterCallback1)
-      after(afterCallback2)
-      after(afterCallback3)
-    })
+    workAsyncStorage.run(workStore, () =>
+      workUnitAsyncStorage.run(createMockWorkUnitStore(), () => {
+        after(afterCallback1)
+        after(afterCallback2)
+        after(afterCallback3)
+      })
+    )
 
     expect(afterCallback1).not.toHaveBeenCalled()
     expect(afterCallback2).not.toHaveBeenCalled()
@@ -557,3 +567,7 @@ const createMockWorkStore = (afterContext: AfterContext): WorkStore => {
     },
   })
 }
+
+const createMockWorkUnitStore = () => {
+  return { phase: 'render' } as WorkUnitStore
+}

packages/next/src/server/after/after-context.ts

@@ -6,6 +6,10 @@ import { isThenable } from '../../shared/lib/is-thenable'
 import { workAsyncStorage } from '../app-render/work-async-storage.external'
 import { withExecuteRevalidates } from './revalidation-utils'
 import { bindSnapshot } from '../app-render/async-local-storage'
+import {
+  workUnitAsyncStorage,
+  type WorkUnitStore,
+} from '../app-render/work-unit-async-storage.external'
 
 export type AfterContextOpts = {
   waitUntil: RequestLifecycleOpts['waitUntil'] | undefined
@@ -18,6 +22,7 @@ export class AfterContext {
 
   private runCallbacksOnClosePromise: Promise<void> | undefined
   private callbackQueue: PromiseQueue
+  private workUnitStores = new Set<WorkUnitStore>()
 
   constructor({ waitUntil, onClose }: AfterContextOpts) {
     this.waitUntil = waitUntil
@@ -55,6 +60,14 @@ export class AfterContext {
       )
     }
 
+    const workUnitStore = workUnitAsyncStorage.getStore()
+    if (!workUnitStore) {
+      throw new InvariantError(
+        'Missing workUnitStore in AfterContext.addCallback'
+      )
+    }
+    this.workUnitStores.add(workUnitStore)
+
     // this should only happen once.
     if (!this.runCallbacksOnClosePromise) {
       this.runCallbacksOnClosePromise = this.runCallbacksOnClose()
@@ -89,9 +102,15 @@ export class AfterContext {
   private async runCallbacks(): Promise<void> {
     if (this.callbackQueue.size === 0) return
 
+    for (const workUnitStore of this.workUnitStores) {
+      workUnitStore.phase = 'after'
+    }
+
     const workStore = workAsyncStorage.getStore()
+    if (!workStore) {
+      throw new InvariantError('Missing workStore in AfterContext.runCallbacks')
+    }
 
-    // TODO(after): Change phase in workUnitStore to disable e.g. `cookies().set()`
     return withExecuteRevalidates(workStore, () => {
       this.callbackQueue.start()
       return this.callbackQueue.onIdle()

packages/next/src/server/after/after.ts

@@ -14,26 +14,29 @@ export function unstable_after<T>(task: AfterTask<T>): void {
   const workStore = workAsyncStorage.getStore()
   const workUnitStore = workUnitAsyncStorage.getStore()
 
-  if (workStore) {
-    const { afterContext } = workStore
-    if (!afterContext) {
-      throw new Error(
-        '`unstable_after()` must be explicitly enabled by setting `experimental.after: true` in your next.config.js.'
-      )
-    }
+  if (!workStore) {
+    // TODO(after): the linked docs page talks about *dynamic* APIs, which unstable_after soon won't be anymore
+    throw new Error(
+      '`unstable_after` was called outside a request scope. Read more: https://nextjs.org/docs/messages/next-dynamic-api-wrong-context'
+    )
+  }
 
-    // TODO: After should not cause dynamic.
-    const callingExpression = 'unstable_after'
-    if (workStore.forceStatic) {
-      throw new StaticGenBailoutError(
-        `Route ${workStore.route} with \`dynamic = "force-static"\` couldn't be rendered statically because it used \`${callingExpression}\`. See more info here: https://nextjs.org/docs/app/building-your-application/rendering/static-and-dynamic#dynamic-rendering`
-      )
-    } else {
-      markCurrentScopeAsDynamic(workStore, workUnitStore, callingExpression)
-    }
+  const { afterContext } = workStore
+  if (!afterContext) {
+    throw new Error(
+      '`unstable_after` must be explicitly enabled by setting `experimental.after: true` in your next.config.js.'
+    )
+  }
 
-    afterContext.after(task)
+  // TODO: After should not cause dynamic.
+  const callingExpression = 'unstable_after'
+  if (workStore.forceStatic) {
+    throw new StaticGenBailoutError(
+      `Route ${workStore.route} with \`dynamic = "force-static"\` couldn't be rendered statically because it used \`${callingExpression}\`. See more info here: https://nextjs.org/docs/app/building-your-application/rendering/static-and-dynamic#dynamic-rendering`
+    )
   } else {
-    // TODO: Error for pages?
+    markCurrentScopeAsDynamic(workStore, workUnitStore, callingExpression)
   }
+
+  afterContext.after(task)
 }

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

@@ -459,6 +459,8 @@ export async function handleAction({
     )
   }
 
+  requestStore.phase = 'action'
+
   // When running actions the default is no-store, you can still `cache: 'force-cache'`
   workStore.fetchCache = 'default-no-store'
 

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

@@ -375,6 +375,7 @@ async function generateDynamicRSCPayload(
     skipFlight: boolean
   }
 ): Promise<RSCPayload> {
+  ctx.requestStore.phase = 'render'
   // Flight data that is going to be passed to the browser.
   // Currently a single item array but in the future multiple patches might be combined in a single request.
 
@@ -1283,6 +1284,7 @@ export const renderToHTMLOrFlight: AppPageRender = (
       req,
       url,
       res,
+      phase: 'render',
       renderOpts,
       isHmrRefresh,
       serverComponentsHmrCache,
@@ -1722,6 +1724,8 @@ async function prerenderToStream(
   workStore: WorkStore,
   tree: LoaderTree
 ): Promise<PrerenderToStreamResult> {
+  ctx.requestStore.phase = 'render'
+
   // When prerendering formState is always null. We still include it
   // because some shared APIs expect a formState value and this is slightly
   // more explicit than making it an optional function argument
@@ -1843,6 +1847,7 @@ async function prerenderToStream(
         const prospectiveRenderPrerenderStore: PrerenderStore =
           (prerenderStore = {
             type: 'prerender',
+            phase: 'render',
             implicitTags: ctx.requestStore.implicitTags,
             renderSignal: flightController.signal,
             cacheSignal,
@@ -1936,6 +1941,7 @@ async function prerenderToStream(
 
         const finalRenderPrerenderStore: PrerenderStore = (prerenderStore = {
           type: 'prerender',
+          phase: 'render',
           implicitTags: ctx.requestStore.implicitTags,
           renderSignal: flightController.signal,
           // During the final prerender we don't need to track cache access so we omit the signal
@@ -1995,6 +2001,7 @@ async function prerenderToStream(
         const SSRController = new AbortController()
         const ssrPrerenderStore: PrerenderStore = {
           type: 'prerender',
+          phase: 'render',
           implicitTags: ctx.requestStore.implicitTags,
           renderSignal: SSRController.signal,
           // For HTML Generation we don't need to track cache reads (RSC only)
@@ -2210,6 +2217,7 @@ async function prerenderToStream(
         const prospectiveRenderPrerenderStore: PrerenderStore =
           (prerenderStore = {
             type: 'prerender',
+            phase: 'render',
             implicitTags: ctx.requestStore.implicitTags,
             renderSignal: flightController.signal,
             cacheSignal,
@@ -2292,6 +2300,7 @@ async function prerenderToStream(
 
         const finalRenderPrerenderStore: PrerenderStore = (prerenderStore = {
           type: 'prerender',
+          phase: 'render',
           implicitTags: ctx.requestStore.implicitTags,
           renderSignal: flightController.signal,
           // During the final prerender we don't need to track cache access so we omit the signal
@@ -2305,6 +2314,7 @@ async function prerenderToStream(
         const SSRController = new AbortController()
         const ssrPrerenderStore: PrerenderStore = {
           type: 'prerender',
+          phase: 'render',
           implicitTags: ctx.requestStore.implicitTags,
           renderSignal: SSRController.signal,
           // For HTML Generation we don't need to track cache reads (RSC only)
@@ -2492,6 +2502,7 @@ async function prerenderToStream(
       )
       const reactServerPrerenderStore: PrerenderStore = (prerenderStore = {
         type: 'prerender-ppr',
+        phase: 'render',
         implicitTags: ctx.requestStore.implicitTags,
         dynamicTracking,
         revalidate: INFINITE_CACHE,
@@ -2520,6 +2531,7 @@ async function prerenderToStream(
 
       const ssrPrerenderStore: PrerenderStore = {
         type: 'prerender-ppr',
+        phase: 'render',
         implicitTags: ctx.requestStore.implicitTags,
         dynamicTracking,
         revalidate: INFINITE_CACHE,
@@ -2686,6 +2698,7 @@ async function prerenderToStream(
     } else {
       const prerenderLegacyStore: PrerenderStore = (prerenderStore = {
         type: 'prerender-legacy',
+        phase: 'render',
         implicitTags: ctx.requestStore.implicitTags,
         revalidate: INFINITE_CACHE,
         tags: [...ctx.requestStore.implicitTags],
@@ -2842,6 +2855,7 @@ async function prerenderToStream(
 
     const prerenderLegacyStore: PrerenderStore = (prerenderStore = {
       type: 'prerender-legacy',
+      phase: 'render',
       implicitTags: ctx.requestStore.implicitTags,
       revalidate: INFINITE_CACHE,
       tags: [...ctx.requestStore.implicitTags],

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

@@ -10,6 +10,13 @@ import type { DynamicTrackingState } from './dynamic-rendering'
 import { workUnitAsyncStorage } from './work-unit-async-storage-instance' with { 'turbopack-transition': 'next-shared' }
 import type { ServerComponentsHmrCache } from '../response-cache'
 
+type WorkUnitPhase = 'action' | 'render' | 'after'
+
+type PhasePartial = {
+  /** NOTE: Will be mutated as phases change */
+  phase: WorkUnitPhase
+}
+
 export type RequestStore = {
   type: 'request'
 
@@ -41,7 +48,7 @@ export type RequestStore = {
 
   // DEV-only
   usedDynamic?: boolean
-}
+} & PhasePartial
 
 /**
  * The Prerender store is for tracking information related to prerenders.
@@ -81,7 +88,7 @@ export type PrerenderStoreModern = {
   // Collected revalidate times and tags for this document during the prerender.
   revalidate: number // in seconds. 0 means dynamic. INFINITE_CACHE and higher means never revalidate.
   tags: null | string[]
-}
+} & PhasePartial
 
 export type PrerenderStorePPR = {
   type: 'prerender-ppr'
@@ -90,15 +97,15 @@ export type PrerenderStorePPR = {
   // Collected revalidate times and tags for this document during the prerender.
   revalidate: number // in seconds. 0 means dynamic. INFINITE_CACHE and higher means never revalidate.
   tags: null | string[]
-}
+} & PhasePartial
 
 export type PrerenderStoreLegacy = {
   type: 'prerender-legacy'
   readonly implicitTags: string[]
   // Collected revalidate times and tags for this document during the prerender.
   revalidate: number // in seconds. 0 means dynamic. INFINITE_CACHE and higher means never revalidate.
   tags: null | string[]
-}
+} & PhasePartial
 
 export type PrerenderStore =
   | PrerenderStoreLegacy
@@ -112,11 +119,11 @@ export type UseCacheStore = {
   revalidate: number // implicit revalidate time from inner caches / fetches
   explicitRevalidate: undefined | number // explicit revalidate time from cacheLife() calls
   tags: null | string[]
-}
+} & PhasePartial
 
 export type UnstableCacheStore = {
   type: 'unstable-cache'
-}
+} & PhasePartial
 
 /**
  * The Cache store is for tracking information inside a "use cache" or unstable_cache context.

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

@@ -64,6 +64,7 @@ export type RequestContext = RequestResponsePair & {
      */
     search?: string
   }
+  phase: RequestStore['phase']
   renderOpts?: WrapperRenderOpts
   isHmrRefresh?: boolean
   serverComponentsHmrCache?: ServerComponentsHmrCache
@@ -111,6 +112,7 @@ export const withRequestStore: WithStore<WorkUnitStore, RequestContext> = <
     req,
     url,
     res,
+    phase,
     renderOpts,
     isHmrRefresh,
     serverComponentsHmrCache,
@@ -133,6 +135,7 @@ export const withRequestStore: WithStore<WorkUnitStore, RequestContext> = <
 
   const store: RequestStore = {
     type: 'request',
+    phase,
     implicitTags: implicitTags ?? [],
     // Rather than just using the whole `url` here, we pull the parts we want
     // to ensure we don't use parts of the URL that we shouldn't. This also