@terreno/api 0.13.2 → 0.14.0

package/src/requestContext.test.ts

@@ -0,0 +1,196 @@
+import {afterEach, beforeEach, describe, expect, it, type Mock} from "bun:test";
+import {Writable} from "node:stream";
+import * as Sentry from "@sentry/bun";
+import express from "express";
+import supertest from "supertest";
+import winston from "winston";
+
+import {logger, setupLogging} from "./logger";
+import {
+  getCurrentLogContext,
+  getCurrentRequestContext,
+  getCurrentRequestContextAttributes,
+  getRequestContextFromAttributes,
+  requestContextMiddleware,
+  runWithRequestContext,
+  runWithRequestContextAttributes,
+  setRequestContext,
+} from "./requestContext";
+
+describe("request context job propagation", () => {
+  beforeEach(() => {
+    const scope = Sentry.getCurrentScope();
+    (scope.setContext as unknown as Mock<(...args: unknown[]) => void>).mockClear();
+    (scope.setTag as unknown as Mock<(...args: unknown[]) => void>).mockClear();
+    (scope.setUser as unknown as Mock<(...args: unknown[]) => void>).mockClear();
+  });
+
+  afterEach(() => {
+    setupLogging({disableFileLogging: true});
+  });
+
+  it("serializes the current context for downstream jobs", () => {
+    runWithRequestContext(
+      {
+        jobId: "job-1",
+        requestId: "request-1",
+        sessionId: "session-1",
+        spanId: "span-1",
+        traceId: "trace-1",
+        traceSampled: false,
+        userId: "user-1",
+      },
+      () => {
+        expect(getCurrentRequestContextAttributes()).toEqual({
+          "x-job-id": "job-1",
+          "x-request-id": "request-1",
+          "x-session-id": "session-1",
+          "x-span-id": "span-1",
+          "x-trace-id": "trace-1",
+          "x-trace-sampled": "false",
+          "x-user-id": "user-1",
+        });
+      }
+    );
+  });
+
+  it("allows downstream jobs to replace only the job id", () => {
+    runWithRequestContext(
+      {jobId: "job-parent", requestId: "request-1", sessionId: "session-1"},
+      () => {
+        expect(getCurrentRequestContextAttributes({jobId: "job-child"})).toEqual({
+          "x-job-id": "job-child",
+          "x-request-id": "request-1",
+          "x-session-id": "session-1",
+        });
+      }
+    );
+  });
+
+  it("restores worker context from message attributes", () => {
+    const scope = Sentry.getCurrentScope();
+    const setContextMock = scope.setContext as unknown as Mock<(...args: unknown[]) => void>;
+    const setTagMock = scope.setTag as unknown as Mock<(...args: unknown[]) => void>;
+    const setUserMock = scope.setUser as unknown as Mock<(...args: unknown[]) => void>;
+
+    runWithRequestContextAttributes(
+      {
+        traceparent: "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01",
+        "x-job-id": "job-worker-1",
+        "x-request-id": "request-1",
+        "x-session-id": "session-1",
+        "x-user-id": "user-1",
+      },
+      () => {
+        expect(getCurrentRequestContext()).toEqual({
+          jobId: "job-worker-1",
+          requestId: "request-1",
+          sessionId: "session-1",
+          spanId: "00f067aa0ba902b7",
+          traceId: "4bf92f3577b34da6a3ce929d0e0e4736",
+          traceSampled: true,
+          userId: "user-1",
+        });
+      }
+    );
+
+    expect(setTagMock).toHaveBeenCalledWith("request_id", "request-1");
+    expect(setTagMock).toHaveBeenCalledWith("session_id", "session-1");
+    expect(setTagMock).toHaveBeenCalledWith("job_id", "job-worker-1");
+    expect(setTagMock).toHaveBeenCalledWith("user_id", "user-1");
+    expect(setTagMock).toHaveBeenCalledWith("trace_id", "4bf92f3577b34da6a3ce929d0e0e4736");
+    expect(setUserMock).toHaveBeenCalledWith({id: "user-1"});
+    expect(setContextMock).toHaveBeenCalledWith("request_context", {
+      jobId: "job-worker-1",
+      requestId: "request-1",
+      sessionId: "session-1",
+      spanId: "00f067aa0ba902b7",
+      traceId: "4bf92f3577b34da6a3ce929d0e0e4736",
+      traceSampled: true,
+      userId: "user-1",
+    });
+  });
+
+  it("sends context updates to Sentry for auth session changes", () => {
+    const scope = Sentry.getCurrentScope();
+    const setTagMock = scope.setTag as unknown as Mock<(...args: unknown[]) => void>;
+    const setUserMock = scope.setUser as unknown as Mock<(...args: unknown[]) => void>;
+
+    runWithRequestContext({requestId: "request-auth-1"}, () => {
+      setTagMock.mockClear();
+      setUserMock.mockClear();
+      setRequestContext({sessionId: "session-auth-1", userId: "user-auth-1"});
+    });
+
+    expect(setTagMock).toHaveBeenCalledWith("session_id", "session-auth-1");
+    expect(setTagMock).toHaveBeenCalledWith("user_id", "user-auth-1");
+    expect(setUserMock).toHaveBeenCalledWith({id: "user-auth-1"});
+  });
+
+  it("uses trace id as request id when attributes do not include request id", () => {
+    const context = getRequestContextFromAttributes({
+      traceparent: "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-00",
+      "x-job-id": "job-worker-1",
+    });
+
+    expect(context).toEqual({
+      jobId: "job-worker-1",
+      requestId: "4bf92f3577b34da6a3ce929d0e0e4736",
+      sessionId: undefined,
+      spanId: "00f067aa0ba902b7",
+      traceId: "4bf92f3577b34da6a3ce929d0e0e4736",
+      traceSampled: false,
+      userId: undefined,
+    });
+  });
+
+  it("accepts job id from HTTP headers", async () => {
+    const app = express();
+    app.use(requestContextMiddleware);
+    app.get("/job", (req, res) => {
+      res.json({context: getCurrentLogContext(), jobId: req.jobId});
+    });
+
+    const res = await supertest(app)
+      .get("/job")
+      .set("X-Job-ID", "job-http-1")
+      .set("X-Request-ID", "request-http-1")
+      .expect(200);
+
+    expect(res.headers["x-job-id"]).toBe("job-http-1");
+    expect(res.body).toEqual({
+      context: {jobId: "job-http-1", requestId: "request-http-1"},
+      jobId: "job-http-1",
+    });
+  });
+
+  it("adds job id to logger context", () => {
+    let output = "";
+    const stream = new Writable({
+      write: (chunk, _encoding, callback): void => {
+        output += chunk.toString();
+        callback();
+      },
+    });
+
+    setupLogging({
+      disableConsoleLogging: true,
+      disableFileLogging: true,
+      transports: [
+        new winston.transports.Stream({
+          format: winston.format.printf((info) => {
+            return `${info.level}: ${info.message} requestId=${info.requestId} jobId=${info.jobId}`;
+          }),
+          stream,
+        }),
+      ],
+    });
+
+    runWithRequestContext({jobId: "job-log-1", requestId: "request-log-1"}, () => {
+      logger.info("worker handled job");
+    });
+
+    expect(output).toContain("requestId=request-log-1");
+    expect(output).toContain("jobId=job-log-1");
+  });
+});

package/src/requestContext.ts

@@ -0,0 +1,368 @@
+import {AsyncLocalStorage} from "node:async_hooks";
+import {randomUUID} from "node:crypto";
+import * as Sentry from "@sentry/bun";
+import type express from "express";
+import type {JwtPayload} from "jsonwebtoken";
+
+const CLOUD_TRACE_CONTEXT_HEADER = "x-cloud-trace-context";
+const JOB_ID_HEADER = "x-job-id";
+const REQUEST_ID_HEADERS = ["x-request-id", "x-correlation-id", "x-transaction-id"];
+const SESSION_ID_HEADER = "x-session-id";
+const SPAN_ID_HEADER = "x-span-id";
+const TRACE_ID_HEADER = "x-trace-id";
+const TRACE_PARENT_HEADER = "traceparent";
+const TRACE_SAMPLED_HEADER = "x-trace-sampled";
+const USER_ID_HEADER = "x-user-id";
+
+export interface RequestContext {
+  jobId?: string;
+  requestId: string;
+  sessionId?: string;
+  spanId?: string;
+  traceId?: string;
+  traceSampled?: boolean;
+  userId?: string;
+}
+
+export type RequestContextAttributes = Record<string, string>;
+
+export const REQUEST_CONTEXT_ATTRIBUTE_NAMES = {
+  jobId: JOB_ID_HEADER,
+  requestId: "x-request-id",
+  sessionId: SESSION_ID_HEADER,
+  spanId: SPAN_ID_HEADER,
+  traceId: TRACE_ID_HEADER,
+  traceParent: TRACE_PARENT_HEADER,
+  traceSampled: TRACE_SAMPLED_HEADER,
+  userId: USER_ID_HEADER,
+} as const;
+
+interface GoogleCloudTraceContext {
+  spanId?: string;
+  traceId?: string;
+  traceSampled?: boolean;
+}
+
+export interface JwtSessionPayload extends JwtPayload {
+  sid?: string;
+  sessionId?: string;
+}
+
+const requestContextStorage = new AsyncLocalStorage<RequestContext>();
+
+const getHeader = (req: express.Request, headerName: string): string | undefined => {
+  const value = req.header(headerName);
+  if (!Array.isArray(value)) {
+    return value;
+  }
+  return value[0];
+};
+
+const parseGoogleCloudTraceContext = (
+  headerValue?: string
+): GoogleCloudTraceContext | undefined => {
+  if (!headerValue) {
+    return undefined;
+  }
+
+  const [traceAndSpan, options] = headerValue.split(";");
+  const [traceId, spanId] = traceAndSpan.split("/");
+  if (!traceId) {
+    return undefined;
+  }
+
+  return {
+    spanId,
+    traceId,
+    traceSampled: options === "o=1",
+  };
+};
+
+const parseTraceParent = (headerValue?: string): GoogleCloudTraceContext | undefined => {
+  if (!headerValue) {
+    return undefined;
+  }
+
+  const [_version, traceId, spanId, flags] = headerValue.split("-");
+  if (!traceId) {
+    return undefined;
+  }
+
+  return {
+    spanId,
+    traceId,
+    traceSampled: flags === "01",
+  };
+};
+
+const parseTraceSampled = (value?: string | boolean): boolean | undefined => {
+  if (typeof value === "boolean") {
+    return value;
+  }
+  if (value === "true" || value === "1") {
+    return true;
+  }
+  if (value === "false" || value === "0") {
+    return false;
+  }
+  return undefined;
+};
+
+const getIncomingRequestId = (
+  req: express.Request,
+  traceContext?: GoogleCloudTraceContext
+): string => {
+  for (const headerName of REQUEST_ID_HEADERS) {
+    const headerValue = getHeader(req, headerName);
+    if (headerValue) {
+      return headerValue;
+    }
+  }
+
+  if (traceContext?.traceId) {
+    return traceContext.traceId;
+  }
+
+  return randomUUID();
+};
+
+export const getSessionIdFromJwtPayload = (
+  payload?: JwtSessionPayload | null
+): string | undefined => {
+  if (!payload) {
+    return undefined;
+  }
+  if (payload.sid) {
+    return payload.sid;
+  }
+  return payload.sessionId;
+};
+
+export const getRequestContextFromAttributes = (
+  attributes: Record<string, string | undefined> = {}
+): RequestContext => {
+  const cloudTraceContext = parseGoogleCloudTraceContext(attributes[CLOUD_TRACE_CONTEXT_HEADER]);
+  const traceParentContext = parseTraceParent(attributes[TRACE_PARENT_HEADER]);
+  const traceContext = cloudTraceContext ?? traceParentContext;
+
+  return {
+    jobId: attributes[JOB_ID_HEADER],
+    requestId:
+      attributes[REQUEST_CONTEXT_ATTRIBUTE_NAMES.requestId] ??
+      attributes["x-correlation-id"] ??
+      attributes["x-transaction-id"] ??
+      traceContext?.traceId ??
+      randomUUID(),
+    sessionId: attributes[SESSION_ID_HEADER],
+    spanId: attributes[SPAN_ID_HEADER] ?? traceContext?.spanId,
+    traceId: attributes[TRACE_ID_HEADER] ?? traceContext?.traceId,
+    traceSampled: parseTraceSampled(attributes[TRACE_SAMPLED_HEADER]) ?? traceContext?.traceSampled,
+    userId: attributes[USER_ID_HEADER],
+  };
+};
+
+export const getCurrentRequestContext = (): RequestContext | undefined => {
+  return requestContextStorage.getStore();
+};
+
+export const getCurrentLogContext = (): Partial<RequestContext> => {
+  const context = getCurrentRequestContext();
+  if (!context) {
+    return {};
+  }
+
+  return {
+    jobId: context.jobId,
+    requestId: context.requestId,
+    sessionId: context.sessionId,
+    spanId: context.spanId,
+    traceId: context.traceId,
+    traceSampled: context.traceSampled,
+    userId: context.userId,
+  };
+};
+
+const setSentryTag = (
+  scope: ReturnType<typeof Sentry.getCurrentScope>,
+  name: string,
+  value?: string | boolean
+): void => {
+  if (typeof value === "undefined") {
+    return;
+  }
+  scope.setTag(name, String(value));
+};
+
+const setSentryContextValue = (
+  context: Record<string, string | boolean>,
+  name: string,
+  value?: string | boolean
+): void => {
+  if (typeof value === "undefined") {
+    return;
+  }
+  context[name] = value;
+};
+
+export const applyRequestContextToSentry = (
+  context: Partial<RequestContext> = getCurrentLogContext()
+): void => {
+  const scope = Sentry.getCurrentScope();
+  setSentryTag(scope, "request_id", context.requestId);
+  setSentryTag(scope, "session_id", context.sessionId);
+  setSentryTag(scope, "job_id", context.jobId);
+  setSentryTag(scope, "user_id", context.userId);
+  setSentryTag(scope, "trace_id", context.traceId);
+  setSentryTag(scope, "span_id", context.spanId);
+  setSentryTag(scope, "trace_sampled", context.traceSampled);
+
+  if (context.userId) {
+    scope.setUser({id: context.userId});
+  }
+
+  const sentryContext: Record<string, string | boolean> = {};
+  setSentryContextValue(sentryContext, "requestId", context.requestId);
+  setSentryContextValue(sentryContext, "sessionId", context.sessionId);
+  setSentryContextValue(sentryContext, "jobId", context.jobId);
+  setSentryContextValue(sentryContext, "userId", context.userId);
+  setSentryContextValue(sentryContext, "traceId", context.traceId);
+  setSentryContextValue(sentryContext, "spanId", context.spanId);
+  setSentryContextValue(sentryContext, "traceSampled", context.traceSampled);
+
+  if (Object.keys(sentryContext).length > 0) {
+    scope.setContext("request_context", sentryContext);
+  }
+};
+
+export const setRequestContext = (updates: Partial<RequestContext>): void => {
+  const context = getCurrentRequestContext();
+  if (!context) {
+    return;
+  }
+  Object.assign(context, updates);
+  applyRequestContextToSentry(context);
+};
+
+const setAttribute = (
+  attributes: RequestContextAttributes,
+  name: string,
+  value?: string | boolean
+): void => {
+  if (typeof value === "undefined") {
+    return;
+  }
+  attributes[name] = String(value);
+};
+
+export const getCurrentRequestContextAttributes = (
+  overrides: Partial<RequestContext> = {}
+): RequestContextAttributes => {
+  const context = {...getCurrentLogContext(), ...overrides};
+  const attributes: RequestContextAttributes = {};
+  setAttribute(attributes, REQUEST_CONTEXT_ATTRIBUTE_NAMES.requestId, context.requestId);
+  setAttribute(attributes, REQUEST_CONTEXT_ATTRIBUTE_NAMES.sessionId, context.sessionId);
+  setAttribute(attributes, REQUEST_CONTEXT_ATTRIBUTE_NAMES.jobId, context.jobId);
+  setAttribute(attributes, REQUEST_CONTEXT_ATTRIBUTE_NAMES.userId, context.userId);
+  setAttribute(attributes, REQUEST_CONTEXT_ATTRIBUTE_NAMES.traceId, context.traceId);
+  setAttribute(attributes, REQUEST_CONTEXT_ATTRIBUTE_NAMES.spanId, context.spanId);
+  setAttribute(attributes, REQUEST_CONTEXT_ATTRIBUTE_NAMES.traceSampled, context.traceSampled);
+  return attributes;
+};
+
+export const runWithRequestContext = <T>(
+  context: Partial<RequestContext>,
+  callback: () => T
+): T => {
+  const nextContext = {
+    ...context,
+    requestId: context.requestId ?? randomUUID(),
+  };
+
+  return requestContextStorage.run(nextContext, () => {
+    applyRequestContextToSentry(nextContext);
+    return callback();
+  });
+};
+
+export const runWithRequestContextAttributes = <T>(
+  attributes: Record<string, string | undefined> = {},
+  callback: () => T
+): T => {
+  return runWithRequestContext(getRequestContextFromAttributes(attributes), callback);
+};
+
+export const updateRequestContextFromRequest = (
+  req: express.Request,
+  res?: express.Response
+): void => {
+  const reqWithContext = req as express.Request & {
+    authTokenPayload?: JwtSessionPayload;
+    betterAuthSession?: {session?: {id?: string}};
+    jobId?: string;
+    requestId?: string;
+    sessionId?: string;
+  };
+  const jobId = reqWithContext.jobId ?? getHeader(req, JOB_ID_HEADER);
+  const sessionId =
+    getSessionIdFromJwtPayload(reqWithContext.authTokenPayload) ??
+    reqWithContext.betterAuthSession?.session?.id ??
+    reqWithContext.sessionId ??
+    getHeader(req, SESSION_ID_HEADER);
+  const user = req.user as {_id?: unknown; id?: string} | undefined;
+  const userId = user?.id ?? (user?._id ? String(user._id) : undefined);
+
+  setRequestContext({jobId, sessionId, userId});
+
+  if (jobId) {
+    reqWithContext.jobId = jobId;
+    res?.setHeader("X-Job-ID", jobId);
+  }
+
+  if (sessionId) {
+    reqWithContext.sessionId = sessionId;
+    res?.setHeader("X-Session-ID", sessionId);
+  }
+};
+
+export const requestContextMiddleware = (
+  req: express.Request,
+  res: express.Response,
+  next: express.NextFunction
+): void => {
+  const cloudTraceContext = parseGoogleCloudTraceContext(
+    getHeader(req, CLOUD_TRACE_CONTEXT_HEADER)
+  );
+  const traceParentContext = parseTraceParent(getHeader(req, TRACE_PARENT_HEADER));
+  const traceContext = cloudTraceContext ?? traceParentContext;
+  const requestId = getIncomingRequestId(req, traceContext);
+  const jobId = getHeader(req, JOB_ID_HEADER);
+  const sessionId = getHeader(req, SESSION_ID_HEADER);
+
+  const context: RequestContext = {
+    jobId,
+    requestId,
+    sessionId,
+    spanId: traceContext?.spanId,
+    traceId: traceContext?.traceId,
+    traceSampled: traceContext?.traceSampled,
+  };
+
+  const reqWithContext = req as express.Request & {
+    jobId?: string;
+    requestId?: string;
+    sessionId?: string;
+  };
+  if (jobId) {
+    reqWithContext.jobId = jobId;
+  }
+  reqWithContext.requestId = requestId;
+  if (sessionId) {
+    reqWithContext.sessionId = sessionId;
+  }
+  res.setHeader("X-Request-ID", requestId);
+
+  requestContextStorage.run(context, () => {
+    updateRequestContextFromRequest(req, res);
+    next();
+  });
+};