fix: preserve multi-page diagrams after export (#895)

* fix: preserve multi-page diagrams after export

* fix: keep chartXML sourced from autosave to preserve multi-page state

The export event's data.xml (xmlsvg format) contains compressed <diagram>
payloads, which would break applyDiagramOperations/replaceNodes consumers
that need plain <root> elements. Instead of writing export results into
chartXML, stop overwriting it entirely: autosave already delivers the full
uncompressed multi-page document, and loadDiagram covers AI-driven updates.

Also stop overwriting chartXMLRef with the page-only export before sending
a chat message, so session persistence never sees single-page XML.

Keep the data.xml preference for .drawio file downloads (compressed pages
are a valid drawio format).

---------

Co-authored-by: dayuan.jiang <jdy.toh@gmail.com>

components/chat-panel.tsx

@@ -830,10 +830,6 @@ export default function ChatPanel({
                 let chartXml = await onFetchChart()
                 chartXml = formatXML(chartXml)
 
-                // Update ref directly to avoid race condition with React's async state update
-                // This ensures edit_diagram has the correct XML before AI responds
-                chartXMLRef.current = chartXml
-
                 // Build user text by concatenating input with pre-extracted text
                 // (Backend only reads first text part, so we must combine them)
                 const parts: any[] = []

contexts/diagram-context.tsx

@@ -2,7 +2,7 @@
 
 import type React from "react"
 import { createContext, useContext, useEffect, useRef, useState } from "react"
-import type { DrawIoEmbedRef } from "react-drawio"
+import type { DrawIoEmbedRef, EventExport } from "react-drawio"
 import { toast } from "sonner"
 import type { ExportFormat } from "@/components/save-dialog"
 import { getApiEndpoint } from "@/lib/base-path"
@@ -22,7 +22,7 @@ interface DiagramContextType {
     handleExportWithoutHistory: () => void
     resolverRef: React.MutableRefObject<((value: string) => void) | null>
     drawioRef: React.MutableRefObject<DrawIoEmbedRef | null>
-    handleDiagramExport: (data: any) => void
+    handleDiagramExport: (data: EventExport) => void
     handleDiagramAutoSave: (data: { xml?: string }) => void
     clearDiagram: () => void
     saveDiagramToFile: (
@@ -83,7 +83,7 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
 
     // Track if we're expecting an export for file save (stores raw export data)
     const saveResolverRef = useRef<{
-        resolver: ((data: string) => void) | null
+        resolver: ((data: string, fullDiagramXML?: string) => void) | null
         format: ExportFormat | null
     }>({ resolver: null, format: null })
 
@@ -204,7 +204,7 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
         return null
     }
 
-    const handleDiagramExport = (data: any) => {
+    const handleDiagramExport = (data: EventExport) => {
         // Handle PNG export for VLM validation
         if (pngResolverRef.current && data.data?.startsWith("data:image/png")) {
             pngResolverRef.current(data.data)
@@ -215,7 +215,7 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
         // Handle save to file if requested (process raw data before extraction)
         if (saveResolverRef.current.resolver) {
             const format = saveResolverRef.current.format
-            saveResolverRef.current.resolver(data.data)
+            saveResolverRef.current.resolver(data.data, data.xml)
             saveResolverRef.current = { resolver: null, format: null }
             // For non-xmlsvg formats, skip XML extraction as it will fail
             // Only drawio (which uses xmlsvg internally) has the content attribute
@@ -225,8 +225,11 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
             }
         }
 
+        // Don't write chartXML here: exports don't change the diagram, and
+        // data.xml from xmlsvg exports has compressed <diagram> payloads that
+        // would break edit_diagram/display_diagram. Autosave keeps chartXML
+        // up to date with the full uncompressed multi-page document (#879).
         const extractedXML = extractDiagramXML(data.data)
-        setChartXML(extractedXML)
         setLatestSvg(data.data)
 
         // Only add to history if this was a user-initiated export
@@ -288,14 +291,16 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
 
         // Set up the resolver before triggering export
         saveResolverRef.current = {
-            resolver: (exportData: string) => {
+            resolver: (exportData: string, fullDiagramXML?: string) => {
                 let fileContent: string | Blob
                 let mimeType: string
                 let extension: string
 
                 if (format === "drawio") {
-                    // Extract XML from SVG for .drawio format
-                    const xml = extractDiagramXML(exportData)
+                    // Prefer the complete document from the export event so all pages are saved.
+                    const xml = fullDiagramXML?.trim()
+                        ? fullDiagramXML
+                        : extractDiagramXML(exportData)
                     let xmlContent = xml
                     if (!xml.includes("<mxfile")) {
                         xmlContent = `<mxfile><diagram name="Page-1" id="page-1">${xml}</diagram></mxfile>`