fix(chat): capture mention chips at submit so they aren't dropped

The composer cleared the live mention atom synchronously on send (via the
editor reset), which raced the async onNew handler that read it — dropping
every @-mention (docs, folders, connectors, and the new chat references)
from the request.

handleSubmit now snapshots the chips before clearing, and onNew consumes
that snapshot (falling back to the live atom for the send-button path),
derives the payload via deriveMentionedPayload, and sends mentioned_thread_ids.
This commit is contained in:
CREDO23 2026-06-23 18:30:39 +02:00
parent 1d5c364e1d
commit 2d1ccb468a
2 changed files with 67 additions and 29 deletions

View file

@ -7,7 +7,7 @@ import {
useExternalStoreRuntime,
} from "@assistant-ui/react";
import { useQueryClient } from "@tanstack/react-query";
import { useAtomValue, useSetAtom } from "jotai";
import { useAtomValue, useSetAtom, useStore } from "jotai";
import dynamic from "next/dynamic";
import { useParams } from "next/navigation";
import { useCallback, useEffect, useMemo, useRef, useState } from "react";
@ -23,9 +23,10 @@ import {
} from "@/atoms/chat/current-thread.atom";
import {
type MentionedDocumentInfo,
mentionedDocumentIdsAtom,
deriveMentionedPayload,
mentionedDocumentsAtom,
messageDocumentsMapAtom,
submittedMentionsAtom,
} from "@/atoms/chat/mentioned-documents.atom";
import { pendingUserImageDataUrlsAtom } from "@/atoms/chat/pending-user-images.atom";
import {
@ -206,7 +207,12 @@ const MentionedDocumentInfoSchema = z.object({
title: z.string(),
document_type: z.string().optional(),
kind: z
.union([z.literal("doc"), z.literal("folder"), z.literal("connector")])
.union([
z.literal("doc"),
z.literal("folder"),
z.literal("connector"),
z.literal("thread"),
])
.optional()
.default("doc"),
connector_type: z.string().optional(),
@ -244,6 +250,13 @@ function extractMentionedDocuments(content: unknown): MentionedDocumentInfo[] {
kind: "folder",
};
}
if (doc.kind === "thread") {
return {
id: doc.id,
title: doc.title,
kind: "thread",
};
}
return {
id: doc.id,
title: doc.title,
@ -433,8 +446,7 @@ export default function NewChatPage() {
// Get disabled tools from the tool toggle UI
const disabledTools = useAtomValue(disabledToolsAtom);
// Get mentioned document IDs from the composer.
const mentionedDocumentIds = useAtomValue(mentionedDocumentIdsAtom);
const jotaiStore = useStore();
const mentionedDocuments = useAtomValue(mentionedDocumentsAtom);
const messageDocumentsMap = useAtomValue(messageDocumentsMapAtom);
const setMentionedDocuments = useSetAtom(mentionedDocumentsAtom);
@ -959,6 +971,16 @@ export default function NewChatPage() {
abortControllerRef.current = null;
}
// Prefer the submit-time snapshot; fall back to the live atom
// for the send-button path.
const submittedSnapshot = jotaiStore.get(submittedMentionsAtom);
jotaiStore.set(submittedMentionsAtom, null);
const activeMentions = submittedSnapshot ?? mentionedDocuments;
const mentionPayload = deriveMentionedPayload(activeMentions);
if (activeMentions.length > 0) {
setMentionedDocuments([]);
}
const urlsSnapshot = [...pendingUserImageUrls];
const { userQuery, userImages } = extractUserTurnForNewChatApi(message, urlsSnapshot);
@ -1060,9 +1082,9 @@ export default function NewChatPage() {
trackChatMessageSent(searchSpaceId, currentThreadId, {
hasAttachments: userImages.length > 0,
hasMentionedDocuments:
mentionedDocumentIds.document_ids.length > 0 ||
mentionedDocumentIds.folder_ids.length > 0 ||
mentionedDocumentIds.connector_ids.length > 0,
mentionPayload.document_ids.length > 0 ||
mentionPayload.folder_ids.length > 0 ||
mentionPayload.connector_ids.length > 0,
messageLength: userQuery.length,
});
@ -1072,7 +1094,7 @@ export default function NewChatPage() {
// can render the correct chip type on reload.
const allMentionedDocs: MentionedDocumentInfo[] = [];
const seenDocKeys = new Set<string>();
for (const doc of mentionedDocuments) {
for (const doc of activeMentions) {
const key = getMentionDocKey(doc);
if (seenDocKeys.has(key)) continue;
seenDocKeys.add(key);
@ -1134,15 +1156,11 @@ export default function NewChatPage() {
})
.filter((m) => m.content.length > 0);
// Get mentioned document IDs for context (separate fields for backend)
const hasDocumentIds = mentionedDocumentIds.document_ids.length > 0;
const hasFolderIds = mentionedDocumentIds.folder_ids.length > 0;
const hasConnectorIds = mentionedDocumentIds.connector_ids.length > 0;
// Clear mentioned documents after capturing them
if (hasDocumentIds || hasFolderIds || hasConnectorIds) {
setMentionedDocuments([]);
}
// Backend expects each mention kind in its own payload bucket.
const hasDocumentIds = mentionPayload.document_ids.length > 0;
const hasFolderIds = mentionPayload.folder_ids.length > 0;
const hasConnectorIds = mentionPayload.connector_ids.length > 0;
const hasThreadIds = mentionPayload.thread_ids.length > 0;
const response = await fetchWithTurnCancellingRetry(() =>
fetch(buildBackendUrl("/api/v1/new_chat"), {
@ -1160,18 +1178,16 @@ export default function NewChatPage() {
local_filesystem_mounts: selection.local_filesystem_mounts,
messages: messageHistory,
mentioned_document_ids: hasDocumentIds
? mentionedDocumentIds.document_ids
? mentionPayload.document_ids
: undefined,
mentioned_folder_ids: hasFolderIds ? mentionedDocumentIds.folder_ids : undefined,
mentioned_folder_ids: hasFolderIds ? mentionPayload.folder_ids : undefined,
mentioned_connector_ids: hasConnectorIds
? mentionedDocumentIds.connector_ids
? mentionPayload.connector_ids
: undefined,
mentioned_connectors: hasConnectorIds ? mentionedDocumentIds.connectors : undefined,
// Full mention metadata (docs + folders, with
// ``kind`` discriminator) so the BE can embed a
// ``mentioned-documents`` ContentPart on the
// persisted user message (replaces the old FE-side
// injection in ``persistUserTurn``).
mentioned_connectors: hasConnectorIds ? mentionPayload.connectors : undefined,
mentioned_thread_ids: hasThreadIds ? mentionPayload.thread_ids : undefined,
// Full mention metadata so the backend can persist a
// ``mentioned-documents`` ContentPart on the user message.
mentioned_documents: allMentionedDocs.length > 0 ? allMentionedDocs : undefined,
disabled_tools: disabledTools.length > 0 ? disabledTools : undefined,
...(userImages.length > 0 ? { user_images: userImages } : {}),
@ -1491,7 +1507,7 @@ export default function NewChatPage() {
threadId,
searchSpaceId,
messages,
mentionedDocumentIds,
jotaiStore,
mentionedDocuments,
setMentionedDocuments,
setMessageDocumentsMap,
@ -2067,6 +2083,9 @@ export default function NewChatPage() {
.filter((d) => d.kind === "folder")
.map((d) => d.id);
const regenerateConnectors = sourceMentionedDocs.filter((d) => d.kind === "connector");
const regenerateThreadIds = sourceMentionedDocs
.filter((d) => d.kind === "thread")
.map((d) => d.id);
const requestBody: Record<string, unknown> = {
search_space_id: searchSpaceId,
@ -2080,6 +2099,8 @@ export default function NewChatPage() {
mentioned_connector_ids:
regenerateConnectors.length > 0 ? regenerateConnectors.map((d) => d.id) : undefined,
mentioned_connectors: regenerateConnectors.length > 0 ? regenerateConnectors : undefined,
mentioned_thread_ids:
regenerateThreadIds.length > 0 ? regenerateThreadIds : undefined,
// Full mention metadata for the regenerate-specific
// source list. Only meaningful for edit (the BE only
// re-persists a user row when ``user_query`` is set);

View file

@ -38,6 +38,7 @@ import { currentThreadAtom } from "@/atoms/chat/current-thread.atom";
import {
type MentionedDocumentInfo,
mentionedDocumentsAtom,
submittedMentionsAtom,
} from "@/atoms/chat/mentioned-documents.atom";
import { pendingUserImageDataUrlsAtom } from "@/atoms/chat/pending-user-images.atom";
import {
@ -446,6 +447,7 @@ const ClipboardChip: FC<{ text: string; onDismiss: () => void }> = ({ text, onDi
const Composer: FC = () => {
const [mentionedDocuments, setMentionedDocuments] = useAtom(mentionedDocumentsAtom);
const setSubmittedMentions = useSetAtom(submittedMentionsAtom);
const [showDocumentPopover, setShowDocumentPopover] = useState(false);
const [showPromptPicker, setShowPromptPicker] = useState(false);
const [mentionQuery, setMentionQuery] = useState("");
@ -575,6 +577,13 @@ const Composer: FC = () => {
kind: "folder",
};
}
if (d.kind === "thread") {
return {
id: d.id,
title: d.title,
kind: "thread",
};
}
return {
id: d.id,
title: d.title,
@ -770,6 +779,10 @@ const Composer: FC = () => {
setClipboardInitialText(undefined);
}
// Capture chips before the reset below clears the live atom, so
// the async ``onNew`` still sees them.
setSubmittedMentions(mentionedDocuments);
aui.composer().send();
editorRef.current?.clear();
setIsComposerInputEmpty(true);
@ -781,6 +794,8 @@ const Composer: FC = () => {
isBlockedByOtherUser,
clipboardInitialText,
aui,
mentionedDocuments,
setSubmittedMentions,
setMentionedDocuments,
]);
@ -788,7 +803,7 @@ const Composer: FC = () => {
(
docId: number,
docType?: string,
kind?: "doc" | "folder" | "connector",
kind?: "doc" | "folder" | "connector" | "thread",
connectorType?: string
) => {
setMentionedDocuments((prev) => {
@ -876,6 +891,8 @@ const Composer: FC = () => {
<DocumentMentionPicker
ref={documentPickerRef}
searchSpaceId={Number(search_space_id)}
enableChatMentions
currentChatId={threadId}
onSelectionChange={handleDocumentsMention}
onDone={() => {
setShowDocumentPopover(false);