fix: after() in edge functions (middleware and route handlers)

packages/next/src/build/webpack/plugins/define-env-plugin.ts

@@ -170,6 +170,7 @@ export function getDefineEnv({
       : '',
     'process.env.NEXT_MINIMAL': '',
     'process.env.__NEXT_PPR': checkIsAppPPREnabled(config.experimental.ppr),
+    'process.env.__NEXT_AFTER': config.experimental.after ?? false,
     'process.env.NEXT_DEPLOYMENT_ID': config.deploymentId || false,
     'process.env.__NEXT_FETCH_CACHE_KEY_PREFIX': fetchCacheKeyPrefix ?? '',
     'process.env.__NEXT_MIDDLEWARE_MATCHERS': middlewareMatchers ?? [],

packages/next/src/export/routes/app-route.ts

@@ -71,7 +71,6 @@ export async function exportAppRoute(
       nextExport: true,
       supportsDynamicHTML: false,
       incrementalCache,
-      // @ts-expect-error TODO(after): fix the typing here
       waitUntil: function noWaitUntilInPrerender() {
         throw new Error('waitUntil cannot be called during prerendering.')
       },

packages/next/src/server/async-storage/request-async-storage-wrapper.ts

@@ -51,7 +51,7 @@ function getMutableCookies(
   return MutableRequestCookiesAdapter.wrap(cookies, onUpdateCookies)
 }
 
-type WrapperRenderOpts = RenderOpts &
+export type WrapperRenderOpts = RenderOpts &
   Partial<RequestLifecycleOpts> &
   RenderOptsForRouteHandlerPartial &
   RenderOptsForWebServerPartial
@@ -64,7 +64,7 @@ type RenderOptsForRouteHandlerPartial = Partial<
 >
 
 type RenderOptsForWebServerPartial = {
-  experimental?: Partial<Pick<RenderOpts['experimental'], 'after'>> // can be undefined in middleware
+  experimental: Pick<RenderOpts['experimental'], 'after'>
 }
 
 export type RequestContext = {

packages/next/src/server/async-storage/static-generation-async-storage-wrapper.ts

@@ -6,6 +6,7 @@ import type { RenderOptsPartial } from '../app-render/types'
 
 import { createPrerenderState } from '../../server/app-render/dynamic-rendering'
 import type { FetchMetric } from '../base-http'
+import type { RequestLifecycleOpts } from '../base-server'
 
 export type StaticGenerationContext = {
   urlPathname: string
@@ -45,7 +46,8 @@ export type StaticGenerationContext = {
     | 'nextExport'
     | 'isDraftMode'
     | 'isDebugPPRSkeleton'
-  >
+  > &
+    Partial<RequestLifecycleOpts>
 }
 
 export const StaticGenerationAsyncStorageWrapper: AsyncStorageWrapper<

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

@@ -2376,7 +2376,6 @@ export default abstract class Server<
               supportsDynamicHTML,
               incrementalCache,
               isRevalidate: isSSG,
-              // @ts-expect-error TODO(after): fix the typing here
               waitUntil: this.getWaitUntil(),
               onClose: res.onClose.bind(res),
             },

packages/next/src/server/future/route-modules/app-route/module.ts

@@ -250,13 +250,11 @@ export class AppRouteRouteModule extends RouteModule<
       req: rawRequest,
     }
 
-    ;(requestContext as any).renderOpts = {
-      // TODO: types for renderOpts should include previewProps
+    requestContext.renderOpts = {
       previewProps: context.prerenderManifest.preview,
-      // @ts-expect-error TODO(after): fix the typing here
       waitUntil: context.renderOpts.waitUntil,
-      // @ts-expect-error TODO(after): fix the typing here
       onClose: context.renderOpts.onClose,
+      // @ts-expect-error: not all required properties are available here for some reason
       experimental: context.renderOpts.experimental,
     }
 

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

@@ -13,11 +13,15 @@ import { stripInternalSearchParams } from '../internal-utils'
 import { normalizeRscURL } from '../../shared/lib/router/utils/app-paths'
 import { FLIGHT_PARAMETERS } from '../../client/components/app-router-headers'
 import { ensureInstrumentationRegistered } from './globals'
-import { RequestAsyncStorageWrapper } from '../async-storage/request-async-storage-wrapper'
+import {
+  RequestAsyncStorageWrapper,
+  type WrapperRenderOpts,
+} from '../async-storage/request-async-storage-wrapper'
 import { requestAsyncStorage } from '../../client/components/request-async-storage.external'
 import { getTracer } from '../lib/trace/tracer'
 import type { TextMapGetter } from 'next/dist/compiled/@opentelemetry/api'
 import { MiddlewareSpan } from '../lib/trace/constants'
+import type { RenderOptsPartial } from '../app-render/types'
 
 export class NextRequestHint extends NextRequest {
   sourcePage: string
@@ -204,19 +208,35 @@ export async function adapter(
   let response
   let cookiesFromResponse
 
-  const requestClosedTarget = new EventTarget()
-  const onRequestClose = (callback: () => void) => {
-    requestClosedTarget.addEventListener('close', callback)
-  }
-  const emitRequestClose = () => {
-    requestClosedTarget.dispatchEvent(new Event('close'))
-  }
-
   response = await propagator(request, () => {
     // we only care to make async storage available for middleware
     const isMiddleware =
       params.page === '/middleware' || params.page === '/src/middleware'
+
     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 onClose: WrapperRenderOpts['onClose'] = undefined
+      let dispatchClose: (() => void) | undefined = undefined
+
+      if (isAfterEnabled) {
+        waitUntil = event.waitUntil.bind(event)
+
+        const requestClosedTarget = new EventTarget()
+        onClose = (callback: () => void) => {
+          requestClosedTarget.addEventListener('close', callback)
+        }
+        dispatchClose = () => {
+          requestClosedTarget.dispatchEvent(new Event('close'))
+        }
+      }
+
       return getTracer().trace(
         MiddlewareSpan.execute,
         {
@@ -235,18 +255,32 @@ export async function adapter(
                 onUpdateCookies: (cookies) => {
                   cookiesFromResponse = cookies
                 },
+                // @ts-expect-error TODO: investigate why previewProps isn't on RenderOpts
                 previewProps: prerenderManifest?.preview || {
                   previewModeId: 'development-id',
                   previewModeEncryptionKey: '',
                   previewModeSigningKey: '',
                 },
-                waitUntil: event.waitUntil.bind(event),
-                onClose: onRequestClose,
-                // @ts-expect-error TODO(after): not sure what to do about this
-                experimental: params.request.nextConfig?.experimental,
+                waitUntil,
+                onClose,
+                experimental: {
+                  after: isAfterEnabled,
+                } as RenderOptsPartial['experimental'],
               },
             },
-            () => params.handler(request, event)
+            async () => {
+              try {
+                return await params.handler(request, event)
+              } finally {
+                // middleware cannot stream, so we can consider the response closed
+                // as soon as the handler returns.
+                if (dispatchClose) {
+                  // we can delay running it until a bit later --
+                  // if it's needed, we'll have a `waitUntil` lock anyway.
+                  setTimeout(dispatchClose, 0)
+                }
+              }
+            }
           )
       )
     }
@@ -258,8 +292,6 @@ export async function adapter(
     throw new TypeError('Expected an instance of Response to be returned')
   }
 
-  emitRequestClose()
-
   if (response && cookiesFromResponse) {
     response.headers.set('set-cookie', cookiesFromResponse)
   }

packages/next/src/server/web/edge-route-module-wrapper.ts

@@ -14,6 +14,7 @@ import type { NextFetchEvent } from './spec-extension/fetch-event'
 import { internal_getCurrentFunctionWaitUntil } from './internal-edge-wait-until'
 import { getUtils } from '../server-utils'
 import { searchParamsToUrlQuery } from '../../shared/lib/router/utils/querystring'
+import type { RequestLifecycleOpts } from '../base-server'
 
 type WrapOptions = Partial<Pick<AdapterOptions, 'page'>>
 
@@ -87,6 +88,23 @@ export class EdgeRouteModuleWrapper {
         ? JSON.parse(self.__PRERENDER_MANIFEST)
         : undefined
 
+    const isAfterEnabled = !!process.env.__NEXT_AFTER
+
+    let waitUntil: RequestLifecycleOpts['waitUntil'] | undefined = undefined
+    let onClose: RequestLifecycleOpts['onClose'] | undefined = undefined
+    let dispatchClose: (() => void) | undefined = undefined
+    if (isAfterEnabled) {
+      waitUntil = evt.waitUntil.bind(evt)
+
+      const requestClosedTarget = new EventTarget()
+      onClose = (callback: () => void) => {
+        requestClosedTarget.addEventListener('close', callback)
+      }
+      dispatchClose = () => {
+        requestClosedTarget.dispatchEvent(new Event('close'))
+      }
+    }
+
     // Create the context for the handler. This contains the params from the
     // match (if any).
     const context: AppRouteRouteHandlerContext = {
@@ -104,22 +122,31 @@ export class EdgeRouteModuleWrapper {
       },
       renderOpts: {
         supportsDynamicHTML: true,
+        waitUntil,
+        onClose,
         experimental: {
-          // @ts-expect-error TODO(after): not sure what to do about this
-          after: undefined,
+          after: isAfterEnabled,
         },
       },
     }
 
-    // Get the response from the handler.
-    const res = await this.routeModule.handle(request, context)
+    try {
+      // Get the response from the handler.
+      const res = await this.routeModule.handle(request, context)
 
-    const waitUntilPromises = [internal_getCurrentFunctionWaitUntil()]
-    if (context.renderOpts.pendingWaitUntil) {
-      waitUntilPromises.push(context.renderOpts.pendingWaitUntil)
-    }
-    evt.waitUntil(Promise.all(waitUntilPromises))
+      const waitUntilPromises = [internal_getCurrentFunctionWaitUntil()]
+      if (context.renderOpts.pendingWaitUntil) {
+        waitUntilPromises.push(context.renderOpts.pendingWaitUntil)
+      }
+      evt.waitUntil(Promise.all(waitUntilPromises))
 
-    return res
+      return res
+    } finally {
+      // TODO(after): this might be a streaming response, in which case this'll run too early.
+      // we should probably do the same thing as `WebNextResponse#onClose` here
+      if (dispatchClose) {
+        setTimeout(dispatchClose, 0)
+      }
+    }
   }
 }