agentic-dev/packages/core-audit/src/trace-id-enriching-audit-log.test.ts

import { describe, it, expect, beforeEach, afterEach, vi } from "vitest";
import { context, trace } from "@opentelemetry/api";
import { AsyncLocalStorageContextManager } from "@opentelemetry/context-async-hooks";
import { BasicTracerProvider, InMemorySpanExporter, SimpleSpanProcessor } from "@opentelemetry/sdk-trace-base";
import { TraceIdEnrichingAuditLog } from "./trace-id-enriching-audit-log";
import type { AuditEntry } from "@repo/core-shared/audit";
import type { IAuditLog } from "./audit-log.interface";

// Register async context manager so startActiveSpan propagates context.
const ctxManager = new AsyncLocalStorageContextManager();
ctxManager.enable();
context.setGlobalContextManager(ctxManager);

function setupProvider(): { exporter: InMemorySpanExporter; provider: BasicTracerProvider } {
  const exporter = new InMemorySpanExporter();
  const provider = new BasicTracerProvider({
    spanProcessors: [new SimpleSpanProcessor(exporter)],
  });
  trace.setGlobalTracerProvider(provider);
  return { exporter, provider };
}

const sample: AuditEntry = {
  actorId: "user_1",
  actorType: "user",
  actorRoles: [],
  action: "VIEW",
  resource: { type: "articles" },
  at: new Date(),
  scope: { feature: "blog", environment: "test", tenant: "default" },
  from: { ipTruncated: "10.0.0.0", userAgent: "test" },
  containsPii: false,
  outcome: "success",
};

function makeInner(): IAuditLog & { records: AuditEntry[] } {
  const records: AuditEntry[] = [];
  return {
    records,
    async record(e) {
      records.push(e);
    },
    eraseSubject: vi.fn(),
  };
}

describe("TraceIdEnrichingAuditLog", () => {
  let exporter: InMemorySpanExporter;
  let provider: BasicTracerProvider;

  beforeEach(() => {
    ({ exporter, provider } = setupProvider());
  });

  afterEach(async () => {
    await provider.shutdown();
    trace.disable();
    void exporter;
  });

  it("passes through when no active span", async () => {
    const inner = makeInner();
    const wrapper = new TraceIdEnrichingAuditLog(inner);
    await wrapper.record(sample);
    expect(inner.records[0]!.correlationId).toBeUndefined();
  });

  it("auto-populates correlationId from active span", async () => {
    const inner = makeInner();
    const wrapper = new TraceIdEnrichingAuditLog(inner);
    const tracer = trace.getTracer("test");
    await new Promise<void>((resolve) => {
      tracer.startActiveSpan("test", async (span) => {
        await wrapper.record(sample);
        const expected = span.spanContext().traceId;
        expect(inner.records[0]!.correlationId).toBe(expected);
        span.end();
        resolve();
      });
    });
  });

  it("explicit correlationId wins over auto-populated", async () => {
    const inner = makeInner();
    const wrapper = new TraceIdEnrichingAuditLog(inner);
    const tracer = trace.getTracer("test");
    await new Promise<void>((resolve) => {
      tracer.startActiveSpan("test", async (span) => {
        await wrapper.record({ ...sample, correlationId: "explicit-trace-id" });
        expect(inner.records[0]!.correlationId).toBe("explicit-trace-id");
        span.end();
        resolve();
      });
    });
  });

  it("eraseSubject passes through unchanged", async () => {
    const eraseSpy = vi.fn();
    const inner: IAuditLog = { record: vi.fn(), eraseSubject: eraseSpy };
    const wrapper = new TraceIdEnrichingAuditLog(inner);
    await wrapper.eraseSubject("user_1", "delete");
    expect(eraseSpy).toHaveBeenCalledWith("user_1", "delete");
  });
});