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 1df9ef06c..6e14cb8e6 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 @@ -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; } | 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 - if (mentionedDocuments.length > 0) { - const docsInfo: MentionedDocumentInfo[] = mentionedDocuments.map((doc) => ({ - id: doc.id, - title: doc.title, - document_type: doc.document_type, - })); + // Combine @-mention chips + sidebar selections for display & persistence + const allMentionedDocs: MentionedDocumentInfo[] = []; + const seenDocKeys = new Set(); + for (const doc of [...mentionedDocuments, ...sidebarDocuments]) { + 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 (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 (mentionedDocuments.length > 0) { + if (allMentionedDocs.length > 0) { persistContent.push({ type: "mentioned-documents", - documents: mentionedDocuments.map((doc) => ({ - id: doc.id, - title: doc.title, - document_type: doc.document_type, - })), + documents: allMentionedDocs, }); } @@ -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, diff --git a/surfsense_web/atoms/chat/mentioned-documents.atom.ts b/surfsense_web/atoms/chat/mentioned-documents.atom.ts index 4487e1732..04d414046 100644 --- a/surfsense_web/atoms/chat/mentioned-documents.atom.ts +++ b/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. - * This persists across component remounts. + * Atom to store the full document objects mentioned via @-mention chips + * in the current chat composer. This persists across component remounts. */ export const mentionedDocumentsAtom = atom[]>([]); +/** + * 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[]>([]); + /** * 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[] ->([]); - -/** - * 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. diff --git a/surfsense_web/components/assistant-ui/inline-mention-editor.tsx b/surfsense_web/components/assistant-ui/inline-mention-editor.tsx index cc3ea44af..8890a9edf 100644 --- a/surfsense_web/components/assistant-ui/inline-mention-editor.tsx +++ b/surfsense_web/components/assistant-ui/inline-mention-editor.tsx @@ -407,13 +407,12 @@ export const InlineMentionEditor = forwardRef { 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(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(); + 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 ( = ({ 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 = ({ )} +
+ {sidebarDocs.length > 0 && ( + + )} + !thread.isRunning}> = ({ +
); }; diff --git a/surfsense_web/components/layout/ui/sidebar/DocumentsSidebar.tsx b/surfsense_web/components/layout/ui/sidebar/DocumentsSidebar.tsx index cb3abf0e3..aa66ff424 100644 --- a/surfsense_web/components/layout/ui/sidebar/DocumentsSidebar.tsx +++ b/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 { - mentionedDocumentsAtom, - pendingDocumentMentionsAtom, - pendingDocumentRemovalsAtom, -} from "@/atoms/chat/mentioned-documents.atom"; +import { sidebarSelectedDocumentsAtom } 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 setPendingMentions = useSetAtom(pendingDocumentMentionsAtom); - const setPendingRemovals = useSetAtom(pendingDocumentRemovalsAtom); + const [sidebarDocs, setSidebarDocs] = useAtom(sidebarSelectedDocumentsAtom); const mentionedDocIds = useMemo( - () => new Set(mentionedDocuments.map((d) => d.id)), - [mentionedDocuments] + () => new Set(sidebarDocs.map((d) => d.id)), + [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) => [ - ...prev, - { id: doc.id, title: doc.title, document_type: doc.document_type as DocumentTypeEnum }, - ]); + setSidebarDocs((prev) => { + if (prev.some((d) => d.id === doc.id)) return prev; + return [...prev, { id: doc.id, title: doc.title, document_type: doc.document_type as DocumentTypeEnum }]; + }); } }, - [setPendingMentions, setPendingRemovals] + [setSidebarDocs] ); const isSearchMode = !!debouncedSearch.trim();