[unstable_after] compat: use waitUntil from NextRequestHint if no bui…

…ltin waitUntil is available

packages/next/src/client/components/lifecycle-async-storage-instance.ts

@@ -0,0 +1,5 @@
+import { createAsyncLocalStorage } from './async-local-storage'
+import type { LifecycleAsyncStorage } from './lifecycle-async-storage.external'
+
+export const lifecycleAsyncStorage: LifecycleAsyncStorage =
+  createAsyncLocalStorage()

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

@@ -0,0 +1,11 @@
+import type { AsyncLocalStorage } from 'async_hooks'
+// Share the instance module in the next-shared layer
+import { lifecycleAsyncStorage } from './lifecycle-async-storage-instance' with { 'turbopack-transition': 'next-shared' }
+
+export interface LifecycleStore {
+  readonly waitUntil: ((promise: Promise<any>) => void) | undefined
+}
+
+export type LifecycleAsyncStorage = AsyncLocalStorage<LifecycleStore>
+
+export { lifecycleAsyncStorage }

packages/next/src/server/base-server.ts

@@ -150,6 +150,7 @@ import {
   getBuiltinRequestContext,
   type WaitUntil,
 } from './after/builtin-request-context'
+import { lifecycleAsyncStorage } from '../client/components/lifecycle-async-storage.external'
 
 export type FindComponentsResult = {
   components: LoadComponentsReturnType
@@ -1679,6 +1680,11 @@ export default abstract class Server<
       return builtinRequestContext.waitUntil
     }
 
+    const lifecycleStore = lifecycleAsyncStorage.getStore()
+    if (lifecycleStore) {
+      return lifecycleStore.waitUntil
+    }
+
     if (process.env.__NEXT_TEST_MODE) {
       // we're in a test, use a no-op.
       return Server.noopWaitUntil

packages/next/src/server/web/adapter.ts

@@ -22,6 +22,8 @@ import { getTracer } from '../lib/trace/tracer'
 import type { TextMapGetter } from 'next/dist/compiled/@opentelemetry/api'
 import { MiddlewareSpan } from '../lib/trace/constants'
 import { CloseController } from './web-on-close'
+import { getBuiltinRequestContext } from '../after/builtin-request-context'
+import { lifecycleAsyncStorage } from '../../client/components/lifecycle-async-storage-instance'
 
 export class NextRequestHint extends NextRequest {
   sourcePage: string
@@ -213,13 +215,14 @@ export async function adapter(
     const isMiddleware =
       params.page === '/middleware' || params.page === '/src/middleware'
 
+    const isAfterEnabled =
+      params.request.nextConfig?.experimental?.after ??
+      !!process.env.__NEXT_AFTER
+
     if (isMiddleware) {
       // if we're in an edge function, we only get a subset of `nextConfig` (no `experimental`),
       // so we have to inject it via DefinePlugin.
       // in `next start` this will be passed normally (see `NextNodeServer.runMiddleware`).
-      const isAfterEnabled =
-        params.request.nextConfig?.experimental?.after ??
-        !!process.env.__NEXT_AFTER
 
       let waitUntil: WrapperRenderOpts['waitUntil'] = undefined
       let closeController: CloseController | undefined = undefined
@@ -279,6 +282,19 @@ export async function adapter(
         }
       )
     }
+
+    if (isAfterEnabled && !getBuiltinRequestContext()) {
+      // FIXME(after): unify on this method of passing waitUntil around,
+      // there's too much variation around this right now
+      //
+      // (here, we need to do this with an ALS because if the handler is a NextWebServer,
+      // there's currently no other way to pass `waitUntil` into it)
+      return lifecycleAsyncStorage.run(
+        { waitUntil: event.waitUntil.bind(event) },
+        () => params.handler(request, event)
+      )
+    }
+
     return params.handler(request, event)
   })