next/after support

Author: conico974

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

@@ -1,7 +1,10 @@
 import type { ReadableStream } from "node:stream/web";
 
 import type { InternalEvent, InternalResult } from "types/open-next";
-import { awaitAllDetachedPromise } from "utils/promise";
+import {
+  awaitAllDetachedPromise,
+  provideNextAfterProvider,
+} from "utils/promise";
 import { emptyReadableStream } from "utils/stream";
 
 // We import it like that so that the edge plugin can replace it
@@ -27,6 +30,7 @@ const defaultHandler = async (
   return globalThis.__als.run(
     { requestId, pendingPromiseRunner, isISRRevalidation },
     async () => {
+      provideNextAfterProvider();
       const host = internalEvent.headers.host
         ? `https://${internalEvent.headers.host}`
         : "http://localhost:3000";

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

@@ -1,5 +1,8 @@
 import type { InternalEvent, Origin } from "types/open-next";
-import { awaitAllDetachedPromise } from "utils/promise";
+import {
+  awaitAllDetachedPromise,
+  provideNextAfterProvider,
+} from "utils/promise";
 
 import { debug } from "../adapters/logger";
 import { createGenericHandler } from "../core/createGenericHandler";
@@ -45,6 +48,7 @@ const defaultHandler = async (internalEvent: InternalEvent) => {
       isISRRevalidation,
     },
     async () => {
+      provideNextAfterProvider();
       const result = await routingHandler(internalEvent);
       if ("internalEvent" in result) {
         debug("Middleware intercepted event", internalEvent);

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

@@ -3,7 +3,10 @@ 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 {
+  awaitAllDetachedPromise,
+  provideNextAfterProvider,
+} from "utils/promise";
 
 import { debug, error, warn } from "../adapters/logger";
 import { generateOpenNextRequestContext } from "../adapters/util";
@@ -40,6 +43,7 @@ export async function openNextHandler(
         internalEvent.headers.host = internalEvent.headers["x-forwarded-host"];
       }
       debug("internalEvent", internalEvent);
+      provideNextAfterProvider();
 
       let preprocessResult: InternalResult | MiddlewareOutputEvent = {
         internalEvent: internalEvent,

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

@@ -69,3 +69,34 @@ export async function awaitAllDetachedPromise() {
   }
   await promisesToAwait;
 }
+
+export function provideNextAfterProvider() {
+  /** This should be considered unstable until `unstable_after` is stablized. */
+  const NEXT_REQUEST_CONTEXT_SYMBOL = Symbol.for("@next/request-context");
+
+  // This is needed by some lib that relies on the vercel request context to properly await stuff.
+  // Remove this when vercel builder is updated to provide '@next/request-context'.
+  const VERCEL_REQUEST_CONTEXT_SYMBOL = Symbol.for("@vercel/request-context");
+
+  const openNextStoreContext = globalThis.__als.getStore();
+
+  const awaiter =
+    globalThis.openNextWaitUntil ??
+    ((promise: Promise<unknown>) =>
+      openNextStoreContext?.pendingPromiseRunner.add(promise));
+
+  const nextAfterContext = {
+    get: () => ({
+      waitUntil: awaiter,
+    }),
+  };
+
+  //@ts-expect-error
+  globalThis[NEXT_REQUEST_CONTEXT_SYMBOL] = nextAfterContext;
+  // We probably want to avoid providing this everytime since some lib may incorrectly think they are running in Vercel
+  // It may break stuff, but at the same time it will allow libs like `@vercel/otel` to work as expected
+  if (process.env.EMULATE_VERCEL_REQUEST_CONTEXT) {
+    //@ts-expect-error
+    globalThis[VERCEL_REQUEST_CONTEXT_SYMBOL] = nextAfterContext;
+  }
+}