refactor: integrate sidebar document selections into chat functionality, enhancing document mention management and improving user experience in the Composer and DocumentsSidebar components

This commit is contained in:
Anish Sarkar 2026-03-06 15:59:45 +05:30
parent f0e4aa6539
commit 3be26429ca
5 changed files with 72 additions and 96 deletions

View file

@ -21,6 +21,7 @@ import {
type MentionedDocumentInfo,
mentionedDocumentIdsAtom,
mentionedDocumentsAtom,
sidebarSelectedDocumentsAtom,
messageDocumentsMapAtom,
} from "@/atoms/chat/mentioned-documents.atom";
import {
@ -180,11 +181,13 @@ export default function NewChatPage() {
interruptData: Record<string, unknown>;
} | null>(null);
// Get mentioned document IDs from the composer
// Get mentioned document IDs from the composer (combines @ 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);
const setCurrentThreadState = useSetAtom(currentThreadAtom);
const setTargetCommentId = useSetAtom(setTargetCommentIdAtom);
@ -528,31 +531,30 @@ export default function NewChatPage() {
messageLength: userQuery.length,
});
// 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,
}));
// Combine @-mention chips + sidebar selections for display & persistence
const allMentionedDocs: MentionedDocumentInfo[] = [];
const seenDocKeys = new Set<string>();
for (const doc of [...mentionedDocuments, ...sidebarDocuments]) {
const key = `${doc.document_type}:${doc.id}`;
if (!seenDocKeys.has(key)) {
seenDocKeys.add(key);
allMentionedDocs.push({ id: doc.id, title: doc.title, document_type: doc.document_type });
}
}
if (allMentionedDocs.length > 0) {
setMessageDocumentsMap((prev) => ({
...prev,
[userMsgId]: docsInfo,
[userMsgId]: allMentionedDocs,
}));
}
// Persist user message with mentioned documents (don't await, fire and forget)
const persistContent: unknown[] = [...message.content];
// Add mentioned documents for persistence
if (mentionedDocuments.length > 0) {
if (allMentionedDocs.length > 0) {
persistContent.push({
type: "mentioned-documents",
documents: mentionedDocuments.map((doc) => ({
id: doc.id,
title: doc.title,
document_type: doc.document_type,
})),
documents: allMentionedDocs,
});
}
@ -623,6 +625,7 @@ export default function NewChatPage() {
document_ids: [],
});
setMentionedDocuments([]);
setSidebarDocuments([]);
}
const response = await fetch(`${backendUrl}/api/v1/new_chat`, {
@ -920,8 +923,10 @@ export default function NewChatPage() {
messages,
mentionedDocumentIds,
mentionedDocuments,
sidebarDocuments,
setMentionedDocumentIds,
setMentionedDocuments,
setSidebarDocuments,
setMessageDocumentsMap,
queryClient,
currentThread,

View file

@ -16,11 +16,17 @@ export const mentionedDocumentIdsAtom = atom<{
});
/**
* Atom to store the full document objects mentioned in the current chat composer.
* This persists across component remounts.
* Atom to store the full document objects mentioned via @-mention chips
* in the current chat composer. This persists across component remounts.
*/
export const mentionedDocumentsAtom = atom<Pick<Document, "id" | "title" | "document_type">[]>([]);
/**
* Atom to store documents selected via the sidebar checkboxes / row clicks.
* These are NOT inserted as chips the composer shows a count badge instead.
*/
export const sidebarSelectedDocumentsAtom = atom<Pick<Document, "id" | "title" | "document_type">[]>([]);
/**
* Simplified document info for display purposes
*/
@ -30,22 +36,6 @@ export interface MentionedDocumentInfo {
document_type: string;
}
/**
* Queue atom for sidebar composer communication (additions).
* The sidebar writes documents here; the Composer picks them up,
* inserts chips, and clears the queue.
*/
export const pendingDocumentMentionsAtom = atom<
Pick<Document, "id" | "title" | "document_type">[]
>([]);
/**
* Queue atom for sidebar composer communication (removals).
* The sidebar writes { id, document_type } here; the Composer removes
* the matching chips and clears the queue.
*/
export const pendingDocumentRemovalsAtom = atom<{ id: number; document_type?: string }[]>([]);
/**
* Atom to store mentioned documents per message ID.
* This allows displaying which documents were mentioned with each user message.

View file

@ -407,13 +407,12 @@ export const InlineMentionEditor = forwardRef<InlineMentionEditorRef, InlineMent
next.delete(chipKey);
return next;
});
onDocumentRemove?.(docId, docType);
const text = getText();
const empty = text.length === 0 && mentionedDocs.size <= 1;
setIsEmpty(empty);
},
[getText, mentionedDocs.size, onDocumentRemove]
[getText, mentionedDocs.size]
);
// Expose methods via ref

View file

@ -32,8 +32,7 @@ import { documentsSidebarOpenAtom } from "@/atoms/documents/ui.atoms";
import {
mentionedDocumentIdsAtom,
mentionedDocumentsAtom,
pendingDocumentMentionsAtom,
pendingDocumentRemovalsAtom,
sidebarSelectedDocumentsAtom,
} from "@/atoms/chat/mentioned-documents.atom";
import { membersAtom } from "@/atoms/members/members-query.atoms";
import {
@ -231,7 +230,7 @@ const ThreadWelcome: FC = () => {
const Composer: FC = () => {
// Document mention state (atoms persist across component remounts)
const [mentionedDocuments, setMentionedDocuments] = useAtom(mentionedDocumentsAtom);
const [pendingMentions, setPendingMentions] = useAtom(pendingDocumentMentionsAtom);
const [sidebarDocs, setSidebarDocs] = useAtom(sidebarSelectedDocumentsAtom);
const [showDocumentPopover, setShowDocumentPopover] = useState(false);
const [mentionQuery, setMentionQuery] = useState("");
const editorRef = useRef<InlineMentionEditorRef>(null);
@ -293,7 +292,7 @@ const Composer: FC = () => {
const assistantIdsKey = useAssistantState(({ thread }) =>
thread.messages
.filter((m) => m.role === "assistant" && m.id?.startsWith("msg-"))
.map((m) => m.id!.replace("msg-", ""))
.map((m) => m.id?.replace("msg-", ""))
.join(",")
);
const assistantDbMessageIds = useMemo(
@ -313,17 +312,25 @@ const Composer: FC = () => {
}
}, [isThreadEmpty]);
// Sync mentioned document IDs to atom for inclusion in chat request payload
// 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: mentionedDocuments
surfsense_doc_ids: deduped
.filter((doc) => doc.document_type === "SURFSENSE_DOCS")
.map((doc) => doc.id),
document_ids: mentionedDocuments
document_ids: deduped
.filter((doc) => doc.document_type !== "SURFSENSE_DOCS")
.map((doc) => doc.id),
});
}, [mentionedDocuments, setMentionedDocumentIds]);
}, [mentionedDocuments, sidebarDocs, setMentionedDocumentIds]);
// Sync editor text with assistant-ui composer runtime
const handleEditorChange = useCallback(
@ -386,6 +393,7 @@ const Composer: FC = () => {
composerRuntime.send();
editorRef.current?.clear();
setMentionedDocuments([]);
setSidebarDocs([]);
setMentionedDocumentIds({
surfsense_doc_ids: [],
document_ids: [],
@ -397,6 +405,7 @@ const Composer: FC = () => {
isBlockedByOtherUser,
composerRuntime,
setMentionedDocuments,
setSidebarDocs,
setMentionedDocumentIds,
]);
@ -453,40 +462,6 @@ const Composer: FC = () => {
[mentionedDocuments, setMentionedDocuments, setMentionedDocumentIds]
);
// Process documents queued from the sidebar (additions)
useEffect(() => {
if (pendingMentions.length === 0) return;
handleDocumentsMention(pendingMentions);
setPendingMentions([]);
}, [pendingMentions, handleDocumentsMention, setPendingMentions]);
// Process documents queued from the sidebar (removals)
const [pendingRemovals, setPendingRemovals] = useAtom(pendingDocumentRemovalsAtom);
useEffect(() => {
if (pendingRemovals.length === 0) return;
for (const { id, document_type } of pendingRemovals) {
editorRef.current?.removeDocumentChip(id, document_type);
}
setMentionedDocuments((prev) => {
const removalKeys = new Set(
pendingRemovals.map((r) => `${r.document_type ?? "UNKNOWN"}:${r.id}`)
);
const updated = prev.filter(
(doc) => !removalKeys.has(`${doc.document_type ?? "UNKNOWN"}:${doc.id}`)
);
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;
});
setPendingRemovals([]);
}, [pendingRemovals, setPendingRemovals, setMentionedDocuments, setMentionedDocumentIds]);
return (
<ComposerPrimitive.Root className="aui-composer-root relative flex w-full flex-col gap-2">
<ChatSessionStatus
@ -551,6 +526,7 @@ const ComposerAction: FC<ComposerActionProps> = ({
isBlockedByOtherUser = false,
}) => {
const mentionedDocuments = useAtomValue(mentionedDocumentsAtom);
const sidebarDocs = useAtomValue(sidebarSelectedDocumentsAtom);
const setDocumentsSidebarOpen = useSetAtom(documentsSidebarOpenAtom);
const isComposerTextEmpty = useAssistantState(({ composer }) => {
@ -603,6 +579,17 @@ const ComposerAction: FC<ComposerActionProps> = ({
</div>
)}
<div className="flex items-center gap-2">
{sidebarDocs.length > 0 && (
<button
type="button"
onClick={() => setDocumentsSidebarOpen(true)}
className="rounded-full border border-border/60 bg-accent/50 px-2.5 py-1 text-xs font-medium text-foreground/80 transition-colors hover:bg-accent"
>
{sidebarDocs.length} {sidebarDocs.length === 1 ? "source" : "sources"} selected
</button>
)}
<AssistantIf condition={({ thread }) => !thread.isRunning}>
<ComposerPrimitive.Send asChild disabled={isSendDisabled}>
<TooltipIconButton
@ -644,6 +631,7 @@ const ComposerAction: FC<ComposerActionProps> = ({
</Button>
</ComposerPrimitive.Cancel>
</AssistantIf>
</div>
</div>
);
};

View file

@ -1,16 +1,12 @@
"use client";
import { useAtomValue, useSetAtom } from "jotai";
import { useAtom, useAtomValue } from "jotai";
import { ChevronLeft } from "lucide-react";
import { useParams } from "next/navigation";
import { useTranslations } from "next-intl";
import { useCallback, useEffect, useMemo, useRef, useState } from "react";
import { toast } from "sonner";
import {
mentionedDocumentsAtom,
pendingDocumentMentionsAtom,
pendingDocumentRemovalsAtom,
} from "@/atoms/chat/mentioned-documents.atom";
import { sidebarSelectedDocumentsAtom } from "@/atoms/chat/mentioned-documents.atom";
import { deleteDocumentMutationAtom } from "@/atoms/documents/document-mutation.atoms";
import { Button } from "@/components/ui/button";
import type { DocumentTypeEnum } from "@/contracts/types/document.types";
@ -44,26 +40,24 @@ export function DocumentsSidebar({ open, onOpenChange }: DocumentsSidebarProps)
const [sortDesc, setSortDesc] = useState(true);
const { mutateAsync: deleteDocumentMutation } = useAtomValue(deleteDocumentMutationAtom);
const mentionedDocuments = useAtomValue(mentionedDocumentsAtom);
const setPendingMentions = useSetAtom(pendingDocumentMentionsAtom);
const setPendingRemovals = useSetAtom(pendingDocumentRemovalsAtom);
const [sidebarDocs, setSidebarDocs] = useAtom(sidebarSelectedDocumentsAtom);
const mentionedDocIds = useMemo(
() => new Set(mentionedDocuments.map((d) => d.id)),
[mentionedDocuments]
() => new Set(sidebarDocs.map((d) => d.id)),
[sidebarDocs]
);
const handleToggleChatMention = useCallback(
(doc: { id: number; title: string; document_type: string }, isMentioned: boolean) => {
if (isMentioned) {
setPendingRemovals((prev) => [...prev, { id: doc.id, document_type: doc.document_type }]);
setSidebarDocs((prev) => prev.filter((d) => d.id !== doc.id));
} else {
setPendingMentions((prev) => [
...prev,
{ id: doc.id, title: doc.title, document_type: doc.document_type as DocumentTypeEnum },
]);
setSidebarDocs((prev) => {
if (prev.some((d) => d.id === doc.id)) return prev;
return [...prev, { id: doc.id, title: doc.title, document_type: doc.document_type as DocumentTypeEnum }];
});
}
},
[setPendingMentions, setPendingRemovals]
[setSidebarDocs]
);
const isSearchMode = !!debouncedSearch.trim();