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"

surfsense_web/components/prompt-kit/chain-of-thought.tsx

@@ -1,24 +1,130 @@
 "use client";
 
 import {
-	Brain,
-	CheckCircle2,
 	ChevronDown,
 	Circle,
-	Lightbulb,
-	Loader2,
-	Search,
-	Sparkles,
+	File,
+	FileAudio,
+	FileCode,
+	FileImage,
+	FileSpreadsheet,
+	FileText,
+	FileVideo,
 } from "lucide-react";
-import React from "react";
+import React, { useEffect, useState } from "react";
 import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible";
 import { cn } from "@/lib/utils";
 
+/**
+ * Custom hook for entrance animation
+ * Returns true after mount to trigger CSS transitions
+ */
+function useEntranceAnimation(delay = 0) {
+	const [isVisible, setIsVisible] = useState(false);
+
+	useEffect(() => {
+		const timer = setTimeout(() => setIsVisible(true), delay);
+		return () => clearTimeout(timer);
+	}, [delay]);
+
+	return isVisible;
+}
+
+/**
+ * Get file icon based on file extension (all icons are muted/gray)
+ */
+function getFileIcon(name: string): React.ReactNode {
+	const ext = name.split(".").pop()?.toLowerCase() || "";
+
+	// PDF / Word documents
+	if (ext === "pdf" || ["doc", "docx"].includes(ext)) {
+		return <FileText className="size-3.5" />;
+	}
+	// Spreadsheets
+	if (["xls", "xlsx", "csv"].includes(ext)) {
+		return <FileSpreadsheet className="size-3.5" />;
+	}
+	// Images
+	if (["png", "jpg", "jpeg", "gif", "webp", "svg"].includes(ext)) {
+		return <FileImage className="size-3.5" />;
+	}
+	// Audio
+	if (["mp3", "wav", "m4a", "ogg", "webm"].includes(ext)) {
+		return <FileAudio className="size-3.5" />;
+	}
+	// Video
+	if (["mp4", "mov", "avi", "mkv"].includes(ext)) {
+		return <FileVideo className="size-3.5" />;
+	}
+	// Code files
+	if (["js", "ts", "tsx", "jsx", "py", "html", "css", "json", "md"].includes(ext)) {
+		return <FileCode className="size-3.5" />;
+	}
+	// Default
+	return <File className="size-3.5" />;
+}
+
+/**
+ * Compact attachment tile component - matches the chat UI style
+ */
+const AttachmentTile: React.FC<{ name: string }> = ({ name }) => {
+	const icon = getFileIcon(name);
+
+	return (
+		<span
+			className="inline-flex items-center gap-1.5 rounded-lg bg-muted px-2 py-1 text-xs text-muted-foreground"
+			title={name}
+		>
+			<span className="shrink-0">{icon}</span>
+			<span className="truncate max-w-[120px]">{name}</span>
+		</span>
+	);
+};
+
+/**
+ * Parse text and render bracketed items (like [filename.pdf]) as styled tiles
+ */
+function parseAndRenderWithBadges(text: string): React.ReactNode {
+	// Match patterns like [filename.ext] or [N files] or [N documents]
+	const regex = /\[([^\]]+)\]/g;
+	const matches = Array.from(text.matchAll(regex));
+
+	if (matches.length === 0) {
+		return text;
+	}
+
+	const parts: React.ReactNode[] = [];
+	let lastIndex = 0;
+
+	for (const match of matches) {
+		const matchIndex = match.index ?? 0;
+
+		// Add text before the match
+		if (matchIndex > lastIndex) {
+			parts.push(text.slice(lastIndex, matchIndex));
+		}
+
+		const content = match[1];
+
+		// Render as a compact tile matching chat UI style with file-type colors
+		parts.push(<AttachmentTile key={`tile-${matchIndex}`} name={content} />);
+
+		lastIndex = matchIndex + match[0].length;
+	}
+
+	// Add remaining text
+	if (lastIndex < text.length) {
+		parts.push(text.slice(lastIndex));
+	}
+
+	return parts;
+}
+
 export type ChainOfThoughtItemProps = React.ComponentProps<"div">;
 
 export const ChainOfThoughtItem = ({ children, className, ...props }: ChainOfThoughtItemProps) => (
-	<div className={cn("text-muted-foreground text-sm", className)} {...props}>
-		{children}
+	<div className={cn("text-muted-foreground text-sm flex flex-wrap items-center gap-1", className)} {...props}>
+		{typeof children === "string" ? parseAndRenderWithBadges(children) : children}
 	</div>
 );
 
@@ -80,9 +186,28 @@ export const ChainOfThoughtContent = ({
 			{...props}
 		>
 			<div className="grid grid-cols-[min-content_minmax(0,1fr)] gap-x-4">
-				<div className="bg-primary/20 ml-1.75 h-full w-px group-data-[last=true]:hidden" />
+				{/* Animated vertical connection line */}
+				<div
+					className={cn(
+						"ml-1.75 w-px bg-primary/20 group-data-[last=true]:hidden",
+						"animate-in fade-in slide-in-from-top-1 duration-300"
+					)}
+				/>
 				<div className="ml-1.75 h-full w-px bg-transparent group-data-[last=false]:hidden" />
-				<div className="mt-2 space-y-2">{children}</div>
+				<div className="mt-2 space-y-1.5">
+					{React.Children.map(children, (child, index) => {
+						const key = React.isValidElement(child) ? child.key : `cot-item-${index}`;
+						return (
+							<div
+								key={key}
+								className="animate-in fade-in slide-in-from-left-2 duration-200"
+								style={{ animationDelay: `${index * 50}ms`, animationFillMode: "backwards" }}
+							>
+								{child}
+							</div>
+						);
+					})}
+				</div>
 			</div>
 		</CollapsibleContent>
 	);
@@ -98,14 +223,19 @@ export function ChainOfThought({ children, className }: ChainOfThoughtProps) {
 
 	return (
 		<div className={cn("space-y-0", className)}>
-			{childrenArray.map((child, index) => (
-				<React.Fragment key={index}>
-					{React.isValidElement(child) &&
-						React.cloneElement(child as React.ReactElement<ChainOfThoughtStepProps>, {
-							isLast: index === childrenArray.length - 1,
-						})}
-				</React.Fragment>
-			))}
+			{childrenArray.map((child, index) => {
+				// React.Children.toArray assigns stable keys to each child
+				const key = React.isValidElement(child) ? child.key : `cot-step-${index}`;
+				return (
+					<React.Fragment key={key}>
+						{React.isValidElement(child) &&
+							React.cloneElement(child as React.ReactElement<ChainOfThoughtStepProps>, {
+								isLast: index === childrenArray.length - 1,
+								stepIndex: index,
+							})}
+					</React.Fragment>
+				);
+			})}
 		</div>
 	);
 }
@@ -114,19 +244,42 @@ export type ChainOfThoughtStepProps = {
 	children: React.ReactNode;
 	className?: string;
 	isLast?: boolean;
+	/** Index of the step for staggered animation */
+	stepIndex?: number;
 };
 
 export const ChainOfThoughtStep = ({
 	children,
 	className,
 	isLast = false,
+	stepIndex = 0,
 	...props
 }: ChainOfThoughtStepProps & React.ComponentProps<typeof Collapsible>) => {
+	// Staggered entrance animation based on step index
+	const isVisible = useEntranceAnimation(stepIndex * 50);
+
 	return (
-		<Collapsible className={cn("group", className)} data-last={isLast} {...props}>
+		<Collapsible
+			className={cn(
+				"group transition-all duration-300 ease-out",
+				// Fade and slide in animation
+				isVisible ? "opacity-100 translate-y-0" : "opacity-0 translate-y-2",
+				className
+			)}
+			data-last={isLast}
+			{...props}
+		>
 			{children}
+			{/* Animated connection line to next step */}
 			<div className="flex justify-start group-data-[last=true]:hidden">
-				<div className="bg-primary/20 ml-1.75 h-4 w-px" />
+				<div
+					className={cn(
+						"ml-1.75 w-px bg-primary/20 transition-all duration-500 ease-out origin-top",
+						// Animate line height from 0 to full
+						isVisible ? "h-4 scale-y-100" : "h-0 scale-y-0"
+					)}
+					style={{ transitionDelay: `${stepIndex * 50 + 150}ms` }}
+				/>
 			</div>
 		</Collapsible>
 	);