promptfoo--promptfoo/test/evaluator/trace-integration.test.ts

import { afterEach, beforeEach, describe, expect, it, type Mock, vi } from 'vitest';
import { evaluate } from '../../src/evaluator';
import { nodeEvaluatorRuntime } from '../../src/node/evaluatorRuntime';
import * as evaluatorTracing from '../../src/tracing/evaluatorTracing';
import { getTraceStore } from '../../src/tracing/store';
import { createMockProvider } from '../factories/provider';

import type { EvaluatorRuntime } from '../../src/evaluator/runtime';
import type Eval from '../../src/models/eval';
import type EvalResult from '../../src/models/evalResult';
import type { EvaluateOptions, TestSuite } from '../../src/types/index';

// Mock dependencies
vi.mock('../../src/tracing/store');
const mockFlushOtel = vi.hoisted(() => vi.fn().mockResolvedValue(undefined));
const mockInitializeOtel = vi.hoisted(() => vi.fn());
const mockShutdownOtel = vi.hoisted(() => vi.fn().mockResolvedValue(undefined));

vi.mock('../../src/tracing/otelSdk', () => ({
  flushOtel: mockFlushOtel,
  initializeOtel: mockInitializeOtel,
  shutdownOtel: mockShutdownOtel,
}));

vi.mock('../../src/tracing/otlpReceiver', () => ({
  startOTLPReceiver: vi.fn(),
  stopOTLPReceiver: vi.fn(),
}));

// Mock evaluatorTracing module
vi.mock('../../src/tracing/evaluatorTracing', () => ({
  generateTraceId: vi.fn(() => 'abcdef1234567890abcdef1234567890'),
  generateSpanId: vi.fn(() => '0123456789abcdef'),
  generateTraceparent: vi.fn((traceId, spanId) => `00-${traceId}-${spanId}-01`),
  generateTraceContextIfNeeded: vi.fn(),
  startOtlpReceiverIfNeeded: vi.fn(),
  stopOtlpReceiverIfNeeded: vi.fn(),
  isOtlpReceiverStarted: vi.fn(() => false),
  isTracingEnabled: vi.fn((test) => test.metadata?.tracingEnabled === true),
}));

describe('evaluator trace integration', () => {
  const mockTraceStore = {
    createTrace: vi.fn(),
    getTrace: vi.fn(),
  };

  const mockEval = {
    id: 'test-eval-id',
    addResult: vi.fn(),
    addPrompts: vi.fn(),
    fetchResultsByTestIdx: vi.fn(),
    setVars: vi.fn(),
    setDurationMs: vi.fn(),
    results: [],
    prompts: [],
    persisted: false,
    config: {
      outputPath: undefined,
    },
  } as unknown as Eval;

  beforeEach(() => {
    vi.resetAllMocks();
    (getTraceStore as Mock).mockReturnValue(mockTraceStore);
  });

  afterEach(() => {
    vi.restoreAllMocks();
  });

  it('should pass traceId through to assertions when tracing is enabled', async () => {
    // Mock trace creation and retrieval
    const testTraceId = 'abcdef1234567890abcdef1234567890';
    mockTraceStore.createTrace.mockResolvedValue(undefined);
    mockTraceStore.getTrace.mockResolvedValue({
      traceId: testTraceId,
      spans: [
        {
          spanId: 'test-span',
          name: 'test.operation',
          startTime: 1000,
          endTime: 2000,
        },
      ],
    });

    // Mock generateTraceContextIfNeeded
    vi.mocked(evaluatorTracing.generateTraceContextIfNeeded).mockResolvedValue({
      traceparent: `00-${testTraceId}-0123456789abcdef-01`,
      evaluationId: 'test-eval-id',
      testCaseId: 'test-case-id',
    });

    const testSuite: TestSuite = {
      providers: [
        createMockProvider({
          id: 'mock-provider',
          response: { output: 'Test response', tokenUsage: {} },
        }),
      ],
      prompts: [{ raw: 'Test prompt', label: 'test' }],
      tests: [
        {
          vars: { input: 'test' },
          metadata: {
            tracingEnabled: true,
            evaluationId: 'test-eval-id',
          },
          assert: [
            {
              type: 'javascript',
              value: `
                // Verify trace data is available
                if (!context.trace) return false;
                return context.trace.spans.length > 0 &&
                       context.trace.spans[0].name === 'test.operation';
              `,
            },
          ],
        },
      ],
      tracing: {
        enabled: true,
        otlp: {
          http: {
            enabled: true,
            port: 4318,
            host: '0.0.0.0',
            acceptFormats: ['protobuf'],
          },
        },
      },
    };

    const options: EvaluateOptions = {
      maxConcurrency: 1,
    };

    // Run evaluation
    await evaluate(testSuite, mockEval, options);

    // Verify trace context was generated
    expect(evaluatorTracing.generateTraceContextIfNeeded).toHaveBeenCalled();
    expect(evaluatorTracing.startOtlpReceiverIfNeeded).toHaveBeenCalledWith(
      testSuite,
      'test-eval-id',
    );

    // Verify trace was fetched for assertion
    expect(mockTraceStore.getTrace).toHaveBeenCalledWith(testTraceId, {
      sanitizeAttributes: false,
    });
    expect(mockFlushOtel).toHaveBeenCalled();
    expect(mockShutdownOtel).toHaveBeenCalledOnce();

    // Verify result was added with passing assertion
    expect(mockEval.addResult).toHaveBeenCalledWith(
      expect.objectContaining({
        success: true,
        score: 1,
      }),
    );
  });

  it('closes writers but skips OTEL shutdown when initialization fails', async () => {
    const writer = {
      write: vi.fn().mockResolvedValue(undefined),
      close: vi.fn().mockResolvedValue(undefined),
    };
    const runtime: EvaluatorRuntime<Eval, EvalResult> = {
      createEvaluationStore: nodeEvaluatorRuntime.createEvaluationStore,
      createResultWriters: vi.fn().mockReturnValue([writer]),
    };
    mockInitializeOtel.mockImplementationOnce(() => {
      throw new Error('otel unavailable');
    });

    await expect(
      evaluate(
        {
          providers: [],
          prompts: [],
          tests: [],
          tracing: { enabled: true },
        },
        mockEval,
        {},
        runtime,
      ),
    ).rejects.toThrow('otel unavailable');

    expect(writer.close).toHaveBeenCalledOnce();
    expect(mockFlushOtel).not.toHaveBeenCalled();
    expect(mockShutdownOtel).not.toHaveBeenCalled();
  });

  it('flushes and shuts down OTEL after successful tracing initialization', async () => {
    await evaluate(
      {
        providers: [
          createMockProvider({
            id: 'mock-provider',
            response: { output: 'Test response', tokenUsage: {} },
          }),
        ],
        prompts: [{ raw: 'Test prompt', label: 'test' }],
        tests: [{}],
        tracing: { enabled: true },
      },
      mockEval,
      {},
    );

    expect(mockInitializeOtel).toHaveBeenCalledOnce();
    expect(mockFlushOtel).toHaveBeenCalledOnce();
    expect(mockShutdownOtel).toHaveBeenCalledOnce();
  });

  it('should handle assertions gracefully when tracing is disabled', async () => {
    // Mock generateTraceContextIfNeeded to return null when tracing is disabled
    vi.mocked(evaluatorTracing.generateTraceContextIfNeeded).mockResolvedValue(null);

    const testSuite: TestSuite = {
      providers: [
        createMockProvider({
          id: 'mock-provider',
          response: { output: 'Test response', tokenUsage: {} },
        }),
      ],
      prompts: [{ raw: 'Test prompt', label: 'test' }],
      tests: [
        {
          vars: { input: 'test' },
          // No tracingEnabled in metadata
          assert: [
            {
              type: 'javascript',
              value: `
                // Should pass when trace is undefined
                return context.trace === undefined && output === 'Test response';
              `,
            },
          ],
        },
      ],
      // Tracing not enabled in test suite
    };

    const options: EvaluateOptions = {
      maxConcurrency: 1,
    };

    // Run evaluation
    await evaluate(testSuite, mockEval, options);

    // Verify trace was NOT created or fetched
    expect(mockTraceStore.createTrace).not.toHaveBeenCalled();
    expect(mockTraceStore.getTrace).not.toHaveBeenCalled();
    expect(mockFlushOtel).not.toHaveBeenCalled();
    expect(mockShutdownOtel).not.toHaveBeenCalled();

    // Verify result was added with passing assertion
    expect(mockEval.addResult).toHaveBeenCalledWith(
      expect.objectContaining({
        success: true,
        score: 1,
      }),
    );
  });

  it('should extract traceId correctly from traceparent header', async () => {
    const testTraceId = '0af7651916cd43dd8448eb211c80319c';
    const testSpanId = 'b7ad6b7169203331';

    // Mock the trace context generation
    vi.mocked(evaluatorTracing.generateTraceContextIfNeeded).mockResolvedValue({
      traceparent: `00-${testTraceId}-${testSpanId}-01`,
      evaluationId: 'test-eval-id',
      testCaseId: 'test-case-id',
    });

    mockTraceStore.createTrace.mockResolvedValue(undefined);
    mockTraceStore.getTrace.mockResolvedValue({
      traceId: testTraceId,
      spans: [
        {
          spanId: 'test-span',
          name: 'extracted.correctly',
          startTime: 1000,
          endTime: 2000,
        },
      ],
    });

    const testSuite: TestSuite = {
      providers: [
        createMockProvider({
          id: 'mock-provider',
          response: { output: 'Test response', tokenUsage: {} },
        }),
      ],
      prompts: [{ raw: 'Test prompt', label: 'test' }],
      tests: [
        {
          vars: { input: 'test' },
          metadata: {
            tracingEnabled: true,
            evaluationId: 'test-eval-id',
          },
          assert: [
            {
              type: 'javascript',
              value: `
                // Verify the extracted traceId matches
                return context.trace && context.trace.traceId === '${testTraceId}';
              `,
            },
          ],
        },
      ],
    };

    const options: EvaluateOptions = {
      maxConcurrency: 1,
    };

    // Run evaluation
    await evaluate(testSuite, mockEval, options);

    // Verify trace was fetched with the correct traceId
    expect(mockTraceStore.getTrace).toHaveBeenCalledWith(testTraceId, {
      sanitizeAttributes: false,
    });

    // Verify result was added with passing assertion
    expect(mockEval.addResult).toHaveBeenCalledWith(
      expect.objectContaining({
        success: true,
      }),
    );
  });

  it('should still run evaluator cleanup when receiver startup is fatal', async () => {
    vi.mocked(evaluatorTracing.startOtlpReceiverIfNeeded).mockRejectedValueOnce(
      new Error('receiver failed'),
    );

    const testSuite: TestSuite = {
      providers: [],
      prompts: [],
      tests: [],
      tracing: {
        enabled: true,
        failOnReceiverStartFailure: true,
        otlp: {
          http: {
            enabled: true,
            port: 4318,
            host: '127.0.0.1',
          },
        },
      },
    };

    await expect(evaluate(testSuite, mockEval, {})).rejects.toThrow('receiver failed');
    expect(evaluatorTracing.stopOtlpReceiverIfNeeded).toHaveBeenCalled();
    expect(mockFlushOtel).not.toHaveBeenCalled();
    expect(mockShutdownOtel).not.toHaveBeenCalled();
  });
});