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

feat(diagram): named styles and compact cells for shorter model output (#969)

Browse Source 
The model defines reused styles once as <mxStyle name="..." value="..."/> and
refers to them by name, and writes shapes as one self-closing mxCell with
x, y, w, h and edges with source and target. style-classes.ts and
compact-cells.ts expand both back into standard draw.io XML after validation
and add draw.io's html=1 / whiteSpace=wrap defaults; the diagram shown to the
model is folded into the same notation. Prompts, tool descriptions and the MCP
drawing guide teach the notation with shared examples.

Measured on the five start-screen examples against main: gpt-6-luna 35% fewer
output tokens, Claude Opus 5.5 41% fewer. Two five-model review rounds fixed
quote-aware cell matching, attribute escaping, edge/vertex inference, root-id
and layer handling, and several auto-fix gaps.
This commit is contained in:
Dayuan Jiang
2026-10-10 21:58:21 +09:00
committed by GitHub
parent 498ee628f1
commit cd5352ca88
16 changed files with 1720 additions and 123 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/unit/compact-cells.test.ts
+370
View File
@@ -0,0 +1,370 @@
import { describe, expect, it } from "vitest"
import { extractCompleteMxCells, isMxCellXmlComplete } from "@/lib/utils"
import {
expandCompactCells,
foldCells,
} from "@/packages/mcp-server/src/compact-cells.ts"
import { editDiagram } from "@/packages/mcp-server/src/edit-diagram.ts"
import { prepareNewDiagram } from "@/packages/mcp-server/src/new-diagram.ts"
import {
normalizeToMxfile,
wrapCellsInModel,
} from "@/packages/mcp-server/src/pages.ts"
const shape =
'<mxCell id="2" value="Start" style="rounded=1;" x="40" y="40" w="120" h="60"/>'
const shapeLong =
'<mxCell id="2" value="Start" style="rounded=1;" vertex="1" parent="1"><mxGeometry x="40" y="40" width="120" height="60" as="geometry"/></mxCell>'
const edge = '<mxCell id="5" style="down;" source="2" target="3"/>'
const edgeLong =
'<mxCell id="5" style="down;" edge="1" parent="1" source="2" target="3"><mxGeometry relative="1" as="geometry"/></mxCell>'
/** Cells compared by what they mean, not by attribute order */
function canonical(xml: string): string {
const doc = new DOMParser().parseFromString(
`<root>${xml}</root>`,
"text/xml",
)
const describe = (el: Element): string => {
const attrs = Array.from(el.attributes)
.map((a) => `${a.name}=${a.value}`)
.sort()
.join(" ")
const children = Array.from(el.children).map(describe).join("")
return `<${el.tagName} ${attrs}>${children}</${el.tagName}>`
}
return Array.from(doc.documentElement.children).map(describe).join("\n")
}
describe("expandCompactCells", () => {
it("gives a compact shape its flags, parent and geometry", () => {
expect(expandCompactCells(shape)).toBe(shapeLong)
})
it("gives a compact edge its flags, parent and relative geometry", () => {
expect(expandCompactCells(edge)).toBe(edgeLong)
})
it("keeps an explicit parent and accepts width/height as names", () => {
const inLane =
'<mxCell id="3" value="Step" style="rounded=1;" parent="lane1" x="20" y="60" width="160" height="40"/>'
expect(expandCompactCells(inLane)).toBe(
'<mxCell id="3" value="Step" style="rounded=1;" vertex="1" parent="lane1"><mxGeometry x="20" y="60" width="160" height="40" as="geometry"/></mxCell>',
)
})
it("fills in a default position and size when some are missing", () => {
expect(expandCompactCells('<mxCell id="2" value="A" w="100"/>')).toBe(
'<mxCell id="2" value="A" vertex="1" parent="1"><mxGeometry x="0" y="0" width="100" height="60" as="geometry"/></mxCell>',
)
})
it("lets an explicit geometry win over compact attributes", () => {
// The stray attributes stay as written; the geometry is the one used
const both =
'<mxCell id="2" x="1" y="1" w="1" h="1"><mxGeometry x="40" y="40" width="120" height="60" as="geometry"/></mxCell>'
expect(expandCompactCells(both)).toBe(
'<mxCell id="2" x="1" y="1" w="1" h="1" vertex="1" parent="1"><mxGeometry x="40" y="40" width="120" height="60" as="geometry"/></mxCell>',
)
})
it("adds the relative geometry to a long-form edge that has none", () => {
const noGeometry =
'<mxCell id="5" style="down;" edge="1" parent="1" source="2" target="3"/>'
expect(expandCompactCells(noGeometry)).toBe(edgeLong)
})
it("leaves root cells, long-form cells and edge labels alone", () => {
const roots = '<mxCell id="0"/><mxCell id="1" parent="0"/>'
const label =
'<mxCell id="9" value="yes" style="edgeLabel;" vertex="1" connectable="0" parent="5"><mxGeometry x="-0.5" relative="1" as="geometry"/></mxCell>'
expect(expandCompactCells(roots + shapeLong + edgeLong + label)).toBe(
roots + shapeLong + edgeLong + label,
)
})
it("works on a cell inside a UserObject wrapper", () => {
const wrapped =
'<UserObject id="4" label="Docs" link="https://example.com"><mxCell style="rounded=1;" x="580" y="40" w="120" h="60"/></UserObject>'
expect(expandCompactCells(wrapped)).toBe(
'<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>',
)
})
})
describe("foldCells", () => {
it("writes plain shapes and edges compactly", () => {
expect(foldCells(shapeLong)).toBe(shape)
expect(foldCells(edgeLong)).toBe(edge)
})
it("keeps a parent other than 1 and leaves special cells as written", () => {
const inLane =
'<mxCell id="3" value="Step" style="rounded=1;" vertex="1" parent="lane1"><mxGeometry x="20" y="60" width="160" height="40" as="geometry"/></mxCell>'
expect(foldCells(inLane)).toBe(
'<mxCell id="3" value="Step" style="rounded=1;" parent="lane1" x="20" y="60" w="160" h="40"/>',
)
const waypoints =
'<mxCell id="6" edge="1" parent="1" source="2" target="3"><mxGeometry relative="1" as="geometry"><Array as="points"><mxPoint x="750" y="80"/></Array></mxGeometry></mxCell>'
const label =
'<mxCell id="9" value="yes" style="edgeLabel;" vertex="1" connectable="0" parent="5"><mxGeometry x="-0.5" relative="1" as="geometry"/></mxCell>'
const sourcePoint =
'<mxCell id="7" edge="1" parent="1" target="3"><mxGeometry relative="1" as="geometry"><mxPoint x="10" y="10" as="sourcePoint"/></mxGeometry></mxCell>'
const roots = '<mxCell id="0"/><mxCell id="1" parent="0"/>'
const wrongAs =
'<mxCell id="8" vertex="1" parent="1"><mxGeometry x="1" y="1" width="2" height="2" as="33"/></mxCell>'
for (const xml of [waypoints, label, sourcePoint, roots, wrongAs]) {
expect(foldCells(xml)).toBe(xml)
}
})
it("round-trips through expandCompactCells", () => {
const page = `<mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/>${shapeLong}${edgeLong}<mxCell id="lane1" value="Lane" style="swimlane;" vertex="1" parent="1"><mxGeometry x="40" y="40" width="200" height="200" as="geometry"/></mxCell><mxCell id="3" value="Step" style="rounded=1;" vertex="1" parent="lane1"><mxGeometry x="20" y="60" width="160" height="40" as="geometry"/></mxCell><mxCell id="6" edge="1" parent="1" source="2" target="3"><mxGeometry relative="1" as="geometry"><Array as="points"><mxPoint x="750" y="80"/></Array></mxGeometry></mxCell></root></mxGraphModel>`
const folded = foldCells(page)
expect(folded).not.toContain("vertex=")
expect(canonical(expandCompactCells(folded))).toBe(canonical(page))
})
})
describe("compact cells through the pipeline", () => {
const geometryOf = (xml: string, id: string) =>
xml.match(
new RegExp(`<mxCell id="${id}"[^>]*>(<mxGeometry[^>]*/>)`),
)?.[1]
it("prepareNewDiagram expands them and still adds the defaults", () => {
const prepared = prepareNewDiagram(`${shape}\n${edge}`)
expect(prepared.ok).toBe(true)
if (!prepared.ok) return
expect(prepared.xml).toContain(
'style="rounded=1;whiteSpace=wrap;html=1;" vertex="1" parent="1"',
)
expect(geometryOf(prepared.xml, "2")).toBe(
'<mxGeometry x="40" y="40" width="120" height="60" as="geometry"/>',
)
expect(prepared.xml).toContain(
'<mxCell id="5" style="down;html=1;" edge="1" parent="1" source="2" target="3"><mxGeometry relative="1" as="geometry"/></mxCell>',
)
})
it("editDiagram accepts a compact cell in new_xml", () => {
const file =
normalizeToMxfile(wrapCellsInModel(shapeLong), {
pageId: "p1",
pageName: "Page-1",
}) ?? ""
const outcome = editDiagram(
file,
[
{
operation: "add",
cell_id: "3",
new_xml:
'<mxCell id="3" value="Next" style="rounded=1;" x="200" y="40" w="120" h="60"/>',
},
{
operation: "update",
cell_id: "2",
new_xml:
'<mxCell id="2" value="Begin" style="rounded=1;" x="40" y="40" w="120" h="60"/>',
},
],
{},
)
expect(outcome.ok).toBe(true)
if (!outcome.ok) return
expect(outcome.xml).toContain('value="Next"')
expect(outcome.xml).toContain('value="Begin"')
expect(geometryOf(outcome.xml, "3")).toBe(
'<mxGeometry x="200" y="40" width="120" height="60" as="geometry"/>',
)
})
})
describe("a compact cell with a slip", () => {
it("repairs a missing closing quote before the slash", () => {
const prepared = prepareNewDiagram(
'<mxCell id="12" value="Coin Flip" style="rounded=1;" x="535" y="276" w="225" h="112/>',
)
expect(prepared.ok).toBe(true)
if (!prepared.ok) return
expect(prepared.xml).toContain(
'<mxGeometry x="535" y="276" width="225" height="112" as="geometry"/>',
)
})
})
describe("review round: cells the model may write", () => {
it("escapes a double quote from a single-quoted value", () => {
const out = expandCompactCells(
`<mxCell id='2' value='Say "hi"' x='1' y='2' w='3' h='4'/>`,
)
expect(out).toContain('value="Say &quot;hi&quot;"')
const prepared = prepareNewDiagram(
`<mxCell id='2' value='a" visible="0' x="1" y="2" w="3" h="4"/>`,
)
expect(prepared.ok).toBe(true)
if (!prepared.ok) return
expect(prepared.xml).toContain('value="a&quot; visible=&quot;0"')
expect(prepared.xml).not.toContain(' visible="0"')
const folded = foldCells(
`<mxCell id="2" value='Say "hi"' vertex="1" parent="1"><mxGeometry x="1" y="2" width="3" height="4" as="geometry"/></mxCell>`,
)
expect(folded).toBe(
'<mxCell id="2" value="Say &quot;hi&quot;" x="1" y="2" w="3" h="4"/>',
)
})
it("accepts labels with > in compact self-closing cells", () => {
const prepared = prepareNewDiagram(
`<mxCell id="2" value="x > 5" style="rounded=1;" x="40" y="40" w="120" h="60"/>
<mxCell id="3" value="A -> B" style="rounded=1;" x="40" y="200" w="120" h="60"/>
<mxCell id="5" source="2" target="3"/>`,
)
expect(prepared.ok).toBe(true)
expect(
isMxCellXmlComplete(
'<mxCell id="2" value="A > B" x="0" y="0" w="120" h="60"/>',
),
).toBe(true)
expect(
extractCompleteMxCells(
'<mxCell id="2" value="A > B" x="0" y="0" w="120" h="60"/>\n<mxCell id="3" value="C" x="0" y="0" w="1',
),
).toBe('<mxCell id="2" value="A > B" x="0" y="0" w="120" h="60"/>')
})
it("infers a vertex for an edge label with a geometry and an edge for a connection with a size", () => {
const label =
'<mxCell id="9" value="yes" style="edgeLabel;" parent="5" connectable="0"><mxGeometry x="-0.5" relative="1" as="geometry"/></mxCell>'
expect(expandCompactCells(label)).toBe(
'<mxCell id="9" value="yes" style="edgeLabel;" connectable="0" vertex="1" parent="5"><mxGeometry x="-0.5" relative="1" as="geometry"/></mxCell>',
)
expect(
expandCompactCells(
'<mxCell id="e" source="2" target="3" x="0" y="0"/>',
),
).toBe(
'<mxCell id="e" edge="1" parent="1" source="2" target="3"><mxGeometry relative="1" as="geometry"/></mxCell>',
)
// A dangling edge has only one end
expect(
expandCompactCells('<mxCell id="e" style="a;" source="2"/>'),
).toBe(
'<mxCell id="e" style="a;" edge="1" parent="1" source="2"><mxGeometry relative="1" as="geometry"/></mxCell>',
)
})
it("does not take a geometry inside custom data for the cell's own", () => {
const out = expandCompactCells(
'<mxCell id="2" x="10" y="20" w="120" h="60"><Object as="payload"><mxGeometry as="backup"/></Object></mxCell>',
)
expect(out).toBe(
'<mxCell id="2" vertex="1" parent="1"><mxGeometry x="10" y="20" width="120" height="60" as="geometry"/><Object as="payload"><mxGeometry as="backup"/></Object></mxCell>',
)
})
it("keeps stray size attributes when the cell has its own geometry", () => {
const xml =
'<mxCell id="2" x="99" vertex="1" parent="1"><mxGeometry x="1" y="2" width="3" height="4" as="geometry"/></mxCell>'
expect(expandCompactCells(xml)).toBe(xml)
expect(foldCells(xml)).toBe(xml)
})
it("ignores prototype names and folds only cells that name their parent", () => {
expect(expandCompactCells('<mxCell id="2" toString="x"/>')).toBe(
'<mxCell id="2" toString="x"/>',
)
const noParent =
'<mxCell id="2" vertex="1"><mxGeometry x="1" y="2" width="3" height="4" as="geometry"/></mxCell>'
expect(foldCells(noParent)).toBe(noParent)
})
it("puts a compact cell on the page's first layer when editing", () => {
const file = `<mxfile><diagram id="p1" name="Page-1"><mxGraphModel><root><mxCell id="0"/><mxCell id="L1" parent="0"/>${shapeLong.replace('parent="1"', 'parent="L1"')}</root></mxGraphModel></diagram></mxfile>`
const outcome = editDiagram(
file,
[
{
operation: "add",
cell_id: "3",
new_xml:
'<mxCell id="3" value="B" x="1" y="2" w="3" h="4"/>',
},
],
{},
)
expect(outcome.ok).toBe(true)
if (!outcome.ok) return
expect(outcome.xml).toContain(
'<mxCell id="3" value="B" vertex="1" parent="L1"',
)
expect(foldCells(file, "L1")).toContain(
'<mxCell id="2" value="Start" style="rounded=1;" x="40" y="40" w="120" h="60"/>',
)
})
it("refuses a compact shape with a root cell id instead of dropping it", () => {
const prepared = prepareNewDiagram(
'<mxCell id="1" value="A" x="0" y="0" w="10" h="10"/>',
)
expect(prepared.ok).toBe(false)
if (!prepared.ok) expect(prepared.error).toContain("root cells")
})
it("repairs the quote slip only on a cell's last numeric attribute", () => {
expect(
isMxCellXmlComplete('<mxCell id="2" x="1" y="2" w="3" h="4/>'),
).toBe(true)
const prepared = prepareNewDiagram(
'<mxCell id="2" value="5/>" x="1" y="2" w="3" h="4"/>',
)
expect(prepared.ok).toBe(true)
if (!prepared.ok) return
expect(prepared.xml).toContain('value="5/>"')
expect(prepared.xml).toContain(
'<mxGeometry x="1" y="2" width="3" height="4" as="geometry"/>',
)
})
it("closes a compact cell written without the slash", () => {
const prepared = prepareNewDiagram(
'<mxCell id="2" value="A" x="1" y="2" w="3" h="4">\n<mxCell id="3" value="B" x="1" y="2" w="3" h="4"/>',
)
expect(prepared.ok).toBe(true)
if (!prepared.ok) return
expect(prepared.xml).toContain('value="A"')
expect(prepared.xml).toContain('value="B"')
expect(prepared.fixes.join()).toContain("without the slash")
})
})
describe("a long-form cell missing its closing tag next to compact cells", () => {
it("is closed before the compact cell that follows", () => {
const prepared = prepareNewDiagram(
`<mxCell id="2" value="A" x="0" y="0" w="80" h="40"/>
<mxCell id="4" edge="1" parent="1" source="2" target="5">
<mxGeometry relative="1" as="geometry"><Array as="points"><mxPoint x="300" y="150"/></Array></mxGeometry>
<mxCell id="5" value="C" x="400" y="0" w="80" h="40"/>`,
)
expect(prepared.ok).toBe(true)
if (!prepared.ok) return
expect(prepared.xml).toContain('<mxPoint x="300" y="150"/>')
expect(prepared.xml).toContain('value="C"')
expect((prepared.xml.match(/<mxCell\b/g) || []).length).toBe(5)
})
it("still flattens a compact cell nested inside an open cell", () => {
const prepared = prepareNewDiagram(
`<mxCell id="2" value="A" vertex="1" parent="1">
<mxGeometry x="0" y="0" width="80" height="40" as="geometry"/>
<mxCell id="3" value="B" x="100" y="0" w="80" h="40"/>
</mxCell>`,
)
expect(prepared.ok).toBe(true)
if (!prepared.ok) return
expect(prepared.xml).toContain('value="A"')
expect(prepared.xml).toContain('value="B"')
expect((prepared.xml.match(/<mxCell\b/g) || []).length).toBe(4)
})
})
tests/unit/style-classes.test.ts
+499
View File
@@ -0,0 +1,499 @@
import { describe, expect, it } from "vitest"
import { isMxCellXmlComplete } from "@/lib/utils"
import { editDiagram } from "@/packages/mcp-server/src/edit-diagram.ts"
import { prepareNewDiagram } from "@/packages/mcp-server/src/new-diagram.ts"
import {
normalizeToMxfile,
wrapCellsInModel,
} from "@/packages/mcp-server/src/pages.ts"
import {
addDefaultStyles,
applyStyleClasses,
expandStyles,
readStyleClasses,
} from "@/packages/mcp-server/src/style-classes.ts"
const geometry =
'<mxGeometry x="0" y="0" width="80" height="40" as="geometry"/>'
const cell = (id: string, style: string, extra = "") =>
`<mxCell id="${id}" value="${id}" style="${style}" vertex="1" parent="1"${extra}>${geometry}</mxCell>`
const edge = (id: string, style: string) =>
`<mxCell id="${id}" style="${style}" edge="1" parent="1" source="2" target="3"><mxGeometry relative="1" as="geometry"/></mxCell>`
const BLUE =
'<mxStyle name="blue" value="fillColor=#dae8fc;strokeColor=#6c8ebf;"/>'
describe("readStyleClasses", () => {
it("reads the definitions and removes them from the XML", () => {
const xml = `${BLUE}\n<mxStyle name="flow" value="edgeStyle=orthogonalEdgeStyle;"></mxStyle>\n${cell("2", "blue;")}`
const { classes, xml: rest } = readStyleClasses(xml)
expect(classes.get("blue")).toBe(
"fillColor=#dae8fc;strokeColor=#6c8ebf",
)
expect(classes.get("flow")).toBe("edgeStyle=orthogonalEdgeStyle")
expect(rest).toBe(cell("2", "blue;"))
})
it("resolves a definition that uses an earlier name", () => {
const xml = `${BLUE}<mxStyle name="pill" value="rounded=1;blue;arcSize=50;"/>`
const { classes } = readStyleClasses(xml)
expect(classes.get("pill")).toBe(
"rounded=1;fillColor=#dae8fc;strokeColor=#6c8ebf;arcSize=50",
)
})
it("leaves a definition that is still streaming in place", () => {
const xml = `${BLUE}<mxStyle name="flow" value="edgeSt`
const { classes, xml: rest } = readStyleClasses(xml)
expect(classes.size).toBe(1)
expect(rest).toBe('<mxStyle name="flow" value="edgeSt')
})
})
describe("applyStyleClasses", () => {
const classes = readStyleClasses(BLUE).classes
it("expands a name in place, so the overrides after it still win", () => {
const out = applyStyleClasses(
cell("2", "rhombus;blue;fontSize=14;"),
classes,
)
expect(out).toBe(
cell(
"2",
"rhombus;fillColor=#dae8fc;strokeColor=#6c8ebf;fontSize=14;",
),
)
})
it("leaves unknown names, key=value tokens and other attributes alone", () => {
const xml = `${cell("2", "text;fontSize=blue;")}<mxCell id="3" vertex="1" parent="1" value="blue">${geometry}</mxCell>`
expect(applyStyleClasses(xml, classes)).toBe(xml)
})
it("works inside a wrapped model and keeps the attribute order", () => {
const xml = `<mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/><mxCell id="2" style="blue;" value="A" vertex="1" parent="1">${geometry}</mxCell></root></mxGraphModel>`
expect(applyStyleClasses(xml, classes)).toBe(
xml.replace(
'style="blue;"',
'style="fillColor=#dae8fc;strokeColor=#6c8ebf;"',
),
)
})
})
describe("addDefaultStyles", () => {
it("adds whiteSpace=wrap and html=1 to a shape", () => {
expect(addDefaultStyles(cell("2", "rounded=1;"))).toBe(
cell("2", "rounded=1;whiteSpace=wrap;html=1;"),
)
})
it("gives a shape without a style attribute one", () => {
const xml = `<mxCell id="2" value="A" vertex="1" parent="1">${geometry}</mxCell>`
expect(addDefaultStyles(xml)).toBe(
`<mxCell id="2" value="A" vertex="1" parent="1" style="whiteSpace=wrap;html=1;">${geometry}</mxCell>`,
)
})
it("keeps html=0 and an existing whiteSpace", () => {
const xml = cell("2", "html=0;whiteSpace=nowrap;")
expect(addDefaultStyles(xml)).toBe(xml)
})
it("adds only html=1 to icon shapes, images and edge labels", () => {
const icon = cell(
"2",
"shape=mxgraph.aws4.resourceIcon;resIcon=mxgraph.aws4.ec2;verticalLabelPosition=bottom;",
)
const image = cell("3", "image=data:image/png,abc;")
const label = cell("4", "edgeLabel;align=center;")
const unconnectable = cell("5", "align=center;", ' connectable="0"')
expect(addDefaultStyles(icon)).toBe(
cell(
"2",
"shape=mxgraph.aws4.resourceIcon;resIcon=mxgraph.aws4.ec2;verticalLabelPosition=bottom;html=1;",
),
)
expect(addDefaultStyles(image)).toBe(
cell("3", "image=data:image/png,abc;html=1;"),
)
expect(addDefaultStyles(label)).toBe(
cell("4", "edgeLabel;align=center;html=1;"),
)
expect(addDefaultStyles(unconnectable)).toBe(
cell("5", "align=center;html=1;", ' connectable="0"'),
)
})
it("adds html=1 to edges and leaves the root cells alone", () => {
const roots = '<mxCell id="0"/><mxCell id="1" parent="0"/>'
expect(addDefaultStyles(roots + edge("4", "endArrow=block;"))).toBe(
roots + edge("4", "endArrow=block;html=1;"),
)
})
})
describe("expandStyles", () => {
it("expands the names, then adds the defaults", () => {
const xml = `${BLUE}${cell("2", "blue;")}`
expect(expandStyles(xml)).toBe(
cell(
"2",
"fillColor=#dae8fc;strokeColor=#6c8ebf;whiteSpace=wrap;html=1;",
),
)
})
})
describe("prepareNewDiagram with named styles", () => {
it("saves plain draw.io XML without the definitions", () => {
const xml = `${BLUE}\n<mxStyle name="flow" value="edgeStyle=orthogonalEdgeStyle;exitX=0.5;exitY=1;entryX=0.5;entryY=0;"/>\n${cell("2", "blue;")}\n${cell("3", "rhombus;blue;")}\n${edge("4", "flow;")}`
const prepared = prepareNewDiagram(xml, { pageId: "p1" })
expect(prepared.ok).toBe(true)
if (!prepared.ok) return
expect(prepared.xml).not.toContain("mxStyle")
expect(prepared.xml).toContain(
'style="fillColor=#dae8fc;strokeColor=#6c8ebf;whiteSpace=wrap;html=1;"',
)
expect(prepared.xml).toContain(
'style="rhombus;fillColor=#dae8fc;strokeColor=#6c8ebf;whiteSpace=wrap;html=1;"',
)
expect(prepared.xml).toContain(
'style="edgeStyle=orthogonalEdgeStyle;exitX=0.5;exitY=1;entryX=0.5;entryY=0;html=1;"',
)
})
})
describe("editDiagram adds the default styles", () => {
it("adds html=1 and whiteSpace=wrap to an added cell", () => {
const file =
normalizeToMxfile(wrapCellsInModel(cell("2", "")), {
pageId: "p1",
pageName: "Page-1",
}) ?? ""
const outcome = editDiagram(
file,
[
{
operation: "add",
cell_id: "3",
new_xml: cell("3", "rounded=1;"),
},
],
{},
)
expect(outcome.ok).toBe(true)
if (!outcome.ok) return
expect(outcome.xml).toContain(
'style="rounded=1;whiteSpace=wrap;html=1;"',
)
})
})
describe("cells written in other shapes", () => {
it("keeps a self-closing cell without a style self-closing", () => {
const edge =
'<mxCell id="4" edge="1" parent="1" source="2" target="3"/>'
expect(addDefaultStyles(edge)).toBe(
'<mxCell id="4" edge="1" parent="1" source="2" target="3" style="html=1;"/>',
)
const shape = '<mxCell id="5" value="A" vertex="1" parent="1" />'
expect(addDefaultStyles(shape)).toBe(
'<mxCell id="5" value="A" vertex="1" parent="1" style="whiteSpace=wrap;html=1;"/>',
)
})
it("expands a name in a cell wrapped in a UserObject", () => {
const xml = `${BLUE}<UserObject id="4" label="Docs" link="https://example.com"><mxCell style="blue;" vertex="1" parent="1">${geometry}</mxCell></UserObject>`
expect(expandStyles(xml)).toBe(
`<UserObject id="4" label="Docs" link="https://example.com"><mxCell style="fillColor=#dae8fc;strokeColor=#6c8ebf;whiteSpace=wrap;html=1;" vertex="1" parent="1">${geometry}</mxCell></UserObject>`,
)
})
it("leaves text that sizes itself unwrapped", () => {
const xml = cell("2", "text;autosize=1;")
expect(addDefaultStyles(xml)).toBe(cell("2", "text;autosize=1;html=1;"))
})
it("reads a definition written without the closing slash", () => {
const { classes, xml } = readStyleClasses(
`<mxStyle name="a" value="x=1">\n${cell("2", "a;")}`,
)
expect(classes.get("a")).toBe("x=1")
expect(xml).toBe(cell("2", "a;"))
})
it("adds the defaults inside a full mxfile", () => {
const file = `<mxfile><diagram id="p" name="Page-1"><mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/>${cell("2", "rounded=1;")}</root></mxGraphModel></diagram></mxfile>`
const prepared = prepareNewDiagram(file)
expect(prepared.ok).toBe(true)
if (!prepared.ok) return
expect(prepared.xml).toContain(
'<mxCell id="0"/><mxCell id="1" parent="0"/>',
)
expect(prepared.xml).toContain(
'style="rounded=1;whiteSpace=wrap;html=1;"',
)
})
})
describe("truncation check with named styles", () => {
it("treats output cut off right after the definitions as incomplete", () => {
expect(isMxCellXmlComplete(BLUE)).toBe(false)
expect(
isMxCellXmlComplete(`${BLUE}<mxStyle name="flow" value="edgeSt`),
).toBe(false)
expect(isMxCellXmlComplete(`${BLUE}${cell("2", "blue;")}`)).toBe(true)
})
})
describe("styles are rewritten as written", () => {
it("keeps entities, a leading semicolon and a missing trailing one", () => {
const quoted = cell(
"2",
"fontFamily=&quot;Times New Roman&quot;;fillColor=red;",
)
expect(addDefaultStyles(quoted)).toBe(
cell(
"2",
"fontFamily=&quot;Times New Roman&quot;;fillColor=red;whiteSpace=wrap;html=1;",
),
)
const noDefaults = cell("3", ";shape=rectangle;strokeColor=#000000")
expect(addDefaultStyles(noDefaults)).toBe(
cell(
"3",
";shape=rectangle;strokeColor=#000000;whiteSpace=wrap;html=1;",
),
)
const classes = readStyleClasses(BLUE).classes
expect(applyStyleClasses(cell("4", "blue"), classes)).toBe(
cell("4", "fillColor=#dae8fc;strokeColor=#6c8ebf"),
)
})
it("does not add wrapping next to html=0, since wrapping turns HTML on", () => {
const xml = cell("2", "html=0;")
expect(addDefaultStyles(xml)).toBe(xml)
})
it("wraps ordinary shape= shapes and leaves library icons alone", () => {
expect(
addDefaultStyles(
cell(
"2",
"shape=parallelogram;perimeter=parallelogramPerimeter;",
),
),
).toBe(
cell(
"2",
"shape=parallelogram;perimeter=parallelogramPerimeter;whiteSpace=wrap;html=1;",
),
)
expect(
addDefaultStyles(
cell(
"3",
"shape=mxgraph.cisco19.router;verticalLabelPosition=bottom;",
),
),
).toBe(
cell(
"3",
"shape=mxgraph.cisco19.router;verticalLabelPosition=bottom;html=1;",
),
)
// Library shapes whose label sits inside wrap like any shape
expect(
addDefaultStyles(cell("5", "shape=mxgraph.flowchart.process;")),
).toBe(
cell(
"5",
"shape=mxgraph.flowchart.process;whiteSpace=wrap;html=1;",
),
)
expect(addDefaultStyles(cell("4", "shape=image;image=a.png;"))).toBe(
cell("4", "shape=image;image=a.png;html=1;"),
)
})
it("escapes a quote that a single-quoted definition brings in", () => {
const xml = `<mxStyle name="a" value='fontFamily="Arial";'/>${cell("2", "a;")}`
expect(expandStyles(xml)).toBe(
cell("2", "fontFamily=&quot;Arial&quot;;whiteSpace=wrap;html=1;"),
)
})
it("resolves references in either order; duplicates and built-in names", () => {
const { classes } = readStyleClasses(
`<mxStyle name="b" value="c;x=1;"/><mxStyle name="c" value="y=2;"/><mxStyle name="c" value="y=3;"/><mxStyle name="loop" value="loop;z=1;"/>`,
)
// A name defined later still resolves; the last definition wins
expect(classes.get("b")).toBe("y=3;x=1")
// A name using itself stays a bare token
expect(classes.get("loop")).toBe("loop;z=1")
// The last definition of a name wins
expect(classes.get("c")).toBe("y=3")
// A definition named like a built-in replaces that built-in's meaning
const shadow = readStyleClasses(
'<mxStyle name="ellipse" value="fillColor=red;"/>',
).classes
expect(applyStyleClasses(cell("2", "ellipse;"), shadow)).toBe(
cell("2", "fillColor=red;"),
)
})
})
describe("the auto-fix runs before the defaults", () => {
it("repairs quotes written as entities in an edit and then adds the defaults", () => {
const file = normalizeToMxfile(wrapCellsInModel(cell("2", "")), {
pageId: "p1",
pageName: "Page-1",
})
const outcome = editDiagram(
file ?? "",
[
{
operation: "add",
cell_id: "3",
new_xml: `<mxCell id="3" value="B" style=&quot;rounded=1;&quot; vertex="1" parent="1">${geometry}</mxCell>`,
},
],
{},
)
expect(outcome.ok).toBe(true)
if (!outcome.ok) return
expect(outcome.xml).toContain(
'style="rounded=1;whiteSpace=wrap;html=1;"',
)
})
it("repairs a new diagram the same way", () => {
const prepared = prepareNewDiagram(
`${BLUE}<mxCell id="2" value="A" style=&quot;blue;&quot; vertex="1" parent="1">${geometry}</mxCell>`,
)
expect(prepared.ok).toBe(true)
if (!prepared.ok) return
expect(prepared.xml).toContain(
'style="fillColor=#dae8fc;strokeColor=#6c8ebf;whiteSpace=wrap;html=1;"',
)
})
})
describe("definitions and cells the model may write imperfectly", () => {
it("escapes & and < in a definition, which the validator never sees", () => {
const xml = `<mxStyle name="f" value="fontFamily=A&B <C;"/>${cell("2", "f;")}`
const prepared = prepareNewDiagram(xml)
expect(prepared.ok).toBe(true)
if (!prepared.ok) return
expect(prepared.xml).toContain(
'style="fontFamily=A&amp;B &lt;C;whiteSpace=wrap;html=1;"',
)
})
it("accepts a lowercase tag, a style attribute and a name ending in ;", () => {
const { classes, xml } = readStyleClasses(
`<mxstyle name="blue;" style="fillColor=red;"/>${cell("2", "blue;")}`,
)
expect(classes.get("blue")).toBe("fillColor=red")
expect(xml).toBe(cell("2", "blue;"))
})
it("reads keys exactly, like draw.io, so a key after a space does not count", () => {
const xml = cell("2", "rounded=1; html=1; whiteSpace=wrap;")
expect(addDefaultStyles(xml)).toBe(
cell(
"2",
"rounded=1; html=1; whiteSpace=wrap;whiteSpace=wrap;html=1;",
),
)
const classes = readStyleClasses(BLUE).classes
expect(
applyStyleClasses(
cell("3", "rounded=1; blue; fontSize=14;"),
classes,
),
).toBe(
cell(
"3",
"rounded=1;fillColor=#dae8fc;strokeColor=#6c8ebf; fontSize=14;",
),
)
})
it("treats a vertex whose parent is an edge as an edge label", () => {
const xml = `${edge("e1", "")}<mxCell id="lbl" value="yes" style="align=center;" vertex="1" parent="e1"><mxGeometry relative="1" as="geometry"/></mxCell>`
expect(addDefaultStyles(xml)).toContain('style="align=center;html=1;"')
})
it("reports a definition that never closed and definitions without cells", () => {
const open = prepareNewDiagram(
`${BLUE}<mxStyle name="flow" value="edgeSt`,
)
expect(open.ok).toBe(false)
if (!open.ok) expect(open.error).toContain("not closed")
const only = prepareNewDiagram(BLUE)
expect(only.ok).toBe(false)
if (!only.ok) expect(only.error).toContain("no cells")
})
})
describe("edge labels and empty definitions", () => {
it("does not wrap a label added to an edge the XML does not show", () => {
const file =
normalizeToMxfile(
wrapCellsInModel(cell("2", "") + cell("3", "") + edge("4", "")),
{
pageId: "p1",
pageName: "Page-1",
},
) ?? ""
const label = `<mxCell id="9" value="yes" style="text;" vertex="1" parent="4"><mxGeometry x="-0.5" relative="1" as="geometry"/></mxCell>`
const outcome = editDiagram(
file,
[{ operation: "add", cell_id: "9", new_xml: label }],
{},
)
expect(outcome.ok).toBe(true)
if (!outcome.ok) return
expect(outcome.xml).toContain('id="9" value="yes" style="text;html=1;"')
// The same through the function the preview uses
expect(addDefaultStyles(label, ["4"])).toContain('style="text;html=1;"')
})
it("treats a cell marked edge and vertex as an edge", () => {
const both = `<mxCell id="5" style="rounded=1;" edge="1" vertex="1" parent="1" source="2" target="3"><mxGeometry relative="1" as="geometry"/></mxCell>`
expect(addDefaultStyles(both)).toContain('style="rounded=1;html=1;"')
})
it("drops a name whose definition is empty instead of leaving a leading ;", () => {
const xml = `<mxStyle name="plain" value=""/>${cell("2", "plain;rounded=1;")}`
expect(expandStyles(xml)).toBe(
cell("2", "rounded=1;whiteSpace=wrap;html=1;"),
)
})
it("tells the model that edit_diagram has no named styles", () => {
const file =
normalizeToMxfile(wrapCellsInModel(cell("2", "")), {
pageId: "p1",
pageName: "Page-1",
}) ?? ""
const outcome = editDiagram(
file,
[
{
operation: "add",
cell_id: "3",
new_xml: `${BLUE}${cell("3", "blue;")}`,
},
],
{},
)
expect(outcome.ok).toBe(false)
if (!outcome.ok)
expect(outcome.errors[0]).toContain("not available in edit_diagram")
})
})
tests/unit/use-diagram-streaming.test.tsx
+87 -1
View File
@@ -11,11 +11,15 @@ const box = (id: string) =>
const BEFORE_FIRST_EDIT = page(box("a"))
const AFTER_FIRST_EDIT = page(box("a") + box("b"))
const mocks = vi.hoisted(() => ({
loadDiagram: vi.fn((_xml: string, ..._rest: unknown[]) => null),
}))
vi.mock("@/contexts/diagram-context", () => ({
useDiagram: () => ({
chartXML: BEFORE_FIRST_EDIT,
chartXMLRef: { current: AFTER_FIRST_EDIT },
loadDiagram: vi.fn(() => null),
loadDiagram: mocks.loadDiagram,
}),
}))
@@ -59,3 +63,85 @@ describe("the streaming preview of a second edit", () => {
)
})
})
describe("the streaming preview of display_diagram", () => {
it("expands the named styles and adds the defaults to the cells so far", () => {
mocks.loadDiagram.mockClear()
// A complete definition and cell, then a cell still being written
const xml = `<mxStyle name="blue" value="fillColor=#dae8fc;strokeColor=#6c8ebf;"/>
<mxCell id="a" value="a" style="rounded=1;blue;" vertex="1" parent="1"><mxGeometry x="0" y="0" width="80" height="40" as="geometry"/></mxCell>
<mxCell id="b" value="b" style="blue;" vertex="1" parent="1"><mxGeometry x="0" y="0" wid`
const messages = [
{
id: "m1",
role: "assistant",
parts: [
{
type: "tool-display_diagram",
toolCallId: "draw-1",
state: "input-streaming",
input: { xml },
},
],
},
] as any
renderHook(() =>
useDiagramStreaming({
messages,
processedToolCallsRef: { current: new Set() },
editDiagramOriginalXmlRef: { current: new Map() },
loadedMessageIdsRef: { current: new Set() },
}),
)
expect(mocks.loadDiagram).toHaveBeenCalledTimes(1)
const loaded = mocks.loadDiagram.mock.calls[0][0]
expect(loaded).toContain(
'style="rounded=1;fillColor=#dae8fc;strokeColor=#6c8ebf;whiteSpace=wrap;html=1;"',
)
expect(loaded).not.toContain("mxStyle")
expect(loaded).not.toContain('id="b"')
})
})
describe("the streaming preview of edit_diagram", () => {
it("adds the default styles to the cell being added, like the final edit", () => {
mocks.loadDiagram.mockClear()
const messages = [
{
id: "m1",
role: "assistant",
parts: [
{
type: "tool-edit_diagram",
toolCallId: "edit-3",
state: "input-streaming",
input: {
operations: [
{
operation: "add",
cell_id: "c",
new_xml: box("c").replace(
'vertex="1"',
'style="rounded=1;" vertex="1"',
),
},
],
},
},
],
},
] as any
renderHook(() =>
useDiagramStreaming({
messages,
processedToolCallsRef: { current: new Set() },
editDiagramOriginalXmlRef: { current: new Map() },
loadedMessageIdsRef: { current: new Set() },
}),
)
expect(mocks.loadDiagram).toHaveBeenCalledTimes(1)
expect(mocks.loadDiagram.mock.calls[0][0]).toContain(
'style="rounded=1;whiteSpace=wrap;html=1;"',
)
})
})