Merge pull request #1536 from CREDO23/feature-mention-chat-in-chat

[Feat] Chat : Reference past chats via @-mention as read-only context
This commit is contained in:
Rohan Verma 2026-06-25 13:32:25 -07:00 committed by GitHub
commit 96e42a1003
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
18 changed files with 835 additions and 74 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 {
@ -214,7 +215,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(),
@ -252,6 +258,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,
@ -441,8 +454,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);
@ -964,6 +976,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);
@ -1061,9 +1083,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,
});
@ -1073,7 +1095,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);
@ -1135,15 +1157,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"), {
@ -1162,18 +1180,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 } : {}),
@ -1493,7 +1509,7 @@ export default function NewChatPage() {
threadId,
searchSpaceId,
messages,
mentionedDocumentIds,
jotaiStore,
mentionedDocuments,
setMentionedDocuments,
setMessageDocumentsMap,
@ -2061,6 +2077,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,
@ -2074,6 +2093,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);