some assertions about ALS in after

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

@@ -61,9 +61,12 @@ export class AfterContext {
     }
 
     const workUnitStore = workUnitAsyncStorage.getStore()
-    if (workUnitStore) {
-      this.workUnitStores.add(workUnitStore)
+    if (!workUnitStore) {
+      throw new InvariantError(
+        'Missing workUnitStore in AfterContext.addCallback'
+      )
     }
+    this.workUnitStores.add(workUnitStore)
 
     // this should only happen once.
     if (!this.runCallbacksOnClosePromise) {
@@ -104,6 +107,9 @@ export class AfterContext {
     }
 
     const workStore = workAsyncStorage.getStore()
+    if (!workStore) {
+      throw new InvariantError('Missing workStore in AfterContext.runCallbacks')
+    }
 
     return withExecuteRevalidates(workStore, () => {
       this.callbackQueue.start()

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)
 }