From fa4ab9805ae96ebee1e29b0af22aab5254df7c0d Mon Sep 17 00:00:00 2001 From: shamAnimates <145093437+shamAnimates@users.noreply.github.com> Date: Mon, 12 Oct 2026 21:44:25 -0700 Subject: [PATCH] fix(memory-graph): deduplicate shared graph entities (#1525) --- .../__tests__/shared-memory-identity.test.tsx | 256 ++++++++++++++++++ .../src/components/memory-graph.tsx | 46 ++-- .../memory-graph/src/hooks/use-graph-data.ts | 21 +- 3 files changed, 294 insertions(+), 29 deletions(-) create mode 100644 packages/memory-graph/src/__tests__/shared-memory-identity.test.tsx diff --git a/packages/memory-graph/src/__tests__/shared-memory-identity.test.tsx b/packages/memory-graph/src/__tests__/shared-memory-identity.test.tsx new file mode 100644 index 00000000..4ba5c728 --- /dev/null +++ b/packages/memory-graph/src/__tests__/shared-memory-identity.test.tsx @@ -0,0 +1,256 @@ +import { cleanup, render, renderHook, waitFor } from "@testing-library/react" +import { afterEach, describe, expect, it, vi } from "vitest" +import { MemoryGraph } from "../components/memory-graph" +import { DEFAULT_COLORS } from "../constants" +import { computeEdges, useGraphData } from "../hooks/use-graph-data" +import type { + GraphApiDocument, + GraphApiMemory, + DocumentNodeData, + GraphEdge, + GraphNode, + MemoryNodeData, +} from "../types" + +const graphCanvasCapture = vi.hoisted(() => ({ + nodes: [] as GraphNode[], + edges: [] as GraphEdge[], +})) + +vi.mock("../components/graph-canvas", () => ({ + GraphCanvas: ({ + nodes, + edges, + }: { + nodes: GraphNode[] + edges: GraphEdge[] + }) => { + graphCanvasCapture.nodes = nodes + graphCanvasCapture.edges = edges + return null + }, +})) + +afterEach(() => { + cleanup() + graphCanvasCapture.nodes = [] + graphCanvasCapture.edges = [] + vi.restoreAllMocks() +}) + +function makeMemory(overrides: Partial = {}): GraphApiMemory { + return { + id: "memory", + memory: "A remembered fact", + isStatic: false, + spaceId: "space", + isLatest: true, + isForgotten: false, + forgetAfter: null, + forgetReason: null, + version: 1, + parentMemoryId: null, + rootMemoryId: null, + createdAt: "2026-01-01T00:00:00Z", + updatedAt: "2026-01-01T00:00:00Z", + memoryRelations: null, + ...overrides, + } +} + +function makeDocument( + id: string, + memories: GraphApiMemory[], +): GraphApiDocument { + return { + id, + title: id, + summary: null, + documentType: "text", + createdAt: "2026-01-01T00:00:00Z", + updatedAt: "2026-01-01T00:00:00Z", + memories, + } +} + +function getMemoryNode( + nodes: GraphNode[], + id: string, +): GraphNode & { data: MemoryNodeData } { + const node = nodes.find((candidate) => candidate.id === id) + expect(node?.type).toBe("memory") + if (!node || !("documentId" in node.data)) { + throw new Error(`Missing memory ${id}`) + } + return node as GraphNode & { data: MemoryNodeData } +} + +function getDocumentNode( + nodes: GraphNode[], + id: string, +): GraphNode & { data: DocumentNodeData } { + const node = nodes.find((candidate) => candidate.id === id) + expect(node?.type).toBe("document") + if (!node || !("memories" in node.data)) { + throw new Error(`Missing document ${id}`) + } + return node as GraphNode & { data: DocumentNodeData } +} + +describe("shared memory identity", () => { + it("materializes the first source occurrence once while retaining every document edge", () => { + const documents = [ + makeDocument("doc-a", [ + makeMemory({ id: "shared-memory", memory: "first source content" }), + ]), + makeDocument("doc-b", [ + makeMemory({ id: "shared-memory", memory: "later source content" }), + ]), + ] + const { result, unmount } = renderHook(() => + useGraphData(documents, null, 800, 600, DEFAULT_COLORS), + ) + + const nodeIds = result.current.nodes.map((node) => node.id) + const sharedNode = getMemoryNode(result.current.nodes, "shared-memory") + const derivesEdges = result.current.edges.filter( + (edge) => edge.edgeType === "document", + ) + + expect(nodeIds).toHaveLength(3) + expect(new Set(nodeIds).size).toBe(nodeIds.length) + expect(sharedNode.data.content).toBe("first source content") + expect(sharedNode.data.documentId).toBe("doc-a") + expect(derivesEdges).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + id: "dm-doc-a-shared-memory", + source: "doc-a", + target: "shared-memory", + }), + expect.objectContaining({ + id: "dm-doc-b-shared-memory", + source: "doc-b", + target: "shared-memory", + }), + ]), + ) + expect(derivesEdges).toHaveLength(2) + + unmount() + }) + + it("preserves cached position when another source occurrence is appended", () => { + const firstDocuments = [ + makeDocument("doc-a", [makeMemory({ id: "shared-memory" })]), + ] + const { result, rerender, unmount } = renderHook( + ({ documents }: { documents: GraphApiDocument[] }) => + useGraphData(documents, null, 800, 600, DEFAULT_COLORS), + { initialProps: { documents: firstDocuments } }, + ) + const firstNode = getMemoryNode(result.current.nodes, "shared-memory") + + rerender({ + documents: [ + ...firstDocuments, + makeDocument("doc-b", [makeMemory({ id: "shared-memory" })]), + ], + }) + + const rerenderedNode = getMemoryNode(result.current.nodes, "shared-memory") + expect({ x: rerenderedNode.x, y: rerenderedNode.y }).toEqual({ + x: firstNode.x, + y: firstNode.y, + }) + expect(rerenderedNode.data.documentId).toBe("doc-a") + + unmount() + }) + + it("unions relation pairs but keeps direction, type, and first occurrence", () => { + const documents = [ + makeDocument("doc-a", [ + makeMemory({ + id: "shared-memory", + memoryRelations: { "target-a": "updates" }, + }), + makeMemory({ id: "target-a" }), + ]), + makeDocument("doc-b", [ + makeMemory({ + id: "shared-memory", + memoryRelations: { + "target-a": "extends", + "target-b": "extends", + }, + }), + makeMemory({ id: "target-b" }), + ]), + ] + + const relationEdges = computeEdges(documents).filter((edge) => + edge.id.startsWith("rel-"), + ) + + expect(relationEdges).toHaveLength(2) + expect(relationEdges).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + id: "rel-target-a-shared-memory", + source: "target-a", + target: "shared-memory", + edgeType: "updates", + }), + expect.objectContaining({ + id: "rel-target-b-shared-memory", + source: "target-b", + target: "shared-memory", + edgeType: "extends", + }), + ]), + ) + }) +}) + +describe("MemoryGraph maxNodes", () => { + it("counts unique memories and retains a later duplicate after saturation", async () => { + vi.spyOn(HTMLElement.prototype, "getBoundingClientRect").mockReturnValue({ + bottom: 600, + height: 600, + left: 0, + right: 800, + top: 0, + width: 800, + x: 0, + y: 0, + toJSON: () => ({}), + }) + const documents = [ + makeDocument("doc-a", [ + makeMemory({ id: "shared-memory" }), + makeMemory({ id: "unique-a" }), + ]), + makeDocument("doc-b", [ + makeMemory({ id: "over-budget" }), + makeMemory({ id: "shared-memory" }), + ]), + ] + + render() + + await waitFor(() => { + expect(graphCanvasCapture.nodes).toHaveLength(4) + }) + expect(graphCanvasCapture.nodes.map((node) => node.id).sort()).toEqual([ + "doc-a", + "doc-b", + "shared-memory", + "unique-a", + ]) + const docB = getDocumentNode(graphCanvasCapture.nodes, "doc-b") + expect(docB.data.memories.map((memory) => memory.id)).toEqual([ + "shared-memory", + ]) + }) +}) diff --git a/packages/memory-graph/src/components/memory-graph.tsx b/packages/memory-graph/src/components/memory-graph.tsx index b0ccc668..218f0bad 100644 --- a/packages/memory-graph/src/components/memory-graph.tsx +++ b/packages/memory-graph/src/components/memory-graph.tsx @@ -114,37 +114,37 @@ export function MemoryGraph({ // Used as a dependency proxy to recalculate popover positions const [viewportVersion, setViewportVersion] = useState(0) - // Limit documents so total node count (documents + their memories) stays under maxNodes + // Cap documents + unique memories at maxNodes; shared memories count once. const limitedDocuments = useMemo(() => { if (!maxNodes || documents.length === 0) return documents let totalNodes = 0 + const seenMemoryIds = new Set() const limited: GraphApiDocument[] = [] - for (let i = 0; i < documents.length; i++) { - const doc = documents[i] - if (!doc) continue + for (const doc of documents) { if (totalNodes >= maxNodes) break - const remainingNodes = maxNodes - totalNodes const memories = doc.memories ?? [] - const docNodes = 1 + memories.length - - if (docNodes <= remainingNodes) { - limited.push(doc) - totalNodes += docNodes - continue - } - - if (remainingNodes > 1) { - limited.push({ - ...doc, - memories: memories.slice(0, remainingNodes - 1), - }) - totalNodes = maxNodes - break - } - - limited.push({ ...doc, memories: [] }) + const includedMemories: GraphApiDocument["memories"] = [] totalNodes += 1 + + for (const memory of memories) { + if (seenMemoryIds.has(memory.id)) { + includedMemories.push(memory) + continue + } + + if (totalNodes >= maxNodes) continue + + seenMemoryIds.add(memory.id) + includedMemories.push(memory) + totalNodes += 1 + } + + limited.push( + includedMemories.length === memories.length + ? doc + : { ...doc, memories: includedMemories }, + ) } return limited }, [documents, maxNodes]) diff --git a/packages/memory-graph/src/hooks/use-graph-data.ts b/packages/memory-graph/src/hooks/use-graph-data.ts index deb59b47..74a291f7 100644 --- a/packages/memory-graph/src/hooks/use-graph-data.ts +++ b/packages/memory-graph/src/hooks/use-graph-data.ts @@ -402,10 +402,8 @@ function getAppendSpatialKey(x: number, y: number): string { return `${x}:${y}` } -/** - * Pure function that computes graph edges from documents. - * Extracted from the hook for testability. - */ +// Pure so it can be tested outside the hook. + export function computeEdges(documents: GraphApiDocument[]): GraphEdge[] { if (!documents || documents.length === 0) return [] @@ -416,10 +414,14 @@ export function computeEdges(documents: GraphApiDocument[]): GraphEdge[] { for (const mem of doc.memories) allNodeIds.add(mem.id) } + const documentEdgeIds = new Set() for (const doc of documents) { for (const mem of doc.memories) { + const edgeId = `dm-${doc.id}-${mem.id}` + if (documentEdgeIds.has(edgeId)) continue + documentEdgeIds.add(edgeId) result.push({ - id: `dm-${doc.id}-${mem.id}`, + id: edgeId, source: doc.id, target: mem.id, visualProps: getEdgeVisualProps("document"), @@ -431,12 +433,16 @@ export function computeEdges(documents: GraphApiDocument[]): GraphEdge[] { // 2. Memory-to-memory relation edges from backend data. // Uses memoryRelations (Record) as primary source, // falls back to parentMemoryId for legacy data. + const relationEdgeIds = new Set() for (const doc of documents) { for (const mem of doc.memories) { const relations = getMemoryRelationTargets(mem) for (const [targetId, relationType] of Object.entries(relations)) { if (!allNodeIds.has(targetId)) continue + const edgeId = `rel-${targetId}-${mem.id}` + if (relationEdgeIds.has(edgeId)) continue + relationEdgeIds.add(edgeId) const edgeType = relationType === "updates" || relationType === "extends" || @@ -444,7 +450,7 @@ export function computeEdges(documents: GraphApiDocument[]): GraphEdge[] { ? relationType : "updates" result.push({ - id: `rel-${targetId}-${mem.id}`, + id: edgeId, source: targetId, target: mem.id, visualProps: getEdgeVisualProps(edgeType), @@ -498,6 +504,7 @@ export function useGraphData( const clusterAssignments = computeClusterAssignments(documents) const result: GraphNode[] = [] + const seenMemoryIds = new Set() // Spiral layout: documents form a compact spiral core, memories orbit // around their parent documents. The force simulation then gently // pushes memories outward to create the constellation/starburst effect. @@ -578,6 +585,8 @@ export function useGraphData( for (let i = 0; i < memCount; i++) { const mem = doc.memories[i] if (!mem) continue + if (seenMemoryIds.has(mem.id)) continue + seenMemoryIds.add(mem.id) const previousMemNode = previousCache.get(mem.id) const memData: MemoryNodeData = { ...mem,