feat(mentions): implement user message rendering with mention chips for referenced documents

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

@@ -1,11 +1,12 @@
 import { ActionBarPrimitive, AuiIf, MessagePrimitive, useAuiState } from "@assistant-ui/react";
 import { useAtomValue } from "jotai";
-import { CheckIcon, CopyIcon, FileText, Pencil } from "lucide-react";
+import { CheckIcon, CopyIcon, Pencil } from "lucide-react";
 import Image from "next/image";
 import { type FC, useState } from "react";
 import { currentThreadAtom } from "@/atoms/chat/current-thread.atom";
 import { messageDocumentsMapAtom } from "@/atoms/chat/mentioned-documents.atom";
 import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
+import { getConnectorIcon } from "@/contracts/enums/connectorIcons";
 
 interface AuthorMetadata {
 	displayName: string | null;
@@ -48,6 +49,19 @@ const UserAvatar: FC<AuthorMetadata> = ({ displayName, avatarUrl }) => {
 
 export const UserMessage: FC = () => {
 	const messageId = useAuiState(({ message }) => message?.id);
+	const messageText = useAuiState(({ message }) =>
+		(message?.content ?? [])
+			.map((part) =>
+				typeof part === "object" &&
+				part !== null &&
+				"type" in part &&
+				(part as { type?: string }).type === "text" &&
+				"text" in part
+					? String((part as { text?: string }).text ?? "")
+					: ""
+			)
+			.join("")
+	);
 	const messageDocumentsMap = useAtomValue(messageDocumentsMapAtom);
 	const mentionedDocs = messageId ? messageDocumentsMap[messageId] : undefined;
 	const metadata = useAuiState(({ message }) => message?.metadata);
@@ -63,22 +77,12 @@ export const UserMessage: FC = () => {
 			<div className="col-start-2 min-w-0">
 				<div className="aui-user-message-content-wrapper flex items-end gap-2">
 					<div className="relative flex-1 min-w-0">
-						{mentionedDocs && mentionedDocs.length > 0 && (
-							<div className="flex flex-wrap items-end gap-2 mb-2 justify-end">
-								{mentionedDocs?.map((doc) => (
-									<span
-										key={`${doc.document_type}:${doc.id}`}
-										className="inline-flex items-center gap-1 px-2 py-0.5 rounded-full bg-primary/10 text-xs font-medium text-primary border border-primary/20"
-										title={doc.title}
-									>
-										<FileText className="size-3" />
-										<span className="max-w-[150px] truncate">{doc.title}</span>
-									</span>
-								))}
-							</div>
-						)}
 						<div className="aui-user-message-content wrap-break-word rounded-2xl bg-muted px-4 py-2.5 text-foreground">
-							<MessagePrimitive.Parts />
+							{mentionedDocs && mentionedDocs.length > 0 ? (
+								<UserMessageWithMentionChips text={messageText} mentionedDocs={mentionedDocs} />
+							) : (
+								<MessagePrimitive.Parts />
+							)}
 						</div>
 						<div className="absolute right-0 top-full mt-1 z-10 opacity-100 pointer-events-auto md:opacity-0 md:pointer-events-none md:transition-opacity md:duration-200 md:delay-300 md:group-hover/user-msg:opacity-100 md:group-hover/user-msg:delay-0 md:group-hover/user-msg:pointer-events-auto">
 							<UserActionBar />
@@ -95,6 +99,64 @@ export const UserMessage: FC = () => {
 	);
 };
 
+const UserMessageWithMentionChips: FC<{
+	text: string;
+	mentionedDocs: { id: number; title: string; document_type: string }[];
+}> = ({ text, mentionedDocs }) => {
+	type Segment =
+		| { type: "text"; value: string; start: number }
+		| { type: "mention"; doc: { id: number; title: string; document_type: string }; start: number };
+
+	const tokens = mentionedDocs
+		.map((doc) => ({ doc, token: `@${doc.title}` }))
+		.sort((a, b) => b.token.length - a.token.length);
+
+	const segments: Segment[] = [];
+	let i = 0;
+	let buffer = "";
+	let bufferStart = 0;
+	while (i < text.length) {
+		const tokenMatch = tokens.find(({ token }) => text.startsWith(token, i));
+		if (tokenMatch) {
+			if (buffer) {
+				segments.push({ type: "text", value: buffer, start: bufferStart });
+				buffer = "";
+			}
+			segments.push({ type: "mention", doc: tokenMatch.doc, start: i });
+			i += tokenMatch.token.length;
+			bufferStart = i;
+			continue;
+		}
+		if (!buffer) bufferStart = i;
+		buffer += text[i];
+		i += 1;
+	}
+	if (buffer) {
+		segments.push({ type: "text", value: buffer, start: bufferStart });
+	}
+
+	return (
+		<span className="whitespace-pre-wrap break-words">
+			{segments.map((segment) =>
+				segment.type === "text" ? (
+					<span key={`txt-${segment.start}`}>{segment.value}</span>
+				) : (
+					<span
+						key={`mention-${segment.doc.document_type}:${segment.doc.id}-${segment.start}`}
+						className="inline-flex items-center gap-1 mx-0.5 px-1 py-0.5 rounded bg-primary/10 text-xs font-bold text-primary/60 select-none align-baseline"
+						title={segment.doc.title}
+					>
+						<span className="flex items-center text-muted-foreground">
+							{getConnectorIcon(segment.doc.document_type ?? "UNKNOWN", "h-3 w-3")}
+						</span>
+						<span className="max-w-[120px] truncate">{segment.doc.title}</span>
+					</span>
+				)
+			)}
+		</span>
+	);
+};
+
 const UserActionBar: FC = () => {
 	const isThreadRunning = useAuiState(({ thread }) => thread.isRunning);