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 3594e15eb..46fa6d320 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 @@ -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 mentionedDocumentIds = useAtomValue(mentionedDocumentIdsAtom); + const jotaiStore = useStore(); 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 || - mentionedDocumentIds.folder_ids.length > 0 || - mentionedDocumentIds.connector_ids.length > 0, + mentionPayload.document_ids.length > 0 || + mentionPayload.folder_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(); - 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) - const hasDocumentIds = mentionedDocumentIds.document_ids.length > 0; - const hasFolderIds = mentionedDocumentIds.folder_ids.length > 0; - const hasConnectorIds = mentionedDocumentIds.connector_ids.length > 0; - - // Clear mentioned documents after capturing them - if (hasDocumentIds || hasFolderIds || hasConnectorIds) { - setMentionedDocuments([]); - } + // Backend expects each mention kind in its own payload bucket. + const hasDocumentIds = mentionPayload.document_ids.length > 0; + const hasFolderIds = mentionPayload.folder_ids.length > 0; + const hasConnectorIds = mentionPayload.connector_ids.length > 0; + const hasThreadIds = mentionPayload.thread_ids.length > 0; 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, - // Full mention metadata (docs + folders, with - // ``kind`` discriminator) so the BE can embed a - // ``mentioned-documents`` ContentPart on the - // persisted user message (replaces the old FE-side - // injection in ``persistUserTurn``). + mentioned_connectors: hasConnectorIds ? mentionPayload.connectors : undefined, + mentioned_thread_ids: hasThreadIds ? mentionPayload.thread_ids : undefined, + // Full mention metadata so the backend can persist a + // ``mentioned-documents`` ContentPart on the user message. 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 = { 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); diff --git a/surfsense_web/components/assistant-ui/thread.tsx b/surfsense_web/components/assistant-ui/thread.tsx index c8da125f4..067c641c6 100644 --- a/surfsense_web/components/assistant-ui/thread.tsx +++ b/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 = () => { { setShowDocumentPopover(false);