feat: enhance chat functionality with improved attachment handling and user experience

- Updated system prompt to clarify usage of the display_image tool, emphasizing URL requirements and restrictions on user-uploaded images.
- Enhanced the streaming chat process to provide more context about user attachments and documents during analysis.
- Implemented state resets when switching between chats to prevent stale data and race conditions.
- Added new components for displaying image previews and document attachments in the chat interface.
- Improved attachment processing to support image data URLs for persistent display after uploads.

surfsense_web/components/assistant-ui/attachment.tsx

@@ -41,13 +41,23 @@ const useAttachmentSrc = () => {
 	const { file, src } = useAssistantState(
 		useShallow(({ attachment }): { file?: File; src?: string } => {
 			if (!attachment || attachment.type !== "image") return {};
+			
+			// First priority: use File object if available (for new uploads)
 			if (attachment.file) return { file: attachment.file };
-			// Only try to filter if content is an array (standard assistant-ui format)
-			// Our custom ChatAttachment has content as a string, so skip this
-			if (Array.isArray(attachment.content)) {
-				const src = attachment.content.filter((c) => c.type === "image")[0]?.image;
-				if (src) return { src };
+			
+			// Second priority: use stored imageDataUrl (for persisted messages)
+			// This is stored in our custom ChatAttachment interface
+			const customAttachment = attachment as { imageDataUrl?: string };
+			if (customAttachment.imageDataUrl) {
+				return { src: customAttachment.imageDataUrl };
 			}
+			
+			// Third priority: try to extract from content array (standard assistant-ui format)
+			if (Array.isArray(attachment.content)) {
+				const contentSrc = attachment.content.filter((c) => c.type === "image")[0]?.image;
+				if (contentSrc) return { src: contentSrc };
+			}
+			
 			return {};
 		})
 	);
@@ -218,11 +228,77 @@ const AttachmentRemove: FC = () => {
 	);
 };
 
+/**
+ * Image attachment with preview thumbnail (click to expand)
+ */
+const MessageImageAttachment: FC = () => {
+	const attachmentName = useAssistantState(({ attachment }) => attachment?.name || "Image");
+	const src = useAttachmentSrc();
+
+	if (!src) return null;
+
+	return (
+		<AttachmentPreviewDialog>
+			<div
+				className="relative group cursor-pointer overflow-hidden rounded-xl border border-border/50 bg-muted transition-all hover:border-primary/30 hover:shadow-md"
+				title={`Click to expand: ${attachmentName}`}
+			>
+				<Image
+					src={src}
+					alt={attachmentName}
+					width={120}
+					height={90}
+					className="object-cover w-[120px] h-[90px] transition-transform group-hover:scale-105"
+				/>
+				{/* Hover overlay with filename */}
+				<div className="absolute inset-0 bg-gradient-to-t from-black/60 via-transparent to-transparent opacity-0 group-hover:opacity-100 transition-opacity">
+					<div className="absolute bottom-1.5 left-1.5 right-1.5">
+						<span className="text-[10px] text-white/90 font-medium truncate block">
+							{attachmentName}
+						</span>
+					</div>
+				</div>
+			</div>
+		</AttachmentPreviewDialog>
+	);
+};
+
+/**
+ * Document/file attachment as chip (similar to mentioned documents)
+ */
+const MessageDocumentAttachment: FC = () => {
+	const attachmentName = useAssistantState(({ attachment }) => attachment?.name || "Attachment");
+
+	return (
+		<AttachmentPreviewDialog>
+			<span
+				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 cursor-pointer hover:bg-primary/20 transition-colors"
+				title={attachmentName}
+			>
+				<FileText className="size-3" />
+				<span className="max-w-[150px] truncate">{attachmentName}</span>
+			</span>
+		</AttachmentPreviewDialog>
+	);
+};
+
+/**
+ * Attachment component for user messages
+ * Shows image preview for images, chip for documents
+ */
+const MessageAttachmentChip: FC = () => {
+	const isImage = useAssistantState(({ attachment }) => attachment?.type === "image");
+
+	if (isImage) {
+		return <MessageImageAttachment />;
+	}
+
+	return <MessageDocumentAttachment />;
+};
+
 export const UserMessageAttachments: FC = () => {
 	return (
-		<div className="aui-user-message-attachments-end col-span-full col-start-1 row-start-1 flex w-full flex-row justify-end gap-2">
-			<MessagePrimitive.Attachments components={{ Attachment: AttachmentUI }} />
-		</div>
+		<MessagePrimitive.Attachments components={{ Attachment: MessageAttachmentChip }} />
 	);
 };
 

surfsense_web/components/assistant-ui/thread.tsx

@@ -426,6 +426,9 @@ const Composer: FC = () => {
 	// Check if thread is empty (new chat)
 	const isThreadEmpty = useAssistantState(({ thread }) => thread.isEmpty);
 
+	// Check if thread is currently running (streaming response)
+	const isThreadRunning = useAssistantState(({ thread }) => thread.isRunning);
+
 	// Auto-focus editor when on new chat page
 	useEffect(() => {
 		if (isThreadEmpty && !hasAutoFocusedRef.current && editorRef.current) {
@@ -504,6 +507,10 @@ const Composer: FC = () => {
 
 	// Handle submit from inline editor (Enter key)
 	const handleSubmit = useCallback(() => {
+		// Prevent sending while a response is still streaming
+		if (isThreadRunning) {
+			return;
+		}
 		if (!showDocumentPopover) {
 			composerRuntime.send();
 			// Clear the editor after sending
@@ -511,7 +518,7 @@ const Composer: FC = () => {
 			setMentionedDocuments([]);
 			setMentionedDocumentIds([]);
 		}
-	}, [showDocumentPopover, composerRuntime, setMentionedDocuments, setMentionedDocumentIds]);
+	}, [showDocumentPopover, isThreadRunning, composerRuntime, setMentionedDocuments, setMentionedDocumentIds]);
 
 	// Handle document removal from inline editor
 	const handleDocumentRemove = useCallback(
@@ -973,19 +980,23 @@ const UserMessage: FC = () => {
 	const messageId = useAssistantState(({ message }) => message?.id);
 	const messageDocumentsMap = useAtomValue(messageDocumentsMapAtom);
 	const mentionedDocs = messageId ? messageDocumentsMap[messageId] : undefined;
+	const hasAttachments = useAssistantState(
+		({ message }) => message?.attachments && message.attachments.length > 0
+	);
 
 	return (
 		<MessagePrimitive.Root
 			className="aui-user-message-root fade-in slide-in-from-bottom-1 mx-auto grid w-full max-w-(--thread-max-width) animate-in auto-rows-auto grid-cols-[minmax(72px,1fr)_auto] content-start gap-y-2 px-2 py-3 duration-150 [&:where(>*)]:col-start-2"
 			data-role="user"
 		>
-			<UserMessageAttachments />
-
 			<div className="aui-user-message-content-wrapper col-start-2 min-w-0">
-				{/* Display mentioned documents as chips */}
-				{mentionedDocs && mentionedDocs.length > 0 && (
-					<div className="flex flex-wrap gap-1.5 mb-2 justify-end">
-						{mentionedDocs.map((doc) => (
+				{/* Display attachments and mentioned documents */}
+				{(hasAttachments || (mentionedDocs && mentionedDocs.length > 0)) && (
+					<div className="flex flex-wrap items-end gap-2 mb-2 justify-end">
+						{/* Attachments (images show as thumbnails, documents as chips) */}
+						<UserMessageAttachments />
+						{/* Mentioned documents as chips */}
+						{mentionedDocs?.map((doc) => (
 							<span
 								key={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"