diff --git a/packages/memory-graph/src/hooks/use-force-simulation.ts b/packages/memory-graph/src/hooks/use-force-simulation.ts index e1d6b19e..3b5bfe97 100644 --- a/packages/memory-graph/src/hooks/use-force-simulation.ts +++ b/packages/memory-graph/src/hooks/use-force-simulation.ts @@ -114,16 +114,17 @@ export function useForceSimulation( // eslint-disable-next-line react-hooks/exhaustive-deps }, [enabled]) - // Update simulation nodes when they change + // Update simulation nodes and edges together to prevent race conditions useEffect(() => { - if (simulationRef.current && nodes.length > 0) { + if (!simulationRef.current) return + + // Update nodes + if (nodes.length > 0) { simulationRef.current.nodes(nodes) } - }, [nodes]) - // Update simulation edges when they change - useEffect(() => { - if (simulationRef.current && edges.length > 0) { + // Update edges + if (edges.length > 0) { const linkForce = simulationRef.current.force< d3.ForceLink >("link") @@ -131,7 +132,7 @@ export function useForceSimulation( linkForce.links(edges) } } - }, [edges]) + }, [nodes, edges]) // Reheat simulation (called on drag start) const reheat = useCallback(() => { diff --git a/packages/memory-graph/src/hooks/use-graph-data.ts b/packages/memory-graph/src/hooks/use-graph-data.ts index d8c7e5b5..f0afab1d 100644 --- a/packages/memory-graph/src/hooks/use-graph-data.ts +++ b/packages/memory-graph/src/hooks/use-graph-data.ts @@ -5,7 +5,7 @@ import { getConnectionVisualProps, getMagicalConnectionColor, } from "@/lib/similarity" -import { useMemo, useRef } from "react" +import { useMemo, useRef, useEffect } from "react" import { colors, LAYOUT_CONSTANTS } from "@/constants" import type { DocumentsResponse, @@ -26,6 +26,27 @@ export function useGraphData( // Cache nodes to preserve d3-force mutations (x, y, vx, vy, fx, fy) const nodeCache = useRef>(new Map()) + // Cleanup nodeCache to prevent memory leak + useEffect(() => { + if (!data?.documents) return + + // Build set of current node IDs + const currentNodeIds = new Set() + data.documents.forEach((doc) => { + currentNodeIds.add(doc.id) + doc.memoryEntries.forEach((mem) => { + currentNodeIds.add(`${mem.id}`) + }) + }) + + // Remove stale nodes from cache + for (const [id] of nodeCache.current.entries()) { + if (!currentNodeIds.has(id)) { + nodeCache.current.delete(id) + } + } + }, [data, selectedSpace]) + return useMemo(() => { if (!data?.documents) return { nodes: [], edges: [] } @@ -296,12 +317,16 @@ export function useGraphData( }) // Document-to-document similarity edges - for (let i = 0; i < filteredDocuments.length; i++) { - const docI = filteredDocuments[i] + // Performance optimization: limit comparisons to prevent O(n²) scaling issues + const MAX_DOCS_FOR_SIMILARITY = 50 + const docsToCompare = filteredDocuments.slice(0, MAX_DOCS_FOR_SIMILARITY) + + for (let i = 0; i < docsToCompare.length; i++) { + const docI = docsToCompare[i] if (!docI) continue - for (let j = i + 1; j < filteredDocuments.length; j++) { - const docJ = filteredDocuments[j] + for (let j = i + 1; j < docsToCompare.length; j++) { + const docJ = docsToCompare[j] if (!docJ) continue const sim = calculateSemanticSimilarity(