From 2f95c614683cf61c7f48d45b9a8d58121861838d Mon Sep 17 00:00:00 2001
From: Janka Uryga <lolzatu2@gmail.com>
Date: Tue, 11 Jun 2024 13:41:29 +0200
Subject: [PATCH] fix: pass waitUntil from NextRequestHint to wrapped
 NextWebServer

---
 .../lifecycle-async-storage-instance.ts       |  5 ++++
 .../lifecycle-async-storage.external.ts       | 11 ++++++++
 packages/next/src/server/base-server.ts       |  6 +++++
 packages/next/src/server/web/adapter.ts       | 27 ++++++++++++++++---
 4 files changed, 46 insertions(+), 3 deletions(-)
 create mode 100644 packages/next/src/client/components/lifecycle-async-storage-instance.ts
 create mode 100644 packages/next/src/client/components/lifecycle-async-storage.external.ts

diff --git a/packages/next/src/client/components/lifecycle-async-storage-instance.ts b/packages/next/src/client/components/lifecycle-async-storage-instance.ts
new file mode 100644
index 0000000000000..738bcb4b85348
--- /dev/null
+++ b/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()
diff --git a/packages/next/src/client/components/lifecycle-async-storage.external.ts b/packages/next/src/client/components/lifecycle-async-storage.external.ts
new file mode 100644
index 0000000000000..ae9c25ad6cf58
--- /dev/null
+++ b/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 as 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 }
diff --git a/packages/next/src/server/base-server.ts b/packages/next/src/server/base-server.ts
index edd043102a176..bc703dba7e8b3 100644
--- a/packages/next/src/server/base-server.ts
+++ b/packages/next/src/server/base-server.ts
@@ -158,6 +158,7 @@ import {
 } from './after/builtin-request-context'
 import { ENCODED_TAGS } from './stream-utils/encodedTags'
 import { NextRequestHint } from './web/adapter'
+import { lifecycleAsyncStorage } from '../client/components/lifecycle-async-storage.external'
 
 export type FindComponentsResult = {
   components: LoadComponentsReturnType
@@ -1737,6 +1738,11 @@ export default abstract class Server<
   }
 
   protected getWaitUntil(): WaitUntil | undefined {
+    const lifecycleStore = lifecycleAsyncStorage.getStore()
+    if (lifecycleStore) {
+      return lifecycleStore.waitUntil
+    }
+
     const builtinRequestContext = getBuiltinRequestContext()
     if (builtinRequestContext) {
       // the platform provided a request context.
diff --git a/packages/next/src/server/web/adapter.ts b/packages/next/src/server/web/adapter.ts
index 9076257aaf662..fde7282a36fd1 100644
--- a/packages/next/src/server/web/adapter.ts
+++ b/packages/next/src/server/web/adapter.ts
@@ -22,6 +22,7 @@ import type { TextMapGetter } from 'next/dist/compiled/@opentelemetry/api'
 import { MiddlewareSpan } from '../lib/trace/constants'
 import { CloseController } from './web-on-close'
 import { getEdgePreviewProps } from './get-edge-preview-props'
+import { lifecycleAsyncStorage } from '../../client/components/lifecycle-async-storage.external'
 
 export class NextRequestHint extends NextRequest {
   sourcePage: string
@@ -207,13 +208,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
@@ -271,6 +273,25 @@ export async function adapter(
         }
       )
     }
+
+    if (isAfterEnabled) {
+      // NOTE:
+      // Currently, `adapter` is expected to return promises passed to `waitUntil`
+      // as part of its result (i.e. a FetchEventResult).
+      // Because of this, we override any outer contexts that might provide a real `waitUntil`,
+      // and provide the `waitUntil` from the NextFetchEvent instead so that we can collect those promises.
+      // This is not ideal, but until we change this calling convention, it's the least surprising thing to do.
+      //
+      // Notably, the only case that currently cares about this ALS is Edge SSR
+      // (i.e. a handler created via `build/webpack/loaders/next-edge-ssr-loader/render.ts`)
+      // Other types of handlers will grab the waitUntil from the passed FetchEvent,
+      // but NextWebServer currently has no interface that'd allow for that.
+      return lifecycleAsyncStorage.run(
+        { waitUntil: event.waitUntil.bind(event) },
+        () => params.handler(request, event)
+      )
+    }
+
     return params.handler(request, event)
   })