diff --git a/src/shared/lib/with-request-context.test.ts b/src/shared/lib/with-request-context.test.ts new file mode 100644 index 0000000..8af340c --- /dev/null +++ b/src/shared/lib/with-request-context.test.ts @@ -0,0 +1,81 @@ +import { describe, it, expect, vi, beforeEach } from "vitest" + +// mock next/headers 的 headers() 函数 +const mockHeaders = vi.fn() +vi.mock("next/headers", () => ({ + headers: () => mockHeaders(), +})) + +// mock node:crypto 的 randomUUID(vi.stubGlobal 无法拦截内置模块导入) +// vitest 4.x 要求 mock 提供与原模块一致的导出结构(含 default) +vi.mock("node:crypto", async (importOriginal) => { + const actual = await importOriginal() + return { + ...actual, + default: { ...actual, randomUUID: () => "generated-uuid" }, + randomUUID: () => "generated-uuid", + } +}) + +import { withRequestContext } from "./with-request-context" +import { getRequestContext } from "./request-context" + +describe("withRequestContext", () => { + beforeEach(() => { + mockHeaders.mockReset() + }) + + it("从 headers 读取 x-request-id 并注入 context", async () => { + mockHeaders.mockReturnValue({ + get: (name: string) => + name === "x-request-id" ? "proxy-injected-id" : null, + }) + + const wrapped = withRequestContext(async () => { + return getRequestContext() + }) + + const result = await wrapped() + expect(result).toEqual({ requestId: "proxy-injected-id" }) + }) + + it("headers 无 x-request-id 时生成新 UUID", async () => { + mockHeaders.mockReturnValue({ + get: () => null, + }) + + const wrapped = withRequestContext(async () => { + return getRequestContext() + }) + + const result = await wrapped() + expect(result).toEqual({ requestId: "generated-uuid" }) + }) + + it("包装函数的参数与返回值正常透传", async () => { + mockHeaders.mockReturnValue({ get: () => null }) + + const wrapped = withRequestContext(async (a: number, b: number) => { + return a + b + }) + + const result = await wrapped(3, 4) + expect(result).toBe(7) + }) + + it("调用栈下游 logger 自动获得 requestId", async () => { + mockHeaders.mockReturnValue({ + get: (name: string) => + name === "x-request-id" ? "req-downstream" : null, + }) + + let capturedCtx: { requestId?: string } = {} + const wrapped = withRequestContext(async () => { + // 模拟 data-access 层调用 + capturedCtx = getRequestContext() + }) + + await wrapped() + expect(capturedCtx.requestId).toBe("req-downstream") + }) +}) diff --git a/src/shared/lib/with-request-context.ts b/src/shared/lib/with-request-context.ts new file mode 100644 index 0000000..dc7db46 --- /dev/null +++ b/src/shared/lib/with-request-context.ts @@ -0,0 +1,34 @@ +import { randomUUID } from "node:crypto" +import { headers } from "next/headers" + +import { requestContextStorage, type RequestContext } from "./request-context" + +/** + * 包装 Server Action / Route Handler,注入 requestId 到 AsyncLocalStorage。 + * + * 工作流程: + * 1. 通过 `headers()` 读取 proxy.ts 注入的 `x-request-id` + * 2. 若请求头无此字段(如直接调用的内部函数),生成新 UUID + * 3. 通过 `requestContextStorage.run()` 注入到 AsyncLocalStorage + * 4. 在调用栈内的所有 logger 调用自动获得 requestId + * + * @example + * ```ts + * export const createUserAction = withRequestContext( + * async (state: ActionState, input: CreateUserInput) => { + * // 此处 logger.info 会自动带 requestId + * return handleAction(...) + * } + * ) + * ``` + */ +export function withRequestContext( + fn: (...args: TArgs) => Promise +): (...args: TArgs) => Promise { + return async (...args: TArgs) => { + const headersList = await headers() + const requestId = headersList.get("x-request-id") ?? randomUUID() + const ctx: RequestContext = { requestId } + return requestContextStorage.run(ctx, () => fn(...args)) + } +}