@softerist/heuristic-mcp 3.0.17 → 3.1.0

package/lib/embedding-worker.js

@@ -18,8 +18,8 @@ function getGlobalCacheDir() {
   return process.env.XDG_CACHE_HOME || path.join(os.homedir(), '.cache');
 }
 
-// Override console.info/warn to write to stderr so we don't break the MCP JSON-RPC protocol on stdout
-// Labels help distinguish worker output from main thread during debugging
+
+
 console.info = (...args) => console.error('[INFO]', ...args);
 console.warn = (...args) => console.error('[WARN]', ...args);
 
@@ -143,10 +143,10 @@ async function isPathInsideWorkspace(targetPath) {
   }
 }
 
-// NOTE: sliceAndNormalize is duplicated here because worker_threads run in a
-// separate context and dynamic imports add latency. The canonical implementation
-// is in lib/slice-normalize.js. Changes should be synchronized.
-// IMPORTANT: test/slice-normalize-sync.test.js verifies both implementations match.
+
+
+
+
 function sliceAndNormalize(vector, targetDim) {
   if (!targetDim || targetDim >= vector.length) {
     return vector;
@@ -165,11 +165,11 @@ function sliceAndNormalize(vector, targetDim) {
   return sliced;
 }
 
-// Get embedding dimension from workerData (null = full dimensions)
+
 const embeddingDimension = workerData.embeddingDimension || null;
 
-// Initialize the embedding model once when worker starts
-// Use a promise to handle concurrent calls to initializeEmbedder safely
+
+
 let embedderPromise = null;
 
 async function initializeEmbedder() {
@@ -199,7 +199,7 @@ async function initializeEmbedder() {
         );
         return model;
       } catch (err) {
-        embedderPromise = null; // Reset promise so we can retry later
+        embedderPromise = null; 
         throw err;
       }
     })();
@@ -237,10 +237,7 @@ function getEmbeddingTensor(output, { requireDimsForBatch = false, batchSize = n
   return { data, hiddenSize };
 }
 
-/**
- * Legacy Protocol: Process chunks with optimized single-text embedding
- * Streams results in batches.
- */
+
 async function processChunks(chunks, batchId) {
   const embedder = await initializeEmbedder();
   let results = [];
@@ -248,11 +245,11 @@ async function processChunks(chunks, batchId) {
   const failFastState = createFailFastState('legacy chunk embedding');
 
   const flush = (done = false) => {
-    // Only flush intermediate results when we have enough for a batch
+    
     if (!done && results.length < RESULT_BATCH_SIZE) return;
 
-    // final batch might be empty if chunks was empty or perfectly divisible by RESULT_BATCH_SIZE
-    // but we still send it to signal we are done.
+    
+    
 
     const payload = {
       type: 'results',
@@ -275,12 +272,12 @@ async function processChunks(chunks, batchId) {
         pooling: 'mean',
         normalize: true,
       });
-      // CRITICAL: Deep copy to release ONNX tensor memory
+      
       const { data } = getEmbeddingTensor(output);
       let vector = new Float32Array(data);
-      // Apply MRL dimension slicing if configured
+      
       vector = sliceAndNormalize(vector, embeddingDimension);
-      // Properly dispose tensor to release ONNX runtime memory
+      
       if (typeof output.dispose === 'function')
         try {
           output.dispose();
@@ -317,27 +314,21 @@ async function processChunks(chunks, batchId) {
 
   flush(true);
 
-  // Force GC if available to free massive tensor buffers immediately
+  
   maybeRunGc();
 }
 
-// =====================================================================
-// SHARED HELPER FUNCTIONS
-// =====================================================================
-
-/**
- * Prepares file content for processing: stat, size check, read content.
- * @param {string} file - File path
- * @param {string|null} providedContent - Pre-provided content (optional)
- * @param {number} maxFileSize - Maximum allowed file size
- * @returns {Promise<{status: string, reason?: string, content?: string, mtimeMs?: number, size?: number}>}
- */
+
+
+
+
+
 async function prepareFileContent(file, providedContent, maxFileSize) {
   let mtimeMs = null;
   let size = null;
   let content = null;
 
-  // Handle pre-provided content
+  
   if (typeof providedContent === 'string') {
     content = providedContent;
     const byteSize = Buffer.byteLength(content, 'utf-8');
@@ -348,16 +339,16 @@ async function prepareFileContent(file, providedContent, maxFileSize) {
     return { status: 'ok', content, mtimeMs, size };
   }
 
-  // Stat the file
+  
   try {
     const st = await fs.stat(file);
     if (st.isDirectory()) {
       return { status: 'skipped', reason: 'is_directory', mtimeMs: st.mtimeMs, size: st.size };
     }
-    // NOTE: TOCTOU race between stat and readFile - file could grow after this check.
-    // Risk: Memory exhaustion if file grows significantly between stat and read.
-    // Mitigation: Process isolation (workers), soft limit (not security boundary),
-    // and Node.js readFile will throw ENOMEM before crashing the process.
+    
+    
+    
+    
     if (st.size > maxFileSize) {
       return { status: 'skipped', reason: 'too_large', mtimeMs: st.mtimeMs, size: st.size };
     }
@@ -367,18 +358,18 @@ async function prepareFileContent(file, providedContent, maxFileSize) {
     return { status: 'skipped', reason: `stat_failed: ${err.message}` };
   }
 
-  // Read content with size limit to mitigate TOCTOU race
+  
   try {
     const handle = await fs.open(file, 'r');
     try {
-      // Defense-in-depth: Allocate enough to detect files that grow beyond limit.
-      // We use maxFileSize + 1 (capped at 64MB) to reliably detect overflow without
-      // excessive memory allocation for small files.
+      
+      
+      
       const bufferSize = Math.min(maxFileSize + 1, Math.max(size + 1024, 64 * 1024));
       const buffer = Buffer.alloc(bufferSize);
       const { bytesRead } = await handle.read(buffer, 0, bufferSize, 0);
       
-      // Reject if file grew beyond limit between stat and read
+      
       if (bytesRead > maxFileSize) {
         return { status: 'skipped', reason: 'too_large_after_read', mtimeMs, size: bytesRead };
       }
@@ -395,23 +386,17 @@ async function prepareFileContent(file, providedContent, maxFileSize) {
   return { status: 'ok', content, mtimeMs, size };
 }
 
-/**
- * Processes file metadata: hash check, call graph, chunking.
- * @param {string} file - File path
- * @param {string} content - File content
- * @param {object} options - { force, expectedHash, callGraphEnabled, chunkConfig, workerLabel }
- * @returns {{status: string, hash: string, callData: object|null, chunks: Array}}
- */
+
 function processFileMetadata(file, content, options) {
   const { force, expectedHash, callGraphEnabled, chunkConfig, workerLabel } = options;
 
-  // Hash and unchanged short-circuit
+  
   const hash = hashContent(content);
   if (!force && expectedHash && expectedHash === hash) {
     return { status: 'unchanged', hash, callData: null, chunks: [] };
   }
 
-  // Call graph extraction (optional)
+  
   let callData = null;
   if (callGraphEnabled) {
     try {
@@ -424,16 +409,13 @@ function processFileMetadata(file, content, options) {
     }
   }
 
-  // Chunking
+  
   const chunks = smartChunk(content, file, chunkConfig);
 
   return { status: 'processing', hash, callData, chunks };
 }
 
-/**
- * New Protocol: Process entire file (read, chunk, embed) in worker.
- * Returns results once processing is complete.
- */
+
 async function processFileTask(message) {
   const embedder = await initializeEmbedder();
   const failFastState = createFailFastState(`file-task ${path.basename(message.file || '')}`);
@@ -442,7 +424,7 @@ async function processFileTask(message) {
   const force = !!message.force;
   const expectedHash = message.expectedHash || null;
 
-  // Check workspace boundary
+  
   if (!(await isPathInsideWorkspace(file))) {
     if (workerData.verbose) {
       console.warn(`[Worker ${workerData.workerId}] Skipping file outside workspace: ${file}`);
@@ -453,14 +435,14 @@ async function processFileTask(message) {
   const maxFileSize = Number.isFinite(workerData.maxFileSize) ? workerData.maxFileSize : Infinity;
   const callGraphEnabled = !!workerData.callGraphEnabled;
 
-  // 1-2) Prepare file content using shared helper
+  
   const prep = await prepareFileContent(file, message.content, maxFileSize);
   if (prep.status !== 'ok') {
     return { status: prep.status, reason: prep.reason, mtimeMs: prep.mtimeMs, size: prep.size };
   }
   const { content, mtimeMs, size } = prep;
 
-  // 3-5) Hash, call graph, chunking using shared helper
+  
   const chunkConfig = {
     ...(workerData.config || {}),
     ...(workerData.chunkConfig || {}),
@@ -482,12 +464,12 @@ async function processFileTask(message) {
 
   const { hash, callData, chunks } = meta;
 
-  // 6) Embed chunks in batches for performance
+  
   const results = [];
   const transferList = [];
 
-  // Batch size for inference (balance between speed and memory)
-  // Configurable via workerData, default 4 balances memory and throughput
+  
+  
   const INFERENCE_BATCH_SIZE = Number.isInteger(workerData.inferenceBatchSize)
     ? workerData.inferenceBatchSize
     : DEFAULT_INFERENCE_BATCH_SIZE;
@@ -498,14 +480,14 @@ async function processFileTask(message) {
     const batchTexts = batchChunks.map((c) => c.text);
 
     try {
-      // Run inference on the batch
+      
       const output = await embedder(batchTexts, {
         pooling: 'mean',
         normalize: true,
       });
 
-      // Output is a Tensor with shape [batch_size, hidden_size]
-      // data is a flat Float32Array
+      
+      
       const { data, hiddenSize } = getEmbeddingTensor(output, {
         requireDimsForBatch: true,
         batchSize: batchChunks.length,
@@ -514,16 +496,16 @@ async function processFileTask(message) {
       for (let j = 0; j < batchChunks.length; j++) {
         const c = batchChunks[j];
 
-        // Slice the flat buffer to get this chunk's vector
-        // specific slice for this element
+        
+        
         const start = j * hiddenSize;
         const end = start + hiddenSize;
         const vectorView =
           typeof data.subarray === 'function' ? data.subarray(start, end) : data.slice(start, end);
 
-        // Deep copy to ensure independent buffer for transfer
+        
         let vector = new Float32Array(vectorView);
-        // Apply MRL dimension slicing if configured
+        
         vector = sliceAndNormalize(vector, embeddingDimension);
 
         results.push({
@@ -534,7 +516,7 @@ async function processFileTask(message) {
         });
         transferList.push(vector.buffer);
       }
-      // Properly dispose tensor to release ONNX runtime memory
+      
       if (typeof output.dispose === 'function')
         try {
           output.dispose();
@@ -549,7 +531,7 @@ async function processFileTask(message) {
         noteEmbeddingFailure(failFastState, err);
         throw err;
       }
-      // Fallback: if batch fails (e.g. OOM), try one by one for this batch
+      
       console.warn(`${workerLabel} Batch inference failed (${err.name}), retrying individually: ${err.message}`);
       noteEmbeddingFailure(failFastState, err);
 
@@ -558,9 +540,9 @@ async function processFileTask(message) {
           const output = await embedder(c.text, { pooling: 'mean', normalize: true });
           const { data } = getEmbeddingTensor(output);
           let vector = new Float32Array(data);
-          // Apply MRL dimension slicing if configured
+          
           vector = sliceAndNormalize(vector, embeddingDimension);
-          // Properly dispose tensor to release ONNX runtime memory
+          
           if (typeof output.dispose === 'function')
             try {
               output.dispose();
@@ -578,9 +560,9 @@ async function processFileTask(message) {
           transferList.push(vector.buffer);
           noteEmbeddingSuccess(failFastState);
         } catch (innerErr) {
-          // Note: No tensor disposal needed - embedder() threw before returning a tensor
+          
           console.warn(`${workerLabel} Chunk embedding failed: ${innerErr.message}`);
-          // We omit this chunk from results, effectively skipping it
+          
           noteEmbeddingFailure(failFastState, innerErr);
           if (isFatalRuntimeEmbeddingError(innerErr)) {
             throw innerErr;
@@ -589,7 +571,7 @@ async function processFileTask(message) {
       }
     }
 
-    // Yield to event loop briefly between batches and trigger GC
+    
     processedSinceGc += batchChunks.length;
     if (chunks.length > INFERENCE_BATCH_SIZE) {
       if (processedSinceGc >= 100) {
@@ -603,7 +585,7 @@ async function processFileTask(message) {
   return { status: 'indexed', hash, mtimeMs, size, callData, results, transferList };
 }
 
-// Listen for messages from main thread
+
 parentPort.on('message', async (message) => {
   try {
     if (!message || typeof message.type !== 'string') return;
@@ -613,14 +595,14 @@ parentPort.on('message', async (message) => {
       return;
     }
 
-    // Unload model to free memory
+    
     if (message.type === 'unload') {
       const wasLoaded = embedderPromise !== null;
       
       if (embedderPromise) {
         try {
           const embedder = await embedderPromise;
-          // Try to dispose the pipeline if it has a dispose method
+          
           if (embedder && typeof embedder.dispose === 'function') {
             try {
               await embedder.dispose();
@@ -633,11 +615,11 @@ parentPort.on('message', async (message) => {
           logInfo(`${workerLabel} Error during model unload: ${err.message}`);
         }
         
-        // Clear references
+        
         embedderPromise = null;
       }
       
-      // Trigger garbage collection if available
+      
       if (explicitGcEnabled && typeof global.gc === 'function') {
         const before = process.memoryUsage();
         global.gc();
@@ -651,7 +633,7 @@ parentPort.on('message', async (message) => {
       return;
     }
 
-    // ---- New protocol: file-level processing (chunking + embedding in worker) ----
+    
     if (message.type === 'processFile') {
       const { id } = message;
       if (!id) {
@@ -661,7 +643,7 @@ parentPort.on('message', async (message) => {
 
       const res = await processFileTask(message);
 
-      // Transfer vectors if present
+      
       if (res && res.transferList && res.transferList.length > 0) {
         const { transferList, ...payload } = res;
         parentPort.postMessage({ id, ...payload }, transferList);
@@ -671,19 +653,42 @@ parentPort.on('message', async (message) => {
       return;
     }
 
-    // ---- Batch file processing ----
-    if (message.type === 'processFiles') {
-      const { files, batchId } = message;
-      const batchTransfer = [];
-      const failFastState = createFailFastState('cross-file batch embedding');
+    
+    if (message.type === 'processFiles') {
+      const { files, batchId } = message;
+      const batchTransfer = [];
+      const failFastState = createFailFastState('cross-file batch embedding');
+      const PROGRESS_HEARTBEAT_MS = 15_000;
+      let lastProgressAt = 0;
+      const sendProgress = (payload = {}, { force = false } = {}) => {
+        const now = Date.now();
+        if (!force && now - lastProgressAt < PROGRESS_HEARTBEAT_MS) {
+          return;
+        }
+        lastProgressAt = now;
+        parentPort.postMessage({
+          type: 'progress',
+          batchId,
+          ...payload,
+        });
+      };
+
+      sendProgress(
+        {
+          stage: 'prepare',
+          filesCompleted: 0,
+          filesTotal: Array.isArray(files) ? files.length : 0,
+        },
+        { force: true }
+      );
 
-      // 1. Pre-process all files: Read, Stat, and Chunk
-      // We do this first to gather a massive list of chunks for batched inference
+      
+      
       const fileTasks = [];
-      const allPendingChunks = []; // { text, fileIndex, chunkIndex, startLine, endLine }
+      const allPendingChunks = []; 
 
-      for (let i = 0; i < files.length; i++) {
-        const fileMsg = files[i];
+      for (let i = 0; i < files.length; i++) {
+        const fileMsg = files[i];
 
         try {
           const file = fileMsg.file;
@@ -694,7 +699,7 @@ parentPort.on('message', async (message) => {
             : Infinity;
           const callGraphEnabled = !!workerData.callGraphEnabled;
 
-          // Check workspace boundary first
+          
           if (!(await isPathInsideWorkspace(file))) {
             fileTasks.push({
               file: fileMsg.file,
@@ -710,7 +715,7 @@ parentPort.on('message', async (message) => {
             continue;
           }
 
-          // Use shared helper for file preparation
+          
           const prep = await prepareFileContent(file, fileMsg.content, maxFileSize);
           if (prep.status !== 'ok') {
             fileTasks.push({
@@ -729,7 +734,7 @@ parentPort.on('message', async (message) => {
 
           const { content, mtimeMs, size } = prep;
 
-          // Use shared helper for metadata processing
+          
           const chunkConfig = {
             ...(workerData.config || {}),
             ...(workerData.chunkConfig || {}),
@@ -763,12 +768,12 @@ parentPort.on('message', async (message) => {
           const { hash, callData, chunks } = meta;
           const chunkCount = chunks.length;
 
-          // Trigger GC every 100 files
+          
           if ((i + 1) % 100 === 0) {
             maybeRunGc();
           }
 
-          // Register chunks for batching
+          
           if (chunks.length > 0) {
             for (const c of chunks) {
               allPendingChunks.push({
@@ -776,43 +781,49 @@ parentPort.on('message', async (message) => {
                 text: c.text,
                 startLine: c.startLine,
                 endLine: c.endLine,
-                vectorBuffer: null, // to be filled
+                vectorBuffer: null, 
               });
             }
           }
 
-          fileTasks.push({
-            file: fileMsg.file,
-            status: 'indexed', // Provisional, pending embedding
-            reason: null,
-            hash,
+          fileTasks.push({
+            file: fileMsg.file,
+            status: 'indexed', 
+            reason: null,
+            hash,
             mtimeMs,
             size,
             callData,
             expectedChunks: chunkCount,
-            results: [], // Will store chunk results
+            results: [], 
           });
         } catch (error) {
-          fileTasks.push({
-            file: fileMsg.file,
-            status: 'error',
-            error: error.message,
-            expectedChunks: 0,
-            results: [],
-          });
-        }
-      }
-
-      // 2. Run Batched Inference on all accumulated chunks
-      if (allPendingChunks.length > 0) {
+          fileTasks.push({
+            file: fileMsg.file,
+            status: 'error',
+            error: error.message,
+            expectedChunks: 0,
+            results: [],
+          });
+        } finally {
+          sendProgress({
+            stage: 'prepare',
+            filesCompleted: i + 1,
+            filesTotal: files.length,
+          });
+        }
+      }
+
+      
+      if (allPendingChunks.length > 0) {
         const embedder = await initializeEmbedder();
         const INFERENCE_BATCH_SIZE = Number.isInteger(workerData.inferenceBatchSize)
           ? workerData.inferenceBatchSize
           : DEFAULT_INFERENCE_BATCH_SIZE;
 
-        for (let i = 0; i < allPendingChunks.length; i += INFERENCE_BATCH_SIZE) {
-          const batchSlice = allPendingChunks.slice(i, i + INFERENCE_BATCH_SIZE);
-          const batchTexts = batchSlice.map((c) => c.text);
+        for (let i = 0; i < allPendingChunks.length; i += INFERENCE_BATCH_SIZE) {
+          const batchSlice = allPendingChunks.slice(i, i + INFERENCE_BATCH_SIZE);
+          const batchTexts = batchSlice.map((c) => c.text);
 
           try {
             const output = await embedder(batchTexts, { pooling: 'mean', normalize: true });
@@ -828,13 +839,13 @@ parentPort.on('message', async (message) => {
                 typeof data.subarray === 'function'
                   ? data.subarray(start, end)
                   : data.slice(start, end);
-              // Deep copy the view to avoid WASM memory issues, then apply MRL slicing
+              
               const vector = sliceAndNormalize(new Float32Array(vectorView), embeddingDimension);
 
               batchSlice[j].vectorBuffer = vector.buffer;
               batchTransfer.push(vector.buffer);
             }
-            // Properly dispose tensor to release ONNX runtime memory
+            
             if (typeof output.dispose === 'function')
               try {
                 output.dispose();
@@ -853,14 +864,14 @@ parentPort.on('message', async (message) => {
               `${workerLabel} Cross-file batch inference failed, retrying individually: ${err.message}`
             );
             noteEmbeddingFailure(failFastState, err);
-            // Fallback: individual embedding for this failed batch
-            for (const item of batchSlice) {
-              try {
-                const output = await embedder(item.text, { pooling: 'mean', normalize: true });
+            
+            for (const item of batchSlice) {
+              try {
+                const output = await embedder(item.text, { pooling: 'mean', normalize: true });
                 const { data } = getEmbeddingTensor(output);
-                // Deep copy and apply MRL slicing
+                
                 const vector = sliceAndNormalize(new Float32Array(data), embeddingDimension);
-                // Properly dispose tensor to release ONNX runtime memory
+                
                 if (typeof output.dispose === 'function')
                   try {
                     output.dispose();
@@ -878,18 +889,24 @@ parentPort.on('message', async (message) => {
                 if (isFatalRuntimeEmbeddingError(innerErr)) {
                   throw innerErr;
                 }
-              }
-            }
-          }
-
-          // Minimal yield to keep event loop breathing (optional, can be removed for max throughput)
-          if (allPendingChunks.length > 50 && i % 50 === 0) {
-            await new Promise((resolve) => setTimeout(resolve, 0));
-          }
+              }
+            }
+          }
+
+          sendProgress({
+            stage: 'embed',
+            chunksCompleted: Math.min(i + batchSlice.length, allPendingChunks.length),
+            chunksTotal: allPendingChunks.length,
+          });
+
+          
+          if (allPendingChunks.length > 50 && i % 50 === 0) {
+            await new Promise((resolve) => setTimeout(resolve, 0));
+          }
         }
       }
 
-      // 3. Reassemble results and validate
+      
       for (const chunkItem of allPendingChunks) {
         if (chunkItem.vectorBuffer) {
           const task = fileTasks[chunkItem.fileIndex];
@@ -913,21 +930,21 @@ parentPort.on('message', async (message) => {
         }
       }
 
-      // 4. Send response
-      // IMPORTANT: Clear vectorBuffer references BEFORE transfer to prevent use-after-transfer.
-      // After postMessage with transfer list, the ArrayBuffers become detached and any
-      // access to them (even for logging) would fail.
+      
+      
+      
+      
       const resultsForTransfer = fileTasks.map((task) => ({
         ...task,
         results: task.results.map((r) => ({
           startLine: r.startLine,
           endLine: r.endLine,
           text: r.text,
-          vectorBuffer: r.vectorBuffer, // Will be transferred
+          vectorBuffer: r.vectorBuffer, 
         })),
       }));
 
-      // Null out original references to avoid accidental access
+      
       for (const task of fileTasks) {
         for (const r of task.results) {
           r.vectorBuffer = null;
@@ -944,13 +961,13 @@ parentPort.on('message', async (message) => {
         batchTransfer
       );
 
-      // Explicitly clear references and trigger GC
+      
       batchTransfer.length = 0;
       maybeRunGc();
       return;
     }
 
-    // ---- Legacy protocol: batch of chunks prepared by main thread ----
+    
     if (message.type === 'process') {
       try {
         await processChunks(message.chunks || [], message.batchId);
@@ -964,10 +981,10 @@ parentPort.on('message', async (message) => {
       return;
     }
 
-    // Unknown type
+    
     parentPort.postMessage({ type: 'error', error: `Unknown message type: ${message.type}` });
   } catch (error) {
-    // If message had an id, respond via RPC style; otherwise legacy error
+    
     if (message && typeof message === 'object' && message.id) {
       parentPort.postMessage({ id: message.id, error: error.message });
     } else {
@@ -976,7 +993,7 @@ parentPort.on('message', async (message) => {
   }
 });
 
-// Signal that worker is ready
+
 initializeEmbedder()
   .then(() => {
     parentPort.postMessage({ type: 'ready' });