[@vercel/flags-core] Add custom waitUntil option

.changeset/quiet-flags-wait.md

@@ -0,0 +1,8 @@
+---
+'@vercel/flags-core': patch
+---
+
+Allow passing a custom `waitUntil` function to `createClient` for background
+usage and exposure reporting. Pending exposure reports are drained by
+`client.shutdown()`. The Next.js conditional export uses `after` from
+`next/server` by default.

packages/vercel-flags-core/CLAUDE.md

@@ -90,6 +90,7 @@ type ControllerOptions = {
   polling?: boolean | { intervalMs: number; initTimeoutMs: number };  // default: true (30s interval, 3s timeout)
   buildStep?: boolean;  // Override build step auto-detection
   metricEnvironment?: string; // Environment attached to ingested evaluation metrics
+  waitUntil?: (promise: Promise<unknown>) => void;  // default: @vercel/functions waitUntil
   sources?: { stream?: StreamSource; polling?: PollingSource; bundled?: BundledSource };  // DI for testing
 };
 ```
@@ -267,7 +268,9 @@ The Controller tags all data with its origin using `tagData(data, origin)` from
 - Sends to `flags.vercel.com/v1/ingest`
 - At runtime: deduplicates by request context (per-instance WeakSet in UsageTracker)
 - During builds: deduplicates all reads to a single event (buildReadTracked flag in Controller), since there is no request context available
-- Uses `waitUntil()` from `@vercel/functions` (wrapped in try/catch for resilience)
+- Uses a custom `waitUntil()` passed to `createClient`, or defaults to `@vercel/functions` (wrapped in try/catch for resilience)
+- The Next.js conditional export defaults to `after()` from `next/server`; an explicit `waitUntil` option always takes precedence
+- Exposure reporting does not block evaluation; `shutdown()` drains pending exposure reports for graceful shutdown in long-lived processes
 - On flush failure, events are re-queued for retry with a max queue size of 500 events (oldest events are dropped when exceeded)
 - `flush()` directly flushes queued events even when no scheduled flush is pending, ensuring events are not lost during `shutdown()`
 

packages/vercel-flags-core/README.md

@@ -33,18 +33,6 @@ export default app;
 
 Outside Vercel, pass an SDK key explicitly: `createClient(process.env.FLAGS)`.
 
-### Initialization and request-scoped OIDC
-
-On Vercel, the OIDC token can be supplied through the current request context (`x-vercel-oidc-token`). It is not guaranteed to be available while modules are loading. Creating the client at module scope is safe, but starting `client.initialize()` there can attempt authentication before a request exists.
-
-Do not store a module-scope initialization promise and await it later in a handler. Calling `initialize()` starts the work immediately; awaiting the promise inside a request does not move that work into the request context. If the promise rejects, every handler awaiting that same promise will reject before reaching evaluation and its fallback handling.
-
-This also applies when definitions are embedded at build time: with OIDC authentication, the client uses the token's `project_id` claim to select the embedded definitions. An importable `@vercel/flags-definitions` module alone is not enough.
-
-For local development, `vercel env pull` writes an OIDC token to `.env.local`. If your local runner loads that file, the environment-variable fallback can make module-scope initialization appear to work, even though request-scoped authentication on a deployment requires deferring it.
-
-Explicit `await client.initialize()` is optional and can be useful when you want to wait for initialization and handle its errors yourself. Only call it once authentication is available: inside a request handler for request-scoped OIDC, or during startup when using an SDK key or an already-available environment token. For normal flag evaluation, prefer calling `evaluate()` or `bulkEvaluate()` directly in the handler.
-
 ## Evaluation Metrics
 
 To associate evaluation metrics with an environment, pass the

packages/vercel-flags-core/src/black-box.test.ts

@@ -3836,7 +3836,7 @@ describe('Controller (black-box)', () => {
       await client.shutdown();
     });
 
-    it('does not block evaluation while reporting an exposure', async () => {
+    it('does not block evaluation and drains exposure reporting on shutdown', async () => {
       let finishReporting: () => void = () => {};
       const reporting = new Promise<void>((resolve) => {
         finishReporting = resolve;
@@ -3856,8 +3856,41 @@ describe('Controller (black-box)', () => {
       ).resolves.toMatchObject({ value: 'treatment-a' });
       expect(reportExposures).toHaveBeenCalledOnce();
 
+      let shutdownComplete = false;
+      const shutdown = Promise.resolve(client.shutdown()).then(() => {
+        shutdownComplete = true;
+      });
+      await Promise.resolve();
+      expect(shutdownComplete).toBe(false);
+
+      finishReporting();
+      await shutdown;
+      expect(shutdownComplete).toBe(true);
+    });
+
+    it('registers exposure reporting with a custom waitUntil', async () => {
+      let finishReporting: () => void = () => {};
+      const reporting = new Promise<void>((resolve) => {
+        finishReporting = resolve;
+      });
+      const waitUntil = vi.fn<(promise: Promise<unknown>) => void>();
+      const client = createClient<typeof entity>(sdkKey, {
+        fetch: fetchMock,
+        stream: false,
+        polling: false,
+        buildStep: true,
+        datafile: makeBundled({ definitions }),
+        experimental_reportExposures: () => reporting,
+        waitUntil,
+      });
+
+      await client.evaluate('flagA', undefined, entity);
+
+      expect(waitUntil).toHaveBeenCalledTimes(2);
+      expect(waitUntil).toHaveBeenNthCalledWith(1, expect.any(Promise));
+      expect(waitUntil).toHaveBeenNthCalledWith(2, expect.any(Promise));
+
       finishReporting();
-      await reporting;
       await client.shutdown();
     });
 
@@ -4060,6 +4093,26 @@ describe('Controller (black-box)', () => {
   // Usage tracking
   // ---------------------------------------------------------------------------
   describe('usage tracking', () => {
+    it('should use a custom waitUntil function for usage tracking', async () => {
+      const cleanupCtx = setRequestContext({ host: 'example.com' });
+      const waitUntil = vi.fn<(promise: Promise<unknown>) => void>();
+      const client = createClient(sdkKey, {
+        datafile: makeBundled(),
+        fetch: fetchMock,
+        polling: false,
+        stream: false,
+        waitUntil,
+      });
+
+      await client.evaluate('flagA');
+
+      expect(waitUntil).toHaveBeenCalledOnce();
+      expect(waitUntil).toHaveBeenCalledWith(expect.any(Promise));
+
+      await client.shutdown();
+      cleanupCtx();
+    });
+
     it('should report counted FLAG_EVALUATION events', async () => {
       const cleanupCtx = setRequestContext({ host: 'example.com' });
       fetchMock.mockImplementation((input) => {

packages/vercel-flags-core/src/controller/normalized-options.ts

@@ -1,8 +1,10 @@
+import { waitUntil as defaultWaitUntil } from '@vercel/functions';
 import type {
   DatafileInput,
   MetricEnvironment,
   PollingOptions,
   StreamOptions,
+  WaitUntil,
 } from '../types';
 import type { Auth } from './auth';
 
@@ -58,6 +60,13 @@ export type ControllerOptions = {
    */
   fetch?: typeof globalThis.fetch;
 
+  /**
+   * Custom function for keeping background work alive after a response has
+   * been sent.
+   * @default waitUntil from `@vercel/functions`
+   */
+  waitUntil?: WaitUntil;
+
   /**
    * Environment included with evaluation metrics sent to the ingest endpoint.
    * Falls back to the `VERCEL_ENV` environment variable when not set.
@@ -84,6 +93,7 @@ export type NormalizedOptions = {
   polling: { enabled: boolean; intervalMs: number; initTimeoutMs: number };
   buildStep: boolean;
   fetch: typeof globalThis.fetch;
+  waitUntil: WaitUntil;
   host: string;
   metricEnvironment: MetricEnvironment | undefined;
   clientName: string | undefined;
@@ -136,6 +146,7 @@ export function normalizeOptions(
     polling,
     buildStep,
     fetch: options.fetch ?? globalThis.fetch,
+    waitUntil: options.waitUntil ?? defaultWaitUntil,
     host: 'https://flags.vercel.com',
     metricEnvironment: options.metricEnvironment,
     clientName: options.clientName,

packages/vercel-flags-core/src/create-raw-client.ts

@@ -1,4 +1,4 @@
-import { waitUntil } from '@vercel/functions';
+import { waitUntil as defaultWaitUntil } from '@vercel/functions';
 import { dequal } from 'dequal/lite';
 import type {
   bulkEvaluate,
@@ -23,6 +23,7 @@ import type {
   FlagsClient,
   Packed,
   Value,
+  WaitUntil,
 } from './types';
 
 let idCount = 0;
@@ -53,12 +54,15 @@ export function createCreateRawClient(fns: {
     controller,
     origin,
     experimental_reportExposures,
+    waitUntil = defaultWaitUntil,
   }: {
     controller: ControllerInterface;
     origin?: { provider: string; sdkKey?: string };
     experimental_reportExposures?: experimental_ReportExposures<Entities>;
+    waitUntil?: WaitUntil;
   }): FlagsClient<Entities> {
     const id = idCount++;
+    const pendingExposureReports = new Set<Promise<void>>();
     controllerInstanceMap.set(id, {
       controller,
       initialized: false,
@@ -81,6 +85,8 @@ export function createCreateRawClient(fns: {
           );
         }
       })();
+      pendingExposureReports.add(pending);
+      void pending.finally(() => pendingExposureReports.delete(pending));
 
       try {
         waitUntil(pending);
@@ -131,6 +137,9 @@ export function createCreateRawClient(fns: {
       },
       shutdown: async () => {
         await fns.shutdown(id);
+        while (pendingExposureReports.size > 0) {
+          await Promise.all(pendingExposureReports);
+        }
         controllerInstanceMap.delete(id);
       },
       getDatafile: async () => {

packages/vercel-flags-core/src/index.common.ts

@@ -28,4 +28,5 @@ export {
   ResolutionReason as Reason,
   type StreamOptions,
   type Value,
+  type WaitUntil,
 } from './types';

packages/vercel-flags-core/src/index.default.ts

@@ -10,6 +10,7 @@
  * We do not need to repeat the JSDoc on the next-js export.
  */
 
+import { waitUntil } from '@vercel/functions';
 import * as fns from './controller-fns';
 import { createCreateRawClient } from './create-raw-client';
 import { make } from './index.make';
@@ -32,4 +33,4 @@ export const {
    * Create a flags client using an SDK key, connection string, or Vercel OIDC.
    */
   createClient,
-} = make(createCreateRawClient(fns));
+} = make(createCreateRawClient(fns), { waitUntil });