feat: Add edge recovery to restore connectivity after node trimming in context building
Addresses cases where BFS with multiple entry points leaves most nodes disconnected after trimming. Discovers edges between already-selected nodes using specific relationship types (calls, extends, implements, references, overrides) to recover inter-node connectivity that would otherwise be lost during the node selection process.
This commit is contained in:
+29
-11
@@ -766,6 +766,8 @@ export class ContextBuilder {
|
|||||||
}
|
}
|
||||||
|
|
||||||
// Trim to max nodes if needed
|
// Trim to max nodes if needed
|
||||||
|
let finalNodes = nodes;
|
||||||
|
let finalEdges = edges;
|
||||||
if (nodes.size > opts.maxNodes) {
|
if (nodes.size > opts.maxNodes) {
|
||||||
// Prioritize entry points and their direct neighbors
|
// Prioritize entry points and their direct neighbors
|
||||||
const priorityIds = new Set(roots);
|
const priorityIds = new Set(roots);
|
||||||
@@ -779,31 +781,47 @@ export class ContextBuilder {
|
|||||||
}
|
}
|
||||||
|
|
||||||
// Keep priority nodes, then fill remaining slots
|
// Keep priority nodes, then fill remaining slots
|
||||||
const trimmedNodes = new Map<string, Node>();
|
finalNodes = new Map<string, Node>();
|
||||||
for (const id of priorityIds) {
|
for (const id of priorityIds) {
|
||||||
const node = nodes.get(id);
|
const node = nodes.get(id);
|
||||||
if (node && trimmedNodes.size < opts.maxNodes) {
|
if (node && finalNodes.size < opts.maxNodes) {
|
||||||
trimmedNodes.set(id, node);
|
finalNodes.set(id, node);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Fill remaining from other nodes
|
// Fill remaining from other nodes
|
||||||
for (const [id, node] of nodes) {
|
for (const [id, node] of nodes) {
|
||||||
if (trimmedNodes.size >= opts.maxNodes) break;
|
if (finalNodes.size >= opts.maxNodes) break;
|
||||||
if (!trimmedNodes.has(id)) {
|
if (!finalNodes.has(id)) {
|
||||||
trimmedNodes.set(id, node);
|
finalNodes.set(id, node);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Filter edges to only include kept nodes
|
// Filter edges to only include kept nodes
|
||||||
const trimmedEdges = edges.filter(
|
finalEdges = edges.filter(
|
||||||
(e) => trimmedNodes.has(e.source) && trimmedNodes.has(e.target)
|
(e) => finalNodes.has(e.source) && finalNodes.has(e.target)
|
||||||
);
|
);
|
||||||
|
|
||||||
return { nodes: trimmedNodes, edges: trimmedEdges, roots };
|
|
||||||
}
|
}
|
||||||
|
|
||||||
return { nodes, edges, roots };
|
// Edge recovery: BFS with many entry points leaves most nodes disconnected.
|
||||||
|
// Discover edges between already-selected nodes to recover connectivity.
|
||||||
|
const recoveryKinds: EdgeKind[] = ['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: finalNodes, edges: finalEdges, roots };
|
||||||
}
|
}
|
||||||
|
|
||||||
/**
|
/**
|
||||||
|
|||||||
@@ -859,6 +859,26 @@ export class QueryBuilder {
|
|||||||
return rows.map(rowToEdge);
|
return rows.map(rowToEdge);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Find all edges where both source and target are in the given node set.
|
||||||
|
* Useful for recovering inter-node connectivity after BFS.
|
||||||
|
*/
|
||||||
|
findEdgesBetweenNodes(nodeIds: string[], kinds?: EdgeKind[]): Edge[] {
|
||||||
|
if (nodeIds.length === 0) return [];
|
||||||
|
|
||||||
|
const idsJson = JSON.stringify(nodeIds);
|
||||||
|
let sql = `SELECT * FROM edges WHERE source IN (SELECT value FROM json_each(?)) AND target IN (SELECT value FROM json_each(?))`;
|
||||||
|
const params: string[] = [idsJson, idsJson];
|
||||||
|
|
||||||
|
if (kinds && kinds.length > 0) {
|
||||||
|
sql += ` AND kind IN (${kinds.map(() => '?').join(',')})`;
|
||||||
|
params.push(...kinds);
|
||||||
|
}
|
||||||
|
|
||||||
|
const rows = this.db.prepare(sql).all(...params) as EdgeRow[];
|
||||||
|
return rows.map(rowToEdge);
|
||||||
|
}
|
||||||
|
|
||||||
// ===========================================================================
|
// ===========================================================================
|
||||||
// File Operations
|
// File Operations
|
||||||
// ===========================================================================
|
// ===========================================================================
|
||||||
|
|||||||
Reference in New Issue
Block a user