code-graph-context 2.0.0 → 2.2.0

package/dist/mcp/workers/chunk-worker-pool.js

@@ -0,0 +1,196 @@
+/**
+ * Chunk Worker Pool
+ * Manages a pool of chunk workers for parallel parsing.
+ * Uses message passing (pull model): workers signal ready, coordinator sends chunks.
+ * Streams results as they complete for pipelined importing.
+ */
+import { cpus } from 'os';
+import { dirname, join } from 'path';
+import { fileURLToPath } from 'url';
+import { Worker } from 'worker_threads';
+import { debugLog } from '../utils.js';
+const __filename = fileURLToPath(import.meta.url);
+const __dirname = dirname(__filename);
+export class ChunkWorkerPool {
+    config;
+    workers = [];
+    chunkQueue = [];
+    totalChunks = 0;
+    completedChunks = 0;
+    totalNodes = 0;
+    totalEdges = 0;
+    totalFiles = 0;
+    startTime = 0;
+    resolve = null;
+    reject = null;
+    onChunkComplete = null;
+    pendingCallbacks = [];
+    isShuttingDown = false;
+    constructor(config) {
+        this.config = config;
+        process.on('exit', () => {
+            this.forceTerminateAll();
+        });
+    }
+    /**
+     * Process chunks in parallel using worker pool.
+     * Calls onChunkComplete for EACH result as it arrives (for pipelined importing).
+     * Returns final stats when all chunks AND all callbacks are complete.
+     */
+    async processChunks(chunks, onChunkComplete) {
+        this.startTime = Date.now();
+        this.totalChunks = chunks.length;
+        this.completedChunks = 0;
+        this.totalNodes = 0;
+        this.totalEdges = 0;
+        this.totalFiles = 0;
+        this.onChunkComplete = onChunkComplete;
+        this.pendingCallbacks = [];
+        this.chunkQueue = chunks.map((files, index) => ({
+            type: 'chunk',
+            chunkIndex: index,
+            totalChunks: chunks.length,
+            files,
+        }));
+        const numWorkers = this.config.numWorkers ?? Math.floor(cpus().length * 0.75);
+        const actualWorkers = Math.min(numWorkers, chunks.length);
+        debugLog(`Spawning ${actualWorkers} chunk workers for ${chunks.length} chunks`);
+        return new Promise((resolve, reject) => {
+            this.resolve = resolve;
+            this.reject = reject;
+            for (let i = 0; i < actualWorkers; i++) {
+                this.spawnWorker();
+            }
+        });
+    }
+    spawnWorker() {
+        const workerPath = join(__dirname, 'chunk.worker.js');
+        const workerConfig = {
+            projectPath: this.config.projectPath,
+            tsconfigPath: this.config.tsconfigPath,
+            projectId: this.config.projectId,
+            projectType: this.config.projectType,
+        };
+        const worker = new Worker(workerPath, {
+            workerData: workerConfig,
+            resourceLimits: {
+                maxOldGenerationSizeMb: 2048,
+                maxYoungGenerationSizeMb: 512,
+            },
+        });
+        const state = { worker, busy: false };
+        this.workers.push(state);
+        worker.on('message', (msg) => {
+            this.handleWorkerMessage(state, msg);
+        });
+        worker.on('error', (error) => {
+            debugLog('Worker error', { error: error.message });
+            this.reject?.(error);
+            void this.shutdown();
+        });
+        worker.on('exit', (code) => {
+            if (code !== 0 && this.completedChunks < this.totalChunks) {
+                this.reject?.(new Error(`Worker exited with code ${code}`));
+                void this.shutdown();
+            }
+        });
+    }
+    handleWorkerMessage(state, msg) {
+        switch (msg.type) {
+            case 'ready':
+                state.busy = false;
+                this.dispatchNextChunk(state);
+                break;
+            case 'result':
+                this.handleResult(msg);
+                break;
+            case 'error':
+                debugLog(`Chunk ${msg.chunkIndex} failed`, { error: msg.error });
+                this.reject?.(new Error(`Chunk ${msg.chunkIndex} failed: ${msg.error}`));
+                void this.shutdown();
+                break;
+        }
+    }
+    handleResult(msg) {
+        this.completedChunks++;
+        this.totalNodes += msg.nodes.length;
+        this.totalEdges += msg.edges.length;
+        this.totalFiles += msg.filesProcessed;
+        const result = {
+            chunkIndex: msg.chunkIndex,
+            nodes: msg.nodes,
+            edges: msg.edges,
+            filesProcessed: msg.filesProcessed,
+            sharedContext: msg.sharedContext,
+            deferredEdges: msg.deferredEdges,
+        };
+        const stats = this.getStats();
+        // Fire callback immediately - enables pipelined importing
+        if (this.onChunkComplete) {
+            const callbackPromise = this.onChunkComplete(result, stats).catch((err) => {
+                debugLog(`Import callback failed for chunk ${msg.chunkIndex}`, {
+                    error: err instanceof Error ? err.message : String(err),
+                });
+                throw err;
+            });
+            this.pendingCallbacks.push(callbackPromise);
+        }
+        // Check if all parsing is done
+        if (this.completedChunks === this.totalChunks) {
+            this.completeWhenCallbacksDone();
+        }
+    }
+    async completeWhenCallbacksDone() {
+        try {
+            await Promise.all(this.pendingCallbacks);
+        }
+        catch (error) {
+            this.reject?.(error instanceof Error ? error : new Error(String(error)));
+            await this.shutdown();
+            return;
+        }
+        await this.shutdown();
+        this.resolve?.(this.getStats());
+    }
+    getStats() {
+        return {
+            totalNodes: this.totalNodes,
+            totalEdges: this.totalEdges,
+            totalFiles: this.totalFiles,
+            chunksCompleted: this.completedChunks,
+            totalChunks: this.totalChunks,
+            elapsedMs: Date.now() - this.startTime,
+        };
+    }
+    dispatchNextChunk(state) {
+        if (this.chunkQueue.length === 0) {
+            return;
+        }
+        const chunk = this.chunkQueue.shift();
+        state.busy = true;
+        state.worker.postMessage(chunk);
+    }
+    /**
+     * Graceful shutdown - lets workers finish cleanup
+     * Call this on normal completion
+     */
+    async shutdown() {
+        if (this.isShuttingDown)
+            return;
+        this.isShuttingDown = true;
+        const exitPromises = this.workers.map(({ worker }) => {
+            return new Promise((resolve) => {
+                worker.on('exit', () => resolve());
+                worker.postMessage({ type: 'terminate' });
+            });
+        });
+        await Promise.race([Promise.all(exitPromises), new Promise((resolve) => setTimeout(resolve, 15000))]);
+        this.forceTerminateAll();
+    }
+    forceTerminateAll() {
+        for (const { worker } of this.workers) {
+            worker.terminate();
+        }
+        this.workers = [];
+    }
+}

package/dist/mcp/workers/chunk-worker.types.js

@@ -0,0 +1,4 @@
+/**
+ * Types for chunk worker communication
+ */
+export {};

package/dist/mcp/workers/chunk.worker.js

@@ -0,0 +1,89 @@
+/**
+ * Chunk Worker
+ * Receives file chunks from coordinator, parses them, returns nodes/edges.
+ * Each worker creates its own parser with lazyLoad=true for memory efficiency.
+ */
+import { dirname, join } from 'path';
+import { fileURLToPath } from 'url';
+import { parentPort, workerData } from 'worker_threads';
+// Load environment variables in worker thread
+import dotenv from 'dotenv';
+const __filename = fileURLToPath(import.meta.url);
+const __dirname = dirname(__filename);
+dotenv.config({ path: join(__dirname, '..', '..', '..', '.env') });
+import { ParserFactory } from '../../core/parsers/parser-factory.js';
+const config = workerData;
+let parser = null;
+const sendReady = () => {
+    const msg = { type: 'ready' };
+    parentPort?.postMessage(msg);
+};
+const sendResult = (result) => {
+    const msg = { type: 'result', ...result };
+    parentPort?.postMessage(msg);
+};
+const sendError = (chunkIndex, error) => {
+    const msg = {
+        type: 'error',
+        chunkIndex,
+        error: error.message,
+        stack: error.stack,
+    };
+    parentPort?.postMessage(msg);
+};
+/**
+ * Initialize parser lazily on first chunk.
+ * Uses lazyLoad=true so parser only loads files we give it.
+ * projectType is already resolved by coordinator (no auto-detection here).
+ */
+const initParser = () => {
+    if (parser)
+        return parser;
+    parser = ParserFactory.createParser({
+        workspacePath: config.projectPath,
+        tsConfigPath: config.tsconfigPath,
+        projectType: config.projectType,
+        projectId: config.projectId,
+        lazyLoad: true, // Critical: only load files we're given
+    });
+    // Defer edge enhancements - coordinator will handle after all chunks complete
+    parser.setDeferEdgeEnhancements(true);
+    return parser;
+};
+const processChunk = async (files, chunkIndex) => {
+    try {
+        const p = initParser();
+        // Clear any accumulated data from previous chunks
+        p.clearParsedData();
+        // Parse chunk - skip deferred edge resolution (coordinator handles that)
+        const { nodes, edges } = await p.parseChunk(files, true);
+        // Get serialized shared context for merging in coordinator
+        const sharedContext = p.getSerializedSharedContext();
+        // Get deferred edges for cross-chunk resolution
+        const deferredEdges = p.getDeferredEdges();
+        sendResult({
+            chunkIndex,
+            nodes,
+            edges,
+            filesProcessed: files.length,
+            sharedContext,
+            deferredEdges,
+        });
+    }
+    catch (error) {
+        sendError(chunkIndex, error instanceof Error ? error : new Error(String(error)));
+    }
+};
+parentPort?.on('message', async (msg) => {
+    switch (msg.type) {
+        case 'chunk':
+            await processChunk(msg.files, msg.chunkIndex);
+            sendReady();
+            break;
+        case 'terminate':
+            parser?.clearParsedData();
+            process.exit(0);
+            break;
+    }
+});
+sendReady();

package/dist/mcp/workers/parse-coordinator.js

@@ -0,0 +1,183 @@
+/**
+ * Parse Coordinator
+ * Runs TypeScript parsing in a separate thread to avoid blocking the MCP server.
+ * For large projects, spawns a worker pool for parallel chunk parsing.
+ */
+import { dirname, join } from 'path';
+import { fileURLToPath } from 'url';
+import { parentPort, workerData } from 'worker_threads';
+// Load environment variables in worker thread
+import dotenv from 'dotenv';
+const __filename = fileURLToPath(import.meta.url);
+const __dirname = dirname(__filename);
+dotenv.config({ path: join(__dirname, '..', '..', '..', '.env') });
+import { EmbeddingsService } from '../../core/embeddings/embeddings.service.js';
+import { ParserFactory } from '../../core/parsers/parser-factory.js';
+import { WorkspaceParser } from '../../core/parsers/workspace-parser.js';
+import { debugLog } from '../../core/utils/file-utils.js';
+import { getProjectName, UPSERT_PROJECT_QUERY, UPDATE_PROJECT_STATUS_QUERY } from '../../core/utils/project-id.js';
+import { WorkspaceDetector } from '../../core/workspace/index.js';
+import { Neo4jService, QUERIES } from '../../storage/neo4j/neo4j.service.js';
+import { PARSING } from '../constants.js';
+import { GraphGeneratorHandler } from '../handlers/graph-generator.handler.js';
+import { ParallelImportHandler } from '../handlers/parallel-import.handler.js';
+import { StreamingImportHandler } from '../handlers/streaming-import.handler.js';
+const sendMessage = (msg) => {
+    parentPort?.postMessage(msg);
+};
+const sendProgress = (phase, filesProcessed, filesTotal, nodesImported, edgesImported, currentChunk, totalChunks) => {
+    sendMessage({
+        type: 'progress',
+        data: { phase, filesProcessed, filesTotal, nodesImported, edgesImported, currentChunk, totalChunks },
+    });
+};
+const runParser = async () => {
+    const config = workerData;
+    const startTime = Date.now();
+    let resolvedProjectId = config.projectId;
+    let neo4jService = null;
+    try {
+        sendProgress('discovery', 0, 0, 0, 0, 0, 0);
+        neo4jService = new Neo4jService();
+        const embeddingsService = new EmbeddingsService();
+        const graphGeneratorHandler = new GraphGeneratorHandler(neo4jService, embeddingsService);
+        const lazyLoad = true;
+        const workspaceDetector = new WorkspaceDetector();
+        await debugLog('Detecting workspace', { projectPath: config.projectPath });
+        const workspaceConfig = await workspaceDetector.detect(config.projectPath);
+        await debugLog('Workspace detection result', {
+            type: workspaceConfig.type,
+            rootPath: workspaceConfig.rootPath,
+            packageCount: workspaceConfig.packages.length,
+        });
+        let detectedProjectType;
+        if (config.projectType === 'auto') {
+            detectedProjectType = await ParserFactory.detectProjectType(config.projectPath);
+            await debugLog('Auto-detected project type', { projectType: detectedProjectType });
+        }
+        else {
+            detectedProjectType = config.projectType;
+        }
+        let parser;
+        if (workspaceConfig.type !== 'single' && workspaceConfig.packages.length > 1) {
+            await debugLog('Using WorkspaceParser', {
+                type: workspaceConfig.type,
+                packageCount: workspaceConfig.packages.length,
+            });
+            parser = new WorkspaceParser(workspaceConfig, config.projectId, lazyLoad, detectedProjectType);
+            resolvedProjectId = parser.getProjectId();
+        }
+        else {
+            await debugLog('Using single project mode');
+            parser = ParserFactory.createParser({
+                workspacePath: config.projectPath,
+                tsConfigPath: config.tsconfigPath,
+                projectType: detectedProjectType,
+                projectId: config.projectId,
+                lazyLoad,
+            });
+            resolvedProjectId = parser.getProjectId();
+        }
+        const sourceFiles = await parser.discoverSourceFiles();
+        const totalFiles = sourceFiles.length;
+        const chunkSize = config.chunkSize > 0 ? config.chunkSize : PARSING.defaultChunkSize;
+        graphGeneratorHandler.setProjectId(resolvedProjectId);
+        await neo4jService.run(QUERIES.CLEAR_PROJECT, { projectId: resolvedProjectId });
+        const projectName = await getProjectName(config.projectPath);
+        await neo4jService.run(UPSERT_PROJECT_QUERY, {
+            projectId: resolvedProjectId,
+            name: projectName,
+            path: config.projectPath,
+            status: 'parsing',
+        });
+        await debugLog('Project node created', { projectId: resolvedProjectId, name: projectName });
+        let totalNodesImported = 0;
+        let totalEdgesImported = 0;
+        const onProgress = async (progress) => {
+            sendProgress(progress.phase, progress.current, progress.total, progress.details?.nodesCreated ?? 0, progress.details?.edgesCreated ?? 0, progress.details?.chunkIndex ?? 0, progress.details?.totalChunks ?? 0);
+        };
+        const useParallel = totalFiles >= PARSING.parallelThreshold;
+        if (useParallel) {
+            await debugLog('Using parallel parsing', { totalFiles });
+            const parallelHandler = new ParallelImportHandler(graphGeneratorHandler);
+            const result = await parallelHandler.importProjectParallel(parser, sourceFiles, {
+                chunkSize,
+                projectId: resolvedProjectId,
+                projectPath: config.projectPath,
+                tsconfigPath: config.tsconfigPath,
+                projectType: detectedProjectType,
+                onProgress,
+            });
+            totalNodesImported = result.nodesImported;
+            totalEdgesImported = result.edgesImported;
+        }
+        else {
+            await debugLog('Using sequential parsing', { totalFiles });
+            const streamingHandler = new StreamingImportHandler(graphGeneratorHandler);
+            const result = await streamingHandler.importProjectStreaming(parser, {
+                chunkSize,
+                projectId: resolvedProjectId,
+                onProgress,
+            });
+            totalNodesImported = result.nodesImported;
+            totalEdgesImported = result.edgesImported;
+        }
+        await neo4jService.run(UPDATE_PROJECT_STATUS_QUERY, {
+            projectId: resolvedProjectId,
+            status: 'complete',
+            nodeCount: totalNodesImported,
+            edgeCount: totalEdgesImported,
+        });
+        await debugLog('Project node updated', {
+            projectId: resolvedProjectId,
+            status: 'complete',
+            nodeCount: totalNodesImported,
+            edgeCount: totalEdgesImported,
+        });
+        sendMessage({
+            type: 'complete',
+            data: {
+                nodesImported: totalNodesImported,
+                edgesImported: totalEdgesImported,
+                elapsedMs: Date.now() - startTime,
+            },
+        });
+    }
+    catch (error) {
+        try {
+            const serviceForUpdate = neo4jService ?? new Neo4jService();
+            await serviceForUpdate.run(UPDATE_PROJECT_STATUS_QUERY, {
+                projectId: resolvedProjectId,
+                status: 'failed',
+                nodeCount: 0,
+                edgeCount: 0,
+            });
+            if (!neo4jService) {
+                await serviceForUpdate.close();
+            }
+        }
+        catch {
+            // Ignore errors updating project status on failure
+        }
+        sendMessage({
+            type: 'error',
+            error: error.message ?? String(error),
+        });
+    }
+    finally {
+        if (neo4jService) {
+            try {
+                await neo4jService.close();
+            }
+            catch {
+                // Ignore cleanup errors
+            }
+        }
+    }
+};
+runParser().catch((err) => {
+    sendMessage({
+        type: 'error',
+        error: err.message ?? String(err),
+    });
+});

package/dist/mcp/workers/worker.pool.js

@@ -0,0 +1,54 @@
+import { Worker } from 'worker_threads';
+export class ParallelPool {
+    workerPath;
+    numWorkers;
+    constructor(workerPath, numWorkers = 2) {
+        this.workerPath = workerPath;
+        this.numWorkers = numWorkers;
+    }
+    async run(items) {
+        const start = Date.now();
+        const indexBuffer = new SharedArrayBuffer(4);
+        const sharedIndex = new Int32Array(indexBuffer);
+        const workerPromises = Array.from({ length: this.numWorkers }, (_, id) => this.spawnWorker(id, items, indexBuffer));
+        const workerResults = await Promise.all(workerPromises);
+        const results = [];
+        const workerTaskCounts = [];
+        for (const { results: map, count } of workerResults) {
+            workerTaskCounts.push(count);
+            for (const [i, result] of map) {
+                results[i] = result;
+            }
+        }
+        return {
+            results,
+            stats: {
+                workerTaskCounts,
+                totalTasks: items.length,
+                totalTimeMs: Date.now() - start,
+            },
+        };
+    }
+    spawnWorker(workerId, items, indexBuffer) {
+        return new Promise((resolve, reject) => {
+            const worker = new Worker(this.workerPath, {
+                workerData: {
+                    items,
+                    indexBuffer,
+                    total: items.length,
+                    workerId,
+                },
+            });
+            worker.on('message', (result) => {
+                worker.terminate();
+                resolve(result);
+            });
+            worker.on('error', reject);
+            worker.on('exit', (code) => {
+                if (code !== 0) {
+                    reject(new Error(`Worker ${workerId} exited with code ${code}`));
+                }
+            });
+        });
+    }
+}