feat(core): Add runtime selection support for worker pools

Add WorkerRuntime type and configurable runtime parameter to createWorkerPool and initTaskRunner functions. This allows choosing between 'worker_threads' and 'child_process' runtimes based on performance requirements.

- Add WorkerRuntime type definition for type safety
- Add optional runtime parameter to createWorkerPool with child_process default
- Add optional runtime parameter to initTaskRunner with child_process default
- Configure fileCollectWorker to use worker_threads for better performance
- Update all test files to use WorkerRuntime type
- Add comprehensive tests for runtime parameter functionality
- Maintain backward compatibility with existing code

The fileCollectWorker now benefits from worker_threads faster startup and shared memory, while other workers continue using child_process for stability.

benchmarks/memory/src/memory-test.ts

@@ -26,7 +26,7 @@ const flags = {
 
 // Extract numeric arguments
 const numericArgs = args.filter((arg) => !arg.startsWith('-') && !Number.isNaN(Number(arg)));
-const iterations = Number(numericArgs[0]) || (flags.full ? 200 : 50);
+const iterations = Number(numericArgs[0]) || (flags.full ? 200 : 100);
 const delay = Number(numericArgs[1]) || (flags.full ? 100 : 50);
 
 // Configuration

src/core/file/fileCollect.ts

@@ -26,6 +26,7 @@ export const collectFiles = async (
   const taskRunner = deps.initTaskRunner<FileCollectTask, FileCollectResult>(
     filePaths.length,
     new URL('./workers/fileCollectWorker.js', import.meta.url).href,
+    'worker_threads',
   );
   const tasks = filePaths.map(
     (filePath) =>

src/shared/processConcurrency.ts

@@ -1,7 +1,9 @@
 import os from 'node:os';
-import { Tinypool } from 'tinypool';
+import { Options, Tinypool } from 'tinypool';
 import { logger } from './logger.js';
 
+export type WorkerRuntime = NonNullable<Options['runtime']>;
+
 // Worker initialization is expensive, so we prefer fewer threads unless there are many files
 const TASKS_PER_THREAD = 100;
 
@@ -23,19 +25,22 @@ export const getWorkerThreadCount = (numOfTasks: number): { minThreads: number;
   };
 };
 
-export const createWorkerPool = (numOfTasks: number, workerPath: string): Tinypool => {
+export const createWorkerPool = (
+  numOfTasks: number,
+  workerPath: string,
+  runtime: WorkerRuntime = 'child_process',
+): Tinypool => {
   const { minThreads, maxThreads } = getWorkerThreadCount(numOfTasks);
 
   logger.trace(
-    `Initializing worker pool with min=${minThreads}, max=${maxThreads} threads. Worker path: ${workerPath}`,
+    `Initializing worker pool with min=${minThreads}, max=${maxThreads} threads, runtime=${runtime}. Worker path: ${workerPath}`,
   );
 
   const startTime = process.hrtime.bigint();
 
   const pool = new Tinypool({
     filename: workerPath,
-    // Use child_process for better memory management
-    runtime: 'child_process',
+    runtime,
     minThreads,
     maxThreads,
     idleTimeout: 5000,
@@ -78,8 +83,12 @@ export interface TaskRunner<T, R> {
   cleanup: () => Promise<void>;
 }
 
-export const initTaskRunner = <T, R>(numOfTasks: number, workerPath: string): TaskRunner<T, R> => {
-  const pool = createWorkerPool(numOfTasks, workerPath);
+export const initTaskRunner = <T, R>(
+  numOfTasks: number,
+  workerPath: string,
+  runtime: WorkerRuntime = 'child_process',
+): TaskRunner<T, R> => {
+  const pool = createWorkerPool(numOfTasks, workerPath, runtime);
   return {
     run: (task: T) => pool.run(task),
     cleanup: () => cleanupWorkerPool(pool),

tests/core/file/fileCollect.test.ts

@@ -9,6 +9,7 @@ import { collectFiles } from '../../../src/core/file/fileCollect.js';
 import type { FileCollectTask } from '../../../src/core/file/workers/fileCollectWorker.js';
 import fileCollectWorker from '../../../src/core/file/workers/fileCollectWorker.js';
 import { logger } from '../../../src/shared/logger.js';
+import type { WorkerRuntime } from '../../../src/shared/processConcurrency.js';
 import { createMockConfig } from '../../testing/testUtils.js';
 
 // Define the max file size constant for tests
@@ -20,7 +21,21 @@ vi.mock('jschardet');
 vi.mock('iconv-lite');
 vi.mock('../../../src/shared/logger');
 
-const mockInitTaskRunner = <T, R>(_numOfTasks: number, _workerPath: string) => {
+interface MockInitTaskRunner {
+  <T, R>(
+    numOfTasks: number,
+    workerPath: string,
+    runtime?: WorkerRuntime,
+  ): {
+    run: (task: T) => Promise<R>;
+    cleanup: () => Promise<void>;
+  };
+  lastRuntime?: WorkerRuntime;
+}
+
+const mockInitTaskRunner = <T, R>(_numOfTasks: number, _workerPath: string, runtime?: WorkerRuntime) => {
+  // Store runtime for verification in tests
+  (mockInitTaskRunner as MockInitTaskRunner).lastRuntime = runtime;
   return {
     run: async (task: T) => {
       return (await fileCollectWorker(task as FileCollectTask)) as R;
@@ -195,4 +210,21 @@ describe('fileCollect', () => {
       expect.any(Error),
     );
   });
+
+  it('should use worker_threads runtime when calling initTaskRunner', async () => {
+    const mockFilePaths = ['test.txt'];
+    const mockRootDir = '/root';
+    const mockConfig = createMockConfig();
+
+    vi.mocked(isBinary).mockReturnValue(false);
+    vi.mocked(fs.readFile).mockResolvedValue(Buffer.from('file content'));
+    vi.mocked(jschardet.detect).mockReturnValue({ encoding: 'utf-8', confidence: 0.99 });
+    vi.mocked(iconv.decode).mockReturnValue('decoded content');
+
+    await collectFiles(mockFilePaths, mockRootDir, mockConfig, () => {}, {
+      initTaskRunner: mockInitTaskRunner,
+    });
+
+    expect((mockInitTaskRunner as MockInitTaskRunner).lastRuntime).toBe('worker_threads');
+  });
 });

tests/core/file/fileProcess.test.ts

@@ -5,6 +5,7 @@ import { processContent } from '../../../src/core/file/fileProcessContent.js';
 import type { RawFile } from '../../../src/core/file/fileTypes.js';
 import type { FileProcessTask } from '../../../src/core/file/workers/fileProcessWorker.js';
 import fileProcessWorker from '../../../src/core/file/workers/fileProcessWorker.js';
+import type { WorkerRuntime } from '../../../src/shared/processConcurrency.js';
 import { createMockConfig } from '../../testing/testUtils.js';
 
 const createMockFileManipulator = (): FileManipulator => ({
@@ -19,7 +20,7 @@ const mockGetFileManipulator = (filePath: string): FileManipulator | null => {
   return null;
 };
 
-const mockInitTaskRunner = <T, R>(_numOfTasks: number, _workerPath: string) => {
+const mockInitTaskRunner = <T, R>(_numOfTasks: number, _workerPath: string, _runtime?: WorkerRuntime) => {
   return {
     run: async (task: T) => {
       return (await fileProcessWorker(task as FileProcessTask)) as R;

tests/core/metrics/calculateOutputMetrics.test.ts

@@ -3,10 +3,11 @@ import { calculateOutputMetrics } from '../../../src/core/metrics/calculateOutpu
 import type { OutputMetricsTask } from '../../../src/core/metrics/workers/outputMetricsWorker.js';
 import outputMetricsWorker from '../../../src/core/metrics/workers/outputMetricsWorker.js';
 import { logger } from '../../../src/shared/logger.js';
+import type { WorkerRuntime } from '../../../src/shared/processConcurrency.js';
 
 vi.mock('../../../src/shared/logger');
 
-const mockInitTaskRunner = <T, R>(_numOfTasks: number, _workerPath: string) => {
+const mockInitTaskRunner = <T, R>(_numOfTasks: number, _workerPath: string, _runtime?: WorkerRuntime) => {
   return {
     run: async (task: T) => {
       return (await outputMetricsWorker(task as OutputMetricsTask)) as R;
@@ -46,7 +47,7 @@ describe('calculateOutputMetrics', () => {
     const encoding = 'o200k_base';
     const mockError = new Error('Worker error');
 
-    const mockErrorTaskRunner = <T, _R>(_numOfTasks: number, _workerPath: string) => {
+    const mockErrorTaskRunner = <T, _R>(_numOfTasks: number, _workerPath: string, _runtime?: WorkerRuntime) => {
       return {
         run: async (_task: T) => {
           throw mockError;
@@ -96,7 +97,7 @@ describe('calculateOutputMetrics', () => {
     const path = 'large-file.txt';
 
     let chunksProcessed = 0;
-    const mockParallelTaskRunner = <T, R>(_numOfTasks: number, _workerPath: string) => {
+    const mockParallelTaskRunner = <T, R>(_numOfTasks: number, _workerPath: string, _runtime?: WorkerRuntime) => {
       return {
         run: async (_task: T) => {
           chunksProcessed++;
@@ -122,7 +123,7 @@ describe('calculateOutputMetrics', () => {
     const encoding = 'o200k_base';
     const mockError = new Error('Parallel processing error');
 
-    const mockErrorTaskRunner = <T, _R>(_numOfTasks: number, _workerPath: string) => {
+    const mockErrorTaskRunner = <T, _R>(_numOfTasks: number, _workerPath: string, _runtime?: WorkerRuntime) => {
       return {
         run: async (_task: T) => {
           throw mockError;
@@ -147,7 +148,7 @@ describe('calculateOutputMetrics', () => {
     const encoding = 'o200k_base';
     const processedChunks: string[] = [];
 
-    const mockChunkTrackingTaskRunner = <T, R>(_numOfTasks: number, _workerPath: string) => {
+    const mockChunkTrackingTaskRunner = <T, R>(_numOfTasks: number, _workerPath: string, _runtime?: WorkerRuntime) => {
       return {
         run: async (task: T) => {
           const outputTask = task as OutputMetricsTask;

tests/core/metrics/calculateSelectiveFileMetrics.test.ts

@@ -3,13 +3,14 @@ import type { ProcessedFile } from '../../../src/core/file/fileTypes.js';
 import { calculateSelectiveFileMetrics } from '../../../src/core/metrics/calculateSelectiveFileMetrics.js';
 import type { FileMetricsTask } from '../../../src/core/metrics/workers/fileMetricsWorker.js';
 import fileMetricsWorker from '../../../src/core/metrics/workers/fileMetricsWorker.js';
+import type { WorkerRuntime } from '../../../src/shared/processConcurrency.js';
 import type { RepomixProgressCallback } from '../../../src/shared/types.js';
 
 vi.mock('../../shared/processConcurrency', () => ({
   getProcessConcurrency: () => 1,
 }));
 
-const mockInitTaskRunner = <T, R>(_numOfTasks: number, _workerPath: string) => {
+const mockInitTaskRunner = <T, R>(_numOfTasks: number, _workerPath: string, _runtime?: WorkerRuntime) => {
   return {
     run: async (task: T) => {
       return (await fileMetricsWorker(task as FileMetricsTask)) as R;

tests/core/security/securityCheck.test.ts

@@ -9,6 +9,7 @@ import type { SecurityCheckTask } from '../../../src/core/security/workers/secur
 import securityCheckWorker from '../../../src/core/security/workers/securityCheckWorker.js';
 import { logger } from '../../../src/shared/logger.js';
 import { repomixLogLevels } from '../../../src/shared/logger.js';
+import type { WorkerRuntime } from '../../../src/shared/processConcurrency.js';
 
 vi.mock('../../../src/shared/logger');
 vi.mock('../../../src/shared/processConcurrency', () => ({
@@ -39,7 +40,7 @@ const mockFiles: RawFile[] = [
   },
 ];
 
-const mockInitTaskRunner = <T, R>(_numOfTasks: number, _workerPath: string) => {
+const mockInitTaskRunner = <T, R>(_numOfTasks: number, _workerPath: string, _runtime?: WorkerRuntime) => {
   return {
     run: async (task: T) => {
       return (await securityCheckWorker(task as SecurityCheckTask)) as R;
@@ -78,7 +79,7 @@ describe('runSecurityCheck', () => {
 
   it('should handle worker errors gracefully', async () => {
     const mockError = new Error('Worker error');
-    const mockErrorTaskRunner = () => {
+    const mockErrorTaskRunner = (_numOfTasks?: number, _workerPath?: string, _runtime?: WorkerRuntime) => {
       return {
         run: async () => {
           throw mockError;

tests/integration-tests/packager.test.ts

@@ -24,13 +24,14 @@ import { copyToClipboardIfEnabled } from '../../src/core/packager/copyToClipboar
 import { writeOutputToDisk } from '../../src/core/packager/writeOutputToDisk.js';
 import { filterOutUntrustedFiles } from '../../src/core/security/filterOutUntrustedFiles.js';
 import { validateFileSafety } from '../../src/core/security/validateFileSafety.js';
+import type { WorkerRuntime } from '../../src/shared/processConcurrency.js';
 import { isWindows } from '../testing/testUtils.js';
 
 const fixturesDir = path.join(__dirname, 'fixtures', 'packager');
 const inputsDir = path.join(fixturesDir, 'inputs');
 const outputsDir = path.join(fixturesDir, 'outputs');
 
-const mockCollectFileInitTaskRunner = <T, R>(_numOfTasks: number, _workerPath: string) => {
+const mockCollectFileInitTaskRunner = <T, R>(_numOfTasks: number, _workerPath: string, _runtime?: WorkerRuntime) => {
   return {
     run: async (task: T) => {
       return (await fileCollectWorker(task as FileCollectTask)) as R;

tests/shared/processConcurrency.test.ts

@@ -87,6 +87,23 @@ describe('processConcurrency', () => {
       });
       expect(tinypool).toBeDefined();
     });
+
+    it('should initialize Tinypool with worker_threads runtime when specified', () => {
+      const workerPath = '/path/to/worker.js';
+      const tinypool = createWorkerPool(500, workerPath, 'worker_threads');
+
+      expect(Tinypool).toHaveBeenCalledWith({
+        filename: workerPath,
+        runtime: 'worker_threads',
+        minThreads: 1,
+        maxThreads: 4, // Math.min(4, 500/100) = 4
+        idleTimeout: 5000,
+        workerData: {
+          logLevel: 2,
+        },
+      });
+      expect(tinypool).toBeDefined();
+    });
   });
 
   describe('initTaskRunner', () => {
@@ -110,5 +127,18 @@ describe('processConcurrency', () => {
       expect(typeof taskRunner.run).toBe('function');
       expect(typeof taskRunner.cleanup).toBe('function');
     });
+
+    it('should pass runtime parameter to createWorkerPool', () => {
+      const workerPath = '/path/to/worker.js';
+      const taskRunner = initTaskRunner(100, workerPath, 'worker_threads');
+
+      expect(Tinypool).toHaveBeenCalledWith(
+        expect.objectContaining({
+          runtime: 'worker_threads',
+        }),
+      );
+      expect(taskRunner).toHaveProperty('run');
+      expect(taskRunner).toHaveProperty('cleanup');
+    });
   });
 });