From 1f6b0e63b568077798d2aa625d9a2a3310a87d38 Mon Sep 17 00:00:00 2001 From: "Benjamin O. Ajayi" Date: Sun, 30 Aug 2026 20:10:48 +0100 Subject: [PATCH] feat(context): add structured request context storage using AsyncLocalStorage Implement a structured, typed request context store backed by AsyncLocalStorage (RequestContext) that captures routing metadata, resolved principal and request-scoped values for the lifetime of each HTTP request. - Add RequestContext store with typed identity/principal getters, scoped data storage and timing recording, propagated through async continuations - Add RequestContextInterceptor that seeds the context on every request, resolving identity from JWT auth or implicit org/agent headers - Register the interceptor globally ahead of the existing agent trace interceptor - Add unit tests for the store and interceptor Closes #80 --- src/app.module.ts | 2 + src/common/context/request-context.spec.ts | 117 ++++++++++++ src/common/context/request-context.ts | 163 +++++++++++++++++ .../request-context.interceptor.spec.ts | 166 ++++++++++++++++++ .../request-context.interceptor.ts | 121 +++++++++++++ 5 files changed, 569 insertions(+) create mode 100644 src/common/context/request-context.spec.ts create mode 100644 src/common/context/request-context.ts create mode 100644 src/common/interceptors/request-context.interceptor.spec.ts create mode 100644 src/common/interceptors/request-context.interceptor.ts diff --git a/src/app.module.ts b/src/app.module.ts index 1b20874..56a4eb4 100644 --- a/src/app.module.ts +++ b/src/app.module.ts @@ -37,6 +37,7 @@ import { AuditModule } from './modules/audit/audit.module'; import { AiModule } from './modules/ai/ai.module'; import { HealthModule } from './modules/health/health.module'; import { AgentTraceInterceptor } from './common/interceptors/agent-trace.interceptor'; +import { RequestContextInterceptor } from './common/interceptors/request-context.interceptor'; /** * Root application module. Wires the global infrastructure (config, logging, @@ -115,6 +116,7 @@ import { AgentTraceInterceptor } from './common/interceptors/agent-trace.interce { provide: APP_GUARD, useClass: JwtAuthGuard }, { provide: APP_GUARD, useClass: RolesGuard }, { provide: APP_GUARD, useClass: AstroidThrottlerGuard }, + { provide: APP_INTERCEPTOR, useClass: RequestContextInterceptor }, { provide: APP_INTERCEPTOR, useClass: AgentTraceInterceptor }, { provide: APP_INTERCEPTOR, useClass: ResponseInterceptor }, { provide: APP_INTERCEPTOR, useClass: AuditInterceptor }, diff --git a/src/common/context/request-context.spec.ts b/src/common/context/request-context.spec.ts new file mode 100644 index 0000000..a692e43 --- /dev/null +++ b/src/common/context/request-context.spec.ts @@ -0,0 +1,117 @@ +import { describe, it, expect } from 'vitest'; +import { RequestContext, RequestContextData } from './request-context'; + +describe('RequestContext', () => { + const baseContext = (): RequestContextData => ({ + identity: { + requestId: 'req_123', + correlationId: 'corr_123', + traceId: 'trace_123', + method: 'POST', + path: '/api/v1/agents/agent-1/execute', + url: '/api/v1/agents/agent-1/execute', + ip: '127.0.0.1', + userAgent: 'test-agent', + startedAt: 1000, + }, + principal: { userId: 'user-1', organizationId: 'org-1', agentId: 'agent-1' }, + timings: {}, + data: {}, + }); + + it('should expose the store within a run() boundary and nothing outside', () => { + RequestContext.run(baseContext(), () => { + expect(RequestContext.getStore()).toEqual(baseContext()); + expect(RequestContext.get()).toEqual(baseContext()); + }); + + expect(RequestContext.getStore()).toBeUndefined(); + }); + + it('should propagate the context across async continuations', async () => { + let captured: RequestContextData | undefined; + + await new Promise((resolve) => { + RequestContext.run(baseContext(), () => { + setTimeout(() => { + captured = RequestContext.getStore(); + resolve(); + }, 5); + }); + }); + + expect(captured?.identity.requestId).toBe('req_123'); + }); + + it('should expose typed identity getters', () => { + RequestContext.run(baseContext(), () => { + expect(RequestContext.getRequestId()).toBe('req_123'); + expect(RequestContext.getCorrelationId()).toBe('corr_123'); + expect(RequestContext.getTraceId()).toBe('trace_123'); + expect(RequestContext.getMethod()).toBe('POST'); + expect(RequestContext.getPath()).toBe('/api/v1/agents/agent-1/execute'); + expect(RequestContext.getStartedAt()).toBe(1000); + }); + }); + + it('should expose typed principal getters', () => { + RequestContext.run(baseContext(), () => { + expect(RequestContext.getPrincipal()).toEqual({ + userId: 'user-1', + organizationId: 'org-1', + agentId: 'agent-1', + }); + expect(RequestContext.getUserId()).toBe('user-1'); + expect(RequestContext.getOrganizationId()).toBe('org-1'); + expect(RequestContext.getAgentId()).toBe('agent-1'); + }); + }); + + it('should return undefined getters when no context is active', () => { + expect(RequestContext.getRequestId()).toBeUndefined(); + expect(RequestContext.getUserId()).toBeUndefined(); + expect(RequestContext.getOrganizationId()).toBeUndefined(); + expect(RequestContext.getAgentId()).toBeUndefined(); + }); + + it('should merge a partial principal via setPrincipal', () => { + RequestContext.run(baseContext(), () => { + RequestContext.setPrincipal({ role: 'ADMIN' as const, authMethod: 'jwt' }); + expect(RequestContext.getPrincipal()).toEqual({ + userId: 'user-1', + organizationId: 'org-1', + agentId: 'agent-1', + role: 'ADMIN', + authMethod: 'jwt', + }); + }); + }); + + it('should store and read arbitrary request-scoped data', () => { + RequestContext.run(baseContext(), () => { + expect(RequestContext.getData('cart')).toBeUndefined(); + RequestContext.setData('cart', { count: 3 }); + expect(RequestContext.getData('cart')).toEqual({ count: 3 }); + }); + }); + + it('should record and read timings measured from request start', () => { + RequestContext.run(baseContext(), () => { + RequestContext.markTiming('db'); + const timing = RequestContext.getTiming('db'); + expect(typeof timing).toBe('number'); + expect(timing!).toBeGreaterThanOrEqual(0); + }); + }); + + it('should not throw when mutating with no active context', () => { + expect(() => { + RequestContext.setPrincipal({ userId: 'x' }); + RequestContext.setData('k', 'v'); + RequestContext.markTiming('nope'); + }).not.toThrow(); + + expect(RequestContext.getData('k')).toBeUndefined(); + expect(RequestContext.getTiming('nope')).toBeUndefined(); + }); +}); diff --git a/src/common/context/request-context.ts b/src/common/context/request-context.ts new file mode 100644 index 0000000..0dc34e0 --- /dev/null +++ b/src/common/context/request-context.ts @@ -0,0 +1,163 @@ +import { AsyncLocalStorage } from 'async_hooks'; + +/** + * Structured principal attached to a request once authentication resolves. + * Fields are optional so the same shape works for public, API-key and webhook + * traffic as well as fully authenticated JWT requests. + */ +export interface RequestPrincipal { + userId?: string; + organizationId?: string; + agentId?: string; + role?: string; + authMethod?: 'jwt' | 'api-key' | 'webhook' | 'service' | 'public'; +} + +/** + * Immutable request-level identifiers and routing metadata captured when the + * request enters the process. Never mutated after seeding. + */ +export interface RequestIdentity { + requestId: string; + correlationId: string; + traceId: string; + method: string; + path: string; + url: string; + ip: string | null; + userAgent: string | null; + startedAt: number; +} + +/** + * The full, structured object stored in AsyncLocalStorage for the lifetime of a + * single request. Grouped so consumers can read identity, routing metadata and + * arbitrary request-scoped values without reaching into raw headers. + */ +export interface RequestContextData { + identity: RequestIdentity; + principal?: RequestPrincipal; + timings: Record; + data: Record; +} + +/** + * Structured request context storage built on `AsyncLocalStorage`. + * + * A single store instance is entered once per HTTP request (see + * `RequestContextInterceptor`) so that any code running in the request's async + * flow — controllers, services, guards, filters, queue producers — can read a + * consistent, typed snapshot of the request without threading arguments + * through the call tree. + * + * The subtle-but-powerful property of `AsyncLocalStorage` is that the context + * is inherited by every async continuation spawned inside `run()`, which makes + * it the correct primitive for tracing and request-scoped state in Node. + */ +export class RequestContext { + private static readonly storage = new AsyncLocalStorage(); + + /** + * Enters a request context for the duration of `fn`, propagating it through + * every asynchronous operation spawned inside. + */ + static run(data: RequestContextData, fn: () => R): R { + return this.storage.run(data, fn); + } + + /** Returns the raw structured store bound to the current async flow. */ + static getStore(): RequestContextData | undefined { + return this.storage.getStore(); + } + + /** Convenience alias for {@link getStore}. */ + static get(): RequestContextData | undefined { + return this.getStore(); + } + + static getRequestId(): string | undefined { + return this.storage.getStore()?.identity.requestId; + } + + static getCorrelationId(): string | undefined { + return this.storage.getStore()?.identity.correlationId; + } + + static getTraceId(): string | undefined { + return this.storage.getStore()?.identity.traceId; + } + + static getMethod(): string | undefined { + return this.storage.getStore()?.identity.method; + } + + static getPath(): string | undefined { + return this.storage.getStore()?.identity.path; + } + + static getStartedAt(): number | undefined { + return this.storage.getStore()?.identity.startedAt; + } + + static getPrincipal(): RequestPrincipal | undefined { + return this.storage.getStore()?.principal; + } + + static getUserId(): string | undefined { + return this.storage.getStore()?.principal?.userId; + } + + static getOrganizationId(): string | undefined { + return this.storage.getStore()?.principal?.organizationId; + } + + static getAgentId(): string | undefined { + return this.storage.getStore()?.principal?.agentId; + } + + /** + * Merges a partial principal into the current store. No-op when no context + * is active so code that runs outside a request never throws. + */ + static setPrincipal(principal: RequestPrincipal): void { + const store = this.storage.getStore(); + if (store) { + store.principal = { ...store.principal, ...principal }; + } + } + + /** + * Reads a request-scoped value previously stored with {@link setData}. + */ + static getData(key: string): T | undefined { + return this.storage.getStore()?.data[key] as T | undefined; + } + + /** + * Stores an arbitrary request-scoped value keyed by `name`. + */ + static setData(name: string, value: unknown): void { + const store = this.storage.getStore(); + if (store) { + store.data[name] = value; + } + } + + /** + * Records the elapsed time (ms) since the request started under `name`, + * then stores it in the context `timings` map for observability. + */ + static markTiming(name: string): void { + const store = this.storage.getStore(); + if (store) { + store.timings[name] = Date.now() - store.identity.startedAt; + } + } + + /** + * Reads a recorded timing value (ms) by name, if present. + */ + static getTiming(name: string): number | undefined { + return this.storage.getStore()?.timings[name]; + } +} diff --git a/src/common/interceptors/request-context.interceptor.spec.ts b/src/common/interceptors/request-context.interceptor.spec.ts new file mode 100644 index 0000000..c1f267c --- /dev/null +++ b/src/common/interceptors/request-context.interceptor.spec.ts @@ -0,0 +1,166 @@ +import { describe, it, expect } from 'vitest'; +import { ExecutionContext, CallHandler } from '@nestjs/common'; +import { of } from 'rxjs'; +import { RequestContextInterceptor } from './request-context.interceptor'; +import { RequestContext } from '../context/request-context'; + +describe('RequestContextInterceptor', () => { + const interceptor = new RequestContextInterceptor(); + + it('should seed structured context from an authenticated request', async () => { + let captured = { + requestId: undefined as string | undefined, + correlationId: undefined as string | undefined, + traceId: undefined as string | undefined, + userId: undefined as string | undefined, + organizationId: undefined as string | undefined, + method: undefined as string | undefined, + path: undefined as string | undefined, + ip: undefined as string | null | undefined, + role: undefined as string | undefined, + }; + + const mockExecutionContext = { + switchToHttp: () => ({ + getRequest: () => ({ + headers: { + 'x-request-id': 'req-abc', + 'x-correlation-id': 'corr-xyz', + 'x-forwarded-for': '203.0.113.9, 70.41.3.18', + 'user-agent': 'curl/8.0', + }, + params: {}, + body: {}, + query: {}, + user: { id: 'user-1', organizationId: 'org-1', role: 'OWNER' }, + method: 'POST', + path: '/api/v1/wallets', + originalUrl: '/api/v1/wallets', + url: '/api/v1/wallets', + socket: { remoteAddress: '127.0.0.1' }, + }), + }), + } as unknown as ExecutionContext; + + const mockCallHandler: CallHandler = { + handle: () => { + captured = { + requestId: RequestContext.getRequestId(), + correlationId: RequestContext.getCorrelationId(), + traceId: RequestContext.getTraceId(), + userId: RequestContext.getUserId(), + organizationId: RequestContext.getOrganizationId(), + method: RequestContext.getMethod(), + path: RequestContext.getPath(), + ip: RequestContext.getStore()?.identity.ip, + role: RequestContext.getPrincipal()?.role, + }; + return of({ success: true }); + }, + }; + + const observable = interceptor.intercept(mockExecutionContext, mockCallHandler); + + await new Promise((resolve, reject) => { + observable.subscribe({ + next: () => resolve(), + error: (err) => reject(err), + }); + }); + + expect(captured.requestId).toBe('req-abc'); + expect(captured.correlationId).toBe('corr-xyz'); + expect(captured.traceId).toBe('corr-xyz'); + expect(captured.userId).toBe('user-1'); + expect(captured.organizationId).toBe('org-1'); + expect(captured.role).toBe('OWNER'); + expect(captured.method).toBe('POST'); + expect(captured.path).toBe('/api/v1/wallets'); + expect(captured.ip).toBe('203.0.113.9'); + }); + + it('should generate a request id and default correlation/trace when none provided', async () => { + let capturedRequestId: string | undefined; + let capturedCorrelation: string | undefined; + + const mockExecutionContext = { + switchToHttp: () => ({ + getRequest: () => ({ + headers: {}, + params: {}, + body: {}, + query: {}, + method: 'GET', + path: '/api/v1/status', + originalUrl: '/api/v1/status', + url: '/api/v1/status', + socket: {}, + }), + }), + } as unknown as ExecutionContext; + + const mockCallHandler: CallHandler = { + handle: () => { + capturedRequestId = RequestContext.getRequestId(); + capturedCorrelation = RequestContext.getCorrelationId(); + return of({ ok: true }); + }, + }; + + const observable = interceptor.intercept(mockExecutionContext, mockCallHandler); + + await new Promise((resolve, reject) => { + observable.subscribe({ + next: () => resolve(), + error: (err) => reject(err), + }); + }); + + expect(capturedRequestId).toBeDefined(); + expect(capturedCorrelation).toBe(capturedRequestId); + }); + + it('should resolve an implicit organization/agent principal for service traffic', async () => { + let capturedOrg: string | undefined; + let capturedAgent: string | undefined; + let capturedAuthMethod: string | undefined; + + const mockExecutionContext = { + switchToHttp: () => ({ + getRequest: () => ({ + headers: { 'x-agent-id': 'agent-9' }, + params: { organizationId: 'org-9' }, + body: { agentId: 'agent-9' }, + query: {}, + method: 'POST', + path: '/api/v1/orgs/org-9/agents/agent-9/execute', + originalUrl: '/api/v1/orgs/org-9/agents/agent-9/execute', + url: '/api/v1/orgs/org-9/agents/agent-9/execute', + socket: {}, + }), + }), + } as unknown as ExecutionContext; + + const mockCallHandler: CallHandler = { + handle: () => { + capturedOrg = RequestContext.getOrganizationId(); + capturedAgent = RequestContext.getAgentId(); + capturedAuthMethod = RequestContext.getPrincipal()?.authMethod; + return of({ ok: true }); + }, + }; + + const observable = interceptor.intercept(mockExecutionContext, mockCallHandler); + + await new Promise((resolve, reject) => { + observable.subscribe({ + next: () => resolve(), + error: (err) => reject(err), + }); + }); + + expect(capturedOrg).toBe('org-9'); + expect(capturedAgent).toBe('agent-9'); + expect(capturedAuthMethod).toBe('service'); + }); +}); diff --git a/src/common/interceptors/request-context.interceptor.ts b/src/common/interceptors/request-context.interceptor.ts new file mode 100644 index 0000000..2d4cee3 --- /dev/null +++ b/src/common/interceptors/request-context.interceptor.ts @@ -0,0 +1,121 @@ +import { + CallHandler, + ExecutionContext, + Injectable, + NestInterceptor, +} from '@nestjs/common'; +import { Observable } from 'rxjs'; +import { Request } from 'express'; +import { v7 as uuidv7 } from 'uuid'; +import { + RequestContext, + RequestContextData, + RequestPrincipal, +} from '../context/request-context'; +import { AuthenticatedUser } from '../interfaces/authenticated-user.interface'; +import { + CORRELATION_ID_HEADER, + REQUEST_ID_HEADER, +} from '../constants/headers'; + +/** + * Seeds the structured request context (see {@link RequestContext}) at the very + * start of every HTTP request, entering the AsyncLocalStorage context around + * the remainder of the request lifecycle. + * + * The context is populated with routing metadata, correlation/trace ids and — + * once authentication has run — the resolved principal (user / organization / + * agent). Because `RequestContext.run` propagates the store through every async + * continuation, downstream controllers and services can read a typed snapshot + * of the request without threading arguments through the call tree. + */ +@Injectable() +export class RequestContextInterceptor implements NestInterceptor { + intercept(context: ExecutionContext, next: CallHandler): Observable { + const http = context.switchToHttp(); + const req = http.getRequest(); + + const contextData = this.seed(req); + + return new Observable((subscriber) => { + RequestContext.run(contextData, () => { + next.handle().subscribe({ + next: (val) => subscriber.next(val), + error: (err) => subscriber.error(err), + complete: () => subscriber.complete(), + }); + }); + }); + } + + private seed(req: Request & { user?: AuthenticatedUser }): RequestContextData { + const requestId = + RequestContext.getRequestId() ?? + (req.headers[REQUEST_ID_HEADER] as string | undefined) ?? + `req_${uuidv7()}`; + + const traceId = + (req.headers[CORRELATION_ID_HEADER] as string | undefined) ?? + requestId; + + const correlationId = + (req.headers[CORRELATION_ID_HEADER] as string | undefined) ?? + requestId; + + const user = req.user; + + const principal: RequestPrincipal | undefined = user + ? { + userId: user.id, + organizationId: user.organizationId, + role: user.role, + authMethod: 'jwt', + } + : this.resolveImplicitPrincipal(req); + + const ip = + (typeof req.headers['x-forwarded-for'] === 'string' + ? req.headers['x-forwarded-for'].split(',')[0]?.trim() + : undefined) ?? + req.socket?.remoteAddress ?? + null; + + const userAgent = (req.headers['user-agent'] as string | undefined) ?? null; + + return { + identity: { + requestId, + correlationId, + traceId, + method: req.method, + path: req.path, + url: req.originalUrl ?? req.url, + ip, + userAgent, + startedAt: Date.now(), + }, + principal, + timings: {}, + data: {}, + }; + } + + private resolveImplicitPrincipal( + req: Request, + ): RequestPrincipal | undefined { + const organizationId = + (req.params?.organizationId as string | undefined) ?? + (req.headers['x-organization-id'] as string | undefined); + const agentId = + (req.params?.agentId as string | undefined) ?? + (req.body?.agentId as string | undefined) ?? + (req.query?.agentId as string | undefined) ?? + (req.headers['x-agent-id'] as string | undefined); + + if (!organizationId && !agentId) { + return undefined; + } + + return { organizationId, agentId, authMethod: 'service' }; + } +}