fix(chat): the AI reads and draws the page the user is viewing (#974)

* fix(chat): the AI reads and draws the page the user is viewing

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.

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

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")