accept direct waitUntil function

Author: s1gr1d

packages/core/src/utils/flushIfServerless.ts

@@ -25,7 +25,8 @@ async function flushWithTimeout(timeout: number): Promise<void> {
  * The function is async, but in environments that support a `waitUntil` mechanism, it will run synchronously.
  *
  * This function is aware of the following serverless platforms:
- * - Cloudflare: If a Cloudflare context is provided, it will use `ctx.waitUntil()` to flush events.
+ * - Cloudflare: If a Cloudflare context is provided, it will use `ctx.waitUntil()` to flush events (keeps the `this` context of `ctx`).
+ *               If a `cloudflareWaitUntil` function is provided, it will use that to flush events (looses the `this` context of `ctx`).
  * - Vercel: It detects the Vercel environment and uses Vercel's `waitUntil` function.
  * - Other Serverless (AWS Lambda, Google Cloud, etc.): It detects the environment via environment variables
  *   and uses a regular `await flush()`.
@@ -34,15 +35,19 @@ async function flushWithTimeout(timeout: number): Promise<void> {
  *  @hidden
  */
 export async function flushIfServerless(
-  params: {
-    timeout?: number;
-    cloudflareCtx?: MinimalCloudflareContext;
-  } = {},
+  params: // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  | { timeout?: number; cloudflareWaitUntil?: (task: Promise<any>) => void }
+    | { timeout?: number; cloudflareCtx?: MinimalCloudflareContext },
 ): Promise<void> {
-  const { timeout = 2000, cloudflareCtx } = params;
+  const { timeout = 2000 } = params;
 
-  if (cloudflareCtx && typeof cloudflareCtx.waitUntil === 'function') {
-    cloudflareCtx.waitUntil(flushWithTimeout(timeout));
+  if ('cloudflareWaitUntil' in params && typeof params?.cloudflareWaitUntil === 'function') {
+    params.cloudflareWaitUntil(flushWithTimeout(timeout));
+    return;
+  }
+
+  if ('cloudflareCtx' in params && typeof params.cloudflareCtx?.waitUntil === 'function') {
+    params.cloudflareCtx.waitUntil(flushWithTimeout(timeout));
     return;
   }
 

packages/core/test/lib/utils/flushIfServerless.test.ts

@@ -16,7 +16,7 @@ describe('flushIfServerless', () => {
     vi.restoreAllMocks();
   });
 
-  test('should bind context (preserve `this`) when calling waitUntil', async () => {
+  test('should bind context (preserve `this`) when calling waitUntil from the Cloudflare execution context', async () => {
     const flushMock = vi.spyOn(flushModule, 'flush').mockResolvedValue(true);
 
     // Mock Cloudflare context with `waitUntil` (which should be called if `this` is bound correctly)
@@ -49,6 +49,18 @@ describe('flushIfServerless', () => {
     expect(flushMock).toHaveBeenCalledWith(5000);
   });
 
+  test('should use cloudflare waitUntil when Cloudflare `waitUntil` is provided', async () => {
+    const flushMock = vi.spyOn(flushModule, 'flush').mockResolvedValue(true);
+    const mockCloudflareCtx = {
+      waitUntil: vi.fn(),
+    };
+
+    await flushIfServerless({ cloudflareWaitUntil: mockCloudflareCtx.waitUntil, timeout: 5000 });
+
+    expect(mockCloudflareCtx.waitUntil).toHaveBeenCalledTimes(1);
+    expect(flushMock).toHaveBeenCalledWith(5000);
+  });
+
   test('should ignore cloudflare context when waitUntil is not a function (and use Vercel waitUntil instead)', async () => {
     const flushMock = vi.spyOn(flushModule, 'flush').mockResolvedValue(true);
     const vercelWaitUntilSpy = vi.spyOn(vercelWaitUntilModule, 'vercelWaitUntil').mockImplementation(() => {});