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}`; +}