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
2026-10-11 12:09:53 +08:00 · 2026-10-10 23:30:51 +09:00
parent cd5352ca88
commit d1ee1594f1
21 changed files with 907 additions and 120 deletions
@@ -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 ({
@@ -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 ({
@@ -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", () => {
@@ -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"] },
+        ])
+    })
+})
@@ -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",
+        )
+    })
+})
@@ -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
+        }
+    })
+})
@@ -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([])
+    })
+})
@@ -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"')
+    })
+})