Move AfterContext to WorkStore

packages/next/src/client/components/request-async-storage.external.ts

@@ -6,7 +6,6 @@ import type { ReadonlyRequestCookies } from '../../server/web/spec-extension/ada
 
 // Share the instance module in the next-shared layer
 import { requestAsyncStorage } from './request-async-storage-instance' with { 'turbopack-transition': 'next-shared' }
-import type { AfterContext } from '../../server/after/after-context'
 import type { ServerComponentsHmrCache } from '../../server/response-cache'
 
 import { cacheAsyncStorage } from '../../server/app-render/cache-async-storage.external'
@@ -34,7 +33,6 @@ export interface RequestStore {
   readonly cookies: ReadonlyRequestCookies
   readonly mutableCookies: ResponseCookies
   readonly draftMode: DraftModeProvider
-  readonly afterContext: AfterContext | undefined
   readonly isHmrRefresh?: boolean
   readonly serverComponentsHmrCache?: ServerComponentsHmrCache
 }

packages/next/src/client/components/work-async-storage.external.ts

@@ -6,6 +6,7 @@ import type { Revalidate } from '../../server/lib/revalidate'
 import type { FallbackRouteParams } from '../../server/request/fallback-params'
 import type { DeepReadonly } from '../../shared/lib/deep-readonly'
 import type { AppSegmentConfig } from '../../build/segment-config/app/app-segment-config'
+import type { AfterContext } from '../../server/after/after-context'
 
 // Share the instance module in the next-shared layer
 import { workAsyncStorage } from './work-async-storage-instance' with { 'turbopack-transition': 'next-shared' }
@@ -43,6 +44,7 @@ export interface WorkStore {
   dynamicShouldError?: boolean
   pendingRevalidates?: Record<string, Promise<any>>
   pendingRevalidateWrites?: Array<Promise<void>> // This is like pendingRevalidates but isn't used for deduping.
+  readonly afterContext: AfterContext | undefined
 
   dynamicUsageDescription?: string
   dynamicUsageStack?: string

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

@@ -1,29 +1,29 @@
 import { DetachedPromise } from '../../lib/detached-promise'
 import { AsyncLocalStorage } from 'async_hooks'
 
-import type { RequestStore } from '../../client/components/request-async-storage.external'
+import type { WorkStore } from '../../client/components/work-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 RASMod =
-    typeof import('../../client/components/request-async-storage.external')
+  type WASMod =
+    typeof import('../../client/components/work-async-storage.external')
   type AfterMod = typeof import('./after')
   type AfterContextMod = typeof import('./after-context')
 
-  let requestAsyncStorage: RASMod['requestAsyncStorage']
+  let workAsyncStorage: WASMod['workAsyncStorage']
   let AfterContext: AfterContextMod['AfterContext']
   let after: AfterMod['unstable_after']
 
   beforeAll(async () => {
     // @ts-expect-error
     globalThis.AsyncLocalStorage = AsyncLocalStorage
 
-    const RASMod = await import(
-      '../../client/components/request-async-storage.external'
+    const WASMod = await import(
+      '../../client/components/work-async-storage.external'
     )
-    requestAsyncStorage = RASMod.requestAsyncStorage
+    workAsyncStorage = WASMod.workAsyncStorage
 
     const AfterContextMod = await import('./after-context')
     AfterContext = AfterContextMod.AfterContext
@@ -33,11 +33,9 @@ describe('AfterContext', () => {
   })
 
   const createRun =
-    (afterContext: AfterContext, requestStore: RequestStore) =>
+    (_afterContext: AfterContext, workStore: WorkStore) =>
     <T>(cb: () => T): T => {
-      return afterContext.run(requestStore, () =>
-        requestAsyncStorage.run(requestStore, cb)
-      )
+      return workAsyncStorage.run(workStore, cb)
     }
 
   it('runs after() callbacks from a run() callback that resolves', async () => {
@@ -54,8 +52,8 @@ describe('AfterContext', () => {
       onClose,
     })
 
-    const requestStore = createMockRequestStore(afterContext)
-    const run = createRun(afterContext, requestStore)
+    const workStore = createMockWorkStore(afterContext)
+    const run = createRun(afterContext, workStore)
 
     // ==================================
 
@@ -120,9 +118,9 @@ describe('AfterContext', () => {
       onClose,
     })
 
-    const requestStore = createMockRequestStore(afterContext)
+    const workStore = createMockWorkStore(afterContext)
 
-    const run = createRun(afterContext, requestStore)
+    const run = createRun(afterContext, workStore)
 
     // ==================================
 
@@ -167,9 +165,9 @@ describe('AfterContext', () => {
       onClose,
     })
 
-    const requestStore = createMockRequestStore(afterContext)
+    const workStore = createMockWorkStore(afterContext)
 
-    const run = createRun(afterContext, requestStore)
+    const run = createRun(afterContext, workStore)
 
     // ==================================
 
@@ -257,8 +255,8 @@ describe('AfterContext', () => {
       onClose,
     })
 
-    const requestStore = createMockRequestStore(afterContext)
-    const run = createRun(afterContext, requestStore)
+    const workStore = createMockWorkStore(afterContext)
+    const run = createRun(afterContext, workStore)
 
     // ==================================
 
@@ -316,9 +314,9 @@ describe('AfterContext', () => {
       onClose,
     })
 
-    const requestStore = createMockRequestStore(afterContext)
+    const workStore = createMockWorkStore(afterContext)
 
-    const run = createRun(afterContext, requestStore)
+    const run = createRun(afterContext, workStore)
 
     // ==================================
 
@@ -353,7 +351,7 @@ describe('AfterContext', () => {
       onClose,
     })
 
-    const requestStore = createMockRequestStore(afterContext)
+    const workStore = createMockWorkStore(afterContext)
 
     // ==================================
 
@@ -367,13 +365,11 @@ describe('AfterContext', () => {
     const promise3 = new DetachedPromise<string>()
     const afterCallback3 = jest.fn(() => promise3.promise)
 
-    requestAsyncStorage.run(requestStore, () =>
-      afterContext.run(requestStore, () => {
-        after(afterCallback1)
-        after(afterCallback2)
-        after(afterCallback3)
-      })
-    )
+    workAsyncStorage.run(workStore, () => {
+      after(afterCallback1)
+      after(afterCallback2)
+      after(afterCallback3)
+    })
 
     expect(afterCallback1).not.toHaveBeenCalled()
     expect(afterCallback2).not.toHaveBeenCalled()
@@ -405,9 +401,9 @@ describe('AfterContext', () => {
       onClose,
     })
 
-    const requestStore = createMockRequestStore(afterContext)
+    const workStore = createMockWorkStore(afterContext)
 
-    const run = createRun(afterContext, requestStore)
+    const run = createRun(afterContext, workStore)
 
     // ==================================
 
@@ -434,9 +430,9 @@ describe('AfterContext', () => {
       onClose,
     })
 
-    const requestStore = createMockRequestStore(afterContext)
+    const workStore = createMockWorkStore(afterContext)
 
-    const run = createRun(afterContext, requestStore)
+    const run = createRun(afterContext, workStore)
 
     // ==================================
 
@@ -452,7 +448,7 @@ describe('AfterContext', () => {
     expect(afterCallback1).not.toHaveBeenCalled()
   })
 
-  it('shadows requestAsyncStorage within after callbacks', async () => {
+  it('does NOT shadow workAsyncStorage within after callbacks', async () => {
     const waitUntil = jest.fn()
 
     let onCloseCallback: (() => void) | undefined = undefined
@@ -465,19 +461,19 @@ describe('AfterContext', () => {
       onClose,
     })
 
-    const requestStore = createMockRequestStore(afterContext)
-    const run = createRun(afterContext, requestStore)
+    const workStore = createMockWorkStore(afterContext)
+    const run = createRun(afterContext, workStore)
 
     // ==================================
 
     const stores = new DetachedPromise<
-      [RequestStore | undefined, RequestStore | undefined]
+      [WorkStore | undefined, WorkStore | undefined]
     >()
 
     await run(async () => {
-      const store1 = requestAsyncStorage.getStore()
+      const store1 = workAsyncStorage.getStore()
       after(() => {
-        const store2 = requestAsyncStorage.getStore()
+        const store2 = workAsyncStorage.getStore()
         stores.resolve([store1, store2])
       })
     })
@@ -486,30 +482,34 @@ describe('AfterContext', () => {
     onCloseCallback!()
 
     const [store1, store2] = await stores.promise
-    // if we use .toBe, the proxy from createMockRequestStore throws because jest checks '$$typeof'
+    // if we use .toBe, the proxy from createMockWorkStore throws because jest checks '$$typeof'
     expect(store1).toBeTruthy()
     expect(store2).toBeTruthy()
-    expect(store1 === requestStore).toBe(true)
-    expect(store2 !== store1).toBe(true)
+    expect(store1 === workStore).toBe(true)
+    expect(store2 === store1).toBe(true)
   })
 })
 
-const createMockRequestStore = (afterContext: AfterContext): RequestStore => {
-  const partialStore: Partial<RequestStore> = {
-    url: { pathname: '/', search: '' },
+const createMockWorkStore = (afterContext: AfterContext): WorkStore => {
+  const partialStore: Partial<WorkStore> = {
     afterContext: afterContext,
-    draftMode: undefined,
-    isHmrRefresh: false,
-    serverComponentsHmrCache: undefined,
+    forceStatic: false,
+    forceDynamic: false,
+    dynamicShouldError: false,
+    isStaticGeneration: false,
+    revalidatedTags: [],
+    pendingRevalidates: undefined,
+    pendingRevalidateWrites: undefined,
+    incrementalCache: undefined,
   }
 
-  return new Proxy(partialStore as RequestStore, {
+  return new Proxy(partialStore as WorkStore, {
     get(target, key) {
       if (key in target) {
         return target[key as keyof typeof target]
       }
       throw new Error(
-        `RequestStore property not mocked: '${typeof key === 'symbol' ? key.toString() : key}'`
+        `WorkStore property not mocked: '${typeof key === 'symbol' ? key.toString() : key}'`
       )
     },
   })

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

@@ -33,11 +33,6 @@ export class AfterContext {
     this.callbackQueue.pause()
   }
 
-  public run<T>(requestStore: RequestStore, callback: () => T): T {
-    this.requestStore = requestStore
-    return callback()
-  }
-
   public after(task: AfterTask): void {
     if (isThenable(task)) {
       task.catch(() => {}) // avoid unhandled rejection crashes
@@ -61,9 +56,10 @@ export class AfterContext {
       errorWaitUntilNotAvailable()
     }
     if (!this.requestStore) {
-      throw new InvariantError(
-        'unstable_after: Expected `AfterContext.requestStore` to be initialized'
-      )
+      // We just stash the first request store we have but this is not sufficient.
+      // TODO: We should store a request store per callback since each callback might
+      // be inside a different store. E.g. inside different batched actions, prerenders or caches.
+      this.requestStore = requestAsyncStorage.getStore()
     }
     if (!this.onClose) {
       throw new InvariantError(
@@ -98,19 +94,27 @@ export class AfterContext {
 
   private async runCallbacksOnClose() {
     await new Promise<void>((resolve) => this.onClose!(resolve))
-    return this.runCallbacks(this.requestStore!)
+    return this.runCallbacks(this.requestStore)
   }
 
-  private async runCallbacks(requestStore: RequestStore): Promise<void> {
+  private async runCallbacks(
+    requestStore: undefined | RequestStore
+  ): Promise<void> {
     if (this.callbackQueue.size === 0) return
 
-    const readonlyRequestStore: RequestStore =
-      wrapRequestStoreForAfterCallbacks(requestStore)
+    const readonlyRequestStore: undefined | RequestStore =
+      requestStore === undefined
+        ? undefined
+        : // TODO: This is not sufficient. It should just be the same store that mutates.
+          wrapRequestStoreForAfterCallbacks(requestStore)
 
     const workStore = workAsyncStorage.getStore()
 
     return withExecuteRevalidates(workStore, () =>
-      requestAsyncStorage.run(readonlyRequestStore, async () => {
+      // Clearing it out or running the first request store.
+      // TODO: This needs to be the request store that was active at the time the
+      // callback was scheduled but p-queue makes this hard so need further refactoring.
+      requestAsyncStorage.run(readonlyRequestStore as any, async () => {
         this.callbackQueue.start()
         await this.callbackQueue.onIdle()
       })
@@ -141,7 +145,6 @@ function wrapRequestStoreForAfterCallbacks(
     },
     // TODO(after): calling a `cookies.set()` in an after() that's in an action doesn't currently error.
     mutableCookies: new ResponseCookies(new Headers()),
-    afterContext: requestStore.afterContext,
     isHmrRefresh: requestStore.isHmrRefresh,
     serverComponentsHmrCache: requestStore.serverComponentsHmrCache,
   }

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

@@ -1,4 +1,3 @@
-import { requestAsyncStorage } from '../../client/components/request-async-storage.external'
 import { workAsyncStorage } from '../../client/components/work-async-storage.external'
 import { cacheAsyncStorage } from '../../server/app-render/cache-async-storage.external'
 import { StaticGenBailoutError } from '../../client/components/static-generation-bailout'
@@ -11,36 +10,30 @@ export type AfterCallback<T = unknown> = () => T | Promise<T>
 /**
  * This function allows you to schedule callbacks to be executed after the current request finishes.
  */
-export function unstable_after<T>(task: AfterTask<T>) {
-  const callingExpression = 'unstable_after'
-
-  // TODO: This is not safe. afterContext should move to WorkStore.
-  const requestStore = requestAsyncStorage.getStore()
-  if (!requestStore) {
-    throw new Error(
-      `\`${callingExpression}\` was called outside a request scope. Read more: https://nextjs.org/docs/messages/next-dynamic-api-wrong-context`
-    )
-  }
-
-  const { afterContext } = requestStore
-  if (!afterContext) {
-    throw new Error(
-      '`unstable_after()` must be explicitly enabled by setting `experimental.after: true` in your next.config.js.'
-    )
-  }
-
+export function unstable_after<T>(task: AfterTask<T>): void {
   const workStore = workAsyncStorage.getStore()
   const cacheStore = cacheAsyncStorage.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.'
+      )
+    }
+
+    // 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, cacheStore, callingExpression)
     }
-  }
 
-  return afterContext.after(task)
+    afterContext.after(task)
+  } else {
+    // TODO: Error for pages?
+  }
 }