fix: include promises from late waitUntil calls in FetchEventResult.w…

…aitUntil

packages/next/src/build/webpack/loaders/next-edge-ssr-loader/render.ts

@@ -157,17 +157,27 @@ export function getRender({
     request: NextRequestHint,
     event?: NextFetchEvent
   ) {
+    const isAfterEnabled = !!process.env.__NEXT_AFTER
+
     const extendedReq = new WebNextRequest(request)
     const extendedRes = new WebNextResponse(
       undefined,
       // tracking onClose adds overhead, so only do it if `experimental.after` is on.
-      !!process.env.__NEXT_AFTER
+      isAfterEnabled
     )
 
     handler(extendedReq, extendedRes)
     const result = await extendedRes.toResponse()
 
     if (event?.waitUntil) {
+      if (isAfterEnabled) {
+        // make sure that NextRequestHint's awaiter stays open long enough
+        // for late `waitUntil`s called during streaming to get picked up.
+        event.waitUntil(
+          new Promise<void>((resolve) => extendedRes.onClose(resolve))
+        )
+      }
+
       // TODO(after):
       // remove `internal_runWithWaitUntil` and the `internal-edge-wait-until` module
       // when consumers switch to `unstable_after`.

packages/next/src/server/lib/awaiter.test.ts

@@ -0,0 +1,81 @@
+import { InvariantError } from '../../shared/lib/invariant-error'
+import { AwaiterMulti, AwaiterOnce } from './awaiter'
+
+describe('AwaiterOnce/AwaiterMulti', () => {
+  describe.each([
+    { name: 'AwaiterMulti', impl: AwaiterMulti },
+    { name: 'AwaiterOnce', impl: AwaiterOnce },
+  ])('$name', ({ impl: AwaiterImpl }) => {
+    it('awaits promises added by other promises', async () => {
+      const awaiter = new AwaiterImpl()
+
+      const MAX_DEPTH = 5
+      const promises: TrackedPromise<unknown>[] = []
+
+      const waitUntil = (promise: Promise<unknown>) => {
+        promises.push(trackPromiseSettled(promise))
+        awaiter.waitUntil(promise)
+      }
+
+      const makeNestedPromise = async () => {
+        if (promises.length >= MAX_DEPTH) {
+          return
+        }
+        await sleep(100)
+        waitUntil(makeNestedPromise())
+      }
+
+      waitUntil(makeNestedPromise())
+
+      await awaiter.awaiting()
+
+      for (const promise of promises) {
+        expect(promise.isSettled).toBe(true)
+      }
+    })
+
+    it('calls onError for rejected promises', async () => {
+      const onError = jest.fn<void, [error: unknown]>()
+      const awaiter = new AwaiterImpl({ onError })
+
+      awaiter.waitUntil(Promise.reject('error 1'))
+      awaiter.waitUntil(
+        sleep(100).then(() => awaiter.waitUntil(Promise.reject('error 2')))
+      )
+
+      await awaiter.awaiting()
+
+      expect(onError).toHaveBeenCalledWith('error 1')
+      expect(onError).toHaveBeenCalledWith('error 2')
+    })
+  })
+})
+
+describe('AwaiterOnce', () => {
+  it("does not allow calling waitUntil after it's been awaited", async () => {
+    const awaiter = new AwaiterOnce()
+    awaiter.waitUntil(Promise.resolve(1))
+    await awaiter.awaiting()
+    expect(() => awaiter.waitUntil(Promise.resolve(2))).toThrow(InvariantError)
+  })
+})
+
+type TrackedPromise<T> = Promise<T> & { isSettled: boolean }
+
+function trackPromiseSettled<T>(promise: Promise<T>): TrackedPromise<T> {
+  const tracked = promise as TrackedPromise<T>
+  tracked.isSettled = false
+  tracked.then(
+    () => {
+      tracked.isSettled = true
+    },
+    () => {
+      tracked.isSettled = true
+    }
+  )
+  return tracked
+}
+
+function sleep(duration: number) {
+  return new Promise<void>((resolve) => setTimeout(resolve, duration))
+}

test/e2e/app-dir/next-after-app/utils/awaiter.ts → packages/next/src/server/lib/awaiter.ts

@@ -1,4 +1,4 @@
-import { InvariantError } from 'next/dist/shared/lib/invariant-error'
+import { InvariantError } from '../../shared/lib/invariant-error'
 
 /**
  * Provides a `waitUntil` implementation which gathers promises to be awaited later (via {@link AwaiterMulti.awaiting}).

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

@@ -422,7 +422,7 @@ export async function adapter(
 
   return {
     response: finalResponse,
-    waitUntil: Promise.all(event[waitUntilSymbol]),
+    waitUntil: event[waitUntilSymbol](),
     fetchMetrics: request.fetchMetrics,
   }
 }

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

@@ -148,6 +148,13 @@ export class EdgeRouteModuleWrapper {
         const trackedBody = trackStreamConsumed(res.body, () =>
           _closeController.dispatchClose()
         )
+
+        // make sure that NextRequestHint's awaiter stays open long enough
+        // for `waitUntil`s called late during streaming to get picked up.
+        evt.waitUntil(
+          new Promise<void>((resolve) => _closeController.onClose(resolve))
+        )
+
         res = new Response(trackedBody, {
           status: res.status,
           statusText: res.statusText,

packages/next/src/server/web/spec-extension/fetch-event.ts

@@ -1,14 +1,22 @@
+import { AwaiterOnce } from '../../lib/awaiter'
 import { PageSignatureError } from '../error'
 import type { NextRequest } from './request'
 
 const responseSymbol = Symbol('response')
 const passThroughSymbol = Symbol('passThrough')
+const awaiterSymbol = Symbol('awaiter')
+
 export const waitUntilSymbol = Symbol('waitUntil')
 
 class FetchEvent {
-  readonly [waitUntilSymbol]: Promise<any>[] = [];
   [responseSymbol]?: Promise<Response>;
-  [passThroughSymbol] = false
+  [passThroughSymbol] = false;
+
+  [awaiterSymbol] = new AwaiterOnce();
+
+  [waitUntilSymbol] = () => {
+    return this[awaiterSymbol].awaiting()
+  }
 
   // eslint-disable-next-line @typescript-eslint/no-useless-constructor
   constructor(_request: Request) {}
@@ -24,7 +32,7 @@ class FetchEvent {
   }
 
   waitUntil(promise: Promise<any>): void {
-    this[waitUntilSymbol].push(promise)
+    this[awaiterSymbol].waitUntil(promise)
   }
 }
 

test/e2e/app-dir/next-after-app/index.test.ts

@@ -380,10 +380,7 @@ describe.each(runtimes)('unstable_after() in %s runtime', (runtimeValue) => {
         return cleanup
       }
 
-      /* eslint-disable jest/no-standalone-expect */
-      const it_failingForEdge = runtimeValue === 'edge' ? it.failing : it
-
-      it_failingForEdge('during render', async () => {
+      it('during render', async () => {
         const cleanup = await setup()
         try {
           const response = await next.fetch('/delay-deep')
@@ -401,27 +398,23 @@ describe.each(runtimes)('unstable_after() in %s runtime', (runtimeValue) => {
         }
       })
 
-      it_failingForEdge(
-        'in a route handler that streams a response',
-        async () => {
-          const cleanup = await setup()
-          try {
-            const response = await next.fetch('/route-streaming')
-            expect(response.status).toBe(200)
-            await response.text()
-            await retry(() => {
-              expect(getLogs()).toContainEqual('simulated-invocation :: end')
-            }, 10_000)
-
-            expect(getLogs()).toContainEqual({
-              source: '[route handler] /route-streaming - after',
-            })
-          } finally {
-            await cleanup()
-          }
+      it('in a route handler that streams a response', async () => {
+        const cleanup = await setup()
+        try {
+          const response = await next.fetch('/route-streaming')
+          expect(response.status).toBe(200)
+          await response.text()
+          await retry(() => {
+            expect(getLogs()).toContainEqual('simulated-invocation :: end')
+          }, 10_000)
+
+          expect(getLogs()).toContainEqual({
+            source: '[route handler] /route-streaming - after',
+          })
+        } finally {
+          await cleanup()
         }
-      )
-      /* eslint-enable jest/no-standalone-expect */
+      })
     })
   }
 

test/e2e/app-dir/next-after-app/utils/simulated-invocation.js

@@ -1,5 +1,5 @@
 import { requestAsyncStorage } from 'next/dist/client/components/request-async-storage.external'
-import { AwaiterOnce } from './awaiter'
+import { AwaiterOnce } from 'next/dist/server/lib/awaiter'
 import { cliLog } from './log'
 
 // replaced in tests