feat(web): add current-thread atom for global state management

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;
 				}

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<CurrentThreadState>(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);
+});