From d9e973cffc65432dcaa0ada64199543dcc5b8a94 Mon Sep 17 00:00:00 2001
From: Colby McHenry <me@colbymchenry.com>
Date: Mon, 6 Apr 2026 16:08:42 -0500
Subject: [PATCH] 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.
---
 src/context/index.ts | 40 +++++++++++++++++++++++++++++-----------
 src/db/queries.ts    | 20 ++++++++++++++++++++
 2 files changed, 49 insertions(+), 11 deletions(-)

diff --git a/src/context/index.ts b/src/context/index.ts
index 87188bc..b2af1c4 100644
--- a/src/context/index.ts
+++ b/src/context/index.ts
@@ -766,6 +766,8 @@ export 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);
@@ -779,31 +781,47 @@ export class ContextBuilder {
       }
 
       // Keep priority nodes, then fill remaining slots
-      const trimmedNodes = new Map<string, Node>();
+      finalNodes = new Map<string, Node>();
       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) break;
-        if (!trimmedNodes.has(id)) {
-          trimmedNodes.set(id, node);
+        if (finalNodes.size >= opts.maxNodes) break;
+        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)
+      finalEdges = edges.filter(
+        (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 };
   }
 
   /**
diff --git a/src/db/queries.ts b/src/db/queries.ts
index 874d7a2..e66a1e5 100644
--- a/src/db/queries.ts
+++ b/src/db/queries.ts
@@ -859,6 +859,26 @@ export class QueryBuilder {
     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
   // ===========================================================================