WRNexusJS/packages/observability/src/trace.ts

import { createTracer, type Context, type Middleware } from "@wrnexus/core";
import { reportObservabilityFailure, type ObservabilityDiagnostic } from "./diagnostics.ts";

export interface TraceContext {
  version: "00";
  traceId: string;
  spanId: string;
  sampled: boolean;
}

export interface SpanRecord {
  name: string;
  traceId: string;
  spanId: string;
  parentSpanId?: string;
  sampled: boolean;
  startTime: number;
  endTime: number;
  durationMs: number;
  status: "ok" | "error";
  attributes: Record<string, string | number | boolean>;
  error?: { name: string; message: string };
}

export interface SpanExporter {
  export(spans: readonly SpanRecord[]): Promise<void> | void;
}

const TRACEPARENT = /^00-([0-9a-f]{32})-([0-9a-f]{16})-([0-9a-f]{2})$/;

export function parseTraceparent(value: string | null | undefined): TraceContext | null {
  if (!value) return null;
  const match = TRACEPARENT.exec(value.trim().toLowerCase());
  if (!match || /^0+$/.test(match[1]!) || /^0+$/.test(match[2]!)) return null;
  return {
    version: "00",
    traceId: match[1]!,
    spanId: match[2]!,
    sampled: (Number.parseInt(match[3]!, 16) & 1) === 1,
  };
}

export function formatTraceparent(context: TraceContext): string {
  return `00-${context.traceId}-${context.spanId}-${context.sampled ? "01" : "00"}`;
}

function randomHex(bytes: number, random: (target: Uint8Array) => Uint8Array): string {
  let value = "";
  while (!value || /^0+$/.test(value)) {
    value = Array.from(random(new Uint8Array(bytes)), (byte) =>
      byte.toString(16).padStart(2, "0"),
    ).join("");
  }
  return value;
}

export interface TraceMiddlewareOptions {
  serviceName?: string;
  sampleRate?: number;
  exporter?: SpanExporter;
  onSpan?: (span: SpanRecord) => void | Promise<void>;
  now?: () => number;
  random?: (target: Uint8Array) => Uint8Array;
  routeName?: (ctx: Context) => string;
  serverTiming?: boolean;
  onExportError?: (error: unknown, span: SpanRecord) => void | Promise<void>;
  diagnostic?: ObservabilityDiagnostic;
}

export function traceMiddleware(options: TraceMiddlewareOptions = {}): Middleware {
  const sampleRate = options.sampleRate ?? 1;
  if (!Number.isFinite(sampleRate) || sampleRate < 0 || sampleRate > 1) {
    throw new Error("Observability trace sampleRate must be between 0 and 1.");
  }
  const now = options.now ?? Date.now;
  const random = options.random ?? ((target: Uint8Array) => crypto.getRandomValues(target));
  return async (ctx, next) => {
    const parent = parseTraceparent(ctx.req.headers.get("traceparent"));
    const sampled = parent?.sampled ?? Math.random() < sampleRate;
    const traceId = parent?.traceId ?? randomHex(16, random);
    const spanId = randomHex(8, random);
    const trace: TraceContext = { version: "00", traceId, spanId, sampled };
    const started = now();
    ctx.locals.traceId = traceId;
    ctx.locals.spanId = spanId;
    ctx.locals.traceparent = formatTraceparent(trace);
    ctx.locals.requestId ??= traceId;
    ctx.tracer ??= createTracer(now);
    const frameworkSpan = ctx.tracer.startSpan("http.request", {
      traceId,
      spanId,
      method: ctx.req.method,
      path: ctx.url.pathname,
    });
    let response: Response | undefined;
    let failure: unknown;
    let failed = false;
    try {
      response = await next();
    } catch (error) {
      failure = error;
      failed = true;
    } finally {
      const ended = now();
      const span: SpanRecord = {
        name: options.routeName?.(ctx) ?? `${ctx.req.method} ${ctx.url.pathname}`,
        traceId,
        spanId,
        ...(parent ? { parentSpanId: parent.spanId } : {}),
        sampled,
        startTime: started,
        endTime: ended,
        durationMs: ended - started,
        status: failed || (response?.status ?? 500) >= 500 ? "error" : "ok",
        attributes: {
          "service.name": options.serviceName ?? "wrnexus",
          "http.request.method": ctx.req.method,
          "url.path": ctx.url.pathname,
          "http.response.status_code": response?.status ?? 500,
        },
        ...(failed
          ? {
              error: {
                name: failure instanceof Error ? failure.name : "Error",
                message: failure instanceof Error ? failure.message : String(failure),
              },
            }
          : {}),
      };
      frameworkSpan.end(span.status, failure);
      if (sampled) {
        try {
          await options.onSpan?.(span);
          await options.exporter?.export([span]);
        } catch (error) {
          if (options.onExportError) await options.onExportError(error, span);
          else reportObservabilityFailure("trace export failed", error, options.diagnostic);
        }
      }
    }
    if (failed) throw failure;
    const headers = new Headers(response!.headers);
    headers.set("traceparent", formatTraceparent(trace));
    headers.set("x-request-id", String(ctx.locals.requestId));
    if (options.serverTiming !== false) {
      headers.append("server-timing", `trace;dur=${Math.max(0, now() - started).toFixed(2)}`);
    }
    return new Response(response!.body, {
      status: response!.status,
      statusText: response!.statusText,
      headers,
    });
  };
}

function otlpValue(value: string | number | boolean) {
  if (typeof value === "boolean") return { boolValue: value };
  if (typeof value === "number")
    return Number.isInteger(value) ? { intValue: String(value) } : { doubleValue: value };
  return { stringValue: value };
}

export function createOtlpTraceExporter(
  endpoint: string,
  options: { headers?: HeadersInit; fetch?: typeof fetch; serviceName?: string } = {},
): SpanExporter {
  const send = options.fetch ?? fetch;
  return {
    async export(spans) {
      if (!spans.length) return;
      const body = {
        resourceSpans: [
          {
            resource: {
              attributes: [
                { key: "service.name", value: { stringValue: options.serviceName ?? "wrnexus" } },
              ],
            },
            scopeSpans: [
              {
                scope: { name: "@wrnexus/observability", version: "0.8.0" },
                spans: spans.map((span) => ({
                  traceId: span.traceId,
                  spanId: span.spanId,
                  ...(span.parentSpanId ? { parentSpanId: span.parentSpanId } : {}),
                  name: span.name,
                  kind: 2,
                  startTimeUnixNano: String(BigInt(Math.trunc(span.startTime)) * 1_000_000n),
                  endTimeUnixNano: String(BigInt(Math.trunc(span.endTime)) * 1_000_000n),
                  attributes: Object.entries(span.attributes).map(([key, value]) => ({
                    key,
                    value: otlpValue(value),
                  })),
                  status: { code: span.status === "error" ? 2 : 1 },
                  ...(span.error
                    ? {
                        events: [
                          {
                            timeUnixNano: String(BigInt(Math.trunc(span.endTime)) * 1_000_000n),
                            name: "exception",
                            attributes: [
                              { key: "exception.type", value: { stringValue: span.error.name } },
                              {
                                key: "exception.message",
                                value: { stringValue: span.error.message },
                              },
                            ],
                          },
                        ],
                      }
                    : {}),
                })),
              },
            ],
          },
        ],
      };
      const response = await send(endpoint, {
        method: "POST",
        headers: {
          "content-type": "application/json",
          ...Object.fromEntries(new Headers(options.headers)),
        },
        body: JSON.stringify(body),
      });
      if (!response.ok) throw new Error(`OTLP trace export failed with ${response.status}.`);
    },
  };
}