mirror of
https://github.com/DayuanJiang/next-ai-draw-io.git
synced 2026-01-02 22:32:27 +08:00
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
This commit is contained in:
Dayuan Jiang
@@ -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
|
||||
const [sessionId, setSessionId] = useState(
|
||||
() => `session-${Date.now()}-${Math.random().toString(36).slice(2, 9)}`,
|
||||
)
|
||||
// Generate a unique session ID for Langfuse tracing (restore from localStorage if available)
|
||||
const [sessionId, setSessionId] = useState(() => {
|
||||
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(
|
||||
`session-${Date.now()}-${Math.random()
|
||||
.toString(36)
|
||||
.slice(2, 9)}`,
|
||||
)
|
||||
const newSessionId = `session-${Date.now()}-${Math.random()
|
||||
.toString(36)
|
||||
.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}
|
||||
|
||||
Reference in New Issue
Block a user