From b1bc1a6dc6eb6734fbd5ad4c92bafc4a58cf94cd Mon Sep 17 00:00:00 2001 From: Dayuan Jiang <34411969+DayuanJiang@users.noreply.github.com> Date: Sun, 7 Dec 2025 01:39:09 +0900 Subject: [PATCH] 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 | 10 +- components/chat-message-display.tsx | 7 +- components/chat-panel.tsx | 203 ++++++++++++++++++++++++++-- contexts/diagram-context.tsx | 17 +++ lib/utils.ts | 4 +- 5 files changed, 223 insertions(+), 18 deletions(-) diff --git a/app/page.tsx b/app/page.tsx index 3fd58b2..09c4007 100644 --- a/app/page.tsx +++ b/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") { - setCloseProtection(false) + // Default to false since auto-save handles persistence + if (saved === "true") { + setCloseProtection(true) } }, []) const chatPanelRef = useRef(null) @@ -107,6 +108,7 @@ export default function Home() { key={drawioUi} ref={drawioRef} onExport={handleDiagramExport} + onLoad={onDrawioLoad} urlParameters={{ ui: drawioUi, spin: true, diff --git a/components/chat-message-display.tsx b/components/chat-message-display.tsx index 04bc1bd..d3b9243 100644 --- a/components/chat-message-display.tsx +++ b/components/chat-message-display.tsx @@ -47,7 +47,7 @@ function EditDiffDisplay({ edits }: { edits: EditPair[] }) {
{edits.map((edit, index) => (
@@ -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) { diff --git a/components/chat-panel.tsx b/components/chat-panel.tsx index 5906c31..6aa66ae 100644 --- a/components/chat-panel.tsx +++ b/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([]) 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>(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(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) => { 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} diff --git a/contexts/diagram-context.tsx b/contexts/diagram-context.tsx index f1f6537..560ea43 100644 --- a/contexts/diagram-context.tsx +++ b/contexts/diagram-context.tsx @@ -22,6 +22,8 @@ interface DiagramContextType { format: ExportFormat, sessionId?: string, ) => void + isDrawioReady: boolean + onDrawioLoad: () => void } const DiagramContext = createContext(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(null) const resolverRef = useRef<((value: string) => void) | null>(null) // Track if we're expecting an export for history (user-initiated) const expectHistoryExportRef = useRef(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} diff --git a/lib/utils.ts b/lib/utils.ts index e280eaf..e179cd0 100644 --- a/lib/utils.ts +++ b/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