From 8be7f2e05c3bd0451da855536d59e2f02c9d27c4 Mon Sep 17 00:00:00 2001 From: Anish Sarkar <104695310+AnishSarkar22@users.noreply.github.com> Date: Wed, 29 Apr 2026 04:19:07 +0530 Subject: [PATCH] refactor(mentions): update document mention handling to use document keys for consistency across components --- surfsense_web/components/assistant-ui/thread.tsx | 11 ++++++++--- .../components/documents/FolderTreeView.tsx | 13 +++++++------ .../layout/ui/sidebar/DocumentsSidebar.tsx | 14 ++++++++++---- 3 files changed, 25 insertions(+), 13 deletions(-) diff --git a/surfsense_web/components/assistant-ui/thread.tsx b/surfsense_web/components/assistant-ui/thread.tsx index f9e5ca7fb..3964d60e5 100644 --- a/surfsense_web/components/assistant-ui/thread.tsx +++ b/surfsense_web/components/assistant-ui/thread.tsx @@ -628,9 +628,14 @@ const Composer: FC = () => { const handleDocumentRemove = useCallback( (docId: number, docType?: string) => { - setMentionedDocuments((prev) => - prev.filter((doc) => !(doc.id === docId && doc.document_type === docType)) - ); + setMentionedDocuments((prev) => { + if (!docType) { + // Defensive fallback: keep UI in sync even when chip type is unavailable. + return prev.filter((doc) => doc.id !== docId); + } + const removedKey = getMentionDocKey({ id: docId, document_type: docType }); + return prev.filter((doc) => getMentionDocKey(doc) !== removedKey); + }); }, [setMentionedDocuments] ); diff --git a/surfsense_web/components/documents/FolderTreeView.tsx b/surfsense_web/components/documents/FolderTreeView.tsx index 9b7a393d8..2063fbee5 100644 --- a/surfsense_web/components/documents/FolderTreeView.tsx +++ b/surfsense_web/components/documents/FolderTreeView.tsx @@ -7,6 +7,7 @@ import { DndProvider } from "react-dnd"; import { HTML5Backend } from "react-dnd-html5-backend"; import { renamingFolderIdAtom } from "@/atoms/documents/folder.atoms"; import type { DocumentTypeEnum } from "@/contracts/types/document.types"; +import { getMentionDocKey } from "@/lib/chat/mention-doc-key"; import { DocumentNode, type DocumentNodeDoc } from "./DocumentNode"; import { type FolderDisplay, FolderNode } from "./FolderNode"; @@ -17,7 +18,7 @@ interface FolderTreeViewProps { documents: DocumentNodeDoc[]; expandedIds: Set; onToggleExpand: (folderId: number) => void; - mentionedDocIds: Set; + mentionedDocKeys: Set; onToggleChatMention: ( doc: { id: number; title: string; document_type: string }, isMentioned: boolean @@ -62,7 +63,7 @@ export function FolderTreeView({ documents, expandedIds, onToggleExpand, - mentionedDocIds, + mentionedDocKeys, onToggleChatMention, onToggleFolderSelect, onRenameFolder, @@ -181,7 +182,7 @@ export function FolderTreeView({ function compute(folderId: number): { selected: number; total: number } { const directDocs = (docsByFolder[folderId] ?? []).filter(isSelectable); - let selected = directDocs.filter((d) => mentionedDocIds.has(d.id)).length; + let selected = directDocs.filter((d) => mentionedDocKeys.has(getMentionDocKey(d))).length; let total = directDocs.length; for (const child of foldersByParent[folderId] ?? []) { @@ -202,7 +203,7 @@ export function FolderTreeView({ if (states[f.id] === undefined) compute(f.id); } return states; - }, [folders, docsByFolder, foldersByParent, mentionedDocIds]); + }, [folders, docsByFolder, foldersByParent, mentionedDocKeys]); const folderMap = useMemo(() => { const map: Record = {}; @@ -276,7 +277,7 @@ export function FolderTreeView({ key={`doc-${d.id}`} doc={d} depth={depth} - isMentioned={mentionedDocIds.has(d.id)} + isMentioned={mentionedDocKeys.has(getMentionDocKey(d))} onToggleChatMention={onToggleChatMention} onPreview={onPreviewDocument} onEdit={onEditDocument} @@ -356,7 +357,7 @@ export function FolderTreeView({ key={`doc-${d.id}`} doc={d} depth={depth} - isMentioned={mentionedDocIds.has(d.id)} + isMentioned={mentionedDocKeys.has(getMentionDocKey(d))} onToggleChatMention={onToggleChatMention} onPreview={onPreviewDocument} onEdit={onEditDocument} diff --git a/surfsense_web/components/layout/ui/sidebar/DocumentsSidebar.tsx b/surfsense_web/components/layout/ui/sidebar/DocumentsSidebar.tsx index 63b6dc1b7..6ff087b9b 100644 --- a/surfsense_web/components/layout/ui/sidebar/DocumentsSidebar.tsx +++ b/surfsense_web/components/layout/ui/sidebar/DocumentsSidebar.tsx @@ -416,7 +416,10 @@ function AuthenticatedDocumentsSidebarBase({ const { mutateAsync: deleteDocumentMutation } = useAtomValue(deleteDocumentMutationAtom); const [sidebarDocs, setSidebarDocs] = useAtom(mentionedDocumentsAtom); - const mentionedDocIds = useMemo(() => new Set(sidebarDocs.map((d) => d.id)), [sidebarDocs]); + const mentionedDocKeys = useMemo( + () => new Set(sidebarDocs.map((d) => getMentionDocKey(d))), + [sidebarDocs] + ); // Folder state const [expandedFolderMap, setExpandedFolderMap] = useAtom(expandedFolderIdsAtom); @@ -1143,7 +1146,7 @@ function AuthenticatedDocumentsSidebarBase({ documents={searchFilteredDocuments} expandedIds={expandedIds} onToggleExpand={toggleFolderExpand} - mentionedDocIds={mentionedDocIds} + mentionedDocKeys={mentionedDocKeys} onToggleChatMention={handleToggleChatMention} onToggleFolderSelect={handleToggleFolderSelect} onRenameFolder={handleRenameFolder} @@ -1572,7 +1575,10 @@ function AnonymousDocumentsSidebar({ const [search, setSearch] = useState(""); const [sidebarDocs, setSidebarDocs] = useAtom(mentionedDocumentsAtom); - const mentionedDocIds = useMemo(() => new Set(sidebarDocs.map((d) => d.id)), [sidebarDocs]); + const mentionedDocKeys = useMemo( + () => new Set(sidebarDocs.map((d) => getMentionDocKey(d))), + [sidebarDocs] + ); const handleToggleChatMention = useCallback( (doc: { id: number; title: string; document_type: string }, isMentioned: boolean) => { @@ -1801,7 +1807,7 @@ function AnonymousDocumentsSidebar({ documents={searchFilteredDocs} expandedIds={new Set()} onToggleExpand={() => {}} - mentionedDocIds={mentionedDocIds} + mentionedDocKeys={mentionedDocKeys} onToggleChatMention={handleToggleChatMention} onToggleFolderSelect={() => {}} onRenameFolder={() => gate("rename folders")}