promptfoo--promptfoo/test/evaluator.sigint.test.ts

import { randomUUID } from 'crypto';
import * as fs from 'fs';
import * as os from 'os';
import * as path from 'path';

import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import { updateSignalFile } from '../src/database/signal';
import { runDbMigrations } from '../src/migrate';
import Eval from '../src/models/eval';
import { JsonlFileWriter } from '../src/util/exportToFile/writeToFile';
import { createMockProvider } from './factories/provider';

import type { ApiProvider, TestSuite } from '../src/types';

// Only mock what we need to verify
vi.mock('../src/database/signal', () => ({
  updateSignalFile: vi.fn(),
  readSignalEvalId: vi.fn(),
}));

import { evaluate } from '../src/evaluator';
import { ResultFailureReason } from '../src/types';
import { createEmptyTokenUsage } from '../src/util/tokenUsageUtils';

describe('evaluate SIGINT/abort handling', () => {
  beforeEach(async () => {
    await runDbMigrations();
    vi.clearAllMocks();
    vi.mocked(updateSignalFile).mockReset();
  });

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

  it('should return early when user aborts (not max-duration timeout)', async () => {
    const abortController = new AbortController();
    let providerCallCount = 0;

    // Provider that aborts after first successful call
    const testProvider = createMockProvider({
      callApi: vi.fn<ApiProvider['callApi']>().mockImplementation(async () => {
        providerCallCount++;
        if (providerCallCount === 1) {
          // First call succeeds, then we trigger abort to simulate user SIGINT
          // Use setImmediate to abort after this call resolves
          setTimeout(() => abortController.abort(), 0);
          return {
            output: 'First response',
            tokenUsage: { total: 10, prompt: 5, completion: 5, cached: 0, numRequests: 1 },
          };
        }
        // Second call will see the abort
        throw new Error('Operation cancelled');
      }),
    });

    const mockAddResult = vi.fn().mockResolvedValue(undefined);
    const mockSetVars = vi.fn();
    const mockAddPrompts = vi.fn().mockResolvedValue(undefined);

    const mockEvalRecord = {
      id: 'test-eval-abort-123',
      results: [],
      prompts: [],
      persisted: false,
      config: {},
      addResult: mockAddResult,
      addPrompts: mockAddPrompts,
      fetchResultsByTestIdx: vi.fn().mockResolvedValue([]),
      getResults: vi.fn().mockResolvedValue([]),
      toEvaluateSummary: vi.fn().mockResolvedValue({
        results: [],
        prompts: [],
        stats: {
          successes: 1,
          failures: 0,
          errors: 0,
          tokenUsage: createEmptyTokenUsage(),
        },
      }),
      save: vi.fn().mockResolvedValue(undefined),
      setVars: mockSetVars,
      setDurationMs: vi.fn(),
    };

    const testSuite: TestSuite = {
      providers: [testProvider],
      prompts: [{ raw: 'Test prompt', label: 'Test prompt' }],
      tests: [{}, {}], // Two tests - second should see abort
    };

    const result = await evaluate(testSuite, mockEvalRecord as unknown as Eval, {
      abortSignal: abortController.signal,
    });

    // Should return the eval record
    expect(result).toBeDefined();
    expect(result.id).toBe('test-eval-abort-123');

    // Non-persisted evals are not visible to database watchers.
    expect(updateSignalFile).not.toHaveBeenCalled();

    // Should persist vars and prompts before early return
    expect(mockSetVars).toHaveBeenCalled();
    expect(mockAddPrompts).toHaveBeenCalled();
  });

  it('should continue normal flow and write timeout rows when per-call timeout triggers', async () => {
    // This test verifies that per-call timeouts (timeoutMs option) write
    // timeout error rows and continue the normal evaluation flow.
    //
    // Key difference from user SIGINT: evaluation continues normally rather than returning
    // early. The evaluator no longer writes the signal file directly — persisted evals emit a
    // refresh via Eval.addPrompts()/save(); a non-persisted eval (this test) writes no signal,
    // which is asserted below.
    let longTimer: NodeJS.Timeout | null = null;

    const slowProvider = createMockProvider({
      id: 'slow-provider',
      cleanup: true,
      callApi: vi
        .fn<ApiProvider['callApi']>()
        .mockImplementation((_prompt, _context, callApiOptions) => {
          // Long-running call that will be interrupted by timeout
          return new Promise((resolve, reject) => {
            const abortSignal = callApiOptions?.abortSignal;
            const onAbort = () => {
              if (longTimer) {
                clearTimeout(longTimer);
                longTimer = null;
              }
              const abortError = new Error('Operation aborted');
              abortError.name = 'AbortError';
              reject(abortError);
            };

            if (abortSignal?.aborted) {
              onAbort();
              return;
            }

            abortSignal?.addEventListener('abort', onAbort, { once: true });
            longTimer = setTimeout(() => {
              abortSignal?.removeEventListener('abort', onAbort);
              resolve({
                output: 'Slow response',
                tokenUsage: { total: 10, prompt: 5, completion: 5, cached: 0, numRequests: 1 },
              });
            }, 5000); // 5 seconds - will be interrupted by 100ms timeout
          });
        }),
    });

    const mockAddResult = vi.fn().mockResolvedValue(undefined);

    const mockEvalRecord = {
      id: 'test-eval-timeout-123',
      results: [],
      prompts: [],
      persisted: false,
      config: {},
      addResult: mockAddResult,
      addPrompts: vi.fn().mockResolvedValue(undefined),
      fetchResultsByTestIdx: vi.fn().mockResolvedValue([]),
      getResults: vi.fn().mockResolvedValue([]),
      toEvaluateSummary: vi.fn().mockResolvedValue({
        results: [],
        prompts: [],
        stats: {
          successes: 0,
          failures: 0,
          errors: 1,
          tokenUsage: createEmptyTokenUsage(),
        },
      }),
      save: vi.fn().mockResolvedValue(undefined),
      setVars: vi.fn(),
      setDurationMs: vi.fn(),
    };

    const testSuite: TestSuite = {
      providers: [slowProvider],
      prompts: [{ raw: 'Test prompt', label: 'Test prompt' }],
      tests: [{}],
    };

    try {
      // Use timeoutMs (per-call timeout) which triggers timeout for individual test cases
      // Unlike maxEvalTimeMs (max-duration), this doesn't abort the entire evaluation
      await evaluate(testSuite, mockEvalRecord as unknown as Eval, {
        timeoutMs: 100,
      });

      // When timeout triggers, should write timeout error row
      expect(mockAddResult).toHaveBeenCalledWith(
        expect.objectContaining({
          error: expect.stringContaining('timed out'),
          success: false,
          failureReason: ResultFailureReason.ERROR,
        }),
      );

      // Timeout should abort the in-flight call without invoking provider-wide cleanup
      expect(vi.mocked(slowProvider.callApi).mock.calls[0]?.[2]?.abortSignal?.aborted).toBe(true);
      expect(slowProvider.cleanup).not.toHaveBeenCalled();
    } finally {
      if (longTimer) {
        clearTimeout(longTimer);
      }
    }
  });

  it('should honor external abortSignal and return partial results', async () => {
    const abortController = new AbortController();
    const resultsAdded: unknown[] = [];

    const testProvider = createMockProvider({
      callApi: vi
        .fn<ApiProvider['callApi']>()
        .mockImplementation(async (_prompt, _context, opts) => {
          // Check if already aborted
          if (opts?.abortSignal?.aborted) {
            throw new Error('Operation cancelled');
          }

          // Abort after returning first result
          if (resultsAdded.length === 1) {
            abortController.abort();
            throw new Error('Operation cancelled');
          }

          return {
            output: `Response ${resultsAdded.length}`,
            tokenUsage: { total: 10, prompt: 5, completion: 5, cached: 0, numRequests: 1 },
          };
        }),
    });

    const mockAddResult = vi.fn().mockImplementation(async (result: unknown) => {
      resultsAdded.push(result);
    });

    const mockEvalRecord = {
      id: 'test-eval-partial-123',
      results: [],
      prompts: [],
      persisted: false,
      config: {},
      addResult: mockAddResult,
      addPrompts: vi.fn().mockResolvedValue(undefined),
      fetchResultsByTestIdx: vi.fn().mockResolvedValue([]),
      getResults: vi.fn().mockResolvedValue([]),
      toEvaluateSummary: vi.fn().mockResolvedValue({
        results: [],
        prompts: [],
        stats: {
          successes: resultsAdded.length,
          failures: 0,
          errors: 0,
          tokenUsage: createEmptyTokenUsage(),
        },
      }),
      save: vi.fn().mockResolvedValue(undefined),
      setVars: vi.fn(),
      setDurationMs: vi.fn(),
    };

    const testSuite: TestSuite = {
      providers: [testProvider],
      prompts: [{ raw: 'Test prompt', label: 'Test prompt' }],
      tests: [{}, {}, {}], // Three tests
    };

    await evaluate(testSuite, mockEvalRecord as unknown as Eval, {
      abortSignal: abortController.signal,
    });

    // Should have added at least one result before abort
    expect(resultsAdded.length).toBeGreaterThanOrEqual(1);

    // Non-persisted evals are not visible to database watchers.
    expect(updateSignalFile).not.toHaveBeenCalled();
  });

  it('should redact streamed JSONL rows when user aborts before final export', async () => {
    const outputPath = path.join(os.tmpdir(), `promptfoo-abort-${randomUUID()}.jsonl`);
    const abortController = new AbortController();
    let providerCallCount = 0;
    const provider = createMockProvider({
      callApi: vi.fn<ApiProvider['callApi']>().mockImplementation(async () => {
        providerCallCount++;
        if (providerCallCount === 1) {
          return {
            output: {
              http: {
                status: 200,
                statusText: 'OK',
                headers: {
                  authorization: 'model output should stay intact',
                },
              },
            },
            metadata: {
              headers: {
                'x-request-id': 'legacy_abort_should_not_persist',
                'x-safe-debug': 'keep-legacy',
              },
              http: {
                status: 200,
                statusText: 'OK',
                headers: {
                  'set-cookie': 'session=secret',
                  'x-request-id': 'req_should_not_persist',
                },
              },
            },
            tokenUsage: createEmptyTokenUsage(),
          };
        }
        throw new Error('Operation cancelled');
      }),
    });
    // Abort right after the first row streams to disk, so finalization never runs and the
    // on-disk artifact is the streamed (write-time-sanitized) row.
    const originalWrite = JsonlFileWriter.prototype.write;
    const writeSpy = vi
      .spyOn(JsonlFileWriter.prototype, 'write')
      .mockImplementation(async function (this: JsonlFileWriter, data) {
        await originalWrite.call(this, data);
        abortController.abort();
      });
    const evalRecord = new Eval({ outputPath });
    const testSuite: TestSuite = {
      providers: [provider],
      prompts: [{ raw: 'Test prompt', label: 'Test prompt' }],
      tests: [{ vars: { apiKey: 'sk-abort-vars-should-not-persist', safe: 'keep-me' } }, {}],
    };

    try {
      await evaluate(testSuite, evalRecord, {
        abortSignal: abortController.signal,
        maxConcurrency: 1,
      });

      const rows = fs
        .readFileSync(outputPath, 'utf8')
        .split(/\r?\n/)
        .filter(Boolean)
        .map((line) => JSON.parse(line));
      expect(rows).toHaveLength(1);
      expect(rows[0].response.metadata.headers).toEqual({
        'x-request-id': '[REDACTED]',
        'x-safe-debug': 'keep-legacy',
      });
      expect(rows[0].response.metadata.http.headers).toEqual({
        'set-cookie': '[REDACTED]',
        'x-request-id': '[REDACTED]',
      });
      expect(rows[0].response.output.http.headers.authorization).toBe(
        'model output should stay intact',
      );
      expect(rows[0].vars).toEqual({
        apiKey: '[REDACTED]',
        safe: 'keep-me',
      });
      expect(rows[0].testCase.vars).toEqual({
        apiKey: '[REDACTED]',
        safe: 'keep-me',
      });
      expect(JSON.stringify(rows[0])).not.toContain('session=secret');
      expect(JSON.stringify(rows[0])).not.toContain('req_should_not_persist');
      expect(JSON.stringify(rows[0])).not.toContain('legacy_abort_should_not_persist');
      expect(JSON.stringify(rows[0])).not.toContain('sk-abort-vars-should-not-persist');
    } finally {
      writeSpy.mockRestore();
      fs.rmSync(outputPath, { force: true });
    }
  });
});