`waitUntil` and `unstable_after` after support for @cloudflare/next-on-pages

wait-until-after.tsx

import { getOptionalRequestContext } from "@cloudflare/next-on-pages";
import { workAsyncStorage } from "next/dist/server/app-render/work-async-storage.external";

export type AfterTask<T = unknown> = Promise<T> | AfterCallback<T>;
export type AfterCallback<T = unknown> = () => T | Promise<T>;

export function waitUntil(promise: Promise<unknown>) {
  const cloudflareCtx = getOptionalRequestContext();

  if (!cloudflareCtx) {
    throw new Error("`waitUntil` was called outside a request");
  }

  cloudflareCtx.waitUntil(promise);
}

/**
 * This function allows you to schedule callbacks to be executed after the current request finishes.
 */
export function unstable_after<T>(task: AfterTask<T>): void {
  const workStore = workAsyncStorage.getStore();

  if (!workStore) {
    // TODO(after): the linked docs page talks about *dynamic* APIs, which unstable_after soon won't be anymore
    throw new Error(
      "`unstable_after` was called outside a request scope. Read more: https://nextjs.org/docs/messages/next-dynamic-api-wrong-context"
    );
  }

  const { afterContext } = workStore;
  if (!afterContext) {
    throw new Error(
      "`unstable_after` must be explicitly enabled by setting `experimental.after: true` in your next.config.js."
    );
  }

  // Inject waitUntil if not already present

  const hasPrivateWaitUntil = Object.prototype.hasOwnProperty.call(
    afterContext,
    "waitUntil"
  );

  // @ts-expect-error: TS doesn't know about the private `waitUntil` property
  const isWaitUntilNotUndefined = typeof afterContext.waitUntil !== "undefined";

  const hasWaitUntil = hasPrivateWaitUntil && isWaitUntilNotUndefined;

  if (!hasWaitUntil) {
    console.warn(
      "`waitUntil` is not available in the current context, injecting from `~/lib/waitUntil` (cloudflare)"
    );
    Object.defineProperty(afterContext, "waitUntil", {
      value: waitUntil,
      writable: false,
      enumerable: true,
      configurable: true,
    });
  }

  afterContext.after(task);
}

Good idea but workAsyncStorage may change in the future.

@micooz yeah probably, guess we'll notice if it does. If you have any suggestions for other ways to inject waitUntil here would be glad to know.