i/next-ai-draw-io
1
0
Fork 0
mirror of https://github.com/DayuanJiang/next-ai-draw-io.git synced 2026-10-11 03:59:58 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
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.
This commit is contained in:
dayuan.jiang
2026-10-10 22:14:50 +09:00
parent cd5352ca88
commit 2757509d4e
21 changed files with 907 additions and 120 deletions
Download Patch File Download Diff File Expand all files Collapse all files
components/canvas/canvas-stage.tsx
+1 -6
View File
@@ -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
+32 -17
View File
@@ -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
+3 -3
View File
@@ -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
+23 -5
View File
@@ -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
+3 -1
View File
@@ -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
+26 -5
View File
@@ -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
+37 -15
View File
@@ -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
+22 -21
View File
@@ -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
+138
View File
@@ -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
+97 -27
View File
@@ -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
+31 -9
View File
@@ -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
+1 -1
View File
@@ -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
+2
View File
@@ -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
+5 -3
View File
@@ -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
+59 -4
View File
@@ -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
+11 -1
View File
@@ -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
+125
View File
@@ -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
+47
View File
@@ -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
+82 -2
View File
@@ -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
+99
View File
@@ -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 { state, errorText } = addToolOutput.mock.lastCall?.[0]
expect(state).toBe("output-error")
expect(errorText).toContain('id="B"')
expect(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
+63
View File
@@ -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"')
})
})