chore: import upstream snapshot with attribution

test/tracing/store.test.ts

@@ -0,0 +1,711 @@
+import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest';
+import { mockGlobal } from '../util/utils';
+
+const mockRandomUUID = vi.fn(() => 'test-uuid');
+const restoreCrypto = mockGlobal('crypto', {
+  ...crypto,
+  randomUUID: mockRandomUUID,
+} as Crypto);
+
+afterAll(() => {
+  restoreCrypto();
+});
+
+// Mock logger
+vi.mock('../../src/logger', () => ({
+  default: {
+    debug: vi.fn(),
+    error: vi.fn(),
+    warn: vi.fn(),
+  },
+}));
+
+// Dynamic import after mocking - initialized in beforeAll
+let TraceStore: typeof import('../../src/tracing/store').TraceStore;
+
+describe('TraceStore', () => {
+  let traceStore: InstanceType<typeof TraceStore>;
+  let mockDb: any;
+  let mockDeleteChain: any;
+
+  beforeAll(async () => {
+    const mod = await import('../../src/tracing/store');
+    TraceStore = mod.TraceStore;
+  });
+
+  beforeEach(async () => {
+    // Reset the UUID mock
+    mockRandomUUID.mockReturnValue('test-uuid');
+
+    // Create mock database methods that properly chain
+    const mockInsertChain = {
+      values: vi.fn().mockReturnThis(),
+      onConflictDoNothing: vi.fn().mockReturnThis(),
+      run: vi.fn().mockResolvedValue(undefined),
+    };
+    const mockSelectChain = {
+      from: vi.fn().mockReturnThis(),
+      where: vi.fn().mockReturnThis(),
+      limit: vi.fn(() => Promise.resolve([])),
+    };
+    mockDeleteChain = {
+      where: vi.fn().mockReturnThis(),
+      run: vi.fn().mockResolvedValue(undefined),
+    };
+
+    mockDb = {
+      insert: vi.fn(() => mockInsertChain),
+      select: vi.fn(() => mockSelectChain),
+      delete: vi.fn(() => mockDeleteChain),
+      transaction: vi.fn(async (callback) => callback(mockDb)),
+    };
+
+    // Create trace store and inject mock DB
+    traceStore = new TraceStore();
+    // Use private property access to inject the mock
+    (traceStore as any).db = mockDb;
+  });
+
+  afterEach(() => {
+    vi.clearAllMocks();
+  });
+
+  describe('createTrace', () => {
+    it('should create a new trace record', async () => {
+      const traceData = {
+        traceId: 'test-trace-id',
+        evaluationId: 'test-eval-id',
+        testCaseId: 'test-case-id',
+        metadata: { test: 'data' },
+      };
+
+      await traceStore.createTrace(traceData);
+
+      expect(mockDb.insert).toHaveBeenCalledWith(expect.anything());
+      expect(mockDb.insert().values).toHaveBeenCalledWith({
+        id: 'test-uuid',
+        traceId: 'test-trace-id',
+        evaluationId: 'test-eval-id',
+        testCaseId: 'test-case-id',
+        createdAt: expect.any(Number),
+        metadata: { test: 'data' },
+      });
+      expect(mockDb.insert().values().onConflictDoNothing).toHaveBeenCalledWith(
+        expect.objectContaining({
+          target: expect.anything(),
+        }),
+      );
+    });
+
+    it('should handle errors when creating trace', async () => {
+      const error = new Error('Database error');
+      mockDb.insert().values().onConflictDoNothing().run.mockRejectedValueOnce(error);
+
+      const traceData = {
+        traceId: 'test-trace-id',
+        evaluationId: 'test-eval-id',
+        testCaseId: 'test-case-id',
+      };
+
+      await expect(traceStore.createTrace(traceData)).rejects.toThrow('Database error');
+    });
+  });
+
+  describe('getTraceMetadata', () => {
+    // This three-way contract is the single source of truth for ingest-time redaction
+    // (otlpReceiver.getRedactAttributePatterns): object => use stored policy, {} => no
+    // policy (no redaction), undefined => no trace row (fall back to receiver default).
+    it('returns the stored metadata object when the trace row has metadata', async () => {
+      mockDb
+        .select()
+        .from()
+        .where()
+        .limit.mockResolvedValueOnce([
+          { metadata: { otlpHttpRedactAttributes: ['authorization'] } },
+        ]);
+
+      await expect(traceStore.getTraceMetadata('trace-1')).resolves.toEqual({
+        otlpHttpRedactAttributes: ['authorization'],
+      });
+    });
+
+    it('returns an empty object (not undefined) when the row exists but metadata is null', async () => {
+      mockDb
+        .select()
+        .from()
+        .where()
+        .limit.mockResolvedValueOnce([{ metadata: null }]);
+
+      await expect(traceStore.getTraceMetadata('trace-2')).resolves.toEqual({});
+    });
+
+    it('returns undefined when no trace row exists', async () => {
+      mockDb.select().from().where().limit.mockResolvedValueOnce([]);
+
+      await expect(traceStore.getTraceMetadata('trace-3')).resolves.toBeUndefined();
+    });
+  });
+
+  describe('addSpans', () => {
+    it('should add spans to an existing trace', async () => {
+      // Mock trace exists check
+      mockDb
+        .select()
+        .from()
+        .where()
+        .limit.mockResolvedValueOnce([{ traceId: 'test-trace-id' }]);
+
+      const spans = [
+        {
+          spanId: 'span-1',
+          name: 'operation-1',
+          startTime: 1000,
+          endTime: 2000,
+          attributes: { key: 'value' },
+        },
+        {
+          spanId: 'span-2',
+          parentSpanId: 'span-1',
+          name: 'operation-2',
+          startTime: 1100,
+          endTime: 1900,
+          statusCode: 0,
+          statusMessage: 'OK',
+        },
+      ];
+
+      await traceStore.addSpans('test-trace-id', spans);
+
+      expect(mockDb.select).toHaveBeenCalledWith();
+      expect(mockDb.select().from).toHaveBeenCalledWith(expect.anything());
+      expect(mockDb.select().from().where).toHaveBeenCalledWith(expect.anything());
+
+      expect(mockDb.insert).toHaveBeenCalledWith(expect.anything());
+      expect(mockDb.insert().values).toHaveBeenCalledWith([
+        {
+          id: 'test-uuid',
+          traceId: 'test-trace-id',
+          spanId: 'span-1',
+          parentSpanId: undefined,
+          name: 'operation-1',
+          startTime: 1000,
+          endTime: 2000,
+          attributes: { key: 'value' },
+          statusCode: undefined,
+          statusMessage: undefined,
+        },
+        {
+          id: 'test-uuid',
+          traceId: 'test-trace-id',
+          spanId: 'span-2',
+          parentSpanId: 'span-1',
+          name: 'operation-2',
+          startTime: 1100,
+          endTime: 1900,
+          attributes: undefined,
+          statusCode: 0,
+          statusMessage: 'OK',
+        },
+      ]);
+    });
+
+    it('should skip spans if trace does not exist', async () => {
+      // Mock trace does not exist
+      mockDb.select().from().where().limit.mockResolvedValueOnce([]);
+
+      const spans = [
+        {
+          spanId: 'span-1',
+          name: 'operation-1',
+          startTime: 1000,
+        },
+      ];
+
+      await traceStore.addSpans('non-existent-trace', spans);
+
+      // Should check for trace existence
+      expect(mockDb.select).toHaveBeenCalledWith();
+
+      // Should not insert spans
+      expect(mockDb.insert).not.toHaveBeenCalled();
+    });
+
+    it('should handle errors when adding spans', async () => {
+      // Mock trace exists
+      mockDb
+        .select()
+        .from()
+        .where()
+        .limit.mockResolvedValueOnce([{ traceId: 'test-trace-id' }]);
+
+      // Mock insert error
+      const error = new Error('Insert failed');
+      mockDb.insert().values().run.mockRejectedValueOnce(error);
+
+      const spans = [
+        {
+          spanId: 'span-1',
+          name: 'operation-1',
+          startTime: 1000,
+        },
+      ];
+
+      await expect(traceStore.addSpans('test-trace-id', spans)).rejects.toThrow('Insert failed');
+    });
+  });
+
+  describe('getTracesByEvaluation', () => {
+    it('should retrieve all traces for an evaluation', async () => {
+      const mockTraces = [
+        { id: '1', traceId: 'trace-1', evaluationId: 'eval-1', testCaseId: 'test-case-1' },
+        { id: '2', traceId: 'trace-2', evaluationId: 'eval-1', testCaseId: 'test-case-2' },
+      ];
+
+      const mockSpans = {
+        'trace-1': [
+          {
+            id: '1',
+            traceId: 'trace-1',
+            spanId: 'span-1-1',
+            parentSpanId: null,
+            name: 'span one',
+            startTime: 1000,
+            endTime: null,
+            attributes: null,
+            statusCode: null,
+            statusMessage: null,
+          },
+          {
+            id: '2',
+            traceId: 'trace-1',
+            spanId: 'span-1-2',
+            parentSpanId: 'span-1-1',
+            name: 'span two',
+            startTime: 2000,
+            endTime: 2500,
+            attributes: { foo: 'bar' },
+            statusCode: 1,
+            statusMessage: 'ok',
+          },
+        ],
+        'trace-2': [
+          {
+            id: '3',
+            traceId: 'trace-2',
+            spanId: 'span-2-1',
+            parentSpanId: null,
+            name: 'span three',
+            startTime: 3000,
+            endTime: null,
+            attributes: null,
+            statusCode: null,
+            statusMessage: null,
+          },
+        ],
+      };
+
+      // Set up the main traces query
+      const tracesSelectChain = {
+        from: vi.fn().mockReturnThis(),
+        where: vi.fn(() => Promise.resolve(mockTraces)),
+      };
+
+      // Set up span queries for each trace
+      const spanQuery1 = {
+        from: vi.fn().mockReturnThis(),
+        where: vi.fn(() => Promise.resolve(mockSpans['trace-1'])),
+      };
+      const spanQuery2 = {
+        from: vi.fn().mockReturnThis(),
+        where: vi.fn(() => Promise.resolve(mockSpans['trace-2'])),
+      };
+
+      // Mock the select calls in sequence: first for traces, then for each trace's spans
+      vi.spyOn(mockDb, 'select')
+        .mockImplementation(() => ({}))
+        .mockReturnValueOnce(tracesSelectChain)
+        .mockReturnValueOnce(spanQuery1)
+        .mockReturnValueOnce(spanQuery2);
+
+      const result = await traceStore.getTracesByEvaluation('eval-1');
+
+      expect(result).toHaveLength(2);
+      expect(result[0]).toEqual({
+        traceId: 'trace-1',
+        evaluationId: 'eval-1',
+        testCaseId: 'test-case-1',
+        metadata: undefined,
+        spans: [
+          {
+            spanId: 'span-1-1',
+            parentSpanId: undefined,
+            name: 'span one',
+            startTime: 1000,
+            endTime: undefined,
+            attributes: undefined,
+            statusCode: undefined,
+            statusMessage: undefined,
+          },
+          {
+            spanId: 'span-1-2',
+            parentSpanId: 'span-1-1',
+            name: 'span two',
+            startTime: 2000,
+            endTime: 2500,
+            attributes: { foo: 'bar' },
+            statusCode: 1,
+            statusMessage: 'ok',
+          },
+        ],
+      });
+      expect(result[1]).toEqual({
+        traceId: 'trace-2',
+        evaluationId: 'eval-1',
+        testCaseId: 'test-case-2',
+        metadata: undefined,
+        spans: [
+          {
+            spanId: 'span-2-1',
+            parentSpanId: undefined,
+            name: 'span three',
+            startTime: 3000,
+            endTime: undefined,
+            attributes: undefined,
+            statusCode: undefined,
+            statusMessage: undefined,
+          },
+        ],
+      });
+    });
+
+    it('should redact sensitive span attributes for an evaluation', async () => {
+      const mockTrace = {
+        id: '1',
+        traceId: 'trace-1',
+        evaluationId: 'eval-1',
+        testCaseId: 'test-case-1',
+        metadata: null,
+      };
+      const mockSpans = [
+        {
+          id: '1',
+          traceId: 'trace-1',
+          spanId: 'span-1',
+          parentSpanId: null,
+          name: 'http request',
+          startTime: 1000,
+          endTime: null,
+          attributes: {
+            authorization: 'Bearer trace-secret',
+            'x-api-key': 'gateway-secret',
+            headers: {
+              'api-key': 'header-secret',
+              set_cookie: 'session=trace-secret',
+              accept: 'application/json',
+            },
+            nested: {
+              api_token: 'nested-secret',
+              safe: 'ok',
+            },
+            'gen_ai.request.max_tokens': 4096,
+            'gen_ai.usage.input_tokens': 100,
+            'gen_ai.usage.output_tokens': 50,
+            'gen_ai.usage.total_tokens': 150,
+          },
+          statusCode: null,
+          statusMessage: null,
+        },
+      ];
+
+      const tracesSelectChain = {
+        from: vi.fn().mockReturnThis(),
+        where: vi.fn(() => Promise.resolve([mockTrace])),
+      };
+      const spanQuery = {
+        from: vi.fn().mockReturnThis(),
+        where: vi.fn(() => Promise.resolve(mockSpans)),
+      };
+
+      vi.spyOn(mockDb, 'select')
+        .mockImplementation(() => ({}))
+        .mockReturnValueOnce(tracesSelectChain)
+        .mockReturnValueOnce(spanQuery);
+
+      const result = await traceStore.getTracesByEvaluation('eval-1');
+
+      expect(result[0].spans[0].attributes).toEqual({
+        authorization: '<redacted>',
+        'x-api-key': '<redacted>',
+        headers: {
+          'api-key': '<redacted>',
+          set_cookie: '<redacted>',
+          accept: 'application/json',
+        },
+        nested: {
+          api_token: '<redacted>',
+          safe: 'ok',
+        },
+        'gen_ai.request.max_tokens': 4096,
+        'gen_ai.usage.input_tokens': 100,
+        'gen_ai.usage.output_tokens': 50,
+        'gen_ai.usage.total_tokens': 150,
+      });
+    });
+
+    it('should allow callers to retrieve raw span attributes for an evaluation', async () => {
+      const mockTrace = {
+        id: '1',
+        traceId: 'trace-1',
+        evaluationId: 'eval-1',
+        testCaseId: 'test-case-1',
+        metadata: null,
+      };
+      const mockSpans = [
+        {
+          id: '1',
+          traceId: 'trace-1',
+          spanId: 'span-1',
+          parentSpanId: null,
+          name: 'http request',
+          startTime: 1000,
+          endTime: null,
+          attributes: {
+            authorization: 'Bearer trace-secret',
+            'x-api-key': 'gateway-secret',
+            headers: {
+              'api-key': 'header-secret',
+              set_cookie: 'session=trace-secret',
+              accept: 'application/json',
+            },
+          },
+          statusCode: null,
+          statusMessage: null,
+        },
+      ];
+
+      const tracesSelectChain = {
+        from: vi.fn().mockReturnThis(),
+        where: vi.fn(() => Promise.resolve([mockTrace])),
+      };
+      const spanQuery = {
+        from: vi.fn().mockReturnThis(),
+        where: vi.fn(() => Promise.resolve(mockSpans)),
+      };
+
+      vi.spyOn(mockDb, 'select')
+        .mockImplementation(() => ({}))
+        .mockReturnValueOnce(tracesSelectChain)
+        .mockReturnValueOnce(spanQuery);
+
+      const result = await traceStore.getTracesByEvaluation('eval-1', {
+        sanitizeAttributes: false,
+      });
+
+      expect(result[0].spans[0].attributes).toEqual(mockSpans[0].attributes);
+    });
+
+    it('should return empty array if no traces found', async () => {
+      mockDb.select().from().where.mockResolvedValueOnce([]);
+
+      const result = await traceStore.getTracesByEvaluation('non-existent-eval');
+
+      expect(result).toEqual([]);
+    });
+  });
+
+  describe('getTrace', () => {
+    it('should retrieve a single trace with spans', async () => {
+      const mockTrace = {
+        id: '1',
+        traceId: 'trace-1',
+        evaluationId: 'eval-1',
+        testCaseId: 'test-case-1',
+        metadata: { source: 'test' },
+      };
+      const mockSpans = [
+        {
+          id: '1',
+          traceId: 'trace-1',
+          spanId: 'span-1',
+          parentSpanId: null,
+          name: 'span one',
+          startTime: 1000,
+          endTime: null,
+          attributes: null,
+          statusCode: null,
+          statusMessage: null,
+        },
+        {
+          id: '2',
+          traceId: 'trace-1',
+          spanId: 'span-2',
+          parentSpanId: 'span-1',
+          name: 'span two',
+          startTime: 2000,
+          endTime: 2500,
+          attributes: {
+            foo: 'bar',
+            authorization: 'Bearer trace-secret',
+            headers: {
+              Cookie: 'session=trace-secret',
+              'x-api-key': 'gateway-secret',
+              accept: 'application/json',
+            },
+            nested: {
+              'api-key': 'nested-header-secret',
+              api_key: 'nested-secret',
+              safe: 'ok',
+            },
+          },
+          statusCode: 1,
+          statusMessage: 'ok',
+        },
+      ];
+
+      // Mock trace query - update the select chain to include limit
+      const traceSelectChain = {
+        from: vi.fn().mockReturnThis(),
+        where: vi.fn().mockReturnThis(),
+        limit: vi.fn(() => Promise.resolve([mockTrace])),
+      };
+
+      // Mock spans query
+      const spanQuery = {
+        from: vi.fn().mockReturnThis(),
+        where: vi.fn(() => Promise.resolve(mockSpans)),
+      };
+
+      vi.spyOn(mockDb, 'select')
+        .mockImplementation(() => ({}))
+        .mockReturnValueOnce(traceSelectChain)
+        .mockReturnValueOnce(spanQuery);
+
+      const result = await traceStore.getTrace('trace-1');
+
+      expect(result).toEqual({
+        traceId: 'trace-1',
+        evaluationId: 'eval-1',
+        testCaseId: 'test-case-1',
+        metadata: { source: 'test' },
+        spans: [
+          {
+            spanId: 'span-1',
+            parentSpanId: undefined,
+            name: 'span one',
+            startTime: 1000,
+            endTime: undefined,
+            attributes: undefined,
+            statusCode: undefined,
+            statusMessage: undefined,
+          },
+          {
+            spanId: 'span-2',
+            parentSpanId: 'span-1',
+            name: 'span two',
+            startTime: 2000,
+            endTime: 2500,
+            attributes: {
+              foo: 'bar',
+              authorization: '<redacted>',
+              headers: {
+                Cookie: '<redacted>',
+                'x-api-key': '<redacted>',
+                accept: 'application/json',
+              },
+              nested: {
+                'api-key': '<redacted>',
+                api_key: '<redacted>',
+                safe: 'ok',
+              },
+            },
+            statusCode: 1,
+            statusMessage: 'ok',
+          },
+        ],
+      });
+    });
+
+    it('should allow callers to retrieve raw span attributes for a single trace', async () => {
+      const mockTrace = {
+        id: '1',
+        traceId: 'trace-1',
+        evaluationId: 'eval-1',
+        testCaseId: 'test-case-1',
+        metadata: null,
+      };
+      const mockSpans = [
+        {
+          id: '1',
+          traceId: 'trace-1',
+          spanId: 'span-1',
+          parentSpanId: null,
+          name: 'tool call',
+          startTime: 1000,
+          endTime: null,
+          attributes: {
+            api_key: 'trace-aware-assertion-secret',
+            token: 'trace-aware-token',
+            arguments: { api_key: 'nested-secret' },
+          },
+          statusCode: null,
+          statusMessage: null,
+        },
+      ];
+
+      const traceSelectChain = {
+        from: vi.fn().mockReturnThis(),
+        where: vi.fn().mockReturnThis(),
+        limit: vi.fn(() => Promise.resolve([mockTrace])),
+      };
+      const spanQuery = {
+        from: vi.fn().mockReturnThis(),
+        where: vi.fn(() => Promise.resolve(mockSpans)),
+      };
+
+      vi.spyOn(mockDb, 'select')
+        .mockImplementation(() => ({}))
+        .mockReturnValueOnce(traceSelectChain)
+        .mockReturnValueOnce(spanQuery);
+
+      const result = await traceStore.getTrace('trace-1', { sanitizeAttributes: false });
+
+      expect(result?.spans[0].attributes).toEqual(mockSpans[0].attributes);
+    });
+
+    it('should return null if trace not found', async () => {
+      // Mock trace query - update the select chain to include limit
+      const traceSelectChain = {
+        from: vi.fn().mockReturnThis(),
+        where: vi.fn().mockReturnThis(),
+        limit: vi.fn(() => Promise.resolve([])),
+      };
+      vi.spyOn(mockDb, 'select')
+        .mockImplementation(() => ({}))
+        .mockReturnValue(traceSelectChain);
+
+      const result = await traceStore.getTrace('non-existent-trace');
+
+      expect(result).toBeNull();
+    });
+  });
+
+  describe('deleteOldTraces', () => {
+    it('should delete dependent spans before traces older than retention period', async () => {
+      const retentionDays = 30;
+
+      await traceStore.deleteOldTraces(retentionDays);
+
+      expect(mockDb.transaction).toHaveBeenCalledWith(expect.any(Function));
+      expect(mockDb.delete).toHaveBeenCalledTimes(2);
+      expect(mockDeleteChain.where).toHaveBeenCalledTimes(2);
+      expect(mockDeleteChain.run).toHaveBeenCalledTimes(2);
+    });
+
+    it('should handle errors when deleting old traces', async () => {
+      const error = new Error('Delete failed');
+      mockDeleteChain.run.mockRejectedValueOnce(error);
+
+      await expect(traceStore.deleteOldTraces(30)).rejects.toThrow('Delete failed');
+    });
+  });
+});