diff --git a/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx b/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx index 827646dd2..c41dd872c 100644 --- a/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx +++ b/surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx @@ -12,6 +12,7 @@ import { useParams } from "next/navigation"; import { useCallback, useEffect, useMemo, useRef, useState } from "react"; import { toast } from "sonner"; import { z } from "zod"; +import { currentThreadAtom } from "@/atoms/chat/current-thread.atom"; import { type MentionedDocumentInfo, mentionedDocumentIdsAtom, @@ -251,6 +252,7 @@ export default function NewChatPage() { const setMentionedDocuments = useSetAtom(mentionedDocumentsAtom); const setMessageDocumentsMap = useSetAtom(messageDocumentsMapAtom); const hydratePlanState = useSetAtom(hydratePlanStateAtom); + const setCurrentThreadState = useSetAtom(currentThreadAtom); // Get current user for author info in shared chats const { data: currentUser } = useAtomValue(currentUserAtom); @@ -365,6 +367,16 @@ export default function NewChatPage() { initializeThread(); }, [initializeThread]); + // Sync current thread state to atom + useEffect(() => { + setCurrentThreadState({ + id: currentThread?.id ?? null, + visibility: currentThread?.visibility ?? null, + hasComments: currentThread?.has_comments ?? false, + addingCommentToMessageId: null, + }); + }, [currentThread, setCurrentThreadState]); + // Cancel ongoing request const cancelRun = useCallback(async () => { if (abortControllerRef.current) { @@ -842,10 +854,32 @@ export default function NewChatPage() { // Persist assistant message (with thinking steps for restoration on refresh) const finalContent = buildContentForPersistence(); if (contentParts.length > 0) { - appendMessage(currentThreadId, { - role: "assistant", - content: finalContent, - }).catch((err) => console.error("Failed to persist assistant message:", err)); + try { + const savedMessage = await appendMessage(currentThreadId, { + role: "assistant", + content: finalContent, + }); + + // Update message ID from temporary to database ID so comments work immediately + const newMsgId = `msg-${savedMessage.id}`; + setMessages((prev) => + prev.map((m) => (m.id === assistantMsgId ? { ...m, id: newMsgId } : m)) + ); + + // Also update thinking steps map with new ID + setMessageThinkingSteps((prev) => { + const steps = prev.get(assistantMsgId); + if (steps) { + const newMap = new Map(prev); + newMap.delete(assistantMsgId); + newMap.set(newMsgId, steps); + return newMap; + } + return prev; + }); + } catch (err) { + console.error("Failed to persist assistant message:", err); + } // Track successful response trackChatResponseReceived(searchSpaceId, currentThreadId); @@ -860,10 +894,20 @@ export default function NewChatPage() { ); if (hasContent && currentThreadId) { const partialContent = buildContentForPersistence(); - appendMessage(currentThreadId, { - role: "assistant", - content: partialContent, - }).catch((err) => console.error("Failed to persist partial assistant message:", err)); + try { + const savedMessage = await appendMessage(currentThreadId, { + role: "assistant", + content: partialContent, + }); + + // Update message ID from temporary to database ID + const newMsgId = `msg-${savedMessage.id}`; + setMessages((prev) => + prev.map((m) => (m.id === assistantMsgId ? { ...m, id: newMsgId } : m)) + ); + } catch (err) { + console.error("Failed to persist partial assistant message:", err); + } } return; } diff --git a/surfsense_web/atoms/chat/current-thread.atom.ts b/surfsense_web/atoms/chat/current-thread.atom.ts new file mode 100644 index 000000000..1231887f8 --- /dev/null +++ b/surfsense_web/atoms/chat/current-thread.atom.ts @@ -0,0 +1,52 @@ +import { atom } from "jotai"; +import type { ChatVisibility } from "@/lib/chat/thread-persistence"; + +// TODO: Update `hasComments` to true when the first comment is created on a thread. +// Currently it only updates on thread load. The gutter still works because +// `addingCommentToMessageId` keeps it open, but the state is technically stale. + +// TODO: Reset `addingCommentToMessageId` to null after a comment is successfully created. +// Currently it stays set until navigation or clicking another message's bubble. +// Not causing issues since panel visibility is driven by per-message comment count. + +// TODO: Consider calling `resetCurrentThreadAtom` when unmounting the chat page +// for explicit cleanup, though React navigation handles this implicitly. + +interface CurrentThreadState { + id: number | null; + visibility: ChatVisibility | null; + hasComments: boolean; + addingCommentToMessageId: number | null; +} + +const initialState: CurrentThreadState = { + id: null, + visibility: null, + hasComments: false, + addingCommentToMessageId: null, +}; + +export const currentThreadAtom = atom(initialState); + +export const commentsEnabledAtom = atom( + (get) => get(currentThreadAtom).visibility === "SEARCH_SPACE" +); + +export const showCommentsGutterAtom = atom((get) => { + const thread = get(currentThreadAtom); + return ( + thread.visibility === "SEARCH_SPACE" && + (thread.hasComments || thread.addingCommentToMessageId !== null) + ); +}); + +export const addingCommentToMessageIdAtom = atom( + (get) => get(currentThreadAtom).addingCommentToMessageId, + (get, set, messageId: number | null) => { + set(currentThreadAtom, { ...get(currentThreadAtom), addingCommentToMessageId: messageId }); + } +); + +export const resetCurrentThreadAtom = atom(null, (_, set) => { + set(currentThreadAtom, initialState); +});