repomix/lib/shared/unifiedWorker.js

A tool to pack repository contents to single file for AI consumption

import { workerData } from 'node:worker_threads';
const handlerCache = new Map();
const loadWorkerHandler = async (workerType) => {
    const cached = handlerCache.get(workerType);
    if (cached) {
        return cached;
    }
    let result;
    switch (workerType) {
        case 'fileProcess': {
            const module = await import('../core/file/workers/fileProcessWorker.js');
            result = { handler: module.default, cleanup: module.onWorkerTermination };
            break;
        }
        case 'securityCheck': {
            const module = await import('../core/security/workers/securityCheckWorker.js');
            result = { handler: module.default, cleanup: module.onWorkerTermination };
            break;
        }
        case 'calculateMetrics': {
            const module = await import('../core/metrics/workers/calculateMetricsWorker.js');
            result = { handler: module.default, cleanup: module.onWorkerTermination };
            break;
        }
        default:
            throw new Error(`Unknown worker type: ${workerType}`);
    }
    handlerCache.set(workerType, result);
    return result;
};
const inferWorkerTypeFromTask = (task) => {
    if (!task || typeof task !== 'object') {
        return null;
    }
    const taskObj = task;
    if ('rawFile' in taskObj && 'config' in taskObj) {
        return 'fileProcess';
    }
    if ('encoding' in taskObj && ('content' in taskObj || 'items' in taskObj)) {
        return 'calculateMetrics';
    }
    if ('items' in taskObj && !('encoding' in taskObj)) {
        return 'securityCheck';
    }
    return null;
};
const getWorkerTypeFromWorkerData = () => {
    if (!workerData) {
        return undefined;
    }
    if (Array.isArray(workerData)) {
        for (const item of workerData) {
            if (item && typeof item === 'object' && 'workerType' in item) {
                return item.workerType;
            }
        }
        return undefined;
    }
    if (typeof workerData === 'object' && 'workerType' in workerData) {
        return workerData.workerType;
    }
    return undefined;
};
export default async (task) => {
    let workerType = getWorkerTypeFromWorkerData() ?? process.env.REPOMIX_WORKER_TYPE;
    const inferredType = inferWorkerTypeFromTask(task);
    if (inferredType) {
        workerType = inferredType;
    }
    if (!workerType) {
        throw new Error('Cannot determine worker type from workerData, env, or task structure');
    }
    const { handler } = await loadWorkerHandler(workerType);
    return handler(task);
};
export const onWorkerTermination = async () => {
    for (const { cleanup } of handlerCache.values()) {
        if (cleanup) {
            await cleanup();
        }
    }
    handlerCache.clear();
};