refactor(mentions): replace sidebarSelectedDocumentsAtom with mentionedDocumentsAtom and introduce getMentionDocKey utility for consistent document key generation

surfsense_web/atoms/chat/mentioned-documents.atom.ts

@@ -9,29 +9,6 @@ import type { Document } from "@/contracts/types/document.types";
  */
 export const mentionedDocumentsAtom = atom<Pick<Document, "id" | "title" | "document_type">[]>([]);
 
-/**
- * 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<Document, "id" | "title" | "document_type">[],
-	[
-		| Pick<Document, "id" | "title" | "document_type">[]
-		| ((
-				prev: Pick<Document, "id" | "title" | "document_type">[]
-		  ) => Pick<Document, "id" | "title" | "document_type">[]),
-	],
-	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.

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<InlineMentionEditorRef, InlineMent
 			onKeyDown,
 			disabled = false,
 			className,
-			initialDocuments = [],
 			initialText,
 		},
 		ref
@@ -117,17 +116,24 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
 		const editorRef = useRef<HTMLDivElement>(null);
 		const [isEmpty, setIsEmpty] = useState(true);
 		const [mentionedDocs, setMentionedDocs] = useState<Map<string, MentionedDocument>>(
-			() => new Map(initialDocuments.map((d) => [`${d.document_type ?? "UNKNOWN"}:${d.id}`, d]))
+			() => new Map()
 		);
 		const isComposingRef = useRef(false);
 		const lastSelectionRangeRef = useRef<Range | null>(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<InlineMentionEditorRef, InlineMent
 		const restoreRememberedSelection = useCallback((): Selection | null => {
 			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<InlineMentionEditorRef, InlineMent
 			return () => 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<InlineMentionEditorRef, InlineMent
 			range.insertNode(anchor);
 			anchor.scrollIntoView({ block: "end" });
 			anchor.remove();
-		}, [initialText, initialDocuments, onChange]);
+		}, [initialText, onChange]);
 
 		// Focus at the end of the editor
 		const focusAtEnd = useCallback(() => {
@@ -284,7 +280,7 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
 					e.preventDefault();
 					e.stopPropagation();
 					chip.remove();
-					const docKey = `${doc.document_type ?? "UNKNOWN"}:${doc.id}`;
+					const docKey = getMentionDocKey(doc);
 					setMentionedDocs((prev) => {
 						const next = new Map(prev);
 						next.delete(docKey);
@@ -358,7 +354,7 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
 				};
 
 				// Add to mentioned docs map using unique key
-				const docKey = `${doc.document_type ?? "UNKNOWN"}:${doc.id}`;
+				const docKey = getMentionDocKey(doc);
 				setMentionedDocs((prev) => new Map(prev).set(docKey, mentionDoc));
 				const nextDocs = new Map(mentionedDocs);
 				nextDocs.set(docKey, mentionDoc);
@@ -367,12 +363,33 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
 				const selection = window.getSelection();
 				const hasActiveSelection = isSelectionInsideEditor(selection);
 				const resolvedSelection = hasActiveSelection ? selection : restoreRememberedSelection();
-				if (!resolvedSelection || resolvedSelection.rangeCount === 0) {
+				if (
-					// No selection, just append
+					!resolvedSelection ||
+					resolvedSelection.rangeCount === 0 ||
+					!isSelectionInsideEditor(resolvedSelection)
+				) {
+					// No valid in-editor selection: deterministically insert at end.
+					editorRef.current.focus();
+					const endSelection = window.getSelection();
+					if (!endSelection) return;
+					const endRange = document.createRange();
+					endRange.selectNodeContents(editorRef.current);
+					endRange.collapse(false);
+					endSelection.removeAllRanges();
+					endSelection.addRange(endRange);
+
 					const chip = createChipElement(mentionDoc);
-					editorRef.current.appendChild(chip);
+					endRange.insertNode(chip);
-					editorRef.current.appendChild(document.createTextNode(" "));
+					endRange.setStartAfter(chip);
-					focusAtEnd();
+					endRange.collapse(true);
+					const space = document.createTextNode(" ");
+					endRange.insertNode(space);
+					endRange.setStartAfter(space);
+					endRange.collapse(true);
+					endSelection.removeAllRanges();
+					endSelection.addRange(endRange);
+
+					syncEditorState(nextDocs);
 					rememberSelection();
 					return;
 				}
@@ -456,7 +473,6 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
 			},
 			[
 				createChipElement,
-				focusAtEnd,
 				isSelectionInsideEditor,
 				mentionedDocs,
 				rememberSelection,
@@ -531,7 +547,7 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
 		const removeDocumentChip = useCallback(
 			(docId: number, docType?: string) => {
 				if (!editorRef.current) return;
-				const chipKey = `${docType ?? "UNKNOWN"}:${docId}`;
+				const chipKey = getMentionDocKey({ id: docId, document_type: docType });
 				const chips = editorRef.current.querySelectorAll<HTMLSpanElement>(
 					`span[${CHIP_DATA_ATTR}="true"]`
 				);
@@ -696,7 +712,10 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
 									const chipDocType = getChipDocType(prevSibling);
 									if (chipId !== null) {
 										prevSibling.remove();
-										const chipKey = `${chipDocType}:${chipId}`;
+										const chipKey = getMentionDocKey({
+											id: chipId,
+											document_type: chipDocType,
+										});
 										setMentionedDocs((prev) => {
 											const next = new Map(prev);
 											next.delete(chipKey);
@@ -734,7 +753,10 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
 									const chipDocType = getChipDocType(prevChild);
 									if (chipId !== null) {
 										prevChild.remove();
-										const chipKey = `${chipDocType}:${chipId}`;
+										const chipKey = getMentionDocKey({
+											id: chipId,
+											document_type: chipDocType,
+										});
 										setMentionedDocs((prev) => {
 											const next = new Map(prev);
 											next.delete(chipKey);

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<InlineMentionEditorRef>(null);
+	const prevMentionedDocsRef = useRef<
+		Map<string, Pick<Document, "id" | "title" | "document_type">>
+	>(new Map());
 	const documentPickerRef = useRef<DocumentMentionPickerRef>(null);
 	const promptPickerRef = useRef<PromptPickerRef>(null);
 	const viewportRef = useRef<Element | null>(null);
@@ -633,51 +637,50 @@ const Composer: FC = () => {
 
 	const handleDocumentsMention = useCallback(
 		(documents: Pick<Document, "id" | "title" | "document_type">[]) => {
-			const existingKeys = new Set(mentionedDocuments.map((d) => `${d.document_type}:${d.id}`));
+			const editorMentionedDocs = editorRef.current?.getMentionedDocuments() ?? [];
-			const newDocs = documents.filter(
+			const editorDocKeys = new Set(editorMentionedDocs.map((doc) => getMentionDocKey(doc)));
-				(doc) => !existingKeys.has(`${doc.document_type}:${doc.id}`)
-			);
 
-			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 existingKeySet = new Set(prev.map((d) => getMentionDocKey(d)));
-				const uniqueNewDocs = documents.filter(
+				const uniqueNewDocs = documents.filter((doc) => !existingKeySet.has(getMentionDocKey(doc)));
-					(doc) => !existingKeySet.has(`${doc.document_type}:${doc.id}`)
-				);
 				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 }) =>
+		if (!editor) {
-			`${doc.document_type ?? "UNKNOWN"}:${doc.id}`;
+			prevMentionedDocsRef.current = nextDocsMap;
-
+			return;
-		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 });
-			}
 		}
 
-		for (const doc of editorDocs) {
+		const editorKeys = new Set(editor.getMentionedDocuments().map(getMentionDocKey));
-			if (!atomKeys.has(toKey(doc))) {
+
+		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 (

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}`));
+				const keysToRemove = new Set(subtreeDocs.map((d) => getMentionDocKey(d)));
-				setSidebarDocs((prev) =>
+				setSidebarDocs((prev) => prev.filter((d) => !keysToRemove.has(getMentionDocKey(d))));
-					prev.filter((d) => !keysToRemove.has(`${d.document_type}:${d.id}`))
-				);
 			}
 		},
 		[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 },

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