From 02527526ba94f7b0fdaf0f424b36518ba693d7bf Mon Sep 17 00:00:00 2001
From: Dayuan Jiang <34411969+DayuanJiang@users.noreply.github.com>
Date: Thu, 1 Jan 2026 15:42:48 +0900
Subject: [PATCH] fix: prevent flash of example panel and animations on page
 refresh (#482)

- Add isRestored state to track when localStorage restoration completes
- Show example panel only after confirming no saved messages exist
- Skip message animations for restored messages
- Default tool calls and reasoning blocks to collapsed for restored messages
---
 components/chat-message-display.tsx | 39 +++++++++++++++++++++++------
 components/chat-panel.tsx           |  4 +++
 2 files changed, 36 insertions(+), 7 deletions(-)

diff --git a/components/chat-message-display.tsx b/components/chat-message-display.tsx
index 519c24c..88fae1f 100644
--- a/components/chat-message-display.tsx
+++ b/components/chat-message-display.tsx
@@ -193,6 +193,7 @@ interface ChatMessageDisplayProps {
     onRegenerate?: (messageIndex: number) => void
     onEditMessage?: (messageIndex: number, newText: string) => void
     status?: "streaming" | "submitted" | "idle" | "error" | "ready"
+    isRestored?: boolean
 }
 
 export function ChatMessageDisplay({
@@ -205,6 +206,7 @@ export function ChatMessageDisplay({
     onRegenerate,
     onEditMessage,
     status = "idle",
+    isRestored = false,
 }: ChatMessageDisplayProps) {
     const dict = useDictionary()
     const { chartXML, loadDiagram: onDisplayChart } = useDiagram()
@@ -250,6 +252,15 @@ export function ChatMessageDisplay({
     const [expandedPdfSections, setExpandedPdfSections] = useState<
         Record<string, boolean>
     >({})
+    // Track message IDs that were restored from localStorage (skip animation for these)
+    const restoredMessageIdsRef = useRef<Set<string> | null>(null)
+
+    // Capture restored message IDs once when isRestored becomes true
+    useEffect(() => {
+        if (isRestored && restoredMessageIdsRef.current === null) {
+            restoredMessageIdsRef.current = new Set(messages.map((m) => m.id))
+        }
+    }, [isRestored, messages])
 
     const setCopyState = (
         messageId: string,
@@ -669,7 +680,8 @@ export function ChatMessageDisplay({
     const renderToolPart = (part: ToolPartLike) => {
         const callId = part.toolCallId
         const { state, input, output } = part
-        const isExpanded = expandedTools[callId] ?? true
+        // Default to collapsed if tool is complete, expanded if still streaming
+        const isExpanded = expandedTools[callId] ?? state !== "output-available"
         const toolName = part.type?.replace("tool-", "")
         const isCopied = copiedToolCallId === callId
 
@@ -859,9 +871,9 @@ export function ChatMessageDisplay({
 
     return (
         <ScrollArea className="h-full w-full scrollbar-thin">
-            {messages.length === 0 ? (
+            {messages.length === 0 && isRestored ? (
                 <ExamplePanel setInput={setInput} setFiles={setFiles} />
-            ) : (
+            ) : messages.length === 0 ? null : (
                 <div className="py-4 px-4 space-y-4">
                     {messages.map((message, messageIndex) => {
                         const userMessageText =
@@ -881,13 +893,23 @@ export function ChatMessageDisplay({
                                     .slice(messageIndex + 1)
                                     .every((m) => m.role !== "user"))
                         const isEditing = editingMessageId === message.id
+                        // Skip animation for restored messages
+                        // If isRestored but ref not set yet, we're in first render after restoration - treat all as restored
+                        const isRestoredMessage =
+                            isRestored &&
+                            (restoredMessageIdsRef.current === null ||
+                                restoredMessageIdsRef.current.has(message.id))
                         return (
                             <div
                                 key={message.id}
-                                className={`flex w-full ${message.role === "user" ? "justify-end" : "justify-start"} animate-message-in`}
-                                style={{
-                                    animationDelay: `${messageIndex * 50}ms`,
-                                }}
+                                className={`flex w-full ${message.role === "user" ? "justify-end" : "justify-start"} ${isRestoredMessage ? "" : "animate-message-in"}`}
+                                style={
+                                    isRestoredMessage
+                                        ? undefined
+                                        : {
+                                              animationDelay: `${messageIndex * 50}ms`,
+                                          }
+                                }
                             >
                                 {message.role === "user" &&
                                     userMessageText &&
@@ -984,6 +1006,9 @@ export function ChatMessageDisplay({
                                                             isStreaming={
                                                                 isStreamingReasoning
                                                             }
+                                                            defaultOpen={
+                                                                !isRestoredMessage
+                                                            }
                                                         >
                                                             <ReasoningTrigger />
                                                             <ReasoningContent>
diff --git a/components/chat-panel.tsx b/components/chat-panel.tsx
index 073b605..b5b2309 100644
--- a/components/chat-panel.tsx
+++ b/components/chat-panel.tsx
@@ -201,6 +201,7 @@ export default function ChatPanel({
 
     // Flag to track if we've restored from localStorage
     const hasRestoredRef = useRef(false)
+    const [isRestored, setIsRestored] = useState(false)
 
     // Ref to track latest chartXML for use in callbacks (avoids stale closure)
     const chartXMLRef = useRef(chartXML)
@@ -457,6 +458,8 @@ export default function ChatPanel({
             localStorage.removeItem(STORAGE_MESSAGES_KEY)
             localStorage.removeItem(STORAGE_XML_SNAPSHOTS_KEY)
             toast.error(dict.errors.sessionCorrupted)
+        } finally {
+            setIsRestored(true)
         }
     }, [setMessages])
 
@@ -1006,6 +1009,7 @@ export default function ChatPanel({
                     onRegenerate={handleRegenerate}
                     status={status}
                     onEditMessage={handleEditMessage}
+                    isRestored={isRestored}
                 />
             </main>