fix(chat): capture mention chips at submit so they aren't dropped

The composer cleared the live mention atom synchronously on send (via the
editor reset), which raced the async onNew handler that read it — dropping
every @-mention (docs, folders, connectors, and the new chat references)
from the request.

handleSubmit now snapshots the chips before clearing, and onNew consumes
that snapshot (falling back to the live atom for the send-button path),
derives the payload via deriveMentionedPayload, and sends mentioned_thread_ids.

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

@@ -7,7 +7,7 @@ import {
 	useExternalStoreRuntime,
 } from "@assistant-ui/react";
 import { useQueryClient } from "@tanstack/react-query";
-import { useAtomValue, useSetAtom } from "jotai";
+import { useAtomValue, useSetAtom, useStore } from "jotai";
 import dynamic from "next/dynamic";
 import { useParams } from "next/navigation";
 import { useCallback, useEffect, useMemo, useRef, useState } from "react";
@@ -23,9 +23,10 @@ import {
 } from "@/atoms/chat/current-thread.atom";
 import {
 	type MentionedDocumentInfo,
-	mentionedDocumentIdsAtom,
+	deriveMentionedPayload,
 	mentionedDocumentsAtom,
 	messageDocumentsMapAtom,
+	submittedMentionsAtom,
 } from "@/atoms/chat/mentioned-documents.atom";
 import { pendingUserImageDataUrlsAtom } from "@/atoms/chat/pending-user-images.atom";
 import {
@@ -206,7 +207,12 @@ const MentionedDocumentInfoSchema = z.object({
 	title: z.string(),
 	document_type: z.string().optional(),
 	kind: z
-		.union([z.literal("doc"), z.literal("folder"), z.literal("connector")])
+		.union([
+			z.literal("doc"),
+			z.literal("folder"),
+			z.literal("connector"),
+			z.literal("thread"),
+		])
 		.optional()
 		.default("doc"),
 	connector_type: z.string().optional(),
@@ -244,6 +250,13 @@ function extractMentionedDocuments(content: unknown): MentionedDocumentInfo[] {
 						kind: "folder",
 					};
 				}
+				if (doc.kind === "thread") {
+					return {
+						id: doc.id,
+						title: doc.title,
+						kind: "thread",
+					};
+				}
 				return {
 					id: doc.id,
 					title: doc.title,
@@ -433,8 +446,7 @@ export default function NewChatPage() {
 	// Get disabled tools from the tool toggle UI
 	const disabledTools = useAtomValue(disabledToolsAtom);
 
-	// Get mentioned document IDs from the composer.
+	const jotaiStore = useStore();
-	const mentionedDocumentIds = useAtomValue(mentionedDocumentIdsAtom);
 	const mentionedDocuments = useAtomValue(mentionedDocumentsAtom);
 	const messageDocumentsMap = useAtomValue(messageDocumentsMapAtom);
 	const setMentionedDocuments = useSetAtom(mentionedDocumentsAtom);
@@ -959,6 +971,16 @@ export default function NewChatPage() {
 				abortControllerRef.current = null;
 			}
 
+			// Prefer the submit-time snapshot; fall back to the live atom
+			// for the send-button path.
+			const submittedSnapshot = jotaiStore.get(submittedMentionsAtom);
+			jotaiStore.set(submittedMentionsAtom, null);
+			const activeMentions = submittedSnapshot ?? mentionedDocuments;
+			const mentionPayload = deriveMentionedPayload(activeMentions);
+			if (activeMentions.length > 0) {
+				setMentionedDocuments([]);
+			}
+
 			const urlsSnapshot = [...pendingUserImageUrls];
 			const { userQuery, userImages } = extractUserTurnForNewChatApi(message, urlsSnapshot);
 
@@ -1060,9 +1082,9 @@ export default function NewChatPage() {
 			trackChatMessageSent(searchSpaceId, currentThreadId, {
 				hasAttachments: userImages.length > 0,
 				hasMentionedDocuments:
-					mentionedDocumentIds.document_ids.length > 0 ||
+					mentionPayload.document_ids.length > 0 ||
-					mentionedDocumentIds.folder_ids.length > 0 ||
+					mentionPayload.folder_ids.length > 0 ||
-					mentionedDocumentIds.connector_ids.length > 0,
+					mentionPayload.connector_ids.length > 0,
 				messageLength: userQuery.length,
 			});
 
@@ -1072,7 +1094,7 @@ export default function NewChatPage() {
 			// can render the correct chip type on reload.
 			const allMentionedDocs: MentionedDocumentInfo[] = [];
 			const seenDocKeys = new Set<string>();
-			for (const doc of mentionedDocuments) {
+			for (const doc of activeMentions) {
 				const key = getMentionDocKey(doc);
 				if (seenDocKeys.has(key)) continue;
 				seenDocKeys.add(key);
@@ -1134,15 +1156,11 @@ export default function NewChatPage() {
 					})
 					.filter((m) => m.content.length > 0);
 
-				// Get mentioned document IDs for context (separate fields for backend)
+				// Backend expects each mention kind in its own payload bucket.
-				const hasDocumentIds = mentionedDocumentIds.document_ids.length > 0;
+				const hasDocumentIds = mentionPayload.document_ids.length > 0;
-				const hasFolderIds = mentionedDocumentIds.folder_ids.length > 0;
+				const hasFolderIds = mentionPayload.folder_ids.length > 0;
-				const hasConnectorIds = mentionedDocumentIds.connector_ids.length > 0;
+				const hasConnectorIds = mentionPayload.connector_ids.length > 0;
-
+				const hasThreadIds = mentionPayload.thread_ids.length > 0;
-				// Clear mentioned documents after capturing them
-				if (hasDocumentIds || hasFolderIds || hasConnectorIds) {
-					setMentionedDocuments([]);
-				}
 
 				const response = await fetchWithTurnCancellingRetry(() =>
 					fetch(buildBackendUrl("/api/v1/new_chat"), {
@@ -1160,18 +1178,16 @@ export default function NewChatPage() {
 							local_filesystem_mounts: selection.local_filesystem_mounts,
 							messages: messageHistory,
 							mentioned_document_ids: hasDocumentIds
-								? mentionedDocumentIds.document_ids
+								? mentionPayload.document_ids
 								: undefined,
-							mentioned_folder_ids: hasFolderIds ? mentionedDocumentIds.folder_ids : undefined,
+							mentioned_folder_ids: hasFolderIds ? mentionPayload.folder_ids : undefined,
 							mentioned_connector_ids: hasConnectorIds
-								? mentionedDocumentIds.connector_ids
+								? mentionPayload.connector_ids
 								: undefined,
-							mentioned_connectors: hasConnectorIds ? mentionedDocumentIds.connectors : undefined,
+							mentioned_connectors: hasConnectorIds ? mentionPayload.connectors : undefined,
-							// Full mention metadata (docs + folders, with
+							mentioned_thread_ids: hasThreadIds ? mentionPayload.thread_ids : undefined,
-							// ``kind`` discriminator) so the BE can embed a
+							// Full mention metadata so the backend can persist a
-							// ``mentioned-documents`` ContentPart on the
+							// ``mentioned-documents`` ContentPart on the user message.
-							// persisted user message (replaces the old FE-side
-							// injection in ``persistUserTurn``).
 							mentioned_documents: allMentionedDocs.length > 0 ? allMentionedDocs : undefined,
 							disabled_tools: disabledTools.length > 0 ? disabledTools : undefined,
 							...(userImages.length > 0 ? { user_images: userImages } : {}),
@@ -1491,7 +1507,7 @@ export default function NewChatPage() {
 			threadId,
 			searchSpaceId,
 			messages,
-			mentionedDocumentIds,
+			jotaiStore,
 			mentionedDocuments,
 			setMentionedDocuments,
 			setMessageDocumentsMap,
@@ -2067,6 +2083,9 @@ export default function NewChatPage() {
 					.filter((d) => d.kind === "folder")
 					.map((d) => d.id);
 				const regenerateConnectors = sourceMentionedDocs.filter((d) => d.kind === "connector");
+				const regenerateThreadIds = sourceMentionedDocs
+					.filter((d) => d.kind === "thread")
+					.map((d) => d.id);
 
 				const requestBody: Record<string, unknown> = {
 					search_space_id: searchSpaceId,
@@ -2080,6 +2099,8 @@ export default function NewChatPage() {
 					mentioned_connector_ids:
 						regenerateConnectors.length > 0 ? regenerateConnectors.map((d) => d.id) : undefined,
 					mentioned_connectors: regenerateConnectors.length > 0 ? regenerateConnectors : undefined,
+					mentioned_thread_ids:
+						regenerateThreadIds.length > 0 ? regenerateThreadIds : undefined,
 					// Full mention metadata for the regenerate-specific
 					// source list. Only meaningful for edit (the BE only
 					// re-persists a user row when ``user_query`` is set);

surfsense_web/components/assistant-ui/thread.tsx

@@ -38,6 +38,7 @@ import { currentThreadAtom } from "@/atoms/chat/current-thread.atom";
 import {
 	type MentionedDocumentInfo,
 	mentionedDocumentsAtom,
+	submittedMentionsAtom,
 } from "@/atoms/chat/mentioned-documents.atom";
 import { pendingUserImageDataUrlsAtom } from "@/atoms/chat/pending-user-images.atom";
 import {
@@ -446,6 +447,7 @@ const ClipboardChip: FC<{ text: string; onDismiss: () => void }> = ({ text, onDi
 
 const Composer: FC = () => {
 	const [mentionedDocuments, setMentionedDocuments] = useAtom(mentionedDocumentsAtom);
+	const setSubmittedMentions = useSetAtom(submittedMentionsAtom);
 	const [showDocumentPopover, setShowDocumentPopover] = useState(false);
 	const [showPromptPicker, setShowPromptPicker] = useState(false);
 	const [mentionQuery, setMentionQuery] = useState("");
@@ -575,6 +577,13 @@ const Composer: FC = () => {
 							kind: "folder",
 						};
 					}
+					if (d.kind === "thread") {
+						return {
+							id: d.id,
+							title: d.title,
+							kind: "thread",
+						};
+					}
 					return {
 						id: d.id,
 						title: d.title,
@@ -770,6 +779,10 @@ const Composer: FC = () => {
 			setClipboardInitialText(undefined);
 		}
 
+		// Capture chips before the reset below clears the live atom, so
+		// the async ``onNew`` still sees them.
+		setSubmittedMentions(mentionedDocuments);
+
 		aui.composer().send();
 		editorRef.current?.clear();
 		setIsComposerInputEmpty(true);
@@ -781,6 +794,8 @@ const Composer: FC = () => {
 		isBlockedByOtherUser,
 		clipboardInitialText,
 		aui,
+		mentionedDocuments,
+		setSubmittedMentions,
 		setMentionedDocuments,
 	]);
 
@@ -788,7 +803,7 @@ const Composer: FC = () => {
 		(
 			docId: number,
 			docType?: string,
-			kind?: "doc" | "folder" | "connector",
+			kind?: "doc" | "folder" | "connector" | "thread",
 			connectorType?: string
 		) => {
 			setMentionedDocuments((prev) => {
@@ -876,6 +891,8 @@ const Composer: FC = () => {
 							<DocumentMentionPicker
 								ref={documentPickerRef}
 								searchSpaceId={Number(search_space_id)}
+								enableChatMentions
+								currentChatId={threadId}
 								onSelectionChange={handleDocumentsMention}
 								onDone={() => {
 									setShowDocumentPopover(false);