refactor: enhance TRPC middleware logging by integrating correlation ID handling and improving error tracking

lib/tracing.test.ts

@@ -0,0 +1,308 @@
+import { type Span, trace } from "@opentelemetry/api";
+import { describe, expect, it, vi } from "vitest";
+import {
+  addSpanAttribute,
+  addSpanAttributes,
+  getCurrentSpan,
+  withChildSpan,
+  withChildSpanSync,
+  withSpan,
+  withSpanSync,
+} from "./tracing";
+
+describe("tracing utilities", () => {
+  describe("withSpan", () => {
+    it("wraps async function with automatic span", async () => {
+      const mockSpan = {
+        startSpan: vi.fn().mockReturnValue({
+          end: vi.fn(),
+          setStatus: vi.fn(),
+          recordException: vi.fn(),
+          setAttribute: vi.fn(),
+        }),
+      };
+
+      vi.spyOn(trace, "getTracer").mockReturnValue(
+        mockSpan as unknown as ReturnType<typeof trace.getTracer>,
+      );
+
+      const testFn = withSpan(
+        { name: "test.operation", attributes: { test: "value" } },
+        async (arg: string) => {
+          return `result-${arg}`;
+        },
+      );
+
+      const result = await testFn("input");
+
+      expect(result).toBe("result-input");
+      expect(mockSpan.startSpan).toHaveBeenCalledWith("test.operation", {
+        attributes: { test: "value" },
+      });
+    });
+
+    it("supports dynamic attributes from function args", async () => {
+      const mockSpan = {
+        end: vi.fn(),
+        setStatus: vi.fn(),
+        recordException: vi.fn(),
+        setAttribute: vi.fn(),
+      };
+
+      vi.spyOn(trace, "getTracer").mockReturnValue({
+        startSpan: vi.fn().mockReturnValue(mockSpan),
+      } as unknown as ReturnType<typeof trace.getTracer>);
+
+      const testFn = withSpan(
+        ([domain]: [string]) => ({
+          name: "dns.lookup",
+          attributes: { domain },
+        }),
+        async (domain: string) => {
+          return { records: [], domain };
+        },
+      );
+
+      const result = await testFn("example.com");
+
+      expect(result.domain).toBe("example.com");
+      expect(trace.getTracer).toHaveBeenCalled();
+    });
+
+    it("records exceptions and sets error status", async () => {
+      const mockSpan = {
+        end: vi.fn(),
+        setStatus: vi.fn(),
+        recordException: vi.fn(),
+        setAttribute: vi.fn(),
+      };
+
+      vi.spyOn(trace, "getTracer").mockReturnValue({
+        startSpan: vi.fn().mockReturnValue(mockSpan),
+      } as unknown as ReturnType<typeof trace.getTracer>);
+
+      const testError = new Error("Test error");
+      const testFn = withSpan({ name: "test.error" }, async () => {
+        throw testError;
+      });
+
+      await expect(testFn()).rejects.toThrow("Test error");
+      expect(mockSpan.recordException).toHaveBeenCalledWith(testError);
+      expect(mockSpan.setStatus).toHaveBeenCalledWith({
+        code: 2, // ERROR
+        message: "Test error",
+      });
+      expect(mockSpan.end).toHaveBeenCalled();
+    });
+
+    it("adds result metadata for arrays", async () => {
+      const mockSpan = {
+        end: vi.fn(),
+        setStatus: vi.fn(),
+        recordException: vi.fn(),
+        setAttribute: vi.fn(),
+      };
+
+      vi.spyOn(trace, "getTracer").mockReturnValue({
+        startSpan: vi.fn().mockReturnValue(mockSpan),
+      } as unknown as ReturnType<typeof trace.getTracer>);
+
+      const testFn = withSpan({ name: "test.array" }, async () => {
+        return [1, 2, 3];
+      });
+
+      await testFn();
+
+      expect(mockSpan.setAttribute).toHaveBeenCalledWith("result.count", 3);
+    });
+  });
+
+  describe("withSpanSync", () => {
+    it("wraps synchronous function with automatic span", () => {
+      const mockSpan = {
+        end: vi.fn(),
+        setStatus: vi.fn(),
+        recordException: vi.fn(),
+        setAttribute: vi.fn(),
+      };
+
+      vi.spyOn(trace, "getTracer").mockReturnValue({
+        startSpan: vi.fn().mockReturnValue(mockSpan),
+      } as unknown as ReturnType<typeof trace.getTracer>);
+
+      const testFn = withSpanSync({ name: "test.sync" }, (n: number) => {
+        return n * 2;
+      });
+
+      const result = testFn(5);
+
+      expect(result).toBe(10);
+      expect(mockSpan.end).toHaveBeenCalled();
+    });
+
+    it("records exceptions in sync functions", () => {
+      const mockSpan = {
+        end: vi.fn(),
+        setStatus: vi.fn(),
+        recordException: vi.fn(),
+        setAttribute: vi.fn(),
+      };
+
+      vi.spyOn(trace, "getTracer").mockReturnValue({
+        startSpan: vi.fn().mockReturnValue(mockSpan),
+      } as unknown as ReturnType<typeof trace.getTracer>);
+
+      const testError = new Error("Sync error");
+      const testFn = withSpanSync({ name: "test.sync.error" }, () => {
+        throw testError;
+      });
+
+      expect(() => testFn()).toThrow("Sync error");
+      expect(mockSpan.recordException).toHaveBeenCalledWith(testError);
+      expect(mockSpan.end).toHaveBeenCalled();
+    });
+  });
+
+  describe("withChildSpan", () => {
+    it("creates child span for scoped operation", async () => {
+      const mockSpan = {
+        end: vi.fn(),
+        setStatus: vi.fn(),
+        recordException: vi.fn(),
+        setAttribute: vi.fn(),
+      };
+
+      vi.spyOn(trace, "getTracer").mockReturnValue({
+        startSpan: vi.fn().mockReturnValue(mockSpan),
+      } as unknown as ReturnType<typeof trace.getTracer>);
+
+      const result = await withChildSpan(
+        { name: "child.operation", attributes: { step: "1" } },
+        async () => {
+          return "child-result";
+        },
+      );
+
+      expect(result).toBe("child-result");
+      expect(mockSpan.end).toHaveBeenCalled();
+    });
+  });
+
+  describe("withChildSpanSync", () => {
+    it("creates child span for synchronous scoped operation", () => {
+      const mockSpan = {
+        end: vi.fn(),
+        setStatus: vi.fn(),
+        recordException: vi.fn(),
+        setAttribute: vi.fn(),
+      };
+
+      vi.spyOn(trace, "getTracer").mockReturnValue({
+        startSpan: vi.fn().mockReturnValue(mockSpan),
+      } as unknown as ReturnType<typeof trace.getTracer>);
+
+      const result = withChildSpanSync({ name: "child.sync" }, () => {
+        return 42;
+      });
+
+      expect(result).toBe(42);
+      expect(mockSpan.end).toHaveBeenCalled();
+    });
+  });
+
+  describe("getCurrentSpan", () => {
+    it("returns undefined when no active span", () => {
+      vi.spyOn(trace, "getSpan").mockReturnValue(undefined);
+      expect(getCurrentSpan()).toBeUndefined();
+    });
+
+    it("returns current active span when available", () => {
+      const mockSpan = {
+        end: vi.fn(),
+        setAttribute: vi.fn(),
+      } as unknown as Span;
+
+      vi.spyOn(trace, "getSpan").mockReturnValue(mockSpan);
+      expect(getCurrentSpan()).toBe(mockSpan);
+    });
+  });
+
+  describe("addSpanAttribute", () => {
+    it("adds attribute to current span when available", () => {
+      const mockSpan = {
+        setAttribute: vi.fn(),
+      } as unknown as Span;
+
+      vi.spyOn(trace, "getSpan").mockReturnValue(mockSpan);
+
+      addSpanAttribute("custom.key", "value");
+      expect(mockSpan.setAttribute).toHaveBeenCalledWith("custom.key", "value");
+    });
+
+    it("does nothing when no active span", () => {
+      vi.spyOn(trace, "getSpan").mockReturnValue(undefined);
+      // Should not throw
+      expect(() => addSpanAttribute("key", "value")).not.toThrow();
+    });
+
+    it("filters out invalid attribute values", () => {
+      const mockSpan = {
+        setAttribute: vi.fn(),
+      } as unknown as Span;
+
+      vi.spyOn(trace, "getSpan").mockReturnValue(mockSpan);
+
+      addSpanAttribute("invalid", null);
+      addSpanAttribute("also.invalid", undefined);
+      addSpanAttribute("valid", "ok");
+
+      expect(mockSpan.setAttribute).toHaveBeenCalledTimes(1);
+      expect(mockSpan.setAttribute).toHaveBeenCalledWith("valid", "ok");
+    });
+  });
+
+  describe("addSpanAttributes", () => {
+    it("adds multiple attributes to current span", () => {
+      const mockSpan = {
+        setAttribute: vi.fn(),
+      } as unknown as Span;
+
+      vi.spyOn(trace, "getSpan").mockReturnValue(mockSpan);
+
+      addSpanAttributes({
+        "dns.domain": "example.com",
+        "dns.records_count": 5,
+        "dns.cache_hit": true,
+      });
+
+      expect(mockSpan.setAttribute).toHaveBeenCalledWith(
+        "dns.domain",
+        "example.com",
+      );
+      expect(mockSpan.setAttribute).toHaveBeenCalledWith(
+        "dns.records_count",
+        5,
+      );
+      expect(mockSpan.setAttribute).toHaveBeenCalledWith("dns.cache_hit", true);
+    });
+
+    it("filters out invalid attributes", () => {
+      const mockSpan = {
+        setAttribute: vi.fn(),
+      } as unknown as Span;
+
+      vi.spyOn(trace, "getSpan").mockReturnValue(mockSpan);
+
+      addSpanAttributes({
+        valid: "ok",
+        invalid: null,
+        alsoInvalid: undefined,
+        number: 123,
+      });
+
+      expect(mockSpan.setAttribute).toHaveBeenCalledTimes(2);
+      expect(mockSpan.setAttribute).toHaveBeenCalledWith("valid", "ok");
+      expect(mockSpan.setAttribute).toHaveBeenCalledWith("number", 123);
+    });
+  });
+});

lib/tracing.ts

@@ -0,0 +1,321 @@
+/**
+ * OpenTelemetry tracing utilities for automatic span creation.
+ *
+ * Provides helpers to wrap service functions with automatic tracing,
+ * reducing boilerplate while maintaining full observability.
+ */
+
+import { type Attributes, context, type Span, trace } from "@opentelemetry/api";
+
+// ============================================================================
+// Types
+// ============================================================================
+
+export type SpanOptions = {
+  /** Name of the span (e.g., 'dns.getDnsRecords') */
+  name: string;
+  /** Additional attributes to attach to the span */
+  attributes?: Attributes;
+};
+
+export type DynamicSpanOptions<TArgs extends unknown[]> = (
+  args: TArgs,
+) => SpanOptions;
+
+// ============================================================================
+// Span Helpers
+// ============================================================================
+
+/**
+ * Wraps an async service function with automatic OpenTelemetry span creation.
+ *
+ * Features:
+ * - Automatic span lifecycle management (start/end)
+ * - Exception recording with proper error status
+ * - Optional dynamic attributes based on function arguments
+ * - Result metadata extraction (e.g., array length)
+ *
+ * @example Static span name and attributes
+ * ```typescript
+ * export const getConfig = withSpan(
+ *   { name: 'config.get', attributes: { 'config.type': 'global' } },
+ *   async () => {
+ *     return await fetchConfig();
+ *   }
+ * );
+ * ```
+ *
+ * @example Dynamic attributes from function arguments
+ * ```typescript
+ * export const getDnsRecords = withSpan(
+ *   ([domain]) => ({
+ *     name: 'dns.getDnsRecords',
+ *     attributes: { 'dns.domain': domain }
+ *   }),
+ *   async (domain: string) => {
+ *     return await lookupDns(domain);
+ *   }
+ * );
+ * ```
+ */
+export function withSpan<TArgs extends unknown[], TReturn>(
+  options: SpanOptions | DynamicSpanOptions<TArgs>,
+  fn: (...args: TArgs) => Promise<TReturn>,
+): (...args: TArgs) => Promise<TReturn> {
+  return async (...args: TArgs): Promise<TReturn> => {
+    // Determine span options (static or dynamic based on args)
+    const spanOptions = typeof options === "function" ? options(args) : options;
+
+    const tracer = trace.getTracer("service-layer");
+    const span = tracer.startSpan(spanOptions.name, {
+      attributes: spanOptions.attributes,
+    });
+
+    try {
+      // Execute the wrapped function
+      const result = await fn(...args);
+
+      // Add result metadata if available
+      addResultMetadata(span, result);
+
+      return result;
+    } catch (error) {
+      // Record exception and set error status
+      span.recordException(error as Error);
+      span.setStatus({
+        code: 2, // ERROR (SpanStatusCode.ERROR)
+        message: error instanceof Error ? error.message : String(error),
+      });
+      throw error;
+    } finally {
+      span.end();
+    }
+  };
+}
+
+/**
+ * Wraps a synchronous service function with automatic span creation.
+ *
+ * Use this for CPU-bound operations or synchronous transforms that
+ * you want to measure independently.
+ *
+ * @example
+ * ```typescript
+ * export const parseConfig = withSpanSync(
+ *   { name: 'config.parse' },
+ *   (raw: string) => {
+ *     return JSON.parse(raw);
+ *   }
+ * );
+ * ```
+ */
+export function withSpanSync<TArgs extends unknown[], TReturn>(
+  options: SpanOptions | DynamicSpanOptions<TArgs>,
+  fn: (...args: TArgs) => TReturn,
+): (...args: TArgs) => TReturn {
+  return (...args: TArgs): TReturn => {
+    const spanOptions = typeof options === "function" ? options(args) : options;
+
+    const tracer = trace.getTracer("service-layer");
+    const span = tracer.startSpan(spanOptions.name, {
+      attributes: spanOptions.attributes,
+    });
+
+    try {
+      const result = fn(...args);
+      addResultMetadata(span, result);
+      return result;
+    } catch (error) {
+      span.recordException(error as Error);
+      span.setStatus({
+        code: 2, // ERROR
+        message: error instanceof Error ? error.message : String(error),
+      });
+      throw error;
+    } finally {
+      span.end();
+    }
+  };
+}
+
+/**
+ * Create a child span for a scoped operation within an already traced function.
+ *
+ * Useful for breaking down a complex operation into multiple instrumented steps
+ * without needing to extract separate functions.
+ *
+ * @example
+ * ```typescript
+ * async function complexOperation(domain: string) {
+ *   const dns = await withChildSpan(
+ *     { name: 'dns.lookup', attributes: { domain } },
+ *     () => lookupDns(domain)
+ *   );
+ *
+ *   const cert = await withChildSpan(
+ *     { name: 'cert.check', attributes: { domain } },
+ *     () => checkCert(domain)
+ *   );
+ *
+ *   return { dns, cert };
+ * }
+ * ```
+ */
+export async function withChildSpan<T>(
+  options: SpanOptions,
+  fn: () => Promise<T>,
+): Promise<T> {
+  const tracer = trace.getTracer("service-layer");
+  const span = tracer.startSpan(options.name, {
+    attributes: options.attributes,
+  });
+
+  try {
+    const result = await fn();
+    addResultMetadata(span, result);
+    return result;
+  } catch (error) {
+    span.recordException(error as Error);
+    span.setStatus({
+      code: 2,
+      message: error instanceof Error ? error.message : String(error),
+    });
+    throw error;
+  } finally {
+    span.end();
+  }
+}
+
+/**
+ * Synchronous version of withChildSpan.
+ */
+export function withChildSpanSync<T>(options: SpanOptions, fn: () => T): T {
+  const tracer = trace.getTracer("service-layer");
+  const span = tracer.startSpan(options.name, {
+    attributes: options.attributes,
+  });
+
+  try {
+    const result = fn();
+    addResultMetadata(span, result);
+    return result;
+  } catch (error) {
+    span.recordException(error as Error);
+    span.setStatus({
+      code: 2,
+      message: error instanceof Error ? error.message : String(error),
+    });
+    throw error;
+  } finally {
+    span.end();
+  }
+}
+
+/**
+ * Get the current active span (if any).
+ * Useful for adding custom attributes to the current span.
+ *
+ * @example
+ * ```typescript
+ * const span = getCurrentSpan();
+ * if (span) {
+ *   span.setAttribute('custom.metric', value);
+ * }
+ * ```
+ */
+export function getCurrentSpan(): Span | undefined {
+  return trace.getSpan(context.active());
+}
+
+/**
+ * Add a custom attribute to the current span (if any).
+ *
+ * @example
+ * ```typescript
+ * addSpanAttribute('cache.hit', true);
+ * addSpanAttribute('db.query.duration_ms', 123);
+ * ```
+ */
+export function addSpanAttribute(key: string, value: unknown): void {
+  const span = getCurrentSpan();
+  if (span && isValidAttributeValue(value)) {
+    span.setAttribute(key, value);
+  }
+}
+
+/**
+ * Add multiple attributes to the current span (if any).
+ */
+export function addSpanAttributes(attributes: Record<string, unknown>): void {
+  const span = getCurrentSpan();
+  if (span) {
+    for (const [key, value] of Object.entries(attributes)) {
+      if (isValidAttributeValue(value)) {
+        span.setAttribute(key, value);
+      }
+    }
+  }
+}
+
+// ============================================================================
+// Internal Helpers
+// ============================================================================
+
+/**
+ * Extract and add metadata from the result to the span.
+ * Adds common patterns like array length, object properties, etc.
+ */
+function addResultMetadata(span: Span, result: unknown): void {
+  if (!result) return;
+
+  // Array length
+  if (Array.isArray(result)) {
+    span.setAttribute("result.count", result.length);
+    return;
+  }
+
+  // Object with length property (array-like)
+  if (
+    typeof result === "object" &&
+    "length" in result &&
+    typeof result.length === "number"
+  ) {
+    span.setAttribute("result.count", result.length);
+    return;
+  }
+
+  // Object with count property
+  if (
+    typeof result === "object" &&
+    "count" in result &&
+    typeof result.count === "number"
+  ) {
+    span.setAttribute("result.count", result.count);
+  }
+}
+
+/**
+ * Check if a value is a valid OpenTelemetry attribute value.
+ * Valid types: string, number, boolean, or arrays of these types.
+ */
+function isValidAttributeValue(
+  value: unknown,
+): value is string | number | boolean | string[] | number[] | boolean[] {
+  if (value === null || value === undefined) return false;
+
+  const type = typeof value;
+  if (type === "string" || type === "number" || type === "boolean") {
+    return true;
+  }
+
+  if (Array.isArray(value)) {
+    return value.every(
+      (item) =>
+        typeof item === "string" ||
+        typeof item === "number" ||
+        typeof item === "boolean",
+    );
+  }
+
+  return false;
+}

trpc/init.ts

@@ -4,7 +4,6 @@ import { after } from "next/server";
 import superjson from "superjson";
 import { updateLastAccessed } from "@/lib/db/repos/domains";
 import { getOrGenerateCorrelationId } from "@/lib/logger/correlation";
-import { setCorrelationId } from "@/lib/logger/server";
 
 const IP_HEADERS = ["x-real-ip", "x-forwarded-for", "cf-connecting-ip"];
 
@@ -37,8 +36,6 @@ export const createContext = async (opts?: { req?: Request }) => {
   try {
     const headerList = await headers();
     correlationId = getOrGenerateCorrelationId(headerList);
-    // Set in AsyncLocalStorage for propagation
-    setCorrelationId(correlationId);
   } catch {
     // headers() not available (tests/scripts)
   }
@@ -63,65 +60,72 @@ export const createCallerFactory = t.createCallerFactory;
  * All logs are structured JSON with OpenTelemetry tracing and correlation IDs.
  * Errors are tracked in PostHog for centralized monitoring.
  */
-const withLogging = t.middleware(async ({ path, type, input, next }) => {
+const withLogging = t.middleware(async ({ path, type, input, next, ctx }) => {
   const start = performance.now();
 
-  // Import logger (dynamic to avoid circular deps)
-  const { logger } = await import("@/lib/logger/server");
+  // Import logger and correlation utilities (dynamic to avoid circular deps)
+  const { logger, withCorrelationId } = await import("@/lib/logger/server");
+  const { generateCorrelationId } = await import("@/lib/logger/correlation");
 
-  // Log procedure start
-  logger.info("procedure start", {
-    source: "trpc",
-    path,
-    type,
-    input: input && typeof input === "object" ? { ...input } : undefined,
-  });
+  // Get correlation ID from context (set in createContext), or generate if missing
+  const correlationId = ctx.correlationId ?? generateCorrelationId();
 
-  try {
-    const result = await next();
-    const durationMs = Math.round(performance.now() - start);
-
-    // Log successful completion
-    logger.info("procedure ok", {
+  // Wrap the entire procedure execution in correlation ID context
+  return withCorrelationId(correlationId, async () => {
+    // Log procedure start
+    logger.info("procedure start", {
       source: "trpc",
       path,
       type,
-      durationMs,
       input: input && typeof input === "object" ? { ...input } : undefined,
     });
 
-    // Track slow requests (>5s threshold) in PostHog
-    if (durationMs > 5000) {
-      logger.warn("slow request", {
+    try {
+      const result = await next();
+      const durationMs = Math.round(performance.now() - start);
+
+      // Log successful completion
+      logger.info("procedure ok", {
+        source: "trpc",
+        path,
+        type,
+        durationMs,
+        input: input && typeof input === "object" ? { ...input } : undefined,
+      });
+
+      // Track slow requests (>5s threshold) in PostHog
+      if (durationMs > 5000) {
+        logger.warn("slow request", {
+          source: "trpc",
+          path,
+          type,
+          durationMs,
+        });
+
+        const { analytics } = await import("@/lib/analytics/server");
+        analytics.track("trpc_slow_request", {
+          path,
+          type,
+          durationMs,
+        });
+      }
+
+      return result;
+    } catch (err) {
+      const durationMs = Math.round(performance.now() - start);
+
+      // Log error with full details
+      logger.error("procedure error", err, {
         source: "trpc",
         path,
         type,
         durationMs,
       });
 
-      const { analytics } = await import("@/lib/analytics/server");
-      analytics.track("trpc_slow_request", {
-        path,
-        type,
-        durationMs,
-      });
+      // Re-throw the error to be handled by the error boundary
+      throw err;
     }
-
-    return result;
-  } catch (err) {
-    const durationMs = Math.round(performance.now() - start);
-
-    // Log error with full details
-    logger.error("procedure error", err, {
-      source: "trpc",
-      path,
-      type,
-      durationMs,
-    });
-
-    // Re-throw the error to be handled by the error boundary
-    throw err;
-  }
+  });
 });
 
 /**