Merge pull request #1439 from AnishSarkar22/fix/mention-documents

feat: improve composer mentions and connector account selection

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

@@ -73,6 +73,7 @@ import {
 	convertToThreadMessage,
 	reconcileInterruptedAssistantMessages,
 } from "@/lib/chat/message-utils";
+import { getMentionDocKey } from "@/lib/chat/mention-doc-key";
 import {
 	isPodcastGenerating,
 	looksLikePodcastRequest,
@@ -206,11 +207,13 @@ function pairBundleToolCallIds(
 const MentionedDocumentInfoSchema = z.object({
 	id: z.number(),
 	title: z.string(),
-	document_type: z.string(),
+	document_type: z.string().optional(),
 	kind: z
-		.union([z.literal("doc"), z.literal("folder")])
+		.union([z.literal("doc"), z.literal("folder"), z.literal("connector")])
 		.optional()
 		.default("doc"),
+	connector_type: z.string().optional(),
+	account_name: z.string().optional(),
 });
 
 const MentionedDocumentsPartSchema = z.object({
@@ -227,7 +230,30 @@ function extractMentionedDocuments(content: unknown): MentionedDocumentInfo[] {
 	for (const part of content) {
 		const result = MentionedDocumentsPartSchema.safeParse(part);
 		if (result.success) {
-			return result.data.documents;
+			return result.data.documents.map<MentionedDocumentInfo>((doc) => {
+				if (doc.kind === "connector") {
+					return {
+						id: doc.id,
+						title: doc.title,
+						kind: "connector",
+						connector_type: doc.connector_type ?? doc.document_type ?? "UNKNOWN",
+						account_name: doc.account_name ?? doc.title,
+					};
+				}
+				if (doc.kind === "folder") {
+					return {
+						id: doc.id,
+						title: doc.title,
+						kind: "folder",
+					};
+				}
+				return {
+					id: doc.id,
+					title: doc.title,
+					document_type: doc.document_type ?? "UNKNOWN",
+					kind: "doc",
+				};
+			});
 		}
 	}
 
@@ -924,28 +950,22 @@ export default function NewChatPage() {
 				hasMentionedDocuments:
 					mentionedDocumentIds.surfsense_doc_ids.length > 0 ||
 					mentionedDocumentIds.document_ids.length > 0 ||
-					mentionedDocumentIds.folder_ids.length > 0,
+					mentionedDocumentIds.folder_ids.length > 0 ||
+					mentionedDocumentIds.connector_ids.length > 0,
 				messageLength: userQuery.length,
 			});
 
 			// Collect unique mention chips for display & persistence.
-			// Dedup key is ``kind:document_type:id`` so a folder and a
-			// doc with the same integer id never collapse into one
-			// entry. The ``kind`` field is forwarded to the backend
+			// The ``kind`` field is forwarded to the backend
 			// so the persisted ``mentioned-documents`` content part
 			// can render the correct chip type on reload.
 			const allMentionedDocs: MentionedDocumentInfo[] = [];
 			const seenDocKeys = new Set<string>();
 			for (const doc of mentionedDocuments) {
-				const key = `${doc.kind}:${doc.document_type}:${doc.id}`;
+				const key = getMentionDocKey(doc);
 				if (seenDocKeys.has(key)) continue;
 				seenDocKeys.add(key);
-				allMentionedDocs.push({
-					id: doc.id,
-					title: doc.title,
-					document_type: doc.document_type,
-					kind: doc.kind,
-				});
+				allMentionedDocs.push(doc);
 			}
 
 			if (allMentionedDocs.length > 0) {
@@ -1008,9 +1028,10 @@ export default function NewChatPage() {
 				const hasDocumentIds = mentionedDocumentIds.document_ids.length > 0;
 				const hasSurfsenseDocIds = mentionedDocumentIds.surfsense_doc_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 || hasSurfsenseDocIds || hasFolderIds) {
+				if (hasDocumentIds || hasSurfsenseDocIds || hasFolderIds || hasConnectorIds) {
 					setMentionedDocuments([]);
 				}
 
@@ -1036,20 +1057,16 @@ export default function NewChatPage() {
 								? mentionedDocumentIds.surfsense_doc_ids
 								: undefined,
 							mentioned_folder_ids: hasFolderIds ? mentionedDocumentIds.folder_ids : undefined,
+							mentioned_connector_ids: hasConnectorIds
+								? mentionedDocumentIds.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_documents:
-								allMentionedDocs.length > 0
-									? allMentionedDocs.map((d) => ({
-											id: d.id,
-											title: d.title,
-											document_type: d.document_type,
-											kind: d.kind,
-										}))
-									: undefined,
+							mentioned_documents: allMentionedDocs.length > 0 ? allMentionedDocs : undefined,
 							disabled_tools: disabledTools.length > 0 ? disabledTools : undefined,
 							...(userImages.length > 0 ? { user_images: userImages } : {}),
 						}),
@@ -1945,6 +1962,7 @@ export default function NewChatPage() {
 				const regenerateFolderIds = sourceMentionedDocs
 					.filter((d) => d.kind === "folder")
 					.map((d) => d.id);
+				const regenerateConnectors = sourceMentionedDocs.filter((d) => d.kind === "connector");
 
 				const requestBody: Record<string, unknown> = {
 					search_space_id: searchSpaceId,
@@ -1957,19 +1975,16 @@ export default function NewChatPage() {
 					mentioned_surfsense_doc_ids:
 						regenerateSurfsenseDocIds.length > 0 ? regenerateSurfsenseDocIds : undefined,
 					mentioned_folder_ids: regenerateFolderIds.length > 0 ? regenerateFolderIds : undefined,
+					mentioned_connector_ids:
+						regenerateConnectors.length > 0 ? regenerateConnectors.map((d) => d.id) : undefined,
+					mentioned_connectors:
+						regenerateConnectors.length > 0 ? regenerateConnectors : 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);
 					// reload reuses the original turn's mentioned_documents.
 					mentioned_documents:
-						sourceMentionedDocs.length > 0
-							? sourceMentionedDocs.map((d) => ({
-									id: d.id,
-									title: d.title,
-									document_type: d.document_type,
-									kind: d.kind,
-								}))
-							: undefined,
+						sourceMentionedDocs.length > 0 ? sourceMentionedDocs : undefined,
 				};
 				if (isEdit) {
 					requestBody.user_images = editExtras?.userImages ?? [];