diff --git a/.changeset/flags-next-request-storage.md b/.changeset/flags-next-request-storage.md new file mode 100644 index 00000000..1393448f --- /dev/null +++ b/.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 }); +}); +``` diff --git a/packages/flags/src/next/index.test.ts b/packages/flags/src/next/index.test.ts index 7d3701db..0350d441 100644 --- a/packages/flags/src/next/index.test.ts +++ b/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({ 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({ 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({ 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({ 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({ + 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({ 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({ 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({ 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(() => { diff --git a/packages/flags/src/next/index.ts b/packages/flags/src/next/index.ts index 57912a93..2d35b149 100644 --- a/packages/flags/src/next/index.ts +++ b/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(); + +/** + * 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, +): (request: TRequest, ...args: TArgs) => Promise { + 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( 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