fix: added loading indicator to ensure chat history is consistent when loading from history (#295)

* fix: keep previous data in chat history

* feat: use react query for restoring conversation messages

* fix: loading issue with chat history

* fix: use state instead of ref for the restoredConversationId

* fix: handle not found scenario on both local and remote restoration

* Revert "fix: handle not found scenario on both local and remote restoration"

This reverts commit d4725134087af047fe18bc6519f5ad5244104544.

* fix: handle conversation not found scenario

* chore: added a loading indicator for the chat history page

* chore: reset restored conversation id state

apps/agent/entrypoints/sidepanel/history/ChatHistory.tsx

@@ -1,5 +1,6 @@
-import { useQueryClient } from '@tanstack/react-query'
+import { keepPreviousData, useQueryClient } from '@tanstack/react-query'
 import type { UIMessage } from 'ai'
+import { Loader2 } from 'lucide-react'
 import type { FC } from 'react'
 import { useMemo } from 'react'
 import { useSessionInfo } from '@/lib/auth/sessionStorage'
@@ -30,6 +31,7 @@ const RemoteChatHistory: FC<{ userId: string }> = ({ userId }) => {
 
   const {
     data: graphqlData,
+    isLoading: isLoadingConversations,
     hasNextPage,
     isFetchingNextPage,
     fetchNextPage,
@@ -43,6 +45,7 @@ const RemoteChatHistory: FC<{ userId: string }> = ({ userId }) => {
         lastPage.conversations?.pageInfo.hasNextPage
           ? lastPage.conversations.pageInfo.endCursor
           : undefined,
+      placeholderData: keepPreviousData,
     },
   )
 
@@ -92,6 +95,14 @@ const RemoteChatHistory: FC<{ userId: string }> = ({ userId }) => {
     [conversations],
   )
 
+  if (!profileId || isLoadingConversations) {
+    return (
+      <div className="flex flex-1 items-center justify-center py-12">
+        <Loader2 className="h-6 w-6 animate-spin text-muted-foreground" />
+      </div>
+    )
+  }
+
   return (
     <ConversationList
       groupedConversations={groupedConversations}

apps/agent/entrypoints/sidepanel/index/Chat.tsx

@@ -1,3 +1,4 @@
+import { Loader2 } from 'lucide-react'
 import { useEffect, useRef, useState } from 'react'
 import { createBrowserOSAction } from '@/lib/chat-actions/types'
 import { SIDEPANEL_AI_TRIGGERED_EVENT } from '@/lib/constants/analyticsEvents'
@@ -27,6 +28,7 @@ export const Chat = () => {
     onClickLike,
     disliked,
     onClickDislike,
+    isRestoringConversation,
   } = useChatSessionContext()
 
   const {
@@ -134,7 +136,11 @@ export const Chat = () => {
   return (
     <>
       <main className="mt-4 flex h-full flex-1 flex-col space-y-4 overflow-y-auto">
-        {messages.length === 0 ? (
+        {isRestoringConversation ? (
+          <div className="flex flex-1 items-center justify-center">
+            <Loader2 className="h-6 w-6 animate-spin text-muted-foreground" />
+          </div>
+        ) : messages.length === 0 ? (
           <ChatEmptyState
             mode={mode}
             mounted={mounted}

apps/agent/entrypoints/sidepanel/index/useChatSession.ts

@@ -19,7 +19,7 @@ import {
   useConversations,
 } from '@/lib/conversations/conversationStorage'
 import { formatConversationHistory } from '@/lib/conversations/formatConversationHistory'
-import { execute } from '@/lib/graphql/execute'
+import { useGraphqlQuery } from '@/lib/graphql/useGraphqlQuery'
 import { useLlmProviders } from '@/lib/llm-providers/useLlmProviders'
 import { track } from '@/lib/metrics/track'
 import { searchActionsStorage } from '@/lib/search-actions/searchActionsStorage'
@@ -299,27 +299,45 @@ export const useChatSession = () => {
     conversationId: conversationIdRef.current,
   })
 
+  const {
+    data: remoteConversationData,
+    isFetched: isRemoteConversationFetched,
+  } = useGraphqlQuery(
+    GetConversationWithMessagesDocument,
+    { conversationId: conversationIdParam ?? '' },
+    {
+      enabled: !!conversationIdParam && isLoggedIn,
+    },
+  )
+
+  const [restoredConversationId, setRestoredConversationId] = useState<
+    string | null
+  >(null)
+
+  // biome-ignore lint/correctness/useExhaustiveDependencies: restore should only run when query data arrives or conversationIdParam changes
   useEffect(() => {
     if (!conversationIdParam) return
+    if (restoredConversationId === conversationIdParam) return
 
-    const restoreConversation = async () => {
-      if (isLoggedIn) {
-        const result = await execute(GetConversationWithMessagesDocument, {
-          conversationId: conversationIdParam,
-        })
+    if (isLoggedIn) {
+      if (!isRemoteConversationFetched) return
 
-        if (result.conversation) {
-          const messages = result.conversation.conversationMessages.nodes
+      if (remoteConversationData?.conversation) {
+        const restoredMessages =
+          remoteConversationData.conversation.conversationMessages.nodes
             .filter((node): node is NonNullable<typeof node> => node !== null)
             .map((node) => node.message as UIMessage)
 
-          setConversationId(
-            conversationIdParam as ReturnType<typeof crypto.randomUUID>,
-          )
-          setMessages(messages)
-          markMessagesAsSaved(conversationIdParam, messages)
-        }
-      } else {
+        setConversationId(
+          conversationIdParam as ReturnType<typeof crypto.randomUUID>,
+        )
+        setMessages(restoredMessages)
+        markMessagesAsSaved(conversationIdParam, restoredMessages)
+      }
+      setRestoredConversationId(conversationIdParam)
+      setSearchParams({}, { replace: true })
+    } else {
+      const restoreLocal = async () => {
         const conversations = await conversationStorage.getValue()
         const conversation = conversations?.find(
           (c) => c.id === conversationIdParam,
@@ -331,19 +349,12 @@ export const useChatSession = () => {
           )
           setMessages(conversation.messages)
         }
+        setRestoredConversationId(conversationIdParam)
+        setSearchParams({}, { replace: true })
       }
-
-      setSearchParams({}, { replace: true })
+      restoreLocal()
     }
-
-    restoreConversation()
-  }, [
-    conversationIdParam,
-    setMessages,
-    setSearchParams,
-    isLoggedIn,
-    markMessagesAsSaved,
-  ])
+  }, [conversationIdParam, remoteConversationData, isLoggedIn])
 
   // biome-ignore lint/correctness/useExhaustiveDependencies: only need to run when messages change
   useEffect(() => {
@@ -414,9 +425,13 @@ export const useChatSession = () => {
     setTextToAction(new Map())
     setLiked({})
     setDisliked({})
+    setRestoredConversationId(null)
     resetRemoteConversation()
   }
 
+  const isRestoringConversation =
+    !!conversationIdParam && restoredConversationId !== conversationIdParam
+
   return {
     mode,
     setMode,
@@ -427,6 +442,7 @@ export const useChatSession = () => {
     providers,
     selectedProvider,
     isLoading: isLoadingProviders || isLoadingAgentUrl,
+    isRestoringConversation,
     agentUrlError,
     chatError,
     handleSelectProvider,