feat: auto-save and restore session state (#135)

- Save and restore chat messages, XML snapshots, session ID, and diagram XML to localStorage
- Restore diagram when DrawIO becomes ready (using new onLoad callback)
- Change close protection default to false since auto-save handles persistence
- Clear localStorage when clearing chat
- Handle edge cases: undefined edit fields, empty chartXML, missing access code header

app/page.tsx

@@ -12,7 +12,7 @@ import {
 import { useDiagram } from "@/contexts/diagram-context"
 
 export default function Home() {
-    const { drawioRef, handleDiagramExport } = useDiagram()
+    const { drawioRef, handleDiagramExport, onDrawioLoad } = useDiagram()
     const [isMobile, setIsMobile] = useState(false)
     const [isChatVisible, setIsChatVisible] = useState(true)
     const [drawioUi, setDrawioUi] = useState<"min" | "sketch">("min")
@@ -26,13 +26,14 @@ export default function Home() {
         }
         setIsThemeLoaded(true)
     }, [])
-    const [closeProtection, setCloseProtection] = useState(true)
+    const [closeProtection, setCloseProtection] = useState(false)
 
     // Load close protection setting from localStorage after mount
     useEffect(() => {
         const saved = localStorage.getItem(STORAGE_CLOSE_PROTECTION_KEY)
-        if (saved === "false") {
+        // Default to false since auto-save handles persistence
-            setCloseProtection(false)
+        if (saved === "true") {
+            setCloseProtection(true)
         }
     }, [])
     const chatPanelRef = useRef<ImperativePanelHandle>(null)
@@ -107,6 +108,7 @@ export default function Home() {
                                     key={drawioUi}
                                     ref={drawioRef}
                                     onExport={handleDiagramExport}
+                                    onLoad={onDrawioLoad}
                                     urlParameters={{
                                         ui: drawioUi,
                                         spin: true,

components/chat-message-display.tsx

@@ -47,7 +47,7 @@ function EditDiffDisplay({ edits }: { edits: EditPair[] }) {
         <div className="space-y-3">
             {edits.map((edit, index) => (
                 <div
-                    key={`${edit.search.slice(0, 50)}-${edit.replace.slice(0, 50)}-${index}`}
+                    key={`${(edit.search || "").slice(0, 50)}-${(edit.replace || "").slice(0, 50)}-${index}`}
                     className="rounded-lg border border-border/50 overflow-hidden bg-background/50"
                 >
                     <div className="px-3 py-1.5 bg-muted/40 border-b border-border/30 flex items-center gap-2">
@@ -177,7 +177,10 @@ export function ChatMessageDisplay({
             const currentXml = xml || ""
             const convertedXml = convertToLegalXml(currentXml)
             if (convertedXml !== previousXML.current) {
-                const replacedXML = replaceNodes(chartXML, convertedXml)
+                // If chartXML is empty, use the converted XML directly
+                const replacedXML = chartXML
+                    ? replaceNodes(chartXML, convertedXml)
+                    : convertedXml
 
                 const validationError = validateMxCellStructure(replacedXML)
                 if (!validationError) {

components/chat-panel.tsx

@@ -14,7 +14,7 @@ import {
 import Image from "next/image"
 import Link from "next/link"
 import type React from "react"
-import { useEffect, useRef, useState } from "react"
+import { useCallback, useEffect, useRef, useState } from "react"
 import { flushSync } from "react-dom"
 import { FaGithub } from "react-icons/fa"
 import { Toaster } from "sonner"
@@ -24,6 +24,13 @@ import {
     SettingsDialog,
     STORAGE_ACCESS_CODE_KEY,
 } from "@/components/settings-dialog"
+
+// localStorage keys for persistence
+const STORAGE_MESSAGES_KEY = "next-ai-draw-io-messages"
+const STORAGE_XML_SNAPSHOTS_KEY = "next-ai-draw-io-xml-snapshots"
+const STORAGE_SESSION_ID_KEY = "next-ai-draw-io-session-id"
+const STORAGE_DIAGRAM_XML_KEY = "next-ai-draw-io-diagram-xml"
+
 import { useDiagram } from "@/contexts/diagram-context"
 import { formatXML, validateMxCellStructure } from "@/lib/utils"
 import { ChatMessageDisplay } from "./chat-message-display"
@@ -52,6 +59,7 @@ export default function ChatPanel({
         resolverRef,
         chartXML,
         clearDiagram,
+        isDrawioReady,
     } = useDiagram()
 
     const onFetchChart = (saveToHistory = true) => {
@@ -83,7 +91,7 @@ export default function ChatPanel({
     const [files, setFiles] = useState<File[]>([])
     const [showHistory, setShowHistory] = useState(false)
     const [showSettingsDialog, setShowSettingsDialog] = useState(false)
-    const [accessCodeRequired, setAccessCodeRequired] = useState(false)
+    const [, setAccessCodeRequired] = useState(false)
     const [input, setInput] = useState("")
 
     // Check if access code is required on mount
@@ -94,14 +102,21 @@ export default function ChatPanel({
             .catch(() => setAccessCodeRequired(false))
     }, [])
 
-    // Generate a unique session ID for Langfuse tracing
+    // Generate a unique session ID for Langfuse tracing (restore from localStorage if available)
-    const [sessionId, setSessionId] = useState(
+    const [sessionId, setSessionId] = useState(() => {
-        () => `session-${Date.now()}-${Math.random().toString(36).slice(2, 9)}`,
+        if (typeof window !== "undefined") {
-    )
+            const saved = localStorage.getItem(STORAGE_SESSION_ID_KEY)
+            if (saved) return saved
+        }
+        return `session-${Date.now()}-${Math.random().toString(36).slice(2, 9)}`
+    })
 
     // Store XML snapshots for each user message (keyed by message index)
     const xmlSnapshotsRef = useRef<Map<number, string>>(new Map())
 
+    // Flag to track if we've restored from localStorage
+    const hasRestoredRef = useRef(false)
+
     // Ref to track latest chartXML for use in callbacks (avoids stale closure)
     const chartXMLRef = useRef(chartXML)
     useEffect(() => {
@@ -253,14 +268,162 @@ Please retry with an adjusted search pattern or use display_diagram if retries a
     // Update stopRef so onToolCall can access it
     stopRef.current = stop
 
+    // Ref to track latest messages for unload persistence
+    const messagesRef = useRef(messages)
+    useEffect(() => {
+        messagesRef.current = messages
+    }, [messages])
+
     const messagesEndRef = useRef<HTMLDivElement>(null)
 
+    // Restore messages and XML snapshots from localStorage on mount
+    useEffect(() => {
+        if (hasRestoredRef.current) return
+        hasRestoredRef.current = true
+
+        try {
+            // Restore messages
+            const savedMessages = localStorage.getItem(STORAGE_MESSAGES_KEY)
+            if (savedMessages) {
+                const parsed = JSON.parse(savedMessages)
+                if (Array.isArray(parsed) && parsed.length > 0) {
+                    setMessages(parsed)
+                }
+            }
+
+            // Restore XML snapshots
+            const savedSnapshots = localStorage.getItem(
+                STORAGE_XML_SNAPSHOTS_KEY,
+            )
+            if (savedSnapshots) {
+                const parsed = JSON.parse(savedSnapshots)
+                xmlSnapshotsRef.current = new Map(parsed)
+            }
+        } catch (error) {
+            console.error("Failed to restore from localStorage:", error)
+        }
+    }, [setMessages])
+
+    // Restore diagram XML when DrawIO becomes ready
+    const hasDiagramRestoredRef = useRef(false)
+    const [canSaveDiagram, setCanSaveDiagram] = useState(false)
+    useEffect(() => {
+        console.log(
+            "[ChatPanel] isDrawioReady:",
+            isDrawioReady,
+            "hasDiagramRestored:",
+            hasDiagramRestoredRef.current,
+        )
+        if (!isDrawioReady || 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,
+                )
+                onDisplayChart(savedDiagramXml)
+                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
+    useEffect(() => {
+        if (!hasRestoredRef.current) return
+        try {
+            localStorage.setItem(STORAGE_MESSAGES_KEY, JSON.stringify(messages))
+        } catch (error) {
+            console.error("Failed to save messages to localStorage:", error)
+        }
+    }, [messages])
+
+    // Save XML snapshots to localStorage whenever they change
+    const saveXmlSnapshots = useCallback(() => {
+        try {
+            const snapshotsArray = Array.from(xmlSnapshotsRef.current.entries())
+            localStorage.setItem(
+                STORAGE_XML_SNAPSHOTS_KEY,
+                JSON.stringify(snapshotsArray),
+            )
+        } catch (error) {
+            console.error(
+                "Failed to save XML snapshots to localStorage:",
+                error,
+            )
+        }
+    }, [])
+
+    // Save session ID to localStorage
+    useEffect(() => {
+        localStorage.setItem(STORAGE_SESSION_ID_KEY, sessionId)
+    }, [sessionId])
+
+    // Save current diagram XML to localStorage whenever it changes
+    // Only save after initial restore is complete and if it's not an empty diagram
+    useEffect(() => {
+        if (!canSaveDiagram) return
+        // Don't save empty diagrams (check for minimal content)
+        if (chartXML && chartXML.length > 300) {
+            console.log(
+                "[ChatPanel] Saving diagram to localStorage, length:",
+                chartXML.length,
+            )
+            localStorage.setItem(STORAGE_DIAGRAM_XML_KEY, chartXML)
+        }
+    }, [chartXML, canSaveDiagram])
+
     useEffect(() => {
         if (messagesEndRef.current) {
             messagesEndRef.current.scrollIntoView({ behavior: "smooth" })
         }
     }, [messages])
 
+    // Save state right before page unload (refresh/close)
+    useEffect(() => {
+        const handleBeforeUnload = () => {
+            try {
+                localStorage.setItem(
+                    STORAGE_MESSAGES_KEY,
+                    JSON.stringify(messagesRef.current),
+                )
+                localStorage.setItem(
+                    STORAGE_XML_SNAPSHOTS_KEY,
+                    JSON.stringify(
+                        Array.from(xmlSnapshotsRef.current.entries()),
+                    ),
+                )
+                const xml = chartXMLRef.current
+                if (xml && xml.length > 300) {
+                    localStorage.setItem(STORAGE_DIAGRAM_XML_KEY, xml)
+                }
+                localStorage.setItem(STORAGE_SESSION_ID_KEY, sessionId)
+            } catch (error) {
+                console.error("Failed to persist state before unload:", error)
+            }
+        }
+
+        window.addEventListener("beforeunload", handleBeforeUnload)
+        return () =>
+            window.removeEventListener("beforeunload", handleBeforeUnload)
+    }, [sessionId])
+
     const onFormSubmit = async (e: React.FormEvent<HTMLFormElement>) => {
         e.preventDefault()
         const isProcessing = status === "streaming" || status === "submitted"
@@ -295,6 +458,7 @@ Please retry with an adjusted search pattern or use display_diagram if retries a
                 // Save XML snapshot for this message (will be at index = current messages.length)
                 const messageIndex = messages.length
                 xmlSnapshotsRef.current.set(messageIndex, chartXml)
+                saveXmlSnapshots()
 
                 const accessCode =
                     localStorage.getItem(STORAGE_ACCESS_CODE_KEY) || ""
@@ -373,6 +537,7 @@ Please retry with an adjusted search pattern or use display_diagram if retries a
                 xmlSnapshotsRef.current.delete(key)
             }
         }
+        saveXmlSnapshots()
 
         // Remove the user message AND assistant message onwards (sendMessage will re-add the user message)
         // Use flushSync to ensure state update is processed synchronously before sending
@@ -382,6 +547,7 @@ Please retry with an adjusted search pattern or use display_diagram if retries a
         })
 
         // Now send the message after state is guaranteed to be updated
+        const accessCode = localStorage.getItem(STORAGE_ACCESS_CODE_KEY) || ""
         sendMessage(
             { parts: userParts },
             {
@@ -389,6 +555,9 @@ Please retry with an adjusted search pattern or use display_diagram if retries a
                     xml: savedXml,
                     sessionId,
                 },
+                headers: {
+                    "x-access-code": accessCode,
+                },
             },
         )
     }
@@ -422,6 +591,7 @@ Please retry with an adjusted search pattern or use display_diagram if retries a
                 xmlSnapshotsRef.current.delete(key)
             }
         }
+        saveXmlSnapshots()
 
         // Create new parts with updated text
         const newParts = message.parts?.map((part: any) => {
@@ -439,6 +609,7 @@ Please retry with an adjusted search pattern or use display_diagram if retries a
         })
 
         // Now send the edited message after state is guaranteed to be updated
+        const accessCode = localStorage.getItem(STORAGE_ACCESS_CODE_KEY) || ""
         sendMessage(
             { parts: newParts },
             {
@@ -446,6 +617,9 @@ Please retry with an adjusted search pattern or use display_diagram if retries a
                     xml: savedXml,
                     sessionId,
                 },
+                headers: {
+                    "x-access-code": accessCode,
+                },
             },
         )
     }
@@ -584,12 +758,19 @@ Please retry with an adjusted search pattern or use display_diagram if retries a
                     onClearChat={() => {
                         setMessages([])
                         clearDiagram()
-                        setSessionId(
+                        const newSessionId = `session-${Date.now()}-${Math.random()
-                            `session-${Date.now()}-${Math.random()
                             .toString(36)
-                                .slice(2, 9)}`,
+                            .slice(2, 9)}`
-                        )
+                        setSessionId(newSessionId)
                         xmlSnapshotsRef.current.clear()
+                        // Clear localStorage
+                        localStorage.removeItem(STORAGE_MESSAGES_KEY)
+                        localStorage.removeItem(STORAGE_XML_SNAPSHOTS_KEY)
+                        localStorage.removeItem(STORAGE_DIAGRAM_XML_KEY)
+                        localStorage.setItem(
+                            STORAGE_SESSION_ID_KEY,
+                            newSessionId,
+                        )
                     }}
                     files={files}
                     onFileChange={handleFileChange}

contexts/diagram-context.tsx

@@ -22,6 +22,8 @@ interface DiagramContextType {
         format: ExportFormat,
         sessionId?: string,
     ) => void
+    isDrawioReady: boolean
+    onDrawioLoad: () => void
 }
 
 const DiagramContext = createContext<DiagramContextType | undefined>(undefined)
@@ -32,10 +34,20 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
     const [diagramHistory, setDiagramHistory] = useState<
         { svg: string; xml: string }[]
     >([])
+    const [isDrawioReady, setIsDrawioReady] = 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)
+
+    const onDrawioLoad = () => {
+        // Only set ready state once to prevent infinite loops
+        if (hasCalledOnLoadRef.current) return
+        hasCalledOnLoadRef.current = true
+        console.log("[DiagramContext] DrawIO loaded, setting ready state")
+        setIsDrawioReady(true)
+    }
     // Track if we're expecting an export for file save (stores raw export data)
     const saveResolverRef = useRef<{
         resolver: ((data: string) => void) | null
@@ -62,6 +74,9 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
     }
 
     const loadDiagram = (chart: string) => {
+        // Keep chartXML in sync even when diagrams are injected (e.g., display_diagram tool)
+        setChartXML(chart)
+
         if (drawioRef.current) {
             drawioRef.current.load({
                 xml: chart,
@@ -220,6 +235,8 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
                 handleDiagramExport,
                 clearDiagram,
                 saveDiagramToFile,
+                isDrawioReady,
+                onDrawioLoad,
             }}
         >
             {children}

lib/utils.ts

@@ -467,7 +467,9 @@ export function replaceXMLParts(
                     i,
                     i + searchLines.length,
                 )
-                const normalizedCandidate = normalizeWs(candidateLines.join(" "))
+                const normalizedCandidate = normalizeWs(
+                    candidateLines.join(" "),
+                )
 
                 if (normalizedCandidate === normalizedSearch) {
                     matchStartLine = i