feat: implement document mention extraction and management in new chat

- Added functionality to extract and manage mentioned documents within the new chat interface.
- Introduced new atoms for storing mentioned documents and their mappings to user messages.
- Enhanced the message persistence logic to include mentioned documents, ensuring they are displayed correctly in the chat.
- Updated the UI components to support document mentions, including a refined document selection interface.
- Improved state management for document mentions to ensure a seamless user experience.

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

@@ -10,7 +10,7 @@ import { useAtomValue, useSetAtom } from "jotai";
 import { useParams, useRouter } from "next/navigation";
 import { useCallback, useEffect, useMemo, useRef, useState } from "react";
 import { toast } from "sonner";
-import { mentionedDocumentIdsAtom } from "@/atoms/chat/mentioned-documents.atom";
+import { mentionedDocumentIdsAtom, mentionedDocumentsAtom, messageDocumentsMapAtom, type MentionedDocumentInfo } from "@/atoms/chat/mentioned-documents.atom";
 import { Thread } from "@/components/assistant-ui/thread";
 import { GeneratePodcastToolUI } from "@/components/tool-ui/generate-podcast";
 import { LinkPreviewToolUI } from "@/components/tool-ui/link-preview";
@@ -48,6 +48,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
@@ -58,13 +75,14 @@ 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(
-			(part: unknown) => 
+			(part: unknown) => {
-				!(typeof part === "object" && 
+				if (typeof part !== "object" || part === null || !("type" in part)) return true;
-				  part !== null && 
+				const partType = (part as { type: string }).type;
-				  "type" in part && 
+				// Filter out thinking-steps and mentioned-documents
-				  (part as { type: string }).type === "thinking-steps")
+				return partType !== "thinking-steps" && partType !== "mentioned-documents";
+			}
 		);
 		content = filteredContent.length > 0 
 			? (filteredContent as ThreadMessageLike["content"])
@@ -111,7 +129,10 @@ export default function NewChatPage() {
 
 	// 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(), []);
@@ -150,6 +171,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);
@@ -157,10 +181,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
@@ -239,10 +272,33 @@ 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
@@ -384,6 +440,7 @@ export default function NewChatPage() {
 				// Clear mentioned documents after capturing them
 				if (mentionedDocumentIds.length > 0) {
 					setMentionedDocumentIds([]);
+					setMentionedDocuments([]);
 				}
 
 				const response = await fetch(`${backendUrl}/api/v1/new_chat`, {
@@ -561,7 +618,7 @@ export default function NewChatPage() {
 				// Note: We no longer clear thinking steps - they persist with the message
 			}
 		},
-		[threadId, searchSpaceId, messages, mentionedDocumentIds, setMentionedDocumentIds]
+		[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

@@ -1,6 +1,7 @@
 "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.
@@ -8,3 +9,24 @@ import { atom } from "jotai";
  */
 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

@@ -18,6 +18,7 @@ import {
 	ChevronRightIcon,
 	CopyIcon,
 	DownloadIcon,
+	FileText,
 	Loader2,
 	PencilIcon,
 	Plug2,
@@ -32,8 +33,8 @@ import { useParams } from "next/navigation";
 import Link from "next/link";
 import { type FC, useState, useRef, useCallback, useEffect, createContext, useContext, useMemo } from "react";
 import { createPortal } from "react-dom";
-import { useAtomValue, useSetAtom } from "jotai";
+import { useAtom, useAtomValue, useSetAtom } from "jotai";
-import { mentionedDocumentIdsAtom } from "@/atoms/chat/mentioned-documents.atom";
+import { mentionedDocumentIdsAtom, mentionedDocumentsAtom, messageDocumentsMapAtom } from "@/atoms/chat/mentioned-documents.atom";
 import { activeSearchSpaceIdAtom } from "@/atoms/search-spaces/search-space-query.atoms";
 import { documentTypeCountsAtom } from "@/atoms/documents/document-query.atoms";
 import { useSearchSourceConnectors } from "@/hooks/use-search-source-connectors";
@@ -55,7 +56,7 @@ import {
 	ChainOfThoughtStep,
 	ChainOfThoughtTrigger,
 } from "@/components/prompt-kit/chain-of-thought";
-import { DocumentsDataTable } from "@/components/new-chat/DocumentsDataTable";
+import { DocumentsDataTable, type DocumentsDataTableRef } from "@/components/new-chat/DocumentsDataTable";
 import { Button } from "@/components/ui/button";
 import type { Document } from "@/contracts/types/document.types";
 import { cn } from "@/lib/utils";
@@ -352,10 +353,12 @@ const ThreadWelcome: FC = () => {
 };
 
 const Composer: FC = () => {
-	// ---- State for document mentions ----
+	// ---- State for document mentions (using atoms to persist across remounts) ----
-	const [mentionedDocuments, setMentionedDocuments] = useState<Document[]>([]);
+	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 { search_space_id } = useParams();
 	const setMentionedDocumentIds = useSetAtom(mentionedDocumentIdsAtom);
 
@@ -364,6 +367,13 @@ const Composer: FC = () => {
 		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);
+	}, []);
+
 	const handleKeyUp = (e: React.KeyboardEvent<HTMLTextAreaElement>) => {
 		const textarea = e.currentTarget;
 		const value = textarea.value;
@@ -371,20 +381,60 @@ const Composer: FC = () => {
 		// Open document picker when user types '@'
 		if (e.key === "@" || (e.key === "2" && e.shiftKey)) {
 			setShowDocumentPopover(true);
+			setMentionQuery("");
+			return;
 		}
 
-		// Close popover if '@' is no longer in the input (user deleted it)
+		// Check if value contains @ and extract query
-		if (showDocumentPopover && !value.includes("@")) {
+		if (value.includes("@")) {
-			setShowDocumentPopover(false);
+			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>) => {
-		// Close popover on Escape
+		// When popup is open, handle navigation keys
-		if (e.key === "Escape" && showDocumentPopover) {
+		if (showDocumentPopover) {
-			e.preventDefault();
+			if (e.key === "ArrowDown") {
-			setShowDocumentPopover(false);
+				e.preventDefault();
-			return;
+				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
@@ -410,25 +460,30 @@ const Composer: FC = () => {
 			return [...prev, ...newDocs];
 		});
 
-		// Clean up the '@' trigger from input if present
+		// Clean up the '@...' mention text from input
 		if (inputRef.current) {
 			const input = inputRef.current;
 			const currentValue = input.value;
-			// Remove trailing @ if it exists
+			const atIndex = currentValue.lastIndexOf("@");
-			if (currentValue.endsWith("@")) {
+			
-				// Use a native input event to properly update the controlled component
+			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, currentValue.slice(0, -1));
+					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) => {
@@ -494,10 +549,15 @@ const Composer: FC = () => {
 							}}
 						>
 							<DocumentsDataTable
+								ref={documentPickerRef}
 								searchSpaceId={Number(search_space_id)}
 								onSelectionChange={handleDocumentsMention}
-								onDone={() => setShowDocumentPopover(false)}
+								onDone={() => {
+									setShowDocumentPopover(false);
+									setMentionQuery("");
+								}}
 								initialSelectedDocuments={mentionedDocuments}
+								externalSearch={mentionQuery}
 							/>
 						</div>
 					</>,
@@ -819,6 +879,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"
@@ -827,6 +891,21 @@ 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>

surfsense_web/components/new-chat/DocumentsDataTable.tsx

@@ -1,21 +1,26 @@
 "use client";
 
 import { useQuery } from "@tanstack/react-query";
-import { FileText, Search } from "lucide-react";
+import { FileText } from "lucide-react";
-import { useCallback, useEffect, useMemo, useRef, useState } from "react";
+import { forwardRef, useCallback, useEffect, useImperativeHandle, useMemo, useRef, useState } from "react";
-import { Input } from "@/components/ui/input";
-import { ScrollArea } from "@/components/ui/scroll-area";
 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) {
@@ -27,190 +32,206 @@ function useDebounced<T>(value: T, delay = 300) {
 	return debounced;
 }
 
-export function DocumentsDataTable({
+export const DocumentsDataTable = forwardRef<DocumentsDataTableRef, DocumentsDataTableProps>(
-	searchSpaceId,
+	function DocumentsDataTable({
-	onSelectionChange,
+		searchSpaceId,
-	onDone,
+		onSelectionChange,
-	initialSelectedDocuments = [],
+		onDone,
-}: DocumentsDataTableProps) {
+		initialSelectedDocuments = [],
-	const [search, setSearch] = useState("");
+		externalSearch = "",
-	const debouncedSearch = useDebounced(search, 300);
+	}, ref) {
-	const [highlightedIndex, setHighlightedIndex] = useState(0);
+		// Use external search
-	const listRef = useRef<HTMLDivElement>(null);
+		const search = externalSearch;
-	const itemRefs = useRef<Map<number, HTMLButtonElement>>(new Map());
+		const debouncedSearch = useDebounced(search, 150);
+		const [highlightedIndex, setHighlightedIndex] = useState(0);
+		const itemRefs = useRef<Map<number, HTMLButtonElement>>(new Map());
 
-	const fetchQueryParams = useMemo(
+		const fetchQueryParams = useMemo(
-		() => ({
+			() => ({
-			search_space_id: searchSpaceId,
+				search_space_id: searchSpaceId,
-			page: 0,
+				page: 0,
-			page_size: 20,
+				page_size: 20,
-		}),
+			}),
-		[searchSpaceId]
+			[searchSpaceId]
-	);
+		);
 
-	const searchQueryParams = useMemo(() => {
+		const searchQueryParams = useMemo(() => {
-		return {
+			return {
-			search_space_id: searchSpaceId,
+				search_space_id: searchSpaceId,
-			page: 0,
+				page: 0,
-			page_size: 20,
+				page_size: 20,
-			title: debouncedSearch,
+				title: debouncedSearch,
-		};
+			};
-	}, [debouncedSearch, searchSpaceId]);
+		}, [debouncedSearch, searchSpaceId]);
 
-	// Use query for fetching documents
+		// Use query for fetching documents
-	const { data: documents, isLoading: isDocumentsLoading } = useQuery({
+		const { data: documents, isLoading: isDocumentsLoading } = useQuery({
-		queryKey: cacheKeys.documents.withQueryParams(fetchQueryParams),
+			queryKey: cacheKeys.documents.withQueryParams(fetchQueryParams),
-		queryFn: () => documentsApiService.getDocuments({ queryParams: fetchQueryParams }),
+			queryFn: () => documentsApiService.getDocuments({ queryParams: fetchQueryParams }),
-		staleTime: 3 * 60 * 1000,
+			staleTime: 3 * 60 * 1000,
-		enabled: !!searchSpaceId && !debouncedSearch.trim(),
+			enabled: !!searchSpaceId && !debouncedSearch.trim(),
-	});
+		});
 
-	// Searching
+		// Searching
-	const { data: searchedDocuments, isLoading: isSearchedDocumentsLoading } = useQuery({
+		const { data: searchedDocuments, isLoading: isSearchedDocumentsLoading } = useQuery({
-		queryKey: cacheKeys.documents.withQueryParams(searchQueryParams),
+			queryKey: cacheKeys.documents.withQueryParams(searchQueryParams),
-		queryFn: () => documentsApiService.searchDocuments({ queryParams: searchQueryParams }),
+			queryFn: () => documentsApiService.searchDocuments({ queryParams: searchQueryParams }),
-		staleTime: 3 * 60 * 1000,
+			staleTime: 3 * 60 * 1000,
-		enabled: !!searchSpaceId && !!debouncedSearch.trim(),
+			enabled: !!searchSpaceId && !!debouncedSearch.trim(),
-	});
+		});
 
-	const actualDocuments = debouncedSearch.trim()
+		const actualDocuments = debouncedSearch.trim()
-		? searchedDocuments?.items || []
+			? searchedDocuments?.items || []
-		: documents?.items || [];
+			: documents?.items || [];
-	const actualLoading = debouncedSearch.trim() ? isSearchedDocumentsLoading : isDocumentsLoading;
+		const actualLoading = debouncedSearch.trim() ? isSearchedDocumentsLoading : isDocumentsLoading;
 
-	// Track already selected document IDs
+		// Track already selected document IDs
-	const selectedIds = useMemo(
+		const selectedIds = useMemo(
-		() => new Set(initialSelectedDocuments.map((d) => d.id)),
+			() => new Set(initialSelectedDocuments.map((d) => d.id)),
-		[initialSelectedDocuments]
+			[initialSelectedDocuments]
-	);
+		);
 
-	// Filter out already selected documents for navigation
+		// Filter out already selected documents for navigation
-	const selectableDocuments = useMemo(
+		const selectableDocuments = useMemo(
-		() => actualDocuments.filter((doc) => !selectedIds.has(doc.id)),
+			() => actualDocuments.filter((doc) => !selectedIds.has(doc.id)),
-		[actualDocuments, selectedIds]
+			[actualDocuments, selectedIds]
-	);
+		);
 
-	const handleSelectDocument = useCallback((doc: Document) => {
+		const handleSelectDocument = useCallback((doc: Document) => {
-		onSelectionChange([...initialSelectedDocuments, doc]);
+			onSelectionChange([...initialSelectedDocuments, doc]);
-		onDone();
+			onDone();
-	}, [initialSelectedDocuments, onSelectionChange, onDone]);
+		}, [initialSelectedDocuments, onSelectionChange, onDone]);
 
-	// Scroll highlighted item into view
+		// Scroll highlighted item into view
-	useEffect(() => {
+		useEffect(() => {
-		const item = itemRefs.current.get(highlightedIndex);
+			const item = itemRefs.current.get(highlightedIndex);
-		if (item) {
+			if (item) {
-			item.scrollIntoView({ block: "nearest", behavior: "smooth" });
+				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);
+			}
 		}
-	}, [highlightedIndex]);
 
-	// Handle keyboard navigation
+		// Expose methods to parent via ref
-	const handleKeyDown = useCallback((e: React.KeyboardEvent) => {
+		useImperativeHandle(ref, () => ({
-		if (selectableDocuments.length === 0) return;
+			selectHighlighted: () => {
-
-		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":
+			moveUp: () => {
-				e.preventDefault();
+				setHighlightedIndex((prev) => 
-				onDone();
+					prev > 0 ? prev - 1 : selectableDocuments.length - 1
-				break;
+				);
-		}
+			},
-	}, [selectableDocuments, highlightedIndex, handleSelectDocument, onDone]);
+			moveDown: () => {
+				setHighlightedIndex((prev) => 
+					prev < selectableDocuments.length - 1 ? prev + 1 : 0
+				);
+			},
+		}), [selectableDocuments, highlightedIndex, handleSelectDocument]);
 
-	return (
+		// Handle keyboard navigation
-		<div 
+		const handleKeyDown = useCallback((e: React.KeyboardEvent) => {
-			className="flex flex-col w-[280px] sm:w-[320px] bg-zinc-900"
+			if (selectableDocuments.length === 0) return;
-			onKeyDown={handleKeyDown}
+
-			role="listbox"
+			switch (e.key) {
-			tabIndex={-1}
+				case "ArrowDown":
-		>
+					e.preventDefault();
-			{/* Search */}
+					setHighlightedIndex((prev) => 
-			<div className="relative p-2 border-b">
+						prev < selectableDocuments.length - 1 ? prev + 1 : 0
-				<Search className="absolute left-4 top-1/2 transform -translate-y-1/2 h-4 w-4 text-muted-foreground" />
+					);
-				<Input
+					break;
-					placeholder="Search documents..."
+				case "ArrowUp":
-					value={search}
+					e.preventDefault();
-					onChange={(e) => {
+					setHighlightedIndex((prev) => 
-						setSearch(e.target.value);
+						prev > 0 ? prev - 1 : selectableDocuments.length - 1
-						setHighlightedIndex(0);
+					);
-					}}
+					break;
-					className="pl-8 h-8 text-sm border-0 focus-visible:ring-0 focus-visible:ring-offset-0"
+				case "Enter":
-					autoFocus
+					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>
-
+		);
-			{/* Document List */}
+	}
-			<ScrollArea className="h-[240px]">
+);
-				{actualLoading ? (
-					<div className="flex items-center justify-center py-6">
-						<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-6 text-center px-4">
-						<FileText className="h-6 w-6 text-muted-foreground/50 mb-2" />
-						<p className="text-sm text-muted-foreground">No documents found</p>
-					</div>
-				) : (
-					<div className="py-1" ref={listRef}>
-						{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>
-				)}
-			</ScrollArea>
-		</div>
-	);
-}