@colbymchenry/codegraph 0.6.8 → 0.7.2

package/dist/context/index.js

@@ -2,7 +2,7 @@
 /**
  * Context Builder
  *
- * Builds rich context for tasks by combining semantic search with graph traversal.
+ * Builds rich context for tasks by combining FTS search with graph traversal.
  * Outputs structured context ready to inject into Claude.
  */
 var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
@@ -42,9 +42,11 @@ Object.defineProperty(exports, "__esModule", { value: true });
 exports.formatContextAsJson = exports.formatContextAsMarkdown = exports.ContextBuilder = void 0;
 exports.createContextBuilder = createContextBuilder;
 const fs = __importStar(require("fs"));
+const path = __importStar(require("path"));
 const formatter_1 = require("./formatter");
 const errors_1 = require("../errors");
 const utils_1 = require("../utils");
+const query_utils_1 = require("../search/query-utils");
 /**
  * Extract likely symbol names from a natural language query
  *
@@ -82,6 +84,13 @@ function extractSymbolsFromQuery(query) {
             symbols.add(match[1]);
         }
     }
+    // Extract ALL_CAPS acronyms (2+ chars, e.g., REST, HTTP, LRU, API)
+    const acronymPattern = /\b([A-Z]{2,})\b/g;
+    while ((match = acronymPattern.exec(query)) !== null) {
+        if (match[1]) {
+            symbols.add(match[1]);
+        }
+    }
     // Extract dot.notation and split into parts (e.g., "app.isPackaged" -> ["app", "isPackaged"])
     const dotPattern = /\b([a-zA-Z][a-zA-Z0-9]*(?:\.[a-zA-Z][a-zA-Z0-9]*)+)\b/g;
     while ((match = dotPattern.exec(query)) !== null) {
@@ -96,13 +105,42 @@ function extractSymbolsFromQuery(query) {
             }
         }
     }
-    // Filter out common English words that might match patterns
+    // Extract plain lowercase identifiers (3+ chars, not already matched)
+    // Catches symbol names like "undo", "redo", "history", "render", "parse"
+    const lowercasePattern = /\b([a-z][a-z0-9]{2,})\b/g;
+    while ((match = lowercasePattern.exec(query)) !== null) {
+        if (match[1]) {
+            symbols.add(match[1]);
+        }
+    }
+    // Filter out common English words that aren't likely symbol names
     const commonWords = new Set([
         'the', 'and', 'for', 'with', 'from', 'this', 'that', 'have', 'been',
         'will', 'would', 'could', 'should', 'does', 'done', 'make', 'made',
         'use', 'used', 'using', 'work', 'works', 'find', 'found', 'show',
         'call', 'called', 'calling', 'get', 'set', 'add', 'all', 'any',
-        'how', 'what', 'when', 'where', 'which', 'who', 'why'
+        'how', 'what', 'when', 'where', 'which', 'who', 'why',
+        'not', 'but', 'are', 'was', 'were', 'has', 'had', 'its',
+        'can', 'did', 'may', 'also', 'into', 'than', 'then', 'them',
+        'each', 'other', 'some', 'such', 'only', 'same', 'about',
+        'after', 'before', 'between', 'through', 'during', 'without',
+        'again', 'further', 'once', 'here', 'there', 'both', 'just',
+        'more', 'most', 'very', 'being', 'having', 'doing',
+        'system', 'need', 'needs', 'want', 'wants', 'like', 'look',
+        'change', 'changes', 'changed', 'changing',
+        // Common English nouns/verbs that match thousands of unrelated code symbols
+        'layer', 'handle', 'handles', 'handling', 'incoming', 'outgoing',
+        'data', 'flow', 'flows', 'level', 'levels', 'request', 'requests',
+        'response', 'responses', 'implement', 'implements', 'implementation',
+        'interface', 'interfaces', 'class', 'classes', 'method', 'methods',
+        'trigger', 'triggers', 'affected', 'affect', 'affects',
+        'else', 'code', 'failing', 'failed', 'silently', 'decide', 'decides',
+        'connect', 'connection', 'connections',
+        'return', 'returns', 'returned', 'take', 'takes', 'taken',
+        'send', 'sends', 'receive', 'receives', 'process', 'processes',
+        'check', 'checks', 'checked', 'create', 'creates', 'created',
+        'read', 'reads', 'write', 'writes', 'written',
+        'start', 'starts', 'stop', 'stops', 'run', 'runs', 'running',
     ]);
     return Array.from(symbols).filter(s => !commonWords.has(s.toLowerCase()));
 }
@@ -154,12 +192,10 @@ class ContextBuilder {
     projectRoot;
     queries;
     traverser;
-    vectorManager;
-    constructor(projectRoot, queries, traverser, vectorManager) {
+    constructor(projectRoot, queries, traverser) {
         this.projectRoot = projectRoot;
         this.queries = queries;
         this.traverser = traverser;
-        this.vectorManager = vectorManager;
     }
     /**
      * Build context for a task
@@ -254,45 +290,136 @@ class ContextBuilder {
         let exactMatches = [];
         if (symbolsFromQuery.length > 0) {
             try {
+                // Get more results so we can apply co-location boosting before trimming
                 exactMatches = this.queries.findNodesByExactName(symbolsFromQuery, {
-                    limit: Math.ceil(opts.searchLimit * 2), // Get more since we'll merge
+                    limit: Math.ceil(opts.searchLimit * 5),
                     kinds: opts.nodeKinds && opts.nodeKinds.length > 0 ? opts.nodeKinds : undefined,
                 });
+                // Co-location boost: when multiple extracted symbols appear in the same file,
+                // those results are much more likely to be what the user is looking for.
+                // E.g., "scrapeLoop" + "run" both in scrape/scrape.go → boost both.
+                if (exactMatches.length > 1) {
+                    // Build a map of files → how many distinct symbol names matched in that file
+                    const fileSymbolCounts = new Map();
+                    for (const r of exactMatches) {
+                        const names = fileSymbolCounts.get(r.node.filePath) || new Set();
+                        names.add(r.node.name.toLowerCase());
+                        fileSymbolCounts.set(r.node.filePath, names);
+                    }
+                    // Boost results in files where multiple query symbols co-occur
+                    exactMatches = exactMatches.map(r => {
+                        const symbolCount = fileSymbolCounts.get(r.node.filePath)?.size || 1;
+                        return {
+                            ...r,
+                            score: symbolCount > 1 ? r.score + (symbolCount - 1) * 20 : r.score,
+                        };
+                    });
+                    exactMatches.sort((a, b) => b.score - a.score);
+                }
+                // Trim back to reasonable size
+                exactMatches = exactMatches.slice(0, Math.ceil(opts.searchLimit * 2));
                 (0, errors_1.logDebug)('Exact symbol matches', { count: exactMatches.length });
             }
             catch (error) {
                 (0, errors_1.logDebug)('Exact symbol lookup failed', { error: String(error) });
             }
         }
-        // Step 3: Try semantic search if vector manager is available
-        let semanticResults = [];
-        if (this.vectorManager && this.vectorManager.isInitialized()) {
-            try {
-                semanticResults = await this.vectorManager.search(query, {
-                    limit: opts.searchLimit,
-                    kinds: opts.nodeKinds && opts.nodeKinds.length > 0 ? opts.nodeKinds : undefined,
-                });
-                (0, errors_1.logDebug)('Semantic search results', { count: semanticResults.length });
-            }
-            catch (error) {
-                (0, errors_1.logDebug)('Semantic search failed, falling back to text search', { query, error: String(error) });
+        // Step 2b: Search for extracted symbols as definition (class/interface) prefixes.
+        // When the user writes "REST", "bulk", or "allocation", they usually mean classes
+        // like RestController, BulkRequest, AllocationService — not nodes named exactly that.
+        // Also tries stem variants: "caching" → "cache" finds Cache, CacheBuilder.
+        if (symbolsFromQuery.length > 0) {
+            const definitionKinds = ['class', 'interface', 'struct', 'trait',
+                'protocol', 'enum', 'type_alias'];
+            // Expand symbols with stem variants for broader definition matching
+            const expandedSymbols = new Set(symbolsFromQuery);
+            for (const sym of symbolsFromQuery) {
+                for (const variant of (0, query_utils_1.getStemVariants)(sym)) {
+                    expandedSymbols.add(variant);
+                }
             }
-        }
-        // Step 4: Fall back to text search if no semantic results
-        if (semanticResults.length === 0 && exactMatches.length === 0) {
-            try {
-                const textResults = this.queries.searchNodes(query, {
-                    limit: opts.searchLimit,
-                    kinds: opts.nodeKinds && opts.nodeKinds.length > 0 ? opts.nodeKinds : undefined,
+            for (const sym of expandedSymbols) {
+                // Title-case the symbol: "REST" → "Rest", "bulk" → "Bulk", "allocation" → "Allocation"
+                const titleCased = sym.charAt(0).toUpperCase() + sym.slice(1).toLowerCase();
+                if (titleCased === sym)
+                    continue; // already title-case (e.g., "Engine") — handled by exact match
+                // Fetch more results since popular prefixes have many matches
+                const prefixResults = this.queries.searchNodes(titleCased, {
+                    limit: 30,
+                    kinds: definitionKinds,
                 });
-                semanticResults = textResults;
+                const matched = [];
+                for (const r of prefixResults) {
+                    if (r.node.name.toLowerCase().startsWith(titleCased.toLowerCase())) {
+                        // Favor shorter names: "AllocationService" (18 chars) over
+                        // "AllocationBalancingRoundMetrics" (31 chars). Core classes tend
+                        // to have concise names; test/helper classes are verbose.
+                        const brevityBonus = Math.max(0, 10 - (r.node.name.length - titleCased.length) / 3);
+                        matched.push({ ...r, score: r.score + 15 + brevityBonus });
+                    }
+                }
+                matched.sort((a, b) => b.score - a.score);
+                for (const r of matched.slice(0, Math.ceil(opts.searchLimit))) {
+                    const existing = exactMatches.find(e => e.node.id === r.node.id);
+                    if (!existing) {
+                        exactMatches.push(r);
+                    }
+                }
             }
-            catch (error) {
-                (0, errors_1.logWarn)('Text search failed', { query, error: String(error) });
-                // Return empty results
+            exactMatches.sort((a, b) => b.score - a.score);
+            exactMatches = exactMatches.slice(0, Math.ceil(opts.searchLimit * 3));
+        }
+        // Step 3: Run text search for natural language term matching
+        // This catches file-name and node-name matches that semantic search may miss,
+        // which is critical for template-heavy codebases (e.g., Liquid/Shopify themes)
+        // where file names are the primary identifiers.
+        let textResults = [];
+        try {
+            const searchTerms = (0, query_utils_1.extractSearchTerms)(query);
+            if (searchTerms.length > 0) {
+                // Search each term individually to get broader coverage,
+                // then boost results that match multiple terms
+                const termResultsMap = new Map();
+                // When no explicit kind filter is set, exclude imports — they flood FTS
+                // results with qualified name matches (e.g., "REST" matches 445K import paths)
+                // but are almost never what exploration queries want.
+                const searchKinds = opts.nodeKinds && opts.nodeKinds.length > 0
+                    ? opts.nodeKinds
+                    : ['file', 'module', 'class', 'struct', 'interface', 'trait', 'protocol',
+                        'function', 'method', 'property', 'field', 'variable', 'constant',
+                        'enum', 'enum_member', 'type_alias', 'namespace', 'export',
+                        'route', 'component'];
+                for (const term of searchTerms) {
+                    const termResults = this.queries.searchNodes(term, {
+                        limit: opts.searchLimit * 2,
+                        kinds: searchKinds,
+                    });
+                    for (const r of termResults) {
+                        const existing = termResultsMap.get(r.node.id);
+                        if (existing) {
+                            existing.termHits++;
+                            existing.result.score = Math.max(existing.result.score, r.score);
+                        }
+                        else {
+                            termResultsMap.set(r.node.id, { result: r, termHits: 1 });
+                        }
+                    }
+                }
+                // Boost results matching multiple terms and sort
+                textResults = Array.from(termResultsMap.values())
+                    .map(({ result, termHits }) => ({
+                    ...result,
+                    score: result.score + (termHits - 1) * 5,
+                }))
+                    .sort((a, b) => b.score - a.score)
+                    .slice(0, opts.searchLimit * 2);
             }
+            (0, errors_1.logDebug)('Text search results', { count: textResults.length });
+        }
+        catch (error) {
+            (0, errors_1.logDebug)('Text search failed', { query, error: String(error) });
         }
-        // Step 5: Merge results, prioritizing exact matches
+        // Step 4: Merge results, prioritizing exact matches, then text (path-boosted)
         const seenIds = new Set();
         let searchResults = [];
         // Add exact matches first (highest priority)
@@ -302,15 +429,217 @@ class ContextBuilder {
                 searchResults.push(result);
             }
         }
-        // Add semantic/text results
-        for (const result of semanticResults) {
+        // Add text search results (includes path relevance scoring from searchNodes)
+        for (const result of textResults) {
             if (!seenIds.has(result.node.id)) {
                 seenIds.add(result.node.id);
                 searchResults.push(result);
             }
         }
-        // Limit total results
-        searchResults = searchResults.slice(0, opts.searchLimit * 2);
+        const queryLower = query.toLowerCase();
+        const isTestQuery = queryLower.includes('test') || queryLower.includes('spec');
+        // Deprioritize test files early so they don't take multi-term boost slots
+        if (!isTestQuery) {
+            for (const result of searchResults) {
+                if ((0, query_utils_1.isTestFile)(result.node.filePath)) {
+                    result.score *= 0.3;
+                }
+            }
+        }
+        // Step 5a: Multi-term co-occurrence re-ranking (applied BEFORE truncation).
+        // For multi-word queries like "search execution from request to shard",
+        // nodes matching 2+ query terms in their name or path are far more relevant
+        // than nodes matching just one generic term. Without this, "ExecutionUtils"
+        // (matches only "execution") fills budget slots meant for "ShardSearchRequest"
+        // (matches "shard" + "search" + "request").
+        const queryTermsForBoost = (0, query_utils_1.extractSearchTerms)(query);
+        if (queryTermsForBoost.length >= 2) {
+            // Group terms that are substrings of each other (stem variants of the same
+            // root word). "indexed", "indexe", "index" should count as ONE concept match,
+            // not three. Without this, stem variants inflate matchCount and give false
+            // multi-term boosts to symbols matching one root word multiple times.
+            const termGroups = [];
+            const sorted = [...queryTermsForBoost].sort((a, b) => b.length - a.length);
+            const assigned = new Set();
+            for (const term of sorted) {
+                if (assigned.has(term))
+                    continue;
+                const group = [term];
+                assigned.add(term);
+                for (const other of sorted) {
+                    if (assigned.has(other))
+                        continue;
+                    if (term.includes(other) || other.includes(term)) {
+                        group.push(other);
+                        assigned.add(other);
+                    }
+                }
+                termGroups.push(group);
+            }
+            for (const result of searchResults) {
+                // Check term matches in name (substring) and path DIRECTORIES (exact).
+                // Directory segments must match exactly — "search" matches directory
+                // "search/" but NOT "elasticsearch/". The class name is checked
+                // separately via substring match on the node name.
+                const nameLower = result.node.name.toLowerCase();
+                const dirSegments = path.dirname(result.node.filePath).toLowerCase().split('/');
+                let matchCount = 0;
+                for (const group of termGroups) {
+                    const groupMatches = group.some(term => {
+                        const inName = nameLower.includes(term);
+                        const inDir = dirSegments.some(seg => seg === term);
+                        return inName || inDir;
+                    });
+                    if (groupMatches)
+                        matchCount++;
+                }
+                if (matchCount >= 2) {
+                    // Multiplicative boost — 2 terms → 2x, 3 terms → 2.5x
+                    result.score *= 1 + matchCount * 0.5;
+                }
+                else {
+                    // Dampen single-term matches — they matched a generic word
+                    // (e.g., "Execution" or "Shard" alone) not the compound concept
+                    result.score *= 0.3;
+                }
+            }
+            searchResults.sort((a, b) => b.score - a.score);
+        }
+        // Step 5b: CamelCase-boundary matching via LIKE query.
+        // FTS can't find "Search" inside "TransportSearchAction" (one FTS token).
+        // LIKE reliably finds these substring matches. Results are appended with
+        // guaranteed slots so they don't compete with higher-scoring prefix matches.
+        if (symbolsFromQuery.length > 0) {
+            const camelDefinitionKinds = ['class', 'interface', 'struct', 'trait',
+                'protocol', 'enum', 'type_alias'];
+            const camelSearchedTerms = new Set();
+            const searchIdSet = new Set(searchResults.map(r => r.node.id));
+            // Track per-node term hits for multi-term boosting
+            const camelNodeTerms = new Map();
+            const maxCamelPerTerm = Math.ceil(opts.searchLimit / 2);
+            for (const sym of symbolsFromQuery) {
+                const titleCased = sym.charAt(0).toUpperCase() + sym.slice(1).toLowerCase();
+                if (titleCased.length < 3)
+                    continue;
+                const termKey = titleCased.toLowerCase();
+                if (camelSearchedTerms.has(termKey))
+                    continue;
+                camelSearchedTerms.add(termKey);
+                // Fetch a large batch — popular terms like "Search" in Elasticsearch
+                // have hundreds of substring matches. The LIKE scan cost is the same
+                // regardless of LIMIT (SQLite scans all matches to sort), so we fetch
+                // generously and let path-relevance scoring pick the best ones.
+                const likeResults = this.queries.findNodesByNameSubstring(titleCased, {
+                    limit: 200,
+                    kinds: camelDefinitionKinds,
+                    excludePrefix: true,
+                });
+                // Filter to CamelCase boundaries, score by path relevance, and take top N
+                const termCandidates = [];
+                for (const r of likeResults) {
+                    const name = r.node.name;
+                    const idx = name.indexOf(titleCased);
+                    if (idx <= 0)
+                        continue;
+                    if (!/[a-z]/.test(name.charAt(idx - 1)))
+                        continue;
+                    if (searchIdSet.has(r.node.id))
+                        continue;
+                    if ((0, query_utils_1.isTestFile)(r.node.filePath) && !isTestQuery)
+                        continue;
+                    const pathScore = (0, query_utils_1.scorePathRelevance)(r.node.filePath, query);
+                    const brevityBonus = Math.max(0, 6 - (name.length - titleCased.length) / 4);
+                    termCandidates.push({ node: r.node, score: 8 + brevityBonus + pathScore });
+                }
+                termCandidates.sort((a, b) => b.score - a.score);
+                // Widen the per-term pool for accumulation so multi-term co-occurrences
+                // can be discovered. A class matching 3 query terms at CamelCase boundaries
+                // is far more relevant than one matching just 1, but it needs to survive
+                // the per-term cut for EACH term to accumulate its count.
+                const accumPerTerm = maxCamelPerTerm * 4;
+                for (const r of termCandidates.slice(0, accumPerTerm)) {
+                    const existing = camelNodeTerms.get(r.node.id);
+                    if (existing) {
+                        existing.termCount++;
+                    }
+                    else {
+                        camelNodeTerms.set(r.node.id, {
+                            result: r,
+                            termCount: 1,
+                        });
+                    }
+                }
+            }
+            // Append CamelCase matches with multi-term boost (guaranteed slots)
+            const camelResults = [];
+            for (const [, info] of camelNodeTerms) {
+                info.result.score += (info.termCount - 1) * 15;
+                camelResults.push(info.result);
+            }
+            camelResults.sort((a, b) => b.score - a.score);
+            const maxCamelTotal = Math.ceil(opts.searchLimit / 2);
+            for (const r of camelResults.slice(0, maxCamelTotal)) {
+                searchResults.push(r);
+                searchIdSet.add(r.node.id);
+            }
+            // Step 5c: Compound term matching — find classes whose name contains 2+
+            // query terms at ANY position (not just CamelCase boundaries).
+            // The CamelCase step above requires idx > 0, which misses classes that
+            // START with a query term (e.g., "SearchShardsRequest" starts with "Search").
+            // For multi-word queries, a class matching multiple query terms in its name
+            // is almost certainly relevant regardless of position.
+            if (symbolsFromQuery.length >= 2) {
+                // Collect ALL LIKE results per term (reusing findNodesByNameSubstring)
+                // but without the CamelCase boundary or prefix exclusion filters.
+                const compoundTermMap = new Map();
+                for (const sym of symbolsFromQuery) {
+                    const titleCased = sym.charAt(0).toUpperCase() + sym.slice(1).toLowerCase();
+                    if (titleCased.length < 3)
+                        continue;
+                    const likeResults = this.queries.findNodesByNameSubstring(titleCased, {
+                        limit: 200,
+                        kinds: camelDefinitionKinds,
+                        excludePrefix: false,
+                    });
+                    for (const r of likeResults) {
+                        if (searchIdSet.has(r.node.id))
+                            continue;
+                        if ((0, query_utils_1.isTestFile)(r.node.filePath) && !isTestQuery)
+                            continue;
+                        const entry = compoundTermMap.get(r.node.id);
+                        if (entry) {
+                            entry.terms.add(titleCased);
+                        }
+                        else {
+                            compoundTermMap.set(r.node.id, { node: r.node, terms: new Set([titleCased]) });
+                        }
+                    }
+                }
+                // Keep only nodes matching 2+ distinct terms
+                const compoundResults = [];
+                for (const [, entry] of compoundTermMap) {
+                    if (entry.terms.size >= 2) {
+                        const pathScore = (0, query_utils_1.scorePathRelevance)(entry.node.filePath, query);
+                        const brevityBonus = Math.max(0, 6 - entry.node.name.length / 8);
+                        compoundResults.push({
+                            node: entry.node,
+                            score: 10 + (entry.terms.size - 1) * 20 + pathScore + brevityBonus,
+                        });
+                    }
+                }
+                compoundResults.sort((a, b) => b.score - a.score);
+                const maxCompound = Math.ceil(opts.searchLimit / 2);
+                for (const r of compoundResults.slice(0, maxCompound)) {
+                    searchResults.push(r);
+                    searchIdSet.add(r.node.id);
+                }
+            }
+        }
+        // Final sort and truncation — all search channels (exact, text, CamelCase,
+        // compound) have now contributed. Sort by score so multi-term matches from
+        // later steps can outrank dampened single-term matches from earlier steps.
+        searchResults.sort((a, b) => b.score - a.score);
+        searchResults = searchResults.slice(0, opts.searchLimit * 3);
         // Filter by minimum score
         let filteredResults = searchResults.filter((r) => r.score >= opts.minScore);
         // Resolve imports/exports to their actual definitions
@@ -322,6 +651,57 @@ class ContextBuilder {
             nodes.set(result.node.id, result.node);
             roots.push(result.node.id);
         }
+        // Expand type hierarchy for class/interface entry points.
+        // BFS often exhausts its per-entry-point budget on contained methods
+        // before reaching extends/implements neighbors. This dedicated step
+        // ensures subclasses and superclasses always appear in results.
+        // Budget: up to maxNodes/4 hierarchy nodes to avoid flooding.
+        const typeHierarchyKinds = new Set(['class', 'interface', 'struct', 'trait', 'protocol']);
+        const maxHierarchyNodes = Math.ceil(opts.maxNodes / 4);
+        let hierarchyNodesAdded = 0;
+        for (const result of filteredResults) {
+            if (hierarchyNodesAdded >= maxHierarchyNodes)
+                break;
+            if (typeHierarchyKinds.has(result.node.kind)) {
+                const hierarchy = this.traverser.getTypeHierarchy(result.node.id);
+                for (const [id, node] of hierarchy.nodes) {
+                    if (!nodes.has(id)) {
+                        nodes.set(id, node);
+                        hierarchyNodesAdded++;
+                    }
+                }
+                for (const edge of hierarchy.edges) {
+                    const exists = edges.some((e) => e.source === edge.source && e.target === edge.target && e.kind === edge.kind);
+                    if (!exists) {
+                        edges.push(edge);
+                    }
+                }
+            }
+        }
+        // Pass 2: expand hierarchy of newly-discovered parent types to find siblings.
+        // E.g., InternalEngine → Engine (parent, from pass 1) → ReadOnlyEngine (sibling).
+        if (hierarchyNodesAdded > 0) {
+            const pass2Candidates = [...nodes.values()].filter(n => typeHierarchyKinds.has(n.kind) && !roots.includes(n.id));
+            for (const candidate of pass2Candidates) {
+                if (hierarchyNodesAdded >= maxHierarchyNodes)
+                    break;
+                const siblingHierarchy = this.traverser.getTypeHierarchy(candidate.id);
+                for (const [id, node] of siblingHierarchy.nodes) {
+                    if (!nodes.has(id) && hierarchyNodesAdded < maxHierarchyNodes) {
+                        nodes.set(id, node);
+                        hierarchyNodesAdded++;
+                    }
+                }
+                for (const edge of siblingHierarchy.edges) {
+                    if (nodes.has(edge.source) && nodes.has(edge.target)) {
+                        const exists = edges.some((e) => e.source === edge.source && e.target === edge.target && e.kind === edge.kind);
+                        if (!exists) {
+                            edges.push(edge);
+                        }
+                    }
+                }
+            }
+        }
         // Traverse from each entry point
         for (const result of filteredResults) {
             const traversalResult = this.traverser.traverseBFS(result.node.id, {
@@ -346,6 +726,8 @@ class ContextBuilder {
             }
         }
         // Trim to max nodes if needed
+        let finalNodes = nodes;
+        let finalEdges = edges;
         if (nodes.size > opts.maxNodes) {
             // Prioritize entry points and their direct neighbors
             const priorityIds = new Set(roots);
@@ -358,26 +740,95 @@ class ContextBuilder {
                 }
             }
             // Keep priority nodes, then fill remaining slots
-            const trimmedNodes = new Map();
+            finalNodes = new Map();
             for (const id of priorityIds) {
                 const node = nodes.get(id);
-                if (node && trimmedNodes.size < opts.maxNodes) {
-                    trimmedNodes.set(id, node);
+                if (node && finalNodes.size < opts.maxNodes) {
+                    finalNodes.set(id, node);
                 }
             }
             // Fill remaining from other nodes
             for (const [id, node] of nodes) {
-                if (trimmedNodes.size >= opts.maxNodes)
+                if (finalNodes.size >= opts.maxNodes)
                     break;
-                if (!trimmedNodes.has(id)) {
-                    trimmedNodes.set(id, node);
+                if (!finalNodes.has(id)) {
+                    finalNodes.set(id, node);
                 }
             }
             // Filter edges to only include kept nodes
-            const trimmedEdges = edges.filter((e) => trimmedNodes.has(e.source) && trimmedNodes.has(e.target));
-            return { nodes: trimmedNodes, edges: trimmedEdges, roots };
+            finalEdges = edges.filter((e) => finalNodes.has(e.source) && finalNodes.has(e.target));
+        }
+        // Per-file diversity cap: prevent any single file from monopolizing the
+        // node budget. When BFS traverses from a method, it follows `contains`
+        // to the parent class, then back down to all sibling methods. With
+        // multiple entry points in the same class, one file can consume 30-40%
+        // of maxNodes. Cap each file to ~20% to ensure cross-file diversity.
+        const maxPerFile = Math.max(5, Math.ceil(opts.maxNodes * 0.2));
+        const fileCounts = new Map();
+        for (const [id, node] of finalNodes) {
+            const ids = fileCounts.get(node.filePath) || [];
+            ids.push(id);
+            fileCounts.set(node.filePath, ids);
+        }
+        const rootSet = new Set(roots);
+        for (const [, nodeIds] of fileCounts) {
+            if (nodeIds.length <= maxPerFile)
+                continue;
+            // Sort: entry points first, then classes/interfaces, then others
+            const kindPriority = {
+                class: 3, interface: 3, struct: 3, trait: 3, protocol: 3, enum: 3,
+                method: 1, function: 1, property: 0, field: 0, variable: 0,
+            };
+            nodeIds.sort((a, b) => {
+                const aRoot = rootSet.has(a) ? 10 : 0;
+                const bRoot = rootSet.has(b) ? 10 : 0;
+                const aKind = kindPriority[finalNodes.get(a).kind] ?? 0;
+                const bKind = kindPriority[finalNodes.get(b).kind] ?? 0;
+                return (bRoot + bKind) - (aRoot + aKind);
+            });
+            // Remove excess nodes (keep the highest-priority ones)
+            for (const id of nodeIds.slice(maxPerFile)) {
+                finalNodes.delete(id);
+            }
+        }
+        // Non-production node cap: limit test/sample/integration/example files to
+        // at most 15% of the budget. Many codebases have dozens of near-identical
+        // test implementations (e.g., 6 Guard classes in integration tests) that
+        // individually survive score dampening but collectively flood the result.
+        // Test entry points are NOT exempt — they should be evicted too.
+        if (!isTestQuery) {
+            const maxNonProd = Math.max(3, Math.ceil(opts.maxNodes * 0.15));
+            const nonProdIds = [];
+            for (const [id, node] of finalNodes) {
+                if ((0, query_utils_1.isTestFile)(node.filePath)) {
+                    nonProdIds.push(id);
+                }
+            }
+            if (nonProdIds.length > maxNonProd) {
+                for (const id of nonProdIds.slice(maxNonProd)) {
+                    finalNodes.delete(id);
+                    // Also remove from roots — test file entry points shouldn't anchor results
+                    const rootIdx = roots.indexOf(id);
+                    if (rootIdx !== -1)
+                        roots.splice(rootIdx, 1);
+                }
+            }
+        }
+        // Re-filter edges after per-file and non-production caps
+        finalEdges = finalEdges.filter((e) => finalNodes.has(e.source) && finalNodes.has(e.target));
+        // Edge recovery: BFS with many entry points leaves most nodes disconnected.
+        // Discover edges between already-selected nodes to recover connectivity.
+        const recoveryKinds = ['calls', 'extends', 'implements', 'references', 'overrides'];
+        const recoveredEdges = this.queries.findEdgesBetweenNodes([...finalNodes.keys()], recoveryKinds);
+        const existingEdgeKeys = new Set(finalEdges.map((e) => `${e.source}:${e.target}:${e.kind}`));
+        for (const edge of recoveredEdges) {
+            const key = `${edge.source}:${edge.target}:${edge.kind}`;
+            if (!existingEdgeKeys.has(key)) {
+                finalEdges.push(edge);
+                existingEdgeKeys.add(key);
+            }
         }
-        return { nodes, edges, roots };
+        return { nodes: finalNodes, edges: finalEdges, roots };
     }
     /**
      * Get the source code for a node
@@ -559,8 +1010,8 @@ exports.ContextBuilder = ContextBuilder;
 /**
  * Create a context builder
  */
-function createContextBuilder(projectRoot, queries, traverser, vectorManager) {
-    return new ContextBuilder(projectRoot, queries, traverser, vectorManager);
+function createContextBuilder(projectRoot, queries, traverser) {
+    return new ContextBuilder(projectRoot, queries, traverser);
 }
 // Re-export formatter
 var formatter_2 = require("./formatter");