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

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

View File

@@ -0,0 +1,219 @@
/**
* The stencil catalog: a name → verbatim draw.io style map.
*
* This is the anti-hallucination layer. The model asks for `icon("s3")`; the engine
* looks the name up here and gets the exact style draw.io ships, including the official
* category colour, the connection points and `aspect=fixed`. A name that is not in the
* catalog fails at build time with a suggestion, rather than becoming an empty square in
* the rendered diagram — which is what happens when a model writes
* `resIcon=mxgraph.aws4.s3_bucket_thing` by hand and nothing checks it.
*
* The styles are verbatim from draw.io's own shape index (via drawio-ai-kit, which
* generated them from jgraph/drawio-mcp's index, Apache-2.0). Nothing here is
* hand-assembled, so there is no chance of a plausible-looking but wrong colour.
*/
import stencils from "./data/aws-stencils.json"
const ICONS = stencils.icons as Record<string, string>
const GROUPS = stencils.groups as Record<string, string>
export interface CatalogEntry {
name: string
kind: "icon" | "group"
style: string
/** Official colour from the style, for showing the model what it is getting. */
color: string | null
}
function colorOf(style: string): string | null {
return style.match(/(?:^|;)fillColor=([^;]+)/)?.[1] ?? null
}
/** Exact lookup. Returns null for an unknown name — never a guess. */
export function lookupStencil(
name: string,
kind?: "icon" | "group",
): CatalogEntry | null {
if (kind !== "group" && ICONS[name])
return {
name,
kind: "icon",
style: ICONS[name],
color: colorOf(ICONS[name]),
}
if (kind !== "icon" && GROUPS[name])
return {
name,
kind: "group",
style: GROUPS[name],
color: colorOf(GROUPS[name]),
}
return null
}
/** The resolver the renderer takes, so the engine itself does not depend on the catalog. */
export function resolveStyle(
name: string,
kind: "icon" | "group",
): string | null {
return lookupStencil(name, kind)?.style ?? null
}
/** Normalise for matching: lowercase, and non-alphanumerics collapsed to single spaces. */
function norm(s: string): string {
return s
.toLowerCase()
.replace(/[^a-z0-9]+/g, " ")
.trim()
}
/**
* Shorthand people type, mapped to words that actually appear in a catalog name.
*
* The direction matters: the target has to exist in the catalog. AWS's own stencil names
* are already abbreviated — EKS is `eks`, not `elastic_kubernetes_service`, and nothing
* in the catalog contains the word "kubernetes" at all — so expanding an abbreviation
* into its full product name finds nothing. These entries go the other way, from a
* spoken-out name or a nickname to the token the catalog uses.
*/
const ALIASES: Record<string, string> = {
k8s: "eks",
kubernetes: "eks",
kube: "eks",
alb: "application load balancer",
nlb: "network load balancer",
elb: "elastic load balancing",
asg: "auto scaling",
apigw: "api gateway",
cf: "cloudfront",
cw: "cloudwatch",
ddb: "dynamodb",
tgw: "transit gateway",
igw: "internet gateway",
r53: "route 53",
iam: "identity and access management",
kms: "key management service",
postgres: "rds",
postgresql: "rds",
mysql: "rds",
aurora: "aurora",
bucket: "s3",
}
/**
* Score one entry against the query tokens. Higher is better; 0 means no match.
*
* The extra-words penalty is what makes "s3" return `s3` rather than
* `backup_aws_backup_support_for_amazon_s3` — both contain the token, so without it the
* winner comes down to iteration order. It counts only the words the query did NOT ask
* for, so a deliberately multi-word query like "nat gateway" is not punished for being
* specific.
*/
function score(name: string, qTokens: string[], qJoined: string): number {
const n = norm(name)
const words = n.split(" ")
let s = 0
if (n === qJoined) s += 100
if (n.replace(/ /g, "") === qJoined.replace(/ /g, "")) s += 60
for (const t of qTokens) {
if (words.includes(t)) s += 25
else if (n.includes(t)) s += 12
}
if (s === 0) return 0
const extra = words.filter((w) => !qTokens.includes(w)).length
return s - Math.min(24, extra * 4)
}
export interface SearchHit {
name: string
kind: "icon" | "group"
color: string | null
}
/**
* Find stencils by keyword.
*
* Returns names and colours only, not styles. The model builds with `icon("<name>")` and
* the engine resolves the style itself, so sending the style — around 600 characters per
* AWS entry, and 20KB+ for an Azure one with an embedded image — would be pure context
* burn.
*/
export function searchStencils(
query: string,
opts: { limit?: number; kind?: "icon" | "group" } = {},
): SearchHit[] {
const limit = opts.limit ?? 8
const tokens = norm(query)
.split(" ")
.filter(Boolean)
.map((t) => ALIASES[t] ?? t)
.flatMap((t) => t.split(" "))
if (tokens.length === 0) return []
const joined = tokens.join(" ")
const pool: [string, string, "icon" | "group"][] = []
if (opts.kind !== "group")
for (const [n, st] of Object.entries(ICONS)) pool.push([n, st, "icon"])
if (opts.kind !== "icon")
for (const [n, st] of Object.entries(GROUPS))
pool.push([n, st, "group"])
return pool
.map(([name, style, kind]) => ({
name,
kind,
color: colorOf(style),
s: score(name, tokens, joined),
}))
.filter((r) => r.s > 0)
.sort((a, b) => b.s - a.s || a.name.length - b.name.length)
.slice(0, limit)
.map(({ name, kind, color }) => ({ name, kind, color }))
}
/**
* Suggest real names for one that does not exist.
*
* Plain search is not quite the right tool here. A model that writes
* `s3_bucket_storage` most likely meant `s3`, but searching that whole phrase ranks
* `s3_storage_lens` first — it matches more of the query. So we also search the
* leading token on its own and put those hits first: an invented name is usually a
* real service name with extra words stuck on the end.
*/
function suggestFor(name: string, kind: "icon" | "group"): string[] {
const words = norm(name.replace(/_/g, " ")).split(" ").filter(Boolean)
const out: string[] = []
const add = (hits: SearchHit[]) => {
for (const h of hits) if (!out.includes(h.name)) out.push(h.name)
}
if (words.length > 1) add(searchStencils(words[0], { limit: 2, kind }))
add(searchStencils(words.join(" "), { limit: 3, kind }))
return out.slice(0, 3)
}
/**
* Validate the icon names in a tree before laying it out, so a bad name is reported as
* a correctable error with suggestions instead of rendering as a blank square — which is
* what an unchecked invented name becomes in draw.io.
*/
export function checkNames(
names: { id: string; name: string; kind: "icon" | "group" }[],
): { id: string; name: string; suggestions: string[] }[] {
const bad: { id: string; name: string; suggestions: string[] }[] = []
for (const n of names) {
if (!n.name || lookupStencil(n.name, n.kind)) continue
bad.push({
id: n.id,
name: n.name,
suggestions: suggestFor(n.name, n.kind),
})
}
return bad
}
/** Total catalog size, for the tool description. */
export const CATALOG_SIZE = {
icons: Object.keys(ICONS).length,
groups: Object.keys(GROUPS).length,
}

File diff suppressed because it is too large Load Diff

120
lib/diagram-engine/index.ts Normal file
View File

@@ -0,0 +1,120 @@
/**
* The engine's entry point: one call takes the current canvas XML plus a list of
* structural operations and returns new canvas XML.
*
* current XML → parse → apply operations → check names → layout → render → new XML
*
* The tree is not stored anywhere between calls. It is re-derived from the canvas every
* time, so a user's manual edits — moving a shape into a different frame, recolouring a
* box, adding an annotation — are simply part of the input to the next layout. There is
* no second copy of the state, and therefore nothing to reconcile.
*/
import { checkNames, resolveStyle } from "./catalog"
import {
applyOperations,
collectNames,
type Operation,
outline,
} from "./operations"
import { parseDiagram } from "./parse"
import { renderDiagram } from "./render"
import type { DiagramTree } from "./types"
export interface RestructureResult {
/** New canvas XML, or null when the request could not be carried out. */
xml: string | null
/** Compact outline of the resulting structure, for the model to read back. */
outline: string
/** Operations that could not be applied, and invented stencil names. */
errors: string[]
/** Non-fatal notes: pages skipped, structure that could not be read cleanly. */
warnings: string[]
}
export interface RestructureOptions {
/** Which page of a multi-page document to work on. */
pageIndex?: number
/** Diagram-wide icon glyph size. */
iconSize?: number
}
/**
* Apply structural operations to whatever is on the canvas.
*
* `currentXml` may be empty — that is how a diagram gets built from scratch.
*
* An invented stencil name is a hard error, not a silent fallback: draw.io renders an
* unknown `resIcon` as a blank square, so a diagram that "worked" would be quietly
* missing icons. The error carries suggestions from the catalog so the model can fix it
* in one more turn.
*/
export function restructureDiagram(
currentXml: string,
ops: Operation[],
opts: RestructureOptions = {},
): RestructureResult {
const warnings: string[] = []
let tree: DiagramTree
if (currentXml.trim()) {
const parsed = parseDiagram(currentXml, opts.pageIndex ?? 0)
tree = parsed.tree
warnings.push(...parsed.warnings)
} else {
tree = { roots: [], links: [], foreign: [] }
}
const applied = applyOperations(tree, ops)
const errors = [...applied.errors]
// Catch invented names before rendering, so the model gets a correctable error
// instead of a diagram with blank squares in it.
for (const bad of checkNames(collectNames(applied.tree))) {
const hint = bad.suggestions.length
? ` Did you mean: ${bad.suggestions.join(", ")}?`
: ""
errors.push(
`"${bad.name}" (node ${bad.id}) is not in the stencil catalog.${hint}`,
)
}
if (errors.length > 0)
return { xml: null, outline: outline(applied.tree), errors, warnings }
const rendered = renderDiagram(applied.tree, {
resolveStyle,
iconSize: opts.iconSize,
})
if (rendered.danglingLinks.length)
warnings.push(
`Dropped edge(s) pointing at missing nodes: ${rendered.danglingLinks.join(", ")}.`,
)
return {
xml: rendered.xml,
outline: outline(applied.tree),
errors: [],
warnings,
}
}
/** Read the current canvas structure without changing it. */
export function describeDiagram(
currentXml: string,
pageIndex = 0,
): { outline: string; warnings: string[]; needsAdoption: boolean } {
if (!currentXml.trim())
return { outline: "(empty canvas)", warnings: [], needsAdoption: false }
const { tree, warnings, needsAdoption } = parseDiagram(
currentXml,
pageIndex,
)
return { outline: outline(tree), warnings, needsAdoption }
}
export { CATALOG_SIZE, lookupStencil, searchStencils } from "./catalog"
export { type Operation, OperationSchema } from "./operations"
export { parseDiagram } from "./parse"
export { renderDiagram } from "./render"
export type { DiagramNode, DiagramTree } from "./types"

View File

@@ -0,0 +1,440 @@
/**
* Structural operations — what the model sends instead of XML.
*
* The tree is re-derived from the canvas on every call, so an operation names existing
* nodes by id and says what to change. Adding one node costs a few dozen tokens; the
* equivalent as raw mxCell XML is hundreds, and re-emitting the whole diagram to add one
* icon costs thousands.
*
* Operations are applied in order, each against the result of the last, so a sequence
* like "add a frame, then move two nodes into it" works in a single call.
*/
import { z } from "zod"
import {
type ContainerNode,
type DiagramNode,
type DiagramTree,
findNode,
findParent,
isContainer,
type LinkSpec,
walkTree,
} from "./types"
export const OperationSchema = z.discriminatedUnion("op", [
z.object({
op: z.literal("add_icon"),
id: z.string().describe("New unique id for this node"),
parent: z
.string()
.optional()
.describe("Container id to add into; omit for top level"),
name: z.string().describe("Catalog stencil name, e.g. 's3' or 'ec2'"),
label: z.string().optional(),
after: z
.string()
.optional()
.describe("Insert after this sibling id; omit to append"),
}),
z.object({
op: z.literal("add_box"),
id: z.string(),
parent: z.string().optional(),
label: z.string(),
after: z.string().optional(),
}),
z.object({
op: z.literal("add_container"),
id: z.string(),
parent: z.string().optional(),
label: z
.string()
.describe("Frame title; empty string means invisible wrapper"),
dir: z.enum(["row", "col"]).describe("How children stack"),
gname: z
.string()
.optional()
.describe(
"Group stencil name, e.g. 'group_vpc'; omit for a plain frame",
),
gap: z.number().optional(),
after: z.string().optional(),
}),
z.object({
op: z.literal("add_grid"),
id: z.string(),
parent: z.string().optional(),
label: z.string(),
cols: z.number().describe("Number of columns"),
gap: z.number().optional(),
after: z.string().optional(),
}),
z.object({
op: z.literal("remove"),
id: z
.string()
.describe("Node to delete; its descendants and edges go too"),
}),
z.object({
op: z.literal("move"),
id: z.string(),
parent: z
.string()
.optional()
.describe("New container id; omit to move to top level"),
after: z.string().optional(),
}),
z.object({
op: z.literal("set_label"),
id: z.string(),
label: z.string(),
}),
z.object({
op: z.literal("set_dir"),
id: z.string().describe("Container to re-orient"),
dir: z.enum(["row", "col"]),
}),
z.object({
op: z.literal("set_gap"),
id: z.string(),
gap: z.number(),
}),
z.object({
op: z.literal("link"),
source: z.string(),
target: z.string(),
label: z.string().optional(),
dashed: z
.boolean()
.optional()
.describe("Dashed line — replication, sync, policy"),
step: z
.number()
.optional()
.describe("Step number, shown as an 'N. ' prefix"),
}),
z.object({
op: z.literal("unlink"),
source: z.string(),
target: z.string(),
}),
z.object({
op: z.literal("set_title"),
title: z.string(),
}),
])
export type Operation = z.infer<typeof OperationSchema>
export interface ApplyResult {
tree: DiagramTree
/** One entry per operation that could not be applied, in order. */
errors: string[]
}
/** Insert into a child list, after a named sibling or at the end. */
function insert(
list: DiagramNode[],
node: DiagramNode,
after: string | undefined,
): void {
if (after) {
const i = list.findIndex((c) => c.id === after)
if (i >= 0) {
list.splice(i + 1, 0, node)
return
}
}
list.push(node)
}
/** Detach a node from wherever it currently sits. Returns it, or null if not found. */
function detach(tree: DiagramTree, id: string): DiagramNode | null {
const rootIdx = tree.roots.findIndex((r) => r.id === id)
if (rootIdx >= 0) return tree.roots.splice(rootIdx, 1)[0]
const parent = findParent(tree, id)
if (!parent) return null
const i = parent.children.findIndex((c) => c.id === id)
return i >= 0 ? parent.children.splice(i, 1)[0] : null
}
/** Would making `id` a descendant of `parentId` create a cycle? */
function wouldCycle(tree: DiagramTree, id: string, parentId: string): boolean {
if (id === parentId) return true
const node = findNode(tree, id)
if (!node || !isContainer(node)) return false
for (const d of walkTree({ ...tree, roots: [node] }))
if (d.id === parentId) return true
return false
}
/**
* Resolve where a new or moved node goes. Returns the child list to insert into, or an
* error string.
*/
function targetList(
tree: DiagramTree,
parentId: string | undefined,
): DiagramNode[] | string {
if (!parentId) return tree.roots
const p = findNode(tree, parentId)
if (!p) return `No node with id "${parentId}"`
if (!isContainer(p))
return `"${parentId}" is a ${p.kind}, not a container — it cannot hold children`
return p.children
}
/**
* Apply operations to a tree, in order.
*
* The input tree is deep-copied first: a partially-applied batch must not leave the
* caller's tree half-mutated when a later operation fails.
*/
export function applyOperations(
input: DiagramTree,
ops: Operation[],
): ApplyResult {
const tree: DiagramTree = structuredClone(input)
const errors: string[] = []
const exists = (id: string) => findNode(tree, id) !== null
for (const op of ops) {
switch (op.op) {
case "add_icon":
case "add_box":
case "add_container":
case "add_grid": {
if (exists(op.id)) {
errors.push(`${op.op}: id "${op.id}" is already taken`)
break
}
const list = targetList(tree, op.parent)
if (typeof list === "string") {
errors.push(`${op.op}: ${list}`)
break
}
let node: DiagramNode
if (op.op === "add_icon")
node = {
kind: "icon",
id: op.id,
name: op.name,
label: op.label ?? "",
}
else if (op.op === "add_box")
node = { kind: "box", id: op.id, label: op.label }
else if (op.op === "add_container")
node = {
kind: "group",
id: op.id,
gname: op.gname ?? null,
label: op.label,
dir: op.dir,
gap: op.gap ?? 20,
children: [],
}
else
node = {
kind: "grid",
id: op.id,
gname: null,
label: op.label,
cols: Math.max(1, op.cols),
gap: op.gap ?? 14,
children: [],
}
insert(list, node, op.after)
break
}
case "remove": {
const node = findNode(tree, op.id)
if (!node) {
errors.push(`remove: no node with id "${op.id}"`)
break
}
// Collect the subtree's ids first — edges touching any of them go too,
// otherwise draw.io renders an arrow pointing at nothing.
const doomed = new Set<string>()
for (const d of walkTree({ ...tree, roots: [node] }))
doomed.add(d.id)
detach(tree, op.id)
tree.links = tree.links.filter(
(l) => !doomed.has(l.source) && !doomed.has(l.target),
)
break
}
case "move": {
if (!exists(op.id)) {
errors.push(`move: no node with id "${op.id}"`)
break
}
if (op.parent && !exists(op.parent)) {
errors.push(`move: no node with id "${op.parent}"`)
break
}
if (op.parent && wouldCycle(tree, op.id, op.parent)) {
errors.push(
`move: cannot move "${op.id}" into "${op.parent}" — that is inside itself`,
)
break
}
const list = targetList(tree, op.parent)
if (typeof list === "string") {
errors.push(`move: ${list}`)
break
}
const node = detach(tree, op.id)
if (!node) {
errors.push(`move: could not detach "${op.id}"`)
break
}
insert(list, node, op.after)
break
}
case "set_label": {
const node = findNode(tree, op.id)
if (!node) {
errors.push(`set_label: no node with id "${op.id}"`)
break
}
if (node.kind === "title") {
errors.push(
`set_label: use set_title to change the page title`,
)
break
}
node.label = op.label
break
}
case "set_dir": {
const node = findNode(tree, op.id)
if (!node || !isContainer(node)) {
errors.push(`set_dir: "${op.id}" is not a container`)
break
}
if (node.kind === "grid") {
errors.push(
`set_dir: "${op.id}" is a grid — change its column count instead`,
)
break
}
node.dir = op.dir
break
}
case "set_gap": {
const node = findNode(tree, op.id)
if (!node || !isContainer(node)) {
errors.push(`set_gap: "${op.id}" is not a container`)
break
}
;(node as ContainerNode).gap = Math.max(0, op.gap)
break
}
case "link": {
if (!exists(op.source)) {
errors.push(`link: no node with id "${op.source}"`)
break
}
if (!exists(op.target)) {
errors.push(`link: no node with id "${op.target}"`)
break
}
const dup = tree.links.some(
(l) => l.source === op.source && l.target === op.target,
)
if (dup) {
errors.push(
`link: "${op.source}" → "${op.target}" already exists`,
)
break
}
const link: LinkSpec = { source: op.source, target: op.target }
if (op.label) link.label = op.label
if (op.dashed) link.dashed = true
if (op.step != null) link.step = op.step
tree.links.push(link)
break
}
case "unlink": {
const before = tree.links.length
tree.links = tree.links.filter(
(l) => !(l.source === op.source && l.target === op.target),
)
if (tree.links.length === before)
errors.push(
`unlink: no edge from "${op.source}" to "${op.target}"`,
)
break
}
case "set_title":
tree.title = op.title
break
}
}
return { tree, errors }
}
/** Every icon/group name in a tree, for validating against the catalog. */
export function collectNames(
tree: DiagramTree,
): { id: string; name: string; kind: "icon" | "group" }[] {
const out: { id: string; name: string; kind: "icon" | "group" }[] = []
for (const n of walkTree(tree)) {
if (n.kind === "icon" && n.name)
out.push({ id: n.id, name: n.name, kind: "icon" })
else if (isContainer(n) && n.gname)
out.push({ id: n.id, name: n.gname, kind: "group" })
}
return out
}
/**
* A compact text outline of the tree, for showing the model what is on the canvas.
*
* Sending the tree as JSON would cost several times more for the same information, and
* the model does not need coordinates — it needs to know what exists and how it nests so
* it can name ids in the next operation.
*/
export function outline(tree: DiagramTree): string {
const lines: string[] = []
if (tree.title) lines.push(`title: ${tree.title}`)
const walk = (n: DiagramNode, depth: number) => {
const pad = " ".repeat(depth)
if (n.kind === "icon")
lines.push(
`${pad}${n.id}: icon ${n.name}${n.label ? ` "${n.label}"` : ""}`,
)
else if (n.kind === "box") lines.push(`${pad}${n.id}: box "${n.label}"`)
else if (n.kind === "title") lines.push(`${pad}${n.id}: title`)
else {
const meta = n.kind === "grid" ? `grid cols=${n.cols}` : n.dir
lines.push(
`${pad}${n.id}: ${meta}${n.label ? ` "${n.label}"` : " (wrapper)"}`,
)
for (const c of n.children) walk(c, depth + 1)
}
}
for (const r of tree.roots) walk(r, 0)
for (const l of tree.links) {
const bits = [l.label, l.dashed ? "dashed" : null]
.filter(Boolean)
.join(", ")
lines.push(`link ${l.source} -> ${l.target}${bits ? ` (${bits})` : ""}`)
}
if (tree.foreign.length)
lines.push(
`${tree.foreign.length} cell(s) kept as-is: ${tree.foreign.map((f) => f.id).join(", ")}`,
)
return lines.join("\n")
}