fix: Prevent DrawIO remount and data loss when resizing window across 768px breakpoint (#306)

* fix: Prevent DrawIO remount and data loss when resizing window across 768px breakpoint

* fix: prevent DrawIO remount and data loss when resizing window

- Move key from ResizablePanelGroup to chat-panel only
- Save diagram to localStorage before breakpoint change
- Restore defaultSize on drawio-panel to prevent layout flash
- Keep save button functionality from main

* fix: reset draw.io ready state on breakpoint change to restore diagram

* fix: skip initial render save and remove console logs

- Add isInitialRenderRef to skip unnecessary save/reset on first render
- Remove console.log statements for production cleanliness
- Add eslint-disable comment explaining loadDiagram dependency

---------

Co-authored-by: dayuan.jiang <jdy.toh@gmail.com>

app/page.tsx

@@ -34,6 +34,7 @@ export default function Home() {
     const chatPanelRef = useRef<ImperativePanelHandle>(null)
     const isSavingRef = useRef(false)
     const mouseOverDrawioRef = useRef(false)
+    const isMobileRef = useRef(false)
 
     // Reset saving flag when dialog closes (with delay to ignore lingering save events from draw.io)
     useEffect(() => {
@@ -102,16 +103,32 @@ export default function Home() {
         resetDrawioReady()
     }
 
-    // Check mobile
+    // Check mobile - save diagram and reset draw.io before crossing breakpoint
+    const isInitialRenderRef = useRef(true)
     useEffect(() => {
         const checkMobile = () => {
-            setIsMobile(window.innerWidth < 768)
+            const newIsMobile = window.innerWidth < 768
+            // If crossing the breakpoint (not initial render), save diagram and reset draw.io
+            if (
+                !isInitialRenderRef.current &&
+                newIsMobile !== isMobileRef.current
+            ) {
+                // Save diagram before remounting (fire and forget)
+                saveDiagramToStorage().catch(() => {
+                    // Ignore timeout errors during resize
+                })
+                // Reset draw.io ready state so onLoad triggers again after remount
+                resetDrawioReady()
+            }
+            isMobileRef.current = newIsMobile
+            isInitialRenderRef.current = false
+            setIsMobile(newIsMobile)
         }
 
         checkMobile()
         window.addEventListener("resize", checkMobile)
         return () => window.removeEventListener("resize", checkMobile)
-    }, [])
+    }, [saveDiagramToStorage, resetDrawioReady])
 
     const toggleChatPanel = () => {
         const panel = chatPanelRef.current
@@ -157,7 +174,6 @@ export default function Home() {
         <div className="h-screen bg-background relative overflow-hidden">
             <ResizablePanelGroup
                 id="main-panel-group"
-                key={isMobile ? "mobile" : "desktop"}
                 direction={isMobile ? "vertical" : "horizontal"}
                 className="h-full"
             >
@@ -209,6 +225,7 @@ export default function Home() {
 
                 {/* Chat Panel */}
                 <ResizablePanel
+                    key={isMobile ? "mobile" : "desktop"}
                     id="chat-panel"
                     ref={chatPanelRef}
                     defaultSize={isMobile ? 50 : 33}

components/chat-panel.tsx

@@ -35,6 +35,9 @@ const STORAGE_XML_SNAPSHOTS_KEY = "next-ai-draw-io-xml-snapshots"
 const STORAGE_SESSION_ID_KEY = "next-ai-draw-io-session-id"
 export const STORAGE_DIAGRAM_XML_KEY = "next-ai-draw-io-diagram-xml"
 
+// sessionStorage keys
+const SESSION_STORAGE_INPUT_KEY = "next-ai-draw-io-input"
+
 // Type for message parts (tool calls and their states)
 interface MessagePart {
     type: string
@@ -106,7 +109,6 @@ export default function ChatPanel({
         resolverRef,
         chartXML,
         clearDiagram,
-        isDrawioReady,
     } = useDiagram()
 
     const onFetchChart = (saveToHistory = true) => {
@@ -148,6 +150,14 @@ export default function ChatPanel({
     const [showNewChatDialog, setShowNewChatDialog] = useState(false)
     const [minimalStyle, setMinimalStyle] = useState(false)
 
+    // Restore input from sessionStorage on mount (when ChatPanel remounts due to key change)
+    useEffect(() => {
+        const savedInput = sessionStorage.getItem(SESSION_STORAGE_INPUT_KEY)
+        if (savedInput) {
+            setInput(savedInput)
+        }
+    }, [])
+
     // Check config on mount
     useEffect(() => {
         fetch("/api/config")
@@ -210,9 +220,6 @@ export default function ChatPanel({
     const localStorageDebounceRef = useRef<ReturnType<
         typeof setTimeout
     > | null>(null)
-    const xmlStorageDebounceRef = useRef<ReturnType<typeof setTimeout> | null>(
-        null,
-    )
     const LOCAL_STORAGE_DEBOUNCE_MS = 1000 // Save at most once per second
 
     const {
@@ -725,47 +732,6 @@ Continue from EXACTLY where you stopped.`,
         }
     }, [setMessages])
 
-    // Restore diagram XML when DrawIO becomes ready
-    const hasDiagramRestoredRef = useRef(false)
-    const [canSaveDiagram, setCanSaveDiagram] = useState(false)
-    useEffect(() => {
-        // Reset restore flag when DrawIO is not ready (e.g., theme/UI change remounts it)
-        if (!isDrawioReady) {
-            hasDiagramRestoredRef.current = false
-            setCanSaveDiagram(false)
-            return
-        }
-        if (hasDiagramRestoredRef.current) return
-        hasDiagramRestoredRef.current = true
-
-        try {
-            const savedDiagramXml = localStorage.getItem(
-                STORAGE_DIAGRAM_XML_KEY,
-            )
-            console.log(
-                "[ChatPanel] Restoring diagram, has saved XML:",
-                !!savedDiagramXml,
-            )
-            if (savedDiagramXml) {
-                console.log(
-                    "[ChatPanel] Loading saved diagram XML, length:",
-                    savedDiagramXml.length,
-                )
-                // Skip validation for trusted saved diagrams
-                onDisplayChart(savedDiagramXml, true)
-                chartXMLRef.current = savedDiagramXml
-            }
-        } catch (error) {
-            console.error("Failed to restore diagram from localStorage:", error)
-        }
-
-        // Allow saving after restore is complete
-        setTimeout(() => {
-            console.log("[ChatPanel] Enabling diagram save")
-            setCanSaveDiagram(true)
-        }, 500)
-    }, [isDrawioReady, onDisplayChart])
-
     // Save messages to localStorage whenever they change (debounced to prevent blocking during streaming)
     useEffect(() => {
         if (!hasRestoredRef.current) return
@@ -795,28 +761,6 @@ Continue from EXACTLY where you stopped.`,
         }
     }, [messages])
 
-    // Save diagram XML to localStorage whenever it changes (debounced)
-    useEffect(() => {
-        if (!canSaveDiagram) return
-        if (!chartXML || chartXML.length <= 300) return
-
-        // Clear any pending save
-        if (xmlStorageDebounceRef.current) {
-            clearTimeout(xmlStorageDebounceRef.current)
-        }
-
-        // Debounce: save after 1 second of no changes
-        xmlStorageDebounceRef.current = setTimeout(() => {
-            localStorage.setItem(STORAGE_DIAGRAM_XML_KEY, chartXML)
-        }, LOCAL_STORAGE_DEBOUNCE_MS)
-
-        return () => {
-            if (xmlStorageDebounceRef.current) {
-                clearTimeout(xmlStorageDebounceRef.current)
-            }
-        }
-    }, [chartXML, canSaveDiagram])
-
     // Save XML snapshots to localStorage whenever they change
     const saveXmlSnapshots = useCallback(() => {
         try {
@@ -916,6 +860,7 @@ Continue from EXACTLY where you stopped.`,
                         },
                     ] as any)
                     setInput("")
+                    sessionStorage.removeItem(SESSION_STORAGE_INPUT_KEY)
                     setFiles([])
                     return
                 }
@@ -963,6 +908,7 @@ Continue from EXACTLY where you stopped.`,
 
                 // Token count is tracked in onFinish with actual server usage
                 setInput("")
+                sessionStorage.removeItem(SESSION_STORAGE_INPUT_KEY)
                 setFiles([])
             } catch (error) {
                 console.error("Error fetching chart data:", error)
@@ -985,6 +931,7 @@ Continue from EXACTLY where you stopped.`,
             localStorage.removeItem(STORAGE_XML_SNAPSHOTS_KEY)
             localStorage.removeItem(STORAGE_DIAGRAM_XML_KEY)
             localStorage.setItem(STORAGE_SESSION_ID_KEY, newSessionId)
+            sessionStorage.removeItem(SESSION_STORAGE_INPUT_KEY)
             toast.success("Started a fresh chat")
         } catch (error) {
             console.error("Failed to clear localStorage:", error)
@@ -999,9 +946,14 @@ Continue from EXACTLY where you stopped.`,
     const handleInputChange = (
         e: React.ChangeEvent<HTMLInputElement | HTMLTextAreaElement>,
     ) => {
+        saveInputToSessionStorage(e.target.value)
         setInput(e.target.value)
     }
 
+    const saveInputToSessionStorage = (input: string) => {
+        sessionStorage.setItem(SESSION_STORAGE_INPUT_KEY, input)
+    }
+
     // Helper functions for message actions (regenerate/edit)
     // Extract previous XML snapshot before a given message index
     const getPreviousXml = (beforeIndex: number): string => {

contexts/diagram-context.tsx

@@ -1,7 +1,7 @@
 "use client"
 
 import type React from "react"
-import { createContext, useContext, useRef, useState } from "react"
+import { createContext, useContext, useEffect, useRef, useState } from "react"
 import type { DrawIoEmbedRef } from "react-drawio"
 import { STORAGE_DIAGRAM_XML_KEY } from "@/components/chat-panel"
 import type { ExportFormat } from "@/components/save-dialog"
@@ -40,12 +40,15 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
         { svg: string; xml: string }[]
     >([])
     const [isDrawioReady, setIsDrawioReady] = useState(false)
+    const [canSaveDiagram, setCanSaveDiagram] = useState(false)
     const [showSaveDialog, setShowSaveDialog] = useState(false)
     const hasCalledOnLoadRef = useRef(false)
     const drawioRef = useRef<DrawIoEmbedRef | null>(null)
     const resolverRef = useRef<((value: string) => void) | null>(null)
     // Track if we're expecting an export for history (user-initiated)
     const expectHistoryExportRef = useRef<boolean>(false)
+    // Track if diagram has been restored from localStorage
+    const hasDiagramRestoredRef = useRef<boolean>(false)
 
     const onDrawioLoad = () => {
         // Only set ready state once to prevent infinite loops
@@ -61,6 +64,48 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
         setIsDrawioReady(false)
     }
 
+    // Restore diagram XML when DrawIO becomes ready
+    // eslint-disable-next-line react-hooks/exhaustive-deps -- loadDiagram uses refs internally and is stable
+    useEffect(() => {
+        // Reset restore flag when DrawIO is not ready (e.g., theme/UI change remounts it)
+        if (!isDrawioReady) {
+            hasDiagramRestoredRef.current = false
+            setCanSaveDiagram(false)
+            return
+        }
+        if (hasDiagramRestoredRef.current) return
+        hasDiagramRestoredRef.current = true
+
+        try {
+            const savedDiagramXml = localStorage.getItem(
+                STORAGE_DIAGRAM_XML_KEY,
+            )
+            if (savedDiagramXml) {
+                // Skip validation for trusted saved diagrams
+                loadDiagram(savedDiagramXml, true)
+            }
+        } catch (error) {
+            console.error("Failed to restore diagram from localStorage:", error)
+        }
+
+        // Allow saving after restore is complete
+        setTimeout(() => {
+            setCanSaveDiagram(true)
+        }, 500)
+    }, [isDrawioReady])
+
+    // Save diagram XML to localStorage whenever it changes (debounced)
+    useEffect(() => {
+        if (!canSaveDiagram) return
+        if (!chartXML || chartXML.length <= 300) return
+
+        const timeoutId = setTimeout(() => {
+            localStorage.setItem(STORAGE_DIAGRAM_XML_KEY, chartXML)
+        }, 1000)
+
+        return () => clearTimeout(timeoutId)
+    }, [chartXML, canSaveDiagram])
+
     // Track if we're expecting an export for file save (stores raw export data)
     const saveResolverRef = useRef<{
         resolver: ((data: string) => void) | null