refactor: integrate sidebar document selections into chat functionality, enhancing document mention management and improving user experience in the Composer and DocumentsSidebar components

surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx

@@ -21,6 +21,7 @@ import {
 	type MentionedDocumentInfo,
 	mentionedDocumentIdsAtom,
 	mentionedDocumentsAtom,
+	sidebarSelectedDocumentsAtom,
 	messageDocumentsMapAtom,
 } from "@/atoms/chat/mentioned-documents.atom";
 import {
@@ -180,11 +181,13 @@ export default function NewChatPage() {
 		interruptData: Record<string, unknown>;
 	} | null>(null);
 
-	// Get mentioned document IDs from the composer
+	// Get mentioned document IDs from the composer (combines @ mentions + sidebar selections)
 	const mentionedDocumentIds = useAtomValue(mentionedDocumentIdsAtom);
 	const mentionedDocuments = useAtomValue(mentionedDocumentsAtom);
+	const sidebarDocuments = useAtomValue(sidebarSelectedDocumentsAtom);
 	const setMentionedDocumentIds = useSetAtom(mentionedDocumentIdsAtom);
 	const setMentionedDocuments = useSetAtom(mentionedDocumentsAtom);
+	const setSidebarDocuments = useSetAtom(sidebarSelectedDocumentsAtom);
 	const setMessageDocumentsMap = useSetAtom(messageDocumentsMapAtom);
 	const setCurrentThreadState = useSetAtom(currentThreadAtom);
 	const setTargetCommentId = useSetAtom(setTargetCommentIdAtom);
@@ -528,31 +531,30 @@ export default function NewChatPage() {
 				messageLength: userQuery.length,
 			});
 
-			// Store mentioned documents with this message for display
+			// Combine @-mention chips + sidebar selections for display & persistence
-			if (mentionedDocuments.length > 0) {
+			const allMentionedDocs: MentionedDocumentInfo[] = [];
-				const docsInfo: MentionedDocumentInfo[] = mentionedDocuments.map((doc) => ({
+			const seenDocKeys = new Set<string>();
-					id: doc.id,
+			for (const doc of [...mentionedDocuments, ...sidebarDocuments]) {
-					title: doc.title,
+				const key = `${doc.document_type}:${doc.id}`;
-					document_type: doc.document_type,
+				if (!seenDocKeys.has(key)) {
-				}));
+					seenDocKeys.add(key);
+					allMentionedDocs.push({ id: doc.id, title: doc.title, document_type: doc.document_type });
+				}
+			}
+
+			if (allMentionedDocs.length > 0) {
 				setMessageDocumentsMap((prev) => ({
 					...prev,
-					[userMsgId]: docsInfo,
+					[userMsgId]: allMentionedDocs,
 				}));
 			}
 
-			// Persist user message with mentioned documents (don't await, fire and forget)
 			const persistContent: unknown[] = [...message.content];
 
-			// Add mentioned documents for persistence
+			if (allMentionedDocs.length > 0) {
-			if (mentionedDocuments.length > 0) {
 				persistContent.push({
 					type: "mentioned-documents",
-					documents: mentionedDocuments.map((doc) => ({
+					documents: allMentionedDocs,
-						id: doc.id,
-						title: doc.title,
-						document_type: doc.document_type,
-					})),
 				});
 			}
 
@@ -623,6 +625,7 @@ export default function NewChatPage() {
 						document_ids: [],
 					});
 					setMentionedDocuments([]);
+					setSidebarDocuments([]);
 				}
 
 				const response = await fetch(`${backendUrl}/api/v1/new_chat`, {
@@ -920,8 +923,10 @@ export default function NewChatPage() {
 			messages,
 			mentionedDocumentIds,
 			mentionedDocuments,
+			sidebarDocuments,
 			setMentionedDocumentIds,
 			setMentionedDocuments,
+			setSidebarDocuments,
 			setMessageDocumentsMap,
 			queryClient,
 			currentThread,

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

@@ -16,11 +16,17 @@ export const mentionedDocumentIdsAtom = atom<{
 });
 
 /**
- * Atom to store the full document objects mentioned in the current chat composer.
+ * Atom to store the full document objects mentioned via @-mention chips
- * This persists across component remounts.
+ * in the current chat composer. This persists across component remounts.
  */
 export const mentionedDocumentsAtom = atom<Pick<Document, "id" | "title" | "document_type">[]>([]);
 
+/**
+ * Atom to store documents selected via the sidebar checkboxes / row clicks.
+ * These are NOT inserted as chips – the composer shows a count badge instead.
+ */
+export const sidebarSelectedDocumentsAtom = atom<Pick<Document, "id" | "title" | "document_type">[]>([]);
+
 /**
  * Simplified document info for display purposes
  */
@@ -30,22 +36,6 @@ export interface MentionedDocumentInfo {
 	document_type: string;
 }
 
-/**
- * Queue atom for sidebar → composer communication (additions).
- * The sidebar writes documents here; the Composer picks them up,
- * inserts chips, and clears the queue.
- */
-export const pendingDocumentMentionsAtom = atom<
-	Pick<Document, "id" | "title" | "document_type">[]
->([]);
-
-/**
- * Queue atom for sidebar → composer communication (removals).
- * The sidebar writes { id, document_type } here; the Composer removes
- * the matching chips and clears the queue.
- */
-export const pendingDocumentRemovalsAtom = atom<{ id: number; document_type?: string }[]>([]);
-
 /**
  * Atom to store mentioned documents per message ID.
  * This allows displaying which documents were mentioned with each user message.

surfsense_web/components/assistant-ui/inline-mention-editor.tsx

@@ -407,13 +407,12 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
 					next.delete(chipKey);
 					return next;
 				});
-				onDocumentRemove?.(docId, docType);
 
 				const text = getText();
 				const empty = text.length === 0 && mentionedDocs.size <= 1;
 				setIsEmpty(empty);
 			},
-			[getText, mentionedDocs.size, onDocumentRemove]
+			[getText, mentionedDocs.size]
 		);
 
 		// Expose methods via ref

surfsense_web/components/assistant-ui/thread.tsx

@@ -32,8 +32,7 @@ import { documentsSidebarOpenAtom } from "@/atoms/documents/ui.atoms";
 import {
 	mentionedDocumentIdsAtom,
 	mentionedDocumentsAtom,
-	pendingDocumentMentionsAtom,
+	sidebarSelectedDocumentsAtom,
-	pendingDocumentRemovalsAtom,
 } from "@/atoms/chat/mentioned-documents.atom";
 import { membersAtom } from "@/atoms/members/members-query.atoms";
 import {
@@ -231,7 +230,7 @@ const ThreadWelcome: FC = () => {
 const Composer: FC = () => {
 	// Document mention state (atoms persist across component remounts)
 	const [mentionedDocuments, setMentionedDocuments] = useAtom(mentionedDocumentsAtom);
-	const [pendingMentions, setPendingMentions] = useAtom(pendingDocumentMentionsAtom);
+	const [sidebarDocs, setSidebarDocs] = useAtom(sidebarSelectedDocumentsAtom);
 	const [showDocumentPopover, setShowDocumentPopover] = useState(false);
 	const [mentionQuery, setMentionQuery] = useState("");
 	const editorRef = useRef<InlineMentionEditorRef>(null);
@@ -293,7 +292,7 @@ const Composer: FC = () => {
 	const assistantIdsKey = useAssistantState(({ thread }) =>
 		thread.messages
 			.filter((m) => m.role === "assistant" && m.id?.startsWith("msg-"))
-			.map((m) => m.id!.replace("msg-", ""))
+			.map((m) => m.id?.replace("msg-", ""))
 			.join(",")
 	);
 	const assistantDbMessageIds = useMemo(
@@ -313,17 +312,25 @@ const Composer: FC = () => {
 		}
 	}, [isThreadEmpty]);
 
-	// Sync mentioned document IDs to atom for inclusion in chat request payload
+	// Combine sidebar selections + @-mention chips → single ID atom for the backend
 	useEffect(() => {
+		const allDocs = [...mentionedDocuments, ...sidebarDocs];
+		const seen = new Set<string>();
+		const deduped = allDocs.filter((d) => {
+			const key = `${d.document_type}:${d.id}`;
+			if (seen.has(key)) return false;
+			seen.add(key);
+			return true;
+		});
 		setMentionedDocumentIds({
-			surfsense_doc_ids: mentionedDocuments
+			surfsense_doc_ids: deduped
 				.filter((doc) => doc.document_type === "SURFSENSE_DOCS")
 				.map((doc) => doc.id),
-			document_ids: mentionedDocuments
+			document_ids: deduped
 				.filter((doc) => doc.document_type !== "SURFSENSE_DOCS")
 				.map((doc) => doc.id),
 		});
-	}, [mentionedDocuments, setMentionedDocumentIds]);
+	}, [mentionedDocuments, sidebarDocs, setMentionedDocumentIds]);
 
 	// Sync editor text with assistant-ui composer runtime
 	const handleEditorChange = useCallback(
@@ -386,6 +393,7 @@ const Composer: FC = () => {
 			composerRuntime.send();
 			editorRef.current?.clear();
 			setMentionedDocuments([]);
+			setSidebarDocs([]);
 			setMentionedDocumentIds({
 				surfsense_doc_ids: [],
 				document_ids: [],
@@ -397,6 +405,7 @@ const Composer: FC = () => {
 		isBlockedByOtherUser,
 		composerRuntime,
 		setMentionedDocuments,
+		setSidebarDocs,
 		setMentionedDocumentIds,
 	]);
 
@@ -453,40 +462,6 @@ const Composer: FC = () => {
 		[mentionedDocuments, setMentionedDocuments, setMentionedDocumentIds]
 	);
 
-	// Process documents queued from the sidebar (additions)
-	useEffect(() => {
-		if (pendingMentions.length === 0) return;
-		handleDocumentsMention(pendingMentions);
-		setPendingMentions([]);
-	}, [pendingMentions, handleDocumentsMention, setPendingMentions]);
-
-	// Process documents queued from the sidebar (removals)
-	const [pendingRemovals, setPendingRemovals] = useAtom(pendingDocumentRemovalsAtom);
-	useEffect(() => {
-		if (pendingRemovals.length === 0) return;
-		for (const { id, document_type } of pendingRemovals) {
-			editorRef.current?.removeDocumentChip(id, document_type);
-		}
-		setMentionedDocuments((prev) => {
-			const removalKeys = new Set(
-				pendingRemovals.map((r) => `${r.document_type ?? "UNKNOWN"}:${r.id}`)
-			);
-			const updated = prev.filter(
-				(doc) => !removalKeys.has(`${doc.document_type ?? "UNKNOWN"}:${doc.id}`)
-			);
-			setMentionedDocumentIds({
-				surfsense_doc_ids: updated
-					.filter((doc) => doc.document_type === "SURFSENSE_DOCS")
-					.map((doc) => doc.id),
-				document_ids: updated
-					.filter((doc) => doc.document_type !== "SURFSENSE_DOCS")
-					.map((doc) => doc.id),
-			});
-			return updated;
-		});
-		setPendingRemovals([]);
-	}, [pendingRemovals, setPendingRemovals, setMentionedDocuments, setMentionedDocumentIds]);
-
 	return (
 		<ComposerPrimitive.Root className="aui-composer-root relative flex w-full flex-col gap-2">
 			<ChatSessionStatus
@@ -551,6 +526,7 @@ const ComposerAction: FC<ComposerActionProps> = ({
 	isBlockedByOtherUser = false,
 }) => {
 	const mentionedDocuments = useAtomValue(mentionedDocumentsAtom);
+	const sidebarDocs = useAtomValue(sidebarSelectedDocumentsAtom);
 	const setDocumentsSidebarOpen = useSetAtom(documentsSidebarOpenAtom);
 
 	const isComposerTextEmpty = useAssistantState(({ composer }) => {
@@ -603,6 +579,17 @@ const ComposerAction: FC<ComposerActionProps> = ({
 				</div>
 			)}
 
+			<div className="flex items-center gap-2">
+				{sidebarDocs.length > 0 && (
+					<button
+						type="button"
+						onClick={() => setDocumentsSidebarOpen(true)}
+						className="rounded-full border border-border/60 bg-accent/50 px-2.5 py-1 text-xs font-medium text-foreground/80 transition-colors hover:bg-accent"
+					>
+						{sidebarDocs.length} {sidebarDocs.length === 1 ? "source" : "sources"} selected
+					</button>
+				)}
+
 			<AssistantIf condition={({ thread }) => !thread.isRunning}>
 				<ComposerPrimitive.Send asChild disabled={isSendDisabled}>
 					<TooltipIconButton
@@ -644,6 +631,7 @@ const ComposerAction: FC<ComposerActionProps> = ({
 					</Button>
 				</ComposerPrimitive.Cancel>
 			</AssistantIf>
+			</div>
 		</div>
 	);
 };

surfsense_web/components/layout/ui/sidebar/DocumentsSidebar.tsx

@@ -1,16 +1,12 @@
 "use client";
 
-import { useAtomValue, useSetAtom } from "jotai";
+import { useAtom, useAtomValue } from "jotai";
 import { ChevronLeft } from "lucide-react";
 import { useParams } from "next/navigation";
 import { useTranslations } from "next-intl";
 import { useCallback, useEffect, useMemo, useRef, useState } from "react";
 import { toast } from "sonner";
-import {
+import { sidebarSelectedDocumentsAtom } from "@/atoms/chat/mentioned-documents.atom";
-	mentionedDocumentsAtom,
-	pendingDocumentMentionsAtom,
-	pendingDocumentRemovalsAtom,
-} from "@/atoms/chat/mentioned-documents.atom";
 import { deleteDocumentMutationAtom } from "@/atoms/documents/document-mutation.atoms";
 import { Button } from "@/components/ui/button";
 import type { DocumentTypeEnum } from "@/contracts/types/document.types";
@@ -44,26 +40,24 @@ export function DocumentsSidebar({ open, onOpenChange }: DocumentsSidebarProps)
 	const [sortDesc, setSortDesc] = useState(true);
 	const { mutateAsync: deleteDocumentMutation } = useAtomValue(deleteDocumentMutationAtom);
 
-	const mentionedDocuments = useAtomValue(mentionedDocumentsAtom);
+	const [sidebarDocs, setSidebarDocs] = useAtom(sidebarSelectedDocumentsAtom);
-	const setPendingMentions = useSetAtom(pendingDocumentMentionsAtom);
-	const setPendingRemovals = useSetAtom(pendingDocumentRemovalsAtom);
 	const mentionedDocIds = useMemo(
-		() => new Set(mentionedDocuments.map((d) => d.id)),
+		() => new Set(sidebarDocs.map((d) => d.id)),
-		[mentionedDocuments]
+		[sidebarDocs]
 	);
 
 	const handleToggleChatMention = useCallback(
 		(doc: { id: number; title: string; document_type: string }, isMentioned: boolean) => {
 			if (isMentioned) {
-				setPendingRemovals((prev) => [...prev, { id: doc.id, document_type: doc.document_type }]);
+				setSidebarDocs((prev) => prev.filter((d) => d.id !== doc.id));
 			} else {
-				setPendingMentions((prev) => [
+				setSidebarDocs((prev) => {
-					...prev,
+					if (prev.some((d) => d.id === doc.id)) return prev;
-					{ id: doc.id, title: doc.title, document_type: doc.document_type as DocumentTypeEnum },
+					return [...prev, { id: doc.id, title: doc.title, document_type: doc.document_type as DocumentTypeEnum }];
-				]);
+				});
 			}
 		},
-		[setPendingMentions, setPendingRemovals]
+		[setSidebarDocs]
 	);
 
 	const isSearchMode = !!debouncedSearch.trim();