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 399cbdf99..5297e275d 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 @@ -51,6 +51,7 @@ import { TokenUsageProvider, } from "@/components/assistant-ui/token-usage-context"; import { Button } from "@/components/ui/button"; +import { Skeleton } from "@/components/ui/skeleton"; import { type HitlDecision, PendingInterruptProvider, @@ -65,6 +66,7 @@ import { } from "@/hooks/use-agent-actions-query"; import { useChatSessionStateSync } from "@/hooks/use-chat-session-state"; import { useMessagesSync } from "@/hooks/use-messages-sync"; +import { useThreadDetail, useThreadMessages } from "@/hooks/use-thread-queries"; import { getAgentFilesystemSelection } from "@/lib/agent-filesystem"; import { documentsApiService } from "@/lib/apis/documents-api.service"; import { getBearerToken } from "@/lib/auth-utils"; @@ -101,8 +103,6 @@ import { appendMessage, createThread, getRegenerateUrl, - getThreadFull, - getThreadMessages, type ThreadListItem, type ThreadListResponse, type ThreadRecord, @@ -120,7 +120,7 @@ import { trackChatMessageSent, trackChatResponseReceived, } from "@/lib/posthog/events"; -import Loading from "../loading"; +import { cacheKeys } from "@/lib/query-client/cache-keys"; const MobileEditorPanel = dynamic( () => @@ -288,11 +288,78 @@ function computeFallbackTurnCancellingRetryDelay(attempt: number): number { return Math.min(raw, TURN_CANCELLING_MAX_DELAY_MS); } +function parseUrlChatId(id: string | string[] | undefined): number { + let parsed = 0; + if (Array.isArray(id) && id.length > 0) { + parsed = Number.parseInt(id[0], 10); + } else if (typeof id === "string") { + parsed = Number.parseInt(id, 10); + } + return Number.isNaN(parsed) ? 0 : parsed; +} + +function ThreadMessagesSkeleton() { + return ( +
+
+
+
+
+ +
+ +
+ + + +
+ +
+ +
+ +
+ + + +
+ +
+ +
+
+ +
+
+ +
+
+
+
+ ); +} + export default function NewChatPage() { const params = useParams(); const queryClient = useQueryClient(); - const [isInitializing, setIsInitializing] = useState(true); - const [threadId, setThreadId] = useState(null); + const urlChatId = useMemo(() => parseUrlChatId(params.chat_id), [params.chat_id]); + const [threadId, setThreadId] = useState(() => (urlChatId > 0 ? urlChatId : null)); + const activeThreadId = urlChatId > 0 ? urlChatId : threadId; + const handledLoadErrorThreadRef = useRef(null); const [currentThread, setCurrentThread] = useState(null); const [messages, setMessages] = useState([]); const [isRunning, setIsRunning] = useState(false); @@ -404,9 +471,11 @@ export default function NewChatPage() { const { data: currentUser } = useAtomValue(currentUserAtom); const { data: agentFlags } = useAtomValue(agentFlagsAtom); const localFilesystemEnabled = agentFlags?.enable_desktop_local_filesystem === true; + const threadDetailQuery = useThreadDetail(activeThreadId); + const threadMessagesQuery = useThreadMessages(activeThreadId); // Live collaboration: sync session state and messages via Zero - useChatSessionStateSync(threadId); + useChatSessionStateSync(activeThreadId); const { data: membersData } = useAtomValue(membersAtom); const handleSyncedMessagesUpdate = useCallback( @@ -467,7 +536,7 @@ export default function NewChatPage() { [isRunning, membersData] ); - useMessagesSync(threadId, handleSyncedMessagesUpdate); + useMessagesSync(activeThreadId, handleSyncedMessagesUpdate); // Extract search_space_id from URL params const searchSpaceId = useMemo(() => { @@ -481,19 +550,7 @@ export default function NewChatPage() { // per-turn Revert button all read). Hydrates from // ``GET /threads/{id}/actions`` and is updated incrementally by the // SSE handlers + revert-batch results below — no atom side-channel. - const { items: agentActionItems } = useAgentActionsQuery(threadId); - - // Extract chat_id from URL params - const urlChatId = useMemo(() => { - const id = params.chat_id; - let parsed = 0; - if (Array.isArray(id) && id.length > 0) { - parsed = Number.parseInt(id[0], 10); - } else if (typeof id === "string") { - parsed = Number.parseInt(id, 10); - } - return Number.isNaN(parsed) ? 0 : parsed; - }, [params.chat_id]); + const { items: agentActionItems } = useAgentActionsQuery(activeThreadId); const handleChatFailure = useCallback( async ({ @@ -632,14 +689,19 @@ export default function NewChatPage() { }); }, []); - // Initialize thread and load messages - // For new chats (no urlChatId), we use lazy creation - thread is created on first message - const initializeThread = useCallback(async () => { - setIsInitializing(true); + const hydratedMessagesRef = useRef<{ + threadId: number | null; + data: typeof threadMessagesQuery.data; + }>({ threadId: null, data: undefined }); - // Reset all state when switching between chats/search spaces to prevent stale data + // Reset thread-local runtime state on route/search-space changes. Data fetching + // is handled by React Query below so the chat shell can render immediately. + useEffect(() => { + const nextThreadId = urlChatId > 0 ? urlChatId : null; + handledLoadErrorThreadRef.current = null; + hydratedMessagesRef.current = { threadId: null, data: undefined }; + setThreadId(nextThreadId); setMessages([]); - setThreadId(null); setCurrentThread(null); setMentionedDocuments([]); tokenUsageStore.clear(); @@ -649,82 +711,96 @@ export default function NewChatPage() { closeEditorPanel(); // Note: agent-action data is keyed by threadId in react-query so // switching threads naturally swaps caches; no explicit reset. - - try { - if (urlChatId > 0) { - // Thread exists - load thread data and messages - setThreadId(urlChatId); - - // Load thread data (for visibility info) and messages in parallel - const [threadData, messagesResponse] = await Promise.all([ - getThreadFull(urlChatId), - getThreadMessages(urlChatId), - ]); - - setCurrentThread(threadData); - - if (messagesResponse.messages && messagesResponse.messages.length > 0) { - const loadedMessages = reconcileInterruptedAssistantMessages( - messagesResponse.messages - ).map(convertToThreadMessage); - setMessages(loadedMessages); - - for (const msg of messagesResponse.messages) { - if (msg.token_usage) { - tokenUsageStore.set(`msg-${msg.id}`, msg.token_usage as TokenUsageData); - } - } - - const restoredDocsMap: Record = {}; - for (const msg of messagesResponse.messages) { - if (msg.role === "user") { - const docs = extractMentionedDocuments(msg.content); - if (docs.length > 0) { - restoredDocsMap[`msg-${msg.id}`] = docs; - } - } - } - if (Object.keys(restoredDocsMap).length > 0) { - setMessageDocumentsMap(restoredDocsMap); - } - } - } - // For new chats (urlChatId === 0), don't create thread yet - // Thread will be created lazily when user sends first message - // This improves UX (instant load) and avoids orphan threads - } catch (error) { - console.error("[NewChatPage] Failed to initialize thread:", error); - if (urlChatId > 0 && error instanceof NotFoundError) { - removeChatTab(urlChatId); - if (typeof window !== "undefined") { - window.history.replaceState(null, "", `/dashboard/${searchSpaceId}/new-chat`); - } - toast.error("This chat was deleted."); - return; - } - // Keep threadId as null - don't use Date.now() as it creates an invalid ID - // that will cause 404 errors on subsequent API calls - setThreadId(null); - setCurrentThread(null); - toast.error("Failed to load chat. Please try again."); - } finally { - setIsInitializing(false); - } }, [ urlChatId, - setMessageDocumentsMap, setMentionedDocuments, + setMessageDocumentsMap, + tokenUsageStore, closeReportPanel, closeEditorPanel, - removeChatTab, - searchSpaceId, + ]); + + useEffect(() => { + if (!activeThreadId) { + setCurrentThread(null); + return; + } + if (threadDetailQuery.data?.id === activeThreadId) { + setCurrentThread(threadDetailQuery.data); + } + }, [activeThreadId, threadDetailQuery.data]); + + useEffect(() => { + const messagesResponse = threadMessagesQuery.data; + if (!activeThreadId || !messagesResponse) return; + + if ( + hydratedMessagesRef.current.threadId === activeThreadId && + hydratedMessagesRef.current.data === messagesResponse + ) { + return; + } + + if (isRunning) { + return; + } + + const loadedMessages = reconcileInterruptedAssistantMessages(messagesResponse.messages).map( + convertToThreadMessage + ); + setMessages(loadedMessages); + + tokenUsageStore.clear(); + const restoredDocsMap: Record = {}; + for (const msg of messagesResponse.messages) { + if (msg.token_usage) { + tokenUsageStore.set(`msg-${msg.id}`, msg.token_usage as TokenUsageData); + } + if (msg.role === "user") { + const docs = extractMentionedDocuments(msg.content); + if (docs.length > 0) { + restoredDocsMap[`msg-${msg.id}`] = docs; + } + } + } + setMessageDocumentsMap(restoredDocsMap); + hydratedMessagesRef.current = { threadId: activeThreadId, data: messagesResponse }; + }, [ + activeThreadId, + isRunning, + setMessageDocumentsMap, + threadMessagesQuery.data, tokenUsageStore, ]); - // Initialize on mount, and re-init when switching search spaces (even if urlChatId is the same) useEffect(() => { - initializeThread(); - }, [initializeThread]); + const loadError = threadDetailQuery.error ?? threadMessagesQuery.error; + if (!activeThreadId || !loadError) return; + if (handledLoadErrorThreadRef.current === activeThreadId) return; + + handledLoadErrorThreadRef.current = activeThreadId; + console.error("[NewChatPage] Failed to load thread:", loadError); + + if (loadError instanceof NotFoundError) { + removeChatTab(activeThreadId); + if (typeof window !== "undefined") { + window.history.replaceState(null, "", `/dashboard/${searchSpaceId}/new-chat`); + } + setThreadId(null); + setCurrentThread(null); + setMessages([]); + toast.error("This chat was deleted."); + return; + } + + toast.error("Failed to load chat. Please try again."); + }, [ + activeThreadId, + removeChatTab, + searchSpaceId, + threadDetailQuery.error, + threadMessagesQuery.error, + ]); // Prefetch document titles for @ mention picker // Runs when user lands on page so data is ready when they type @ @@ -752,7 +828,7 @@ export default function NewChatPage() { const readAndApplyCommentId = () => { const params = new URLSearchParams(window.location.search); const raw = params.get("commentId"); - if (raw && !isInitializing) { + if (raw && activeThreadId) { const commentId = Number.parseInt(raw, 10); if (!Number.isNaN(commentId)) { setTargetCommentId(commentId); @@ -770,11 +846,14 @@ export default function NewChatPage() { window.removeEventListener("popstate", readAndApplyCommentId); clearTargetCommentId(); }; - }, [isInitializing, setTargetCommentId, clearTargetCommentId]); + }, [activeThreadId, setTargetCommentId, clearTargetCommentId]); // Sync current thread state to atom useEffect(() => { if (!currentThread) { + if (activeThreadId) { + return; + } setCurrentThreadMetadata({ id: null, visibility: null, @@ -794,6 +873,7 @@ export default function NewChatPage() { hasComments: currentThread.has_comments ?? false, }); }, [ + activeThreadId, currentThread, currentThreadState.id, currentThreadState.visibility, @@ -882,6 +962,8 @@ export default function NewChatPage() { setThreadId(currentThreadId); // Set currentThread so share button in header appears immediately setCurrentThread(newThread); + queryClient.setQueryData(cacheKeys.threads.detail(newThread.id), newThread); + queryClient.setQueryData(cacheKeys.threads.messages(newThread.id), { messages: [] }); // Track chat creation trackChatCreated(searchSpaceId, currentThreadId); @@ -1389,6 +1471,14 @@ export default function NewChatPage() { } finally { setIsRunning(false); abortControllerRef.current = null; + if (currentThreadId) { + void queryClient.invalidateQueries({ + queryKey: cacheKeys.threads.messages(currentThreadId), + }); + void queryClient.invalidateQueries({ + queryKey: cacheKeys.threads.detail(currentThreadId), + }); + } } }, [ @@ -1737,6 +1827,12 @@ export default function NewChatPage() { } finally { setIsRunning(false); abortControllerRef.current = null; + void queryClient.invalidateQueries({ + queryKey: cacheKeys.threads.messages(resumeThreadId), + }); + void queryClient.invalidateQueries({ + queryKey: cacheKeys.threads.detail(resumeThreadId), + }); } }, [ @@ -2230,6 +2326,12 @@ export default function NewChatPage() { } finally { setIsRunning(false); abortControllerRef.current = null; + void queryClient.invalidateQueries({ + queryKey: cacheKeys.threads.messages(threadId), + }); + void queryClient.invalidateQueries({ + queryKey: cacheKeys.threads.detail(threadId), + }); } }, [ @@ -2416,22 +2518,25 @@ export default function NewChatPage() { onCancel: cancelRun, }); - // Show loading state only when loading an existing thread - if (isInitializing) { - return ; - } + const threadLoadError = activeThreadId + ? (threadDetailQuery.error ?? threadMessagesQuery.error) + : null; + const shouldShowThreadLoadError = + !!threadLoadError && !!activeThreadId && !currentThread && messages.length === 0; + const isThreadMessagesLoading = + !!activeThreadId && + threadMessagesQuery.isPending && + messages.length === 0 && + !threadMessagesQuery.error; - // Show error state only if we tried to load an existing thread but failed - // For new chats (urlChatId === 0), threadId being null is expected (lazy creation) - if (!threadId && urlChatId > 0) { + if (shouldShowThreadLoadError) { return (
Failed to load chat