From 960f761c6cb52fa0179f47ad68a0f0913608b80c Mon Sep 17 00:00:00 2001 From: Anish Sarkar <104695310+AnishSarkar22@users.noreply.github.com> Date: Tue, 28 Apr 2026 17:50:21 +0530 Subject: [PATCH 1/6] feat(mentions): add sidebar mention event atom and enhance document mention handling in the editor --- .../atoms/chat/mentioned-documents.atom.ts | 14 ++- .../assistant-ui/inline-mention-editor.tsx | 96 +++++++++++++--- .../components/assistant-ui/thread.tsx | 38 +++++++ .../layout/ui/sidebar/DocumentsSidebar.tsx | 106 ++++++++++++++++-- 4 files changed, 227 insertions(+), 27 deletions(-) diff --git a/surfsense_web/atoms/chat/mentioned-documents.atom.ts b/surfsense_web/atoms/chat/mentioned-documents.atom.ts index ee93a409a..27a205af2 100644 --- a/surfsense_web/atoms/chat/mentioned-documents.atom.ts +++ b/surfsense_web/atoms/chat/mentioned-documents.atom.ts @@ -11,12 +11,24 @@ export const mentionedDocumentsAtom = atom[] >([]); +export interface SidebarMentionEvent { + kind: "add" | "remove"; + docs: Pick[]; + nonce: number; +} + +/** + * Event atom used to tell the composer that documents were selected/unselected + * from sidebar checkboxes, so chips can be inserted/removed in-editor. + */ +export const sidebarMentionEventAtom = atom(null); + /** * Derived read-only atom that merges @-mention chips and sidebar selections * into a single deduplicated set of document IDs for the backend. diff --git a/surfsense_web/components/assistant-ui/inline-mention-editor.tsx b/surfsense_web/components/assistant-ui/inline-mention-editor.tsx index 45ad219dd..31b2a0566 100644 --- a/surfsense_web/components/assistant-ui/inline-mention-editor.tsx +++ b/surfsense_web/components/assistant-ui/inline-mention-editor.tsx @@ -44,7 +44,10 @@ export interface InlineMentionEditorRef { setText: (text: string) => void; getText: () => string; getMentionedDocuments: () => MentionedDocument[]; - insertDocumentChip: (doc: Pick) => void; + insertDocumentChip: ( + doc: Pick, + options?: { removeTriggerText?: boolean } + ) => void; removeDocumentChip: (docId: number, docType?: string) => void; setDocumentChipStatus: ( docId: number, @@ -129,6 +132,40 @@ export const InlineMentionEditor = forwardRef new Map(initialDocuments.map((d) => [`${d.document_type ?? "UNKNOWN"}:${d.id}`, d])) ); const isComposingRef = useRef(false); + const lastSelectionRangeRef = useRef(null); + const isSelectionInsideEditor = useCallback( + (selection: Selection | null): selection is Selection => { + if (!selection || selection.rangeCount === 0 || !editorRef.current) return false; + const range = selection.getRangeAt(0); + return editorRef.current.contains(range.startContainer); + }, + [] + ); + + const rememberSelection = useCallback(() => { + const selection = window.getSelection(); + if (!isSelectionInsideEditor(selection)) return; + lastSelectionRangeRef.current = selection.getRangeAt(0).cloneRange(); + }, [isSelectionInsideEditor]); + + const restoreRememberedSelection = useCallback((): Selection | null => { + const selection = window.getSelection(); + if (!selection) return null; + if (!lastSelectionRangeRef.current) return selection; + selection.removeAllRanges(); + selection.addRange(lastSelectionRangeRef.current.cloneRange()); + return selection; + }, []); + + useEffect(() => { + const handleSelectionChange = () => { + if (document.activeElement !== editorRef.current) return; + rememberSelection(); + }; + document.addEventListener("selectionchange", handleSelectionChange); + return () => document.removeEventListener("selectionchange", handleSelectionChange); + }, [rememberSelection]); + // Sync initial documents useEffect(() => { @@ -157,7 +194,7 @@ export const InlineMentionEditor = forwardRef { @@ -299,8 +336,12 @@ export const InlineMentionEditor = forwardRef) => { + ( + doc: Pick, + options?: { removeTriggerText?: boolean } + ) => { if (!editorRef.current) return; + const removeTriggerText = options?.removeTriggerText ?? true; // Validate required fields for type safety if (typeof doc.id !== "number" || typeof doc.title !== "string") { @@ -320,20 +361,23 @@ export const InlineMentionEditor = forwardRef - {/** biome-ignore lint/a11y/useSemanticElements: */}
{ // Document mention state (atoms persist across component remounts) const [mentionedDocuments, setMentionedDocuments] = useAtom(mentionedDocumentsAtom); const setSidebarDocs = useSetAtom(sidebarSelectedDocumentsAtom); + const [sidebarMentionEvent, setSidebarMentionEvent] = useAtom(sidebarMentionEventAtom); const [showDocumentPopover, setShowDocumentPopover] = useState(false); const [showPromptPicker, setShowPromptPicker] = useState(false); const [mentionQuery, setMentionQuery] = useState(""); @@ -660,6 +662,42 @@ const Composer: FC = () => { [mentionedDocuments, setMentionedDocuments] ); + useEffect(() => { + if (!sidebarMentionEvent) return; + + const eventDocs = sidebarMentionEvent.docs; + if (eventDocs.length === 0) { + setSidebarMentionEvent(null); + return; + } + + const docKey = (doc: Pick) => + `${doc.document_type}:${doc.id}`; + const mentionedKeys = new Set(mentionedDocuments.map(docKey)); + + if (sidebarMentionEvent.kind === "add") { + const docsToAdd = eventDocs.filter((doc) => !mentionedKeys.has(docKey(doc))); + for (const doc of docsToAdd) { + editorRef.current?.insertDocumentChip(doc, { removeTriggerText: false }); + } + if (docsToAdd.length > 0) { + setMentionedDocuments((prev) => { + const existing = new Set(prev.map(docKey)); + const uniqueAdds = docsToAdd.filter((doc) => !existing.has(docKey(doc))); + return uniqueAdds.length > 0 ? [...prev, ...uniqueAdds] : prev; + }); + } + } else { + const removeKeys = new Set(eventDocs.map(docKey)); + for (const doc of eventDocs) { + editorRef.current?.removeDocumentChip(doc.id, doc.document_type); + } + setMentionedDocuments((prev) => prev.filter((doc) => !removeKeys.has(docKey(doc)))); + } + + setSidebarMentionEvent(null); + }, [sidebarMentionEvent, mentionedDocuments, setMentionedDocuments, setSidebarMentionEvent]); + return ( new Set(sidebarDocs.map((d) => d.id)), [sidebarDocs]); // Folder state @@ -857,19 +861,42 @@ function AuthenticatedDocumentsSidebarBase({ const handleToggleChatMention = useCallback( (doc: { id: number; title: string; document_type: string }, isMentioned: boolean) => { + const key = `${doc.document_type}:${doc.id}`; if (isMentioned) { - setSidebarDocs((prev) => prev.filter((d) => d.id !== doc.id)); + setSidebarDocs((prev) => prev.filter((d) => `${d.document_type}:${d.id}` !== key)); + setSidebarMentionEvent({ + kind: "remove", + docs: [ + { + id: doc.id, + title: doc.title, + document_type: doc.document_type as DocumentTypeEnum, + }, + ], + nonce: Date.now(), + }); } else { setSidebarDocs((prev) => { - if (prev.some((d) => d.id === doc.id)) return prev; + if (prev.some((d) => `${d.document_type}:${d.id}` === key)) return prev; return [ ...prev, { id: doc.id, title: doc.title, document_type: doc.document_type as DocumentTypeEnum }, ]; }); + setSidebarMentionEvent({ + kind: "add", + docs: [ + { + id: doc.id, + title: doc.title, + document_type: doc.document_type as DocumentTypeEnum, + }, + ], + nonce: Date.now(), + }); } }, - [setSidebarDocs] + [setSidebarDocs, setSidebarMentionEvent] ); const handleToggleFolderSelect = useCallback( @@ -891,10 +918,18 @@ function AuthenticatedDocumentsSidebarBase({ if (subtreeDocs.length === 0) return; if (selectAll) { + const existingKeys = new Set(sidebarDocs.map((d) => `${d.document_type}:${d.id}`)); + const docsToAdd = subtreeDocs + .filter((d) => !existingKeys.has(`${d.document_type}:${d.id}`)) + .map((d) => ({ + id: d.id, + title: d.title, + document_type: d.document_type as DocumentTypeEnum, + })); setSidebarDocs((prev) => { - const existingIds = new Set(prev.map((d) => d.id)); + const existingDocKeys = new Set(prev.map((d) => `${d.document_type}:${d.id}`)); const newDocs = subtreeDocs - .filter((d) => !existingIds.has(d.id)) + .filter((d) => !existingDocKeys.has(`${d.document_type}:${d.id}`)) .map((d) => ({ id: d.id, title: d.title, @@ -902,12 +937,35 @@ function AuthenticatedDocumentsSidebarBase({ })); return newDocs.length > 0 ? [...prev, ...newDocs] : prev; }); + if (docsToAdd.length > 0) { + setSidebarMentionEvent({ + kind: "add", + docs: docsToAdd, + nonce: Date.now(), + }); + } } else { - const idsToRemove = new Set(subtreeDocs.map((d) => d.id)); - setSidebarDocs((prev) => prev.filter((d) => !idsToRemove.has(d.id))); + const keysToRemove = new Set(subtreeDocs.map((d) => `${d.document_type}:${d.id}`)); + const docsToRemove = sidebarDocs + .filter((d) => keysToRemove.has(`${d.document_type}:${d.id}`)) + .map((d) => ({ + id: d.id, + title: d.title, + document_type: d.document_type as DocumentTypeEnum, + })); + setSidebarDocs((prev) => + prev.filter((d) => !keysToRemove.has(`${d.document_type}:${d.id}`)) + ); + if (docsToRemove.length > 0) { + setSidebarMentionEvent({ + kind: "remove", + docs: docsToRemove, + nonce: Date.now(), + }); + } } }, - [treeDocuments, foldersByParent, setSidebarDocs] + [treeDocuments, foldersByParent, sidebarDocs, setSidebarDocs, setSidebarMentionEvent] ); const searchFilteredDocuments = useMemo(() => { @@ -1568,23 +1626,47 @@ function AnonymousDocumentsSidebar({ const [search, setSearch] = useState(""); const [sidebarDocs, setSidebarDocs] = useAtom(sidebarSelectedDocumentsAtom); + const setSidebarMentionEvent = useSetAtom(sidebarMentionEventAtom); const mentionedDocIds = useMemo(() => new Set(sidebarDocs.map((d) => d.id)), [sidebarDocs]); const handleToggleChatMention = useCallback( (doc: { id: number; title: string; document_type: string }, isMentioned: boolean) => { + const key = `${doc.document_type}:${doc.id}`; if (isMentioned) { - setSidebarDocs((prev) => prev.filter((d) => d.id !== doc.id)); + setSidebarDocs((prev) => prev.filter((d) => `${d.document_type}:${d.id}` !== key)); + setSidebarMentionEvent({ + kind: "remove", + docs: [ + { + id: doc.id, + title: doc.title, + document_type: doc.document_type as DocumentTypeEnum, + }, + ], + nonce: Date.now(), + }); } else { setSidebarDocs((prev) => { - if (prev.some((d) => d.id === doc.id)) return prev; + if (prev.some((d) => `${d.document_type}:${d.id}` === key)) return prev; return [ ...prev, { id: doc.id, title: doc.title, document_type: doc.document_type as DocumentTypeEnum }, ]; }); + setSidebarMentionEvent({ + kind: "add", + docs: [ + { + id: doc.id, + title: doc.title, + document_type: doc.document_type as DocumentTypeEnum, + }, + ], + nonce: Date.now(), + }); } }, - [setSidebarDocs] + [setSidebarDocs, setSidebarMentionEvent] ); const uploadedDoc = anonMode.isAnonymous ? anonMode.uploadedDoc : null; From 1427809119d82c0335eb5f5e600acff816d9d3f4 Mon Sep 17 00:00:00 2001 From: Anish Sarkar <104695310+AnishSarkar22@users.noreply.github.com> Date: Tue, 28 Apr 2026 18:20:53 +0530 Subject: [PATCH 2/6] refactor(mentions): consolidate sidebar document handling into mentionedDocumentsAtom and remove sidebarSelectedDocumentsAtom references --- .../new-chat/[[...chat_id]]/page.tsx | 21 ++--- .../atoms/chat/mentioned-documents.atom.ts | 43 +++++----- .../assistant-ui/inline-mention-editor.tsx | 24 ++---- .../components/assistant-ui/thread.tsx | 65 ++++----------- .../layout/ui/sidebar/DocumentsSidebar.tsx | 82 +------------------ 5 files changed, 54 insertions(+), 181 deletions(-) diff --git a/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx b/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx index 06f3bf79f..c363d69d3 100644 --- a/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx +++ b/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx @@ -24,7 +24,6 @@ import { mentionedDocumentIdsAtom, mentionedDocumentsAtom, messageDocumentsMapAtom, - sidebarSelectedDocumentsAtom, } from "@/atoms/chat/mentioned-documents.atom"; import { clearPlanOwnerRegistry, @@ -216,12 +215,10 @@ export default function NewChatPage() { // Get disabled tools from the tool toggle UI const disabledTools = useAtomValue(disabledToolsAtom); - // Get mentioned document IDs from the composer (derived from @ mentions + sidebar selections) + // Get mentioned document IDs from the composer. const mentionedDocumentIds = useAtomValue(mentionedDocumentIdsAtom); const mentionedDocuments = useAtomValue(mentionedDocumentsAtom); - const sidebarDocuments = useAtomValue(sidebarSelectedDocumentsAtom); const setMentionedDocuments = useSetAtom(mentionedDocumentsAtom); - const setSidebarDocuments = useSetAtom(sidebarSelectedDocumentsAtom); const setMessageDocumentsMap = useSetAtom(messageDocumentsMapAtom); const setCurrentThreadState = useSetAtom(currentThreadAtom); const setTargetCommentId = useSetAtom(setTargetCommentIdAtom); @@ -319,7 +316,6 @@ export default function NewChatPage() { setCurrentThread(null); setMentionedDocuments([]); tokenUsageStore.clear(); - setSidebarDocuments([]); setMessageDocumentsMap({}); clearPlanOwnerRegistry(); closeReportPanel(); @@ -387,7 +383,6 @@ export default function NewChatPage() { urlChatId, setMessageDocumentsMap, setMentionedDocuments, - setSidebarDocuments, closeReportPanel, closeEditorPanel, removeChatTab, @@ -578,15 +573,14 @@ export default function NewChatPage() { messageLength: userQuery.length, }); - // Combine @-mention chips + sidebar selections for display & persistence + // Collect unique mentioned docs for display & persistence const allMentionedDocs: MentionedDocumentInfo[] = []; const seenDocKeys = new Set(); - for (const doc of [...mentionedDocuments, ...sidebarDocuments]) { + for (const doc of mentionedDocuments) { const key = `${doc.document_type}:${doc.id}`; - if (!seenDocKeys.has(key)) { - seenDocKeys.add(key); - allMentionedDocs.push({ id: doc.id, title: doc.title, document_type: doc.document_type }); - } + if (seenDocKeys.has(key)) continue; + seenDocKeys.add(key); + allMentionedDocs.push({ id: doc.id, title: doc.title, document_type: doc.document_type }); } if (allMentionedDocs.length > 0) { @@ -689,7 +683,6 @@ export default function NewChatPage() { // Clear mentioned documents after capturing them if (hasDocumentIds || hasSurfsenseDocIds) { setMentionedDocuments([]); - setSidebarDocuments([]); } const response = await fetch(`${backendUrl}/api/v1/new_chat`, { @@ -979,9 +972,7 @@ export default function NewChatPage() { messages, mentionedDocumentIds, mentionedDocuments, - sidebarDocuments, setMentionedDocuments, - setSidebarDocuments, setMessageDocumentsMap, setAgentCreatedDocuments, queryClient, diff --git a/surfsense_web/atoms/chat/mentioned-documents.atom.ts b/surfsense_web/atoms/chat/mentioned-documents.atom.ts index 27a205af2..47401995d 100644 --- a/surfsense_web/atoms/chat/mentioned-documents.atom.ts +++ b/surfsense_web/atoms/chat/mentioned-documents.atom.ts @@ -10,33 +10,34 @@ import type { Document } from "@/contracts/types/document.types"; export const mentionedDocumentsAtom = atom[]>([]); /** - * Atom to store documents selected via the sidebar checkboxes / row clicks. - * These power the selected-sources badge and backend doc filters. + * Back-compat alias for sidebar checkbox selection. + * This now points to mentionedDocumentsAtom so the app has a single source + * of truth for mentioned/selected documents. */ export const sidebarSelectedDocumentsAtom = atom< - Pick[] ->([]); - -export interface SidebarMentionEvent { - kind: "add" | "remove"; - docs: Pick[]; - nonce: number; -} + Pick[], + [ + | Pick[] + | (( + prev: Pick[] + ) => Pick[]), + ], + void +>( + (get) => get(mentionedDocumentsAtom), + (get, set, update) => { + const prev = get(mentionedDocumentsAtom); + const next = typeof update === "function" ? update(prev) : update; + set(mentionedDocumentsAtom, next); + } +); /** - * Event atom used to tell the composer that documents were selected/unselected - * from sidebar checkboxes, so chips can be inserted/removed in-editor. - */ -export const sidebarMentionEventAtom = atom(null); - -/** - * Derived read-only atom that merges @-mention chips and sidebar selections - * into a single deduplicated set of document IDs for the backend. + * Derived read-only atom that maps deduplicated mentioned docs + * into backend payload fields. */ export const mentionedDocumentIdsAtom = atom((get) => { - const chipDocs = get(mentionedDocumentsAtom); - const sidebarDocs = get(sidebarSelectedDocumentsAtom); - const allDocs = [...chipDocs, ...sidebarDocs]; + const allDocs = get(mentionedDocumentsAtom); const seen = new Set(); const deduped = allDocs.filter((d) => { const key = `${d.document_type}:${d.id}`; diff --git a/surfsense_web/components/assistant-ui/inline-mention-editor.tsx b/surfsense_web/components/assistant-ui/inline-mention-editor.tsx index 31b2a0566..81d6cbd77 100644 --- a/surfsense_web/components/assistant-ui/inline-mention-editor.tsx +++ b/surfsense_web/components/assistant-ui/inline-mention-editor.tsx @@ -11,25 +11,13 @@ import { useRef, useState, } from "react"; -import { flushSync } from "react-dom"; -import { createRoot } from "react-dom/client"; +import { renderToStaticMarkup } from "react-dom/server"; import { getConnectorIcon } from "@/contracts/enums/connectorIcons"; import type { Document } from "@/contracts/types/document.types"; import { cn } from "@/lib/utils"; -// Render a React element to an HTML string on the client without pulling -// `react-dom/server` into the bundle. `createRoot` + `flushSync` use the -// same `react-dom` package React itself imports, so this adds zero new -// runtime weight. function renderElementToHTML(element: ReactElement): string { - const container = document.createElement("div"); - const root = createRoot(container); - flushSync(() => { - root.render(element); - }); - const html = container.innerHTML; - root.unmount(); - return html; + return renderToStaticMarkup(element); } export interface MentionedDocument { @@ -182,7 +170,7 @@ export const InlineMentionEditor = forwardRef { @@ -779,6 +767,7 @@ export const InlineMentionEditor = forwardRef + {/* biome-ignore lint/a11y/noStaticElementInteractions: contenteditable mention editor requires a div for inline chips */}
{/* Placeholder with fade animation on change */} {isEmpty && ( diff --git a/surfsense_web/components/assistant-ui/thread.tsx b/surfsense_web/components/assistant-ui/thread.tsx index 111b5f3cc..dcc068bd1 100644 --- a/surfsense_web/components/assistant-ui/thread.tsx +++ b/surfsense_web/components/assistant-ui/thread.tsx @@ -38,12 +38,9 @@ import { import { chatSessionStateAtom } from "@/atoms/chat/chat-session-state.atom"; import { mentionedDocumentsAtom, - sidebarMentionEventAtom, - sidebarSelectedDocumentsAtom, } from "@/atoms/chat/mentioned-documents.atom"; import { connectorDialogOpenAtom } from "@/atoms/connector-dialog/connector-dialog.atoms"; import { connectorsAtom } from "@/atoms/connectors/connector-query.atoms"; -import { documentsSidebarOpenAtom } from "@/atoms/documents/ui.atoms"; import { membersAtom } from "@/atoms/members/members-query.atoms"; import { globalNewLLMConfigsAtom, @@ -336,8 +333,6 @@ const ClipboardChip: FC<{ text: string; onDismiss: () => void }> = ({ text, onDi const Composer: FC = () => { // Document mention state (atoms persist across component remounts) const [mentionedDocuments, setMentionedDocuments] = useAtom(mentionedDocumentsAtom); - const setSidebarDocs = useSetAtom(sidebarSelectedDocumentsAtom); - const [sidebarMentionEvent, setSidebarMentionEvent] = useAtom(sidebarMentionEventAtom); const [showDocumentPopover, setShowDocumentPopover] = useState(false); const [showPromptPicker, setShowPromptPicker] = useState(false); const [mentionQuery, setMentionQuery] = useState(""); @@ -578,7 +573,6 @@ const Composer: FC = () => { aui.composer().send(); editorRef.current?.clear(); setMentionedDocuments([]); - setSidebarDocs([]); // With turnAnchor="top", ViewportSlack adds min-height to the last // assistant message so that scrolling-to-bottom actually positions the @@ -625,7 +619,6 @@ const Composer: FC = () => { clipboardInitialText, aui, setMentionedDocuments, - setSidebarDocs, threadViewportStore, ]); @@ -663,40 +656,29 @@ const Composer: FC = () => { ); useEffect(() => { - if (!sidebarMentionEvent) return; + const editor = editorRef.current; + if (!editor) return; - const eventDocs = sidebarMentionEvent.docs; - if (eventDocs.length === 0) { - setSidebarMentionEvent(null); - return; + const toKey = (doc: { id: number; document_type?: string }) => + `${doc.document_type ?? "UNKNOWN"}:${doc.id}`; + + const atomDocs = mentionedDocuments; + const editorDocs = editor.getMentionedDocuments(); + const atomKeys = new Set(atomDocs.map(toKey)); + const editorKeys = new Set(editorDocs.map(toKey)); + + for (const doc of atomDocs) { + if (!editorKeys.has(toKey(doc))) { + editor.insertDocumentChip(doc, { removeTriggerText: false }); + } } - const docKey = (doc: Pick) => - `${doc.document_type}:${doc.id}`; - const mentionedKeys = new Set(mentionedDocuments.map(docKey)); - - if (sidebarMentionEvent.kind === "add") { - const docsToAdd = eventDocs.filter((doc) => !mentionedKeys.has(docKey(doc))); - for (const doc of docsToAdd) { - editorRef.current?.insertDocumentChip(doc, { removeTriggerText: false }); + for (const doc of editorDocs) { + if (!atomKeys.has(toKey(doc))) { + editor.removeDocumentChip(doc.id, doc.document_type); } - if (docsToAdd.length > 0) { - setMentionedDocuments((prev) => { - const existing = new Set(prev.map(docKey)); - const uniqueAdds = docsToAdd.filter((doc) => !existing.has(docKey(doc))); - return uniqueAdds.length > 0 ? [...prev, ...uniqueAdds] : prev; - }); - } - } else { - const removeKeys = new Set(eventDocs.map(docKey)); - for (const doc of eventDocs) { - editorRef.current?.removeDocumentChip(doc.id, doc.document_type); - } - setMentionedDocuments((prev) => prev.filter((doc) => !removeKeys.has(docKey(doc)))); } - - setSidebarMentionEvent(null); - }, [sidebarMentionEvent, mentionedDocuments, setMentionedDocuments, setSidebarMentionEvent]); + }, [mentionedDocuments]); return ( @@ -775,8 +757,6 @@ interface ComposerActionProps { const ComposerAction: FC = ({ isBlockedByOtherUser = false }) => { const mentionedDocuments = useAtomValue(mentionedDocumentsAtom); - const sidebarDocs = useAtomValue(sidebarSelectedDocumentsAtom); - const setDocumentsSidebarOpen = useSetAtom(documentsSidebarOpenAtom); const setConnectorDialogOpen = useSetAtom(connectorDialogOpenAtom); const [toolsPopoverOpen, setToolsPopoverOpen] = useState(false); const isDesktop = useMediaQuery("(min-width: 640px)"); @@ -1222,15 +1202,6 @@ const ComposerAction: FC = ({ isBlockedByOtherUser = false )} - {sidebarDocs.length > 0 && ( - - )}
{!hasModelConfigured && (
diff --git a/surfsense_web/components/layout/ui/sidebar/DocumentsSidebar.tsx b/surfsense_web/components/layout/ui/sidebar/DocumentsSidebar.tsx index 00cc2d4ef..3c5a64b0e 100644 --- a/surfsense_web/components/layout/ui/sidebar/DocumentsSidebar.tsx +++ b/surfsense_web/components/layout/ui/sidebar/DocumentsSidebar.tsx @@ -24,7 +24,6 @@ import type React from "react"; import { useCallback, useEffect, useMemo, useRef, useState } from "react"; import { toast } from "sonner"; import { - sidebarMentionEventAtom, sidebarSelectedDocumentsAtom, } from "@/atoms/chat/mentioned-documents.atom"; import { connectorDialogOpenAtom } from "@/atoms/connector-dialog/connector-dialog.atoms"; @@ -416,7 +415,6 @@ function AuthenticatedDocumentsSidebarBase({ const { mutateAsync: deleteDocumentMutation } = useAtomValue(deleteDocumentMutationAtom); const [sidebarDocs, setSidebarDocs] = useAtom(sidebarSelectedDocumentsAtom); - const setSidebarMentionEvent = useSetAtom(sidebarMentionEventAtom); const mentionedDocIds = useMemo(() => new Set(sidebarDocs.map((d) => d.id)), [sidebarDocs]); // Folder state @@ -864,17 +862,6 @@ function AuthenticatedDocumentsSidebarBase({ const key = `${doc.document_type}:${doc.id}`; if (isMentioned) { setSidebarDocs((prev) => prev.filter((d) => `${d.document_type}:${d.id}` !== key)); - setSidebarMentionEvent({ - kind: "remove", - docs: [ - { - id: doc.id, - title: doc.title, - document_type: doc.document_type as DocumentTypeEnum, - }, - ], - nonce: Date.now(), - }); } else { setSidebarDocs((prev) => { if (prev.some((d) => `${d.document_type}:${d.id}` === key)) return prev; @@ -883,20 +870,9 @@ function AuthenticatedDocumentsSidebarBase({ { id: doc.id, title: doc.title, document_type: doc.document_type as DocumentTypeEnum }, ]; }); - setSidebarMentionEvent({ - kind: "add", - docs: [ - { - id: doc.id, - title: doc.title, - document_type: doc.document_type as DocumentTypeEnum, - }, - ], - nonce: Date.now(), - }); } }, - [setSidebarDocs, setSidebarMentionEvent] + [setSidebarDocs] ); const handleToggleFolderSelect = useCallback( @@ -918,14 +894,6 @@ function AuthenticatedDocumentsSidebarBase({ if (subtreeDocs.length === 0) return; if (selectAll) { - const existingKeys = new Set(sidebarDocs.map((d) => `${d.document_type}:${d.id}`)); - const docsToAdd = subtreeDocs - .filter((d) => !existingKeys.has(`${d.document_type}:${d.id}`)) - .map((d) => ({ - id: d.id, - title: d.title, - document_type: d.document_type as DocumentTypeEnum, - })); setSidebarDocs((prev) => { const existingDocKeys = new Set(prev.map((d) => `${d.document_type}:${d.id}`)); const newDocs = subtreeDocs @@ -937,35 +905,14 @@ function AuthenticatedDocumentsSidebarBase({ })); return newDocs.length > 0 ? [...prev, ...newDocs] : prev; }); - if (docsToAdd.length > 0) { - setSidebarMentionEvent({ - kind: "add", - docs: docsToAdd, - nonce: Date.now(), - }); - } } else { const keysToRemove = new Set(subtreeDocs.map((d) => `${d.document_type}:${d.id}`)); - const docsToRemove = sidebarDocs - .filter((d) => keysToRemove.has(`${d.document_type}:${d.id}`)) - .map((d) => ({ - id: d.id, - title: d.title, - document_type: d.document_type as DocumentTypeEnum, - })); setSidebarDocs((prev) => prev.filter((d) => !keysToRemove.has(`${d.document_type}:${d.id}`)) ); - if (docsToRemove.length > 0) { - setSidebarMentionEvent({ - kind: "remove", - docs: docsToRemove, - nonce: Date.now(), - }); - } } }, - [treeDocuments, foldersByParent, sidebarDocs, setSidebarDocs, setSidebarMentionEvent] + [treeDocuments, foldersByParent, setSidebarDocs] ); const searchFilteredDocuments = useMemo(() => { @@ -1626,7 +1573,6 @@ function AnonymousDocumentsSidebar({ const [search, setSearch] = useState(""); const [sidebarDocs, setSidebarDocs] = useAtom(sidebarSelectedDocumentsAtom); - const setSidebarMentionEvent = useSetAtom(sidebarMentionEventAtom); const mentionedDocIds = useMemo(() => new Set(sidebarDocs.map((d) => d.id)), [sidebarDocs]); const handleToggleChatMention = useCallback( @@ -1634,17 +1580,6 @@ function AnonymousDocumentsSidebar({ const key = `${doc.document_type}:${doc.id}`; if (isMentioned) { setSidebarDocs((prev) => prev.filter((d) => `${d.document_type}:${d.id}` !== key)); - setSidebarMentionEvent({ - kind: "remove", - docs: [ - { - id: doc.id, - title: doc.title, - document_type: doc.document_type as DocumentTypeEnum, - }, - ], - nonce: Date.now(), - }); } else { setSidebarDocs((prev) => { if (prev.some((d) => `${d.document_type}:${d.id}` === key)) return prev; @@ -1653,20 +1588,9 @@ function AnonymousDocumentsSidebar({ { id: doc.id, title: doc.title, document_type: doc.document_type as DocumentTypeEnum }, ]; }); - setSidebarMentionEvent({ - kind: "add", - docs: [ - { - id: doc.id, - title: doc.title, - document_type: doc.document_type as DocumentTypeEnum, - }, - ], - nonce: Date.now(), - }); } }, - [setSidebarDocs, setSidebarMentionEvent] + [setSidebarDocs] ); const uploadedDoc = anonMode.isAnonymous ? anonMode.uploadedDoc : null; From 294c719965f9e83867ec8831994bb0ae67caac29 Mon Sep 17 00:00:00 2001 From: Anish Sarkar <104695310+AnishSarkar22@users.noreply.github.com> Date: Tue, 28 Apr 2026 18:36:49 +0530 Subject: [PATCH 3/6] feat(mentions): implement user message rendering with mention chips for referenced documents --- .../components/assistant-ui/user-message.tsx | 94 +++++++++++++++---- 1 file changed, 78 insertions(+), 16 deletions(-) diff --git a/surfsense_web/components/assistant-ui/user-message.tsx b/surfsense_web/components/assistant-ui/user-message.tsx index 86863a501..fb7212119 100644 --- a/surfsense_web/components/assistant-ui/user-message.tsx +++ b/surfsense_web/components/assistant-ui/user-message.tsx @@ -1,11 +1,12 @@ import { ActionBarPrimitive, AuiIf, MessagePrimitive, useAuiState } from "@assistant-ui/react"; import { useAtomValue } from "jotai"; -import { CheckIcon, CopyIcon, FileText, Pencil } from "lucide-react"; +import { CheckIcon, CopyIcon, Pencil } from "lucide-react"; import Image from "next/image"; import { type FC, useState } from "react"; import { currentThreadAtom } from "@/atoms/chat/current-thread.atom"; import { messageDocumentsMapAtom } from "@/atoms/chat/mentioned-documents.atom"; import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button"; +import { getConnectorIcon } from "@/contracts/enums/connectorIcons"; interface AuthorMetadata { displayName: string | null; @@ -48,6 +49,19 @@ const UserAvatar: FC = ({ displayName, avatarUrl }) => { export const UserMessage: FC = () => { const messageId = useAuiState(({ message }) => message?.id); + const messageText = useAuiState(({ message }) => + (message?.content ?? []) + .map((part) => + typeof part === "object" && + part !== null && + "type" in part && + (part as { type?: string }).type === "text" && + "text" in part + ? String((part as { text?: string }).text ?? "") + : "" + ) + .join("") + ); const messageDocumentsMap = useAtomValue(messageDocumentsMapAtom); const mentionedDocs = messageId ? messageDocumentsMap[messageId] : undefined; const metadata = useAuiState(({ message }) => message?.metadata); @@ -63,22 +77,12 @@ export const UserMessage: FC = () => {
- {mentionedDocs && mentionedDocs.length > 0 && ( -
- {mentionedDocs?.map((doc) => ( - - - {doc.title} - - ))} -
- )}
- + {mentionedDocs && mentionedDocs.length > 0 ? ( + + ) : ( + + )}
@@ -95,6 +99,64 @@ export const UserMessage: FC = () => { ); }; +const UserMessageWithMentionChips: FC<{ + text: string; + mentionedDocs: { id: number; title: string; document_type: string }[]; +}> = ({ text, mentionedDocs }) => { + type Segment = + | { type: "text"; value: string; start: number } + | { type: "mention"; doc: { id: number; title: string; document_type: string }; start: number }; + + const tokens = mentionedDocs + .map((doc) => ({ doc, token: `@${doc.title}` })) + .sort((a, b) => b.token.length - a.token.length); + + const segments: Segment[] = []; + let i = 0; + let buffer = ""; + let bufferStart = 0; + while (i < text.length) { + const tokenMatch = tokens.find(({ token }) => text.startsWith(token, i)); + if (tokenMatch) { + if (buffer) { + segments.push({ type: "text", value: buffer, start: bufferStart }); + buffer = ""; + } + segments.push({ type: "mention", doc: tokenMatch.doc, start: i }); + i += tokenMatch.token.length; + bufferStart = i; + continue; + } + if (!buffer) bufferStart = i; + buffer += text[i]; + i += 1; + } + if (buffer) { + segments.push({ type: "text", value: buffer, start: bufferStart }); + } + + return ( + + {segments.map((segment) => + segment.type === "text" ? ( + {segment.value} + ) : ( + + + {getConnectorIcon(segment.doc.document_type ?? "UNKNOWN", "h-3 w-3")} + + {segment.doc.title} + + ) + )} + + ); +}; + const UserActionBar: FC = () => { const isThreadRunning = useAuiState(({ thread }) => thread.isRunning); From 282510f93ce16bf7b76779046960baa44ed7d9fb Mon Sep 17 00:00:00 2001 From: Anish Sarkar <104695310+AnishSarkar22@users.noreply.github.com> Date: Tue, 28 Apr 2026 18:47:57 +0530 Subject: [PATCH 4/6] feat(mentions): add syncEditorState function to manage editor state and mentioned documents --- .../assistant-ui/inline-mention-editor.tsx | 57 ++++++++++--------- 1 file changed, 31 insertions(+), 26 deletions(-) diff --git a/surfsense_web/components/assistant-ui/inline-mention-editor.tsx b/surfsense_web/components/assistant-ui/inline-mention-editor.tsx index 81d6cbd77..e75a840c0 100644 --- a/surfsense_web/components/assistant-ui/inline-mention-editor.tsx +++ b/surfsense_web/components/assistant-ui/inline-mention-editor.tsx @@ -236,6 +236,19 @@ export const InlineMentionEditor = forwardRef) => { + const docs = docsOverride + ? Array.from(docsOverride.values()) + : Array.from(mentionedDocs.values()); + const text = getText(); + const empty = text.length === 0 && docs.length === 0; + setIsEmpty(empty); + onChange?.(text, docs); + }, + [getText, mentionedDocs, onChange] + ); + // Create a chip element for a document const createChipElement = useCallback( (doc: MentionedDocument): HTMLSpanElement => { @@ -275,6 +288,7 @@ export const InlineMentionEditor = forwardRef { const next = new Map(prev); next.delete(docKey); + syncEditorState(next); return next; }); onDocumentRemove?.(doc.id, doc.document_type); @@ -319,7 +333,7 @@ export const InlineMentionEditor = forwardRef new Map(prev).set(docKey, mentionDoc)); + const nextDocs = new Map(mentionedDocs); + nextDocs.set(docKey, mentionDoc); // Find and remove the @query text const selection = window.getSelection(); @@ -436,25 +452,16 @@ export const InlineMentionEditor = forwardRef { - onChange(getText(), getMentionedDocuments()); - }, 0); - } + syncEditorState(nextDocs); }, [ createChipElement, focusAtEnd, - getText, - getMentionedDocuments, isSelectionInsideEditor, - onChange, + mentionedDocs, rememberSelection, restoreRememberedSelection, + syncEditorState, ] ); @@ -462,22 +469,21 @@ export const InlineMentionEditor = forwardRef { if (editorRef.current) { editorRef.current.innerHTML = ""; - setIsEmpty(true); - setMentionedDocs(new Map()); + const emptyDocs = new Map(); + setMentionedDocs(emptyDocs); + syncEditorState(emptyDocs); } - }, []); + }, [syncEditorState]); // Replace editor content with plain text and place cursor at end const setText = useCallback( (text: string) => { if (!editorRef.current) return; editorRef.current.innerText = text; - const empty = text.length === 0; - setIsEmpty(empty); - onChange?.(text, Array.from(mentionedDocs.values())); + syncEditorState(); focusAtEnd(); }, - [focusAtEnd, onChange, mentionedDocs] + [focusAtEnd, syncEditorState] ); const setDocumentChipStatus = useCallback( @@ -538,14 +544,11 @@ export const InlineMentionEditor = forwardRef { const next = new Map(prev); next.delete(chipKey); + syncEditorState(next); return next; }); - - const text = getText(); - const empty = text.length === 0 && mentionedDocs.size <= 1; - setIsEmpty(empty); }, - [getText, mentionedDocs.size] + [syncEditorState] ); // Expose methods via ref @@ -697,6 +700,7 @@ export const InlineMentionEditor = forwardRef { const next = new Map(prev); next.delete(chipKey); + syncEditorState(next); return next; }); // Notify parent that a document was removed @@ -734,6 +738,7 @@ export const InlineMentionEditor = forwardRef { const next = new Map(prev); next.delete(chipKey); + syncEditorState(next); return next; }); // Notify parent that a document was removed @@ -745,7 +750,7 @@ export const InlineMentionEditor = forwardRef Date: Wed, 29 Apr 2026 04:12:42 +0530 Subject: [PATCH 5/6] refactor(mentions): replace sidebarSelectedDocumentsAtom with mentionedDocumentsAtom and introduce getMentionDocKey utility for consistent document key generation --- .../atoms/chat/mentioned-documents.atom.ts | 23 ------ .../assistant-ui/inline-mention-editor.tsx | 82 ++++++++++++------- .../components/assistant-ui/thread.tsx | 53 ++++++------ .../layout/ui/sidebar/DocumentsSidebar.tsx | 29 ++++--- surfsense_web/lib/chat/mention-doc-key.ts | 8 ++ 5 files changed, 102 insertions(+), 93 deletions(-) create mode 100644 surfsense_web/lib/chat/mention-doc-key.ts diff --git a/surfsense_web/atoms/chat/mentioned-documents.atom.ts b/surfsense_web/atoms/chat/mentioned-documents.atom.ts index 47401995d..9c4546237 100644 --- a/surfsense_web/atoms/chat/mentioned-documents.atom.ts +++ b/surfsense_web/atoms/chat/mentioned-documents.atom.ts @@ -9,29 +9,6 @@ import type { Document } from "@/contracts/types/document.types"; */ export const mentionedDocumentsAtom = atom[]>([]); -/** - * Back-compat alias for sidebar checkbox selection. - * This now points to mentionedDocumentsAtom so the app has a single source - * of truth for mentioned/selected documents. - */ -export const sidebarSelectedDocumentsAtom = atom< - Pick[], - [ - | Pick[] - | (( - prev: Pick[] - ) => Pick[]), - ], - void ->( - (get) => get(mentionedDocumentsAtom), - (get, set, update) => { - const prev = get(mentionedDocumentsAtom); - const next = typeof update === "function" ? update(prev) : update; - set(mentionedDocumentsAtom, next); - } -); - /** * Derived read-only atom that maps deduplicated mentioned docs * into backend payload fields. diff --git a/surfsense_web/components/assistant-ui/inline-mention-editor.tsx b/surfsense_web/components/assistant-ui/inline-mention-editor.tsx index e75a840c0..05277f508 100644 --- a/surfsense_web/components/assistant-ui/inline-mention-editor.tsx +++ b/surfsense_web/components/assistant-ui/inline-mention-editor.tsx @@ -14,6 +14,7 @@ import { import { renderToStaticMarkup } from "react-dom/server"; import { getConnectorIcon } from "@/contracts/enums/connectorIcons"; import type { Document } from "@/contracts/types/document.types"; +import { getMentionDocKey } from "@/lib/chat/mention-doc-key"; import { cn } from "@/lib/utils"; function renderElementToHTML(element: ReactElement): string { @@ -57,7 +58,6 @@ interface InlineMentionEditorProps { onKeyDown?: (e: React.KeyboardEvent) => void; disabled?: boolean; className?: string; - initialDocuments?: MentionedDocument[]; initialText?: string; } @@ -109,7 +109,6 @@ export const InlineMentionEditor = forwardRef(null); const [isEmpty, setIsEmpty] = useState(true); const [mentionedDocs, setMentionedDocs] = useState>( - () => new Map(initialDocuments.map((d) => [`${d.document_type ?? "UNKNOWN"}:${d.id}`, d])) + () => new Map() ); const isComposingRef = useRef(false); const lastSelectionRangeRef = useRef(null); + const isRangeInsideEditor = useCallback((range: Range | null): range is Range => { + if (!range || !editorRef.current) return false; + return ( + editorRef.current.contains(range.startContainer) && + editorRef.current.contains(range.endContainer) + ); + }, []); const isSelectionInsideEditor = useCallback( (selection: Selection | null): selection is Selection => { if (!selection || selection.rangeCount === 0 || !editorRef.current) return false; const range = selection.getRangeAt(0); - return editorRef.current.contains(range.startContainer); + return isRangeInsideEditor(range); }, - [] + [isRangeInsideEditor] ); const rememberSelection = useCallback(() => { @@ -139,11 +145,11 @@ export const InlineMentionEditor = forwardRef { const selection = window.getSelection(); if (!selection) return null; - if (!lastSelectionRangeRef.current) return selection; + if (!isRangeInsideEditor(lastSelectionRangeRef.current)) return null; selection.removeAllRanges(); selection.addRange(lastSelectionRangeRef.current.cloneRange()); return selection; - }, []); + }, [isRangeInsideEditor]); useEffect(() => { const handleSelectionChange = () => { @@ -154,23 +160,13 @@ export const InlineMentionEditor = forwardRef document.removeEventListener("selectionchange", handleSelectionChange); }, [rememberSelection]); - - // Sync initial documents - useEffect(() => { - if (initialDocuments.length > 0) { - setMentionedDocs( - new Map(initialDocuments.map((d) => [`${d.document_type ?? "UNKNOWN"}:${d.id}`, d])) - ); - } - }, [initialDocuments]); - useEffect(() => { if (!initialText || !editorRef.current) return; editorRef.current.innerText = initialText; editorRef.current.appendChild(document.createElement("br")); editorRef.current.appendChild(document.createElement("br")); setIsEmpty(false); - onChange?.(initialText, initialDocuments); + onChange?.(initialText, []); editorRef.current.focus(); const sel = window.getSelection(); const range = document.createRange(); @@ -182,7 +178,7 @@ export const InlineMentionEditor = forwardRef { @@ -284,7 +280,7 @@ export const InlineMentionEditor = forwardRef { const next = new Map(prev); next.delete(docKey); @@ -358,7 +354,7 @@ export const InlineMentionEditor = forwardRef new Map(prev).set(docKey, mentionDoc)); const nextDocs = new Map(mentionedDocs); nextDocs.set(docKey, mentionDoc); @@ -367,12 +363,33 @@ export const InlineMentionEditor = forwardRef { if (!editorRef.current) return; - const chipKey = `${docType ?? "UNKNOWN"}:${docId}`; + const chipKey = getMentionDocKey({ id: docId, document_type: docType }); const chips = editorRef.current.querySelectorAll( `span[${CHIP_DATA_ATTR}="true"]` ); @@ -696,7 +712,10 @@ export const InlineMentionEditor = forwardRef { const next = new Map(prev); next.delete(chipKey); @@ -734,7 +753,10 @@ export const InlineMentionEditor = forwardRef { const next = new Map(prev); next.delete(chipKey); diff --git a/surfsense_web/components/assistant-ui/thread.tsx b/surfsense_web/components/assistant-ui/thread.tsx index dcc068bd1..f9e5ca7fb 100644 --- a/surfsense_web/components/assistant-ui/thread.tsx +++ b/surfsense_web/components/assistant-ui/thread.tsx @@ -87,6 +87,7 @@ import { useBatchCommentsPreload } from "@/hooks/use-comments"; import { useCommentsSync } from "@/hooks/use-comments-sync"; import { useMediaQuery } from "@/hooks/use-media-query"; import { useElectronAPI } from "@/hooks/use-platform"; +import { getMentionDocKey } from "@/lib/chat/mention-doc-key"; import { SLIDEOUT_PANEL_OPENED_EVENT } from "@/lib/layout-events"; import { cn } from "@/lib/utils"; @@ -338,6 +339,9 @@ const Composer: FC = () => { const [mentionQuery, setMentionQuery] = useState(""); const [actionQuery, setActionQuery] = useState(""); const editorRef = useRef(null); + const prevMentionedDocsRef = useRef< + Map> + >(new Map()); const documentPickerRef = useRef(null); const promptPickerRef = useRef(null); const viewportRef = useRef(null); @@ -633,51 +637,50 @@ const Composer: FC = () => { const handleDocumentsMention = useCallback( (documents: Pick[]) => { - const existingKeys = new Set(mentionedDocuments.map((d) => `${d.document_type}:${d.id}`)); - const newDocs = documents.filter( - (doc) => !existingKeys.has(`${doc.document_type}:${doc.id}`) - ); + const editorMentionedDocs = editorRef.current?.getMentionedDocuments() ?? []; + const editorDocKeys = new Set(editorMentionedDocs.map((doc) => getMentionDocKey(doc))); - for (const doc of newDocs) { + for (const doc of documents) { + const key = getMentionDocKey(doc); + if (editorDocKeys.has(key)) continue; editorRef.current?.insertDocumentChip(doc); } setMentionedDocuments((prev) => { - const existingKeySet = new Set(prev.map((d) => `${d.document_type}:${d.id}`)); - const uniqueNewDocs = documents.filter( - (doc) => !existingKeySet.has(`${doc.document_type}:${doc.id}`) - ); + const existingKeySet = new Set(prev.map((d) => getMentionDocKey(d))); + const uniqueNewDocs = documents.filter((doc) => !existingKeySet.has(getMentionDocKey(doc))); return [...prev, ...uniqueNewDocs]; }); setMentionQuery(""); }, - [mentionedDocuments, setMentionedDocuments] + [setMentionedDocuments] ); useEffect(() => { const editor = editorRef.current; - if (!editor) return; + const nextDocsMap = new Map(mentionedDocuments.map((doc) => [getMentionDocKey(doc), doc])); + const prevDocsMap = prevMentionedDocsRef.current; - const toKey = (doc: { id: number; document_type?: string }) => - `${doc.document_type ?? "UNKNOWN"}:${doc.id}`; - - const atomDocs = mentionedDocuments; - const editorDocs = editor.getMentionedDocuments(); - const atomKeys = new Set(atomDocs.map(toKey)); - const editorKeys = new Set(editorDocs.map(toKey)); - - for (const doc of atomDocs) { - if (!editorKeys.has(toKey(doc))) { - editor.insertDocumentChip(doc, { removeTriggerText: false }); - } + if (!editor) { + prevMentionedDocsRef.current = nextDocsMap; + return; } - for (const doc of editorDocs) { - if (!atomKeys.has(toKey(doc))) { + const editorKeys = new Set(editor.getMentionedDocuments().map(getMentionDocKey)); + + for (const [key, doc] of nextDocsMap) { + if (prevDocsMap.has(key) || editorKeys.has(key)) continue; + editor.insertDocumentChip(doc, { removeTriggerText: false }); + } + + for (const [key, doc] of prevDocsMap) { + if (!nextDocsMap.has(key)) { editor.removeDocumentChip(doc.id, doc.document_type); } } + + prevMentionedDocsRef.current = nextDocsMap; }, [mentionedDocuments]); return ( diff --git a/surfsense_web/components/layout/ui/sidebar/DocumentsSidebar.tsx b/surfsense_web/components/layout/ui/sidebar/DocumentsSidebar.tsx index 3c5a64b0e..63b6dc1b7 100644 --- a/surfsense_web/components/layout/ui/sidebar/DocumentsSidebar.tsx +++ b/surfsense_web/components/layout/ui/sidebar/DocumentsSidebar.tsx @@ -24,7 +24,7 @@ import type React from "react"; import { useCallback, useEffect, useMemo, useRef, useState } from "react"; import { toast } from "sonner"; import { - sidebarSelectedDocumentsAtom, + mentionedDocumentsAtom, } from "@/atoms/chat/mentioned-documents.atom"; import { connectorDialogOpenAtom } from "@/atoms/connector-dialog/connector-dialog.atoms"; import { connectorsAtom } from "@/atoms/connectors/connector-query.atoms"; @@ -74,6 +74,7 @@ import type { DocumentTypeEnum } from "@/contracts/types/document.types"; import { useDebouncedValue } from "@/hooks/use-debounced-value"; import { useMediaQuery } from "@/hooks/use-media-query"; import { usePlatform, useElectronAPI } from "@/hooks/use-platform"; +import { getMentionDocKey } from "@/lib/chat/mention-doc-key"; import { anonymousChatApiService } from "@/lib/apis/anonymous-chat-api.service"; import { documentsApiService } from "@/lib/apis/documents-api.service"; import { foldersApiService } from "@/lib/apis/folders-api.service"; @@ -414,7 +415,7 @@ function AuthenticatedDocumentsSidebarBase({ }, [refreshWatchedIds]); const { mutateAsync: deleteDocumentMutation } = useAtomValue(deleteDocumentMutationAtom); - const [sidebarDocs, setSidebarDocs] = useAtom(sidebarSelectedDocumentsAtom); + const [sidebarDocs, setSidebarDocs] = useAtom(mentionedDocumentsAtom); const mentionedDocIds = useMemo(() => new Set(sidebarDocs.map((d) => d.id)), [sidebarDocs]); // Folder state @@ -859,12 +860,12 @@ function AuthenticatedDocumentsSidebarBase({ const handleToggleChatMention = useCallback( (doc: { id: number; title: string; document_type: string }, isMentioned: boolean) => { - const key = `${doc.document_type}:${doc.id}`; + const key = getMentionDocKey(doc); if (isMentioned) { - setSidebarDocs((prev) => prev.filter((d) => `${d.document_type}:${d.id}` !== key)); + setSidebarDocs((prev) => prev.filter((d) => getMentionDocKey(d) !== key)); } else { setSidebarDocs((prev) => { - if (prev.some((d) => `${d.document_type}:${d.id}` === key)) return prev; + if (prev.some((d) => getMentionDocKey(d) === key)) return prev; return [ ...prev, { id: doc.id, title: doc.title, document_type: doc.document_type as DocumentTypeEnum }, @@ -895,9 +896,9 @@ function AuthenticatedDocumentsSidebarBase({ if (selectAll) { setSidebarDocs((prev) => { - const existingDocKeys = new Set(prev.map((d) => `${d.document_type}:${d.id}`)); + const existingDocKeys = new Set(prev.map((d) => getMentionDocKey(d))); const newDocs = subtreeDocs - .filter((d) => !existingDocKeys.has(`${d.document_type}:${d.id}`)) + .filter((d) => !existingDocKeys.has(getMentionDocKey(d))) .map((d) => ({ id: d.id, title: d.title, @@ -906,10 +907,8 @@ function AuthenticatedDocumentsSidebarBase({ return newDocs.length > 0 ? [...prev, ...newDocs] : prev; }); } else { - const keysToRemove = new Set(subtreeDocs.map((d) => `${d.document_type}:${d.id}`)); - setSidebarDocs((prev) => - prev.filter((d) => !keysToRemove.has(`${d.document_type}:${d.id}`)) - ); + const keysToRemove = new Set(subtreeDocs.map((d) => getMentionDocKey(d))); + setSidebarDocs((prev) => prev.filter((d) => !keysToRemove.has(getMentionDocKey(d)))); } }, [treeDocuments, foldersByParent, setSidebarDocs] @@ -1572,17 +1571,17 @@ function AnonymousDocumentsSidebar({ const [isUploading, setIsUploading] = useState(false); const [search, setSearch] = useState(""); - const [sidebarDocs, setSidebarDocs] = useAtom(sidebarSelectedDocumentsAtom); + const [sidebarDocs, setSidebarDocs] = useAtom(mentionedDocumentsAtom); const mentionedDocIds = useMemo(() => new Set(sidebarDocs.map((d) => d.id)), [sidebarDocs]); const handleToggleChatMention = useCallback( (doc: { id: number; title: string; document_type: string }, isMentioned: boolean) => { - const key = `${doc.document_type}:${doc.id}`; + const key = getMentionDocKey(doc); if (isMentioned) { - setSidebarDocs((prev) => prev.filter((d) => `${d.document_type}:${d.id}` !== key)); + setSidebarDocs((prev) => prev.filter((d) => getMentionDocKey(d) !== key)); } else { setSidebarDocs((prev) => { - if (prev.some((d) => `${d.document_type}:${d.id}` === key)) return prev; + if (prev.some((d) => getMentionDocKey(d) === key)) return prev; return [ ...prev, { id: doc.id, title: doc.title, document_type: doc.document_type as DocumentTypeEnum }, diff --git a/surfsense_web/lib/chat/mention-doc-key.ts b/surfsense_web/lib/chat/mention-doc-key.ts new file mode 100644 index 000000000..5dfa11ea3 --- /dev/null +++ b/surfsense_web/lib/chat/mention-doc-key.ts @@ -0,0 +1,8 @@ +type MentionKeyInput = { + id: number; + document_type?: string | null; +}; + +export function getMentionDocKey(doc: MentionKeyInput): string { + return `${doc.document_type ?? "UNKNOWN"}:${doc.id}`; +} 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 6/6] 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")}