SurfSense/surfsense_web/components/assistant-ui/composer.tsx

import { ComposerPrimitive, useAssistantState, useComposerRuntime } from "@assistant-ui/react";
import { useAtom, useSetAtom } from "jotai";
import { useParams } from "next/navigation";
import type { FC } from "react";
import { useCallback, useEffect, useRef, useState } from "react";
import { createPortal } from "react-dom";
import {
	mentionedDocumentIdsAtom,
	mentionedDocumentsAtom,
} from "@/atoms/chat/mentioned-documents.atom";
import { ComposerAddAttachment, ComposerAttachments } from "@/components/assistant-ui/attachment";
import { ComposerAction } from "@/components/assistant-ui/composer-action";
import {
	InlineMentionEditor,
	type InlineMentionEditorRef,
} from "@/components/assistant-ui/inline-mention-editor";
import {
	DocumentMentionPicker,
	type DocumentMentionPickerRef,
} from "@/components/new-chat/document-mention-picker";
import type { Document } from "@/contracts/types/document.types";

export const Composer: FC = () => {
	// ---- State for document mentions (using atoms to persist across remounts) ----
	const [mentionedDocuments, setMentionedDocuments] = useAtom(mentionedDocumentsAtom);
	const [showDocumentPopover, setShowDocumentPopover] = useState(false);
	const [mentionQuery, setMentionQuery] = useState("");
	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]);

	// Sync mentioned document IDs to atom for use in chat request
	useEffect(() => {
		setMentionedDocumentIds(mentionedDocuments.map((doc) => doc.id));
	}, [mentionedDocuments, setMentionedDocumentIds]);

	// 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);
	}, []);

	// Handle mention close
	const handleMentionClose = useCallback(() => {
		if (showDocumentPopover) {
			setShowDocumentPopover(false);
			setMentionQuery("");
		}
	}, [showDocumentPopover]);

	// 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;
				}
			}
		},
		[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,
	]);

	// 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]
	);

	// 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 />
				{/* -------- Inline Mention Editor -------- */}
				<div ref={editorContainerRef} className="aui-composer-input-wrapper px-3 pt-3 pb-6">
					<InlineMentionEditor
						ref={editorRef}
						placeholder="Ask SurfSense or @mention docs"
						onMentionTrigger={handleMentionTrigger}
						onMentionClose={handleMentionClose}
						onChange={handleEditorChange}
						onDocumentRemove={handleDocumentRemove}
						onSubmit={handleSubmit}
						onKeyDown={handleKeyDown}
						className="min-h-[24px]"
					/>
				</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 bg-popover"
								style={{
									zIndex: 9999,
									bottom: editorContainerRef.current
										? `${window.innerHeight - editorContainerRef.current.getBoundingClientRect().top + 8}px`
										: "200px",
									left: editorContainerRef.current
										? `${editorContainerRef.current.getBoundingClientRect().left}px`
										: "50%",
								}}
							>
								<DocumentMentionPicker
									ref={documentPickerRef}
									searchSpaceId={Number(search_space_id)}
									onSelectionChange={handleDocumentsMention}
									onDone={() => {
										setShowDocumentPopover(false);
										setMentionQuery("");
									}}
									initialSelectedDocuments={mentionedDocuments}
									externalSearch={mentionQuery}
								/>
							</div>
						</>,
						document.body
					)}
				<ComposerAction />
			</ComposerPrimitive.AttachmentDropzone>
		</ComposerPrimitive.Root>
	);
};