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

feat(diagram-engine): wire up restructure_diagram + stencil catalog

Browse Source 
Closes the loop: the model can now build and edit AWS architecture diagrams by
declaring structure, and never writes an mxCell again.

catalog.ts — 983 AWS icon and 19 group stencils as a name→style map, generated from
drawio-ai-kit's catalog (itself generated from jgraph's draw.io shape index). Styles
are verbatim, so the official category colours, connection points and aspect=fixed
come along for free and nothing is hand-assembled. An invented name is rejected with
suggestions instead of rendering as a blank square, which is what draw.io does with
an unknown resIcon today.

operations.ts — what the model actually sends: add_icon / add_container / move /
link / set_dir and so on, applied in order against the tree. Guards the things that
break a diagram quietly: duplicate ids (draw.io drops one of the two cells), edges
left pointing at a removed node, and moving a container inside itself.

index.ts — the entry point. current XML → parse → apply ops → check names → layout →
render → new XML. The tree is not stored between calls; it is re-derived from the
canvas every time, so a user's manual edits are input to the next layout rather than
state to reconcile.

Token cost, measured with Claude's tokenizer rather than estimated:
  - build a VPC diagram:  515 tok as operations vs 3180 as XML   (6.2x)
  - add one icon:          27 tok as an operation vs 3823 re-emitting (142x)
  - read current state:   216 tok as an outline vs 3180 as XML   (14.7x)

The 142x is the one that matters day to day: "add a Redis" is one operation, not a
rewrite of the whole diagram.

Routing in the system prompt sends AWS architecture through this path and leaves
flowcharts, BPMN, sequence diagrams, mind maps and Azure/GCP on display_diagram —
the layout engine's primitives (nested rows, columns, grids) do not model a sequence
diagram's lifelines or a mind map's radial spread, and pretending otherwise would
make those worse rather than better.

Also: added a NOTICE recording the MIT port and the AWS Architecture Icons terms,
and a narrow .gitignore exception so the generated catalog is tracked while the
root data/ directory (admin settings, contains secrets) stays ignored.

403 unit tests + 3 new e2e. Verified in the real app: a structural tool call renders
with real stencils and container markers; a second call adds one node and keeps
everything from the first; an invented name is refused and nothing is drawn. The 13
existing diagram e2e tests still pass.
This commit is contained in:
dayuan.jiang
2026-08-09 12:13:54 +09:00
parent a2f892ca82
commit cd1df1eb6a
14 changed files with 3475 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

325
tests/unit/diagram-engine-index.test.ts Normal file
View File

@@ -0,0 +1,325 @@
import { countTokens } from "@anthropic-ai/tokenizer"
import { describe, expect, it } from "vitest"
import {
describeDiagram,
type Operation,
restructureDiagram,
} from "@/lib/diagram-engine"
import { parseDiagram } from "@/lib/diagram-engine/parse"
import { findNode, findParent } from "@/lib/diagram-engine/types"
/** The operations that build a 3-tier VPC diagram from nothing. */
const BUILD_3TIER: Operation[] = [
{ op: "set_title", title: "VPC Multi-AZ 3-tier" },
{ op: "add_box", id: "users", label: "Users / Internet" },
{
op: "add_container",
id: "region",
label: "Region (ap-southeast-1)",
dir: "row",
gname: "group_region",
},
{
op: "add_container",
id: "vpc",
parent: "region",
label: "VPC 10.0.0.0/16",
dir: "col",
gname: "group_vpc",
},
{
op: "add_icon",
id: "igw",
parent: "vpc",
name: "internet_gateway",
label: "Internet Gateway",
},
{
op: "add_icon",
id: "alb",
parent: "vpc",
name: "application_load_balancer",
label: "ALB",
},
{ op: "add_container", id: "azs", parent: "vpc", label: "", dir: "row" },
{
op: "add_container",
id: "az_a",
parent: "azs",
label: "AZ-a",
dir: "col",
gname: "group_availability_zone",
},
{
op: "add_container",
id: "pub_a",
parent: "az_a",
label: "Public Subnet",
dir: "col",
gname: "group_subnet",
},
{
op: "add_icon",
id: "nat_a",
parent: "pub_a",
name: "nat_gateway",
label: "NAT",
},
{
op: "add_container",
id: "app_a",
parent: "az_a",
label: "Private Subnet (App)",
dir: "col",
gname: "group_subnet",
},
{ op: "add_icon", id: "ec2_a", parent: "app_a", name: "ec2", label: "EC2" },
{
op: "add_container",
id: "db_a",
parent: "az_a",
label: "Private Subnet (Data)",
dir: "col",
gname: "group_subnet",
},
{
op: "add_icon",
id: "rds_a",
parent: "db_a",
name: "rds",
label: "RDS (Primary)",
},
{ op: "link", source: "users", target: "igw", label: "HTTPS", step: 1 },
{ op: "link", source: "igw", target: "alb", label: "forward", step: 2 },
{ op: "link", source: "alb", target: "ec2_a", label: "route", step: 3 },
{ op: "link", source: "ec2_a", target: "rds_a", label: "query", step: 4 },
]
describe("restructureDiagram builds from an empty canvas", () => {
const r = restructureDiagram("", BUILD_3TIER)
it("succeeds", () => {
expect(r.errors).toEqual([])
expect(r.xml).not.toBeNull()
})
it("produces XML draw.io can load", () => {
expect(r.xml).toContain("<mxfile")
expect(r.xml).toContain('<mxCell id="0"/>')
expect(r.xml).toContain('<mxCell id="1" parent="0"/>')
})
it("resolves catalog names to verbatim stencil styles, with official colours", () => {
// #ED7100 is AWS Compute orange; nothing here is hand-assembled.
expect(r.xml).toContain("resIcon=mxgraph.aws4.ec2")
expect(r.xml).toContain("#ED7100")
expect(r.xml).toContain("grIcon=mxgraph.aws4.group_vpc")
})
it("stamps container=1 so a user can drag shapes between frames", () => {
const vpcStyle =
r.xml?.match(/<mxCell id="vpc"[^>]*style="([^"]*)"/)?.[1] ?? ""
expect(vpcStyle).toContain("container=1")
})
it("reads back the structure it was asked to build", () => {
const { tree } = parseDiagram(r.xml as string)
expect(findParent(tree, "vpc")?.id).toBe("region")
expect(findParent(tree, "az_a")?.id).toBe("azs")
expect(findParent(tree, "nat_a")?.id).toBe("pub_a")
expect(tree.links).toHaveLength(4)
expect(tree.title).toBe("VPC Multi-AZ 3-tier")
})
})
describe("restructureDiagram edits an existing canvas", () => {
const built = restructureDiagram("", BUILD_3TIER).xml as string
it("adds one node with one small operation", () => {
const r = restructureDiagram(built, [
{
op: "add_icon",
id: "cache_a",
parent: "app_a",
name: "elasticache",
label: "Redis",
},
])
expect(r.errors).toEqual([])
const { tree } = parseDiagram(r.xml as string)
expect(findParent(tree, "cache_a")?.id).toBe("app_a")
// everything else is still there
expect(findNode(tree, "rds_a")).not.toBeNull()
expect(tree.links).toHaveLength(4)
})
it("re-orients a container without touching anything else", () => {
const r = restructureDiagram(built, [
{ op: "set_dir", id: "vpc", dir: "row" },
])
expect(r.errors).toEqual([])
const { tree } = parseDiagram(r.xml as string)
expect((findNode(tree, "vpc") as { dir: string }).dir).toBe("row")
})
it("removes a subtree and the edges that pointed into it", () => {
const r = restructureDiagram(built, [{ op: "remove", id: "db_a" }])
const { tree } = parseDiagram(r.xml as string)
expect(findNode(tree, "db_a")).toBeNull()
expect(findNode(tree, "rds_a")).toBeNull()
// the ec2 → rds edge went with it
expect(tree.links).toHaveLength(3)
})
it("keeps a colour the user changed by hand", () => {
const edited = built.replace(
/(<mxCell id="users"[^>]*style="[^"]*)"/,
'$1fillColor=#FF0000;"',
)
const r = restructureDiagram(edited, [
{ op: "set_label", id: "users", label: "Clients" },
])
expect(r.xml).toContain("fillColor=#FF0000")
})
it("keeps a shape the user added by hand", () => {
const edited = built.replace(
"</root>",
'<mxCell id="mynote" value="Note" style="shape=note;whiteSpace=wrap;html=1;" vertex="1" parent="1"><mxGeometry x="1200" y="60" width="140" height="80" as="geometry"/></mxCell></root>',
)
const r = restructureDiagram(edited, [
{
op: "add_icon",
id: "s3",
parent: "vpc",
name: "s3",
label: "S3",
},
])
expect(r.xml).toContain('id="mynote"')
})
it("respects where the user dragged a node to", () => {
// The user moved the EC2 icon from the app subnet into the public subnet.
const moved = built.replace(
/(<mxCell id="ec2_a"[^>]*)parent="app_a"/,
'$1parent="pub_a"',
)
const r = restructureDiagram(moved, [
{ op: "set_label", id: "ec2_a", label: "EC2 (moved)" },
])
const { tree } = parseDiagram(r.xml as string)
expect(findParent(tree, "ec2_a")?.id).toBe("pub_a")
})
})
describe("invented stencil names are rejected, not rendered blank", () => {
it("fails the whole call and suggests real names", () => {
const r = restructureDiagram("", [
{
op: "add_icon",
id: "x",
name: "s3_bucket_storage",
label: "Bucket",
},
])
expect(r.xml).toBeNull()
expect(r.errors[0]).toContain("not in the stencil catalog")
expect(r.errors[0]).toContain("s3")
})
it("rejects an invented group stencil too", () => {
const r = restructureDiagram("", [
{
op: "add_container",
id: "g",
label: "X",
dir: "row",
gname: "group_made_up",
},
])
expect(r.xml).toBeNull()
expect(r.errors[0]).toContain("not in the stencil catalog")
})
it("still returns the outline so the model can see what it built", () => {
const r = restructureDiagram("", [
{ op: "add_icon", id: "x", name: "nope_not_real" },
])
expect(r.outline).toContain("x: icon nope_not_real")
})
it("reports a failed operation without rendering a half-built diagram", () => {
const r = restructureDiagram("", [
{ op: "add_icon", id: "a", name: "s3" },
{ op: "move", id: "ghost", parent: "a" },
])
expect(r.xml).toBeNull()
expect(r.errors.some((e) => e.includes("ghost"))).toBe(true)
})
})
describe("describeDiagram", () => {
it("reports an empty canvas plainly", () => {
expect(describeDiagram("").outline).toBe("(empty canvas)")
})
it("outlines what is on the canvas without changing it", () => {
const built = restructureDiagram("", BUILD_3TIER).xml as string
const d = describeDiagram(built)
expect(d.outline).toContain("vpc: col")
expect(d.outline).toContain("ec2_a: icon ec2")
expect(d.outline).toContain("link users -> igw")
expect(d.needsAdoption).toBe(false)
})
it("flags a diagram that did not come from the engine", () => {
const foreign = `<mxfile><diagram name="Page-1" id="p"><mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/><mxCell id="a" value="X" style="rounded=0;" vertex="1" parent="1"><mxGeometry x="0" y="0" width="100" height="50" as="geometry"/></mxCell></root></mxGraphModel></diagram></mxfile>`
expect(describeDiagram(foreign).needsAdoption).toBe(true)
})
})
describe("token cost: operations versus raw XML", () => {
// The reason for the whole exercise. Measured with Claude's own tokenizer, which the
// repo already depends on, rather than a characters/4 estimate.
const built = restructureDiagram("", BUILD_3TIER).xml as string
it("building a diagram costs far fewer tokens as operations than as XML", () => {
const asOps = countTokens(JSON.stringify(BUILD_3TIER))
const asXml = countTokens(built)
console.log(
`build: ${asOps} tok as operations vs ${asXml} tok as XML (${(asXml / asOps).toFixed(1)}x)`,
)
expect(asOps).toBeLessThan(asXml / 2)
})
it("adding one icon costs a fraction of re-emitting the diagram", () => {
const oneOp: Operation[] = [
{
op: "add_icon",
id: "cache",
parent: "app_a",
name: "elasticache",
label: "Redis",
},
]
const opTokens = countTokens(JSON.stringify(oneOp))
const xmlTokens = countTokens(
restructureDiagram(built, oneOp).xml as string,
)
console.log(
`add one icon: ${opTokens} tok as an operation vs ${xmlTokens} tok re-emitting the XML (${Math.round(xmlTokens / opTokens)}x)`,
)
expect(opTokens).toBeLessThan(60)
expect(opTokens).toBeLessThan(xmlTokens / 20)
})
it("the outline the model reads back is much cheaper than the XML", () => {
const outlineTokens = countTokens(describeDiagram(built).outline)
const xmlTokens = countTokens(built)
console.log(
`read current state: ${outlineTokens} tok as an outline vs ${xmlTokens} tok as XML (${(xmlTokens / outlineTokens).toFixed(1)}x)`,
)
expect(outlineTokens).toBeLessThan(xmlTokens / 3)
})
})