feat(after): allow request APIs in after (actions/handlers) (#73345)

This PR relaxes some of the limitations around calling request APIs
(`headers()`, `cookies()`, `connection()`) inside `unstable_after`
(which was previously blanket-forbidden since #71231).
We're now allowing it in server actions and route handlers[1]:
```tsx
// route.ts ✅
export function POST() {
  unstable_after(async () => console.log(await headers()));
}
```
```tsx
// action.ts ✅
export function myAction() {
  "use server"
  unstable_after(async () => console.log(await headers()));
}
```

however using it in server components remains banned:
```tsx
// page.tsx ❌
export default function Page() {
  unstable_after(async () => console.log(await headers()));
}
```

### Implementation notes

The way to detect actions and route handlers is to check if
`workUnitStore.phase === "action"` [2]. But we can't use the existing
`workUnitStore.phase` for this because it's deliberately mutable, so
it'd have changed by the time we get to executing the callbacks and the
actual `headers()` call.

To solve this, I'm introducing a new ALS store - `afterTaskStore`. It
currently only stores the phase that the (topmost) `after` was called
in. But this info is also specific to each "task"/callback (or rather,
task chain -- nested afters inherit the context of their parents), so we
can't use any of the existing stores.

I 'm also going to need task-specific info for upcoming devtools work
(more useful callstacks, see
https://github.com/vercel/next.js/blob/54465dbe782f505fd0105258d2d053a6439457ed/packages/next/src/server/app-render/after-task-async-storage.external.ts#L14),
so i think introducing a new store is justified.

---
[1] - Note that if it's in a prerendered GET handler (with `dynamic =
'error'`), `headers()` & co. will still error, but that's because
they're dynamic, not because they were called within `after`!
[2] - `actionAsyncStorage` is a no-go, because `actionStore.isAction`
can remain true in the render triggered by an action (or in promises
passed into it), which was the reason we introduced `phase` in the first
place.

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

@@ -10,6 +10,7 @@ import {
   workUnitAsyncStorage,
   type WorkUnitStore,
 } from '../app-render/work-unit-async-storage.external'
+import { afterTaskAsyncStorage } from '../app-render/after-task-async-storage.external'
 
 export type AfterContextOpts = {
   isEnabled: boolean
@@ -70,6 +71,16 @@ export class AfterContext {
       this.workUnitStores.add(workUnitStore)
     }
 
+    const afterTaskStore = afterTaskAsyncStorage.getStore()
+
+    // This is used for checking if request APIs can be called inside `after`.
+    // Note that we need to check the phase in which the *topmost* `after` was called (which should be "action"),
+    // not the current phase (which might be "after" if we're in a nested after).
+    // Otherwise, we might allow `after(() => headers())`, but not `after(() => after(() => headers()))`.
+    const rootTaskSpawnPhase = afterTaskStore
+      ? afterTaskStore.rootTaskSpawnPhase // nested after
+      : workUnitStore?.phase // topmost after
+
     // this should only happen once.
     if (!this.runCallbacksOnClosePromise) {
       this.runCallbacksOnClosePromise = this.runCallbacksOnClose()
@@ -83,7 +94,9 @@ export class AfterContext {
     //   await x()
     const wrappedCallback = bindSnapshot(async () => {
       try {
-        await callback()
+        await afterTaskAsyncStorage.run({ rootTaskSpawnPhase }, () =>
+          callback()
+        )
       } catch (error) {
         this.reportTaskError(error)
       }

packages/next/src/server/app-render/after-task-async-storage-instance.ts

@@ -0,0 +1,5 @@
+import type { AfterTaskAsyncStorage } from './after-task-async-storage.external'
+import { createAsyncLocalStorage } from './async-local-storage'
+
+export const afterTaskAsyncStorageInstance: AfterTaskAsyncStorage =
+  createAsyncLocalStorage()

packages/next/src/server/app-render/after-task-async-storage.external.ts

@@ -0,0 +1,18 @@
+import type { AsyncLocalStorage } from 'async_hooks'
+
+// Share the instance module in the next-shared layer
+import { afterTaskAsyncStorageInstance as afterTaskAsyncStorage } from './after-task-async-storage-instance' with { 'turbopack-transition': 'next-shared' }
+import type { WorkUnitStore } from './work-unit-async-storage.external'
+
+export interface AfterTaskStore {
+  /** The phase in which the topmost `unstable_after` was called.
+   *
+   * NOTE: Can be undefined when running `generateStaticParams`,
+   * where we only have a `workStore`, no `workUnitStore`.
+   */
+  readonly rootTaskSpawnPhase: WorkUnitStore['phase'] | undefined
+}
+
+export type AfterTaskAsyncStorage = AsyncLocalStorage<AfterTaskStore>
+
+export { afterTaskAsyncStorage }

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

@@ -7,6 +7,7 @@ import {
 } from '../app-render/dynamic-rendering'
 import { StaticGenBailoutError } from '../../client/components/static-generation-bailout'
 import { makeHangingPromise } from '../dynamic-rendering-utils'
+import { isRequestAPICallableInsideAfter } from './utils'
 
 /**
  * This function allows you to indicate that you require an actual user Request before continuing.
@@ -18,7 +19,11 @@ export function connection(): Promise<void> {
   const workUnitStore = workUnitAsyncStorage.getStore()
 
   if (workStore) {
-    if (workUnitStore && workUnitStore.phase === 'after') {
+    if (
+      workUnitStore &&
+      workUnitStore.phase === 'after' &&
+      !isRequestAPICallableInsideAfter()
+    ) {
       throw new Error(
         `Route ${workStore.route} used "connection" inside "unstable_after(...)". The \`connection()\` function is used to indicate the subsequent code must only run when there is an actual Request, but "unstable_after(...)" executes after the request, so this function is not allowed in this scope. See more info here: https://nextjs.org/docs/canary/app/api-reference/functions/unstable_after`
       )

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

@@ -22,6 +22,7 @@ import { StaticGenBailoutError } from '../../client/components/static-generation
 import { makeHangingPromise } from '../dynamic-rendering-utils'
 import { createDedupedByCallsiteServerErrorLoggerDev } from '../create-deduped-by-callsite-server-error-logger'
 import { scheduleImmediate } from '../../lib/scheduler'
+import { isRequestAPICallableInsideAfter } from './utils'
 
 /**
  * In this version of Next.js `cookies()` returns a Promise however you can still reference the properties of the underlying cookies object
@@ -52,7 +53,11 @@ export function cookies(): Promise<ReadonlyRequestCookies> {
   const workUnitStore = workUnitAsyncStorage.getStore()
 
   if (workStore) {
-    if (workUnitStore && workUnitStore.phase === 'after') {
+    if (
+      workUnitStore &&
+      workUnitStore.phase === 'after' &&
+      !isRequestAPICallableInsideAfter()
+    ) {
       throw new Error(
         // TODO(after): clarify that this only applies to pages?
         `Route ${workStore.route} used "cookies" inside "unstable_after(...)". This is not supported. If you need this data inside an "unstable_after" callback, use "cookies" outside of the callback. See more info here: https://nextjs.org/docs/canary/app/api-reference/functions/unstable_after`

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

@@ -19,6 +19,7 @@ import { StaticGenBailoutError } from '../../client/components/static-generation
 import { makeHangingPromise } from '../dynamic-rendering-utils'
 import { createDedupedByCallsiteServerErrorLoggerDev } from '../create-deduped-by-callsite-server-error-logger'
 import { scheduleImmediate } from '../../lib/scheduler'
+import { isRequestAPICallableInsideAfter } from './utils'
 
 /**
  * In this version of Next.js `headers()` returns a Promise however you can still reference the properties of the underlying Headers instance
@@ -57,7 +58,11 @@ export function headers(): Promise<ReadonlyHeaders> {
   const workUnitStore = workUnitAsyncStorage.getStore()
 
   if (workStore) {
-    if (workUnitStore && workUnitStore.phase === 'after') {
+    if (
+      workUnitStore &&
+      workUnitStore.phase === 'after' &&
+      !isRequestAPICallableInsideAfter()
+    ) {
       throw new Error(
         `Route ${workStore.route} used "headers" inside "unstable_after(...)". This is not supported. If you need this data inside an "unstable_after" callback, use "headers" outside of the callback. See more info here: https://nextjs.org/docs/canary/app/api-reference/functions/unstable_after`
       )

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

@@ -1,4 +1,5 @@
 import { StaticGenBailoutError } from '../../client/components/static-generation-bailout'
+import { afterTaskAsyncStorage } from '../app-render/after-task-async-storage.external'
 
 // This regex will have fast negatives meaning valid identifiers may not pass
 // this test. However this is only used during static generation to provide hints
@@ -40,6 +41,11 @@ export function throwWithStaticGenerationBailoutErrorWithDynamicError(
   )
 }
 
+export function isRequestAPICallableInsideAfter() {
+  const afterTaskStore = afterTaskAsyncStorage.getStore()
+  return afterTaskStore?.rootTaskSpawnPhase === 'action'
+}
+
 export const wellKnownProperties = new Set([
   'hasOwnProperty',
   'isPrototypeOf',

test/e2e/app-dir/next-after-app-api-usage/app/request-apis/helpers.js

@@ -1,31 +1,64 @@
 import { cookies, headers } from 'next/headers'
 import { unstable_after as after, connection } from 'next/server'
 
-export function testRequestAPIs() {
+export function testRequestAPIs(/** @type {string} */ route) {
   after(async () => {
     try {
       await headers()
-      console.log('headers(): ok')
+      console.log(`[${route}] headers(): ok`)
     } catch (err) {
-      console.error(err)
+      console.error(`[${route}] headers(): error:`, err)
     }
   })
 
+  after(() =>
+    after(async () => {
+      try {
+        await headers()
+        console.log(`[${route}] nested headers(): ok`)
+      } catch (err) {
+        console.error(`[${route}] nested headers(): error:`, err)
+      }
+    })
+  )
+
   after(async () => {
     try {
       await cookies()
-      console.log('cookies(): ok')
+      console.log(`[${route}] cookies(): ok`)
     } catch (err) {
-      console.error(err)
+      console.error(`[${route}] cookies(): error:`, err)
     }
   })
 
+  after(() =>
+    after(async () => {
+      try {
+        await cookies()
+        console.log(`[${route}] nested cookies(): ok`)
+      } catch (err) {
+        console.error(`[${route}] nested cookies(): error:`, err)
+      }
+    })
+  )
+
   after(async () => {
     try {
       await connection()
-      console.log('connection(): ok')
+      console.log(`[${route}] connection(): ok`)
     } catch (err) {
-      console.error(err)
+      console.error(`[${route}] connection(): error:`, err)
     }
   })
+
+  after(() =>
+    after(async () => {
+      try {
+        await connection()
+        console.log(`[${route}] nested connection(): ok`)
+      } catch (err) {
+        console.error(`[${route}] nested connection(): error:`, err)
+      }
+    })
+  )
 }

test/e2e/app-dir/next-after-app-api-usage/app/request-apis/page-dynamic-error/page.js

@@ -3,6 +3,6 @@ import { testRequestAPIs } from '../helpers'
 export const dynamic = 'error'
 
 export default async function Page() {
-  testRequestAPIs()
+  testRequestAPIs('/request-apis/page-dynamic-error')
   return null
 }

test/e2e/app-dir/next-after-app-api-usage/app/request-apis/page-dynamic/page.js

@@ -3,6 +3,6 @@ import { testRequestAPIs } from '../helpers'
 
 export default async function Page() {
   await connection()
-  testRequestAPIs()
+  testRequestAPIs('/request-apis/page-dynamic')
   return null
 }

test/e2e/app-dir/next-after-app-api-usage/app/request-apis/page-force-static/page.js

@@ -3,6 +3,6 @@ import { testRequestAPIs } from '../helpers'
 export const dynamic = 'force-static'
 
 export default async function Page() {
-  testRequestAPIs()
+  testRequestAPIs('/request-apis/page-force-static')
   return null
 }

test/e2e/app-dir/next-after-app-api-usage/app/request-apis/route-handler-dynamic-error/route.js

@@ -3,6 +3,6 @@ import { testRequestAPIs } from '../helpers'
 export const dynamic = 'error'
 
 export async function GET() {
-  testRequestAPIs()
+  testRequestAPIs('/request-apis/route-handler-dynamic-error')
   return new Response()
 }

test/e2e/app-dir/next-after-app-api-usage/app/request-apis/route-handler-dynamic/route.js

@@ -1,6 +1,6 @@
 import { testRequestAPIs } from '../helpers'
 
 export async function GET() {
-  testRequestAPIs()
+  testRequestAPIs('/request-apis/route-handler-dynamic')
   return new Response()
 }

test/e2e/app-dir/next-after-app-api-usage/app/request-apis/route-handler-force-static/route.js

@@ -3,6 +3,6 @@ import { testRequestAPIs } from '../helpers'
 export const dynamic = 'force-static'
 
 export async function GET() {
-  testRequestAPIs()
+  testRequestAPIs('/request-apis/route-handler-force-static')
   return new Response()
 }

test/e2e/app-dir/next-after-app-api-usage/app/request-apis/server-action/page.js

@@ -5,7 +5,7 @@ export default async function Page() {
     <form
       action={async () => {
         'use server'
-        testRequestAPIs()
+        testRequestAPIs('/request-apis/server-action')
       }}
     >
       <button type="submit">Submit</button>