use already existing trackBackgroundWork helper from request context to handle next/after

Author: pieh

src/run/handlers/server.ts

@@ -16,10 +16,12 @@ import { nextResponseProxy } from '../revalidate.js'
 
 import { createRequestContext, getLogger, getRequestContext } from './request-context.cjs'
 import { getTracer } from './tracer.cjs'
-import { setWaitUntil } from './wait-until.cjs'
+import { setupWaitUntil } from './wait-until.cjs'
 
 const nextImportPromise = import('../next.cjs')
 
+setupWaitUntil()
+
 let nextHandler: WorkerRequestHandler, nextConfig: NextConfigComplete
 
 /**
@@ -53,7 +55,7 @@ interface FutureContext extends Context {
 
 export default async (request: Request, context: FutureContext) => {
   const tracer = getTracer()
-  setWaitUntil(context)
+
   if (!nextHandler) {
     await tracer.withActiveSpan('initialize next server', async () => {
       // set the server config

src/run/handlers/wait-until.cts

@@ -1,3 +1,5 @@
+import { getRequestContext } from './request-context.cjs'
+
 /**
  *  @see https://github.com/vercel/next.js/blob/canary/packages/next/src/server/after/builtin-request-context.ts
  */
@@ -14,15 +16,11 @@ type GlobalThisWithRequestContext = typeof globalThis & {
 /**
  * Registers a `waitUntil` to be used by Next.js for next/after
  */
-
-export function setWaitUntil({ waitUntil }: { waitUntil?: (promise: Promise<unknown>) => void }) {
-  if (!waitUntil) {
-    return
-  }
+export function setupWaitUntil() {
   // eslint-disable-next-line @typescript-eslint/no-extra-semi
   ;(globalThis as GlobalThisWithRequestContext)[NEXT_REQUEST_CONTEXT_SYMBOL] = {
     get() {
-      return { waitUntil }
+      return { waitUntil: getRequestContext()?.trackBackgroundWork }
     },
   }
 }