feat: implement document mention extraction and management in new chat

- Added functionality to extract and manage mentioned documents within the new chat interface.
- Introduced new atoms for storing mentioned documents and their mappings to user messages.
- Enhanced the message persistence logic to include mentioned documents, ensuring they are displayed correctly in the chat.
- Updated the UI components to support document mentions, including a refined document selection interface.
- Improved state management for document mentions to ensure a seamless user experience.

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

@@ -10,7 +10,7 @@ import { useAtomValue, useSetAtom } from "jotai";
 import { useParams, useRouter } from "next/navigation";
 import { useCallback, useEffect, useMemo, useRef, useState } from "react";
 import { toast } from "sonner";
-import { mentionedDocumentIdsAtom } from "@/atoms/chat/mentioned-documents.atom";
+import { mentionedDocumentIdsAtom, mentionedDocumentsAtom, messageDocumentsMapAtom, type MentionedDocumentInfo } from "@/atoms/chat/mentioned-documents.atom";
 import { Thread } from "@/components/assistant-ui/thread";
 import { GeneratePodcastToolUI } from "@/components/tool-ui/generate-podcast";
 import { LinkPreviewToolUI } from "@/components/tool-ui/link-preview";
@@ -48,6 +48,23 @@ function extractThinkingSteps(content: unknown): ThinkingStep[] {
 	return thinkingPart?.steps || [];
 }
 
+/**
+ * Extract mentioned documents from message content
+ */
+function extractMentionedDocuments(content: unknown): MentionedDocumentInfo[] {
+	if (!Array.isArray(content)) return [];
+	
+	const docsPart = content.find(
+		(part: unknown) => 
+			typeof part === "object" && 
+			part !== null && 
+			"type" in part && 
+			(part as { type: string }).type === "mentioned-documents"
+	) as { type: "mentioned-documents"; documents: MentionedDocumentInfo[] } | undefined;
+	
+	return docsPart?.documents || [];
+}
+
 /**
  * Convert backend message to assistant-ui ThreadMessageLike format
  * Filters out 'thinking-steps' part as it's handled separately
@@ -58,13 +75,14 @@ function convertToThreadMessage(msg: MessageRecord): ThreadMessageLike {
 	if (typeof msg.content === "string") {
 		content = [{ type: "text", text: msg.content }];
 	} else if (Array.isArray(msg.content)) {
-		// Filter out thinking-steps part - it's handled separately via messageThinkingSteps
+		// Filter out custom metadata parts - they're handled separately
 		const filteredContent = msg.content.filter(
-			(part: unknown) => 
-				!(typeof part === "object" && 
-				  part !== null && 
-				  "type" in part && 
-				  (part as { type: string }).type === "thinking-steps")
+			(part: unknown) => {
+				if (typeof part !== "object" || part === null || !("type" in part)) return true;
+				const partType = (part as { type: string }).type;
+				// Filter out thinking-steps and mentioned-documents
+				return partType !== "thinking-steps" && partType !== "mentioned-documents";
+			}
 		);
 		content = filteredContent.length > 0 
 			? (filteredContent as ThreadMessageLike["content"])
@@ -111,7 +129,10 @@ export default function NewChatPage() {
 
 	// Get mentioned document IDs from the composer
 	const mentionedDocumentIds = useAtomValue(mentionedDocumentIdsAtom);
+	const mentionedDocuments = useAtomValue(mentionedDocumentsAtom);
 	const setMentionedDocumentIds = useSetAtom(mentionedDocumentIdsAtom);
+	const setMentionedDocuments = useSetAtom(mentionedDocumentsAtom);
+	const setMessageDocumentsMap = useSetAtom(messageDocumentsMapAtom);
 
 	// Create the attachment adapter for file processing
 	const attachmentAdapter = useMemo(() => createAttachmentAdapter(), []);
@@ -150,6 +171,9 @@ export default function NewChatPage() {
 					
 					// Extract and restore thinking steps from persisted messages
 					const restoredThinkingSteps = new Map<string, ThinkingStep[]>();
+					// Extract and restore mentioned documents from persisted messages
+					const restoredDocsMap: Record<string, MentionedDocumentInfo[]> = {};
+					
 					for (const msg of response.messages) {
 						if (msg.role === "assistant") {
 							const steps = extractThinkingSteps(msg.content);
@@ -157,10 +181,19 @@ export default function NewChatPage() {
 								restoredThinkingSteps.set(`msg-${msg.id}`, steps);
 							}
 						}
+						if (msg.role === "user") {
+							const docs = extractMentionedDocuments(msg.content);
+							if (docs.length > 0) {
+								restoredDocsMap[`msg-${msg.id}`] = docs;
+							}
+						}
 					}
 					if (restoredThinkingSteps.size > 0) {
 						setMessageThinkingSteps(restoredThinkingSteps);
 					}
+					if (Object.keys(restoredDocsMap).length > 0) {
+						setMessageDocumentsMap(restoredDocsMap);
+					}
 				}
 			} else {
 				// Create new thread
@@ -239,10 +272,33 @@ export default function NewChatPage() {
 			};
 			setMessages((prev) => [...prev, userMessage]);
 
-			// Persist user message (don't await, fire and forget)
+			// 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,
+				}));
+				setMessageDocumentsMap((prev) => ({
+					...prev,
+					[userMsgId]: docsInfo,
+				}));
+			}
+
+			// Persist user message with mentioned documents (don't await, fire and forget)
+			const persistContent = mentionedDocuments.length > 0
+				? [
+					...message.content,
+					{ type: "mentioned-documents", documents: mentionedDocuments.map((doc) => ({
+						id: doc.id,
+						title: doc.title,
+						document_type: doc.document_type,
+					})) },
+				]
+				: message.content;
 			appendMessage(threadId, {
 				role: "user",
-				content: message.content,
+				content: persistContent,
 			}).catch((err) => console.error("Failed to persist user message:", err));
 
 			// Start streaming response
@@ -384,6 +440,7 @@ export default function NewChatPage() {
 				// Clear mentioned documents after capturing them
 				if (mentionedDocumentIds.length > 0) {
 					setMentionedDocumentIds([]);
+					setMentionedDocuments([]);
 				}
 
 				const response = await fetch(`${backendUrl}/api/v1/new_chat`, {
@@ -561,7 +618,7 @@ export default function NewChatPage() {
 				// Note: We no longer clear thinking steps - they persist with the message
 			}
 		},
-		[threadId, searchSpaceId, messages, mentionedDocumentIds, setMentionedDocumentIds]
+		[threadId, searchSpaceId, messages, mentionedDocumentIds, mentionedDocuments, setMentionedDocumentIds, setMentionedDocuments, setMessageDocumentsMap]
 	);
 
 	// Convert message (pass through since already in correct format)