Merge pull request #619 from MODSetter/dev

feat: in chat document mentions

surfsense_backend/app/routes/new_chat_routes.py

@@ -705,6 +705,7 @@ async def handle_new_chat(
                 session=session,
                 llm_config_id=llm_config_id,
                 attachments=request.attachments,
+                mentioned_document_ids=request.mentioned_document_ids,
             ),
             media_type="text/event-stream",
             headers={

surfsense_backend/app/schemas/new_chat.py

@@ -160,3 +160,6 @@ class NewChatRequest(BaseModel):
     attachments: list[ChatAttachment] | None = (
         None  # Optional attachments with extracted content
     )
+    mentioned_document_ids: list[int] | None = (
+        None  # Optional document IDs mentioned with @ in the chat
+    )

surfsense_backend/app/tasks/chat/stream_new_chat.py

@@ -14,6 +14,7 @@ from collections.abc import AsyncGenerator
 
 from langchain_core.messages import HumanMessage
 from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.future import select
 
 from app.agents.new_chat.chat_deepagent import create_surfsense_deep_agent
 from app.agents.new_chat.checkpointer import get_checkpointer
@@ -24,6 +25,7 @@ from app.agents.new_chat.llm_config import (
     load_agent_config,
     load_llm_config_from_yaml,
 )
+from app.db import Document
 from app.schemas.new_chat import ChatAttachment
 from app.services.connector_service import ConnectorService
 from app.services.new_streaming_service import VercelStreamingService
@@ -46,6 +48,27 @@ def format_attachments_as_context(attachments: list[ChatAttachment]) -> str:
     return "\n".join(context_parts)
 
 
+def format_mentioned_documents_as_context(documents: list[Document]) -> str:
+    """Format mentioned documents as context for the agent."""
+    if not documents:
+        return ""
+
+    context_parts = ["<mentioned_documents>"]
+    context_parts.append(
+        "The user has explicitly mentioned the following documents from their knowledge base. "
+        "These documents are directly relevant to the query and should be prioritized as primary sources."
+    )
+    for i, doc in enumerate(documents, 1):
+        context_parts.append(
+            f"<document index='{i}' id='{doc.id}' title='{doc.title}' type='{doc.document_type.value}'>"
+        )
+        context_parts.append(f"<![CDATA[{doc.content}]]>")
+        context_parts.append("</document>")
+    context_parts.append("</mentioned_documents>")
+
+    return "\n".join(context_parts)
+
+
 async def stream_new_chat(
     user_query: str,
     search_space_id: int,
@@ -53,6 +76,7 @@ async def stream_new_chat(
     session: AsyncSession,
     llm_config_id: int = -1,
     attachments: list[ChatAttachment] | None = None,
+    mentioned_document_ids: list[int] | None = None,
 ) -> AsyncGenerator[str, None]:
     """
     Stream chat responses from the new SurfSense deep agent.
@@ -68,6 +92,8 @@ async def stream_new_chat(
         session: The database session
         llm_config_id: The LLM configuration ID (default: -1 for first global config)
         messages: Optional chat history from frontend (list of ChatMessage)
+        attachments: Optional attachments with extracted content
+        mentioned_document_ids: Optional list of document IDs mentioned with @ in the chat
 
     Yields:
         str: SSE formatted response strings
@@ -136,13 +162,32 @@ async def stream_new_chat(
         # Build input with message history from frontend
         langchain_messages = []
 
-        # Format the user query with attachment context if any
+        # Fetch mentioned documents if any
-        final_query = user_query
+        mentioned_documents: list[Document] = []
-        if attachments:
+        if mentioned_document_ids:
-            attachment_context = format_attachments_as_context(attachments)
+            result = await session.execute(
-            final_query = (
+                select(Document).filter(
-                f"{attachment_context}\n\n<user_query>{user_query}</user_query>"
+                    Document.id.in_(mentioned_document_ids),
+                    Document.search_space_id == search_space_id,
+                )
             )
+            mentioned_documents = list(result.scalars().all())
+
+        # Format the user query with context (attachments + mentioned documents)
+        final_query = user_query
+        context_parts = []
+
+        if attachments:
+            context_parts.append(format_attachments_as_context(attachments))
+
+        if mentioned_documents:
+            context_parts.append(
+                format_mentioned_documents_as_context(mentioned_documents)
+            )
+
+        if context_parts:
+            context = "\n\n".join(context_parts)
+            final_query = f"{context}\n\n<user_query>{user_query}</user_query>"
 
         # if messages:
         #     # Convert frontend messages to LangChain format

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

@@ -6,9 +6,16 @@ import {
 	type ThreadMessageLike,
 	useExternalStoreRuntime,
 } from "@assistant-ui/react";
+import { useAtomValue, useSetAtom } from "jotai";
 import { useParams, useRouter } from "next/navigation";
 import { useCallback, useEffect, useMemo, useRef, useState } from "react";
 import { toast } from "sonner";
+import {
+	type MentionedDocumentInfo,
+	mentionedDocumentIdsAtom,
+	mentionedDocumentsAtom,
+	messageDocumentsMapAtom,
+} from "@/atoms/chat/mentioned-documents.atom";
 import { Thread } from "@/components/assistant-ui/thread";
 import { ChatHeader } from "@/components/new-chat/chat-header";
 import type { ThinkingStep } from "@/components/tool-ui/deepagent-thinking";
@@ -47,6 +54,23 @@ function extractThinkingSteps(content: unknown): ThinkingStep[] {
 	return thinkingPart?.steps || [];
 }
 
+/**
+ * Extract mentioned documents from message content
+ */
+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;
+
+	return docsPart?.documents || [];
+}
+
 /**
  * Convert backend message to assistant-ui ThreadMessageLike format
  * Filters out 'thinking-steps' part as it's handled separately via messageThinkingSteps
@@ -57,16 +81,13 @@ function convertToThreadMessage(msg: MessageRecord): ThreadMessageLike {
 	if (typeof msg.content === "string") {
 		content = [{ type: "text", text: msg.content }];
 	} else if (Array.isArray(msg.content)) {
-		// Filter out thinking-steps part - it's handled separately via messageThinkingSteps
+		// Filter out custom metadata parts - they're handled separately
-		const filteredContent = msg.content.filter(
+		const filteredContent = msg.content.filter((part: unknown) => {
-			(part: unknown) =>
+			if (typeof part !== "object" || part === null || !("type" in part)) return true;
-				!(
+			const partType = (part as { type: string }).type;
-					typeof part === "object" &&
+			// Filter out thinking-steps and mentioned-documents
-					part !== null &&
+			return partType !== "thinking-steps" && partType !== "mentioned-documents";
-					"type" in part &&
+		});
-					(part as { type: string }).type === "thinking-steps"
-				)
-		);
 		content =
 			filteredContent.length > 0
 				? (filteredContent as ThreadMessageLike["content"])
@@ -117,6 +138,13 @@ export default function NewChatPage() {
 	);
 	const abortControllerRef = useRef<AbortController | null>(null);
 
+	// Get mentioned document IDs from the composer
+	const mentionedDocumentIds = useAtomValue(mentionedDocumentIdsAtom);
+	const mentionedDocuments = useAtomValue(mentionedDocumentsAtom);
+	const setMentionedDocumentIds = useSetAtom(mentionedDocumentIdsAtom);
+	const setMentionedDocuments = useSetAtom(mentionedDocumentsAtom);
+	const setMessageDocumentsMap = useSetAtom(messageDocumentsMapAtom);
+
 	// Create the attachment adapter for file processing
 	const attachmentAdapter = useMemo(() => createAttachmentAdapter(), []);
 
@@ -154,6 +182,9 @@ export default function NewChatPage() {
 
 					// Extract and restore thinking steps from persisted messages
 					const restoredThinkingSteps = new Map<string, ThinkingStep[]>();
+					// Extract and restore mentioned documents from persisted messages
+					const restoredDocsMap: Record<string, MentionedDocumentInfo[]> = {};
+
 					for (const msg of response.messages) {
 						if (msg.role === "assistant") {
 							const steps = extractThinkingSteps(msg.content);
@@ -161,10 +192,19 @@ export default function NewChatPage() {
 								restoredThinkingSteps.set(`msg-${msg.id}`, steps);
 							}
 						}
+						if (msg.role === "user") {
+							const docs = extractMentionedDocuments(msg.content);
+							if (docs.length > 0) {
+								restoredDocsMap[`msg-${msg.id}`] = docs;
+							}
+						}
 					}
 					if (restoredThinkingSteps.size > 0) {
 						setMessageThinkingSteps(restoredThinkingSteps);
 					}
+					if (Object.keys(restoredDocsMap).length > 0) {
+						setMessageDocumentsMap(restoredDocsMap);
+					}
 				}
 			} else {
 				// Create new thread
@@ -181,7 +221,7 @@ export default function NewChatPage() {
 		} finally {
 			setIsInitializing(false);
 		}
-	}, [urlChatId, searchSpaceId, router]);
+	}, [urlChatId, searchSpaceId, router, setMessageDocumentsMap]);
 
 	// Initialize on mount
 	useEffect(() => {
@@ -243,10 +283,37 @@ export default function NewChatPage() {
 			};
 			setMessages((prev) => [...prev, userMessage]);
 
-			// Persist user message (don't await, fire and forget)
+			// Store mentioned documents with this message for display
+			if (mentionedDocuments.length > 0) {
+				const docsInfo: MentionedDocumentInfo[] = mentionedDocuments.map((doc) => ({
+					id: doc.id,
+					title: doc.title,
+					document_type: doc.document_type,
+				}));
+				setMessageDocumentsMap((prev) => ({
+					...prev,
+					[userMsgId]: docsInfo,
+				}));
+			}
+
+			// Persist user message with mentioned documents (don't await, fire and forget)
+			const persistContent =
+				mentionedDocuments.length > 0
+					? [
+							...message.content,
+							{
+								type: "mentioned-documents",
+								documents: mentionedDocuments.map((doc) => ({
+									id: doc.id,
+									title: doc.title,
+									document_type: doc.document_type,
+								})),
+							},
+						]
+					: message.content;
 			appendMessage(threadId, {
 				role: "user",
-				content: message.content,
+				content: persistContent,
 			}).catch((err) => console.error("Failed to persist user message:", err));
 
 			// Start streaming response
@@ -385,6 +452,15 @@ export default function NewChatPage() {
 				// Extract attachment content to send with the request
 				const attachments = extractAttachmentContent(messageAttachments);
 
+				// Get mentioned document IDs for context
+				const documentIds = mentionedDocumentIds.length > 0 ? [...mentionedDocumentIds] : undefined;
+
+				// Clear mentioned documents after capturing them
+				if (mentionedDocumentIds.length > 0) {
+					setMentionedDocumentIds([]);
+					setMentionedDocuments([]);
+				}
+
 				const response = await fetch(`${backendUrl}/api/v1/new_chat`, {
 					method: "POST",
 					headers: {
@@ -397,6 +473,7 @@ export default function NewChatPage() {
 						search_space_id: searchSpaceId,
 						messages: messageHistory,
 						attachments: attachments.length > 0 ? attachments : undefined,
+						mentioned_document_ids: documentIds,
 					}),
 					signal: controller.signal,
 				});
@@ -558,7 +635,16 @@ export default function NewChatPage() {
 				// Note: We no longer clear thinking steps - they persist with the message
 			}
 		},
-		[threadId, searchSpaceId, messages]
+		[
+			threadId,
+			searchSpaceId,
+			messages,
+			mentionedDocumentIds,
+			mentionedDocuments,
+			setMentionedDocumentIds,
+			setMentionedDocuments,
+			setMessageDocumentsMap,
+		]
 	);
 
 	// Convert message (pass through since already in correct format)

surfsense_web/atoms/chat/mentioned-documents.atom.ts

@@ -0,0 +1,31 @@
+"use client";
+
+import { atom } from "jotai";
+import type { Document } from "@/contracts/types/document.types";
+
+/**
+ * Atom to store the IDs of documents mentioned in the current chat composer.
+ * This is used to pass document context to the backend when sending a message.
+ */
+export const mentionedDocumentIdsAtom = atom<number[]>([]);
+
+/**
+ * Atom to store the full document objects mentioned in the current chat composer.
+ * This persists across component remounts.
+ */
+export const mentionedDocumentsAtom = atom<Document[]>([]);
+
+/**
+ * Simplified document info for display purposes
+ */
+export interface MentionedDocumentInfo {
+	id: number;
+	title: string;
+	document_type: string;
+}
+
+/**
+ * Atom to store mentioned documents per message ID.
+ * This allows displaying which documents were mentioned with each user message.
+ */
+export const messageDocumentsMapAtom = atom<Record<string, MentionedDocumentInfo[]>>({});

surfsense_web/components/assistant-ui/thread.tsx

@@ -8,8 +8,9 @@ import {
 	ThreadPrimitive,
 	useAssistantState,
 	useMessage,
+	useThreadViewport,
 } from "@assistant-ui/react";
-import { useAtomValue } from "jotai";
+import { useAtom, useAtomValue, useSetAtom } from "jotai";
 import {
 	AlertCircle,
 	ArrowDownIcon,
@@ -21,6 +22,7 @@ import {
 	ChevronRightIcon,
 	CopyIcon,
 	DownloadIcon,
+	FileText,
 	Loader2,
 	PencilIcon,
 	Plug2,
@@ -29,10 +31,27 @@ import {
 	Search,
 	Sparkles,
 	SquareIcon,
+	X,
 } from "lucide-react";
 import Link from "next/link";
-import { type FC, useCallback, useEffect, useMemo, useRef, useState } from "react";
+import { useParams } from "next/navigation";
+import {
+	createContext,
+	type FC,
+	useCallback,
+	useContext,
+	useEffect,
+	useMemo,
+	useRef,
+	useState,
+} from "react";
+import { createPortal } from "react-dom";
 import { getDocumentTypeLabel } from "@/app/dashboard/[search_space_id]/documents/(manage)/components/DocumentTypeIcon";
+import {
+	mentionedDocumentIdsAtom,
+	mentionedDocumentsAtom,
+	messageDocumentsMapAtom,
+} from "@/atoms/chat/mentioned-documents.atom";
 import { documentTypeCountsAtom } from "@/atoms/documents/document-query.atoms";
 import {
 	globalNewLLMConfigsAtom,
@@ -49,6 +68,10 @@ import {
 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";
 import {
 	ChainOfThought,
 	ChainOfThoughtContent,
@@ -60,6 +83,7 @@ import type { ThinkingStep } from "@/components/tool-ui/deepagent-thinking";
 import { Button } from "@/components/ui/button";
 import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover";
 import { getConnectorIcon } from "@/contracts/enums/connectorIcons";
+import type { Document } from "@/contracts/types/document.types";
 import { useSearchSourceConnectors } from "@/hooks/use-search-source-connectors";
 import { cn } from "@/lib/utils";
 
@@ -73,8 +97,6 @@ interface ThreadProps {
 }
 
 // Context to pass thinking steps to AssistantMessage
-import { createContext, useContext } from "react";
-
 const ThinkingStepsContext = createContext<Map<string, ThinkingStep[]>>(new Map());
 
 /**
@@ -213,6 +235,56 @@ const ThinkingStepsDisplay: FC<{ steps: ThinkingStep[]; isThreadRunning?: boolea
 	);
 };
 
+/**
+ * Component that handles auto-scroll when thinking steps update.
+ * Uses useThreadViewport to scroll to bottom when thinking steps change,
+ * ensuring the user always sees the latest content during streaming.
+ */
+const ThinkingStepsScrollHandler: FC = () => {
+	const thinkingStepsMap = useContext(ThinkingStepsContext);
+	const viewport = useThreadViewport();
+	const isRunning = useAssistantState(({ thread }) => thread.isRunning);
+	// Track the serialized state to detect any changes
+	const prevStateRef = useRef<string>("");
+
+	useEffect(() => {
+		// Only act during streaming
+		if (!isRunning) {
+			prevStateRef.current = "";
+			return;
+		}
+
+		// Serialize the thinking steps state to detect any changes
+		// This catches new steps, status changes, and item additions
+		let stateString = "";
+		thinkingStepsMap.forEach((steps, msgId) => {
+			steps.forEach((step) => {
+				stateString += `${msgId}:${step.id}:${step.status}:${step.items?.length || 0};`;
+			});
+		});
+
+		// If state changed at all during streaming, scroll
+		if (stateString !== prevStateRef.current && stateString !== "") {
+			prevStateRef.current = stateString;
+
+			// Multiple attempts to ensure scroll happens after DOM updates
+			const scrollAttempt = () => {
+				try {
+					viewport.scrollToBottom();
+				} catch {
+					// Ignore errors - viewport might not be ready
+				}
+			};
+
+			// Delayed attempts to handle async DOM updates
+			requestAnimationFrame(scrollAttempt);
+			setTimeout(scrollAttempt, 100);
+		}
+	}, [thinkingStepsMap, viewport, isRunning]);
+
+	return null; // This component doesn't render anything
+};
+
 export const Thread: FC<ThreadProps> = ({ messageThinkingSteps = new Map(), header }) => {
 	return (
 		<ThinkingStepsContext.Provider value={messageThinkingSteps}>
@@ -316,12 +388,15 @@ const getTimeBasedGreeting = (userEmail?: string): string => {
 const ThreadWelcome: FC = () => {
 	const { data: user } = useAtomValue(currentUserAtom);
 
+	// Memoize greeting so it doesn't change on re-renders (only on user change)
+	const greeting = useMemo(() => getTimeBasedGreeting(user?.email), [user?.email]);
+
 	return (
 		<div className="aui-thread-welcome-root mx-auto flex w-full max-w-(--thread-max-width) grow flex-col items-center px-4 relative">
 			{/* Greeting positioned above the composer - fixed position */}
 			<div className="aui-thread-welcome-message absolute bottom-[calc(50%+5rem)] left-0 right-0 flex flex-col items-center text-center">
 				<h1 className="aui-thread-welcome-message-inner fade-in slide-in-from-bottom-2 animate-in text-5xl delay-100 duration-500 ease-out fill-mode-both">
-					{getTimeBasedGreeting(user?.email)}
+					{greeting}
 				</h1>
 			</div>
 			{/* Composer - top edge fixed, expands downward only */}
@@ -333,42 +408,226 @@ const ThreadWelcome: FC = () => {
 };
 
 const Composer: FC = () => {
-	// Check if a model is configured - needed to disable input
+	// ---- State for document mentions (using atoms to persist across remounts) ----
-	const { data: userConfigs } = useAtomValue(newLLMConfigsAtom);
+	const [mentionedDocuments, setMentionedDocuments] = useAtom(mentionedDocumentsAtom);
-	const { data: globalConfigs } = useAtomValue(globalNewLLMConfigsAtom);
+	const [showDocumentPopover, setShowDocumentPopover] = useState(false);
-	const { data: preferences } = useAtomValue(llmPreferencesAtom);
+	const [mentionQuery, setMentionQuery] = useState("");
+	const inputRef = useRef<HTMLTextAreaElement | null>(null);
+	const documentPickerRef = useRef<DocumentsDataTableRef>(null);
+	const { search_space_id } = useParams();
+	const setMentionedDocumentIds = useSetAtom(mentionedDocumentIdsAtom);
 
-	const hasModelConfigured = useMemo(() => {
+	// Sync mentioned document IDs to atom for use in chat request
-		if (!preferences) return false;
+	useEffect(() => {
-		const agentLlmId = preferences.agent_llm_id;
+		setMentionedDocumentIds(mentionedDocuments.map((doc) => doc.id));
-		if (agentLlmId === null || agentLlmId === undefined) return false;
+	}, [mentionedDocuments, setMentionedDocumentIds]);
 
-		// Check if the configured model actually exists
+	// Extract mention query (text after @)
-		if (agentLlmId < 0) {
+	const extractMentionQuery = useCallback((value: string): string => {
-			return globalConfigs?.some((c) => c.id === agentLlmId) ?? false;
+		const atIndex = value.lastIndexOf("@");
+		if (atIndex === -1) return "";
+		return value.slice(atIndex + 1);
+	}, []);
+
+	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;
 		}
-		return userConfigs?.some((c) => c.id === agentLlmId) ?? false;
+
-	}, [preferences, globalConfigs, userConfigs]);
+		// 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
+		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;
+			}
+		}
+
+		// 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 }));
+				}
+			}
+			// Focus the input so user can continue typing
+			input.focus();
+		}
+
+		// Reset mention query
+		setMentionQuery("");
+	};
+
+	const handleRemoveDocument = (docId: number) => {
+		setMentionedDocuments((prev) => prev.filter((doc) => doc.id !== docId));
+	};
 
 	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 />
-				<ComposerPrimitive.Input
+				{/* -------- Input field with inline document chips -------- */}
-					placeholder={
+				<div className="aui-composer-input-wrapper flex flex-wrap items-center gap-1.5 px-3 pt-2 pb-6">
-						hasModelConfigured
+					{/* Inline document chips */}
-							? "Ask SurfSense"
+					{mentionedDocuments.map((doc) => (
-							: "Select a model from the header to start chatting..."
+						<span
-					}
+							key={doc.id}
-					className={cn(
+							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"
-						"aui-composer-input mb-1 max-h-32 min-h-14 w-full resize-none bg-transparent px-4 pt-2 pb-3 text-sm outline-none placeholder:text-muted-foreground focus-visible:ring-0",
+							title={doc.title}
-						!hasModelConfigured && "cursor-not-allowed opacity-60"
+						>
+							<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}
+						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"
+					/>
+				</div>
+
+				{/* -------- Document mention popover (rendered via portal) -------- */}
+				{showDocumentPopover &&
+					typeof document !== "undefined" &&
+					createPortal(
+						<>
+							{/* Backdrop */}
+							<button
+								type="button"
+								className="fixed inset-0 cursor-default"
+								style={{ zIndex: 9998 }}
+								onClick={() => setShowDocumentPopover(false)}
+								aria-label="Close document picker"
+							/>
+							{/* Popover positioned above input */}
+							<div
+								className="fixed shadow-2xl rounded-lg border border-border overflow-hidden"
+								style={{
+									zIndex: 9999,
+									backgroundColor: "#18181b",
+									bottom: inputRef.current
+										? `${window.innerHeight - inputRef.current.getBoundingClientRect().top + 8}px`
+										: "200px",
+									left: inputRef.current
+										? `${inputRef.current.getBoundingClientRect().left}px`
+										: "50%",
+								}}
+							>
+								<DocumentsDataTable
+									ref={documentPickerRef}
+									searchSpaceId={Number(search_space_id)}
+									onSelectionChange={handleDocumentsMention}
+									onDone={() => {
+										setShowDocumentPopover(false);
+										setMentionQuery("");
+									}}
+									initialSelectedDocuments={mentionedDocuments}
+									externalSearch={mentionQuery}
+								/>
+							</div>
+						</>,
+						document.body
 					)}
-					rows={1}
-					autoFocus={hasModelConfigured}
-					aria-label="Message input"
-					disabled={!hasModelConfigured}
-				/>
 				<ComposerAction />
 			</ComposerPrimitive.AttachmentDropzone>
 		</ComposerPrimitive.Root>
@@ -469,7 +728,7 @@ const ConnectorIndicator: FC = () => {
 						</div>
 						<div className="flex flex-wrap gap-2">
 							{/* Document types from the search space */}
-							{activeDocumentTypes.map(([docType, count]) => (
+							{activeDocumentTypes.map(([docType]) => (
 								<div
 									key={docType}
 									className="flex items-center gap-1.5 rounded-md bg-muted/80 px-2.5 py-1.5 text-xs border border-border/50"
@@ -642,7 +901,7 @@ const ThinkingStepsPart: FC = () => {
 	const thinkingStepsMap = useContext(ThinkingStepsContext);
 
 	// Get the current message ID to look up thinking steps
-	const messageId = useMessage((m) => m.id);
+	const messageId = useAssistantState(({ message }) => message?.id);
 	const thinkingSteps = thinkingStepsMap.get(messageId) || [];
 
 	// Check if thread is still running (for stopping the spinner when cancelled)
@@ -725,6 +984,10 @@ const AssistantActionBar: FC = () => {
 };
 
 const UserMessage: FC = () => {
+	const messageId = useAssistantState(({ message }) => message?.id);
+	const messageDocumentsMap = useAtomValue(messageDocumentsMapAtom);
+	const mentionedDocs = messageId ? messageDocumentsMap[messageId] : undefined;
+
 	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"
@@ -733,10 +996,25 @@ const UserMessage: FC = () => {
 			<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) => (
+							<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"
+								title={doc.title}
+							>
+								<FileText className="size-3" />
+								<span className="max-w-[150px] truncate">{doc.title}</span>
+							</span>
+						))}
+					</div>
+				)}
 				<div className="aui-user-message-content wrap-break-word rounded-2xl bg-muted px-4 py-2.5 text-foreground">
 					<MessagePrimitive.Parts />
 				</div>
-				<div className="aui-user-action-bar-wrapper -translate-x-full -translate-y-1/2 absolute top-1/2 left-0 pr-2">
+				<div className="aui-user-action-bar-wrapper -translate-x-full -translate-y-full absolute top-full left-0 pr-2">
 					<UserActionBar />
 				</div>
 			</div>

surfsense_web/components/new-chat/DocumentsDataTable.tsx

@@ -0,0 +1,248 @@
+"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>
+		);
+	}
+);