Merge pull request #368 from HelpCode-ai/fix/kg-observational-oom-2

fix(kg): bound observational ingest memory + CPU (per-page flush, walk cap)

Author: keysersoft

packages/backend/src/knowledge-graph/identifier.ts

@@ -46,15 +46,23 @@ export interface IdentifierOccurrence {
   value: string;
 }
 
+/** Max JSON nodes to walk per payload (bounds CPU on huge bulk responses). */
+const MAX_WALK_NODES = 5000;
+
 /**
  * Walk an arbitrary JSON value and yield (leaf-field, identifier-value) pairs.
  * Uses the leaf key as the field name (FK fields are leaf keys like customer_id).
  */
 export function extractIdentifiers(input: unknown): IdentifierOccurrence[] {
   const out: IdentifierOccurrence[] = [];
   const seen = new Set<string>();
+  // Bound work on pathologically large payloads (e.g. a bulk DB dump): a few
+  // thousand nodes is plenty to capture representative identifiers, and stops a
+  // single huge response from pinning the event loop.
+  let visited = 0;
 
   const walk = (node: unknown, key: string | null): void => {
+    if (visited++ > MAX_WALK_NODES) return;
     if (node === null || node === undefined) return;
     if (Array.isArray(node)) {
       for (const item of node) walk(item, key);
@@ -89,7 +97,9 @@ export function extractIdentifiers(input: unknown): IdentifierOccurrence[] {
  */
 export function extractFieldNames(input: unknown): string[] {
   const out = new Set<string>();
+  let visited = 0;
   const walk = (node: unknown): void => {
+    if (visited++ > MAX_WALK_NODES) return;
     if (node === null || node === undefined) return;
     if (Array.isArray(node)) {
       for (const item of node) walk(item);

packages/backend/src/knowledge-graph/kg-observational.service.ts

@@ -81,15 +81,9 @@ export class KgObservationalService {
       s.add(r.entity);
     }
 
-    const newHashes = new Set<string>();
-    const valueRows: Array<{
-      organizationId: string;
-      connectorId: string;
-      valueHash: string;
-      entity: string;
-      field: string;
-      direction: string;
-    }> = [];
+    // Edges accumulate across pages; the raw value rows are flushed per page
+    // (below) so the in-memory set never grows to the size of the whole batch.
+    let edges = 0;
     // references edges mined from response field names: key -> details.
     const refBumps = new Map<
       string,
@@ -103,7 +97,7 @@ export class KgObservationalService {
     // Page through invocations so we never hold more than CHUNK full input/output
     // payloads in memory at once. A busy org's payloads can be megabytes each;
     // loading thousands together previously OOM'd the backend.
-    const CHUNK = 200;
+    const CHUNK = 100;
     let processed = 0;
     while (processed < MAX_INVOCATIONS_PER_RUN) {
       const page = await this.prisma.toolInvocation.findMany({
@@ -124,6 +118,17 @@ export class KgObservationalService {
       if (page.length === 0) break;
       processed += page.length;
 
+      // Per-page value occurrences, flushed at the end of each page.
+      const newHashes = new Set<string>();
+      const valueRows: Array<{
+        organizationId: string;
+        connectorId: string;
+        valueHash: string;
+        entity: string;
+        field: string;
+        direction: string;
+      }> = [];
+
       for (const inv of page) {
       const connectorId = inv.connectorId!;
       // Skip invocations already covered by this connector's watermark.
@@ -185,16 +190,19 @@ export class KgObservationalService {
         }
       }
       } // for (const inv of page)
-      if (page.length < CHUNK) break;
-    } // while pages
 
-    if (valueRows.length) {
-      await this.prisma.kgValueSeen.createMany({ data: valueRows });
-    }
+      // Flush this page's value occurrences and correlate immediately. correlate
+      // reads kgValueSeen (which now includes earlier pages), so cross-page
+      // produces_consumes / same_identity links are still found.
+      if (valueRows.length) {
+        await this.prisma.kgValueSeen.createMany({ data: valueRows });
+      }
+      if (newHashes.size) {
+        edges += await this.correlate(organizationId, [...newHashes]);
+      }
 
-    let edges = newHashes.size
-      ? await this.correlate(organizationId, [...newHashes])
-      : 0;
+      if (page.length < CHUNK) break;
+    } // while pages
 
     // Apply references edges mined from response shapes.
     const refCache = new Map<string, string>();