feat(web): enhance mention handling to support connectors and improve document key management

surfsense_web/components/assistant-ui/inline-mention-editor.tsx

@@ -20,7 +20,6 @@ import {
 	useMemo,
 	useRef,
 } from "react";
-import { FOLDER_MENTION_DOCUMENT_TYPE } from "@/atoms/chat/mentioned-documents.atom";
 import { Button } from "@/components/ui/button";
 import { getConnectorIcon } from "@/contracts/enums/connectorIcons";
 import type { Document } from "@/contracts/types/document.types";
@@ -40,8 +39,6 @@ export interface MentionedDocument {
 
 /**
  * Input shape for inserting a chip. ``kind`` defaults to ``"doc"``.
- * Folder chips default ``document_type`` to ``FOLDER_MENTION_DOCUMENT_TYPE``
- * so the dedup key never collides with a doc chip sharing the same id.
  */
 export type MentionChipInput = {
 	id: number;
@@ -78,7 +75,12 @@ export interface InlineMentionEditorRef {
 		doc: Pick<Document, "id" | "title" | "document_type">,
 		options?: { removeTriggerText?: boolean }
 	) => void;
-	removeDocumentChip: (docId: number, docType?: string) => void;
+	removeDocumentChip: (
+		docId: number,
+		docType?: string,
+		kind?: MentionKind,
+		connectorType?: string
+	) => void;
 	setDocumentChipStatus: (
 		docId: number,
 		docType: string | undefined,
@@ -95,7 +97,7 @@ interface InlineMentionEditorProps {
 	onActionClose?: () => void;
 	onSubmit?: () => void;
 	onChange?: (text: string, docs: MentionedDocument[]) => void;
-	onDocumentRemove?: (docId: number, docType?: string) => void;
+	onDocumentRemove?: (docId: number, docType?: string, kind?: MentionKind, connectorType?: string) => void;
 	onKeyDown?: (e: React.KeyboardEvent) => void;
 	disabled?: boolean;
 	className?: string;
@@ -135,7 +137,12 @@ const EMPTY_VALUE: ComposerValue = [{ type: "p", children: [{ text: "" }] }];
  * the X button and Backspace go through the same call site.
  */
 type MentionEditorContextValue = {
-	removeChip: (docId: number, docType: string | undefined) => void;
+	removeChip: (
+		docId: number,
+		docType: string | undefined,
+		kind: MentionKind | undefined,
+		connectorType: string | undefined
+	) => void;
 };
 const MentionEditorContext = createContext<MentionEditorContextValue | null>(null);
 
@@ -181,7 +188,12 @@ const MentionElement: FC<PlateElementProps<MentionElementNode>> = ({
 								onMouseDown={(e) => e.preventDefault()}
 								onClick={(e) => {
 									e.stopPropagation();
-									ctx.removeChip(element.id, element.document_type);
+									ctx.removeChip(
+										element.id,
+										element.document_type,
+										element.kind,
+										element.connector_type
+									);
 								}}
 								className="absolute inset-0 size-3 rounded-sm p-0 opacity-0 transition-opacity hover:bg-transparent hover:text-primary focus-visible:opacity-100 focus-visible:outline-none focus-visible:ring-0 group-hover:opacity-100 [&_svg]:size-3"
 							>
@@ -456,18 +468,11 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
 
 				const removeTriggerText = options?.removeTriggerText ?? true;
 				const kind: MentionKind = mention.kind ?? "doc";
-				const document_type =
-					mention.document_type ??
-					(kind === "folder"
-						? FOLDER_MENTION_DOCUMENT_TYPE
-						: kind === "connector"
-							? mention.connector_type
-							: undefined);
 				const mentionNode: MentionElementNode = {
 					type: MENTION_TYPE,
 					id: mention.id,
 					title: mention.title,
-					document_type,
+					document_type: mention.document_type,
 					kind,
 					connector_type: mention.connector_type,
 					account_name: mention.account_name,
@@ -526,17 +531,33 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
 			[insertMentionChip]
 		);
 
-		// Remove chip(s) matching (id, document_type). Iterates in
+		// Remove chip(s) matching the mention identity. Iterates in
 		// descending path order so removing one entry can't invalidate
 		// later paths. Chips are deduped today, so this typically runs
 		// at most once.
 		const removeDocumentChip = useCallback(
-			(docId: number, docType?: string) => {
+			(docId: number, docType?: string, kind?: MentionKind, connectorType?: string) => {
 				const match = (n: unknown) => {
 					if (!n || typeof n !== "object" || !("type" in n)) return false;
 					const node = n as MentionElementNode;
 					if (node.type !== MENTION_TYPE) return false;
 					if (node.id !== docId) return false;
+					if (kind) {
+						return (
+							getMentionDocKey({
+								id: node.id,
+								kind: node.kind ?? "doc",
+								document_type: node.document_type,
+								connector_type: node.connector_type,
+							}) ===
+							getMentionDocKey({
+								id: docId,
+								kind,
+								document_type: docType,
+								connector_type: connectorType,
+							})
+						);
+					}
 					return (node.document_type ?? "UNKNOWN") === (docType ?? "UNKNOWN");
 				};
 
@@ -554,9 +575,14 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
 		// Single removal call site for Backspace and the X button so the
 		// two can never diverge (e.g. one forgetting to notify the parent).
 		const removeChip = useCallback(
-			(docId: number, docType: string | undefined) => {
-				removeDocumentChip(docId, docType);
-				onDocumentRemove?.(docId, docType);
+			(
+				docId: number,
+				docType: string | undefined,
+				kind: MentionKind | undefined,
+				connectorType: string | undefined
+			) => {
+				removeDocumentChip(docId, docType, kind, connectorType);
+				onDocumentRemove?.(docId, docType, kind, connectorType);
 			},
 			[onDocumentRemove, removeDocumentChip]
 		);
@@ -679,7 +705,7 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
 				if (!isMentionNode(prev)) return;
 
 				e.preventDefault();
-				removeChip(prev.id, prev.document_type);
+				removeChip(prev.id, prev.document_type, prev.kind, prev.connector_type);
 			},
 			[editor.selection, getCurrentValue, onKeyDown, onSubmit, removeChip]
 		);

surfsense_web/components/assistant-ui/thread.tsx

@@ -36,7 +36,6 @@ import {
 import { chatSessionStateAtom } from "@/atoms/chat/chat-session-state.atom";
 import { currentThreadAtom } from "@/atoms/chat/current-thread.atom";
 import {
-	FOLDER_MENTION_DOCUMENT_TYPE,
 	type MentionedDocumentInfo,
 	mentionedDocumentsAtom,
 } from "@/atoms/chat/mentioned-documents.atom";
@@ -544,14 +543,12 @@ const Composer: FC = () => {
 					}
 				}
 				return docs.map<MentionedDocumentInfo>((d) => {
-					const documentType = d.document_type ?? "UNKNOWN";
 					if (d.kind === "connector") {
 						return {
 							id: d.id,
 							title: d.title,
-							document_type: documentType,
 							kind: "connector",
-							connector_type: d.connector_type ?? documentType,
+							connector_type: d.connector_type ?? "UNKNOWN",
 							account_name: d.account_name ?? d.title,
 						};
 					}
@@ -559,17 +556,13 @@ const Composer: FC = () => {
 						return {
 							id: d.id,
 							title: d.title,
-							document_type: FOLDER_MENTION_DOCUMENT_TYPE,
 							kind: "folder",
 						};
 					}
 					return {
 						id: d.id,
 						title: d.title,
-						// Atom requires a string; ``"UNKNOWN"`` matches the
-						// sentinel ``getMentionDocKey`` and the editor's
-						// match predicates use.
-						document_type: documentType,
+						document_type: d.document_type ?? "UNKNOWN",
 						kind: "doc",
 					};
 				});
@@ -760,13 +753,14 @@ const Composer: FC = () => {
 	]);
 
 	const handleDocumentRemove = useCallback(
-		(docId: number, docType?: string) => {
+		(docId: number, docType?: string, kind?: "doc" | "folder" | "connector", connectorType?: string) => {
 			setMentionedDocuments((prev) => {
-				if (!docType) {
-					// Fallback when chip type is unavailable.
-					return prev.filter((doc) => doc.id !== docId);
-				}
-				const removedKey = getMentionDocKey({ id: docId, document_type: docType });
+				const removedKey = getMentionDocKey({
+					id: docId,
+					document_type: docType,
+					kind,
+					connector_type: connectorType,
+				});
 				return prev.filter((doc) => getMentionDocKey(doc) !== removedKey);
 			});
 		},
@@ -810,7 +804,12 @@ const Composer: FC = () => {
 
 		for (const [key, doc] of prevDocsMap) {
 			if (!nextDocsMap.has(key)) {
-				editor.removeDocumentChip(doc.id, doc.document_type);
+				editor.removeDocumentChip(
+					doc.id,
+					doc.kind === "doc" ? doc.document_type : undefined,
+					doc.kind,
+					doc.kind === "connector" ? doc.connector_type : undefined
+				);
 			}
 		}
 

surfsense_web/components/assistant-ui/user-message.tsx

@@ -104,7 +104,7 @@ const UserTextPart: FC = () => {
 				const icon = isFolder ? (
 					<FolderIcon className="size-3.5" />
 				) : isConnector ? (
-					getConnectorIcon(segment.doc.connector_type ?? segment.doc.document_type, "size-3.5") ?? (
+					getConnectorIcon(segment.doc.connector_type, "size-3.5") ?? (
 						<Plug className="size-3.5" />
 					)
 				) : (

surfsense_web/components/documents/FolderTreeView.tsx

@@ -190,7 +190,6 @@ export function FolderTreeView({
 		for (const f of folders) {
 			const folderMentionKey = getMentionDocKey({
 				id: f.id,
-				document_type: "FOLDER",
 				kind: "folder",
 			});
 			states[f.id] = mentionedDocKeys.has(folderMentionKey) ? "all" : "none";

surfsense_web/components/new-chat/document-mention-picker.tsx

@@ -20,10 +20,7 @@ import {
 	useState,
 } from "react";
 import type * as React from "react";
-import {
-	FOLDER_MENTION_DOCUMENT_TYPE,
-	type MentionedDocumentInfo,
-} from "@/atoms/chat/mentioned-documents.atom";
+import type { MentionedDocumentInfo } from "@/atoms/chat/mentioned-documents.atom";
 import { useAtomValue } from "jotai";
 import { connectorsAtom } from "@/atoms/connectors/connector-query.atoms";
 import {
@@ -78,6 +75,10 @@ type ResourceNodeValue =
 	| { kind: "view"; view: BrowseView }
 	| { kind: "mention"; mention: MentionedDocumentInfo };
 
+function isConnectorActive(connector: SearchSourceConnector) {
+	return connector.is_active !== false;
+}
+
 function useDebounced<T>(value: T, delay = DEBOUNCE_MS) {
 	const [debounced, setDebounced] = useState(value);
 	const timeoutRef = useRef<ReturnType<typeof setTimeout> | undefined>(undefined);
@@ -115,22 +116,24 @@ function makeDocMention(doc: Pick<Document, "id" | "title" | "document_type">):
 	};
 }
 
-function makeFolderMention(folder: { id: number; title: string }): MentionedDocumentInfo {
+function makeFolderMention(
+	folder: { id: number; title: string }
+): Extract<MentionedDocumentInfo, { kind: "folder" }> {
 	return {
 		id: folder.id,
 		title: folder.title,
-		document_type: FOLDER_MENTION_DOCUMENT_TYPE,
 		kind: "folder",
 	};
 }
 
-function makeConnectorMention(connector: SearchSourceConnector): MentionedDocumentInfo {
+function makeConnectorMention(
+	connector: SearchSourceConnector
+): Extract<MentionedDocumentInfo, { kind: "connector" }> {
 	const accountName = getConnectorDisplayName(connector.name);
 	const connectorTitle = titleForConnectorType(connector.connector_type);
 	return {
 		id: connector.id,
 		title: `${connectorTitle}: ${accountName}`,
-		document_type: connector.connector_type,
 		kind: "connector",
 		connector_type: connector.connector_type,
 		account_name: accountName,
@@ -140,8 +143,8 @@ function makeConnectorMention(connector: SearchSourceConnector): MentionedDocume
 function mentionMatchesSearch(mention: MentionedDocumentInfo, searchLower: string) {
 	return [
 		mention.title,
-		mention.document_type,
 		mention.kind,
+		mention.kind === "doc" ? mention.document_type : "",
 		mention.kind === "connector" ? mention.connector_type : "",
 		mention.kind === "connector" ? mention.account_name : "",
 	].some((value) => value.toLowerCase().includes(searchLower));
@@ -171,6 +174,7 @@ export const DocumentMentionPicker = forwardRef<
 
 	const [zeroFolders] = useZeroQuery(queries.folders.bySpace({ searchSpaceId }));
 	const { data: connectors = [], isLoading: isConnectorsLoading } = useAtomValue(connectorsAtom);
+	const activeConnectors = useMemo(() => connectors.filter(isConnectorActive), [connectors]);
 	const paginationScopeKey = useMemo(
 		() => `${searchSpaceId}:${debouncedSearch}`,
 		[searchSpaceId, debouncedSearch]
@@ -307,8 +311,8 @@ export const DocumentMentionPicker = forwardRef<
 	}, [zeroFolders, debouncedSearch, deferredSearch, isSingleCharSearch, hasSearch]);
 
 	const connectorMentions = useMemo(
-		() => connectors.filter((c) => c.is_active).map(makeConnectorMention),
-		[connectors]
+		() => activeConnectors.map(makeConnectorMention),
+		[activeConnectors]
 	);
 
 	const selectedKeys = useMemo(
@@ -345,16 +349,16 @@ export const DocumentMentionPicker = forwardRef<
 			{
 				id: "connectors",
 				label: "Connectors",
-				subtitle: connectors.length
+				subtitle: activeConnectors.length
 					? "Choose the exact account for tool use"
 					: "No connected accounts yet",
 				icon: <Plug className="size-4" />,
 				type: "branch",
-				disabled: connectors.length === 0,
+				disabled: activeConnectors.length === 0,
 				value: { kind: "view", view: { kind: "connectors" } },
 			},
 		],
-		[connectors.length]
+		[activeConnectors.length]
 	);
 
 	const searchNodes = useMemo<ComposerSuggestionNode<ResourceNodeValue>[]>(() => {
@@ -385,7 +389,7 @@ export const DocumentMentionPicker = forwardRef<
 				id: getMentionDocKey(mention),
 				label: mention.title,
 				subtitle: "Connector account",
-				icon: getConnectorIcon(mention.document_type, "size-4") ?? <Plug className="size-4" />,
+				icon: getConnectorIcon(mention.connector_type, "size-4") ?? <Plug className="size-4" />,
 				type: "item" as const,
 				disabled: selectedKeys.has(getMentionDocKey(mention)),
 				value: { kind: "mention" as const, mention },
@@ -404,7 +408,7 @@ export const DocumentMentionPicker = forwardRef<
 
 	const connectorTypeEntries = useMemo(() => {
 		const byType = new Map<string, SearchSourceConnector[]>();
-		for (const connector of connectors.filter((c) => c.is_active)) {
+		for (const connector of activeConnectors) {
 			const list = byType.get(connector.connector_type) ?? [];
 			list.push(connector);
 			byType.set(connector.connector_type, list);
@@ -412,7 +416,7 @@ export const DocumentMentionPicker = forwardRef<
 		return Array.from(byType.entries()).sort(([a], [b]) =>
 			titleForConnectorType(a).localeCompare(titleForConnectorType(b))
 		);
-	}, [connectors]);
+	}, [activeConnectors]);
 
 	const browseNodes = useMemo<ComposerSuggestionNode<ResourceNodeValue>[]>(() => {
 		if (view.kind === "root") return rootNodes;
@@ -469,8 +473,8 @@ export const DocumentMentionPicker = forwardRef<
 				},
 			}));
 		}
-		return connectors
-			.filter((connector) => connector.is_active && connector.connector_type === view.connectorType)
+		return activeConnectors
+			.filter((connector) => connector.connector_type === view.connectorType)
 			.map((connector) => {
 				const mention = makeConnectorMention(connector);
 				return {
@@ -484,7 +488,7 @@ export const DocumentMentionPicker = forwardRef<
 				};
 			});
 	}, [
-		connectors,
+		activeConnectors,
 		connectorTypeEntries,
 		folderMentions,
 		rootNodes,