refactor: streamline document mention handling by consolidating atom logic and removing redundant state updates, enhancing code clarity and maintainability in chat components

This commit is contained in:
Anish Sarkar 2026-03-06 23:33:51 +05:30
parent 8d5d8e490c
commit c1ba3a9b6d
6 changed files with 47 additions and 90 deletions

View file

@ -169,7 +169,7 @@ export function DocumentsFilters({
className="peer h-9 w-full pl-9 pr-9 text-sm bg-sidebar border-border/60 focus-visible:ring-1 focus-visible:ring-ring/30 select-none focus:select-text"
value={searchValue}
onChange={(e) => onSearch(e.target.value)}
placeholder="Search"
placeholder="Search docs"
type="text"
aria-label={t("filter_placeholder")}
/>

View file

@ -180,11 +180,10 @@ export default function NewChatPage() {
interruptData: Record<string, unknown>;
} | null>(null);
// Get mentioned document IDs from the composer (combines @ mentions + sidebar selections)
// Get mentioned document IDs from the composer (derived from @ mentions + sidebar selections)
const mentionedDocumentIds = useAtomValue(mentionedDocumentIdsAtom);
const mentionedDocuments = useAtomValue(mentionedDocumentsAtom);
const sidebarDocuments = useAtomValue(sidebarSelectedDocumentsAtom);
const setMentionedDocumentIds = useSetAtom(mentionedDocumentIdsAtom);
const setMentionedDocuments = useSetAtom(mentionedDocumentsAtom);
const setSidebarDocuments = useSetAtom(sidebarSelectedDocumentsAtom);
const setMessageDocumentsMap = useSetAtom(messageDocumentsMapAtom);
@ -278,11 +277,8 @@ export default function NewChatPage() {
setThreadId(null);
setCurrentThread(null);
setMessageThinkingSteps(new Map());
setMentionedDocumentIds({
surfsense_doc_ids: [],
document_ids: [],
});
setMentionedDocuments([]);
setSidebarDocuments([]);
setMessageDocumentsMap({});
clearPlanOwnerRegistry(); // Reset plan ownership for new chat
closeReportPanel(); // Close report panel when switching chats
@ -347,8 +343,8 @@ export default function NewChatPage() {
}, [
urlChatId,
setMessageDocumentsMap,
setMentionedDocumentIds,
setMentionedDocuments,
setSidebarDocuments,
closeReportPanel,
]);
@ -619,10 +615,6 @@ export default function NewChatPage() {
// Clear mentioned documents after capturing them
if (hasDocumentIds || hasSurfsenseDocIds) {
setMentionedDocumentIds({
surfsense_doc_ids: [],
document_ids: [],
});
setMentionedDocuments([]);
setSidebarDocuments([]);
}
@ -914,7 +906,6 @@ export default function NewChatPage() {
mentionedDocumentIds,
mentionedDocuments,
sidebarDocuments,
setMentionedDocumentIds,
setMentionedDocuments,
setSidebarDocuments,
setMessageDocumentsMap,

View file

@ -3,18 +3,6 @@
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<{
surfsense_doc_ids: number[];
document_ids: number[];
}>({
surfsense_doc_ids: [],
document_ids: [],
});
/**
* Atom to store the full document objects mentioned via @-mention chips
* in the current chat composer. This persists across component remounts.
@ -27,6 +15,31 @@ export const mentionedDocumentsAtom = atom<Pick<Document, "id" | "title" | "docu
*/
export const sidebarSelectedDocumentsAtom = atom<Pick<Document, "id" | "title" | "document_type">[]>([]);
/**
* Derived read-only atom that merges @-mention chips and sidebar selections
* into a single deduplicated set of document IDs for the backend.
*/
export const mentionedDocumentIdsAtom = atom((get) => {
const chipDocs = get(mentionedDocumentsAtom);
const sidebarDocs = get(sidebarSelectedDocumentsAtom);
const allDocs = [...chipDocs, ...sidebarDocs];
const seen = new Set<string>();
const deduped = allDocs.filter((d) => {
const key = `${d.document_type}:${d.id}`;
if (seen.has(key)) return false;
seen.add(key);
return true;
});
return {
surfsense_doc_ids: deduped
.filter((doc) => doc.document_type === "SURFSENSE_DOCS")
.map((doc) => doc.id),
document_ids: deduped
.filter((doc) => doc.document_type !== "SURFSENSE_DOCS")
.map((doc) => doc.id),
};
});
/**
* Simplified document info for display purposes
*/

View file

@ -30,7 +30,6 @@ import { chatSessionStateAtom } from "@/atoms/chat/chat-session-state.atom";
import { showCommentsGutterAtom } from "@/atoms/chat/current-thread.atom";
import { documentsSidebarOpenAtom } from "@/atoms/documents/ui.atoms";
import {
mentionedDocumentIdsAtom,
mentionedDocumentsAtom,
sidebarSelectedDocumentsAtom,
} from "@/atoms/chat/mentioned-documents.atom";
@ -227,14 +226,13 @@ const ThreadWelcome: FC = () => {
const Composer: FC = () => {
// Document mention state (atoms persist across component remounts)
const [mentionedDocuments, setMentionedDocuments] = useAtom(mentionedDocumentsAtom);
const [sidebarDocs, setSidebarDocs] = useAtom(sidebarSelectedDocumentsAtom);
const setSidebarDocs = useSetAtom(sidebarSelectedDocumentsAtom);
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, chat_id } = useParams();
const setMentionedDocumentIds = useSetAtom(mentionedDocumentIdsAtom);
const composerRuntime = useComposerRuntime();
const hasAutoFocusedRef = useRef(false);
@ -309,26 +307,6 @@ const Composer: FC = () => {
}
}, [isThreadEmpty]);
// Combine sidebar selections + @-mention chips → single ID atom for the backend
useEffect(() => {
const allDocs = [...mentionedDocuments, ...sidebarDocs];
const seen = new Set<string>();
const deduped = allDocs.filter((d) => {
const key = `${d.document_type}:${d.id}`;
if (seen.has(key)) return false;
seen.add(key);
return true;
});
setMentionedDocumentIds({
surfsense_doc_ids: deduped
.filter((doc) => doc.document_type === "SURFSENSE_DOCS")
.map((doc) => doc.id),
document_ids: deduped
.filter((doc) => doc.document_type !== "SURFSENSE_DOCS")
.map((doc) => doc.id),
});
}, [mentionedDocuments, sidebarDocs, setMentionedDocumentIds]);
// Sync editor text with assistant-ui composer runtime
const handleEditorChange = useCallback(
(text: string) => {
@ -391,10 +369,6 @@ const Composer: FC = () => {
editorRef.current?.clear();
setMentionedDocuments([]);
setSidebarDocs([]);
setMentionedDocumentIds({
surfsense_doc_ids: [],
document_ids: [],
});
}
}, [
showDocumentPopover,
@ -403,29 +377,17 @@ const Composer: FC = () => {
composerRuntime,
setMentionedDocuments,
setSidebarDocs,
setMentionedDocumentIds,
]);
// Remove document from mentions and sync IDs to atom
const handleDocumentRemove = useCallback(
(docId: number, docType?: string) => {
setMentionedDocuments((prev) => {
const updated = prev.filter((doc) => !(doc.id === docId && doc.document_type === docType));
setMentionedDocumentIds({
surfsense_doc_ids: updated
.filter((doc) => doc.document_type === "SURFSENSE_DOCS")
.map((doc) => doc.id),
document_ids: updated
.filter((doc) => doc.document_type !== "SURFSENSE_DOCS")
.map((doc) => doc.id),
});
return updated;
});
setMentionedDocuments((prev) =>
prev.filter((doc) => !(doc.id === docId && doc.document_type === docType))
);
},
[setMentionedDocuments, setMentionedDocumentIds]
[setMentionedDocuments]
);
// Add selected documents from picker, insert chips, and sync IDs to atom
const handleDocumentsMention = useCallback(
(documents: Pick<Document, "id" | "title" | "document_type">[]) => {
const existingKeys = new Set(mentionedDocuments.map((d) => `${d.document_type}:${d.id}`));
@ -442,21 +404,12 @@ const Composer: FC = () => {
const uniqueNewDocs = documents.filter(
(doc) => !existingKeySet.has(`${doc.document_type}:${doc.id}`)
);
const updated = [...prev, ...uniqueNewDocs];
setMentionedDocumentIds({
surfsense_doc_ids: updated
.filter((doc) => doc.document_type === "SURFSENSE_DOCS")
.map((doc) => doc.id),
document_ids: updated
.filter((doc) => doc.document_type !== "SURFSENSE_DOCS")
.map((doc) => doc.id),
});
return updated;
return [...prev, ...uniqueNewDocs];
});
setMentionQuery("");
},
[mentionedDocuments, setMentionedDocuments, setMentionedDocumentIds]
[mentionedDocuments, setMentionedDocuments]
);
return (

View file

@ -178,12 +178,12 @@ export function ChatShareButton({ thread, onVisibilityChange, className }: ChatS
</Tooltip>
<PopoverContent
className="w-[280px] md:w-[320px] p-0 rounded-lg shadow-lg border-border/60"
className="w-[280px] md:w-[320px] p-0 rounded-lg shadow-lg border-border/60 dark:bg-muted dark:border dark:border-neutral-700 select-none"
align="end"
sideOffset={8}
onCloseAutoFocus={(e) => e.preventDefault()}
>
<div className="p-1.5 space-y-1 select-none">
<div className="p-1.5 space-y-1">
{/* Visibility Options */}
{visibilityOptions.map((option) => {
const isSelected = currentVisibility === option.value;
@ -196,27 +196,27 @@ export function ChatShareButton({ thread, onVisibilityChange, className }: ChatS
onClick={() => handleVisibilityChange(option.value)}
className={cn(
"w-full flex items-center gap-2.5 px-2.5 py-2 rounded-md transition-all",
"hover:bg-accent/50 cursor-pointer",
"hover:bg-accent/50 dark:hover:bg-white/10 cursor-pointer",
"focus:outline-none",
isSelected && "bg-accent/80"
isSelected && "bg-accent/80 dark:bg-white/10"
)}
>
<div
className={cn(
"size-7 rounded-md shrink-0 grid place-items-center",
isSelected ? "bg-primary/10" : "bg-muted"
isSelected ? "bg-primary/10 dark:bg-white/10" : "bg-muted dark:bg-white/5"
)}
>
<Icon
className={cn(
"size-4 block",
isSelected ? "text-primary" : "text-muted-foreground"
isSelected ? "text-primary dark:text-white" : "text-muted-foreground"
)}
/>
</div>
<div className="flex-1 text-left min-w-0">
<div className="flex items-center gap-1.5">
<span className={cn("text-sm font-medium", isSelected && "text-primary")}>
<span className={cn("text-sm font-medium", isSelected && "text-primary dark:text-white")}>
{option.label}
</span>
</div>
@ -231,7 +231,7 @@ export function ChatShareButton({ thread, onVisibilityChange, className }: ChatS
{canCreatePublicLink && (
<>
{/* Divider */}
<div className="border-t border-border my-1" />
<div className="border-t border-border dark:border-white/5 my-1" />
{/* Public Link Option */}
<button
@ -240,12 +240,12 @@ export function ChatShareButton({ thread, onVisibilityChange, className }: ChatS
disabled={isCreatingSnapshot}
className={cn(
"w-full flex items-center gap-2.5 px-2.5 py-2 rounded-md transition-all",
"hover:bg-accent/50 cursor-pointer",
"hover:bg-accent/50 dark:hover:bg-white/10 cursor-pointer",
"focus:outline-none",
"disabled:opacity-50 disabled:cursor-not-allowed"
)}
>
<div className="size-7 rounded-md shrink-0 grid place-items-center bg-muted">
<div className="size-7 rounded-md shrink-0 grid place-items-center bg-muted dark:bg-white/5">
<Earth className="size-4 block text-muted-foreground" />
</div>
<div className="flex-1 text-left min-w-0">

View file

@ -264,7 +264,7 @@ export function ModelSelector({
)}
{/* Divider */}
<div className="h-4 w-px bg-border/60 mx-0.5" />
<div className="h-4 w-px bg-border/60 dark:bg-white/10 mx-0.5" />
{/* Image section */}
{currentImageConfig ? (