@softerist/heuristic-mcp 2.1.47 → 3.0.0

package/features/hybrid-search.js

@@ -1,7 +1,7 @@
-import path from "path";
-import fs from "fs/promises";
-import { dotSimilarity } from "../lib/utils.js";
-import { extractSymbolsFromContent } from "../lib/call-graph.js";
+import path from 'path';
+import fs from 'fs/promises';
+import { dotSimilarity } from '../lib/utils.js';
+import { extractSymbolsFromContent } from '../lib/call-graph.js';
 
 export class HybridSearch {
   constructor(embedder, cache, config) {
@@ -11,6 +11,14 @@ export class HybridSearch {
     this.fileModTimes = new Map(); // Cache for file modification times
   }
 
+  async getChunkContent(chunkOrIndex) {
+    return await this.cache.getChunkContent(chunkOrIndex);
+  }
+
+  getChunkVector(chunk) {
+    return this.cache.getChunkVector(chunk);
+  }
+
   getAnnCandidateCount(maxResults, totalChunks) {
     const minCandidates = this.config.annMinCandidates ?? 0;
     const maxCandidates = this.config.annMaxCandidates ?? totalChunks;
@@ -26,7 +34,13 @@ export class HybridSearch {
 
     for (const file of uniqueFiles) {
       if (!this.fileModTimes.has(file)) {
-        missing.push(file);
+        // Try to get from cache metadata first (fast)
+        const meta = this.cache.getFileMeta(file);
+        if (meta && typeof meta.mtimeMs === 'number') {
+          this.fileModTimes.set(file, meta.mtimeMs);
+        } else {
+          missing.push(file);
+        }
       }
     }
 
@@ -34,17 +48,31 @@ export class HybridSearch {
       return;
     }
 
-    const BATCH_SIZE = 200;
-    for (let i = 0; i < missing.length; i += BATCH_SIZE) {
-      const batch = missing.slice(i, i + BATCH_SIZE);
-      await Promise.all(batch.map(async file => {
+    // Concurrency-limited execution to avoid EMFILE
+    const CONCURRENCY = 50;
+    let index = 0;
+    
+    const worker = async () => {
+      while (index < missing.length) {
+        const file = missing[index++];
+        if (!file) break; // Safety check
         try {
           const stats = await fs.stat(file);
           this.fileModTimes.set(file, stats.mtimeMs);
         } catch {
           this.fileModTimes.set(file, null);
         }
-      }));
+      }
+    };
+
+    await Promise.all(Array.from({ length: Math.min(CONCURRENCY, missing.length) }, worker));
+
+    // Prevent unbounded growth (simple eviction)
+    if (this.fileModTimes.size > 5000) {
+      for (const [key] of this.fileModTimes) {
+        this.fileModTimes.delete(key);
+        if (this.fileModTimes.size <= 4000) break;
+      }
     }
   }
 
@@ -54,82 +82,190 @@ export class HybridSearch {
   }
 
   async search(query, maxResults) {
-    const vectorStore = this.cache.getVectorStore();
-
-    if (vectorStore.length === 0) {
-      return {
-        results: [],
-        message: "No code has been indexed yet. Please wait for initial indexing to complete."
-      };
-    }
+    try {
+      if (typeof this.cache.ensureLoaded === 'function') {
+        await this.cache.ensureLoaded();
+      }
+      this.cache.startRead();
+      
+      const storeSize = this.cache.getStoreSize();
+
+      if (storeSize === 0) {
+        return {
+          results: [],
+          message: 'No code has been indexed yet. Please wait for initial indexing to complete.',
+        };
+      }
 
-    // Generate query embedding
-    const queryEmbed = await this.embedder(query, { pooling: "mean", normalize: true });
-    const queryVector = Array.from(queryEmbed.data);
-    const queryVectorTyped = queryEmbed.data;
+      // Generate query embedding
+      console.info(`[Search] Query: "${query}"`);
+      const queryEmbed = await this.embedder(query, {
+        pooling: 'mean',
+        normalize: true,
+      });
+      const queryVector = queryEmbed.data; // Keep as Float32Array for performance
+      const queryVectorTyped = queryVector;
 
-    let candidates = vectorStore;
+      let candidateIndices = null; // null implies full scan of all chunks
     let usedAnn = false;
+
     if (this.config.annEnabled) {
-      const candidateCount = this.getAnnCandidateCount(maxResults, vectorStore.length);
+      const candidateCount = this.getAnnCandidateCount(maxResults, storeSize);
       const annLabels = await this.cache.queryAnn(queryVectorTyped, candidateCount);
       if (annLabels && annLabels.length >= maxResults) {
         usedAnn = true;
-        const seen = new Set();
-        candidates = annLabels
-          .map((index) => {
-            if (seen.has(index)) return null;
-            seen.add(index);
-            return vectorStore[index];
-          })
-          .filter(Boolean);
+        console.info(`[Search] Using ANN index (${annLabels.length} candidates)`);
+        candidateIndices = Array.from(new Set(annLabels)); // dedupe
       }
     }
 
-    if (usedAnn && candidates.length < maxResults) {
-      candidates = vectorStore;
-      usedAnn = false;
+    if (!usedAnn) {
+      console.info(`[Search] Using full scan (${storeSize} chunks)`);
     }
 
-    if (this.config.recencyBoost > 0) {
-      await this.populateFileModTimes(candidates.map(chunk => chunk.file));
+    if (usedAnn && candidateIndices && candidateIndices.length < maxResults) {
+      console.info(`[Search] ANN returned fewer results (${candidateIndices.length}) than requested (${maxResults}), augmenting with full scan...`);
+      candidateIndices = null; // Fallback to full scan to ensure we don't miss anything relevant
+      usedAnn = false;
     }
 
-    // Score all chunks (synchronous map now, much faster)
-    const scoredChunks = candidates.map(chunk => {
-      // Semantic similarity (vectors are normalized)
-      let score = dotSimilarity(queryVector, chunk.vector) * this.config.semanticWeight;
+    const lowerQuery = query.toLowerCase();
+    const queryWords =
+      lowerQuery.length > 1 ? lowerQuery.split(/\s+/).filter((word) => word.length > 2) : [];
+    const queryWordCount = queryWords.length;
+
+    if (usedAnn && candidateIndices && lowerQuery.length > 1) {
+      let exactMatchCount = 0;
+      for (const index of candidateIndices) {
+        const content = await this.getChunkContent(index);
+        if (content && content.toLowerCase().includes(lowerQuery)) {
+          exactMatchCount++;
+        }
+      }
 
-      // Exact match boost
-      const lowerQuery = query.toLowerCase();
-      const lowerContent = chunk.content.toLowerCase();
+      if (exactMatchCount < maxResults) {
+        // Fallback to full scan if keyword constraint isn't met in candidates
+        // Note: This is expensive as it iterates everything.
+        // Optimization: Only do this for small-ish codebases to avoid UI freeze
+        const MAX_FULL_SCAN_SIZE = 2000;
+        
+        if (storeSize <= MAX_FULL_SCAN_SIZE) {
+          const seen = new Set(candidateIndices);
+          
+          // Full scan logic for keyword augmentation
+          // Iterate by index with yielding
+          const FALLBACK_BATCH = 100;
+          for (let i = 0; i < storeSize; i += FALLBACK_BATCH) {
+             if (i > 0) await new Promise(r => setTimeout(r, 0)); // Yield
+
+             const limit = Math.min(storeSize, i + FALLBACK_BATCH);
+             for (let j = i; j < limit; j++) {
+                if (seen.has(j)) continue;
+                
+                // Lazy load content only if needed (this might be slow for huge repo)
+                // But `getChunkContent` should use cache.
+                const content = await this.getChunkContent(j);
+                if (content && content.toLowerCase().includes(lowerQuery)) {
+                    seen.add(j);
+                    candidateIndices.push(j);
+                }
+             }
+          }
+        } else {
+          console.info(`[Search] Skipping full scan fallback (store size ${storeSize} > ${MAX_FULL_SCAN_SIZE})`);
+        }
+      }
+    }
 
-      if (lowerContent.includes(lowerQuery)) {
-        score += this.config.exactMatchBoost;
+    // Recency pre-processing
+    let recencyBoostEnabled = this.config.recencyBoost > 0;
+    let now = Date.now();
+    let recencyDecayMs = (this.config.recencyDecayDays || 30) * 24 * 60 * 60 * 1000;
+    let semanticWeight = this.config.semanticWeight;
+    let exactMatchBoost = this.config.exactMatchBoost;
+    let recencyBoost = this.config.recencyBoost;
+
+    if (recencyBoostEnabled) {
+      const candidates = candidateIndices
+        ? candidateIndices.map((idx) => this.cache.getChunk(idx)).filter(Boolean)
+        : Array.from({ length: storeSize }, (_, i) => this.cache.getChunk(i)).filter(Boolean);
+      // optimization: avoid IO storm during full scan fallbacks
+      // For large candidate sets, we strictly rely on cached metadata
+      // For small sets, we allow best-effort fs.stat
+      if (candidates.length <= 1000) {
+        await this.populateFileModTimes(candidates.map((chunk) => chunk.file));
       } else {
-        // Partial word matching
-        const queryWords = lowerQuery.split(/\s+/);
-        const matchedWords = queryWords.filter(word =>
-          word.length > 2 && lowerContent.includes(word)
-        ).length;
-        score += (matchedWords / queryWords.length) * 0.3;
+        // Bulk pre-populate from cache only (no syscalls)
+        for (const chunk of candidates) {
+          if (!this.fileModTimes.has(chunk.file)) {
+            const meta = this.cache.getFileMeta(chunk.file);
+            if (meta && typeof meta.mtimeMs === 'number') {
+              this.fileModTimes.set(chunk.file, meta.mtimeMs);
+            }
+          }
+        }
       }
+    }
 
-      // Recency boost - recently modified files rank higher
-      if (this.config.recencyBoost > 0) {
-        const mtime = this.fileModTimes.get(chunk.file);
-        if (typeof mtime === "number") {
-          const daysSinceModified = (Date.now() - mtime) / (1000 * 60 * 60 * 24);
-          const decayDays = this.config.recencyDecayDays || 30;
+    // Score all chunks (batched to prevent blocking event loop)
+    const BATCH_SIZE = 500;
+    const scoredChunks = [];
 
-          // Linear decay: full boost at 0 days, no boost after decayDays
-          const recencyScore = Math.max(0, 1 - (daysSinceModified / decayDays));
-          score += recencyScore * this.config.recencyBoost;
-        }
+    // Process in batches
+    // Candidates is now implicitly range 0..storeSize OR candidateIndices
+    const totalCandidates = candidateIndices ? candidateIndices.length : storeSize;
+
+    for (let i = 0; i < totalCandidates; i += BATCH_SIZE) {
+      // Allow event loop to tick between batches
+      if (i > 0) {
+        await new Promise((resolve) => setTimeout(resolve, 0));
       }
 
-      return { ...chunk, score };
-    });
+      const limit = Math.min(totalCandidates, i + BATCH_SIZE);
+      
+      for (let j = i; j < limit; j++) {
+        const idx = candidateIndices ? candidateIndices[j] : j;
+        
+        // Lazy load keys
+        const vector = this.cache.getVector(idx);
+        if (!vector) continue;
+        
+        // Ensure vector compatibility (dotSimilarity now checks length too)
+        let score = dotSimilarity(queryVector, vector) * semanticWeight;
+
+        // Exact match boost
+        const content = await this.getChunkContent(idx);
+        const lowerContent = content ? content.toLowerCase() : '';
+
+        if (lowerContent && lowerContent.includes(lowerQuery)) {
+          score += exactMatchBoost;
+        } else if (lowerContent && queryWordCount > 0) {
+          // Partial word matching (optimized)
+          let matchedWords = 0;
+          for (let k = 0; k < queryWordCount; k++) {
+            if (lowerContent.includes(queryWords[k])) matchedWords++;
+          }
+          score += (matchedWords / queryWordCount) * 0.3;
+        }
+
+        // Needs chunk info for result
+        const chunkInfo = this.cache.getChunk(idx);
+        
+        // Recency boost
+        if (recencyBoostEnabled && chunkInfo) {
+              const mtime = this.fileModTimes.get(chunkInfo.file);
+              if (typeof mtime === 'number') {
+                const ageMs = now - mtime;
+                const recencyFactor = Math.max(0, 1 - ageMs / recencyDecayMs);
+                score += recencyFactor * recencyBoost;
+              }
+        }
+        
+        if (chunkInfo) {
+            scoredChunks.push({ ...chunkInfo, score, content });
+        }
+      }
+    }
 
     // Sort by initial score
     scoredChunks.sort((a, b) => b.score - a.score);
@@ -140,7 +276,8 @@ export class HybridSearch {
       const topN = Math.min(5, scoredChunks.length);
       const symbolsFromTop = new Set();
       for (let i = 0; i < topN; i++) {
-        const symbols = extractSymbolsFromContent(scoredChunks[i].content);
+        const content = await this.getChunkContent(scoredChunks[i]);
+        const symbols = extractSymbolsFromContent(content || '');
         for (const sym of symbols) {
           symbolsFromTop.add(sym);
         }
@@ -157,62 +294,85 @@ export class HybridSearch {
             chunk.score += proximity * this.config.callGraphBoost;
           }
         }
-
         // Re-sort after applying call graph boost
         scoredChunks.sort((a, b) => b.score - a.score);
       }
     }
 
     // Get top results
-    const results = scoredChunks.slice(0, maxResults);
+    const results = await Promise.all(scoredChunks.slice(0, maxResults).map(async (chunk) => {
+      if (chunk.content === undefined || chunk.content === null) {
+        return { ...chunk, content: await this.getChunkContent(chunk) };
+      }
+      return chunk;
+    }));
+
+    if (results.length > 0) {
+      console.info(`[Search] Found ${results.length} results. Top score: ${results[0].score.toFixed(4)}`);
+    } else {
+      console.info('[Search] No results found.');
+    }
 
     return { results, message: null };
+    } finally {
+      this.cache.endRead();
+    }
   }
 
-  formatResults(results) {
+  async formatResults(results) {
     if (results.length === 0) {
-      return "No matching code found for your query.";
+      return 'No matching code found for your query.';
     }
 
-    return results.map((r, idx) => {
+    const formatted = await Promise.all(results.map(async (r, idx) => {
       const relPath = path.relative(this.config.searchDirectory, r.file);
-      return `## Result ${idx + 1} (Relevance: ${(r.score * 100).toFixed(1)}%)\n` +
-             `**File:** \`${relPath}\`\n` +
-             `**Lines:** ${r.startLine}-${r.endLine}\n\n` +
-             "```" + path.extname(r.file).slice(1) + "\n" +
-             r.content + "\n" +
-             "```\n";
-    }).join("\n");
+      const content = r.content ?? await this.getChunkContent(r);
+      return (
+        `## Result ${idx + 1} (Relevance: ${(r.score * 100).toFixed(1)}%)\n` +
+        `**File:** \`${relPath}\`\n` +
+        `**Lines:** ${r.startLine}-${r.endLine}\n\n` +
+        '```' +
+        path.extname(r.file).slice(1) +
+        '\n' +
+        content +
+        '\n' +
+        '```\n'
+      );
+    }));
+
+    return formatted.join('\n');
   }
 }
 
 // MCP Tool definition for this feature
 export function getToolDefinition(config) {
   return {
-    name: "a_semantic_search",
-    description: "Performs intelligent hybrid code search combining semantic understanding with exact text matching. Ideal for finding code by meaning (e.g., 'authentication logic', 'database queries') even with typos or variations. Returns the most relevant code snippets with file locations and line numbers.",
+    name: 'a_semantic_search',
+    description:
+      "Performs intelligent hybrid code search combining semantic understanding with exact text matching. Ideal for finding code by meaning (e.g., 'authentication logic', 'database queries') even with typos or variations. Returns the most relevant code snippets with file locations and line numbers.",
     inputSchema: {
-      type: "object",
+      type: 'object',
       properties: {
         query: {
-          type: "string",
-          description: "Search query - can be natural language (e.g., 'where do we handle user login') or specific terms"
+          type: 'string',
+          description:
+            "Search query - can be natural language (e.g., 'where do we handle user login') or specific terms",
         },
         maxResults: {
-          type: "number",
-          description: "Maximum number of results to return (default: from config)",
-          default: config.maxResults
-        }
+          type: 'number',
+          description: 'Maximum number of results to return (default: from config)',
+          default: config.maxResults,
+        },
       },
-      required: ["query"]
+      required: ['query'],
     },
     annotations: {
-      title: "Semantic Code Search",
+      title: 'Semantic Code Search',
       readOnlyHint: true,
       destructiveHint: false,
       idempotentHint: true,
-      openWorldHint: false
-    }
+      openWorldHint: false,
+    },
   };
 }
 
@@ -225,13 +385,13 @@ export async function handleToolCall(request, hybridSearch) {
 
   if (message) {
     return {
-      content: [{ type: "text", text: message }]
+      content: [{ type: 'text', text: message }],
     };
   }
 
-  const formattedText = hybridSearch.formatResults(results);
+  const formattedText = await hybridSearch.formatResults(results);
 
   return {
-    content: [{ type: "text", text: formattedText }]
+    content: [{ type: 'text', text: formattedText }],
   };
 }