add requestStorage

.changeset/flags-next-request-storage.md

@@ -0,0 +1,31 @@
+---
+'flags': minor
+---
+
+Add `requestStorage` and `attach` to `flags/next` for passing the request
+into flag evaluation via `AsyncLocalStorage`. Works with both App Router
+Route Handlers (Web `Request`) and Pages Router API Routes (Node
+`IncomingMessage`).
+
+When a request is set on `requestStorage` (directly via `requestStorage.run`
+or implicitly through the `attach()` wrapper), `flag()` reads headers and
+cookies from the request and skips the dynamic `import('next/headers')` and
+`headers()` / `cookies()` calls entirely.
+
+```ts
+// App Router Route Handler
+import { attach } from 'flags/next';
+
+export const GET = attach(async (request) => {
+  const value = await someFlag();
+  return Response.json({ value });
+});
+
+// Pages Router API Route
+import { attach } from 'flags/next';
+
+export default attach(async (req, res) => {
+  const value = await someFlag();
+  res.json({ value });
+});
+```

packages/flags/src/next/index.test.ts

@@ -2,9 +2,24 @@ import { IncomingMessage } from 'node:http';
 import type { Socket } from 'node:net';
 import { Readable } from 'node:stream';
 import type { NextApiRequestCookies } from 'next/dist/server/api-utils';
-import { afterEach, beforeAll, describe, expect, it, vi } from 'vitest';
+import {
+  afterEach,
+  beforeAll,
+  beforeEach,
+  describe,
+  expect,
+  it,
+  vi,
+} from 'vitest';
 import { type Adapter, encryptOverrides } from '..';
-import { clearDedupeCacheForCurrentRequest, dedupe, flag, precompute } from '.';
+import {
+  attach,
+  clearDedupeCacheForCurrentRequest,
+  dedupe,
+  flag,
+  precompute,
+  requestStorage,
+} from '.';
 
 const mocks = vi.hoisted(() => {
   return {
@@ -84,6 +99,13 @@ describe('exports', () => {
   it('should export clearDedupeCacheForCurrentRequest', () => {
     expect(typeof clearDedupeCacheForCurrentRequest).toBe('function');
   });
+  it('should export attach', () => {
+    expect(typeof attach).toBe('function');
+  });
+  it('should export requestStorage', () => {
+    expect(typeof requestStorage.run).toBe('function');
+    expect(typeof requestStorage.getStore).toBe('function');
+  });
 });
 
 describe('flag on app router', () => {
@@ -668,6 +690,193 @@ describe('flag on pages router', () => {
   });
 });
 
+describe('flag with requestStorage', () => {
+  beforeAll(() => {
+    // a random secret for testing purposes
+    process.env.FLAGS_SECRET = 'yuhyxaVI0Zue85SguKlMIUQojvJyBPzm95fFYvOa4Rc';
+  });
+
+  beforeEach(() => {
+    mocks.headers.mockClear();
+    mocks.cookies.mockClear();
+  });
+
+  it('reads headers from the stored Request and never calls next/headers', async () => {
+    const decide = vi.fn(
+      ({ headers }: { headers: Headers }) => headers.get('x-test') === 'on',
+    );
+    const f = flag<boolean>({ key: 'first-flag', decide });
+
+    const request = new Request('http://localhost/test', {
+      headers: { 'x-test': 'on' },
+    });
+
+    await requestStorage.run(request, async () => {
+      await expect(f()).resolves.toEqual(true);
+    });
+
+    expect(decide).toHaveBeenCalledTimes(1);
+    expect(mocks.headers).not.toHaveBeenCalled();
+    expect(mocks.cookies).not.toHaveBeenCalled();
+  });
+
+  it('attach() forwards the request and additional args', async () => {
+    const decide = vi.fn(
+      ({ headers }: { headers: Headers }) => headers.get('x-test') === 'on',
+    );
+    const f = flag<boolean>({ key: 'first-flag', decide });
+
+    const handler = attach(async (request, ctx: { params: { id: string } }) => {
+      const value = await f();
+      return { value, id: ctx.params.id, url: request.url };
+    });
+
+    const request = new Request('http://localhost/test', {
+      headers: { 'x-test': 'on' },
+    });
+
+    await expect(handler(request, { params: { id: '42' } })).resolves.toEqual({
+      value: true,
+      id: '42',
+      url: 'http://localhost/test',
+    });
+
+    expect(mocks.headers).not.toHaveBeenCalled();
+    expect(mocks.cookies).not.toHaveBeenCalled();
+  });
+
+  it('caches across multiple flag calls within the same scope', async () => {
+    let i = 0;
+    const decide = vi.fn(() => i++);
+    const f = flag<number>({ key: 'first-flag', decide });
+
+    const request = new Request('http://localhost/test');
+
+    await requestStorage.run(request, async () => {
+      await expect(f()).resolves.toEqual(0);
+      await expect(f()).resolves.toEqual(0);
+    });
+
+    expect(decide).toHaveBeenCalledTimes(1);
+
+    // a second request gets a fresh evaluation
+    await requestStorage.run(new Request('http://localhost/test'), async () => {
+      await expect(f()).resolves.toEqual(1);
+    });
+
+    expect(decide).toHaveBeenCalledTimes(2);
+  });
+
+  it('respects overrides parsed from the request cookie header', async () => {
+    const decide = vi.fn(() => false);
+    const f = flag<boolean>({ key: 'first-flag', decide });
+    const override = await encryptOverrides({ 'first-flag': true });
+
+    const request = new Request('http://localhost/test', {
+      headers: { cookie: `vercel-flag-overrides=${override}` },
+    });
+
+    await requestStorage.run(request, async () => {
+      await expect(f()).resolves.toEqual(true);
+    });
+
+    expect(decide).not.toHaveBeenCalled();
+    expect(mocks.cookies).not.toHaveBeenCalled();
+  });
+
+  it('passes sealed headers and cookies to identify', async () => {
+    const identify = vi.fn(({ headers, cookies }) => ({
+      user: headers.get('x-user'),
+      session: cookies.get('session')?.value,
+    }));
+    const decide = vi.fn(
+      ({ entities }: { entities?: { user: string | null } }) =>
+        entities?.user === 'alice',
+    );
+    const f = flag<boolean, { user: string | null; session: string }>({
+      key: 'first-flag',
+      identify,
+      decide,
+    });
+
+    const request = new Request('http://localhost/test', {
+      headers: { 'x-user': 'alice', cookie: 'session=abc' },
+    });
+
+    await requestStorage.run(request, async () => {
+      await expect(f()).resolves.toEqual(true);
+      // second call hits the per-request decide cache
+      await expect(f()).resolves.toEqual(true);
+    });
+
+    expect(decide).toHaveBeenCalledTimes(1);
+    expect(identify).toHaveBeenCalled();
+    const [params] = identify.mock.calls[0] ?? [];
+    expect(params.headers.get('x-user')).toBe('alice');
+    expect(params.cookies.get('session')?.value).toBe('abc');
+  });
+
+  it('reads headers from a stored IncomingMessage (Pages Router)', async () => {
+    const decide = vi.fn(
+      ({ headers }: { headers: Headers }) => headers.get('x-test') === 'on',
+    );
+    const f = flag<boolean>({ key: 'first-flag', decide });
+
+    const [request, socket] = createRequest();
+    request.headers['x-test'] = 'on';
+
+    await requestStorage.run(request, async () => {
+      await expect(f()).resolves.toEqual(true);
+    });
+
+    expect(decide).toHaveBeenCalledTimes(1);
+    expect(mocks.headers).not.toHaveBeenCalled();
+    expect(mocks.cookies).not.toHaveBeenCalled();
+    socket.destroy();
+  });
+
+  it('attach() works with a Pages Router IncomingMessage and forwards res', async () => {
+    const decide = vi.fn(
+      ({ headers }: { headers: Headers }) => headers.get('x-test') === 'on',
+    );
+    const f = flag<boolean>({ key: 'first-flag', decide });
+
+    const [request, socket] = createRequest();
+    request.headers['x-test'] = 'on';
+
+    const res = { json: vi.fn() };
+    const handler = attach(
+      async (req: typeof request, response: typeof res) => {
+        const value = await f();
+        response.json({ value });
+        return value;
+      },
+    );
+
+    await expect(handler(request, res)).resolves.toEqual(true);
+    expect(res.json).toHaveBeenCalledWith({ value: true });
+    expect(mocks.headers).not.toHaveBeenCalled();
+    socket.destroy();
+  });
+
+  it('respects overrides parsed from the cookie header of a stored IncomingMessage', async () => {
+    const decide = vi.fn(() => false);
+    const f = flag<boolean>({ key: 'first-flag', decide });
+    const override = await encryptOverrides({ 'first-flag': true });
+
+    const [request, socket] = createRequest({
+      'vercel-flag-overrides': override,
+    });
+
+    await requestStorage.run(request, async () => {
+      await expect(f()).resolves.toEqual(true);
+    });
+
+    expect(decide).not.toHaveBeenCalled();
+    socket.destroy();
+  });
+});
+
 describe('dynamic io', () => {
   it('should re-throw dynamic usage erorrs even when a defaultValue is present', async () => {
     const mockDecide = vi.fn(() => {

packages/flags/src/next/index.ts

@@ -1,4 +1,5 @@
-import type { IncomingHttpHeaders } from 'node:http';
+import { AsyncLocalStorage } from 'node:async_hooks';
+import type { IncomingHttpHeaders, IncomingMessage } from 'node:http';
 import { RequestCookies } from '@edge-runtime/cookies';
 import {
   type FlagDefinitionsType,
@@ -41,6 +42,57 @@ export {
 } from './precompute';
 export type { Flag } from './types';
 
+/**
+ * Either a Web `Request` (App Router Route Handler) or a Node
+ * `IncomingMessage` (Pages Router API Route).
+ */
+type StoredRequest = Request | IncomingMessage;
+
+/**
+ * AsyncLocalStorage holding the current request. When a request is set here,
+ * `flag()` reads headers and cookies from it instead of importing and calling
+ * `next/headers`. Accepts either a Web `Request` (App Router Route Handler)
+ * or a Node `IncomingMessage` (Pages Router API Route).
+ */
+export const requestStorage = new AsyncLocalStorage<StoredRequest>();
+
+/**
+ * Wraps a Route Handler or API Route so its execution runs inside
+ * `requestStorage`. Accepts either a Web `Request` (App Router) or a Node
+ * `IncomingMessage` (Pages Router) as the first argument; remaining arguments
+ * (e.g. `{ params }` in App Router, `res` in Pages Router) are forwarded.
+ *
+ * @example App Router Route Handler
+ * ```ts
+ * import { attach } from 'flags/next';
+ *
+ * export const GET = attach(async (request) => {
+ *   const value = await someFlag();
+ *   return Response.json({ value });
+ * });
+ * ```
+ *
+ * @example Pages Router API Route
+ * ```ts
+ * import { attach } from 'flags/next';
+ *
+ * export default attach(async (req, res) => {
+ *   const value = await someFlag();
+ *   res.json({ value });
+ * });
+ * ```
+ */
+export function attach<
+  TRequest extends StoredRequest,
+  TArgs extends unknown[],
+  TReturn,
+>(
+  handler: (request: TRequest, ...args: TArgs) => TReturn | Promise<TReturn>,
+): (request: TRequest, ...args: TArgs) => Promise<TReturn> {
+  return async (request, ...args) =>
+    requestStorage.run(request, async () => handler(request, ...args));
+}
+
 // a map of (headers, flagKey, entitiesKey) => value
 const evaluationCache = new WeakMap<
   Headers | IncomingHttpHeaders,
@@ -231,12 +283,21 @@ function getRun<ValueType, EntitiesType>(
     let readonlyCookies: ReadonlyRequestCookies;
     let dedupeCacheKey: Headers | IncomingHttpHeaders;
 
-    if (options.request) {
-      // pages router
-      const headers = transformToHeaders(options.request.headers);
-      readonlyHeaders = sealHeaders(headers);
-      readonlyCookies = sealCookies(headers);
-      dedupeCacheKey = options.request.headers;
+    const storedRequest = options.request ?? requestStorage.getStore();
+    if (storedRequest) {
+      if (storedRequest instanceof Request) {
+        // app router route handler with a Web Request via AsyncLocalStorage
+        readonlyHeaders = sealHeaders(storedRequest.headers);
+        readonlyCookies = sealCookies(storedRequest.headers);
+        dedupeCacheKey = storedRequest.headers;
+      } else {
+        // pages router api route — either passed explicitly to the flag or
+        // set on requestStorage
+        const headers = transformToHeaders(storedRequest.headers);
+        readonlyHeaders = sealHeaders(headers);
+        readonlyCookies = sealCookies(headers);
+        dedupeCacheKey = storedRequest.headers;
+      }
     } else {
       // app router