add support for a global waitUntil

Author: conico974

packages/open-next/src/adapters/edge-adapter.ts

@@ -1,6 +1,7 @@
 import type { ReadableStream } from "node:stream/web";
 
 import type { InternalEvent, InternalResult } from "types/open-next";
+import { awaitAllDetachedPromise } from "utils/promise";
 import { emptyReadableStream } from "utils/stream";
 
 // We import it like that so that the edge plugin can replace it
@@ -64,6 +65,8 @@ const defaultHandler = async (
       const body =
         (response.body as ReadableStream<Uint8Array>) ?? emptyReadableStream();
 
+      await awaitAllDetachedPromise();
+
       return {
         type: "core",
         statusCode: response.status,

packages/open-next/src/adapters/middleware.ts

@@ -1,4 +1,5 @@
 import type { InternalEvent, Origin } from "types/open-next";
+import { awaitAllDetachedPromise } from "utils/promise";
 
 import { debug } from "../adapters/logger";
 import { createGenericHandler } from "../core/createGenericHandler";
@@ -51,6 +52,7 @@ const defaultHandler = async (internalEvent: InternalEvent) => {
         if (!result.isExternalRewrite) {
           origin = await originResolver.resolve(result.internalEvent.rawPath);
         }
+        await awaitAllDetachedPromise();
         return {
           type: "middleware",
           internalEvent: result.internalEvent,

packages/open-next/src/core/requestHandler.ts

@@ -3,6 +3,7 @@ import { AsyncLocalStorage } from "node:async_hooks";
 import type { OpenNextNodeResponse, StreamCreator } from "http/index.js";
 import { IncomingMessage } from "http/index.js";
 import type { InternalEvent, InternalResult } from "types/open-next";
+import { awaitAllDetachedPromise } from "utils/promise";
 
 import { debug, error, warn } from "../adapters/logger";
 import { generateOpenNextRequestContext } from "../adapters/util";
@@ -150,7 +151,7 @@ export async function openNextHandler(
       // reset lastModified. We need to do this to avoid memory leaks
       delete globalThis.lastModified[requestId];
 
-      await pendingPromiseRunner.await();
+        await awaitAllDetachedPromise();
 
       return internalResult;
     },

packages/open-next/src/overrides/wrappers/cloudflare.ts

@@ -11,13 +11,22 @@ const cfPropNameToHeaderName = {
   longitude: "x-open-next-longitude",
 };
 
+interface WorkerContext {
+  waitUntil: (promise: Promise<unknown>) => void;
+}
+
 const handler: WrapperHandler<
   InternalEvent,
   InternalResult | ({ type: "middleware" } & MiddlewareOutputEvent)
 > =
   async (handler, converter) =>
-  async (request: Request, env: Record<string, string>): Promise<Response> => {
+  async (
+    request: Request,
+    env: Record<string, string>,
+    ctx: WorkerContext,
+  ): Promise<Response> => {
     globalThis.process = process;
+    globalThis.openNextWaitUntil = ctx.waitUntil;
 
     // Set the environment variables
     // Cloudflare suggests to not override the process.env object but instead apply the values to it

packages/open-next/src/types/global.ts

@@ -138,12 +138,19 @@ declare global {
 
   /**
    * The AsyncLocalStorage instance that is used to store the request context.
-   * Only available in main functions.
+   * Only available in main, middleware and edge functions.
    * TODO: should be available everywhere in the future.
-   * Defined in `requestHandler.ts`.
+   * Defined in `requestHandler.ts`, `middleware.ts` and `edge-adapter.ts`.
    */
   var __als: AsyncLocalStorage<OpenNextRequestContext>;
 
+  /**
+   * The function that is used to run background tasks even after the response has been sent.
+   * This one is defined by the wrapper function as most of them don't need or support this feature.
+   * If not present, all the awaiting promises will be resolved before sending the response.
+   */
+  var openNextWaitUntil: ((promise: Promise<void>) => void) | undefined;
+
   /**
    * The entries object that contains the functions that are available in the function.
    * Only available in edge runtime functions.

packages/open-next/src/utils/promise.ts

@@ -58,3 +58,14 @@ export class DetachedPromiseRunner {
     });
   }
 }
+
+export async function awaitAllDetachedPromise() {
+  const promisesToAwait =
+    globalThis.__als.getStore()?.pendingPromiseRunner.await() ??
+    Promise.resolve();
+  if (globalThis.openNextWaitUntil) {
+    globalThis.openNextWaitUntil(promisesToAwait);
+    return;
+  }
+  await promisesToAwait;
+}