feat(web): enhance chat context and mention handling with connector support

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

@@ -1,6 +1,6 @@
 "use client";
 
-import { Folder as FolderIcon, X as XIcon } from "lucide-react";
+import { Folder as FolderIcon, Plug as PlugIcon, X as XIcon } from "lucide-react";
 import type { NodeEntry, TElement } from "platejs";
 import type { PlateElementProps } from "platejs/react";
 import {
@@ -27,13 +27,15 @@ import type { Document } from "@/contracts/types/document.types";
 import { getMentionDocKey } from "@/lib/chat/mention-doc-key";
 import { cn } from "@/lib/utils";
 
-export type MentionKind = "doc" | "folder";
+export type MentionKind = "doc" | "folder" | "connector";
 
 export interface MentionedDocument {
 	id: number;
 	title: string;
 	document_type?: string;
 	kind: MentionKind;
+	connector_type?: string;
+	account_name?: string;
 }
 
 /**
@@ -46,6 +48,8 @@ export type MentionChipInput = {
 	title: string;
 	document_type?: string;
 	kind?: MentionKind;
+	connector_type?: string;
+	account_name?: string;
 };
 
 export type SuggestionAnchorRect = {
@@ -107,6 +111,8 @@ type MentionElementNode = {
 	document_type?: string;
 	/** Discriminator; defaults to ``"doc"`` for legacy nodes. */
 	kind?: MentionKind;
+	connector_type?: string;
+	account_name?: string;
 	statusLabel?: string | null;
 	statusKind?: MentionStatusKind;
 	children: [{ text: "" }];
@@ -146,6 +152,7 @@ const MentionElement: FC<PlateElementProps<MentionElementNode>> = ({
 				: "text-amber-700";
 
 	const isFolder = element.kind === "folder";
+	const isConnector = element.kind === "connector";
 	const ctx = useContext(MentionEditorContext);
 
 	return (
@@ -156,6 +163,10 @@ const MentionElement: FC<PlateElementProps<MentionElementNode>> = ({
 						<span className="flex items-center justify-center transition-opacity group-hover:opacity-0">
 							{isFolder ? (
 								<FolderIcon className="h-3 w-3" />
+							) : isConnector ? (
+								getConnectorIcon(element.connector_type ?? element.document_type ?? "UNKNOWN", "h-3 w-3") ?? (
+									<PlugIcon className="h-3 w-3" />
+								)
 							) : (
 								getConnectorIcon(element.document_type ?? "UNKNOWN", "h-3 w-3")
 							)}
@@ -242,6 +253,8 @@ function getMentionedDocuments(value: ComposerValue): MentionedDocument[] {
 				title: node.title,
 				document_type: node.document_type,
 				kind,
+				connector_type: node.connector_type,
+				account_name: node.account_name,
 			};
 			map.set(getMentionDocKey(doc), doc);
 		}
@@ -444,13 +457,20 @@ 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 : undefined);
+					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,
 					kind,
+					connector_type: mention.connector_type,
+					account_name: mention.account_name,
 					children: [{ text: "" }],
 				};
 

surfsense_web/components/assistant-ui/thread.tsx

@@ -36,6 +36,7 @@ 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";
@@ -71,7 +72,7 @@ import { ComposerSuggestionPopoverContent } from "@/components/new-chat/composer
 import {
 	DocumentMentionPicker,
 	type DocumentMentionPickerRef,
-} from "@/components/new-chat/document-mention-picker";
+} from "../new-chat/document-mention-picker";
 import { PromptPicker, type PromptPickerRef } from "@/components/new-chat/prompt-picker";
 import { Avatar, AvatarFallback, AvatarGroup } from "@/components/ui/avatar";
 import { Button } from "@/components/ui/button";
@@ -542,15 +543,36 @@ const Composer: FC = () => {
 						return prev;
 					}
 				}
-				return docs.map<MentionedDocumentInfo>((d) => ({
-					id: d.id,
-					title: d.title,
-					// Atom requires a string; ``"UNKNOWN"`` matches the
-					// sentinel ``getMentionDocKey`` and the editor's
-					// match predicates use.
-					document_type: d.document_type ?? "UNKNOWN",
-					kind: d.kind,
-				}));
+				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,
+							account_name: d.account_name ?? d.title,
+						};
+					}
+					if (d.kind === "folder") {
+						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,
+						kind: "doc",
+					};
+				});
 			});
 		},
 		[aui, setMentionedDocuments]
@@ -700,6 +722,9 @@ const Composer: FC = () => {
 				}
 				if (e.key === "Escape") {
 					e.preventDefault();
+					if (documentPickerRef.current?.goBack()) {
+						return;
+					}
 					setShowDocumentPopover(false);
 					setMentionQuery("");
 					setSuggestionAnchorPoint(null);

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

@@ -6,7 +6,7 @@ import {
 	useMessagePartText,
 } from "@assistant-ui/react";
 import { useAtomValue, useSetAtom } from "jotai";
-import { CheckIcon, CopyIcon, Folder as FolderIcon, Pencil } from "lucide-react";
+import { CheckIcon, CopyIcon, Folder as FolderIcon, Pencil, Plug } from "lucide-react";
 import Image from "next/image";
 import { useParams } from "next/navigation";
 import { type FC, useCallback, useState } from "react";
@@ -100,8 +100,13 @@ const UserTextPart: FC = () => {
 					return <span key={`txt-${segment.start}`}>{segment.value}</span>;
 				}
 				const isFolder = segment.doc.kind === "folder";
+				const isConnector = segment.doc.kind === "connector";
 				const icon = isFolder ? (
 					<FolderIcon className="size-3.5" />
+				) : isConnector ? (
+					getConnectorIcon(segment.doc.connector_type ?? segment.doc.document_type, "size-3.5") ?? (
+						<Plug className="size-3.5" />
+					)
 				) : (
 					getConnectorIcon(segment.doc.document_type ?? "UNKNOWN", "size-3.5")
 				);
@@ -110,8 +115,16 @@ const UserTextPart: FC = () => {
 						key={`mention-${getMentionDocKey(segment.doc)}-${segment.start}`}
 						icon={icon}
 						label={segment.doc.title}
-						tooltip={isFolder ? `Folder: ${segment.doc.title}` : segment.doc.title}
-						onClick={isFolder ? undefined : () => handleOpenDoc(segment.doc.id, segment.doc.title)}
+						tooltip={
+							isFolder
+								? `Folder: ${segment.doc.title}`
+								: isConnector
+									? `Connector account: ${segment.doc.title}`
+									: segment.doc.title
+						}
+						onClick={
+							isFolder || isConnector ? undefined : () => handleOpenDoc(segment.doc.id, segment.doc.title)
+						}
 						className="mx-0.5"
 					/>
 				);