Merge pull request #627 from MODSetter/dev

feat: various ux improvements

surfsense_backend/app/agents/new_chat/system_prompt.py

@@ -64,14 +64,18 @@ You have access to the following tools:
   - The preview card will automatically be displayed in the chat.
 
 4. display_image: Display an image in the chat with metadata.
-  - Use this tool when you want to show an image to the user.
+  - Use this tool when you want to show an image from a URL to the user.
   - This displays the image with an optional title, description, and source attribution.
   - Common use cases:
     * Showing an image from a URL mentioned in the conversation
     * Displaying a diagram, chart, or illustration you're referencing
     * Showing visual examples when explaining concepts
+  - IMPORTANT: Do NOT use this tool for user-uploaded image attachments!
+    * User attachments are already visible in the chat UI - the user can see them
+    * This tool requires a valid HTTP/HTTPS URL, not a local file path
+    * When a user uploads an image, just analyze it and respond - don't try to display it again
   - Args:
-    - src: The URL of the image to display (must be a valid HTTP/HTTPS image URL)
+    - src: The URL of the image to display (must be a valid HTTP/HTTPS image URL, not a local path)
     - alt: Alternative text describing the image (for accessibility)
     - title: Optional title to display below the image
     - description: Optional description providing context about the image

surfsense_backend/app/agents/new_chat/tools/link_preview.py

@@ -6,13 +6,19 @@ Open Graph image, etc.) to display rich link previews in the chat UI.
 """
 
 import hashlib
+import logging
 import re
 from typing import Any
 from urllib.parse import urlparse
 
 import httpx
+import trafilatura
+from fake_useragent import UserAgent
+from langchain_community.document_loaders import AsyncChromiumLoader
 from langchain_core.tools import tool
 
+logger = logging.getLogger(__name__)
+
 
 def extract_domain(url: str) -> str:
     """Extract the domain from a URL."""
@@ -138,6 +144,96 @@ def generate_preview_id(url: str) -> str:
     return f"link-preview-{hash_val}"
 
 
+def _unescape_html(text: str) -> str:
+    """Unescape common HTML entities."""
+    return (
+        text.replace("&", "&")
+        .replace("&lt;", "<")
+        .replace("&gt;", ">")
+        .replace("&quot;", '"')
+        .replace("&#39;", "'")
+        .replace("&apos;", "'")
+    )
+
+
+def _make_absolute_url(image_url: str, base_url: str) -> str:
+    """Convert a relative image URL to an absolute URL."""
+    if image_url.startswith(("http://", "https://")):
+        return image_url
+    if image_url.startswith("//"):
+        return f"https:{image_url}"
+    if image_url.startswith("/"):
+        parsed = urlparse(base_url)
+        return f"{parsed.scheme}://{parsed.netloc}{image_url}"
+    return image_url
+
+
+async def fetch_with_chromium(url: str) -> dict[str, Any] | None:
+    """
+    Fetch page content using headless Chromium browser.
+    Used as a fallback when simple HTTP requests are blocked (403, etc.).
+
+    Args:
+        url: URL to fetch
+
+    Returns:
+        Dict with title, description, image, and raw_html, or None if failed
+    """
+    try:
+        logger.info(f"[link_preview] Falling back to Chromium for {url}")
+
+        # Generate a realistic User-Agent to avoid bot detection
+        ua = UserAgent()
+        user_agent = ua.random
+
+        # Use AsyncChromiumLoader to fetch the page
+        crawl_loader = AsyncChromiumLoader(
+            urls=[url], headless=True, user_agent=user_agent
+        )
+        documents = await crawl_loader.aload()
+
+        if not documents:
+            logger.warning(f"[link_preview] Chromium returned no documents for {url}")
+            return None
+
+        doc = documents[0]
+        raw_html = doc.page_content
+
+        if not raw_html or len(raw_html.strip()) == 0:
+            logger.warning(f"[link_preview] Chromium returned empty content for {url}")
+            return None
+
+        # Extract metadata using Trafilatura
+        trafilatura_metadata = trafilatura.extract_metadata(raw_html)
+
+        # Extract OG image from raw HTML (trafilatura doesn't extract this)
+        image = extract_image(raw_html)
+
+        result = {
+            "title": None,
+            "description": None,
+            "image": image,
+            "raw_html": raw_html,
+        }
+
+        if trafilatura_metadata:
+            result["title"] = trafilatura_metadata.title
+            result["description"] = trafilatura_metadata.description
+
+        # If trafilatura didn't get the title/description, try OG tags
+        if not result["title"]:
+            result["title"] = extract_title(raw_html)
+        if not result["description"]:
+            result["description"] = extract_description(raw_html)
+
+        logger.info(f"[link_preview] Successfully fetched {url} via Chromium")
+        return result
+
+    except Exception as e:
+        logger.error(f"[link_preview] Chromium fallback failed for {url}: {e}")
+        return None
+
+
 def create_link_preview_tool():
     """
     Factory function to create the link_preview tool.
@@ -184,13 +280,20 @@ def create_link_preview_tool():
             url = f"https://{url}"
 
         try:
+            # Use a browser-like User-Agent to fetch Open Graph metadata.
+            # This is the same approach used by Slack, Discord, Twitter, etc. for link previews.
+            # We're only fetching publicly available metadata (title, description, thumbnail)
+            # that websites intentionally expose via OG tags for link preview purposes.
             async with httpx.AsyncClient(
                 timeout=10.0,
                 follow_redirects=True,
                 headers={
-                    "User-Agent": "Mozilla/5.0 (compatible; SurfSenseBot/1.0; +https://surfsense.net)",
-                    "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8",
-                    "Accept-Language": "en-US,en;q=0.5",
+                    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36",
+                    "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8",
+                    "Accept-Language": "en-US,en;q=0.9",
+                    "Accept-Encoding": "gzip, deflate, br",
+                    "Cache-Control": "no-cache",
+                    "Pragma": "no-cache",
                 },
             ) as client:
                 response = await client.get(url)
@@ -218,32 +321,14 @@ def create_link_preview_tool():
                 image = extract_image(html)
 
                 # Make sure image URL is absolute
-                if image and not image.startswith(("http://", "https://")):
-                    if image.startswith("//"):
-                        image = f"https:{image}"
-                    elif image.startswith("/"):
-                        parsed = urlparse(url)
-                        image = f"{parsed.scheme}://{parsed.netloc}{image}"
+                if image:
+                    image = _make_absolute_url(image, url)
 
                 # Clean up title and description (unescape HTML entities)
                 if title:
-                    title = (
-                        title.replace("&amp;", "&")
-                        .replace("&lt;", "<")
-                        .replace("&gt;", ">")
-                        .replace("&quot;", '"')
-                        .replace("&#39;", "'")
-                        .replace("&apos;", "'")
-                    )
+                    title = _unescape_html(title)
                 if description:
-                    description = (
-                        description.replace("&amp;", "&")
-                        .replace("&lt;", "<")
-                        .replace("&gt;", ">")
-                        .replace("&quot;", '"')
-                        .replace("&#39;", "'")
-                        .replace("&apos;", "'")
-                    )
+                    description = _unescape_html(description)
                     # Truncate long descriptions
                     if len(description) > 200:
                         description = description[:197] + "..."
@@ -260,6 +345,39 @@ def create_link_preview_tool():
                 }
 
         except httpx.TimeoutException:
+            # Timeout - try Chromium fallback
+            logger.warning(
+                f"[link_preview] Timeout for {url}, trying Chromium fallback"
+            )
+            chromium_result = await fetch_with_chromium(url)
+            if chromium_result:
+                title = chromium_result.get("title") or domain
+                description = chromium_result.get("description")
+                image = chromium_result.get("image")
+
+                # Clean up and truncate
+                if title:
+                    title = _unescape_html(title)
+                if description:
+                    description = _unescape_html(description)
+                    if len(description) > 200:
+                        description = description[:197] + "..."
+
+                # Make sure image URL is absolute
+                if image:
+                    image = _make_absolute_url(image, url)
+
+                return {
+                    "id": preview_id,
+                    "assetId": url,
+                    "kind": "link",
+                    "href": url,
+                    "title": title,
+                    "description": description,
+                    "thumb": image,
+                    "domain": domain,
+                }
+
             return {
                 "id": preview_id,
                 "assetId": url,
@@ -270,6 +388,42 @@ def create_link_preview_tool():
                 "error": "Request timed out",
             }
         except httpx.HTTPStatusError as e:
+            status_code = e.response.status_code
+
+            # For 403 (Forbidden) and similar bot-detection errors, try Chromium fallback
+            if status_code in (403, 401, 406, 429):
+                logger.warning(
+                    f"[link_preview] HTTP {status_code} for {url}, trying Chromium fallback"
+                )
+                chromium_result = await fetch_with_chromium(url)
+                if chromium_result:
+                    title = chromium_result.get("title") or domain
+                    description = chromium_result.get("description")
+                    image = chromium_result.get("image")
+
+                    # Clean up and truncate
+                    if title:
+                        title = _unescape_html(title)
+                    if description:
+                        description = _unescape_html(description)
+                        if len(description) > 200:
+                            description = description[:197] + "..."
+
+                    # Make sure image URL is absolute
+                    if image:
+                        image = _make_absolute_url(image, url)
+
+                    return {
+                        "id": preview_id,
+                        "assetId": url,
+                        "kind": "link",
+                        "href": url,
+                        "title": title,
+                        "description": description,
+                        "thumb": image,
+                        "domain": domain,
+                    }
+
             return {
                 "id": preview_id,
                 "assetId": url,
@@ -277,11 +431,11 @@ def create_link_preview_tool():
                 "href": url,
                 "title": domain or "Link",
                 "domain": domain,
-                "error": f"HTTP {e.response.status_code}",
+                "error": f"HTTP {status_code}",
             }
         except Exception as e:
             error_message = str(e)
-            print(f"[link_preview] Error fetching {url}: {error_message}")
+            logger.error(f"[link_preview] Error fetching {url}: {error_message}")
             return {
                 "id": preview_id,
                 "assetId": url,

surfsense_backend/app/tasks/chat/stream_new_chat.py

@@ -255,13 +255,59 @@ async def stream_new_chat(
         # Initial thinking step - analyzing the request
         analyze_step_id = next_thinking_step_id()
         last_active_step_id = analyze_step_id
-        last_active_step_title = "Understanding your request"
-        last_active_step_items = [
-            f"Processing: {user_query[:80]}{'...' if len(user_query) > 80 else ''}"
-        ]
+
+        # Determine step title and action verb based on context
+        if attachments and mentioned_documents:
+            last_active_step_title = "Analyzing your content"
+            action_verb = "Reading"
+        elif attachments:
+            last_active_step_title = "Reading your content"
+            action_verb = "Reading"
+        elif mentioned_documents:
+            last_active_step_title = "Analyzing referenced content"
+            action_verb = "Analyzing"
+        else:
+            last_active_step_title = "Understanding your request"
+            action_verb = "Processing"
+
+        # Build the message with inline context about attachments/documents
+        processing_parts = []
+
+        # Add the user query
+        query_text = user_query[:80] + ("..." if len(user_query) > 80 else "")
+        processing_parts.append(query_text)
+
+        # Add file attachment names inline
+        if attachments:
+            attachment_names = []
+            for attachment in attachments:
+                name = attachment.name
+                if len(name) > 30:
+                    name = name[:27] + "..."
+                attachment_names.append(name)
+            if len(attachment_names) == 1:
+                processing_parts.append(f"[{attachment_names[0]}]")
+            else:
+                processing_parts.append(f"[{len(attachment_names)} files]")
+
+        # Add mentioned document names inline
+        if mentioned_documents:
+            doc_names = []
+            for doc in mentioned_documents:
+                title = doc.title
+                if len(title) > 30:
+                    title = title[:27] + "..."
+                doc_names.append(title)
+            if len(doc_names) == 1:
+                processing_parts.append(f"[{doc_names[0]}]")
+            else:
+                processing_parts.append(f"[{len(doc_names)} documents]")
+
+        last_active_step_items = [f"{action_verb}: {' '.join(processing_parts)}"]
+
         yield streaming_service.format_thinking_step(
             step_id=analyze_step_id,
-            title="Understanding your request",
+            title=last_active_step_title,
             status="in_progress",
             items=last_active_step_items,
         )
@@ -369,13 +415,13 @@ async def stream_new_chat(
                         if isinstance(tool_input, dict)
                         else ""
                     )
-                    last_active_step_title = "Displaying image"
+                    last_active_step_title = "Analyzing the image"
                     last_active_step_items = [
-                        f"Image: {title[:50] if title else src[:50]}{'...' if len(title or src) > 50 else ''}"
+                        f"Analyzing: {title[:50] if title else src[:50]}{'...' if len(title or src) > 50 else ''}"
                     ]
                     yield streaming_service.format_thinking_step(
                         step_id=tool_step_id,
-                        title="Displaying image",
+                        title="Analyzing the image",
                         status="in_progress",
                         items=last_active_step_items,
                     )
@@ -471,7 +517,7 @@ async def stream_new_chat(
                         else str(tool_input)
                     )
                     yield streaming_service.format_terminal_info(
-                        f"Displaying image: {src[:60]}{'...' if len(src) > 60 else ''}",
+                        f"Analyzing image: {src[:60]}{'...' if len(src) > 60 else ''}",
                         "info",
                     )
                 elif tool_name == "scrape_webpage":
@@ -575,20 +621,20 @@ async def stream_new_chat(
                         items=completed_items,
                     )
                 elif tool_name == "display_image":
-                    # Build completion items for image display
+                    # Build completion items for image analysis
                     if isinstance(tool_output, dict):
                         title = tool_output.get("title", "")
                         alt = tool_output.get("alt", "Image")
                         display_name = title or alt
                         completed_items = [
                             *last_active_step_items,
-                            f"Showing: {display_name[:50]}{'...' if len(display_name) > 50 else ''}",
+                            f"Analyzed: {display_name[:50]}{'...' if len(display_name) > 50 else ''}",
                         ]
                     else:
-                        completed_items = [*last_active_step_items, "Image displayed"]
+                        completed_items = [*last_active_step_items, "Image analyzed"]
                     yield streaming_service.format_thinking_step(
                         step_id=original_step_id,
-                        title="Displaying image",
+                        title="Analyzing the image",
                         status="completed",
                         items=completed_items,
                     )
@@ -744,7 +790,7 @@ async def stream_new_chat(
                             "alt", "Image"
                         )
                         yield streaming_service.format_terminal_info(
-                            f"Image displayed: {title[:40]}{'...' if len(title) > 40 else ''}",
+                            f"Image analyzed: {title[:40]}{'...' if len(title) > 40 else ''}",
                             "success",
                         )
                 elif tool_name == "scrape_webpage":

surfsense_web/app/dashboard/[search_space_id]/editor/[documentId]/page.tsx

@@ -25,17 +25,31 @@ import { Separator } from "@/components/ui/separator";
 import { notesApiService } from "@/lib/apis/notes-api.service";
 import { authenticatedFetch, getBearerToken, redirectToLogin } from "@/lib/auth-utils";
 
+// BlockNote types
+type BlockNoteInlineContent =
+	| string
+	| { text?: string; type?: string; styles?: Record<string, unknown> };
+
+interface BlockNoteBlock {
+	type: string;
+	content?: BlockNoteInlineContent[];
+	children?: BlockNoteBlock[];
+	props?: Record<string, unknown>;
+}
+
+type BlockNoteDocument = BlockNoteBlock[] | null | undefined;
+
 interface EditorContent {
 	document_id: number;
 	title: string;
 	document_type?: string;
-	blocknote_document: any;
+	blocknote_document: BlockNoteDocument;
 	updated_at: string | null;
 }
 
 // Helper function to extract title from BlockNote document
 // Takes the text content from the first block (should be a heading for notes)
-function extractTitleFromBlockNote(blocknoteDocument: any[] | null | undefined): string {
+function extractTitleFromBlockNote(blocknoteDocument: BlockNoteDocument): string {
 	if (!blocknoteDocument || !Array.isArray(blocknoteDocument) || blocknoteDocument.length === 0) {
 		return "Untitled";
 	}
@@ -49,9 +63,9 @@ function extractTitleFromBlockNote(blocknoteDocument: any[] | null | undefined):
 	// BlockNote blocks have a content array with inline content
 	if (firstBlock.content && Array.isArray(firstBlock.content)) {
 		const textContent = firstBlock.content
-			.map((item: any) => {
+			.map((item: BlockNoteInlineContent) => {
 				if (typeof item === "string") return item;
-				if (item?.text) return item.text;
+				if (typeof item === "object" && item?.text) return item.text;
 				return "";
 			})
 			.join("")
@@ -73,7 +87,7 @@ export default function EditorPage() {
 	const [document, setDocument] = useState<EditorContent | null>(null);
 	const [loading, setLoading] = useState(true);
 	const [saving, setSaving] = useState(false);
-	const [editorContent, setEditorContent] = useState<any>(null);
+	const [editorContent, setEditorContent] = useState<BlockNoteDocument>(null);
 	const [error, setError] = useState<string | null>(null);
 	const [hasUnsavedChanges, setHasUnsavedChanges] = useState(false);
 	const [showUnsavedDialog, setShowUnsavedDialog] = useState(false);
@@ -323,7 +337,7 @@ export default function EditorPage() {
 		if (hasUnsavedChanges) {
 			setShowUnsavedDialog(true);
 		} else {
-			router.push(`/dashboard/${searchSpaceId}/researcher`);
+			router.push(`/dashboard/${searchSpaceId}/new-chat`);
 		}
 	};
 
@@ -333,12 +347,12 @@ export default function EditorPage() {
 		setGlobalHasUnsavedChanges(false);
 		setHasUnsavedChanges(false);
 
-		// If there's a pending navigation (from sidebar), use that; otherwise go back to researcher
+		// If there's a pending navigation (from sidebar), use that; otherwise go back to chat
 		if (pendingNavigation) {
 			router.push(pendingNavigation);
 			setPendingNavigation(null);
 		} else {
-			router.push(`/dashboard/${searchSpaceId}/researcher`);
+			router.push(`/dashboard/${searchSpaceId}/new-chat`);
 		}
 	};
 
@@ -379,7 +393,7 @@ export default function EditorPage() {
 						</CardHeader>
 						<CardContent>
 							<Button
-								onClick={() => router.push(`/dashboard/${searchSpaceId}/researcher`)}
+								onClick={() => router.push(`/dashboard/${searchSpaceId}/new-chat`)}
 								variant="outline"
 								className="gap-2"
 							>
@@ -410,7 +424,7 @@ export default function EditorPage() {
 		<motion.div
 			initial={{ opacity: 0 }}
 			animate={{ opacity: 1 }}
-			className="flex flex-col h-full w-full"
+			className="flex flex-col min-h-screen w-full"
 		>
 			{/* Toolbar */}
 			<div className="sticky top-0 z-40 flex h-16 shrink-0 items-center gap-4 border-b bg-background/95 backdrop-blur supports-backdrop-filter:bg-background/60 px-6">
@@ -444,7 +458,7 @@ export default function EditorPage() {
 			</div>
 
 			{/* Editor Container */}
-			<div className="flex-1 overflow-visible relative">
+			<div className="flex-1 min-h-0 overflow-hidden relative">
 				<div className="h-full w-full overflow-auto p-6">
 					{error && (
 						<motion.div

surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx

@@ -7,9 +7,10 @@ import {
 	useExternalStoreRuntime,
 } from "@assistant-ui/react";
 import { useAtomValue, useSetAtom } from "jotai";
-import { useParams, useRouter } from "next/navigation";
+import { useParams } from "next/navigation";
 import { useCallback, useEffect, useMemo, useRef, useState } from "react";
 import { toast } from "sonner";
+import { z } from "zod";
 import {
 	type MentionedDocumentInfo,
 	mentionedDocumentIdsAtom,
@@ -55,20 +56,33 @@ function extractThinkingSteps(content: unknown): ThinkingStep[] {
 }
 
 /**
- * Extract mentioned documents from message content
+ * Zod schema for mentioned document info (for type-safe parsing)
+ */
+const MentionedDocumentInfoSchema = z.object({
+	id: z.number(),
+	title: z.string(),
+	document_type: z.string(),
+});
+
+const MentionedDocumentsPartSchema = z.object({
+	type: z.literal("mentioned-documents"),
+	documents: z.array(MentionedDocumentInfoSchema),
+});
+
+/**
+ * Extract mentioned documents from message content (type-safe with Zod)
  */
 function extractMentionedDocuments(content: unknown): MentionedDocumentInfo[] {
 	if (!Array.isArray(content)) return [];
 
-	const docsPart = content.find(
-		(part: unknown) =>
-			typeof part === "object" &&
-			part !== null &&
-			"type" in part &&
-			(part as { type: string }).type === "mentioned-documents"
-	) as { type: "mentioned-documents"; documents: MentionedDocumentInfo[] } | undefined;
+	for (const part of content) {
+		const result = MentionedDocumentsPartSchema.safeParse(part);
+		if (result.success) {
+			return result.data.documents;
+		}
+	}
 
-	return docsPart?.documents || [];
+	return [];
 }
 
 /**
@@ -126,7 +140,6 @@ interface ThinkingStepData {
 
 export default function NewChatPage() {
 	const params = useParams();
-	const router = useRouter();
 	const [isInitializing, setIsInitializing] = useState(true);
 	const [threadId, setThreadId] = useState<number | null>(null);
 	const [messages, setMessages] = useState<ThreadMessageLike[]>([]);
@@ -168,9 +181,18 @@ export default function NewChatPage() {
 	}, [params.chat_id]);
 
 	// Initialize thread and load messages
+	// For new chats (no urlChatId), we use lazy creation - thread is created on first message
 	const initializeThread = useCallback(async () => {
 		setIsInitializing(true);
 
+		// Reset all state when switching between chats to prevent stale data
+		setMessages([]);
+		setThreadId(null);
+		setMessageThinkingSteps(new Map());
+		setMentionedDocumentIds([]);
+		setMentionedDocuments([]);
+		setMessageDocumentsMap({});
+
 		try {
 			if (urlChatId > 0) {
 				// Thread exists - load messages
@@ -206,22 +228,20 @@ export default function NewChatPage() {
 						setMessageDocumentsMap(restoredDocsMap);
 					}
 				}
-			} else {
-				// Create new thread
-				const newThread = await createThread(searchSpaceId, "New Chat");
-				setThreadId(newThread.id);
-				router.replace(`/dashboard/${searchSpaceId}/new-chat/${newThread.id}`);
 			}
+			// For new chats (urlChatId === 0), don't create thread yet
+			// Thread will be created lazily when user sends first message
+			// This improves UX (instant load) and avoids orphan threads
 		} catch (error) {
 			console.error("[NewChatPage] Failed to initialize thread:", error);
 			// Keep threadId as null - don't use Date.now() as it creates an invalid ID
 			// that will cause 404 errors on subsequent API calls
 			setThreadId(null);
-			toast.error("Failed to initialize chat. Please try again.");
+			toast.error("Failed to load chat. Please try again.");
 		} finally {
 			setIsInitializing(false);
 		}
-	}, [urlChatId, searchSpaceId, router, setMessageDocumentsMap]);
+	}, [urlChatId, setMessageDocumentsMap, setMentionedDocumentIds, setMentionedDocuments]);
 
 	// Initialize on mount
 	useEffect(() => {
@@ -240,7 +260,12 @@ export default function NewChatPage() {
 	// Handle new message from user
 	const onNew = useCallback(
 		async (message: AppendMessage) => {
-			if (!threadId) return;
+			// Abort any previous streaming request to prevent race conditions
+			// when user sends a second query while the first is still streaming
+			if (abortControllerRef.current) {
+				abortControllerRef.current.abort();
+				abortControllerRef.current = null;
+			}
 
 			// Extract user query text from content parts
 			let userQuery = "";
@@ -273,6 +298,27 @@ export default function NewChatPage() {
 				return;
 			}
 
+			// Lazy thread creation: create thread on first message if it doesn't exist
+			let currentThreadId = threadId;
+			if (!currentThreadId) {
+				try {
+					const newThread = await createThread(searchSpaceId, "New Chat");
+					currentThreadId = newThread.id;
+					setThreadId(currentThreadId);
+					// Update URL silently using browser API (not router.replace) to avoid
+					// interrupting the ongoing fetch/streaming with React navigation
+					window.history.replaceState(
+						null,
+						"",
+						`/dashboard/${searchSpaceId}/new-chat/${currentThreadId}`
+					);
+				} catch (error) {
+					console.error("[NewChatPage] Failed to create thread:", error);
+					toast.error("Failed to start chat. Please try again.");
+					return;
+				}
+			}
+
 			// Add user message to state
 			const userMsgId = `msg-user-${Date.now()}`;
 			const userMessage: ThreadMessageLike = {
@@ -280,6 +326,8 @@ export default function NewChatPage() {
 				role: "user",
 				content: message.content,
 				createdAt: new Date(),
+				// Include attachments so they can be displayed
+				attachments: message.attachments || [],
 			};
 			setMessages((prev) => [...prev, userMessage]);
 
@@ -311,7 +359,7 @@ export default function NewChatPage() {
 							},
 						]
 					: message.content;
-			appendMessage(threadId, {
+			appendMessage(currentThreadId, {
 				role: "user",
 				content: persistContent,
 			}).catch((err) => console.error("Failed to persist user message:", err));
@@ -468,7 +516,7 @@ export default function NewChatPage() {
 						Authorization: `Bearer ${token}`,
 					},
 					body: JSON.stringify({
-						chat_id: threadId,
+						chat_id: currentThreadId,
 						user_query: userQuery.trim(),
 						search_space_id: searchSpaceId,
 						messages: messageHistory,
@@ -601,7 +649,7 @@ export default function NewChatPage() {
 				// Persist assistant message (with thinking steps for restoration on refresh)
 				const finalContent = buildContentForPersistence();
 				if (contentParts.length > 0) {
-					appendMessage(threadId, {
+					appendMessage(currentThreadId, {
 						role: "assistant",
 						content: finalContent,
 					}).catch((err) => console.error("Failed to persist assistant message:", err));
@@ -678,7 +726,7 @@ export default function NewChatPage() {
 		},
 	});
 
-	// Show loading state
+	// Show loading state only when loading an existing thread
 	if (isInitializing) {
 		return (
 			<div className="flex h-[calc(100vh-64px)] items-center justify-center">
@@ -687,11 +735,12 @@ export default function NewChatPage() {
 		);
 	}
 
-	// Show error state if thread initialization failed
-	if (!threadId) {
+	// Show error state only if we tried to load an existing thread but failed
+	// For new chats (urlChatId === 0), threadId being null is expected (lazy creation)
+	if (!threadId && urlChatId > 0) {
 		return (
 			<div className="flex h-[calc(100vh-64px)] flex-col items-center justify-center gap-4">
-				<div className="text-destructive">Failed to initialize chat</div>
+				<div className="text-destructive">Failed to load chat</div>
 				<button
 					type="button"
 					onClick={() => {

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,14 +228,78 @@ const AttachmentRemove: FC = () => {
 	);
 };
 
-export const UserMessageAttachments: 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 (
-		<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>
+		<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 <MessagePrimitive.Attachments components={{ Attachment: MessageAttachmentChip }} />;
+};
+
 export const ComposerAttachments: FC = () => {
 	return (
 		<div className="aui-composer-attachments mb-2 flex w-full flex-row items-center gap-2 overflow-x-auto px-1.5 pt-0.5 pb-1 empty:hidden">

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

@@ -0,0 +1,527 @@
+"use client";
+
+import { X } from "lucide-react";
+import {
+	createElement,
+	forwardRef,
+	useCallback,
+	useEffect,
+	useImperativeHandle,
+	useRef,
+	useState,
+} from "react";
+import ReactDOMServer from "react-dom/server";
+import type { Document } from "@/contracts/types/document.types";
+import { cn } from "@/lib/utils";
+
+export interface MentionedDocument {
+	id: number;
+	title: string;
+	document_type?: string;
+}
+
+export interface InlineMentionEditorRef {
+	focus: () => void;
+	clear: () => void;
+	getText: () => string;
+	getMentionedDocuments: () => MentionedDocument[];
+	insertDocumentChip: (doc: Document) => void;
+}
+
+interface InlineMentionEditorProps {
+	placeholder?: string;
+	onMentionTrigger?: (query: string) => void;
+	onMentionClose?: () => void;
+	onSubmit?: () => void;
+	onChange?: (text: string, docs: MentionedDocument[]) => void;
+	onDocumentRemove?: (docId: number) => void;
+	onKeyDown?: (e: React.KeyboardEvent) => void;
+	disabled?: boolean;
+	className?: string;
+	initialDocuments?: MentionedDocument[];
+}
+
+// Unique data attribute to identify chip elements
+const CHIP_DATA_ATTR = "data-mention-chip";
+const CHIP_ID_ATTR = "data-mention-id";
+
+/**
+ * Type guard to check if a node is a chip element
+ */
+function isChipElement(node: Node | null): node is HTMLSpanElement {
+	return (
+		node !== null &&
+		node.nodeType === Node.ELEMENT_NODE &&
+		(node as Element).hasAttribute(CHIP_DATA_ATTR)
+	);
+}
+
+/**
+ * Safely parse chip ID from element attribute
+ */
+function getChipId(element: Element): number | null {
+	const idStr = element.getAttribute(CHIP_ID_ATTR);
+	if (!idStr) return null;
+	const id = parseInt(idStr, 10);
+	return Number.isNaN(id) ? null : id;
+}
+
+export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMentionEditorProps>(
+	(
+		{
+			placeholder = "Type @ to mention documents...",
+			onMentionTrigger,
+			onMentionClose,
+			onSubmit,
+			onChange,
+			onDocumentRemove,
+			onKeyDown,
+			disabled = false,
+			className,
+			initialDocuments = [],
+		},
+		ref
+	) => {
+		const editorRef = useRef<HTMLDivElement>(null);
+		const [isEmpty, setIsEmpty] = useState(true);
+		const [mentionedDocs, setMentionedDocs] = useState<Map<number, MentionedDocument>>(
+			() => new Map(initialDocuments.map((d) => [d.id, d]))
+		);
+		const isComposingRef = useRef(false);
+
+		// Sync initial documents
+		useEffect(() => {
+			if (initialDocuments.length > 0) {
+				setMentionedDocs(new Map(initialDocuments.map((d) => [d.id, d])));
+			}
+		}, [initialDocuments]);
+
+		// Focus at the end of the editor
+		const focusAtEnd = useCallback(() => {
+			if (!editorRef.current) return;
+			editorRef.current.focus();
+			const selection = window.getSelection();
+			const range = document.createRange();
+			range.selectNodeContents(editorRef.current);
+			range.collapse(false);
+			selection?.removeAllRanges();
+			selection?.addRange(range);
+		}, []);
+
+		// Get plain text content (excluding chips)
+		const getText = useCallback((): string => {
+			if (!editorRef.current) return "";
+
+			let text = "";
+			const walker = document.createTreeWalker(
+				editorRef.current,
+				NodeFilter.SHOW_TEXT | NodeFilter.SHOW_ELEMENT,
+				{
+					acceptNode: (node) => {
+						// Skip chip elements entirely
+						if (node.nodeType === Node.ELEMENT_NODE) {
+							const el = node as Element;
+							if (el.hasAttribute(CHIP_DATA_ATTR)) {
+								return NodeFilter.FILTER_REJECT; // Skip this subtree
+							}
+							return NodeFilter.FILTER_SKIP; // Continue into children
+						}
+						return NodeFilter.FILTER_ACCEPT;
+					},
+				}
+			);
+
+			let node: Node | null = walker.nextNode();
+			while (node) {
+				if (node.nodeType === Node.TEXT_NODE) {
+					text += node.textContent;
+				}
+				node = walker.nextNode();
+			}
+
+			return text.trim();
+		}, []);
+
+		// Get all mentioned documents
+		const getMentionedDocuments = useCallback((): MentionedDocument[] => {
+			return Array.from(mentionedDocs.values());
+		}, [mentionedDocs]);
+
+		// Create a chip element for a document
+		const createChipElement = useCallback(
+			(doc: MentionedDocument): HTMLSpanElement => {
+				const chip = document.createElement("span");
+				chip.setAttribute(CHIP_DATA_ATTR, "true");
+				chip.setAttribute(CHIP_ID_ATTR, String(doc.id));
+				chip.contentEditable = "false";
+				chip.className =
+					"inline-flex items-center gap-0.5 mx-0.5 pl-1 pr-0.5 py-0.5 rounded bg-primary/10 text-xs font-bold text-primary border border-primary/10 select-none";
+				chip.style.userSelect = "none";
+				chip.style.verticalAlign = "baseline";
+
+				const titleSpan = document.createElement("span");
+				titleSpan.className = "max-w-[80px] truncate";
+				titleSpan.textContent = doc.title;
+				titleSpan.title = doc.title;
+
+				const removeBtn = document.createElement("button");
+				removeBtn.type = "button";
+				removeBtn.className =
+					"size-3 flex items-center justify-center rounded-full hover:bg-primary/20 transition-colors ml-0.5";
+				removeBtn.innerHTML = ReactDOMServer.renderToString(
+					createElement(X, { className: "h-2.5 w-2.5", strokeWidth: 2.5 })
+				);
+				removeBtn.onclick = (e) => {
+					e.preventDefault();
+					e.stopPropagation();
+					chip.remove();
+					setMentionedDocs((prev) => {
+						const next = new Map(prev);
+						next.delete(doc.id);
+						return next;
+					});
+					// Notify parent that a document was removed
+					onDocumentRemove?.(doc.id);
+					focusAtEnd();
+				};
+
+				chip.appendChild(titleSpan);
+				chip.appendChild(removeBtn);
+
+				return chip;
+			},
+			[focusAtEnd, onDocumentRemove]
+		);
+
+		// Insert a document chip at the current cursor position
+		const insertDocumentChip = useCallback(
+			(doc: Document) => {
+				if (!editorRef.current) return;
+
+				// Validate required fields for type safety
+				if (typeof doc.id !== "number" || typeof doc.title !== "string") {
+					console.warn("[InlineMentionEditor] Invalid document passed to insertDocumentChip:", doc);
+					return;
+				}
+
+				const mentionDoc: MentionedDocument = {
+					id: doc.id,
+					title: doc.title,
+					document_type: doc.document_type,
+				};
+
+				// Add to mentioned docs map
+				setMentionedDocs((prev) => new Map(prev).set(doc.id, mentionDoc));
+
+				// Find and remove the @query text
+				const selection = window.getSelection();
+				if (!selection || selection.rangeCount === 0) {
+					// No selection, just append
+					const chip = createChipElement(mentionDoc);
+					editorRef.current.appendChild(chip);
+					editorRef.current.appendChild(document.createTextNode(" "));
+					focusAtEnd();
+					return;
+				}
+
+				// Find the @ symbol before the cursor and remove it along with any query text
+				const range = selection.getRangeAt(0);
+				const textNode = range.startContainer;
+
+				if (textNode.nodeType === Node.TEXT_NODE) {
+					const text = textNode.textContent || "";
+					const cursorPos = range.startOffset;
+
+					// Find the @ symbol before cursor
+					let atIndex = -1;
+					for (let i = cursorPos - 1; i >= 0; i--) {
+						if (text[i] === "@") {
+							atIndex = i;
+							break;
+						}
+					}
+
+					if (atIndex !== -1) {
+						// Remove @query and insert chip
+						const beforeAt = text.slice(0, atIndex);
+						const afterCursor = text.slice(cursorPos);
+
+						// Create chip
+						const chip = createChipElement(mentionDoc);
+
+						// Replace text node content
+						const parent = textNode.parentNode;
+						if (parent) {
+							const beforeNode = document.createTextNode(beforeAt);
+							const afterNode = document.createTextNode(` ${afterCursor}`);
+
+							parent.insertBefore(beforeNode, textNode);
+							parent.insertBefore(chip, textNode);
+							parent.insertBefore(afterNode, textNode);
+							parent.removeChild(textNode);
+
+							// Set cursor after the chip
+							const newRange = document.createRange();
+							newRange.setStart(afterNode, 1);
+							newRange.collapse(true);
+							selection.removeAllRanges();
+							selection.addRange(newRange);
+						}
+					} else {
+						// No @ found, just insert at cursor
+						const chip = createChipElement(mentionDoc);
+						range.insertNode(chip);
+						range.setStartAfter(chip);
+						range.collapse(true);
+
+						// Add space after chip
+						const space = document.createTextNode(" ");
+						range.insertNode(space);
+						range.setStartAfter(space);
+						range.collapse(true);
+					}
+				} else {
+					// Not in a text node, append to editor
+					const chip = createChipElement(mentionDoc);
+					editorRef.current.appendChild(chip);
+					editorRef.current.appendChild(document.createTextNode(" "));
+					focusAtEnd();
+				}
+
+				// Update empty state
+				setIsEmpty(false);
+
+				// Trigger onChange
+				if (onChange) {
+					setTimeout(() => {
+						onChange(getText(), getMentionedDocuments());
+					}, 0);
+				}
+			},
+			[createChipElement, focusAtEnd, getText, getMentionedDocuments, onChange]
+		);
+
+		// Clear the editor
+		const clear = useCallback(() => {
+			if (editorRef.current) {
+				editorRef.current.innerHTML = "";
+				setIsEmpty(true);
+				setMentionedDocs(new Map());
+			}
+		}, []);
+
+		// Expose methods via ref
+		useImperativeHandle(ref, () => ({
+			focus: () => editorRef.current?.focus(),
+			clear,
+			getText,
+			getMentionedDocuments,
+			insertDocumentChip,
+		}));
+
+		// Handle input changes
+		const handleInput = useCallback(() => {
+			if (!editorRef.current) return;
+
+			const text = getText();
+			const empty = text.length === 0 && mentionedDocs.size === 0;
+			setIsEmpty(empty);
+
+			// Check for @ mentions
+			const selection = window.getSelection();
+			let shouldTriggerMention = false;
+			let mentionQuery = "";
+
+			if (selection && selection.rangeCount > 0) {
+				const range = selection.getRangeAt(0);
+				const textNode = range.startContainer;
+
+				if (textNode.nodeType === Node.TEXT_NODE) {
+					const textContent = textNode.textContent || "";
+					const cursorPos = range.startOffset;
+
+					// Look for @ before cursor
+					let atIndex = -1;
+					for (let i = cursorPos - 1; i >= 0; i--) {
+						if (textContent[i] === "@") {
+							atIndex = i;
+							break;
+						}
+						// Stop if we hit a space (@ must be at word boundary)
+						if (textContent[i] === " " || textContent[i] === "\n") {
+							break;
+						}
+					}
+
+					if (atIndex !== -1) {
+						const query = textContent.slice(atIndex + 1, cursorPos);
+						// Only trigger if query doesn't start with space
+						if (!query.startsWith(" ")) {
+							shouldTriggerMention = true;
+							mentionQuery = query;
+						}
+					}
+				}
+			}
+
+			// If no @ found before cursor, check if text contains @ at all
+			// If text is empty or doesn't contain @, close the mention
+			if (!shouldTriggerMention) {
+				if (text.length === 0 || !text.includes("@")) {
+					onMentionClose?.();
+				} else {
+					// Text contains @ but not before cursor, close mention
+					onMentionClose?.();
+				}
+			} else {
+				onMentionTrigger?.(mentionQuery);
+			}
+
+			// Notify parent of change
+			onChange?.(text, Array.from(mentionedDocs.values()));
+		}, [getText, mentionedDocs, onChange, onMentionTrigger, onMentionClose]);
+
+		// Handle keydown
+		const handleKeyDown = useCallback(
+			(e: React.KeyboardEvent<HTMLDivElement>) => {
+				// Let parent handle navigation keys when mention popover is open
+				if (onKeyDown) {
+					onKeyDown(e);
+					if (e.defaultPrevented) return;
+				}
+
+				// Handle Enter for submit (without shift)
+				if (e.key === "Enter" && !e.shiftKey) {
+					e.preventDefault();
+					onSubmit?.();
+					return;
+				}
+
+				// Handle backspace on chips
+				if (e.key === "Backspace") {
+					const selection = window.getSelection();
+					if (selection && selection.rangeCount > 0) {
+						const range = selection.getRangeAt(0);
+						if (range.collapsed) {
+							// Check if cursor is right after a chip
+							const node = range.startContainer;
+							const offset = range.startOffset;
+
+							if (node.nodeType === Node.TEXT_NODE && offset === 0) {
+								// Check previous sibling using type guard
+								const prevSibling = node.previousSibling;
+								if (isChipElement(prevSibling)) {
+									e.preventDefault();
+									const chipId = getChipId(prevSibling);
+									if (chipId !== null) {
+										prevSibling.remove();
+										setMentionedDocs((prev) => {
+											const next = new Map(prev);
+											next.delete(chipId);
+											return next;
+										});
+										// Notify parent that a document was removed
+										onDocumentRemove?.(chipId);
+									}
+									return;
+								}
+								// Check if we're about to delete @ at the start
+								const textContent = node.textContent || "";
+								if (textContent.length > 0 && textContent[0] === "@") {
+									// Will delete @, close mention popover
+									setTimeout(() => {
+										onMentionClose?.();
+									}, 0);
+								}
+							} else if (node.nodeType === Node.TEXT_NODE && offset > 0) {
+								// Check if we're about to delete @
+								const textContent = node.textContent || "";
+								if (textContent[offset - 1] === "@") {
+									// Will delete @, close mention popover
+									setTimeout(() => {
+										onMentionClose?.();
+									}, 0);
+								}
+							} else if (node.nodeType === Node.ELEMENT_NODE && offset > 0) {
+								// Check if previous child is a chip using type guard
+								const prevChild = (node as Element).childNodes[offset - 1];
+								if (isChipElement(prevChild)) {
+									e.preventDefault();
+									const chipId = getChipId(prevChild);
+									if (chipId !== null) {
+										prevChild.remove();
+										setMentionedDocs((prev) => {
+											const next = new Map(prev);
+											next.delete(chipId);
+											return next;
+										});
+										// Notify parent that a document was removed
+										onDocumentRemove?.(chipId);
+									}
+								}
+							}
+						}
+					}
+				}
+			},
+			[onKeyDown, onSubmit, onDocumentRemove, onMentionClose]
+		);
+
+		// Handle paste - strip formatting
+		const handlePaste = useCallback((e: React.ClipboardEvent) => {
+			e.preventDefault();
+			const text = e.clipboardData.getData("text/plain");
+			document.execCommand("insertText", false, text);
+		}, []);
+
+		// Handle composition (for IME input)
+		const handleCompositionStart = useCallback(() => {
+			isComposingRef.current = true;
+		}, []);
+
+		const handleCompositionEnd = useCallback(() => {
+			isComposingRef.current = false;
+			handleInput();
+		}, [handleInput]);
+
+		return (
+			<div className="relative w-full">
+				{/** biome-ignore lint/a11y/useSemanticElements: <not important> */}
+				<div
+					ref={editorRef}
+					contentEditable={!disabled}
+					suppressContentEditableWarning
+					tabIndex={disabled ? -1 : 0}
+					onInput={handleInput}
+					onKeyDown={handleKeyDown}
+					onPaste={handlePaste}
+					onCompositionStart={handleCompositionStart}
+					onCompositionEnd={handleCompositionEnd}
+					className={cn(
+						"min-h-[24px] max-h-32 overflow-y-auto",
+						"text-sm outline-none",
+						"whitespace-pre-wrap break-words",
+						disabled && "opacity-50 cursor-not-allowed",
+						className
+					)}
+					style={{ wordBreak: "break-word" }}
+					data-placeholder={placeholder}
+					aria-label="Message input with inline mentions"
+					role="textbox"
+					aria-multiline="true"
+				/>
+				{/* Placeholder */}
+				{isEmpty && (
+					<div
+						className="absolute top-0 left-0 pointer-events-none text-muted-foreground text-sm"
+						aria-hidden="true"
+					>
+						{placeholder}
+					</div>
+				)}
+			</div>
+		);
+	}
+);
+
+InlineMentionEditor.displayName = "InlineMentionEditor";

surfsense_web/components/assistant-ui/thread.tsx

@@ -7,7 +7,7 @@ import {
 	MessagePrimitive,
 	ThreadPrimitive,
 	useAssistantState,
-	useMessage,
+	useComposerRuntime,
 	useThreadViewport,
 } from "@assistant-ui/react";
 import { useAtom, useAtomValue, useSetAtom } from "jotai";
@@ -31,7 +31,6 @@ import {
 	Search,
 	Sparkles,
 	SquareIcon,
-	X,
 } from "lucide-react";
 import Link from "next/link";
 import { useParams } from "next/navigation";
@@ -65,13 +64,17 @@ import {
 	ComposerAttachments,
 	UserMessageAttachments,
 } from "@/components/assistant-ui/attachment";
+import {
+	InlineMentionEditor,
+	type InlineMentionEditorRef,
+} from "@/components/assistant-ui/inline-mention-editor";
 import { MarkdownText } from "@/components/assistant-ui/markdown-text";
 import { ToolFallback } from "@/components/assistant-ui/tool-fallback";
 import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
 import {
-	DocumentsDataTable,
-	type DocumentsDataTableRef,
-} from "@/components/new-chat/DocumentsDataTable";
+	DocumentMentionPicker,
+	type DocumentMentionPickerRef,
+} from "@/components/new-chat/document-mention-picker";
 import {
 	ChainOfThought,
 	ChainOfThoughtContent,
@@ -144,15 +147,6 @@ const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?: boolea
 		return step.status;
 	};
 
-	// Check if any step is effectively in progress
-	const hasInProgressStep = steps.some((step) => getEffectiveStatus(step) === "in_progress");
-
-	// Find the last completed step index (using effective status)
-	const lastCompletedIndex = steps
-		.map((s, i) => (getEffectiveStatus(s) === "completed" ? i : -1))
-		.filter((i) => i !== -1)
-		.pop();
-
 	// Clear manual overrides when a step's status changes
 	useEffect(() => {
 		const currentStatuses: Record<string, string> = {};
@@ -172,7 +166,7 @@ const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?: boolea
 
 	if (steps.length === 0) return null;
 
-	const getStepOpenState = (step: ThinkingStep, index: number): boolean => {
+	const getStepOpenState = (step: ThinkingStep): boolean => {
 		const effectiveStatus = getEffectiveStatus(step);
 		// If user has manually toggled, respect that
 		if (manualOverrides[step.id] !== undefined) {
@@ -182,11 +176,7 @@ const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?: boolea
 		if (effectiveStatus === "in_progress") {
 			return true;
 		}
-		// Auto behavior: keep last completed step open if no in-progress step
-		if (!hasInProgressStep && index === lastCompletedIndex) {
-			return true;
-		}
-		// Default: collapsed
+		// Default: collapsed (all steps collapse when processing is done)
 		return false;
 	};
 
@@ -200,10 +190,10 @@ const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?: boolea
 	return (
 		<div className="mx-auto w-full max-w-(--thread-max-width) px-2 py-2">
 			<ChainOfThought>
-				{steps.map((step, index) => {
+				{steps.map((step) => {
 					const effectiveStatus = getEffectiveStatus(step);
 					const icon = getStepIcon(effectiveStatus, step.title);
-					const isOpen = getStepOpenState(step, index);
+					const isOpen = getStepOpenState(step);
 					return (
 						<ChainOfThoughtStep
 							key={step.id}
@@ -240,7 +230,7 @@ const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?: boolea
  * Uses useThreadViewport to scroll to bottom when thinking steps change,
  * ensuring the user always sees the latest content during streaming.
  */
-const ThinkingStepsScrollHandler: FC = () => {
+const _ThinkingStepsScrollHandler: FC = () => {
 	const thinkingStepsMap = useContext(ThinkingStepsContext);
 	const viewport = useThreadViewport();
 	const isRunning = useAssistantState(({ thread }) => thread.isRunning);
@@ -351,7 +341,7 @@ const getTimeBasedGreeting = (userEmail?: string): string => {
 		: null;
 
 	// Array of greeting variations for each time period
-	const morningGreetings = ["Good morning", "Rise and shine", "Morning", "Hey there"];
+	const morningGreetings = ["Good morning", "Fresh start today", "Morning", "Hey there"];
 
 	const afternoonGreetings = ["Good afternoon", "Afternoon", "Hey there", "Hi there"];
 
@@ -359,7 +349,7 @@ const getTimeBasedGreeting = (userEmail?: string): string => {
 
 	const nightGreetings = ["Good night", "Evening", "Hey there", "Winding down"];
 
-	const lateNightGreetings = ["Still up", "Night owl mode", "The night is young", "Hi there"];
+	const lateNightGreetings = ["Still up", "Night owl mode", "Up past bedtime", "Hi there"];
 
 	// Select a random greeting based on time
 	let greeting: string;
@@ -412,177 +402,173 @@ const Composer: FC = () => {
 	const [mentionedDocuments, setMentionedDocuments] = useAtom(mentionedDocumentsAtom);
 	const [showDocumentPopover, setShowDocumentPopover] = useState(false);
 	const [mentionQuery, setMentionQuery] = useState("");
-	const inputRef = useRef<HTMLTextAreaElement | null>(null);
-	const documentPickerRef = useRef<DocumentsDataTableRef>(null);
+	const editorRef = useRef<InlineMentionEditorRef>(null);
+	const editorContainerRef = useRef<HTMLDivElement>(null);
+	const documentPickerRef = useRef<DocumentMentionPickerRef>(null);
 	const { search_space_id } = useParams();
 	const setMentionedDocumentIds = useSetAtom(mentionedDocumentIdsAtom);
+	const composerRuntime = useComposerRuntime();
+	const hasAutoFocusedRef = useRef(false);
+
+	// 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) {
+			// Small delay to ensure the editor is fully mounted
+			const timeoutId = setTimeout(() => {
+				editorRef.current?.focus();
+				hasAutoFocusedRef.current = true;
+			}, 100);
+			return () => clearTimeout(timeoutId);
+		}
+	}, [isThreadEmpty]);
+
+	// Reset auto-focus flag when thread becomes non-empty (user sent a message)
+	useEffect(() => {
+		if (!isThreadEmpty) {
+			hasAutoFocusedRef.current = false;
+		}
+	}, [isThreadEmpty]);
 
 	// Sync mentioned document IDs to atom for use in chat request
 	useEffect(() => {
 		setMentionedDocumentIds(mentionedDocuments.map((doc) => doc.id));
 	}, [mentionedDocuments, setMentionedDocumentIds]);
 
-	// Extract mention query (text after @)
-	const extractMentionQuery = useCallback((value: string): string => {
-		const atIndex = value.lastIndexOf("@");
-		if (atIndex === -1) return "";
-		return value.slice(atIndex + 1);
+	// Handle text change from inline editor - sync with assistant-ui composer
+	const handleEditorChange = useCallback(
+		(text: string) => {
+			composerRuntime.setText(text);
+		},
+		[composerRuntime]
+	);
+
+	// Handle @ mention trigger from inline editor
+	const handleMentionTrigger = useCallback((query: string) => {
+		setShowDocumentPopover(true);
+		setMentionQuery(query);
 	}, []);
 
-	const handleKeyUp = (e: React.KeyboardEvent<HTMLTextAreaElement>) => {
-		const textarea = e.currentTarget;
-		const value = textarea.value;
-
-		// Open document picker when user types '@'
-		if (e.key === "@" || (e.key === "2" && e.shiftKey)) {
-			setShowDocumentPopover(true);
-			setMentionQuery("");
-			return;
-		}
-
-		// Check if value contains @ and extract query
-		if (value.includes("@")) {
-			const query = extractMentionQuery(value);
-
-			// Close popup if query starts with space (user typed "@ ")
-			if (query.startsWith(" ")) {
-				setShowDocumentPopover(false);
-				setMentionQuery("");
-				return;
-			}
-
-			// Reopen popup if @ is present and query doesn't start with space
-			// (handles case where user deleted the space after @)
-			if (!showDocumentPopover) {
-				setShowDocumentPopover(true);
-			}
-			setMentionQuery(query);
-		} else {
-			// Close popover if '@' is no longer in the input (user deleted it)
-			if (showDocumentPopover) {
-				setShowDocumentPopover(false);
-				setMentionQuery("");
-			}
-		}
-	};
-
-	const handleKeyDown = (e: React.KeyboardEvent<HTMLTextAreaElement>) => {
-		// When popup is open, handle navigation keys
+	// Handle mention close
+	const handleMentionClose = useCallback(() => {
 		if (showDocumentPopover) {
-			if (e.key === "ArrowDown") {
-				e.preventDefault();
-				documentPickerRef.current?.moveDown();
-				return;
-			}
-			if (e.key === "ArrowUp") {
-				e.preventDefault();
-				documentPickerRef.current?.moveUp();
-				return;
-			}
-			if (e.key === "Enter") {
-				e.preventDefault();
-				documentPickerRef.current?.selectHighlighted();
-				return;
-			}
-			if (e.key === "Escape") {
-				e.preventDefault();
-				setShowDocumentPopover(false);
-				setMentionQuery("");
-				return;
-			}
+			setShowDocumentPopover(false);
+			setMentionQuery("");
 		}
+	}, [showDocumentPopover]);
 
-		// Remove last document chip when pressing backspace at the beginning of input
-		if (e.key === "Backspace" && mentionedDocuments.length > 0) {
-			const textarea = e.currentTarget;
-			const selectionStart = textarea.selectionStart;
-			const selectionEnd = textarea.selectionEnd;
-
-			// Only remove chip if cursor is at position 0 and nothing is selected
-			if (selectionStart === 0 && selectionEnd === 0) {
-				e.preventDefault();
-				// Remove the last document chip
-				setMentionedDocuments((prev) => prev.slice(0, -1));
-			}
-		}
-	};
-
-	const handleDocumentsMention = (documents: Document[]) => {
-		// Update mentioned documents (merge with existing, avoid duplicates)
-		setMentionedDocuments((prev) => {
-			const existingIds = new Set(prev.map((d) => d.id));
-			const newDocs = documents.filter((doc) => !existingIds.has(doc.id));
-			return [...prev, ...newDocs];
-		});
-
-		// Clean up the '@...' mention text from input
-		if (inputRef.current) {
-			const input = inputRef.current;
-			const currentValue = input.value;
-			const atIndex = currentValue.lastIndexOf("@");
-
-			if (atIndex !== -1) {
-				// Remove @ and everything after it
-				const newValue = currentValue.slice(0, atIndex);
-				const nativeInputValueSetter = Object.getOwnPropertyDescriptor(
-					window.HTMLTextAreaElement.prototype,
-					"value"
-				)?.set;
-				if (nativeInputValueSetter) {
-					nativeInputValueSetter.call(input, newValue);
-					input.dispatchEvent(new Event("input", { bubbles: true }));
+	// Handle keyboard navigation when popover is open
+	const handleKeyDown = useCallback(
+		(e: React.KeyboardEvent) => {
+			if (showDocumentPopover) {
+				if (e.key === "ArrowDown") {
+					e.preventDefault();
+					documentPickerRef.current?.moveDown();
+					return;
+				}
+				if (e.key === "ArrowUp") {
+					e.preventDefault();
+					documentPickerRef.current?.moveUp();
+					return;
+				}
+				if (e.key === "Enter") {
+					e.preventDefault();
+					documentPickerRef.current?.selectHighlighted();
+					return;
+				}
+				if (e.key === "Escape") {
+					e.preventDefault();
+					setShowDocumentPopover(false);
+					setMentionQuery("");
+					return;
 				}
 			}
-			// Focus the input so user can continue typing
-			input.focus();
+		},
+		[showDocumentPopover]
+	);
+
+	// 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
+			editorRef.current?.clear();
+			setMentionedDocuments([]);
+			setMentionedDocumentIds([]);
+		}
+	}, [
+		showDocumentPopover,
+		isThreadRunning,
+		composerRuntime,
+		setMentionedDocuments,
+		setMentionedDocumentIds,
+	]);
 
-		// Reset mention query
-		setMentionQuery("");
-	};
+	// Handle document removal from inline editor
+	const handleDocumentRemove = useCallback(
+		(docId: number) => {
+			setMentionedDocuments((prev) => {
+				const updated = prev.filter((doc) => doc.id !== docId);
+				// Immediately sync document IDs to avoid race conditions
+				setMentionedDocumentIds(updated.map((doc) => doc.id));
+				return updated;
+			});
+		},
+		[setMentionedDocuments, setMentionedDocumentIds]
+	);
 
-	const handleRemoveDocument = (docId: number) => {
-		setMentionedDocuments((prev) => prev.filter((doc) => doc.id !== docId));
-	};
+	// Handle document selection from picker
+	const handleDocumentsMention = useCallback(
+		(documents: Document[]) => {
+			// Insert chips into the inline editor for each new document
+			const existingIds = new Set(mentionedDocuments.map((d) => d.id));
+			const newDocs = documents.filter((doc) => !existingIds.has(doc.id));
+
+			for (const doc of newDocs) {
+				editorRef.current?.insertDocumentChip(doc);
+			}
+
+			// Update mentioned documents state
+			setMentionedDocuments((prev) => {
+				const existingIdSet = new Set(prev.map((d) => d.id));
+				const uniqueNewDocs = documents.filter((doc) => !existingIdSet.has(doc.id));
+				const updated = [...prev, ...uniqueNewDocs];
+				// Immediately sync document IDs to avoid race conditions
+				setMentionedDocumentIds(updated.map((doc) => doc.id));
+				return updated;
+			});
+
+			// Reset mention query but keep popover open for more selections
+			setMentionQuery("");
+		},
+		[mentionedDocuments, setMentionedDocuments, setMentionedDocumentIds]
+	);
 
 	return (
 		<ComposerPrimitive.Root className="aui-composer-root relative flex w-full flex-col">
 			<ComposerPrimitive.AttachmentDropzone className="aui-composer-attachment-dropzone flex w-full flex-col rounded-2xl border-input bg-muted px-1 pt-2 outline-none transition-shadow data-[dragging=true]:border-ring data-[dragging=true]:border-dashed data-[dragging=true]:bg-accent/50">
 				<ComposerAttachments />
-				{/* -------- Input field with inline document chips -------- */}
-				<div className="aui-composer-input-wrapper flex flex-wrap items-center gap-1.5 px-3 pt-2 pb-6">
-					{/* Inline document chips */}
-					{mentionedDocuments.map((doc) => (
-						<span
-							key={doc.id}
-							className="inline-flex items-center gap-1 pl-2 pr-1 py-0.5 rounded-full bg-primary/10 text-xs font-medium text-primary border border-primary/20 shrink-0"
-							title={doc.title}
-						>
-							<span className="max-w-[120px] truncate">{doc.title}</span>
-							<button
-								type="button"
-								onClick={() => handleRemoveDocument(doc.id)}
-								className="size-4 flex items-center justify-center rounded-full hover:bg-primary/20 transition-colors"
-								aria-label={`Remove ${doc.title}`}
-							>
-								<X className="size-3" />
-							</button>
-						</span>
-					))}
-					{/* Text input */}
-					<ComposerPrimitive.Input
-						ref={inputRef}
-						onKeyUp={handleKeyUp}
+				{/* -------- Inline Mention Editor -------- */}
+				<div ref={editorContainerRef} className="aui-composer-input-wrapper px-3 pt-3 pb-6">
+					<InlineMentionEditor
+						ref={editorRef}
+						placeholder="Ask SurfSense (type @ to mention docs)"
+						onMentionTrigger={handleMentionTrigger}
+						onMentionClose={handleMentionClose}
+						onChange={handleEditorChange}
+						onDocumentRemove={handleDocumentRemove}
+						onSubmit={handleSubmit}
 						onKeyDown={handleKeyDown}
-						placeholder={
-							mentionedDocuments.length > 0
-								? "Ask about these documents..."
-								: "Ask SurfSense (type @ to mention docs)"
-						}
-						className="aui-composer-input flex-1 min-w-[120px] max-h-32 resize-none bg-transparent text-sm outline-none placeholder:text-muted-foreground focus-visible:ring-0 py-1"
-						rows={1}
-						autoFocus
-						aria-label="Message input"
+						className="min-h-[24px]"
 					/>
 				</div>
 
@@ -601,19 +587,18 @@ const Composer: FC = () => {
 							/>
 							{/* Popover positioned above input */}
 							<div
-								className="fixed shadow-2xl rounded-lg border border-border overflow-hidden"
+								className="fixed shadow-2xl rounded-lg border border-border overflow-hidden bg-popover"
 								style={{
 									zIndex: 9999,
-									backgroundColor: "#18181b",
-									bottom: inputRef.current
-										? `${window.innerHeight - inputRef.current.getBoundingClientRect().top + 8}px`
+									bottom: editorContainerRef.current
+										? `${window.innerHeight - editorContainerRef.current.getBoundingClientRect().top + 8}px`
 										: "200px",
-									left: inputRef.current
-										? `${inputRef.current.getBoundingClientRect().left}px`
+									left: editorContainerRef.current
+										? `${editorContainerRef.current.getBoundingClientRect().left}px`
 										: "50%",
 								}}
 							>
-								<DocumentsDataTable
+								<DocumentMentionPicker
 									ref={documentPickerRef}
 									searchSpaceId={Number(search_space_id)}
 									onSelectionChange={handleDocumentsMention}
@@ -987,19 +972,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 relative 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) => (
+			<div className="aui-user-message-content-wrapper col-start-2 min-w-0">
+				{/* 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"
@@ -1011,11 +1000,14 @@ const UserMessage: FC = () => {
 						))}
 					</div>
 				)}
-				<div className="aui-user-message-content wrap-break-word rounded-2xl bg-muted px-4 py-2.5 text-foreground">
-					<MessagePrimitive.Parts />
-				</div>
-				<div className="aui-user-action-bar-wrapper -translate-x-full -translate-y-full absolute top-full left-0 pr-2">
-					<UserActionBar />
+				{/* Message bubble with action bar positioned relative to it */}
+				<div className="relative">
+					<div className="aui-user-message-content wrap-break-word rounded-2xl bg-muted px-4 py-2.5 text-foreground">
+						<MessagePrimitive.Parts />
+					</div>
+					<div className="aui-user-action-bar-wrapper absolute top-1/2 right-full -translate-y-1/2 pr-1">
+						<UserActionBar />
+					</div>
 				</div>
 			</div>
 

surfsense_web/components/dashboard-breadcrumb.tsx

@@ -165,14 +165,11 @@ export function DashboardBreadcrumb() {
 					}
 
 					// Handle new-chat sub-sections (thread IDs)
+					// Don't show thread ID in breadcrumb - users identify chats by content, not by ID
 					if (section === "new-chat") {
 						breadcrumbs.push({
 							label: t("chat") || "Chat",
-							href: `/dashboard/${segments[1]}/new-chat`,
 						});
-						if (subSection) {
-							breadcrumbs.push({ label: `Thread ${subSection}` });
-						}
 						return breadcrumbs;
 					}
 

surfsense_web/components/new-chat/DocumentsDataTable.tsx

@@ -1,248 +0,0 @@
-"use client";
-
-import { useQuery } from "@tanstack/react-query";
-import { FileText } from "lucide-react";
-import {
-	forwardRef,
-	useCallback,
-	useEffect,
-	useImperativeHandle,
-	useMemo,
-	useRef,
-	useState,
-} from "react";
-import { getConnectorIcon } from "@/contracts/enums/connectorIcons";
-import type { Document } from "@/contracts/types/document.types";
-import { documentsApiService } from "@/lib/apis/documents-api.service";
-import { cacheKeys } from "@/lib/query-client/cache-keys";
-import { cn } from "@/lib/utils";
-
-export interface DocumentsDataTableRef {
-	selectHighlighted: () => void;
-	moveUp: () => void;
-	moveDown: () => void;
-}
-
-interface DocumentsDataTableProps {
-	searchSpaceId: number;
-	onSelectionChange: (documents: Document[]) => void;
-	onDone: () => void;
-	initialSelectedDocuments?: Document[];
-	externalSearch?: string;
-}
-
-function useDebounced<T>(value: T, delay = 300) {
-	const [debounced, setDebounced] = useState(value);
-	useEffect(() => {
-		const t = setTimeout(() => setDebounced(value), delay);
-		return () => clearTimeout(t);
-	}, [value, delay]);
-	return debounced;
-}
-
-export const DocumentsDataTable = forwardRef<DocumentsDataTableRef, DocumentsDataTableProps>(
-	function DocumentsDataTable(
-		{
-			searchSpaceId,
-			onSelectionChange,
-			onDone,
-			initialSelectedDocuments = [],
-			externalSearch = "",
-		},
-		ref
-	) {
-		// Use external search
-		const search = externalSearch;
-		const debouncedSearch = useDebounced(search, 150);
-		const [highlightedIndex, setHighlightedIndex] = useState(0);
-		const itemRefs = useRef<Map<number, HTMLButtonElement>>(new Map());
-
-		const fetchQueryParams = useMemo(
-			() => ({
-				search_space_id: searchSpaceId,
-				page: 0,
-				page_size: 20,
-			}),
-			[searchSpaceId]
-		);
-
-		const searchQueryParams = useMemo(() => {
-			return {
-				search_space_id: searchSpaceId,
-				page: 0,
-				page_size: 20,
-				title: debouncedSearch,
-			};
-		}, [debouncedSearch, searchSpaceId]);
-
-		// Use query for fetching documents
-		const { data: documents, isLoading: isDocumentsLoading } = useQuery({
-			queryKey: cacheKeys.documents.withQueryParams(fetchQueryParams),
-			queryFn: () => documentsApiService.getDocuments({ queryParams: fetchQueryParams }),
-			staleTime: 3 * 60 * 1000,
-			enabled: !!searchSpaceId && !debouncedSearch.trim(),
-		});
-
-		// Searching
-		const { data: searchedDocuments, isLoading: isSearchedDocumentsLoading } = useQuery({
-			queryKey: cacheKeys.documents.withQueryParams(searchQueryParams),
-			queryFn: () => documentsApiService.searchDocuments({ queryParams: searchQueryParams }),
-			staleTime: 3 * 60 * 1000,
-			enabled: !!searchSpaceId && !!debouncedSearch.trim(),
-		});
-
-		const actualDocuments = debouncedSearch.trim()
-			? searchedDocuments?.items || []
-			: documents?.items || [];
-		const actualLoading = debouncedSearch.trim() ? isSearchedDocumentsLoading : isDocumentsLoading;
-
-		// Track already selected document IDs
-		const selectedIds = useMemo(
-			() => new Set(initialSelectedDocuments.map((d) => d.id)),
-			[initialSelectedDocuments]
-		);
-
-		// Filter out already selected documents for navigation
-		const selectableDocuments = useMemo(
-			() => actualDocuments.filter((doc) => !selectedIds.has(doc.id)),
-			[actualDocuments, selectedIds]
-		);
-
-		const handleSelectDocument = useCallback(
-			(doc: Document) => {
-				onSelectionChange([...initialSelectedDocuments, doc]);
-				onDone();
-			},
-			[initialSelectedDocuments, onSelectionChange, onDone]
-		);
-
-		// Scroll highlighted item into view
-		useEffect(() => {
-			const item = itemRefs.current.get(highlightedIndex);
-			if (item) {
-				item.scrollIntoView({ block: "nearest", behavior: "smooth" });
-			}
-		}, [highlightedIndex]);
-
-		// Reset highlighted index when external search changes
-		const prevSearchRef = useRef(search);
-		if (prevSearchRef.current !== search) {
-			prevSearchRef.current = search;
-			if (highlightedIndex !== 0) {
-				setHighlightedIndex(0);
-			}
-		}
-
-		// Expose methods to parent via ref
-		useImperativeHandle(
-			ref,
-			() => ({
-				selectHighlighted: () => {
-					if (selectableDocuments[highlightedIndex]) {
-						handleSelectDocument(selectableDocuments[highlightedIndex]);
-					}
-				},
-				moveUp: () => {
-					setHighlightedIndex((prev) => (prev > 0 ? prev - 1 : selectableDocuments.length - 1));
-				},
-				moveDown: () => {
-					setHighlightedIndex((prev) => (prev < selectableDocuments.length - 1 ? prev + 1 : 0));
-				},
-			}),
-			[selectableDocuments, highlightedIndex, handleSelectDocument]
-		);
-
-		// Handle keyboard navigation
-		const handleKeyDown = useCallback(
-			(e: React.KeyboardEvent) => {
-				if (selectableDocuments.length === 0) return;
-
-				switch (e.key) {
-					case "ArrowDown":
-						e.preventDefault();
-						setHighlightedIndex((prev) => (prev < selectableDocuments.length - 1 ? prev + 1 : 0));
-						break;
-					case "ArrowUp":
-						e.preventDefault();
-						setHighlightedIndex((prev) => (prev > 0 ? prev - 1 : selectableDocuments.length - 1));
-						break;
-					case "Enter":
-						e.preventDefault();
-						if (selectableDocuments[highlightedIndex]) {
-							handleSelectDocument(selectableDocuments[highlightedIndex]);
-						}
-						break;
-					case "Escape":
-						e.preventDefault();
-						onDone();
-						break;
-				}
-			},
-			[selectableDocuments, highlightedIndex, handleSelectDocument, onDone]
-		);
-
-		return (
-			<div
-				className="flex flex-col w-[280px] sm:w-[320px] bg-zinc-900 rounded-lg"
-				onKeyDown={handleKeyDown}
-				role="listbox"
-				tabIndex={-1}
-			>
-				{/* Document List */}
-				<div className="max-h-[280px] overflow-y-auto">
-					{actualLoading ? (
-						<div className="flex items-center justify-center py-4">
-							<div className="animate-spin h-5 w-5 border-2 border-primary border-t-transparent rounded-full" />
-						</div>
-					) : actualDocuments.length === 0 ? (
-						<div className="flex flex-col items-center justify-center py-4 text-center px-4">
-							<FileText className="h-5 w-5 text-muted-foreground/50 mb-1" />
-							<p className="text-sm text-muted-foreground">No documents found</p>
-						</div>
-					) : (
-						<div className="py-1">
-							{actualDocuments.map((doc) => {
-								const isAlreadySelected = selectedIds.has(doc.id);
-								const selectableIndex = selectableDocuments.findIndex((d) => d.id === doc.id);
-								const isHighlighted = !isAlreadySelected && selectableIndex === highlightedIndex;
-
-								return (
-									<button
-										key={doc.id}
-										ref={(el) => {
-											if (el && selectableIndex >= 0) {
-												itemRefs.current.set(selectableIndex, el);
-											}
-										}}
-										type="button"
-										onClick={() => !isAlreadySelected && handleSelectDocument(doc)}
-										onMouseEnter={() => {
-											if (!isAlreadySelected && selectableIndex >= 0) {
-												setHighlightedIndex(selectableIndex);
-											}
-										}}
-										disabled={isAlreadySelected}
-										className={cn(
-											"w-full flex items-center gap-2 px-3 py-2 text-left transition-colors",
-											isAlreadySelected ? "opacity-50 cursor-not-allowed" : "cursor-pointer",
-											isHighlighted && "bg-accent"
-										)}
-									>
-										{/* Type icon */}
-										<span className="flex-shrink-0 text-muted-foreground text-sm">
-											{getConnectorIcon(doc.document_type)}
-										</span>
-										{/* Title */}
-										<span className="flex-1 text-sm truncate" title={doc.title}>
-											{doc.title}
-										</span>
-									</button>
-								);
-							})}
-						</div>
-					)}
-				</div>
-			</div>
-		);
-	}
-);

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

@@ -0,0 +1,243 @@
+"use client";
+
+import { useQuery } from "@tanstack/react-query";
+import { FileText } from "lucide-react";
+import {
+	forwardRef,
+	useCallback,
+	useEffect,
+	useImperativeHandle,
+	useMemo,
+	useRef,
+	useState,
+} from "react";
+import { getConnectorIcon } from "@/contracts/enums/connectorIcons";
+import type { Document } from "@/contracts/types/document.types";
+import { documentsApiService } from "@/lib/apis/documents-api.service";
+import { cacheKeys } from "@/lib/query-client/cache-keys";
+import { cn } from "@/lib/utils";
+
+export interface DocumentMentionPickerRef {
+	selectHighlighted: () => void;
+	moveUp: () => void;
+	moveDown: () => void;
+}
+
+interface DocumentMentionPickerProps {
+	searchSpaceId: number;
+	onSelectionChange: (documents: Document[]) => void;
+	onDone: () => void;
+	initialSelectedDocuments?: Document[];
+	externalSearch?: string;
+}
+
+function useDebounced<T>(value: T, delay = 300) {
+	const [debounced, setDebounced] = useState(value);
+	useEffect(() => {
+		const t = setTimeout(() => setDebounced(value), delay);
+		return () => clearTimeout(t);
+	}, [value, delay]);
+	return debounced;
+}
+
+export const DocumentMentionPicker = forwardRef<
+	DocumentMentionPickerRef,
+	DocumentMentionPickerProps
+>(function DocumentMentionPicker(
+	{ searchSpaceId, onSelectionChange, onDone, initialSelectedDocuments = [], externalSearch = "" },
+	ref
+) {
+	// Use external search
+	const search = externalSearch;
+	const debouncedSearch = useDebounced(search, 150);
+	const [highlightedIndex, setHighlightedIndex] = useState(0);
+	const itemRefs = useRef<Map<number, HTMLButtonElement>>(new Map());
+
+	const fetchQueryParams = useMemo(
+		() => ({
+			search_space_id: searchSpaceId,
+			page: 0,
+			page_size: 20,
+		}),
+		[searchSpaceId]
+	);
+
+	const searchQueryParams = useMemo(() => {
+		return {
+			search_space_id: searchSpaceId,
+			page: 0,
+			page_size: 20,
+			title: debouncedSearch,
+		};
+	}, [debouncedSearch, searchSpaceId]);
+
+	// Use query for fetching documents
+	const { data: documents, isLoading: isDocumentsLoading } = useQuery({
+		queryKey: cacheKeys.documents.withQueryParams(fetchQueryParams),
+		queryFn: () => documentsApiService.getDocuments({ queryParams: fetchQueryParams }),
+		staleTime: 3 * 60 * 1000,
+		enabled: !!searchSpaceId && !debouncedSearch.trim(),
+	});
+
+	// Searching
+	const { data: searchedDocuments, isLoading: isSearchedDocumentsLoading } = useQuery({
+		queryKey: cacheKeys.documents.withQueryParams(searchQueryParams),
+		queryFn: () => documentsApiService.searchDocuments({ queryParams: searchQueryParams }),
+		staleTime: 3 * 60 * 1000,
+		enabled: !!searchSpaceId && !!debouncedSearch.trim(),
+	});
+
+	const actualDocuments = debouncedSearch.trim()
+		? searchedDocuments?.items || []
+		: documents?.items || [];
+	const actualLoading = debouncedSearch.trim() ? isSearchedDocumentsLoading : isDocumentsLoading;
+
+	// Track already selected document IDs
+	const selectedIds = useMemo(
+		() => new Set(initialSelectedDocuments.map((d) => d.id)),
+		[initialSelectedDocuments]
+	);
+
+	// Filter out already selected documents for navigation
+	const selectableDocuments = useMemo(
+		() => actualDocuments.filter((doc) => !selectedIds.has(doc.id)),
+		[actualDocuments, selectedIds]
+	);
+
+	const handleSelectDocument = useCallback(
+		(doc: Document) => {
+			onSelectionChange([...initialSelectedDocuments, doc]);
+			onDone();
+		},
+		[initialSelectedDocuments, onSelectionChange, onDone]
+	);
+
+	// Scroll highlighted item into view
+	useEffect(() => {
+		const item = itemRefs.current.get(highlightedIndex);
+		if (item) {
+			item.scrollIntoView({ block: "nearest", behavior: "smooth" });
+		}
+	}, [highlightedIndex]);
+
+	// Reset highlighted index when external search changes
+	const prevSearchRef = useRef(search);
+	if (prevSearchRef.current !== search) {
+		prevSearchRef.current = search;
+		if (highlightedIndex !== 0) {
+			setHighlightedIndex(0);
+		}
+	}
+
+	// Expose methods to parent via ref
+	useImperativeHandle(
+		ref,
+		() => ({
+			selectHighlighted: () => {
+				if (selectableDocuments[highlightedIndex]) {
+					handleSelectDocument(selectableDocuments[highlightedIndex]);
+				}
+			},
+			moveUp: () => {
+				setHighlightedIndex((prev) => (prev > 0 ? prev - 1 : selectableDocuments.length - 1));
+			},
+			moveDown: () => {
+				setHighlightedIndex((prev) => (prev < selectableDocuments.length - 1 ? prev + 1 : 0));
+			},
+		}),
+		[selectableDocuments, highlightedIndex, handleSelectDocument]
+	);
+
+	// Handle keyboard navigation
+	const handleKeyDown = useCallback(
+		(e: React.KeyboardEvent) => {
+			if (selectableDocuments.length === 0) return;
+
+			switch (e.key) {
+				case "ArrowDown":
+					e.preventDefault();
+					setHighlightedIndex((prev) => (prev < selectableDocuments.length - 1 ? prev + 1 : 0));
+					break;
+				case "ArrowUp":
+					e.preventDefault();
+					setHighlightedIndex((prev) => (prev > 0 ? prev - 1 : selectableDocuments.length - 1));
+					break;
+				case "Enter":
+					e.preventDefault();
+					if (selectableDocuments[highlightedIndex]) {
+						handleSelectDocument(selectableDocuments[highlightedIndex]);
+					}
+					break;
+				case "Escape":
+					e.preventDefault();
+					onDone();
+					break;
+			}
+		},
+		[selectableDocuments, highlightedIndex, handleSelectDocument, onDone]
+	);
+
+	return (
+		<div
+			className="flex flex-col w-[280px] sm:w-[320px] bg-popover rounded-lg"
+			onKeyDown={handleKeyDown}
+			role="listbox"
+			tabIndex={-1}
+		>
+			{/* Document List */}
+			<div className="max-h-[280px] overflow-y-auto">
+				{actualLoading ? (
+					<div className="flex items-center justify-center py-4">
+						<div className="animate-spin h-5 w-5 border-2 border-primary border-t-transparent rounded-full" />
+					</div>
+				) : actualDocuments.length === 0 ? (
+					<div className="flex flex-col items-center justify-center py-4 text-center px-4">
+						<FileText className="h-5 w-5 text-muted-foreground/50 mb-1" />
+						<p className="text-sm text-muted-foreground">No documents found</p>
+					</div>
+				) : (
+					<div className="py-1">
+						{actualDocuments.map((doc) => {
+							const isAlreadySelected = selectedIds.has(doc.id);
+							const selectableIndex = selectableDocuments.findIndex((d) => d.id === doc.id);
+							const isHighlighted = !isAlreadySelected && selectableIndex === highlightedIndex;
+
+							return (
+								<button
+									key={doc.id}
+									ref={(el) => {
+										if (el && selectableIndex >= 0) {
+											itemRefs.current.set(selectableIndex, el);
+										}
+									}}
+									type="button"
+									onClick={() => !isAlreadySelected && handleSelectDocument(doc)}
+									onMouseEnter={() => {
+										if (!isAlreadySelected && selectableIndex >= 0) {
+											setHighlightedIndex(selectableIndex);
+										}
+									}}
+									disabled={isAlreadySelected}
+									className={cn(
+										"w-full flex items-center gap-2 px-3 py-2 text-left transition-colors",
+										isAlreadySelected ? "opacity-50 cursor-not-allowed" : "cursor-pointer",
+										isHighlighted && "bg-accent"
+									)}
+								>
+									{/* Type icon */}
+									<span className="flex-shrink-0 text-muted-foreground text-sm">
+										{getConnectorIcon(doc.document_type)}
+									</span>
+									{/* Title */}
+									<span className="flex-1 text-sm truncate" title={doc.title}>
+										{doc.title}
+									</span>
+								</button>
+							);
+						})}
+					</div>
+				)}
+			</div>
+		</div>
+	);
+});

surfsense_web/components/new-chat/model-selector.tsx

@@ -175,7 +175,7 @@ export function ModelSelector({ onEdit, onAddNew, className }: ModelSelectorProp
 					role="combobox"
 					aria-expanded={open}
 					className={cn(
-						"h-9 gap-2 px-3 rounded-xl border border-border/30 bg-background/50 backdrop-blur-sm",
+						"h-9 gap-2 px-3 rounded-xl border border-border/80 bg-background/50 backdrop-blur-sm",
 						"hover:bg-muted/80 hover:border-border/30 transition-all duration-200",
 						"text-sm font-medium text-foreground",
 						"focus-visible:ring-0 focus-visible:ring-offset-0",

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

@@ -1,39 +1,206 @@
 "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";
 
-export type ChainOfThoughtItemProps = React.ComponentProps<"div">;
+// ============================================================================
+// Constants
+// ============================================================================
 
-export const ChainOfThoughtItem = ({ children, className, ...props }: ChainOfThoughtItemProps) => (
-	<div className={cn("text-muted-foreground text-sm", className)} {...props}>
-		{children}
+/** Animation timing constants (in milliseconds) */
+const ANIMATION = {
+	/** Delay between each step appearing */
+	STAGGER_DELAY_MS: 50,
+	/** Additional delay for connection line animation */
+	CONNECTION_LINE_DELAY_MS: 150,
+} as const;
+
+/** File extension categories for icon mapping */
+const FILE_EXTENSIONS = {
+	DOCUMENT: ["pdf", "doc", "docx"] as const,
+	SPREADSHEET: ["xls", "xlsx", "csv"] as const,
+	IMAGE: ["png", "jpg", "jpeg", "gif", "webp", "svg"] as const,
+	AUDIO: ["mp3", "wav", "m4a", "ogg", "webm"] as const,
+	VIDEO: ["mp4", "mov", "avi", "mkv"] as const,
+	CODE: ["js", "ts", "tsx", "jsx", "py", "html", "css", "json", "md"] as const,
+} as const;
+
+/** Type for file extension categories */
+type FileExtensionCategory = keyof typeof FILE_EXTENSIONS;
+
+/** Icon size class for file icons */
+const FILE_ICON_SIZE_CLASS = "size-3.5" as const;
+
+// ============================================================================
+// Hooks
+// ============================================================================
+
+/**
+ * Custom hook for entrance animation
+ * Returns true after mount to trigger CSS transitions
+ */
+function useEntranceAnimation(delay = 0): boolean {
+	const [isVisible, setIsVisible] = useState(false);
+
+	useEffect(() => {
+		const timer = setTimeout(() => setIsVisible(true), delay);
+		return () => clearTimeout(timer);
+	}, [delay]);
+
+	return isVisible;
+}
+
+// ============================================================================
+// File Icon Utilities
+// ============================================================================
+
+/**
+ * Check if an extension belongs to a specific category
+ */
+function isExtensionInCategory(ext: string, category: FileExtensionCategory): boolean {
+	return (FILE_EXTENSIONS[category] as readonly string[]).includes(ext);
+}
+
+/**
+ * Get file icon based on file extension (all icons are muted/gray)
+ */
+function getFileIcon(name: string): React.ReactNode {
+	const ext = name.split(".").pop()?.toLowerCase() ?? "";
+
+	if (isExtensionInCategory(ext, "DOCUMENT")) {
+		return <FileText className={FILE_ICON_SIZE_CLASS} />;
+	}
+	if (isExtensionInCategory(ext, "SPREADSHEET")) {
+		return <FileSpreadsheet className={FILE_ICON_SIZE_CLASS} />;
+	}
+	if (isExtensionInCategory(ext, "IMAGE")) {
+		return <FileImage className={FILE_ICON_SIZE_CLASS} />;
+	}
+	if (isExtensionInCategory(ext, "AUDIO")) {
+		return <FileAudio className={FILE_ICON_SIZE_CLASS} />;
+	}
+	if (isExtensionInCategory(ext, "VIDEO")) {
+		return <FileVideo className={FILE_ICON_SIZE_CLASS} />;
+	}
+	if (isExtensionInCategory(ext, "CODE")) {
+		return <FileCode className={FILE_ICON_SIZE_CLASS} />;
+	}
+	return <File className={FILE_ICON_SIZE_CLASS} />;
+}
+
+// ============================================================================
+// Attachment Components
+// ============================================================================
+
+interface AttachmentTileProps {
+	/** File name to display */
+	name: string;
+}
+
+/**
+ * Compact attachment tile component - matches the chat UI style
+ */
+const AttachmentTile: React.FC<AttachmentTileProps> = ({ 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;
+}
+
+// ============================================================================
+// Chain of Thought Components
+// ============================================================================
+
+export interface ChainOfThoughtItemProps extends React.HTMLAttributes<HTMLDivElement> {
+	children: React.ReactNode;
+}
+
+export const ChainOfThoughtItem: React.FC<ChainOfThoughtItemProps> = ({
+	children,
+	className,
+	...props
+}) => (
+	<div
+		className={cn("text-muted-foreground text-sm flex flex-wrap items-center gap-1", className)}
+		{...props}
+	>
+		{typeof children === "string" ? parseAndRenderWithBadges(children) : children}
 	</div>
 );
 
-export type ChainOfThoughtTriggerProps = React.ComponentProps<typeof CollapsibleTrigger> & {
+export interface ChainOfThoughtTriggerProps
+	extends React.ComponentProps<typeof CollapsibleTrigger> {
+	/** Optional icon to display on the left side */
 	leftIcon?: React.ReactNode;
+	/** Whether to swap the icon with chevron on hover */
 	swapIconOnHover?: boolean;
-};
+}
 
-export const ChainOfThoughtTrigger = ({
+export const ChainOfThoughtTrigger: React.FC<ChainOfThoughtTriggerProps> = ({
 	children,
 	className,
 	leftIcon,
 	swapIconOnHover = true,
 	...props
-}: ChainOfThoughtTriggerProps) => (
+}) => (
 	<CollapsibleTrigger
 		className={cn(
 			"group text-muted-foreground hover:text-foreground flex cursor-pointer items-center justify-start gap-1 text-left text-sm transition-colors",
@@ -64,13 +231,14 @@ export const ChainOfThoughtTrigger = ({
 	</CollapsibleTrigger>
 );
 
-export type ChainOfThoughtContentProps = React.ComponentProps<typeof CollapsibleContent>;
+export interface ChainOfThoughtContentProps
+	extends React.ComponentProps<typeof CollapsibleContent> {}
 
-export const ChainOfThoughtContent = ({
+export const ChainOfThoughtContent: React.FC<ChainOfThoughtContentProps> = ({
 	children,
 	className,
 	...props
-}: ChainOfThoughtContentProps) => {
+}) => {
 	return (
 		<CollapsibleContent
 			className={cn(
@@ -80,53 +248,109 @@ 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 * ANIMATION.STAGGER_DELAY_MS}ms`,
+									animationFillMode: "backwards",
+								}}
+							>
+								{child}
+							</div>
+						);
+					})}
+				</div>
 			</div>
 		</CollapsibleContent>
 	);
 };
 
-export type ChainOfThoughtProps = {
+export interface ChainOfThoughtProps {
 	children: React.ReactNode;
 	className?: string;
-};
+}
 
-export function ChainOfThought({ children, className }: ChainOfThoughtProps) {
+export const ChainOfThought: React.FC<ChainOfThoughtProps> = ({ children, className }) => {
 	const childrenArray = React.Children.toArray(children);
 
 	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>
 	);
-}
-
-export type ChainOfThoughtStepProps = {
-	children: React.ReactNode;
-	className?: string;
-	isLast?: boolean;
 };
 
-export const ChainOfThoughtStep = ({
+export interface ChainOfThoughtStepProps
+	extends Omit<React.ComponentProps<typeof Collapsible>, "children"> {
+	children: React.ReactNode;
+	className?: string;
+	/** Whether this is the last step (hides connection line) */
+	isLast?: boolean;
+	/** Index of the step for staggered animation timing */
+	stepIndex?: number;
+}
+
+export const ChainOfThoughtStep: React.FC<ChainOfThoughtStepProps> = ({
 	children,
 	className,
 	isLast = false,
+	stepIndex = 0,
 	...props
-}: ChainOfThoughtStepProps & React.ComponentProps<typeof Collapsible>) => {
+}) => {
+	// Staggered entrance animation based on step index
+	const isVisible = useEntranceAnimation(stepIndex * ANIMATION.STAGGER_DELAY_MS);
+
+	// Calculate connection line delay: step delay + additional offset
+	const connectionLineDelay =
+		stepIndex * ANIMATION.STAGGER_DELAY_MS + ANIMATION.CONNECTION_LINE_DELAY_MS;
+
 	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: `${connectionLineDelay}ms` }}
+				/>
 			</div>
 		</Collapsible>
 	);

surfsense_web/components/tool-ui/article/index.tsx

@@ -19,20 +19,20 @@ import { cn } from "@/lib/utils";
  */
 const SerializableArticleSchema = z.object({
 	id: z.string().default("article-unknown"),
-	assetId: z.string().optional(),
-	kind: z.literal("article").optional(),
+	assetId: z.string().nullish(),
+	kind: z.literal("article").nullish(),
 	title: z.string().default("Untitled Article"),
-	description: z.string().optional(),
-	content: z.string().optional(),
-	href: z.string().url().optional(),
-	domain: z.string().optional(),
-	author: z.string().optional(),
-	date: z.string().optional(),
-	word_count: z.number().optional(),
-	wordCount: z.number().optional(),
-	was_truncated: z.boolean().optional(),
-	wasTruncated: z.boolean().optional(),
-	error: z.string().optional(),
+	description: z.string().nullish(),
+	content: z.string().nullish(),
+	href: z.string().url().nullish(),
+	domain: z.string().nullish(),
+	author: z.string().nullish(),
+	date: z.string().nullish(),
+	word_count: z.number().nullish(),
+	wordCount: z.number().nullish(),
+	was_truncated: z.boolean().nullish(),
+	wasTruncated: z.boolean().nullish(),
+	error: z.string().nullish(),
 });
 
 /**

surfsense_web/components/tool-ui/audio.tsx

@@ -188,19 +188,6 @@ export function Audio({ id, src, title, description, artwork, durationMs, classN
 								<Volume2Icon className="size-8 text-primary/50" />
 							</div>
 						)}
-						{/* Play overlay on artwork */}
-						<button
-							type="button"
-							onClick={togglePlayPause}
-							className="absolute inset-0 flex items-center justify-center bg-black/0 opacity-0 transition-all group-hover:bg-black/30 group-hover:opacity-100"
-							aria-label={isPlaying ? "Pause" : "Play"}
-						>
-							{isPlaying ? (
-								<PauseIcon className="size-8 text-white drop-shadow-lg" />
-							) : (
-								<PlayIcon className="size-8 text-white drop-shadow-lg" />
-							)}
-						</button>
 					</div>
 				</div>
 
@@ -254,17 +241,29 @@ export function Audio({ id, src, title, description, artwork, durationMs, classN
 					</Button>
 
 					{/* Volume control */}
-					<div className="flex items-center gap-2">
+					<div className="flex items-center gap-1.5">
 						<Button variant="ghost" size="icon" onClick={toggleMute} className="size-8">
 							{isMuted ? <VolumeXIcon className="size-4" /> : <Volume2Icon className="size-4" />}
 						</Button>
-						<Slider
-							value={[isMuted ? 0 : volume]}
-							max={1}
-							step={0.01}
-							onValueChange={handleVolumeChange}
-							className="w-20"
-						/>
+						{/* Custom volume bar - visually distinct from progress slider */}
+						<div className="relative flex h-6 w-16 items-center">
+							<div className="relative h-1 w-full rounded-full bg-muted-foreground/20">
+								<div
+									className="absolute left-0 top-0 h-full rounded-full bg-muted-foreground/60 transition-all"
+									style={{ width: `${(isMuted ? 0 : volume) * 100}%` }}
+								/>
+							</div>
+							<input
+								type="range"
+								min={0}
+								max={1}
+								step={0.01}
+								value={isMuted ? 0 : volume}
+								onChange={(e) => handleVolumeChange([Number.parseFloat(e.target.value)])}
+								className="absolute inset-0 h-full w-full cursor-pointer opacity-0"
+								aria-label="Volume"
+							/>
+						</div>
 					</div>
 				</div>
 

surfsense_web/components/tool-ui/deepagent-thinking.tsx

@@ -2,7 +2,8 @@
 
 import { makeAssistantToolUI } from "@assistant-ui/react";
 import { Brain, CheckCircle2, Loader2, Search, Sparkles } from "lucide-react";
-import { useEffect, useMemo, useRef, useState } from "react";
+import type { FC, ReactNode } from "react";
+import { useCallback, useEffect, useMemo, useRef, useState } from "react";
 import { z } from "zod";
 import {
 	ChainOfThought,
@@ -13,34 +14,96 @@ import {
 } from "@/components/prompt-kit/chain-of-thought";
 import { cn } from "@/lib/utils";
 
-/**
- * Zod schemas for runtime validation
- */
+// ============================================================================
+// Constants
+// ============================================================================
+
+/** Step status values */
+const STEP_STATUS = {
+	PENDING: "pending",
+	IN_PROGRESS: "in_progress",
+	COMPLETED: "completed",
+} as const;
+
+/** Agent thinking status values */
+const THINKING_STATUS = {
+	THINKING: "thinking",
+	SEARCHING: "searching",
+	SYNTHESIZING: "synthesizing",
+	COMPLETED: "completed",
+} as const;
+
+/** Keywords for icon detection */
+const STEP_KEYWORDS = {
+	SEARCH: ["search", "knowledge"] as const,
+	ANALYSIS: ["analy", "understand"] as const,
+} as const;
+
+/** Icon size class */
+const ICON_SIZE_CLASS = "size-4" as const;
+
+/** Status text mapping */
+const STATUS_TEXT_MAP: Record<string, string> = {
+	[THINKING_STATUS.SEARCHING]: "Searching knowledge base...",
+	[THINKING_STATUS.SYNTHESIZING]: "Synthesizing response...",
+	[THINKING_STATUS.THINKING]: "Thinking...",
+} as const;
+
+// ============================================================================
+// Type Definitions
+// ============================================================================
+
+type StepStatus = (typeof STEP_STATUS)[keyof typeof STEP_STATUS];
+type ThinkingStatus = (typeof THINKING_STATUS)[keyof typeof THINKING_STATUS];
+
+// ============================================================================
+// Zod Schemas
+// ============================================================================
+
 const ThinkingStepSchema = z.object({
 	id: z.string(),
 	title: z.string(),
 	items: z.array(z.string()).default([]),
-	status: z.enum(["pending", "in_progress", "completed"]).default("pending"),
+	status: z
+		.enum([STEP_STATUS.PENDING, STEP_STATUS.IN_PROGRESS, STEP_STATUS.COMPLETED])
+		.default(STEP_STATUS.PENDING),
 });
 
 const DeepAgentThinkingArgsSchema = z.object({
-	query: z.string().optional(),
-	context: z.string().optional(),
+	query: z.string().nullish(),
+	context: z.string().nullish(),
 });
 
 const DeepAgentThinkingResultSchema = z.object({
-	steps: z.array(ThinkingStepSchema).optional(),
-	status: z.enum(["thinking", "searching", "synthesizing", "completed"]).optional(),
-	summary: z.string().optional(),
+	steps: z.array(ThinkingStepSchema).nullish(),
+	status: z
+		.enum([
+			THINKING_STATUS.THINKING,
+			THINKING_STATUS.SEARCHING,
+			THINKING_STATUS.SYNTHESIZING,
+			THINKING_STATUS.COMPLETED,
+		])
+		.nullish(),
+	summary: z.string().nullish(),
 });
 
-/**
- * Types derived from Zod schemas
- */
+/** Types derived from Zod schemas */
 type ThinkingStep = z.infer<typeof ThinkingStepSchema>;
 type DeepAgentThinkingArgs = z.infer<typeof DeepAgentThinkingArgsSchema>;
 type DeepAgentThinkingResult = z.infer<typeof DeepAgentThinkingResultSchema>;
 
+// ============================================================================
+// Parser Functions
+// ============================================================================
+
+/** Default fallback step when parsing fails */
+const DEFAULT_FALLBACK_STEP: ThinkingStep = {
+	id: "unknown",
+	title: "Processing...",
+	items: [],
+	status: STEP_STATUS.PENDING,
+} as const;
+
 /**
  * Parse and validate a single thinking step
  */
@@ -48,13 +111,7 @@ export function parseThinkingStep(data: unknown): ThinkingStep {
 	const result = ThinkingStepSchema.safeParse(data);
 	if (!result.success) {
 		console.warn("Invalid thinking step data:", result.error.issues);
-		// Return a fallback step
-		return {
-			id: "unknown",
-			title: "Processing...",
-			items: [],
-			status: "pending",
-		};
+		return DEFAULT_FALLBACK_STEP;
 	}
 	return result.data;
 }
@@ -71,55 +128,69 @@ export function parseThinkingResult(data: unknown): DeepAgentThinkingResult {
 	return result.data;
 }
 
+// ============================================================================
+// Icon Utilities
+// ============================================================================
+
 /**
- * Get icon based on step status and type
+ * Check if title contains any of the keywords
  */
-function getStepIcon(status: "pending" | "in_progress" | "completed", title: string) {
-	// Check for specific step types based on title keywords
+function titleContainsKeywords(title: string, keywords: readonly string[]): boolean {
 	const titleLower = title.toLowerCase();
+	return keywords.some((keyword) => titleLower.includes(keyword));
+}
 
-	if (status === "in_progress") {
-		return <Loader2 className="size-4 animate-spin text-primary" />;
+/**
+ * Get icon based on step status and title
+ */
+function getStepIcon(status: StepStatus, title: string): ReactNode {
+	if (status === STEP_STATUS.IN_PROGRESS) {
+		return <Loader2 className={cn(ICON_SIZE_CLASS, "animate-spin text-primary")} />;
 	}
 
-	if (status === "completed") {
-		return <CheckCircle2 className="size-4 text-emerald-500" />;
+	if (status === STEP_STATUS.COMPLETED) {
+		return <CheckCircle2 className={cn(ICON_SIZE_CLASS, "text-emerald-500")} />;
 	}
 
-	// Default icons based on step type
-	if (titleLower.includes("search") || titleLower.includes("knowledge")) {
-		return <Search className="size-4 text-muted-foreground" />;
+	// Default icons based on step type keywords
+	if (titleContainsKeywords(title, STEP_KEYWORDS.SEARCH)) {
+		return <Search className={cn(ICON_SIZE_CLASS, "text-muted-foreground")} />;
 	}
 
-	if (titleLower.includes("analy") || titleLower.includes("understand")) {
-		return <Brain className="size-4 text-muted-foreground" />;
+	if (titleContainsKeywords(title, STEP_KEYWORDS.ANALYSIS)) {
+		return <Brain className={cn(ICON_SIZE_CLASS, "text-muted-foreground")} />;
 	}
 
-	return <Sparkles className="size-4 text-muted-foreground" />;
+	return <Sparkles className={cn(ICON_SIZE_CLASS, "text-muted-foreground")} />;
+}
+
+// ============================================================================
+// Sub-Components
+// ============================================================================
+
+interface ThinkingStepDisplayProps {
+	step: ThinkingStep;
+	isOpen: boolean;
+	onToggle: () => void;
 }
 
 /**
  * Component to display a single thinking step with controlled open state
  */
-function ThinkingStepDisplay({
-	step,
-	isOpen,
-	onToggle,
-}: {
-	step: ThinkingStep;
-	isOpen: boolean;
-	onToggle: () => void;
-}) {
+const ThinkingStepDisplay: FC<ThinkingStepDisplayProps> = ({ step, isOpen, onToggle }) => {
 	const icon = useMemo(() => getStepIcon(step.status, step.title), [step.status, step.title]);
 
+	const isInProgress = step.status === STEP_STATUS.IN_PROGRESS;
+	const isCompleted = step.status === STEP_STATUS.COMPLETED;
+
 	return (
 		<ChainOfThoughtStep open={isOpen} onOpenChange={onToggle}>
 			<ChainOfThoughtTrigger
 				leftIcon={icon}
-				swapIconOnHover={step.status !== "in_progress"}
+				swapIconOnHover={!isInProgress}
 				className={cn(
-					step.status === "in_progress" && "text-foreground font-medium",
-					step.status === "completed" && "text-muted-foreground"
+					isInProgress && "text-foreground font-medium",
+					isCompleted && "text-muted-foreground"
 				)}
 			>
 				{step.title}
@@ -131,22 +202,21 @@ function ThinkingStepDisplay({
 			</ChainOfThoughtContent>
 		</ChainOfThoughtStep>
 	);
+};
+
+interface ThinkingLoadingStateProps {
+	status?: ThinkingStatus | string;
 }
 
 /**
  * Loading state with animated thinking indicator
  */
-function ThinkingLoadingState({ status }: { status?: string }) {
+const ThinkingLoadingState: FC<ThinkingLoadingStateProps> = ({ status }) => {
 	const statusText = useMemo(() => {
-		switch (status) {
-			case "searching":
-				return "Searching knowledge base...";
-			case "synthesizing":
-				return "Synthesizing response...";
-			case "thinking":
-			default:
-				return "Thinking...";
+		if (status && status in STATUS_TEXT_MAP) {
+			return STATUS_TEXT_MAP[status];
 		}
+		return STATUS_TEXT_MAP[THINKING_STATUS.THINKING];
 	}, [status]);
 
 	return (
@@ -161,33 +231,35 @@ function ThinkingLoadingState({ status }: { status?: string }) {
 			<span className="text-sm text-muted-foreground">{statusText}</span>
 		</div>
 	);
+};
+
+interface SmartChainOfThoughtProps {
+	steps: ThinkingStep[];
 }
 
+/** Type for tracking step override states */
+type StepOverrides = Record<string, boolean>;
+
+/** Type for tracking step status history */
+type StepStatusHistory = Record<string, StepStatus>;
+
 /**
  * Smart chain of thought renderer with state management
  */
-function SmartChainOfThought({ steps }: { steps: ThinkingStep[] }) {
+const SmartChainOfThought: FC<SmartChainOfThoughtProps> = ({ steps }) => {
 	// Track which steps the user has manually toggled
-	const [manualOverrides, setManualOverrides] = useState<Record<string, boolean>>({});
+	const [manualOverrides, setManualOverrides] = useState<StepOverrides>({});
 	// Track previous step statuses to detect changes
-	const prevStatusesRef = useRef<Record<string, string>>({});
-
-	// Check if any step is currently in progress
-	const hasInProgressStep = steps.some((step) => step.status === "in_progress");
-
-	// Find the last completed step index
-	const lastCompletedIndex = steps
-		.map((s, i) => (s.status === "completed" ? i : -1))
-		.filter((i) => i !== -1)
-		.pop();
+	const prevStatusesRef = useRef<StepStatusHistory>({});
 
 	// Clear manual overrides when a step's status changes
 	useEffect(() => {
-		const currentStatuses: Record<string, string> = {};
+		const currentStatuses: StepStatusHistory = {};
 		steps.forEach((step) => {
 			currentStatuses[step.id] = step.status;
 			// If status changed, clear any manual override for this step
-			if (prevStatusesRef.current[step.id] && prevStatusesRef.current[step.id] !== step.status) {
+			const prevStatus = prevStatusesRef.current[step.id];
+			if (prevStatus && prevStatus !== step.status) {
 				setManualOverrides((prev) => {
 					const next = { ...prev };
 					delete next[step.id];
@@ -198,34 +270,33 @@ function SmartChainOfThought({ steps }: { steps: ThinkingStep[] }) {
 		prevStatusesRef.current = currentStatuses;
 	}, [steps]);
 
-	const getStepOpenState = (step: ThinkingStep, index: number): boolean => {
-		// If user has manually toggled, respect that
-		if (manualOverrides[step.id] !== undefined) {
-			return manualOverrides[step.id];
-		}
-		// Auto behavior: open if in progress
-		if (step.status === "in_progress") {
-			return true;
-		}
-		// Auto behavior: keep last completed step open if no in-progress step
-		if (!hasInProgressStep && index === lastCompletedIndex) {
-			return true;
-		}
-		// Default: collapsed
-		return false;
-	};
+	const getStepOpenState = useCallback(
+		(step: ThinkingStep): boolean => {
+			// If user has manually toggled, respect that
+			if (manualOverrides[step.id] !== undefined) {
+				return manualOverrides[step.id];
+			}
+			// Auto behavior: open if in progress
+			if (step.status === STEP_STATUS.IN_PROGRESS) {
+				return true;
+			}
+			// Default: collapsed (all steps collapse when processing is done)
+			return false;
+		},
+		[manualOverrides]
+	);
 
-	const handleToggle = (stepId: string, currentOpen: boolean) => {
+	const handleToggle = useCallback((stepId: string, currentOpen: boolean) => {
 		setManualOverrides((prev) => ({
 			...prev,
 			[stepId]: !currentOpen,
 		}));
-	};
+	}, []);
 
 	return (
 		<ChainOfThought>
-			{steps.map((step, index) => {
-				const isOpen = getStepOpenState(step, index);
+			{steps.map((step) => {
+				const isOpen = getStepOpenState(step);
 				return (
 					<ThinkingStepDisplay
 						key={step.id}
@@ -237,7 +308,7 @@ function SmartChainOfThought({ steps }: { steps: ThinkingStep[] }) {
 			})}
 		</ChainOfThought>
 	);
-}
+};
 
 /**
  * DeepAgent Thinking Tool UI Component
@@ -254,7 +325,7 @@ export const DeepAgentThinkingToolUI = makeAssistantToolUI<
 	render: function DeepAgentThinkingUI({ result, status }) {
 		// Loading state - tool is still running
 		if (status.type === "running" || status.type === "requires-action") {
-			return <ThinkingLoadingState status={result?.status} />;
+			return <ThinkingLoadingState status={result?.status ?? undefined} />;
 		}
 
 		// Incomplete/cancelled state
@@ -281,21 +352,30 @@ export const DeepAgentThinkingToolUI = makeAssistantToolUI<
 	},
 });
 
+// ============================================================================
+// Public Components
+// ============================================================================
+
+export interface InlineThinkingDisplayProps {
+	/** The thinking steps to display */
+	steps: ThinkingStep[];
+	/** Whether content is currently streaming */
+	isStreaming?: boolean;
+	/** Additional CSS class names */
+	className?: string;
+}
+
 /**
  * Inline Thinking Display Component
  *
  * A simpler version that can be used inline with the message content
  * for displaying reasoning without the full tool UI infrastructure.
  */
-export function InlineThinkingDisplay({
+export const InlineThinkingDisplay: FC<InlineThinkingDisplayProps> = ({
 	steps,
 	isStreaming = false,
 	className,
-}: {
-	steps: ThinkingStep[];
-	isStreaming?: boolean;
-	className?: string;
-}) {
+}) => {
 	if (steps.length === 0 && !isStreaming) {
 		return null;
 	}
@@ -309,6 +389,18 @@ export function InlineThinkingDisplay({
 			)}
 		</div>
 	);
-}
+};
 
-export type { ThinkingStep, DeepAgentThinkingArgs, DeepAgentThinkingResult };
+// ============================================================================
+// Exports
+// ============================================================================
+
+export type {
+	ThinkingStep,
+	DeepAgentThinkingArgs,
+	DeepAgentThinkingResult,
+	StepStatus,
+	ThinkingStatus,
+};
+
+export { STEP_STATUS, THINKING_STATUS };

surfsense_web/components/tool-ui/display-image.tsx

@@ -23,7 +23,7 @@ interface DisplayImageResult {
 	id: string;
 	assetId: string;
 	src: string;
-	alt: string;
+	alt?: string; // Made optional - parseSerializableImage provides fallback
 	title?: string;
 	description?: string;
 	domain?: string;

surfsense_web/components/tool-ui/generate-podcast.tsx

@@ -14,27 +14,27 @@ import { clearActivePodcastTaskId, setActivePodcastTaskId } from "@/lib/chat/pod
  */
 const GeneratePodcastArgsSchema = z.object({
 	source_content: z.string(),
-	podcast_title: z.string().optional(),
-	user_prompt: z.string().optional(),
+	podcast_title: z.string().nullish(),
+	user_prompt: z.string().nullish(),
 });
 
 const GeneratePodcastResultSchema = z.object({
 	status: z.enum(["processing", "already_generating", "success", "error"]),
-	task_id: z.string().optional(),
-	podcast_id: z.number().optional(),
-	title: z.string().optional(),
-	transcript_entries: z.number().optional(),
-	message: z.string().optional(),
-	error: z.string().optional(),
+	task_id: z.string().nullish(),
+	podcast_id: z.number().nullish(),
+	title: z.string().nullish(),
+	transcript_entries: z.number().nullish(),
+	message: z.string().nullish(),
+	error: z.string().nullish(),
 });
 
 const TaskStatusResponseSchema = z.object({
 	status: z.enum(["processing", "success", "error"]),
-	podcast_id: z.number().optional(),
-	title: z.string().optional(),
-	transcript_entries: z.number().optional(),
-	state: z.string().optional(),
-	error: z.string().optional(),
+	podcast_id: z.number().nullish(),
+	title: z.string().nullish(),
+	transcript_entries: z.number().nullish(),
+	state: z.string().nullish(),
+	error: z.string().nullish(),
 });
 
 const PodcastTranscriptEntrySchema = z.object({
@@ -43,7 +43,7 @@ const PodcastTranscriptEntrySchema = z.object({
 });
 
 const PodcastDetailsSchema = z.object({
-	podcast_transcript: z.array(PodcastTranscriptEntrySchema).optional(),
+	podcast_transcript: z.array(PodcastTranscriptEntrySchema).nullish(),
 });
 
 /**
@@ -75,7 +75,9 @@ function parsePodcastDetails(data: unknown): { podcast_transcript?: PodcastTrans
 		console.warn("Invalid podcast details:", result.error.issues);
 		return {};
 	}
-	return result.data;
+	return {
+		podcast_transcript: result.data.podcast_transcript ?? undefined,
+	};
 }
 
 /**

surfsense_web/components/tool-ui/image/index.tsx

@@ -11,26 +11,26 @@ import { cn } from "@/lib/utils";
 /**
  * Zod schemas for runtime validation
  */
-const AspectRatioSchema = z.enum(["1:1", "4:3", "16:9", "9:16", "auto"]);
+const AspectRatioSchema = z.enum(["1:1", "4:3", "16:9", "9:16", "21:9", "auto"]);
 const ImageFitSchema = z.enum(["cover", "contain"]);
 
 const ImageSourceSchema = z.object({
 	label: z.string(),
-	iconUrl: z.string().optional(),
-	url: z.string().optional(),
+	iconUrl: z.string().nullish(),
+	url: z.string().nullish(),
 });
 
 const SerializableImageSchema = z.object({
 	id: z.string(),
 	assetId: z.string(),
 	src: z.string(),
-	alt: z.string(),
-	title: z.string().optional(),
-	description: z.string().optional(),
-	href: z.string().optional(),
-	domain: z.string().optional(),
-	ratio: AspectRatioSchema.optional(),
-	source: ImageSourceSchema.optional(),
+	alt: z.string().nullish(), // Made optional - will use fallback if missing
+	title: z.string().nullish(),
+	description: z.string().nullish(),
+	href: z.string().nullish(),
+	domain: z.string().nullish(),
+	ratio: AspectRatioSchema.nullish(),
+	source: ImageSourceSchema.nullish(),
 });
 
 /**
@@ -48,7 +48,7 @@ export interface ImageProps {
 	id: string;
 	assetId: string;
 	src: string;
-	alt: string;
+	alt?: string; // Optional with default fallback
 	title?: string;
 	description?: string;
 	href?: string;
@@ -62,18 +62,45 @@ export interface ImageProps {
 
 /**
  * Parse and validate serializable image from tool result
+ * Returns a valid SerializableImage with fallback values for missing optional fields
  */
-export function parseSerializableImage(result: unknown): SerializableImage {
+export function parseSerializableImage(result: unknown): SerializableImage & { alt: string } {
 	const parsed = SerializableImageSchema.safeParse(result);
 
 	if (!parsed.success) {
 		console.warn("Invalid image data:", parsed.error.issues);
-		// Try to extract basic info for error display
+
+		// Try to extract basic info and return a fallback object
 		const obj = (result && typeof result === "object" ? result : {}) as Record<string, unknown>;
+
+		// If we have at least id, assetId, and src, we can still render the image
+		if (
+			typeof obj.id === "string" &&
+			typeof obj.assetId === "string" &&
+			typeof obj.src === "string"
+		) {
+			return {
+				id: obj.id,
+				assetId: obj.assetId,
+				src: obj.src,
+				alt: typeof obj.alt === "string" ? obj.alt : "Image",
+				title: typeof obj.title === "string" ? obj.title : undefined,
+				description: typeof obj.description === "string" ? obj.description : undefined,
+				href: typeof obj.href === "string" ? obj.href : undefined,
+				domain: typeof obj.domain === "string" ? obj.domain : undefined,
+				ratio: undefined, // Use default ratio
+				source: undefined,
+			};
+		}
+
 		throw new Error(`Invalid image: ${parsed.error.issues.map((i) => i.message).join(", ")}`);
 	}
 
-	return parsed.data;
+	// Provide fallback for alt if it's null/undefined
+	return {
+		...parsed.data,
+		alt: parsed.data.alt ?? "Image",
+	};
 }
 
 /**
@@ -89,6 +116,8 @@ function getAspectRatioClass(ratio?: AspectRatio): string {
 			return "aspect-video";
 		case "9:16":
 			return "aspect-[9/16]";
+		case "21:9":
+			return "aspect-[21/9]";
 		case "auto":
 		default:
 			return "aspect-[4/3]";
@@ -172,7 +201,7 @@ export function ImageLoading({ title = "Loading image..." }: { title?: string })
 export function Image({
 	id,
 	src,
-	alt,
+	alt = "Image",
 	title,
 	description,
 	href,

surfsense_web/components/tool-ui/media-card/index.tsx

@@ -13,27 +13,27 @@ import { cn } from "@/lib/utils";
 /**
  * Zod schemas for runtime validation
  */
-const AspectRatioSchema = z.enum(["1:1", "4:3", "16:9", "21:9", "auto"]);
+const AspectRatioSchema = z.enum(["1:1", "4:3", "16:9", "9:16", "21:9", "auto"]);
 const MediaCardKindSchema = z.enum(["link", "image", "video", "audio"]);
 
 const ResponseActionSchema = z.object({
 	id: z.string(),
 	label: z.string(),
-	variant: z.enum(["default", "secondary", "outline", "destructive", "ghost"]).optional(),
-	confirmLabel: z.string().optional(),
+	variant: z.enum(["default", "secondary", "outline", "destructive", "ghost"]).nullish(),
+	confirmLabel: z.string().nullish(),
 });
 
 const SerializableMediaCardSchema = z.object({
 	id: z.string(),
 	assetId: z.string(),
 	kind: MediaCardKindSchema,
-	href: z.string().optional(),
-	src: z.string().optional(),
+	href: z.string().nullish(),
+	src: z.string().nullish(),
 	title: z.string(),
-	description: z.string().optional(),
-	thumb: z.string().optional(),
-	ratio: AspectRatioSchema.optional(),
-	domain: z.string().optional(),
+	description: z.string().nullish(),
+	thumb: z.string().nullish(),
+	ratio: AspectRatioSchema.nullish(),
+	domain: z.string().nullish(),
 });
 
 /**
@@ -90,6 +90,8 @@ function getAspectRatioClass(ratio?: AspectRatio): string {
 			return "aspect-[4/3]";
 		case "16:9":
 			return "aspect-video";
+		case "9:16":
+			return "aspect-[9/16]";
 		case "21:9":
 			return "aspect-[21/9]";
 		case "auto":

surfsense_web/lib/chat/attachment-adapter.ts

@@ -60,9 +60,11 @@ interface ProcessAttachmentResponse {
 /**
  * Extended CompleteAttachment with our custom extractedContent field
  * We store the extracted text in a custom field so we can access it in onNew
+ * For images, we also store the data URL so it can be displayed after persistence
  */
 export interface ChatAttachment extends CompleteAttachment {
 	extractedContent: string;
+	imageDataUrl?: string; // Base64 data URL for images (persists across page reloads)
 }
 
 /**
@@ -118,6 +120,21 @@ async function processAttachment(file: File): Promise<ProcessAttachmentResponse>
 // Store processed results for the send() method
 const processedAttachments = new Map<string, ProcessAttachmentResponse>();
 
+// Store image data URLs for attachments (so they persist after File objects are lost)
+const imageDataUrls = new Map<string, string>();
+
+/**
+ * Convert a File to a data URL (base64) for images
+ */
+async function fileToDataUrl(file: File): Promise<string> {
+	return new Promise((resolve, reject) => {
+		const reader = new FileReader();
+		reader.onload = () => resolve(reader.result as string);
+		reader.onerror = reject;
+		reader.readAsDataURL(file);
+	});
+}
+
 /**
  * Create the attachment adapter for assistant-ui
  *
@@ -170,6 +187,12 @@ export function createAttachmentAdapter(): AttachmentAdapter {
 			} as PendingAttachment;
 
 			try {
+				// For images, convert to data URL so we can display them after persistence
+				if (attachmentType === "image") {
+					const dataUrl = await fileToDataUrl(file);
+					imageDataUrls.set(id, dataUrl);
+				}
+
 				// Process the file through the backend ETL service
 				const result = await processAttachment(file);
 
@@ -204,10 +227,14 @@ export function createAttachmentAdapter(): AttachmentAdapter {
 		 */
 		async send(pendingAttachment: PendingAttachment): Promise<ChatAttachment> {
 			const result = processedAttachments.get(pendingAttachment.id);
+			const imageDataUrl = imageDataUrls.get(pendingAttachment.id);
 
 			if (result) {
 				// Clean up stored result
 				processedAttachments.delete(pendingAttachment.id);
+				if (imageDataUrl) {
+					imageDataUrls.delete(pendingAttachment.id);
+				}
 
 				return {
 					id: result.id,
@@ -222,6 +249,7 @@ export function createAttachmentAdapter(): AttachmentAdapter {
 						},
 					],
 					extractedContent: result.content,
+					imageDataUrl, // Store data URL for images so they can be displayed after persistence
 				};
 			}
 
@@ -238,6 +266,7 @@ export function createAttachmentAdapter(): AttachmentAdapter {
 				status: { type: "complete" },
 				content: [],
 				extractedContent: "",
+				imageDataUrl, // Still include data URL if available
 			};
 		},