test: read the failed edit's output without unsafe optional chaining

The model always got the first page and display_diagram replaced the whole
file with one page, so drawing on page 3 of a multi-page file deleted the
other pages (#241), and edits went to page 1 while the user looked at
another page.

Now a turn records the page the user is viewing when the message is sent.
That page's model goes to the model (also as previousXml, on regenerate and
in edit_diagram error messages), selected shapes are sent on any page,
display_diagram and append_diagram replace only that page (a drawing that
brings several pages still replaces the file), edit_diagram targets it,
and the streaming preview draws on it. Compact cells and default styles
take the layer and the edges of that page, in the preview and in
editDiagram. The editor bridge replaces the current page in place when the
other pages are unchanged, so Ctrl+Z undoes an AI change on a multi-page
file in one step; when a full load is needed, the user's page is shown
again afterwards. Version summaries and the compare dialog count the page
the change was made on.

components/canvas/canvas-stage.tsx

@@ -46,11 +46,6 @@ export function CanvasStage({
     const engine = useChatEngine()
     const isDark = useSettingsStore((s) => s.isDark)
     const hasEditor = useCanvasStore((s) => s.hasEditor)
-    // The model sees the first page only, so asking about shapes elsewhere
-    // would change the wrong page
-    const onFirstPage = useCanvasStore(
-        (s) => s.pages.length === 0 || s.pages[0]?.id === s.currentPageId,
-    )
     // Canvas width when the chat panel starts sliding in or out
     const slideStartWidthRef = useRef(0)
 
@@ -168,7 +163,7 @@ export function CanvasStage({
                 <SelectionAsk
                     width={size.width}
                     height={size.height}
-                    hidden={engine.isBusy || !onFirstPage}
+                    hidden={engine.isBusy}
                 />
             )}
         </div>

components/chat/chat-engine.tsx

@@ -42,6 +42,7 @@ import { useValidateDiagram } from "@/hooks/use-validate-diagram"
 import { getApiEndpoint } from "@/lib/base-path"
 import { findCachedResponse } from "@/lib/cached-responses"
 import { buildChatHeaders } from "@/lib/chat-request"
+import { pageModelXml, placeOnPage } from "@/lib/diagram-pages"
 import { EMPTY_DRAWIO_DOCUMENT } from "@/lib/drawio/drawio-config"
 import { formatMessage } from "@/lib/i18n/utils"
 import { isPdfFile, isTextFile } from "@/lib/pdf-utils"
@@ -90,15 +91,20 @@ function hasToolErrors(messages: UIMessage[]): boolean {
 }
 
 /**
- * Snapshots keep the full multi-page document, but the model only sees and
- * edits the first page, so give it the first page's mxGraphModel.
- * Older snapshots already hold a single mxGraphModel and are returned as is.
+ * Snapshots keep the full multi-page document, but the model sees and
+ * edits one page, the one the user is viewing (the first when unknown), so
+ * give it that page's mxGraphModel. Older snapshots already hold a single
+ * mxGraphModel and are returned as is.
  */
-function getFirstPageXml(xml: string): string {
+function getPageXml(xml: string, pageId: string | null): string {
     if (!xml.includes("<mxfile")) return xml
-    const doc = new DOMParser().parseFromString(xml, "text/xml")
-    const model = doc.querySelector("diagram")?.querySelector("mxGraphModel")
-    return model ? formatXML(new XMLSerializer().serializeToString(model)) : xml
+    const model = pageModelXml(xml, pageId)
+    return model ? formatXML(model) : xml
+}
+
+/** The page the user is viewing; null with an external draw.io */
+function viewedPageId(): string | null {
+    return useCanvasStore.getState().currentPageId
 }
 
 // Shapes sent with a user message (also kept in its metadata)
@@ -374,6 +380,9 @@ export function ChatEngineProvider({
 
     // XML snapshot taken before each user message (keyed by message index)
     const xmlSnapshotsRef = useRef<Map<number, string>>(new Map())
+    // The page the running turn's model reads and writes: the one the user
+    // viewed when the message was sent (null: the first page)
+    const turnPageIdRef = useRef<string | null>(null)
     // Index of the user message whose turn is running (versions belong to it)
     const currentTurnRef = useRef(0)
     // The chat (its generation) the running turn was sent in: a reply that
@@ -450,6 +459,7 @@ export function ChatEngineProvider({
         validationRetryCountRef,
         // A preview undone just before the tool call is in this one already
         chartXMLRef: liveChartXMLRef,
+        turnPageIdRef,
         onDisplayChart,
         onFetchChart,
         captureValidationPng,
@@ -748,6 +758,7 @@ export function ChatEngineProvider({
         processedToolCallsRef,
         editDiagramOriginalXmlRef,
         loadedMessageIdsRef,
+        turnPageIdRef,
     })
 
     const { restoreVersion, undoVersion } = useVersions({
@@ -762,14 +773,10 @@ export function ChatEngineProvider({
     // ---------------------------------------------------------------------
 
     const selection = useCanvasStore((s) => s.selection)
-    // The model sees and edits the first page only
-    const onFirstPage = useCanvasStore(
-        (s) => s.pages.length === 0 || s.pages[0]?.id === s.currentPageId,
-    )
     const [dismissedSelectionKey, setDismissedSelectionKey] = useState("")
     const selectionKey = selection.map((c) => c.id).join(",")
     const chatSelection =
-        !onFirstPage || selectionKey === dismissedSelectionKey ? [] : selection
+        selectionKey === dismissedSelectionKey ? [] : selection
     // Once nothing is selected, picking the same shapes again attaches them
     useEffect(() => {
         if (selection.length === 0) setDismissedSelectionKey("")
@@ -1178,8 +1185,9 @@ export function ChatEngineProvider({
             .filter((k) => k < beforeIndex)
             .sort((a, b) => b - a)
         return snapshotKeys.length > 0
-            ? getFirstPageXml(
+            ? getPageXml(
                   xmlSnapshotsRef.current.get(snapshotKeys[0]) || "",
+                  viewedPageId(),
               )
             : ""
     }
@@ -1198,6 +1206,7 @@ export function ChatEngineProvider({
         continuationOriginalRef.current = null
         stoppedRef.current = false
         currentTurnRef.current = turnIndex
+        turnPageIdRef.current = viewedPageId()
         turnChatRef.current = getChatGeneration()
         // Busy from now on, before the next render says so
         busyRef.current = true
@@ -1254,7 +1263,8 @@ export function ChatEngineProvider({
         const turnIndex = messagesRef.current.length
         const previousXml = getPreviousXml(turnIndex)
         // Snapshot the full multi-page document (kept fresh by autosave) so
-        // regenerate/edit can restore every page; the model gets page 1 only
+        // regenerate/edit can restore every page; the model gets the page
+        // the user is viewing
         xmlSnapshotsRef.current.set(turnIndex, chartXMLRef.current || chartXml)
         sendChatMessage(parts, chartXml, previousXml, turnIndex, selectedCells)
         return true
@@ -1328,11 +1338,16 @@ export function ChatEngineProvider({
                     pageName: "Page-1",
                 })
                 if (prepared.ok) {
+                    const pageId = viewedPageId()
+                    const canvasXml = chartXMLRef.current || ""
                     onDisplayChart(
-                        keepFileVars(prepared.xml, chartXMLRef.current || ""),
+                        keepFileVars(
+                            placeOnPage(prepared.xml, canvasXml, pageId),
+                            canvasXml,
+                        ),
                         true,
                         "commit",
-                        { toolCallId },
+                        { toolCallId, pageId },
                     )
                 }
                 clearComposer()
@@ -1492,7 +1507,7 @@ export function ChatEngineProvider({
         })
         sendChatMessage(
             parts,
-            getFirstPageXml(savedXml),
+            getPageXml(savedXml, viewedPageId()),
             previousXml,
             index,
             selectionOf(messages[index]),

components/chat/compare-dialog.tsx

@@ -56,11 +56,11 @@ export function CompareDialog() {
         () => !version || isSameDocument(version.xml, chartXML),
         [version, chartXML],
     )
-    // What the canvas has changed since this version on the first page
-    // (null: nothing there)
+    // What the canvas has changed since this version on the page it was
+    // made on (null: nothing there)
     const changes = useMemo(() => {
         if (!version || same) return null
-        const { summary } = diffDiagrams(version.xml, chartXML)
+        const { summary } = diffDiagrams(version.xml, chartXML, version.pageId)
         return Object.values(summary).some((n) => n > 0) ? summary : null
     }, [version, chartXML, same])
 

components/chat/use-diagram-streaming.ts

@@ -5,6 +5,7 @@ import type { RefObject } from "react"
 import { useCallback, useEffect, useRef } from "react"
 import type { DiagramOperation, ToolPartLike } from "@/components/chat/types"
 import { useDiagram } from "@/contexts/diagram-context"
+import { pageSelectorFor } from "@/lib/diagram-pages"
 import {
     convertToLegalXml,
     extractCompleteMxCells,
@@ -15,6 +16,7 @@ import {
     expandCompactCells,
 } from "@/packages/mcp-server/src/compact-cells.ts"
 import { applyDiagramOperations } from "@/packages/mcp-server/src/diagram-operations.ts"
+import { targetPageXml } from "@/packages/mcp-server/src/edit-diagram.ts"
 import { BLANK_MXFILE } from "@/packages/mcp-server/src/pages.ts"
 import {
     addDefaultStyles,
@@ -44,6 +46,8 @@ interface UseDiagramStreamingParams {
     processedToolCallsRef: RefObject<Set<string>>
     editDiagramOriginalXmlRef: RefObject<Map<string, string>>
     loadedMessageIdsRef: RefObject<Set<string>>
+    // The page the model draws on in this turn; null: the first page
+    turnPageIdRef?: RefObject<string | null>
 }
 
 /**
@@ -60,6 +64,7 @@ export function useDiagramStreaming({
     processedToolCallsRef,
     editDiagramOriginalXmlRef,
     loadedMessageIdsRef,
+    turnPageIdRef,
 }: UseDiagramStreamingParams) {
     const { chartXML, chartXMLRef, loadDiagram } = useDiagram()
     const previousXML = useRef<string>("")
@@ -102,10 +107,15 @@ export function useDiagramStreaming({
             if (testDoc.querySelector("parsererror")) return
 
             try {
-                // Replace the first page's cells so other pages stay intact.
-                // An empty canvas gets a default mxfile to put the cells in.
+                // Replace the model's page's cells so other pages stay
+                // intact. An empty canvas gets a default mxfile to put the
+                // cells in.
                 const baseXML = chartXML || BLANK_MXFILE
-                const replacedXML = replaceNodes(baseXML, convertedXml)
+                const replacedXML = replaceNodes(
+                    baseXML,
+                    convertedXml,
+                    turnPageIdRef?.current,
+                )
                 previousXML.current = convertedXml
                 loadDiagram(replacedXML, true, "preview")
             } catch (error) {
@@ -220,8 +230,15 @@ export function useDiagramStreaming({
             if (lastProcessedXmlRef.current.get(opsKey) === ops) continue
             lastProcessedXmlRef.current.set(opsKey, ops)
             try {
-                const edges = edgeIdsOf(originalXml)
-                const layer = defaultLayerOf(originalXml)
+                // The model's page: the edges on it and its first layer,
+                // as editDiagram reads them
+                const selector = pageSelectorFor(
+                    originalXml,
+                    turnPageIdRef?.current,
+                )
+                const page = targetPageXml(originalXml, selector)
+                const edges = edgeIdsOf(page)
+                const layer = defaultLayerOf(page)
                 const { result } = applyDiagramOperations(
                     originalXml,
                     completeOps.map((op) =>
@@ -235,6 +252,7 @@ export function useDiagramStreaming({
                               }
                             : op,
                     ),
+                    selector,
                 )
                 if (lastProcessedXmlRef.current.get(resultKey) === result) {
                     continue

components/chat/use-versions.ts

@@ -46,15 +46,17 @@ export function useVersions({
     // (effect cleanups run before the streaming hook commits). The handler
     // only reads refs and stable functions.
     useEffect(() => {
-        setCommitHandler(({ beforeXml, afterXml, toolCallId }) => {
+        setCommitHandler(({ beforeXml, afterXml, toolCallId, pageId }) => {
             if (restoringRef.current) return
             const { summary, touchedIds, fromScratch } = diffDiagrams(
                 beforeXml,
                 afterXml,
+                pageId,
             )
             const id = useVersionsStore.getState().addVersion({
                 xml: afterXml,
                 beforeXml,
+                ...(pageId && { pageId }),
                 turnIndex: currentTurnRef.current,
                 toolCallId,
                 summary,

contexts/diagram-context.tsx

@@ -23,12 +23,14 @@ import {
     previewDiagram,
     resetPreview,
     revertPreview,
+    selectPage,
 } from "@/lib/drawio/editor-bridge"
 import {
     BLANK_MXFILE,
     normalizeToMxfile,
 } from "@/packages/mcp-server/src/pages.ts"
 import { validateAndFixXml } from "@/packages/mcp-server/src/xml-validation.ts"
+import { useCanvasStore } from "@/stores/canvas-store"
 import { extractDiagramXML, isRealDiagram } from "../lib/utils"
 
 /**
@@ -54,6 +56,8 @@ export interface DiagramCommit {
     beforeXml: string
     afterXml: string
     toolCallId?: string
+    /** The page the change was made on; null or absent: the first page */
+    pageId?: string | null
 }
 
 interface DiagramContextType {
@@ -65,12 +69,12 @@ interface DiagramContextType {
         chart: string,
         skipValidation?: boolean,
         mode?: LoadMode,
-        meta?: { toolCallId?: string },
+        meta?: { toolCallId?: string; pageId?: string | null },
     ) => string | null
     // Returns the export's tag (empty when draw.io is not there yet)
     handleExport: () => string
-    // Pending exports by tag; a plain export's resolver gets the first
-    // page's XML
+    // Pending exports by tag; a plain export's resolver gets the XML of the
+    // page on screen
     exportResolversRef: React.MutableRefObject<
         Record<string, (data: string, xml?: string) => void>
     >
@@ -134,17 +138,28 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
         null,
     )
 
+    // The page to show again once a full load is done: draw.io shows the
+    // first page of a loaded document
+    const pageAfterLoadRef = useRef<string | null>(null)
+
     // Sends a full load. draw.io runs it when its message arrives, and then
     // reports "load": until then the editor shows the diagram from before,
     // so later changes go the same way and keep their order
     const fullLoad = (xml: string) => {
         if (!drawioRef.current) return
+        const pageId = useCanvasStore.getState().currentPageId
+        pageAfterLoadRef.current =
+            pageId && xml.includes(`id="${pageId}"`) ? pageId : null
         pendingLoadsRef.current++
         drawioRef.current.load({ xml })
     }
 
     const onDrawioLoad = () => {
         pendingLoadsRef.current = Math.max(0, pendingLoadsRef.current - 1)
+        // Back to the page the user was on
+        const pageId = pageAfterLoadRef.current
+        pageAfterLoadRef.current = null
+        if (pageId) selectPage(pageId)
         // Only set ready state once to prevent infinite loops
         if (hasCalledOnLoadRef.current) return
         hasCalledOnLoadRef.current = true
@@ -241,7 +256,7 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
         chart: string,
         skipValidation?: boolean,
         mode: LoadMode = "load",
-        meta?: { toolCallId?: string },
+        meta?: { toolCallId?: string; pageId?: string | null },
     ): string | null => {
         // The editor bridge is shared: a page that is gone (another language
         // mounted a new one) must not change the new page's canvas
@@ -308,6 +323,7 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
                 beforeXml,
                 afterXml: xmlToLoad,
                 toolCallId: meta?.toolCallId,
+                pageId: meta?.pageId,
             })
         }
 
@@ -332,7 +348,12 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
         // 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)
+        // The chat gets the page the user is viewing (the first one with
+        // an external draw.io, which cannot tell us)
+        const extractedXML = extractDiagramXML(
+            data.data,
+            useCanvasStore.getState().currentPageId,
+        )
         setLatestSvg(data.data)
 
         // The chat's own export (onFetchChart), not another one in flight

hooks/use-diagram-tool-handlers.ts

@@ -5,6 +5,7 @@ import type {
     ValidationStatus,
 } from "@/components/chat/ValidationCard"
 import type { LoadMode } from "@/contexts/diagram-context"
+import { pageModelXml, pageSelectorFor, placeOnPage } from "@/lib/diagram-pages"
 import type { ValidationResult } from "@/lib/diagram-validator"
 import { formatValidationFeedback } from "@/lib/diagram-validator"
 import { isMxCellXmlComplete } from "@/lib/utils"
@@ -15,7 +16,8 @@ import { hasCells } from "@/packages/mcp-server/src/pages.ts"
 
 const DEBUG = process.env.NODE_ENV === "development"
 
-// display_diagram replaces the document with this one page
+// The page a drawing is wrapped in; on the canvas it replaces the page the
+// model works on, so the canvas page keeps its own id and name
 const NEW_PAGE = { pageId: "page-1", pageName: "Page-1" }
 
 /**
@@ -83,11 +85,14 @@ interface UseDiagramToolHandlersParams {
     // Failed VLM validations in the current user turn (reset on each user message)
     validationRetryCountRef: RefObject<number>
     chartXMLRef: RefObject<string>
+    // The page the model reads and writes in this turn (the one the user
+    // viewed when it began); null, or absent, means the first page
+    turnPageIdRef?: RefObject<string | null>
     onDisplayChart: (
         xml: string,
         skipValidation?: boolean,
         mode?: LoadMode,
-        meta?: { toolCallId?: string },
+        meta?: { toolCallId?: string; pageId?: string | null },
     ) => string | null
     onFetchChart: () => Promise<string>
     captureValidationPng?: () => Promise<string | null>
@@ -119,6 +124,7 @@ export function useDiagramToolHandlers({
     processedToolCallsRef,
     validationRetryCountRef,
     chartXMLRef,
+    turnPageIdRef,
     onDisplayChart,
     onFetchChart,
     captureValidationPng,
@@ -182,11 +188,29 @@ export function useDiagramToolHandlers({
     const commit = (xml: string, toolCallId: string) => {
         const original = continuationOriginalRef.current
         if (original !== null) onDisplayChart(original, true, "revert")
-        const error = onDisplayChart(xml, true, "commit", { toolCallId })
+        const error = onDisplayChart(xml, true, "commit", {
+            toolCallId,
+            pageId: aiPageId(),
+        })
         if (!error) continuationOriginalRef.current = null
         return error
     }
 
+    const aiPageId = () => turnPageIdRef?.current ?? null
+
+    // The diagram as the model sees it in the chat request: its page only
+    const shownToModel = (xml: string) =>
+        foldCells(pageModelXml(xml, aiPageId()) ?? xml)
+
+    // A drawn page goes onto the model's page of the canvas file, which
+    // keeps its other pages; a drawing that brings several pages replaces
+    // the file, and gets its variables
+    const onAiPage = (drawnXml: string) =>
+        keepFileVars(
+            placeOnPage(drawnXml, chartXMLRef.current, aiPageId()),
+            chartXMLRef.current,
+        )
+
     const takeOriginals = (): string | undefined => {
         const [originalXml] = editDiagramOriginalXmlRef.current.values()
         for (const id of editDiagramOriginalXmlRef.current.keys()) {
@@ -265,10 +289,7 @@ NEXT STEP: Call append_diagram with the continuation XML.
         // finds the diagram already in place
         const prepared = prepareNewDiagram(finalXml, NEW_PAGE)
         const validationError = prepared.ok
-            ? commit(
-                  keepFileVars(prepared.xml, chartXMLRef.current),
-                  toolCall.toolCallId,
-              )
+            ? commit(onAiPage(prepared.xml), toolCall.toolCallId)
             : prepared.error
 
         if (validationError) {
@@ -503,8 +524,12 @@ ${finalXml}
             }
 
             // All or nothing, checked like the MCP server's edit_diagram.
-            // The model sees the first page, so edits target it.
-            const outcome = editDiagram(currentXml, operations, {})
+            // The model sees its page, so edits target it.
+            const outcome = editDiagram(
+                currentXml,
+                operations,
+                pageSelectorFor(currentXml, aiPageId()),
+            )
             if (!outcome.ok) {
                 const reason = outcome.pageError
                     ? outcome.errors[0]
@@ -518,7 +543,7 @@ ${finalXml}
 
 Current diagram XML:
 \`\`\`xml
-${foldCells(currentXml)}
+${shownToModel(currentXml)}
 \`\`\`
 
 Please check the cell IDs and retry.`,
@@ -547,7 +572,7 @@ Please check the cell IDs and retry.`,
 
 Current diagram XML:
 \`\`\`xml
-${currentXml ? foldCells(currentXml) : "No XML available"}
+${currentXml ? shownToModel(currentXml) : "No XML available"}
 \`\`\`
 
 Please check cell IDs and retry, or use display_diagram to regenerate.`,
@@ -620,10 +645,7 @@ Start your continuation with the NEXT character after where it stopped.`,
             // originals, so the preview code undoes none of them later
             const originalXml = takeOriginals()
             const validationError = prepared.ok
-                ? commit(
-                      keepFileVars(prepared.xml, chartXMLRef.current),
-                      toolCall.toolCallId,
-                  )
+                ? commit(onAiPage(prepared.xml), toolCall.toolCallId)
                 : prepared.error
 
             if (validationError) {

lib/diagram-diff.ts

@@ -1,6 +1,6 @@
-import { decompressPageContent } from "@/packages/mcp-server/src/load-diagram.ts"
+import { modelOfPage, pageElement } from "@/lib/diagram-pages"
 
-/** What changed on the first page between two versions of a diagram */
+/** What changed on one page between two versions of a diagram */
 export interface ChangeSummary {
     shapesAdded: number
     shapesRemoved: number
@@ -32,27 +32,20 @@ interface CellInfo {
     signature: string
 }
 
-/** A page's mxGraphModel element, inflated when the page is compressed */
-function modelOfPage(diagram: Element): Element | null {
-    const model = diagram.querySelector("mxGraphModel")
-    if (model) return model
-    const inflated = decompressPageContent(diagram.textContent || "")
-    if (!inflated) return null
-    const inner = new DOMParser().parseFromString(inflated, "text/xml")
-    return inner.querySelector("mxGraphModel")
-}
-
 function parse(xml: string): Document | null {
     if (!xml?.trim()) return null
     const doc = new DOMParser().parseFromString(xml, "text/xml")
     return doc.querySelector("parsererror") ? null : doc
 }
 
-/** The first page's mxGraphModel element, whatever wrapper the XML has */
-function firstPageModel(xml: string): Element | null {
+/**
+ * A page's mxGraphModel element, whatever wrapper the XML has: the page
+ * with this id, else the first
+ */
+function pageModel(xml: string, pageId?: string | null): Element | null {
     const doc = parse(xml)
     if (!doc) return null
-    const diagram = doc.querySelector("diagram")
+    const diagram = pageElement(doc, pageId)
     if (diagram) return modelOfPage(diagram)
     return doc.querySelector("mxGraphModel") ?? doc.documentElement
 }
@@ -133,10 +126,13 @@ function cellsIn(model: Element | null) {
     })
 }
 
-/** Shapes and connectors of the first page (not the root and its layers) */
-function collectCells(xml: string): Map<string, CellInfo> {
+/** Shapes and connectors of one page (not the root and its layers) */
+function collectCells(
+    xml: string,
+    pageId?: string | null,
+): Map<string, CellInfo> {
     const cells = new Map<string, CellInfo>()
-    for (const { id, cell, node } of cellsIn(firstPageModel(xml))) {
+    for (const { id, cell, node } of cellsIn(pageModel(xml, pageId))) {
         if (id === "0" || cell.getAttribute("parent") === "0") continue
         cells.set(id, {
             isEdge: cell.getAttribute("edge") === "1",
@@ -146,9 +142,14 @@ function collectCells(xml: string): Map<string, CellInfo> {
     return cells
 }
 
-export function diffDiagrams(beforeXml: string, afterXml: string): DiagramDiff {
-    const before = collectCells(beforeXml)
-    const after = collectCells(afterXml)
+/** What changed on one page (the first when no id is given) */
+export function diffDiagrams(
+    beforeXml: string,
+    afterXml: string,
+    pageId?: string | null,
+): DiagramDiff {
+    const before = collectCells(beforeXml, pageId)
+    const after = collectCells(afterXml, pageId)
     const summary: ChangeSummary = { ...EMPTY_SUMMARY }
     const touchedIds: string[] = []
 

lib/diagram-pages.ts

@@ -0,0 +1,138 @@
+/**
+ * The page the AI works on is the page the user is viewing. These helpers
+ * read that page out of a multi-page document and write a drawing into it
+ * while the other pages stay as they are. A page id that matches no page
+ * (or none at all, as with an external draw.io that cannot tell us) means
+ * the first page, which is what the app always used before.
+ */
+
+import { decompressPageContent } from "@/packages/mcp-server/src/load-diagram.ts"
+import {
+    BLANK_MXFILE,
+    type PageSelector,
+} from "@/packages/mcp-server/src/pages.ts"
+
+function parse(xml: string): Document | null {
+    if (!xml?.trim()) return null
+    const doc = new DOMParser().parseFromString(xml, "text/xml")
+    return doc.querySelector("parsererror") ? null : doc
+}
+
+/** The <diagram> the AI works on: the one with this id, else the first */
+export function pageElement(
+    doc: Document,
+    pageId: string | null | undefined,
+): Element | null {
+    const diagrams = Array.from(doc.getElementsByTagName("diagram"))
+    if (diagrams.length === 0) return null
+    return (
+        (pageId &&
+            diagrams.find(
+                (diagram) => diagram.getAttribute("id") === pageId,
+            )) ||
+        diagrams[0]
+    )
+}
+
+/** A page's mxGraphModel element, inflated when the page is compressed */
+export function modelOfPage(diagram: Element): Element | null {
+    const model = diagram.querySelector("mxGraphModel")
+    if (model) return model
+    const inflated = decompressPageContent(diagram.textContent || "")
+    if (!inflated) return null
+    return parse(inflated)?.querySelector("mxGraphModel") ?? null
+}
+
+/**
+ * Selector for editDiagram and applyDiagramOperations: the page with this
+ * id when the document has it, else their default (the first page)
+ */
+export function pageSelectorFor(
+    xml: string,
+    pageId: string | null | undefined,
+): PageSelector {
+    if (!pageId) return {}
+    const doc = parse(xml)
+    const found = doc
+        ? Array.from(doc.getElementsByTagName("diagram")).some(
+              (diagram) => diagram.getAttribute("id") === pageId,
+          )
+        : false
+    return found ? { page_id: pageId } : {}
+}
+
+/**
+ * The page's mxGraphModel as XML. A bare mxGraphModel is returned as it
+ * is; null when the XML does not parse or has no model.
+ */
+export function pageModelXml(
+    xml: string,
+    pageId: string | null | undefined,
+): string | null {
+    const doc = parse(xml)
+    if (!doc) return null
+    const root = doc.documentElement
+    if (root.nodeName === "mxGraphModel") return xml
+    if (root.nodeName !== "mxfile") return null
+    const diagram = pageElement(doc, pageId)
+    const model = diagram ? modelOfPage(diagram) : null
+    return model ? new XMLSerializer().serializeToString(model) : null
+}
+
+/**
+ * The document with one page's content replaced by this model. The page
+ * keeps its id and name, the other pages and the file's attributes (its
+ * variables) stay. Without a document to put the page in, the blank
+ * one-page file is used.
+ */
+export function replacePageModel(
+    fileXml: string,
+    pageId: string | null | undefined,
+    modelXml: string,
+): string {
+    const model = parse(modelXml)
+    if (model?.documentElement.nodeName !== "mxGraphModel") {
+        throw new Error("replacePageModel needs an <mxGraphModel>")
+    }
+    let doc = parse(fileXml)
+    if (doc?.documentElement.nodeName !== "mxfile") doc = parse(BLANK_MXFILE)
+    let diagram = pageElement(doc as Document, pageId)
+    if (!diagram) {
+        doc = parse(BLANK_MXFILE)
+        diagram = pageElement(doc as Document, null)
+    }
+    const target = diagram as Element
+    while (target.firstChild) target.removeChild(target.firstChild)
+    target.appendChild(
+        (doc as Document).importNode(model.documentElement, true),
+    )
+    return new XMLSerializer().serializeToString(doc as Document)
+}
+
+/**
+ * Where a drawn diagram goes. One page's worth (bare cells wrapped by
+ * prepareNewDiagram, or a one-page file) replaces the AI's page of the
+ * canvas file; file variables the drawn file sets replace the canvas
+ * file's. A file with several pages replaces the whole document: the model
+ * wrote the pages on purpose.
+ */
+export function placeOnPage(
+    drawnXml: string,
+    canvasXml: string,
+    pageId: string | null | undefined,
+): string {
+    const doc = parse(drawnXml)
+    const drawnFile =
+        doc?.documentElement.nodeName === "mxfile" ? doc.documentElement : null
+    if (drawnFile && drawnFile.getElementsByTagName("diagram").length > 1) {
+        return drawnXml
+    }
+    const model = pageModelXml(drawnXml, null)
+    if (!model) return drawnXml
+    const placed = replacePageModel(canvasXml, pageId, model)
+    const vars = drawnFile?.getAttribute("vars")
+    if (vars === null || vars === undefined) return placed
+    const result = parse(placed) as Document
+    result.documentElement.setAttribute("vars", vars)
+    return new XMLSerializer().serializeToString(result)
+}

lib/drawio/editor-bridge.ts

@@ -9,7 +9,7 @@
  * With an external draw.io (cross-origin) none of this is available and the
  * app uses the postMessage protocol only.
  */
-import { sameFileVars } from "@/lib/diagram-diff"
+import { isSameDocument, sameFileVars } from "@/lib/diagram-diff"
 import { hasCells } from "@/packages/mcp-server/src/pages.ts"
 import { type SelectedCell, useCanvasStore } from "@/stores/canvas-store"
 
@@ -152,23 +152,18 @@ function withoutUndo(fn: () => void) {
     }
 }
 
-/**
- * replaceDiagramData replaces the current page with one <mxGraphModel>. Only
- * single-page documents qualify; multi-page ones fall back to a full load.
- */
-function toSinglePageModel(xml: string): string | null {
-    const doc = new DOMParser().parseFromString(xml, "text/xml")
-    if (doc.querySelector("parsererror")) return null
-    const root = doc.documentElement
-    if (root.nodeName === "mxGraphModel") return xml
-    if (root.nodeName !== "mxfile") return null
-    const diagrams = root.getElementsByTagName("diagram")
-    if (diagrams.length !== 1) return null
-    const model = diagrams[0].getElementsByTagName("mxGraphModel")[0]
+/** Id of the page on the canvas, when draw.io has pages */
+function currentPageId(): string | null {
+    const id = ui?.currentPage?.getId?.()
+    return id === undefined || id === null ? null : String(id)
+}
+
+/** A page's <mxGraphModel> as XML, inflated when the page is compressed */
+function modelXmlOf(diagram: Element): string | null {
+    const model = diagram.getElementsByTagName("mxGraphModel")[0]
     if (model) return new XMLSerializer().serializeToString(model)
-    // Compressed page
     try {
-        const text = diagrams[0].textContent?.trim()
+        const text = diagram.textContent?.trim()
         const inflated = text ? win?.Graph?.decompress?.(text) : null
         return typeof inflated === "string" &&
             inflated.includes("<mxGraphModel")
@@ -179,8 +174,62 @@ function toSinglePageModel(xml: string): string | null {
     }
 }
 
+/**
+ * replaceDiagramData replaces the current page with one <mxGraphModel>:
+ * the document's page with the canvas page's id, or its only page
+ */
+function pageModelOf(xml: string): string | null {
+    const doc = new DOMParser().parseFromString(xml, "text/xml")
+    if (doc.querySelector("parsererror")) return null
+    const root = doc.documentElement
+    if (root.nodeName === "mxGraphModel") return xml
+    if (root.nodeName !== "mxfile") return null
+    const diagrams = Array.from(root.getElementsByTagName("diagram"))
+    const id = currentPageId()
+    const diagram =
+        diagrams.length === 1
+            ? diagrams[0]
+            : diagrams.find((d) => d.getAttribute("id") === id)
+    return diagram ? modelXmlOf(diagram) : null
+}
+
+/**
+ * The document without the current page, for comparing the other pages.
+ * null when the XML is not a document; "" for a single page.
+ */
+function withoutCurrentPage(xml: string | null): string | null {
+    if (!xml) return null
+    const doc = new DOMParser().parseFromString(xml, "text/xml")
+    if (doc.querySelector("parsererror")) return null
+    const root = doc.documentElement
+    if (root.nodeName === "mxGraphModel") return ""
+    if (root.nodeName !== "mxfile") return null
+    const diagrams = Array.from(root.getElementsByTagName("diagram"))
+    if (diagrams.length <= 1) return ""
+    const id = currentPageId()
+    for (const diagram of diagrams) {
+        if (diagram.getAttribute("id") === id) root.removeChild(diagram)
+    }
+    return new XMLSerializer().serializeToString(root)
+}
+
+/**
+ * The document changes the current page only: its other pages are the
+ * canvas's (same names, cells and page settings; draw.io fills in settings
+ * a loaded file left out, so the text can differ)
+ */
+function otherPagesSame(xml: string): boolean {
+    const theirs = withoutCurrentPage(xml)
+    if (theirs === null) return false
+    const pageCount = Array.isArray(ui.pages) ? ui.pages.length : 1
+    if (pageCount <= 1 && theirs === "") return true
+    const ours = withoutCurrentPage(currentFileXml())
+    if (ours === null || ours === "" || theirs === "") return false
+    return isSameDocument(theirs, ours)
+}
+
 function replace(xml: string) {
-    const model = toSinglePageModel(xml)
+    const model = pageModelOf(xml)
     if (!model || typeof ui?.replaceDiagramData !== "function") {
         throw new Error("Diagram can't be replaced in place")
     }
@@ -202,7 +251,7 @@ function isEmptyModel(): boolean {
 
 /** Whether a page has shapes; layers (cells under the root) are none */
 function hasShapes(xml: string): boolean {
-    const model = toSinglePageModel(xml)
+    const model = pageModelOf(xml)
     if (model === null) return hasCells(xml)
     const cells = new DOMParser()
         .parseFromString(model, "text/xml")
@@ -256,8 +305,9 @@ export function canReplaceDiagram(xml: string): boolean {
         warnOnce("replace", "replaceDiagramData not found, using full loads")
         return false
     }
-    const pageCount = Array.isArray(ui.pages) ? ui.pages.length : 1
-    if (pageCount > 1) return false
+    // Replacing changes the current page only: a document whose other
+    // pages differ from the canvas's loads in full
+    if (!otherPagesSame(xml)) return false
     // Replacing the page keeps the file's variables: other ones, or none
     // over a file with some, load in full
     if (
@@ -268,7 +318,7 @@ export function canReplaceDiagram(xml: string): boolean {
     ) {
         return false
     }
-    const model = toSinglePageModel(xml)
+    const model = pageModelOf(xml)
     // A document with them, or replacing one with them, loads in full
     return (
         model !== null &&
@@ -285,13 +335,21 @@ export function previewDiagram(xml: string) {
     if (wasEmpty) fitDiagram()
 }
 
-/** Name and id of the page in a single-page mxfile, if it has them */
+/**
+ * Name and id the document gives the canvas page: of its page with that
+ * id, or of its only page
+ */
 function pageOf(xml: string): { name: string | null; id: string | null } {
     const doc = new DOMParser().parseFromString(xml, "text/xml")
-    const diagram =
+    const diagrams =
         doc.documentElement?.nodeName === "mxfile"
-            ? doc.getElementsByTagName("diagram")[0]
-            : undefined
+            ? Array.from(doc.getElementsByTagName("diagram"))
+            : []
+    const id = currentPageId()
+    const diagram =
+        diagrams.length === 1
+            ? diagrams[0]
+            : diagrams.find((d) => d.getAttribute("id") === id)
     return {
         name: diagram?.getAttribute("name") || null,
         id: diagram?.getAttribute("id") || null,
@@ -322,7 +380,7 @@ function changePageId(page: any, id: string) {
 export function commitDiagram(xml: string) {
     const wasEmpty =
         isEmptyModel() || (previewBase !== null && !hasShapes(previewBase))
-    const base = previewBase ? toSinglePageModel(previewBase) : null
+    const base = previewBase ? pageModelOf(previewBase) : null
     previewBase = null
     // Undo goes back to the diagram before streaming started. draw.io's
     // ReplaceDiagram change keeps the document it replaced for undo: hand it
@@ -337,7 +395,7 @@ export function commitDiagram(xml: string) {
     const page = ui?.currentPage
     model?.beginUpdate()
     try {
-        const next = toSinglePageModel(xml)
+        const next = pageModelOf(xml)
         if (direct && next) {
             const change = new ReplaceDiagram(ui, parse(next))
             model.execute(change)
@@ -368,6 +426,18 @@ export function resetPreview() {
     previewBase = null
 }
 
+/** Show the page with this id; false when draw.io has no such page */
+export function selectPage(pageId: string): boolean {
+    try {
+        const page = ui?.getPageById?.(pageId)
+        if (!page || typeof ui.selectPage !== "function") return false
+        if (ui.currentPage !== page) ui.selectPage(page, true)
+        return true
+    } catch {
+        return false
+    }
+}
+
 // When the app last fitted the diagram on its own; a canvas resize right
 // after (the chat panel sliding in) fits again, unless the user zoomed since
 let lastAutoFitAt = 0

lib/utils.ts

@@ -227,12 +227,18 @@ export function convertToLegalXml(xmlString: string): string {
 }
 
 /**
- * Replace nodes in a Draw.io XML diagram
+ * Replace the cells of one page of a Draw.io XML document
  * @param currentXML - The original Draw.io XML string
  * @param nodes - The XML string containing new nodes to replace in the diagram
+ * @param pageId - The page whose cells are replaced; the first page when
+ * missing or not found
  * @returns The updated XML string with replaced nodes
  */
-export function replaceNodes(currentXML: string, nodes: string): string {
+export function replaceNodes(
+    currentXML: string,
+    nodes: string,
+    pageId?: string | null,
+): string {
     // Check for valid inputs
     if (!currentXML || !nodes) {
         throw new Error("Both currentXML and nodes must be provided")
@@ -251,16 +257,23 @@ export function replaceNodes(currentXML: string, nodes: string): string {
 
         const nodesDoc = parser.parseFromString(nodesString, "text/xml")
 
+        // The page to draw on, when the document has it
+        const page = pageId
+            ? Array.from(currentDoc.getElementsByTagName("diagram")).find(
+                  (diagram) => diagram.getAttribute("id") === pageId,
+              )
+            : undefined
+        const scope: ParentNode = page ?? currentDoc
         // Find the root element in the current document
-        let currentRoot = currentDoc.querySelector("mxGraphModel > root")
+        let currentRoot = scope.querySelector("mxGraphModel > root")
         if (!currentRoot) {
             // If no root element is found, create the proper structure
             const mxGraphModel =
-                currentDoc.querySelector("mxGraphModel") ||
+                scope.querySelector("mxGraphModel") ||
                 currentDoc.createElement("mxGraphModel")
 
             if (!currentDoc.contains(mxGraphModel)) {
-                currentDoc.appendChild(mxGraphModel)
+                ;(page ?? currentDoc).appendChild(mxGraphModel)
             }
 
             currentRoot = currentDoc.createElement("root")
@@ -330,10 +343,13 @@ export function replaceNodes(currentXML: string, nodes: string): string {
 
 /**
  * Decode an xmlsvg export (SVG data URL) into uncompressed diagram XML.
- * Only the first page is returned; for the full multi-page document use the
- * autosaved chartXML instead.
+ * One page is returned: the one with the given id, else the first; for the
+ * full multi-page document use the autosaved chartXML instead.
  */
-export function extractDiagramXML(xml_svg_string: string): string {
+export function extractDiagramXML(
+    xml_svg_string: string,
+    pageId?: string | null,
+): string {
     try {
         // 1. Parse the SVG string (using built-in DOMParser in a browser-like environment)
         const svgString = atob(xml_svg_string.slice(26))
@@ -357,7 +373,13 @@ export function extractDiagramXML(xml_svg_string: string): string {
 
         // 4. Parse the XML content
         const xmlDoc = parser.parseFromString(xmlContent, "text/xml")
-        const diagramElement = xmlDoc.querySelector("diagram")
+        const diagrams = Array.from(xmlDoc.getElementsByTagName("diagram"))
+        const diagramElement =
+            (pageId &&
+                diagrams.find(
+                    (diagram) => diagram.getAttribute("id") === pageId,
+                )) ||
+            diagrams[0]
 
         if (!diagramElement) {
             throw new Error("No diagram element found")

packages/mcp-server/src/edit-diagram.ts

@@ -53,7 +53,7 @@ export function editDiagram(
     // Edges already on the page, so a label added to one is not wrapped, and
     // the layer a compact cell without a parent goes on
     const page = targetPageXml(xml, selector)
-    const edges = edgeIdsOf(xml)
+    const edges = edgeIdsOf(page)
     const layer = defaultLayerOf(page)
     for (const op of operations) {
         if (op.operation === "delete" || !op.new_xml) {

stores/versions-store.ts

@@ -10,6 +10,8 @@ export interface DiagramVersion {
     number: number
     /** Full multi-page document after the change */
     xml: string
+    /** The page the change was made on; older sessions have none (the first page) */
+    pageId?: string
     /** Diagram before the change; kept only on the latest version */
     beforeXml?: string
     /** SVG data URL thumbnail; "" when none could be made */

tests/e2e/diagram-content.spec.ts

@@ -131,7 +131,7 @@ const NEW_CELLS =
     cell("3", "R&D", 400) +
     `<UserObject id="4" label="Docs" link="https://example.com"><mxCell style="rounded=1;" vertex="1" parent="1"><mxGeometry x="580" y="40" width="120" height="60" as="geometry"/></mxCell></UserObject>`
 
-test("display_diagram replaces the document with the fixed diagram", async ({
+test("display_diagram redraws the page on screen and keeps the other pages", async ({
     page: p,
 }) => {
     const canvas = await mockReplies(p, [
@@ -155,9 +155,11 @@ test("display_diagram replaces the document with the fixed diagram", async ({
             timeout: 15000,
         })
     }
-    // The old pages are gone
+    // The first page was redrawn; the second page is still there
     await expect(canvas.getByText("Old A")).toHaveCount(0)
-    await expect(pageTabs.getByText("Second", { exact: true })).toHaveCount(0)
+    await expect(pageTabs.getByText("Second", { exact: true })).toBeVisible()
+    await pageTabs.getByText("Second", { exact: true }).click()
+    await expect(canvas.getByText("Old B")).toBeVisible()
 })
 
 test("an edit with a fixable cell is fixed, not rejected", async ({

tests/e2e/workspace.spec.ts

@@ -1001,7 +1001,7 @@ test.describe("Workspace", () => {
         await expect(canvas.getByText("Old A")).toBeVisible()
         await sendMessage(page, "Create a flowchart")
         await waitForComplete(page)
-        // Multi-page documents are loaded in full: the card is the way back
+        // The card is one way back
         await page.waitForTimeout(2000)
         await page.locator('[data-testid="version-undo"]').click()
         await expect(canvas.getByText("Old A")).toBeVisible()
@@ -1533,7 +1533,7 @@ test.describe("Workspace", () => {
         await expect.poll(savedPageId, { timeout: 10000 }).toBe("orig")
     })
 
-    test("Ctrl+Z after an AI drawing gives the page its id back", async ({
+    test("an AI drawing keeps the page's id, for links to the page", async ({
         page,
     }) => {
         await page.route("**/api/chat", drawThenEdit())
@@ -1567,7 +1567,7 @@ test.describe("Workspace", () => {
                     .querySelector("diagram")
                     ?.getAttribute("id")
             })
-        await expect.poll(savedPageId, { timeout: 10000 }).toBe("page-1")
+        await expect.poll(savedPageId, { timeout: 10000 }).toBe("orig")
         await canvas
             .locator(".geDiagramContainer")
             .click({ position: { x: 10, y: 10 } })
@@ -1576,7 +1576,62 @@ test.describe("Workspace", () => {
         await expect.poll(savedPageId, { timeout: 10000 }).toBe("orig")
         await page.keyboard.press("ControlOrMeta+Shift+z")
         await expect(canvas.getByText("Process", { exact: true })).toBeVisible()
-        await expect.poll(savedPageId, { timeout: 10000 }).toBe("page-1")
+        await expect.poll(savedPageId, { timeout: 10000 }).toBe("orig")
+    })
+
+    test("the AI reads and draws the page the user is viewing", async ({
+        page,
+    }) => {
+        const bodies: any[] = []
+        await page.route(
+            "**/api/chat",
+            drawThenEdit((body) => bodies.push(body)),
+        )
+        await openApp(page)
+        await openDrawioFile(page, "two.drawio", TWO_PAGE_FILE)
+        const canvas = getIframeContent(page)
+        await expect(canvas.getByText("Old A")).toBeVisible()
+        // The second page, through draw.io's own tabs
+        await drawioTabs(page).getByText("P1", { exact: true }).click()
+        await expect(canvas.getByText("Old B")).toBeVisible()
+        await sendMessage(page, "Create a flowchart")
+        await waitForComplete(page)
+        // The model got the page on screen
+        expect(bodies[0].xml).toContain("Old B")
+        expect(bodies[0].xml).not.toContain("Old A")
+        // Its drawing replaced that page, and only that page
+        await expect(canvas.getByText("Process", { exact: true })).toBeVisible()
+        await expect(canvas.getByText("Old B")).toHaveCount(0)
+        const pagesIn = (xml: string) =>
+            Object.fromEntries(
+                Array.from(
+                    xml.matchAll(
+                        /<diagram[^>]*\bid="([^"]+)"[^>]*>([\s\S]*?)<\/diagram>/g,
+                    ),
+                    (m) => [m[1], m[2]],
+                ),
+            )
+        await expect
+            .poll(
+                async () => {
+                    const pages = pagesIn(await savedDiagram(page))
+                    return {
+                        p0: pages.p0?.includes("Old A"),
+                        p1: pages.p1?.includes("Process"),
+                        oldB: pages.p1?.includes("Old B"),
+                    }
+                },
+                { timeout: 10000 },
+            )
+            .toEqual({ p0: true, p1: true, oldB: false })
+        // One undo step takes the drawing back, as on a one-page file
+        await canvas
+            .locator(".geDiagramContainer")
+            .click({ position: { x: 10, y: 10 } })
+        await page.keyboard.press("ControlOrMeta+z")
+        await expect(canvas.getByText("Old B")).toBeVisible()
+        await drawioTabs(page).getByText("P0", { exact: true }).click()
+        await expect(canvas.getByText("Old A")).toBeVisible()
     })
 
     test("page settings the AI writes in single quotes are applied too", async ({

tests/unit/diagram-diff.test.ts

@@ -153,9 +153,19 @@ describe("diffDiagrams", () => {
         expect(touchedIds).toEqual(["u1"])
     })
 
-    it("only looks at the first page", () => {
+    it("looks at the first page unless told which page", () => {
         const twoPages = file(pageOf("p1", box("a")), pageOf("p2", box("z")))
         expect(diffDiagrams("", twoPages).summary.shapesAdded).toBe(1)
+        const changedP2 = file(
+            pageOf("p1", box("a")),
+            pageOf("p2", box("z") + box("y")),
+        )
+        expect(diffDiagrams(twoPages, changedP2, "p2").touchedIds).toEqual([
+            "y",
+        ])
+        expect(diffDiagrams(twoPages, changedP2, "p1").touchedIds).toEqual([])
+        // An unknown page is the first one
+        expect(diffDiagrams(twoPages, changedP2, "p9").touchedIds).toEqual([])
     })
 
     it("reads compressed pages", () => {

tests/unit/diagram-pages.test.ts

@@ -0,0 +1,125 @@
+import pako from "pako"
+import { describe, expect, it } from "vitest"
+import {
+    pageModelXml,
+    pageSelectorFor,
+    placeOnPage,
+    replacePageModel,
+} from "@/lib/diagram-pages"
+
+const box = (id: string) =>
+    `<mxCell id="${id}" value="${id}" vertex="1" parent="1"><mxGeometry x="0" y="0" width="80" height="40" as="geometry"/></mxCell>`
+const model = (cells: string) =>
+    `<mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/>${cells}</root></mxGraphModel>`
+const page = (id: string, cells: string) =>
+    `<diagram id="${id}" name="Page ${id}">${model(cells)}</diagram>`
+const packedPage = (id: string, cells: string) => {
+    const packed = Buffer.from(
+        pako.deflateRaw(encodeURIComponent(model(cells))),
+    ).toString("base64")
+    return `<diagram id="${id}" name="Page ${id}">${packed}</diagram>`
+}
+const TWO_PAGES = `<mxfile vars="{&quot;x&quot;:&quot;1&quot;}">${page("a", box("A"))}${page("b", box("B"))}</mxfile>`
+
+const pagesOf = (xml: string) =>
+    Array.from(
+        new DOMParser()
+            .parseFromString(xml, "text/xml")
+            .getElementsByTagName("diagram"),
+    ).map((d) => ({
+        id: d.getAttribute("id"),
+        name: d.getAttribute("name"),
+        cells: Array.from(d.getElementsByTagName("mxCell"))
+            .map((c) => c.getAttribute("id"))
+            .filter((id) => id !== "0" && id !== "1"),
+    }))
+
+describe("pageModelXml", () => {
+    it("gives the page with the id, and the first page without one", () => {
+        expect(pageModelXml(TWO_PAGES, "b")).toContain('id="B"')
+        expect(pageModelXml(TWO_PAGES, "b")).not.toContain('id="A"')
+        expect(pageModelXml(TWO_PAGES, null)).toContain('id="A"')
+        expect(pageModelXml(TWO_PAGES, "missing")).toContain('id="A"')
+    })
+
+    it("inflates a compressed page and passes a bare model through", () => {
+        const file = `<mxfile>${page("a", box("A"))}${packedPage("b", box("B"))}</mxfile>`
+        expect(pageModelXml(file, "b")).toContain('id="B"')
+        expect(pageModelXml(model(box("Z")), "b")).toBe(model(box("Z")))
+        expect(pageModelXml("<broken", "a")).toBeNull()
+    })
+})
+
+describe("pageSelectorFor", () => {
+    it("selects the page only when the document has it", () => {
+        expect(pageSelectorFor(TWO_PAGES, "b")).toEqual({ page_id: "b" })
+        expect(pageSelectorFor(TWO_PAGES, "missing")).toEqual({})
+        expect(pageSelectorFor(TWO_PAGES, null)).toEqual({})
+        expect(pageSelectorFor(model(box("Z")), "b")).toEqual({})
+    })
+})
+
+describe("replacePageModel", () => {
+    it("replaces one page and keeps the others, the names and the variables", () => {
+        const result = replacePageModel(TWO_PAGES, "b", model(box("N")))
+        expect(pagesOf(result)).toEqual([
+            { id: "a", name: "Page a", cells: ["A"] },
+            { id: "b", name: "Page b", cells: ["N"] },
+        ])
+        expect(result).toContain('vars="{&quot;x&quot;:&quot;1&quot;}"')
+    })
+
+    it("falls back to the first page for an unknown or missing id", () => {
+        for (const id of ["missing", null]) {
+            const result = replacePageModel(TWO_PAGES, id, model(box("N")))
+            expect(pagesOf(result).map((p) => p.cells)).toEqual([["N"], ["B"]])
+        }
+    })
+
+    it("starts from the blank file when there is no document", () => {
+        for (const base of ["", "<mxGraphModel/>", "<broken"]) {
+            const result = replacePageModel(base, "any", model(box("N")))
+            expect(pagesOf(result)).toEqual([
+                { id: "page-1", name: "Page-1", cells: ["N"] },
+            ])
+        }
+    })
+
+    it("refuses anything but a model", () => {
+        expect(() => replacePageModel(TWO_PAGES, "a", "<root/>")).toThrow()
+    })
+})
+
+describe("placeOnPage", () => {
+    it("puts one page's worth of drawing on the model's page", () => {
+        const drawn = `<mxfile><diagram id="page-1" name="Page-1">${model(box("N"))}</diagram></mxfile>`
+        const result = placeOnPage(drawn, TWO_PAGES, "b")
+        expect(pagesOf(result)).toEqual([
+            { id: "a", name: "Page a", cells: ["A"] },
+            { id: "b", name: "Page b", cells: ["N"] },
+        ])
+    })
+
+    it("applies the drawn file's variables, and keeps the canvas's otherwise", () => {
+        const withVars = `<mxfile vars="{&quot;team&quot;:&quot;New&quot;}"><diagram id="p" name="P">${model(box("N"))}</diagram></mxfile>`
+        expect(placeOnPage(withVars, TWO_PAGES, "b")).toContain(
+            'vars="{&quot;team&quot;:&quot;New&quot;}"',
+        )
+        const without = `<mxfile><diagram id="p" name="P">${model(box("N"))}</diagram></mxfile>`
+        expect(placeOnPage(without, TWO_PAGES, "b")).toContain(
+            'vars="{&quot;x&quot;:&quot;1&quot;}"',
+        )
+    })
+
+    it("lets a drawing with several pages replace the document", () => {
+        const drawn = `<mxfile>${page("x", box("X"))}${page("y", box("Y"))}</mxfile>`
+        expect(placeOnPage(drawn, TWO_PAGES, "b")).toBe(drawn)
+    })
+
+    it("draws on the blank file when the canvas is empty", () => {
+        const drawn = `<mxfile><diagram id="page-1" name="Page-1">${model(box("N"))}</diagram></mxfile>`
+        expect(pagesOf(placeOnPage(drawn, "", null))).toEqual([
+            { id: "page-1", name: "Page-1", cells: ["N"] },
+        ])
+    })
+})

tests/unit/mcp-core.test.ts

@@ -263,3 +263,50 @@ describe("applyDiagramOperations with wrapped cells", () => {
         expect(errors[0]?.message).toContain("not found")
     })
 })
+
+describe("edit_diagram on one page of a multi-page file", () => {
+    // Page A has an edge "g"; on page B, "g" is a shape on layer "L"
+    const twoPages = `<mxfile><diagram id="pa" name="A"><mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/>${box("s")}${box("t")}<mxCell id="g" edge="1" parent="1" source="s" target="t"><mxGeometry relative="1" as="geometry"/></mxCell></root></mxGraphModel></diagram><diagram id="pb" name="B"><mxGraphModel><root><mxCell id="0"/><mxCell id="L" parent="0"/>${box("g", "L")}</root></mxGraphModel></diagram></mxfile>`
+    const cellOn = (xml: string, pageId: string, id: string) =>
+        new DOMParser()
+            .parseFromString(xml, "text/xml")
+            .querySelector(`diagram[id="${pageId}"] mxCell[id="${id}"]`)
+
+    it("puts a compact cell without a parent on that page's layer", () => {
+        const outcome = editDiagram(
+            twoPages,
+            [
+                {
+                    operation: "add",
+                    cell_id: "c",
+                    new_xml:
+                        '<mxCell id="c" value="c" x="10" y="10" w="80" h="40"/>',
+                },
+            ],
+            { page_id: "pb" },
+        )
+        if (!outcome.ok) throw new Error(outcome.errors.join("; "))
+        expect(cellOn(outcome.xml, "pb", "c")?.getAttribute("parent")).toBe("L")
+        expect(cellOn(outcome.xml, "pa", "c")).toBeNull()
+    })
+
+    it("reads the edges of that page only", () => {
+        // A shape inside the group "g" is no edge label on page B
+        const outcome = editDiagram(
+            twoPages,
+            [
+                {
+                    operation: "add",
+                    cell_id: "c",
+                    new_xml:
+                        '<mxCell id="c" value="c" parent="g" x="10" y="10" w="80" h="40"/>',
+                },
+            ],
+            { page_id: "pb" },
+        )
+        if (!outcome.ok) throw new Error(outcome.errors.join("; "))
+        expect(cellOn(outcome.xml, "pb", "c")?.getAttribute("style")).toContain(
+            "whiteSpace=wrap",
+        )
+    })
+})

tests/unit/use-diagram-streaming.test.tsx

@@ -13,12 +13,14 @@ const AFTER_FIRST_EDIT = page(box("a") + box("b"))
 
 const mocks = vi.hoisted(() => ({
     loadDiagram: vi.fn((_xml: string, ..._rest: unknown[]) => null),
+    // The canvas for a test that needs another one
+    canvas: null as string | null,
 }))
 
 vi.mock("@/contexts/diagram-context", () => ({
     useDiagram: () => ({
-        chartXML: BEFORE_FIRST_EDIT,
-        chartXMLRef: { current: AFTER_FIRST_EDIT },
+        chartXML: mocks.canvas ?? BEFORE_FIRST_EDIT,
+        chartXMLRef: { current: mocks.canvas ?? AFTER_FIRST_EDIT },
         loadDiagram: mocks.loadDiagram,
     }),
 }))
@@ -145,3 +147,81 @@ describe("the streaming preview of edit_diagram", () => {
         )
     })
 })
+
+describe("the streaming preview on the model's page", () => {
+    const box2 = (id: string, parent = "1") =>
+        `<mxCell id="${id}" value="${id}" vertex="1" parent="${parent}"><mxGeometry x="0" y="0" width="80" height="40" as="geometry"/></mxCell>`
+    // Page A has an edge "g"; on page B, "g" is a shape on layer "L"
+    const twoPages = `<mxfile><diagram id="pa" name="A"><mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/>${box2("s")}${box2("t")}<mxCell id="g" edge="1" parent="1" source="s" target="t"><mxGeometry relative="1" as="geometry"/></mxCell></root></mxGraphModel></diagram><diagram id="pb" name="B"><mxGraphModel><root><mxCell id="0"/><mxCell id="L" parent="0"/>${box2("g", "L")}</root></mxGraphModel></diagram></mxfile>`
+    const cellOn = (xml: string, pageId: string, id: string) =>
+        new DOMParser()
+            .parseFromString(xml, "text/xml")
+            .querySelector(`diagram[id="${pageId}"] mxCell[id="${id}"]`)
+    const stream = (part: object) =>
+        renderHook(() =>
+            useDiagramStreaming({
+                messages: [
+                    { id: "m1", role: "assistant", parts: [part] },
+                ] as any,
+                processedToolCallsRef: { current: new Set() },
+                editDiagramOriginalXmlRef: { current: new Map() },
+                loadedMessageIdsRef: { current: new Set() },
+                turnPageIdRef: { current: "pb" },
+            }),
+        )
+
+    it("applies an edit to that page, with its layer and its edges", () => {
+        mocks.loadDiagram.mockClear()
+        mocks.canvas = twoPages
+        try {
+            stream({
+                type: "tool-edit_diagram",
+                toolCallId: "edit-p",
+                state: "input-streaming",
+                input: {
+                    operations: [
+                        {
+                            operation: "add",
+                            cell_id: "c",
+                            new_xml:
+                                '<mxCell id="c" value="c" x="10" y="10" w="80" h="40"/>',
+                        },
+                        {
+                            operation: "add",
+                            cell_id: "d",
+                            new_xml:
+                                '<mxCell id="d" value="d" parent="g" x="10" y="10" w="80" h="40"/>',
+                        },
+                    ],
+                },
+            })
+            const loaded = mocks.loadDiagram.mock.calls.at(-1)?.[0] as string
+            expect(cellOn(loaded, "pb", "c")?.getAttribute("parent")).toBe("L")
+            expect(cellOn(loaded, "pb", "d")?.getAttribute("style")).toContain(
+                "whiteSpace=wrap",
+            )
+            expect(cellOn(loaded, "pa", "c")).toBeNull()
+        } finally {
+            mocks.canvas = null
+        }
+    })
+
+    it("draws display_diagram's cells on that page", () => {
+        mocks.loadDiagram.mockClear()
+        mocks.canvas = twoPages
+        try {
+            stream({
+                type: "tool-display_diagram",
+                toolCallId: "draw-p",
+                state: "input-streaming",
+                input: { xml: box2("n") },
+            })
+            const loaded = mocks.loadDiagram.mock.calls.at(-1)?.[0] as string
+            expect(cellOn(loaded, "pb", "n")).not.toBeNull()
+            expect(cellOn(loaded, "pb", "g")).toBeNull()
+            expect(cellOn(loaded, "pa", "g")).not.toBeNull()
+        } finally {
+            mocks.canvas = null
+        }
+    })
+})

tests/unit/use-diagram-tool-handlers.test.tsx

@@ -276,3 +276,102 @@ describe("a cut off drawing", () => {
         expect(refs.continuationOriginalRef.current).toBeNull()
     })
 })
+
+describe("the model's page", () => {
+    const page = (id: string, cells: string) =>
+        `<diagram id="${id}" name="Page ${id}"><mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/>${cells}</root></mxGraphModel></diagram>`
+    const twoPages = `<mxfile>${page("a", box("A"))}${page("b", box("B"))}</mxfile>`
+    /** Ids of the shapes on a page (not the root cells) */
+    const shapesOn = (xml: string, pageId: string) =>
+        Array.from(
+            new DOMParser()
+                .parseFromString(xml, "text/xml")
+                .querySelectorAll(`diagram[id="${pageId}"] mxCell`),
+        )
+            .map((cell) => cell.getAttribute("id"))
+            .filter((id) => id !== "0" && id !== "1")
+
+    function setupOn(pageId: string | null) {
+        const onDisplayChart = vi.fn(
+            (
+                _xml: string,
+                _skipValidation?: boolean,
+                _mode?: string,
+                _meta?: object,
+            ): string | null => null,
+        )
+        const { result } = renderHook(() =>
+            useDiagramToolHandlers({
+                partialXmlRef: { current: "" },
+                continuationOriginalRef: { current: null },
+                editDiagramOriginalXmlRef: { current: new Map() },
+                processedToolCallsRef: { current: new Set() },
+                validationRetryCountRef: { current: 0 },
+                chartXMLRef: { current: twoPages },
+                turnPageIdRef: { current: pageId },
+                onDisplayChart,
+                onFetchChart: async () => twoPages,
+                enableVlmValidation: false,
+            }),
+        )
+        const addToolOutput = vi.fn()
+        const call = (toolName: string, input: object) =>
+            result.current.handleToolCall(
+                { toolCall: { toolCallId: "call-1", toolName, input } },
+                addToolOutput,
+            )
+        return { onDisplayChart, addToolOutput, call }
+    }
+
+    it("display_diagram draws on that page and keeps the others", async () => {
+        const { onDisplayChart, call } = setupOn("b")
+        await call("display_diagram", { xml: box("N") })
+        const [xml, , mode, meta] = onDisplayChart.mock.calls[0]
+        expect(mode).toBe("commit")
+        expect(meta).toEqual({ toolCallId: "call-1", pageId: "b" })
+        expect(shapesOn(xml, "a")).toEqual(["A"])
+        expect(shapesOn(xml, "b")).toEqual(["N"])
+        // The page keeps its own name
+        expect(xml).toContain('name="Page b"')
+    })
+
+    it("edit_diagram changes that page", async () => {
+        const { onDisplayChart, addToolOutput, call } = setupOn("b")
+        await call("edit_diagram", {
+            operations: [{ operation: "delete", cell_id: "B" }],
+        })
+        expect(addToolOutput.mock.lastCall?.[0].state).toBeUndefined()
+        const [xml] = onDisplayChart.mock.calls.at(-1) as [string]
+        expect(shapesOn(xml, "a")).toEqual(["A"])
+        expect(shapesOn(xml, "b")).toEqual([])
+    })
+
+    it("shows the model only its page when an edit fails", async () => {
+        const { addToolOutput, call } = setupOn("b")
+        await call("edit_diagram", {
+            operations: [{ operation: "delete", cell_id: "missing" }],
+        })
+        const output = addToolOutput.mock.lastCall?.[0]
+        expect(output?.state).toBe("output-error")
+        expect(output?.errorText).toContain('id="B"')
+        expect(output?.errorText).not.toContain('id="A"')
+    })
+
+    it("draws on the first page when the page is not known", async () => {
+        const { onDisplayChart, call } = setupOn(null)
+        await call("display_diagram", { xml: box("N") })
+        const [xml, , , meta] = onDisplayChart.mock.calls[0]
+        expect(meta).toEqual({ toolCallId: "call-1", pageId: null })
+        expect(shapesOn(xml, "a")).toEqual(["N"])
+        expect(shapesOn(xml, "b")).toEqual(["B"])
+    })
+
+    it("lets a drawing with several pages replace the document", async () => {
+        const { onDisplayChart, call } = setupOn("b")
+        const drawn = `<mxfile>${page("x", box("X"))}${page("y", box("Y"))}</mxfile>`
+        await call("display_diagram", { xml: drawn })
+        const [xml] = onDisplayChart.mock.calls[0]
+        expect(shapesOn(xml, "x")).toEqual(["X"])
+        expect(shapesOn(xml, "a")).toEqual([])
+    })
+})

tests/unit/utils.test.ts

@@ -1,9 +1,12 @@
+import pako from "pako"
 import { describe, expect, it } from "vitest"
 import {
     cn,
     extractCompleteMxCells,
+    extractDiagramXML,
     isMxCellXmlComplete,
     isRealDiagram,
+    replaceNodes,
 } from "@/lib/utils"
 import { BLANK_MXFILE } from "@/packages/mcp-server/src/pages.ts"
 
@@ -134,3 +137,63 @@ describe("extractCompleteMxCells", () => {
         )
     })
 })
+
+const pageBox = (id: string) => `<mxCell id="${id}" vertex="1" parent="1"/>`
+const pageModel = (cells: string) =>
+    `<mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/>${cells}</root></mxGraphModel>`
+const pageOf = (id: string, cells: string) =>
+    `<diagram id="${id}" name="${id}">${pageModel(cells)}</diagram>`
+/** Ids of the shapes on a page (not the root cells) */
+const shapesOn = (xml: string, pageId: string) =>
+    Array.from(
+        new DOMParser()
+            .parseFromString(xml, "text/xml")
+            .querySelectorAll(`diagram[id="${pageId}"] mxCell`),
+    )
+        .map((cell) => cell.getAttribute("id"))
+        .filter((id) => id !== "0" && id !== "1")
+
+describe("replaceNodes", () => {
+    const two = `<mxfile>${pageOf("a", pageBox("A"))}${pageOf("b", pageBox("B"))}</mxfile>`
+
+    it("replaces the first page's cells by default", () => {
+        const out = replaceNodes(two, pageBox("N"))
+        expect(shapesOn(out, "a")).toEqual(["N"])
+        expect(shapesOn(out, "b")).toEqual(["B"])
+    })
+
+    it("replaces the given page's cells and keeps the others", () => {
+        const out = replaceNodes(two, pageBox("N"), "b")
+        expect(shapesOn(out, "a")).toEqual(["A"])
+        expect(shapesOn(out, "b")).toEqual(["N"])
+    })
+
+    it("uses the first page when the given page is not there", () => {
+        const out = replaceNodes(two, pageBox("N"), "zzz")
+        expect(shapesOn(out, "a")).toEqual(["N"])
+        expect(shapesOn(out, "b")).toEqual(["B"])
+    })
+})
+
+describe("extractDiagramXML", () => {
+    const packed = (model: string) =>
+        Buffer.from(pako.deflateRaw(encodeURIComponent(model))).toString(
+            "base64",
+        )
+    const file = `<mxfile><diagram id="a" name="a">${packed(pageModel(pageBox("A")))}</diagram><diagram id="b" name="b">${packed(pageModel(pageBox("B")))}</diagram></mxfile>`
+    const escapeAttr = (s: string) =>
+        s
+            .replace(/&/g, "&amp;")
+            .replace(/"/g, "&quot;")
+            .replace(/</g, "&lt;")
+            .replace(/>/g, "&gt;")
+    const svg = `<svg xmlns="http://www.w3.org/2000/svg" content="${escapeAttr(file)}"></svg>`
+    const dataUrl = `data:image/svg+xml;base64,${Buffer.from(svg).toString("base64")}`
+
+    it("returns the first page by default and a page by its id", () => {
+        expect(extractDiagramXML(dataUrl)).toContain('id="A"')
+        expect(extractDiagramXML(dataUrl, "b")).toContain('id="B"')
+        expect(extractDiagramXML(dataUrl, "b")).not.toContain('id="A"')
+        expect(extractDiagramXML(dataUrl, "zzz")).toContain('id="A"')
+    })
+})