chore: import upstream snapshot with attribution

test/main.test.ts

@@ -0,0 +1,654 @@
+import fs from 'node:fs';
+import os from 'node:os';
+import path from 'node:path';
+import { pathToFileURL } from 'node:url';
+
+import { Command } from 'commander';
+import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest';
+
+// Hoisted mocks for shutdown tests
+const mockSetupEnv = vi.hoisted(() => vi.fn());
+const mockSetLogLevel = vi.hoisted(() => vi.fn());
+const mockTelemetryRecord = vi.hoisted(() => vi.fn());
+const mockTelemetryInitialize = vi.hoisted(() => vi.fn());
+const mockTelemetryShutdown = vi.hoisted(() => vi.fn().mockResolvedValue(undefined));
+const mockCloseLogger = vi.hoisted(() => vi.fn().mockResolvedValue(undefined));
+const mockCloseDbIfOpen = vi.hoisted(() => vi.fn());
+const mockDispatcherDestroy = vi.hoisted(() => vi.fn().mockResolvedValue(undefined));
+const mockGetGlobalDispatcher = vi.hoisted(() =>
+  vi.fn().mockReturnValue({ destroy: mockDispatcherDestroy }),
+);
+// Import actual undici to preserve other exports (Agent, ProxyAgent, setGlobalDispatcher, etc.)
+const actualUndici = vi.hoisted(() => vi.importActual<typeof import('undici')>('undici'));
+
+// Mock the dependencies
+vi.mock('../src/util', () => ({
+  setupEnv: mockSetupEnv,
+}));
+
+vi.mock('../src/logger', () => ({
+  __esModule: true,
+  default: { debug: vi.fn(), warn: vi.fn(), error: vi.fn(), info: vi.fn() },
+  setLogLevel: mockSetLogLevel,
+  closeLogger: mockCloseLogger,
+}));
+
+vi.mock('../src/telemetry', () => ({
+  __esModule: true,
+  default: {
+    initialize: mockTelemetryInitialize,
+    record: mockTelemetryRecord,
+    shutdown: mockTelemetryShutdown,
+  },
+}));
+
+vi.mock('../src/database/index', () => ({
+  closeDbIfOpen: mockCloseDbIfOpen,
+}));
+
+vi.mock('undici', async () => ({
+  ...(await actualUndici),
+  getGlobalDispatcher: mockGetGlobalDispatcher,
+}));
+
+// Mock code scan commands to avoid ESM import issues with execa
+vi.mock('../src/codeScan', () => ({
+  codeScansCommand: vi.fn(),
+}));
+
+let addCommonOptionsRecursively: typeof import('../src/mainUtils').addCommonOptionsRecursively;
+let isMainModule: typeof import('../src/mainUtils').isMainModule;
+let shouldSkipDefaultConfigLoading: typeof import('../src/mainUtils').shouldSkipDefaultConfigLoading;
+let setupEnvFilesFromArgv: typeof import('../src/mainUtils').setupEnvFilesFromArgv;
+let shutdownGracefully: typeof import('../src/mainUtils').shutdownGracefully;
+
+async function loadMainModule() {
+  vi.resetModules();
+  ({
+    addCommonOptionsRecursively,
+    isMainModule,
+    shouldSkipDefaultConfigLoading,
+    setupEnvFilesFromArgv,
+    shutdownGracefully,
+  } = await import('../src/mainUtils'));
+}
+
+describe('setupEnvFilesFromArgv', () => {
+  beforeEach(async () => {
+    await loadMainModule();
+    mockSetupEnv.mockReset();
+    mockTelemetryInitialize.mockReset();
+  });
+
+  it('should load env files before command actions run', () => {
+    setupEnvFilesFromArgv(['eval', '--env-file', '.env.local']);
+
+    expect(mockSetupEnv).toHaveBeenCalledWith('.env.local', { refreshConfigDirectory: true });
+    expect(mockTelemetryInitialize).toHaveBeenCalledOnce();
+  });
+
+  it('should support repeated and comma-separated env file args', () => {
+    setupEnvFilesFromArgv(['eval', '--env-file', '.env.one', '--env-path=.env.two,.env.three']);
+
+    expect(mockSetupEnv).toHaveBeenCalledWith(['.env.one', '.env.two', '.env.three'], {
+      refreshConfigDirectory: true,
+    });
+  });
+
+  it('should ignore flags after --', () => {
+    setupEnvFilesFromArgv(['eval', '--', '--env-file', '.env.local']);
+
+    expect(mockSetupEnv).not.toHaveBeenCalled();
+    expect(mockTelemetryInitialize).toHaveBeenCalledOnce();
+  });
+
+  it('should recognize the --env-path alias', () => {
+    setupEnvFilesFromArgv(['eval', '--env-path', '.env.staging']);
+
+    expect(mockSetupEnv).toHaveBeenCalledWith('.env.staging', { refreshConfigDirectory: true });
+  });
+
+  it('should be a no-op when no env flags are present', () => {
+    setupEnvFilesFromArgv(['eval', '--verbose', '--no-cache']);
+
+    expect(mockSetupEnv).not.toHaveBeenCalled();
+  });
+
+  it('should ignore --env-file= with empty value', () => {
+    setupEnvFilesFromArgv(['eval', '--env-file=']);
+
+    expect(mockSetupEnv).not.toHaveBeenCalled();
+  });
+
+  it('should skip --env-file when next arg starts with -', () => {
+    setupEnvFilesFromArgv(['eval', '--env-file', '--verbose']);
+
+    expect(mockSetupEnv).not.toHaveBeenCalled();
+  });
+});
+
+describe('shouldSkipDefaultConfigLoading', () => {
+  beforeEach(async () => {
+    await loadMainModule();
+  });
+
+  it('skips unrelated default config discovery for code-scans commands', () => {
+    expect(shouldSkipDefaultConfigLoading(['code-scans', 'run'])).toBe(true);
+    expect(shouldSkipDefaultConfigLoading(['--verbose', 'code-scans', 'run', '--help'])).toBe(true);
+    expect(
+      shouldSkipDefaultConfigLoading(['--env-file', '.env.local', 'code-scans', 'run', '--help']),
+    ).toBe(true);
+    expect(shouldSkipDefaultConfigLoading(['--env-path=.env.local', 'code-scans', 'run'])).toBe(
+      true,
+    );
+  });
+
+  it('keeps default config discovery for other commands and post-separator arguments', () => {
+    expect(shouldSkipDefaultConfigLoading(['eval', '--help'])).toBe(false);
+    expect(shouldSkipDefaultConfigLoading(['--', 'code-scans', 'run'])).toBe(false);
+  });
+});
+
+describe('addCommonOptionsRecursively', () => {
+  const originalExit = process.exit;
+  let program: Command;
+  let subCommand: Command;
+
+  beforeAll(() => {
+    process.exit = vi.fn() as any;
+  });
+
+  beforeEach(async () => {
+    await loadMainModule();
+    program = new Command('promptfoo');
+    program.action(() => {});
+    subCommand = program.command('subcommand');
+    subCommand.action(() => {});
+    mockTelemetryRecord.mockReset();
+    vi.clearAllMocks();
+  });
+
+  afterAll(() => {
+    process.exit = originalExit;
+  });
+
+  it('should add both verbose and env-file options to a command', () => {
+    addCommonOptionsRecursively(program);
+
+    const hasVerboseOption = program.options.some(
+      (option) => option.short === '-v' || option.long === '--verbose',
+    );
+    const hasEnvFileOption = program.options.some(
+      (option) => option.long === '--env-file' || option.long === '--env-path',
+    );
+
+    expect(hasVerboseOption).toBe(true);
+    expect(hasEnvFileOption).toBe(true);
+  });
+
+  it('should not add duplicate options if they already exist', () => {
+    program.option('--env-file, --env-path <path>', 'Path to .env file');
+    program.option('-v, --verbose', 'Show debug logs', false);
+
+    // Count options before
+    const envFileOptionsBefore = program.options.filter(
+      (option) => option.long === '--env-file' || option.long === '--env-path',
+    ).length;
+
+    const verboseOptionsBefore = program.options.filter(
+      (option) => option.short === '-v' || option.long === '--verbose',
+    ).length;
+
+    addCommonOptionsRecursively(program);
+
+    // Count options after
+    const envFileOptionsAfter = program.options.filter(
+      (option) => option.long === '--env-file' || option.long === '--env-path',
+    ).length;
+
+    const verboseOptionsAfter = program.options.filter(
+      (option) => option.short === '-v' || option.long === '--verbose',
+    ).length;
+
+    // Should still have the same number of options
+    expect(envFileOptionsAfter).toBe(envFileOptionsBefore);
+    expect(verboseOptionsAfter).toBe(verboseOptionsBefore);
+  });
+
+  it('should add options to subcommands', () => {
+    addCommonOptionsRecursively(program);
+
+    const hasSubcommandVerboseOption = subCommand.options.some(
+      (option) => option.short === '-v' || option.long === '--verbose',
+    );
+    const hasSubcommandEnvFileOption = subCommand.options.some(
+      (option) => option.long === '--env-file' || option.long === '--env-path',
+    );
+
+    expect(hasSubcommandVerboseOption).toBe(true);
+    expect(hasSubcommandEnvFileOption).toBe(true);
+  });
+
+  it('should add options to nested subcommands', () => {
+    // Create a deeper command structure
+    const subSubCommand = subCommand.command('subsubcommand');
+    subSubCommand.action(() => {});
+    const level3Command = subSubCommand.command('subsubsubcommand');
+    level3Command.action(() => {});
+
+    addCommonOptionsRecursively(program);
+
+    // Check all levels of commands have the options
+    const hasMainVerboseOption = program.options.some(
+      (option) => option.short === '-v' || option.long === '--verbose',
+    );
+    const hasMainEnvFileOption = program.options.some(
+      (option) => option.long === '--env-file' || option.long === '--env-path',
+    );
+
+    const hasSubCommandVerboseOption = subCommand.options.some(
+      (option) => option.short === '-v' || option.long === '--verbose',
+    );
+    const hasSubCommandEnvFileOption = subCommand.options.some(
+      (option) => option.long === '--env-file' || option.long === '--env-path',
+    );
+
+    const hasSubSubCommandVerboseOption = subSubCommand.options.some(
+      (option) => option.short === '-v' || option.long === '--verbose',
+    );
+    const hasSubSubCommandEnvFileOption = subSubCommand.options.some(
+      (option) => option.long === '--env-file' || option.long === '--env-path',
+    );
+
+    const hasSubSubSubCommandVerboseOption = level3Command.options.some(
+      (option) => option.short === '-v' || option.long === '--verbose',
+    );
+    const hasSubSubSubCommandEnvFileOption = level3Command.options.some(
+      (option) => option.long === '--env-file' || option.long === '--env-path',
+    );
+
+    expect(hasMainVerboseOption).toBe(true);
+    expect(hasMainEnvFileOption).toBe(true);
+    expect(hasSubCommandVerboseOption).toBe(true);
+    expect(hasSubCommandEnvFileOption).toBe(true);
+    expect(hasSubSubCommandVerboseOption).toBe(true);
+    expect(hasSubSubCommandEnvFileOption).toBe(true);
+    expect(hasSubSubSubCommandVerboseOption).toBe(true);
+    expect(hasSubSubSubCommandEnvFileOption).toBe(true);
+  });
+
+  it('should register a single hook that handles both options', () => {
+    // Create a fake action that manually mocks the Commander hook system
+    const mockHookRegister = vi.fn();
+    (program as any).hook = mockHookRegister;
+
+    // Apply common options
+    addCommonOptionsRecursively(program);
+
+    // Verify the hook was registered only once
+    expect(mockHookRegister).toHaveBeenCalledTimes(1);
+    expect(mockHookRegister).toHaveBeenCalledWith('preAction', expect.any(Function));
+
+    // Get the hook function
+    const preActionFn = mockHookRegister.mock.calls[0][1];
+
+    // Create mock command object with name() and parent for telemetry
+    const createMockCommand = (opts: Record<string, unknown>, commandName = 'test') => ({
+      opts: () => opts,
+      name: () => commandName,
+      parent: null,
+    });
+
+    // Test verbose option
+    preActionFn(createMockCommand({ verbose: true }));
+    expect(mockSetLogLevel).toHaveBeenCalledWith('debug');
+
+    // Test env-file option
+    preActionFn(createMockCommand({ envFile: '.env.test' }));
+    expect(mockSetupEnv).toHaveBeenCalledWith('.env.test');
+
+    // Test both options together
+    preActionFn(createMockCommand({ verbose: true, envFile: '.env.combined' }));
+    expect(mockSetLogLevel).toHaveBeenCalledWith('debug');
+    expect(mockSetupEnv).toHaveBeenCalledWith('.env.combined');
+  });
+
+  it('keeps structured code-scan output muted even when --verbose is present', () => {
+    const mockHookRegister = vi.fn();
+    (program as any).hook = mockHookRegister;
+
+    addCommonOptionsRecursively(program);
+
+    const preActionFn = mockHookRegister.mock.calls[0][1];
+    const originalArgv = process.argv;
+
+    process.argv = ['node', 'promptfoo', 'code-scans', 'run', '.', '--json', '--verbose'];
+
+    try {
+      preActionFn({
+        opts: () => ({ verbose: true }),
+        name: () => 'run',
+        parent: {
+          name: () => 'code-scans',
+          parent: null,
+        },
+      });
+    } finally {
+      process.argv = originalArgv;
+    }
+
+    expect(mockSetLogLevel).not.toHaveBeenCalledWith('debug');
+  });
+
+  it('should parse --env-file without consuming positional subcommand arguments', async () => {
+    const action = vi.fn();
+    const documentCommand = program.command('scan-model').argument('<model>').action(action);
+
+    addCommonOptionsRecursively(program);
+
+    await program.parseAsync(['scan-model', '--env-file', '.env.local', 'llama-3'], {
+      from: 'user',
+    });
+
+    expect(action).toHaveBeenCalledWith(
+      'llama-3',
+      expect.objectContaining({ verbose: false }),
+      documentCommand,
+    );
+    expect(mockSetupEnv).toHaveBeenCalledWith('.env.local');
+  });
+
+  it('should support repeated and comma-separated --env-file values', async () => {
+    addCommonOptionsRecursively(program);
+
+    await program.parseAsync(
+      ['subcommand', '--env-file', '.env.one', '--env-file', '.env.two,.env.three'],
+      { from: 'user' },
+    );
+
+    expect(mockSetupEnv).toHaveBeenCalledWith(['.env.one', '.env.two', '.env.three']);
+  });
+
+  it('should record telemetry once for nested subcommands', async () => {
+    const nestedParent = program.command('redteam');
+    const nestedAction = vi.fn();
+    nestedParent.command('run').action(nestedAction);
+
+    addCommonOptionsRecursively(program);
+
+    await program.parseAsync(['redteam', 'run'], { from: 'user' });
+
+    expect(nestedAction).toHaveBeenCalledTimes(1);
+    expect(mockTelemetryRecord).toHaveBeenCalledTimes(1);
+    expect(mockTelemetryRecord).toHaveBeenCalledWith('command_used', { name: 'redteam run' });
+  });
+});
+
+describe('isMainModule', () => {
+  let tempDir: string;
+  let realFilePath: string;
+  let symlinkPath: string;
+
+  beforeAll(async () => {
+    await loadMainModule();
+    // Create a temporary directory with a real file and a symlink
+    tempDir = fs.mkdtempSync(path.join(os.tmpdir(), 'main-test-'));
+    realFilePath = path.join(tempDir, 'real-file.js');
+    symlinkPath = path.join(tempDir, 'symlink-file.js');
+
+    // Create a real file
+    fs.writeFileSync(realFilePath, '// test file');
+
+    // Create a symlink pointing to the real file
+    fs.symlinkSync(realFilePath, symlinkPath);
+  });
+
+  afterAll(() => {
+    // Clean up temporary files
+    try {
+      fs.unlinkSync(symlinkPath);
+      fs.unlinkSync(realFilePath);
+      fs.rmdirSync(tempDir);
+    } catch {
+      // Ignore cleanup errors
+    }
+  });
+
+  it('should return false when processArgv1 is undefined', () => {
+    const result = isMainModule('file:///some/path/main.js', undefined);
+    expect(result).toBe(false);
+  });
+
+  it('should return false when processArgv1 is empty string', () => {
+    const result = isMainModule('file:///some/path/main.js', '');
+    expect(result).toBe(false);
+  });
+
+  it('should return true when paths match directly', () => {
+    const fileUrl = pathToFileURL(realFilePath).href;
+    const result = isMainModule(fileUrl, realFilePath);
+    expect(result).toBe(true);
+  });
+
+  it('should return true when processArgv1 is a symlink pointing to the module', () => {
+    // This is the key test case for npm global bin symlinks
+    const fileUrl = pathToFileURL(realFilePath).href;
+    const result = isMainModule(fileUrl, symlinkPath);
+    expect(result).toBe(true);
+  });
+
+  it('should return false when paths do not match', () => {
+    const fileUrl = pathToFileURL(realFilePath).href;
+    const otherPath = path.join(tempDir, 'other-file.js');
+    fs.writeFileSync(otherPath, '// other file');
+
+    try {
+      const result = isMainModule(fileUrl, otherPath);
+      expect(result).toBe(false);
+    } finally {
+      fs.unlinkSync(otherPath);
+    }
+  });
+
+  it('should return false when processArgv1 points to non-existent file', () => {
+    const fileUrl = pathToFileURL(realFilePath).href;
+    const nonExistentPath = path.join(tempDir, 'non-existent.js');
+    const result = isMainModule(fileUrl, nonExistentPath);
+    expect(result).toBe(false);
+  });
+
+  it('should handle relative paths correctly', () => {
+    // Save current directory
+    const originalCwd = process.cwd();
+
+    try {
+      // Change to temp directory
+      process.chdir(tempDir);
+
+      const fileUrl = pathToFileURL(realFilePath).href;
+      const result = isMainModule(fileUrl, './real-file.js');
+      expect(result).toBe(true);
+    } finally {
+      // Restore original directory
+      process.chdir(originalCwd);
+    }
+  });
+
+  it('should handle symlink to symlink chains', () => {
+    const secondSymlinkPath = path.join(tempDir, 'second-symlink.js');
+
+    try {
+      // Create a symlink pointing to the first symlink
+      fs.symlinkSync(symlinkPath, secondSymlinkPath);
+
+      const fileUrl = pathToFileURL(realFilePath).href;
+      const result = isMainModule(fileUrl, secondSymlinkPath);
+      expect(result).toBe(true);
+    } finally {
+      try {
+        fs.unlinkSync(secondSymlinkPath);
+      } catch {
+        // Ignore cleanup errors
+      }
+    }
+  });
+});
+
+describe('shutdownGracefully', () => {
+  const originalExit = process.exit;
+
+  beforeEach(async () => {
+    await loadMainModule();
+    vi.useFakeTimers();
+    process.exit = vi.fn() as never;
+    // Reset all mocks to default behavior
+    mockSetupEnv.mockReset();
+    mockSetLogLevel.mockReset();
+    mockTelemetryShutdown.mockReset().mockResolvedValue(undefined);
+    mockCloseLogger.mockReset().mockResolvedValue(undefined);
+    mockCloseDbIfOpen.mockReset();
+    mockDispatcherDestroy.mockReset().mockResolvedValue(undefined);
+    mockGetGlobalDispatcher.mockReset().mockReturnValue({ destroy: mockDispatcherDestroy });
+  });
+
+  afterEach(() => {
+    vi.useRealTimers();
+    vi.resetAllMocks();
+    process.exit = originalExit;
+  });
+
+  it('should complete gracefully when all operations succeed quickly', async () => {
+    const shutdownPromise = shutdownGracefully();
+
+    // Advance past any internal timeouts
+    await vi.runAllTimersAsync();
+
+    await shutdownPromise;
+
+    expect(mockTelemetryShutdown).toHaveBeenCalled();
+    expect(mockCloseLogger).toHaveBeenCalled();
+    expect(mockCloseDbIfOpen).toHaveBeenCalled();
+    expect(mockDispatcherDestroy).toHaveBeenCalled();
+  });
+
+  it('should handle telemetry shutdown timeout', async () => {
+    // Make telemetry.shutdown() hang forever
+    mockTelemetryShutdown.mockImplementation(() => new Promise(() => {}));
+
+    const shutdownPromise = shutdownGracefully();
+
+    // Advance time to trigger the individual operation timeout (1s)
+    await vi.advanceTimersByTimeAsync(1100);
+
+    // Advance remaining timers
+    await vi.runAllTimersAsync();
+
+    await shutdownPromise;
+
+    // Should still have called other cleanup operations
+    expect(mockCloseLogger).toHaveBeenCalled();
+    expect(mockCloseDbIfOpen).toHaveBeenCalled();
+  });
+
+  it('should handle closeLogger timeout', async () => {
+    // Make closeLogger() hang forever
+    mockCloseLogger.mockImplementation(() => new Promise(() => {}));
+
+    const shutdownPromise = shutdownGracefully();
+
+    // Advance time to trigger all timeouts
+    await vi.runAllTimersAsync();
+
+    await shutdownPromise;
+
+    // Other operations should still complete
+    expect(mockTelemetryShutdown).toHaveBeenCalled();
+    expect(mockCloseDbIfOpen).toHaveBeenCalled();
+    expect(mockDispatcherDestroy).toHaveBeenCalled();
+  });
+
+  it('should handle dispatcher.destroy() timeout', async () => {
+    // Make dispatcher.destroy() hang forever
+    mockDispatcherDestroy.mockImplementation(() => new Promise(() => {}));
+
+    const shutdownPromise = shutdownGracefully();
+
+    // Advance time to trigger all timeouts
+    await vi.runAllTimersAsync();
+
+    await shutdownPromise;
+
+    // Other operations should still complete
+    expect(mockTelemetryShutdown).toHaveBeenCalled();
+    expect(mockCloseLogger).toHaveBeenCalled();
+    expect(mockCloseDbIfOpen).toHaveBeenCalled();
+  });
+
+  it('should handle telemetry shutdown throwing error', async () => {
+    mockTelemetryShutdown.mockRejectedValue(new Error('Telemetry failed'));
+
+    const shutdownPromise = shutdownGracefully();
+
+    await vi.runAllTimersAsync();
+
+    // Should not throw - should handle error gracefully
+    await expect(shutdownPromise).resolves.toBeUndefined();
+
+    // Other operations should still complete
+    expect(mockCloseLogger).toHaveBeenCalled();
+    expect(mockCloseDbIfOpen).toHaveBeenCalled();
+  });
+
+  it('should handle dispatcher.destroy() throwing error', async () => {
+    mockDispatcherDestroy.mockRejectedValue(new Error('Dispatcher failed'));
+
+    const shutdownPromise = shutdownGracefully();
+
+    await vi.runAllTimersAsync();
+
+    // Should not throw
+    await expect(shutdownPromise).resolves.toBeUndefined();
+  });
+
+  it('should force exit when all cleanup operations hang', async () => {
+    // Make all operations hang
+    mockTelemetryShutdown.mockImplementation(() => new Promise(() => {}));
+    mockCloseLogger.mockImplementation(() => new Promise(() => {}));
+    mockDispatcherDestroy.mockImplementation(() => new Promise(() => {}));
+
+    // Start shutdown but don't await yet
+    void shutdownGracefully();
+
+    // The force exit timeout is 3000ms
+    await vi.advanceTimersByTimeAsync(3000);
+
+    expect(process.exit).toHaveBeenCalledWith(0);
+  });
+
+  it('should clear force exit timeout when cleanup completes normally', async () => {
+    const setTimeoutSpy = vi.spyOn(global, 'setTimeout');
+    const clearTimeoutSpy = vi.spyOn(global, 'clearTimeout');
+    const shutdownPromise = shutdownGracefully();
+
+    await vi.runAllTimersAsync();
+
+    await shutdownPromise;
+
+    const forceExitTimeout = setTimeoutSpy.mock.results[0]?.value;
+    expect(clearTimeoutSpy).toHaveBeenCalledWith(forceExitTimeout);
+    expect(process.exit).toHaveBeenCalled();
+  });
+
+  it('should handle getGlobalDispatcher throwing', async () => {
+    mockGetGlobalDispatcher.mockImplementation(() => {
+      throw new Error('No dispatcher');
+    });
+
+    const shutdownPromise = shutdownGracefully();
+
+    await vi.runAllTimersAsync();
+
+    // Should complete without throwing
+    await expect(shutdownPromise).resolves.toBeUndefined();
+  });
+});