Fix modal memory graph scoping by document and space (supermemoryai#948)

Author: ishaanxgupta

apps/web/components/document-modal/graph-list-memories.tsx

@@ -2,6 +2,7 @@ import { useState } from "react"
 import { cn } from "@lib/utils"
 import { Tabs, TabsList, TabsTrigger, TabsContent } from "@ui/components/tabs"
 import { MemoryGraph } from "../memory-graph"
+import { useProject } from "@/stores"
 
 export interface MemoryEntry {
 	id: string
@@ -172,6 +173,7 @@ export function GraphListMemories({
 	memoryEntries: MemoryEntry[]
 	documentId?: string
 }) {
+	const { effectiveContainerTags } = useProject()
 	const [expandedMemories, setExpandedMemories] = useState<Set<string>>(
 		new Set(),
 	)
@@ -263,7 +265,10 @@ export function GraphListMemories({
 				<TabsContent value="graph" className="flex-1 min-h-0 mt-3">
 					<div className="size-full rounded-lg overflow-hidden">
 						<MemoryGraph
+							containerTags={effectiveContainerTags}
 							documentIds={documentId ? [documentId] : undefined}
+							highlightDocumentIds={documentId ? [documentId] : undefined}
+							highlightsVisible
 							variant="consumer"
 							maxNodes={50}
 						/>

apps/web/components/memory-graph/hooks/use-graph-api.ts

@@ -13,6 +13,7 @@ const PAGE_SIZE = 100
 
 interface UseGraphApiOptions {
 	containerTags?: string[]
+	documentIds?: string[]
 	enabled?: boolean
 }
 
@@ -81,20 +82,36 @@ function toGraphMemory(mem: ApiMemoryEntry): GraphApiMemory {
 	}
 }
 
-function toGraphDocument(doc: ApiDocument): GraphApiDocument {
+function toGraphDocument(
+	doc: ApiDocument,
+	containerTags?: string[],
+): GraphApiDocument {
+	const allowedContainerTags = new Set(containerTags?.filter(Boolean) ?? [])
+	const memoryEntries =
+		allowedContainerTags.size > 0
+			? doc.memoryEntries.filter(
+					(mem) =>
+						mem.spaceContainerTag != null &&
+						allowedContainerTags.has(mem.spaceContainerTag),
+				)
+			: doc.memoryEntries
+
 	return {
 		id: doc.id,
 		title: doc.title,
 		summary: doc.summary ?? null,
 		documentType: doc.type,
 		createdAt: doc.createdAt,
 		updatedAt: doc.updatedAt,
-		memories: doc.memoryEntries.map(toGraphMemory),
+		memories: memoryEntries.map(toGraphMemory),
 	}
 }
 
 export function useGraphApi(options: UseGraphApiOptions = {}) {
-	const { containerTags, enabled = true } = options
+	const { containerTags, documentIds, enabled = true } = options
+	const filteredDocumentIds = documentIds?.filter(Boolean)
+	const hasDocumentIds =
+		filteredDocumentIds != null && filteredDocumentIds.length > 0
 
 	const {
 		data,
@@ -104,19 +121,33 @@ export function useGraphApi(options: UseGraphApiOptions = {}) {
 		hasNextPage,
 		fetchNextPage,
 	} = useInfiniteQuery<ApiDocumentsResponse, Error>({
-		queryKey: ["documents-with-memories", containerTags, []],
+		queryKey: [
+			"documents-with-memories",
+			containerTags,
+			[],
+			filteredDocumentIds,
+		],
 		initialPageParam: 1,
 		queryFn: async ({ pageParam }) => {
-			const response = await $fetch("@post/documents/documents", {
-				body: {
-					page: pageParam as number,
-					limit: PAGE_SIZE,
-					sort: "createdAt",
-					order: "desc",
-					containerTags,
-				},
-				disableValidation: true,
-			})
+			const response = hasDocumentIds
+				? await $fetch("@post/documents/documents/by-ids", {
+						body: {
+							ids: filteredDocumentIds,
+							by: "id",
+							containerTags,
+						},
+						disableValidation: true,
+					})
+				: await $fetch("@post/documents/documents", {
+						body: {
+							page: pageParam as number,
+							limit: PAGE_SIZE,
+							sort: "createdAt",
+							order: "desc",
+							containerTags,
+						},
+						disableValidation: true,
+					})
 
 			if (response.error) {
 				throw new Error(response.error?.message || "Failed to fetch documents")
@@ -134,8 +165,10 @@ export function useGraphApi(options: UseGraphApiOptions = {}) {
 
 	const documents = useMemo(() => {
 		if (!data?.pages) return []
-		return data.pages.flatMap((page) => page.documents.map(toGraphDocument))
-	}, [data])
+		return data.pages.flatMap((page) =>
+			page.documents.map((doc) => toGraphDocument(doc, containerTags)),
+		)
+	}, [data, containerTags])
 
 	const totalCount = data?.pages[0]?.pagination.totalItems ?? 0
 

apps/web/components/memory-graph/memory-graph-wrapper.tsx

@@ -29,6 +29,7 @@ export function MemoryGraph({
 	error: externalError = null,
 	variant = "console",
 	containerTags,
+	documentIds,
 	maxNodes,
 	canvasRef,
 	...rest
@@ -57,6 +58,7 @@ export function MemoryGraph({
 		totalCount,
 	} = useGraphApi({
 		containerTags,
+		documentIds,
 		enabled: containerSize.width > 0 && containerSize.height > 0,
 	})