fix(memory-graph): deduplicate shared graph entities (#1525)

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> = {}): 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(<MemoryGraph documents={documents} maxNodes={4} />)
+
+		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",
+		])
+	})
+})

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<string>()
 		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])

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<string>()
 	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<targetId, relationType>) as primary source,
 	//    falls back to parentMemoryId for legacy data.
+	const relationEdgeIds = new Set<string>()
 	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<string>()
 		// 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,