fix: restore diagram when chartXML changes after iframe ready

The previous logic only restored on iframe ready, but chartXML might
be set AFTER the iframe is ready (session loaded after iframe).

Now we track the last restored XML and load whenever chartXML changes
to a new real diagram while iframe is ready.

app/[lang]/page.tsx

@@ -11,10 +11,16 @@ import {
 } from "@/components/ui/resizable"
 import { useDiagram } from "@/contexts/diagram-context"
 import { i18n, type Locale } from "@/lib/i18n/config"
+import { isIndexedDBUsable } from "@/lib/session-storage"
 
 export default function Home() {
-    const { drawioRef, handleDiagramExport, onDrawioLoad, resetDrawioReady } =
-        useDiagram()
+    const {
+        drawioRef,
+        handleDiagramExport,
+        handleDiagramAutoSave,
+        onDrawioLoad,
+        resetDrawioReady,
+    } = useDiagram()
     const router = useRouter()
     const pathname = usePathname()
     // Extract current language from pathname (e.g., "/zh/about" → "zh")
@@ -26,6 +32,8 @@ export default function Home() {
     const [isLoaded, setIsLoaded] = useState(false)
     const [isDrawioReady, setIsDrawioReady] = useState(false)
     const [isElectron, setIsElectron] = useState(false)
+    const [canPersist, setCanPersist] = useState(false)
+    const [canPersistChecked, setCanPersistChecked] = useState(false)
     const [drawioBaseUrl, setDrawioBaseUrl] = useState(
         process.env.NEXT_PUBLIC_DRAWIO_BASE_URL || "https://embed.diagrams.net",
     )
@@ -76,6 +84,11 @@ export default function Home() {
             setDrawioBaseUrl(`${window.location.origin}/drawio/index.html`)
         }
 
+        void (async () => {
+            const usable = await isIndexedDBUsable()
+            setCanPersist(usable)
+            setCanPersistChecked(true)
+        })()
         setIsLoaded(true)
     }, [pathname, router])
 
@@ -84,6 +97,17 @@ export default function Home() {
         onDrawioLoad()
     }, [onDrawioLoad])
 
+    const handleDrawioAutoSave = useCallback(
+        (data: { xml?: string }) => {
+            handleDiagramAutoSave(data)
+            // Only suppress modified state when persistence is available
+            if (canPersist) {
+                drawioRef.current?.status({ message: "", modified: false })
+            }
+        },
+        [canPersist, drawioRef, handleDiagramAutoSave],
+    )
+
     const handleDarkModeChange = () => {
         const newValue = !darkMode
         setDarkMode(newValue)
@@ -167,20 +191,32 @@ export default function Home() {
                         }`}
                     >
                         <div className="h-full rounded-xl overflow-hidden shadow-soft-lg border border-border/30 relative">
-                            {isLoaded && (
+                            {isLoaded && canPersistChecked && (
                                 <div
                                     className={`h-full w-full ${isDrawioReady ? "" : "invisible absolute inset-0"}`}
                                 >
                                     <DrawIoEmbed
                                         key={`${drawioUi}-${darkMode}-${currentLang}-${isElectron}`}
                                         ref={drawioRef}
+                                        autosave
+                                        onAutoSave={handleDrawioAutoSave}
                                         onExport={handleDiagramExport}
                                         onLoad={handleDrawioLoad}
                                         baseUrl={drawioBaseUrl}
+                                        configuration={
+                                            canPersist
+                                                ? { confirmExit: false }
+                                                : undefined
+                                        }
                                         urlParameters={{
                                             ui: drawioUi,
                                             spin: false,
                                             libraries: false,
+                                            // Disable modified tracking only when persistence is available
+                                            ...(canPersist && {
+                                                modified: false,
+                                                keepmodified: false,
+                                            }),
                                             saveAndExit: false,
                                             noSaveBtn: true,
                                             noExitBtn: true,

components/chat-panel.tsx

@@ -137,26 +137,22 @@ export default function ChatPanel({
     const onFetchChart = (saveToHistory = true) => {
         return Promise.race([
             new Promise<string>((resolve) => {
-                if (resolverRef && "current" in resolverRef) {
-                    resolverRef.current = resolve
-                }
+                resolverRef.current = resolve
                 if (saveToHistory) {
                     onExport()
                 } else {
                     handleExportWithoutHistory()
                 }
             }),
-            new Promise<string>((_, reject) =>
-                setTimeout(
-                    () =>
-                        reject(
-                            new Error(
-                                "Chart export timed out after 10 seconds",
-                            ),
-                        ),
-                    10000,
-                ),
-            ),
+            new Promise<string>((_, reject) => {
+                const currentResolver = resolverRef.current
+                setTimeout(() => {
+                    if (resolverRef.current === currentResolver) {
+                        resolverRef.current = null
+                    }
+                    reject(new Error("Chart export timed out after 10 seconds"))
+                }, 10000)
+            }),
         ])
     }
 
@@ -585,7 +581,7 @@ export default function ChatPanel({
 
         try {
             const currentSession = sessionManager.currentSession
-            if (currentSession) {
+            if (currentSession && currentSession.messages.length > 0) {
                 // Restore from session manager (IndexedDB)
                 justLoadedSessionRef.current = true
                 syncUIWithSession(currentSession)
@@ -622,7 +618,7 @@ export default function ChatPanel({
         lastSyncedSessionIdRef.current = newSessionId
 
         // Sync UI with new session
-        if (newSession) {
+        if (newSession && newSession.messages.length > 0) {
             justLoadedSessionRef.current = true
             syncUIWithSession(newSession)
         } else if (!newSession) {
@@ -677,7 +673,7 @@ export default function ChatPanel({
         // Debounce: save after 1 second of no changes
         localStorageDebounceRef.current = setTimeout(async () => {
             try {
-                if (messages.length > 0) {
+                if (messages.length > 0 || hasDiagramNow) {
                     const sessionData = await buildSessionData({
                         // Only capture thumbnail if there was a diagram AND this isn't a no-diagram session
                         withThumbnail: hasDiagramNow && !isNodiagramSession,
@@ -699,6 +695,7 @@ export default function ChatPanel({
             }
         }
     }, [
+        chartXML,
         messages,
         status,
         sessionIsAvailable,
@@ -729,7 +726,8 @@ export default function ChatPanel({
         const handleVisibilityChange = async () => {
             if (
                 document.visibilityState === "hidden" &&
-                messagesRef.current.length > 0
+                (messagesRef.current.length > 0 ||
+                    isRealDiagram(chartXMLRef.current))
             ) {
                 try {
                     // Attempt to save session - browser may not wait for completion

contexts/diagram-context.tsx

@@ -20,9 +20,10 @@ interface DiagramContextType {
     loadDiagram: (chart: string, skipValidation?: boolean) => string | null
     handleExport: () => void
     handleExportWithoutHistory: () => void
-    resolverRef: React.Ref<((value: string) => void) | null>
-    drawioRef: React.Ref<DrawIoEmbedRef | null>
+    resolverRef: React.MutableRefObject<((value: string) => void) | null>
+    drawioRef: React.MutableRefObject<DrawIoEmbedRef | null>
     handleDiagramExport: (data: any) => void
+    handleDiagramAutoSave: (data: { xml?: string }) => void
     clearDiagram: () => void
     saveDiagramToFile: (
         filename: string,
@@ -56,8 +57,6 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
     const pngResolverRef = 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 after DrawIO remount (e.g., theme change)
-    const hasDiagramRestoredRef = useRef<boolean>(false)
     // Track latest chartXML for restoration after remount
     const chartXMLRef = useRef<string>("")
 
@@ -79,22 +78,19 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
     }, [chartXML])
 
     // Restore diagram when DrawIO becomes ready after remount (e.g., theme/UI change)
+    // Also restore when chartXML changes while DrawIO is ready (e.g., session loaded after iframe ready)
+    const lastRestoredXmlRef = useRef<string>("")
     useEffect(() => {
-        // Reset restore flag when DrawIO is not ready (preparing for next restore cycle)
-        if (!isDrawioReady) {
-            hasDiagramRestoredRef.current = false
-            return
+        if (!isDrawioReady || !drawioRef.current) return
+        // Only load if we have a real diagram and it's different from what we already loaded
+        if (
+            isRealDiagram(chartXML) &&
+            chartXML !== lastRestoredXmlRef.current
+        ) {
+            lastRestoredXmlRef.current = chartXML
+            drawioRef.current.load({ xml: chartXML })
         }
-        // Only restore once per ready cycle
-        if (hasDiagramRestoredRef.current) return
-        hasDiagramRestoredRef.current = true
-
-        // Restore diagram from ref if we have one
-        const xmlToRestore = chartXMLRef.current
-        if (isRealDiagram(xmlToRestore) && drawioRef.current) {
-            drawioRef.current.load({ xml: xmlToRestore })
-        }
-    }, [isDrawioReady])
+    }, [isDrawioReady, chartXML])
 
     // Track if we're expecting an export for file save (stores raw export data)
     const saveResolverRef = useRef<{
@@ -267,6 +263,11 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
         }
     }
 
+    const handleDiagramAutoSave = (data: { xml?: string }) => {
+        if (!data?.xml) return
+        setChartXML(data.xml)
+    }
+
     const clearDiagram = () => {
         const emptyDiagram = `<mxfile><diagram name="Page-1" id="page-1"><mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/></root></mxGraphModel></diagram></mxfile>`
         // Skip validation for trusted internal template (loadDiagram also sets chartXML)
@@ -391,6 +392,7 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
                 resolverRef,
                 drawioRef,
                 handleDiagramExport,
+                handleDiagramAutoSave,
                 clearDiagram,
                 saveDiagramToFile,
                 getThumbnailSvg,

lib/session-storage.ts

@@ -47,33 +47,6 @@ interface ChatSessionDB extends DBSchema {
 
 // Database singleton
 let dbPromise: Promise<IDBPDatabase<ChatSessionDB>> | null = null
-const resetDBPromise = () => {
-    dbPromise = null
-}
-
-const isClosingError = (error: unknown): boolean => {
-    return (
-        error instanceof DOMException &&
-        error.name === "InvalidStateError" &&
-        /closing/i.test(error.message)
-    )
-}
-
-const withDB = async <T>(
-    action: (db: IDBPDatabase<ChatSessionDB>) => Promise<T>,
-): Promise<T> => {
-    try {
-        const db = await getDB()
-        return await action(db)
-    } catch (error) {
-        if (isClosingError(error)) {
-            resetDBPromise()
-            const db = await getDB()
-            return await action(db)
-        }
-        throw error
-    }
-}
 
 async function getDB(): Promise<IDBPDatabase<ChatSessionDB>> {
     if (!dbPromise) {
@@ -87,23 +60,7 @@ async function getDB(): Promise<IDBPDatabase<ChatSessionDB>> {
                 }
                 // Future migrations: if (oldVersion < 2) { ... }
             },
-            terminated() {
-                resetDBPromise()
-            },
         })
-        dbPromise
-            .then((db) => {
-                db.onversionchange = () => {
-                    db.close()
-                    resetDBPromise()
-                }
-                db.onclose = () => {
-                    resetDBPromise()
-                }
-            })
-            .catch(() => {
-                resetDBPromise()
-            })
     }
     return dbPromise
 }
@@ -118,33 +75,44 @@ export function isIndexedDBAvailable(): boolean {
     }
 }
 
+// Check if IndexedDB is actually usable (not just present).
+// Note: Do NOT close the db here - getDB() returns a shared singleton connection
+// that other code depends on.
+export async function isIndexedDBUsable(): Promise<boolean> {
+    if (!isIndexedDBAvailable()) return false
+    try {
+        await getDB()
+        return true
+    } catch {
+        return false
+    }
+}
+
 // CRUD Operations
 export async function getAllSessionMetadata(): Promise<SessionMetadata[]> {
     if (!isIndexedDBAvailable()) return []
     try {
-        return await withDB(async (db) => {
-            const tx = db.transaction(STORE_NAME, "readonly")
-            const index = tx.store.index("by-updated")
-            const metadata: SessionMetadata[] = []
+        const db = await getDB()
+        const tx = db.transaction(STORE_NAME, "readonly")
+        const index = tx.store.index("by-updated")
+        const metadata: SessionMetadata[] = []
 
-            // Use cursor to read only metadata fields (avoids loading full messages/XML)
-            let cursor = await index.openCursor(null, "prev") // newest first
-            while (cursor) {
-                const s = cursor.value
-                metadata.push({
-                    id: s.id,
-                    title: s.title,
-                    createdAt: s.createdAt,
-                    updatedAt: s.updatedAt,
-                    messageCount: s.messages.length,
-                    hasDiagram:
-                        !!s.diagramXml && s.diagramXml.trim().length > 0,
-                    thumbnailDataUrl: s.thumbnailDataUrl,
-                })
-                cursor = await cursor.continue()
-            }
-            return metadata
-        })
+        // Use cursor to read only metadata fields (avoids loading full messages/XML)
+        let cursor = await index.openCursor(null, "prev") // newest first
+        while (cursor) {
+            const s = cursor.value
+            metadata.push({
+                id: s.id,
+                title: s.title,
+                createdAt: s.createdAt,
+                updatedAt: s.updatedAt,
+                messageCount: s.messages.length,
+                hasDiagram: !!s.diagramXml && s.diagramXml.trim().length > 0,
+                thumbnailDataUrl: s.thumbnailDataUrl,
+            })
+            cursor = await cursor.continue()
+        }
+        return metadata
     } catch (error) {
         console.error("Failed to get session metadata:", error)
         return []
@@ -154,9 +122,8 @@ export async function getAllSessionMetadata(): Promise<SessionMetadata[]> {
 export async function getSession(id: string): Promise<ChatSession | null> {
     if (!isIndexedDBAvailable()) return null
     try {
-        return await withDB(async (db) => {
-            return (await db.get(STORE_NAME, id)) || null
-        })
+        const db = await getDB()
+        return (await db.get(STORE_NAME, id)) || null
     } catch (error) {
         console.error("Failed to get session:", error)
         return null
@@ -166,9 +133,8 @@ export async function getSession(id: string): Promise<ChatSession | null> {
 export async function saveSession(session: ChatSession): Promise<boolean> {
     if (!isIndexedDBAvailable()) return false
     try {
-        await withDB(async (db) => {
-            await db.put(STORE_NAME, session)
-        })
+        const db = await getDB()
+        await db.put(STORE_NAME, session)
         return true
     } catch (error) {
         // Handle quota exceeded
@@ -180,9 +146,8 @@ export async function saveSession(session: ChatSession): Promise<boolean> {
             await deleteOldestSession()
             // Retry once
             try {
-                await withDB(async (db) => {
-                    await db.put(STORE_NAME, session)
-                })
+                const db = await getDB()
+                await db.put(STORE_NAME, session)
                 return true
             } catch (retryError) {
                 console.error(
@@ -201,9 +166,8 @@ export async function saveSession(session: ChatSession): Promise<boolean> {
 export async function deleteSession(id: string): Promise<void> {
     if (!isIndexedDBAvailable()) return
     try {
-        await withDB(async (db) => {
-            await db.delete(STORE_NAME, id)
-        })
+        const db = await getDB()
+        await db.delete(STORE_NAME, id)
     } catch (error) {
         console.error("Failed to delete session:", error)
     }
@@ -212,9 +176,8 @@ export async function deleteSession(id: string): Promise<void> {
 export async function getSessionCount(): Promise<number> {
     if (!isIndexedDBAvailable()) return 0
     try {
-        return await withDB(async (db) => {
-            return await db.count(STORE_NAME)
-        })
+        const db = await getDB()
+        return await db.count(STORE_NAME)
     } catch (error) {
         console.error("Failed to get session count:", error)
         return 0
@@ -224,15 +187,14 @@ export async function getSessionCount(): Promise<number> {
 export async function deleteOldestSession(): Promise<void> {
     if (!isIndexedDBAvailable()) return
     try {
-        await withDB(async (db) => {
-            const tx = db.transaction(STORE_NAME, "readwrite")
-            const index = tx.store.index("by-updated")
-            const cursor = await index.openCursor()
-            if (cursor) {
-                await cursor.delete()
-            }
-            await tx.done
-        })
+        const db = await getDB()
+        const tx = db.transaction(STORE_NAME, "readwrite")
+        const index = tx.store.index("by-updated")
+        const cursor = await index.openCursor()
+        if (cursor) {
+            await cursor.delete()
+        }
+        await tx.done
     } catch (error) {
         console.error("Failed to delete oldest session:", error)
     }