i/next-ai-draw-io
1
0
Fork 0
mirror of https://github.com/DayuanJiang/next-ai-draw-io.git synced 2026-01-02 14:22:28 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
- 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
2025-12-07 01:39:09 +09:00
committed by GitHub
parent 8b578a456e
commit b1bc1a6dc6
5 changed files with 223 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
app/page.tsx
View File

@@ -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") {
setCloseProtection(false)
// Default to false since auto-save handles persistence
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,

7
components/chat-message-display.tsx
View File

@@ -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) {

201
components/chat-panel.tsx
View File

@@ -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()
const newSessionId = `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}

17
contexts/diagram-context.tsx
View File

@@ -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}

4
lib/utils.ts
View File

@@ -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