From 98f2a4fda6081084f451b9a29278974560b0af10 Mon Sep 17 00:00:00 2001 From: "dayuan.jiang" Date: Sat, 10 Oct 2026 16:28:42 +0900 Subject: [PATCH] feat(diagram): named styles the model defines once and uses by name The model may now write before its cells and refer to the name inside a cell's style, the way draw.io resolves its own named styles. style-classes.ts expands the names after validation (so repaired cells get them too) and adds draw.io's defaults html=1 and whiteSpace=wrap where they belong, so the prompts can tell the model never to write them. Wired into display_diagram, append_diagram, create_new_diagram, add_page, edit_diagram (defaults only) and the streaming preview. Prompts (system prompt, tool description, MCP drawing guide) teach the notation with a shared example, drop the draw.io defaults from every example and tell the model how labels, line breaks and reserved names work. Measured on the five start-screen examples, same model and prompt, old vs new prompt: gpt-6-luna 19% fewer output tokens, Opus 5.5 19% fewer; flow and architecture diagrams 15% to 46%, creative drawings about equal. --- app/api/chat/route.ts | 16 +- components/chat/use-diagram-streaming.ts | 25 +- hooks/use-diagram-tool-handlers.ts | 4 + lib/system-prompts.ts | 32 +- lib/utils.ts | 5 +- packages/mcp-server/src/drawing-guide.ts | 27 +- packages/mcp-server/src/edit-diagram.ts | 16 +- packages/mcp-server/src/index.ts | 34 +- packages/mcp-server/src/new-diagram.ts | 51 ++- packages/mcp-server/src/style-classes.ts | 232 ++++++++++ packages/mcp-server/src/xml-examples.ts | 17 +- tests/unit/style-classes.test.ts | 499 ++++++++++++++++++++++ tests/unit/use-diagram-streaming.test.tsx | 88 +++- 13 files changed, 995 insertions(+), 51 deletions(-) create mode 100644 packages/mcp-server/src/style-classes.ts create mode 100644 tests/unit/style-classes.test.ts diff --git a/app/api/chat/route.ts b/app/api/chat/route.ts index 486fdc35..c6256960 100644 --- a/app/api/chat/route.ts +++ b/app/api/chat/route.ts @@ -60,7 +60,10 @@ import { getShapeLibrary, SHAPE_LIBRARY_LIST, } from "@/packages/mcp-server/src/shape-library.ts" -import { SWIMLANE_EXAMPLE } from "@/packages/mcp-server/src/xml-examples.ts" +import { + STYLE_CLASS_EXAMPLE, + SWIMLANE_EXAMPLE, +} from "@/packages/mcp-server/src/xml-examples.ts" // No explicit cap: a reasoning model can spend minutes planning before it emits // the tool call, so take whatever the host allows. Vercel's own default is 300s, @@ -659,10 +662,10 @@ IMPORTANT: The "Current diagram XML" is the SINGLE SOURCE OF TRUTH for what's on tools: { // Client-side tool that will be executed on the client display_diagram: { - description: `Display a diagram on draw.io. Pass ONLY the mxCell elements - wrapper tags and root cells are added automatically. + description: `Display a diagram on draw.io. Pass ONLY the mxStyle definitions and the mxCell elements - wrapper tags and root cells are added automatically. VALIDATION RULES (XML will be rejected if violated): -1. Generate ONLY mxCell elements - NO wrapper tags (, , ) +1. Generate ONLY mxStyle definitions and mxCell elements - NO wrapper tags (, , ) 2. Do NOT include root cells (id="0" or id="1") - they are added automatically 3. All mxCell elements must be siblings - never nested 4. Every mxCell needs a unique id (start from "2") @@ -672,6 +675,9 @@ VALIDATION RULES (XML will be rejected if violated): Example (generate ONLY this - no wrapper tags): ${SWIMLANE_EXAMPLE} +Styles: define a style used by several cells ONCE with before the cells and use the name in style like a CSS class; overrides after the name win. Name only styles that two or more cells share; names must not be draw.io's own style names (text, ellipse, rhombus, swimlane, label, image, blue, green, red, gray, yellow, orange, purple, pink). A definition applies to the call it is in. html=1 and whiteSpace=wrap are added automatically, never write them. Labels are HTML: use <br> for a line break, never \\n; a literal < or > is written &lt; or &gt;. +${STYLE_CLASS_EXAMPLE} + Notes: - For AWS diagrams, use **AWS 2025 icons**. - For animated connectors, add "flowAnimation=1" to edge style. @@ -690,12 +696,12 @@ Operations: - add: Add a new cell. Provide cell_id (new unique id) and new_xml. - delete: Remove a cell. Cascade is automatic: children AND edges (source/target) are auto-deleted. Only specify ONE cell_id. -For update/add, new_xml must be a complete mxCell element including mxGeometry. +For update/add, new_xml must be a complete mxCell element including mxGeometry, with its complete style (named styles are not available in edit_diagram). ⚠️ JSON ESCAPING: Every " inside new_xml MUST be escaped as \\". Example: id=\\"5\\" value=\\"Label\\" Example - Add a rectangle: -{"operations": [{"operation": "add", "cell_id": "rect-1", "new_xml": ""}]} +{"operations": [{"operation": "add", "cell_id": "rect-1", "new_xml": ""}]} Example - Delete container (children & edges auto-deleted): {"operations": [{"operation": "delete", "cell_id": "2"}]}`, diff --git a/components/chat/use-diagram-streaming.ts b/components/chat/use-diagram-streaming.ts index dffd9962..c7381dab 100644 --- a/components/chat/use-diagram-streaming.ts +++ b/components/chat/use-diagram-streaming.ts @@ -12,6 +12,12 @@ import { } from "@/lib/utils" import { applyDiagramOperations } from "@/packages/mcp-server/src/diagram-operations.ts" import { BLANK_MXFILE } from "@/packages/mcp-server/src/pages.ts" +import { + addDefaultStyles, + applyStyleClasses, + edgeIdsOf, + readStyleClasses, +} from "@/packages/mcp-server/src/style-classes.ts" // Helper to extract complete operations from streaming input function getCompleteOperations( @@ -71,9 +77,14 @@ export function useDiagramStreaming({ // so far. The tool handler validates and loads the final diagram. const handleDisplayChart = useCallback( (xml: string) => { - const completeCells = extractCompleteMxCells(xml || "") + // Named styles come before the cells, so the ones written so + // far are known by the time their cells stream in + const { classes, xml: cellsXml } = readStyleClasses(xml || "") + const completeCells = extractCompleteMxCells(cellsXml) if (!completeCells) return - const convertedXml = convertToLegalXml(completeCells) + const convertedXml = addDefaultStyles( + applyStyleClasses(convertToLegalXml(completeCells), classes), + ) if (convertedXml === previousXML.current) return // Skip this update while the cells written so far don't parse @@ -202,9 +213,17 @@ export function useDiagramStreaming({ if (lastProcessedXmlRef.current.get(opsKey) === ops) continue lastProcessedXmlRef.current.set(opsKey, ops) try { + const edges = edgeIdsOf(originalXml) const { result } = applyDiagramOperations( originalXml, - completeOps, + completeOps.map((op) => + op.new_xml + ? { + ...op, + new_xml: addDefaultStyles(op.new_xml, edges), + } + : op, + ), ) if (lastProcessedXmlRef.current.get(resultKey) === result) { continue diff --git a/hooks/use-diagram-tool-handlers.ts b/hooks/use-diagram-tool-handlers.ts index 22a4ffcb..4f2c5ec8 100644 --- a/hooks/use-diagram-tool-handlers.ts +++ b/hooks/use-diagram-tool-handlers.ts @@ -10,6 +10,7 @@ import { formatValidationFeedback } from "@/lib/diagram-validator" 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 { hasCells } from "@/packages/mcp-server/src/pages.ts" const DEBUG = process.env.NODE_ENV === "development" @@ -575,7 +576,10 @@ Use display_diagram to create the complete diagram, or edit_diagram to change th // Detect if LLM incorrectly started fresh instead of continuing // LLM should only output bare mxCells now, so wrapper tags indicate error const trimmed = xml.trim() + // Definitions at the start are a restart only once cells were written const isFreshStart = + (trimmed.startsWith("\\n \\n"}]} - Example delete: {"operations": [{"operation": "delete", "cell_id": "5"}]} @@ -111,7 +113,7 @@ When using edit_diagram tool: ## Draw.io XML Structure Reference -**IMPORTANT:** You only generate the mxCell elements. The wrapper structure and root cells (id="0", id="1") are added automatically. +**IMPORTANT:** You only generate the named styles and the mxCell elements. The wrapper structure and root cells (id="0", id="1") are added automatically. A named style is written before the cells as ; a cell uses it by putting the name among its style tokens (see Styles). Example - generate ONLY this: \`\`\`xml @@ -121,7 +123,7 @@ Example - generate ONLY this: \`\`\` CRITICAL RULES: -1. Generate ONLY mxCell elements - NO wrapper tags (, , ) +1. Generate ONLY mxStyle definitions and mxCell elements - NO wrapper tags (, , ) 2. Do NOT include root cells (id="0" or id="1") - they are added automatically 3. ALL mxCell elements must be siblings - NEVER nest mxCell inside another mxCell 4. Use unique sequential IDs starting from "2" @@ -129,14 +131,14 @@ CRITICAL RULES: Shape (vertex) example: \`\`\`xml - + \`\`\` Connector (edge) example: \`\`\`xml - + @@ -153,7 +155,7 @@ When creating edges/connectors, you MUST follow these rules to avoid overlapping **Rule 3: Always specify exitX, exitY, entryX, entryY explicitly** - Every edge MUST have these 4 attributes set in the style -- Example: style="edgeStyle=orthogonalEdgeStyle;exitX=1;exitY=0.3;entryX=0;entryY=0.3;endArrow=classic;" +- Example: style="edgeStyle=orthogonalEdgeStyle;exitX=1;exitY=0.3;entryX=0;entryY=0.3;" **Rule 4: Route edges AROUND intermediate shapes (obstacle avoidance) - CRITICAL!** - Before creating an edge, identify ALL shapes positioned between source and target @@ -195,10 +197,14 @@ When creating edges/connectors, you MUST follow these rules to avoid overlapping // Style instructions - only included when minimalStyle is false const STYLE_INSTRUCTIONS = ` -Common styles: -- Shapes: rounded=1 (rounded corners), fillColor=#hex, strokeColor=#hex -- Edges: endArrow=classic/block/open/none, startArrow=none/classic, curved=1, edgeStyle=orthogonalEdgeStyle -- Text: fontSize=14, fontStyle=1 (bold), align=center/left/right +## Styles +Define each style used by several cells ONCE, as a named style before the cells, and use the name in the cells like a CSS class. A cell's style can combine a shape token, a name and overrides; later pairs win. Name only styles that two or more cells share; a style used by one cell stays inline. Names must not be draw.io's own style names: shapes such as text, ellipse, rhombus, swimlane, label, image, and colors such as blue, green, red, gray, yellow, orange, purple, pink. A definition applies to the call it is in: each display_diagram call defines the names it uses. The app expands the names, so the saved file is standard draw.io XML. +\`\`\`xml +${STYLE_CLASS_EXAMPLE} +\`\`\` +- NEVER write html=1 or whiteSpace=wrap: the app adds html=1 to every cell and whiteSpace=wrap to shapes. Labels are HTML: use <br> for a line break and <b> for bold, never \\n; a literal < or > in a label is written &lt; or &gt;. +- Do NOT repeat what draw.io already uses. For a plain shape: rounded=0, align=center, verticalAlign=middle, fontSize=12, strokeWidth=1, fillColor=#ffffff, strokeColor=#000000, fontColor=#000000. For an edge: endArrow=classic, strokeColor=#000000. Writing one of them is right only when it overrides what a name or the shape sets: an edge is rounded by default, so rounded=0 on an edge is a real setting, and a text cell is left/top aligned by default, so there align=center or verticalAlign=middle are real settings. +- Keys: shapes rounded=1, fillColor=#hex, strokeColor=#hex; edges endArrow=block/open/none, startArrow=classic, curved=1, dashed=1, edgeStyle=orthogonalEdgeStyle; text fontSize=14, fontStyle=1 (bold), align=center/right. ` // Minimal style instruction - skip styling and focus on layout (prepended to prompt for emphasis) @@ -208,13 +214,13 @@ const MINIMAL_STYLE_INSTRUCTION = ` ### No Styling - Plain Black/White Only - NO fillColor, NO strokeColor, NO rounded, NO fontSize, NO fontStyle - NO color attributes (no hex colors like #ff69b4) -- Style: "whiteSpace=wrap;html=1;" for shapes, "html=1;endArrow=classic;" for edges +- Shapes: no style, or only the shape (ellipse, rhombus). Edges: edgeStyle=orthogonalEdgeStyle plus the exit/entry points from the Edge Routing Rules, nothing else. html=1 and whiteSpace=wrap are added automatically. - IGNORE all color/style examples below ### Container/Group Shapes - MUST be Transparent - For container shapes (boxes that contain other shapes): use "fillColor=none;" to make background transparent - This prevents containers from covering child elements -- Example: style="whiteSpace=wrap;html=1;fillColor=none;" for container rectangles +- Example: style="fillColor=none;" for container rectangles ### Focus on Layout Quality Since we skip styling, STRICTLY follow the "Edge Routing Rules" section below: @@ -235,7 +241,7 @@ const EXTENDED_ADDITIONS = ` ### display_diagram Details **VALIDATION RULES** (XML will be rejected if violated): -1. Generate ONLY mxCell elements - wrapper tags and root cells are added automatically +1. Generate ONLY mxStyle definitions and mxCell elements - wrapper tags and root cells are added automatically 2. All mxCell elements must be siblings - never nested inside other mxCell elements 3. Every mxCell needs a unique id attribute (start from "2") 4. Every mxCell needs a valid parent attribute (use "1" for top-level, or container-id for grouped) @@ -312,7 +318,7 @@ ${TWO_EDGES_EXAMPLE} ### Edge with single waypoint (simple detour): \`\`\`xml - + diff --git a/lib/utils.ts b/lib/utils.ts index ce27d4bb..e3ea8d13 100644 --- a/lib/utils.ts +++ b/lib/utils.ts @@ -2,6 +2,7 @@ import { type ClassValue, clsx } from "clsx" import * as pako from "pako" import { twMerge } from "tailwind-merge" import { hasCells } from "@/packages/mcp-server/src/pages.ts" +import { readStyleClasses } from "@/packages/mcp-server/src/style-classes.ts" export function cn(...inputs: ClassValue[]) { return twMerge(clsx(inputs)) @@ -41,7 +42,9 @@ export function isRealDiagram(xml: string | undefined | null): boolean { * @returns true if XML appears complete, false if truncated or empty */ export function isMxCellXmlComplete(xml: string | undefined | null): boolean { - const trimmed = xml?.trim() || "" + // Named style definitions before the cells are not cells: output cut off + // right after them is incomplete + const trimmed = readStyleClasses(xml || "").xml.trim() if (!trimmed) return false // Find position of last complete mxCell ending (either /> or ) diff --git a/packages/mcp-server/src/drawing-guide.ts b/packages/mcp-server/src/drawing-guide.ts index 4265a64e..32b7cbd4 100644 --- a/packages/mcp-server/src/drawing-guide.ts +++ b/packages/mcp-server/src/drawing-guide.ts @@ -9,6 +9,7 @@ import { indent, + STYLE_CLASS_EXAMPLE, SWIMLANE_EXAMPLE, TWO_EDGES_EXAMPLE, WAYPOINT_EXAMPLE, @@ -26,12 +27,12 @@ export const DRAWING_GUIDE = `# Draw.io drawing guide - The preview page has History (it saves a snapshot before every AI change and can restore any of the last 20 versions) and Download. You can make changes freely; nothing is lost. ## The XML you send -Single page (create_new_diagram, add_page): send ONLY the mxCell elements. The server adds , , and the root cells id="0" and id="1". +Single page (create_new_diagram, add_page): send ONLY the named styles and the mxCell elements. The server adds , , and the root cells id="0" and id="1", expands named styles (see Styles), adds html=1 to every cell and whiteSpace=wrap to shapes. - + - + @@ -59,7 +60,7 @@ ${indent(SWIMLANE_EXAMPLE)} ## Edge routing rules Rule 1: Never let two edges share a path. Two edges between the same nodes exit and enter at different points (exitY=0.3 for the first, exitY=0.7 for the second, not both 0.5). Rule 2: For bidirectional connections (A to B and B to A), use opposite sides: A exits right (exitX=1) into the left of B (entryX=0); B exits left (exitX=0) into the right of A (entryX=1). -Rule 3: Always set exitX, exitY, entryX and entryY in the edge style, e.g. style="edgeStyle=orthogonalEdgeStyle;exitX=1;exitY=0.3;entryX=0;entryY=0.3;endArrow=classic;". +Rule 3: Always set exitX, exitY, entryX and entryY in the edge style, e.g. style="edgeStyle=orthogonalEdgeStyle;exitX=1;exitY=0.3;entryX=0;entryY=0.3;". Rule 4: Route edges AROUND shapes in the way. Before drawing an edge, find every shape between source and target; if one is in the path, add waypoints. Route diagonal connections along the outside of the diagram, not through the middle. Keep 20-30px clearance from shapes. An edge must never cross another shape's box. Rule 5: Plan the layout first. Organize shapes into rows or columns following the flow, space them 150-200px apart so edges have room, and prefer one flow direction (left to right or top to bottom). Rule 6: Use 2-3 waypoints for L-shaped or U-shaped paths. Each change of direction needs a waypoint, and segments should be horizontal or vertical. @@ -80,20 +81,26 @@ Waypoints go inside in the edge geometry. Example: Hotfix (r ${indent(WAYPOINT_EXAMPLE)} ## Styles -- Shapes: rounded=1, fillColor=#hex, strokeColor=#hex, whiteSpace=wrap;html=1; -- Edges: endArrow=classic, block, open or none; startArrow=none or classic; curved=1; edgeStyle=orthogonalEdgeStyle -- Text: fontSize=14, fontStyle=1 (bold), align=center, left or right +Define each style used by several cells ONCE, as a named style before the cells, and use the name in the cells like a CSS class. A cell's style can combine a shape token, a name and overrides; later pairs win. Name only styles that two or more cells share; a style used by one cell stays inline. Names must not be draw.io's own style names: shapes such as text, ellipse, rhombus, swimlane, label, image, and colors such as blue, green, red, gray, yellow, orange, purple, pink. A definition applies to the call it is in: each create_new_diagram or add_page call defines the names it uses. The server expands the names, so the saved file is standard draw.io XML. + +${indent(STYLE_CLASS_EXAMPLE)} + +- NEVER write html=1 or whiteSpace=wrap: the server adds html=1 to every cell and whiteSpace=wrap to shapes. Labels are HTML: use <br> for a line break and <b> for bold, never \\n; a literal < or > in a label is written &lt; or &gt;. +- Do NOT repeat what draw.io already uses. For a plain shape: rounded=0, align=center, verticalAlign=middle, fontSize=12, strokeWidth=1, fillColor=#ffffff, strokeColor=#000000, fontColor=#000000. For an edge: endArrow=classic, strokeColor=#000000. Writing one of them is right only when it overrides what a name or the shape sets: an edge is rounded by default, so rounded=0 on an edge is a real setting, and a text cell is left/top aligned by default, so there align=center or verticalAlign=middle are real settings. +- Shapes: rounded=1, fillColor=#hex, strokeColor=#hex +- Edges: endArrow=block, open or none; startArrow=classic; curved=1; dashed=1; edgeStyle=orthogonalEdgeStyle +- Text: fontSize=14, fontStyle=1 (bold), align=center or right - Animated connectors: add flowAnimation=1 to the edge style. ## Minimal style When the user asks for a minimal, plain, black-and-white or unstyled diagram, use these rules instead of the styles above: - No fillColor, strokeColor, rounded, fontSize, fontStyle or hex colors. -- Shapes use style "whiteSpace=wrap;html=1;", edges use "html=1;endArrow=classic;". -- Containers that hold other shapes use "whiteSpace=wrap;html=1;fillColor=none;" so they do not cover their children. +- Shapes have no style, or only the shape (ellipse, rhombus); edges have edgeStyle=orthogonalEdgeStyle plus the exit and entry points from the routing rules, nothing else. +- Containers that hold other shapes use "fillColor=none;" so they do not cover their children. - Keep at least 50px between elements, and follow all edge routing rules strictly. ## Editing with edit_diagram -- update replaces a cell: send the complete mxCell including mxGeometry, with the same id as cell_id. +- update replaces a cell: send the complete mxCell including mxGeometry, with the same id as cell_id. Named styles are not available here: write each cell's complete style. - add inserts a new cell with a new id. One cell per operation. - delete removes a cell. Its children and every edge connected to it are deleted too, so give only the container's id. - All-or-nothing: if any operation fails, nothing is applied. A rejected call includes the current XML of the page; rebuild your operations on it and retry. diff --git a/packages/mcp-server/src/edit-diagram.ts b/packages/mcp-server/src/edit-diagram.ts index a77e8e92..b81a4926 100644 --- a/packages/mcp-server/src/edit-diagram.ts +++ b/packages/mcp-server/src/edit-diagram.ts @@ -12,6 +12,7 @@ import { type DiagramOperation, } from "./diagram-operations.ts" import { type PageSelector, projectPage } from "./pages.ts" +import { addDefaultStyles, edgeIdsOf } from "./style-classes.ts" import { validateAndFixXml, validateMxCellStructure } from "./xml-validation.ts" export type EditOutcome = @@ -48,11 +49,19 @@ export function editDiagram( const fixes: string[] = [] const prepared: DiagramOperation[] = [] + // Edges already on the page, so a label added to one is not wrapped + const edges = edgeIdsOf(xml) for (const op of operations) { if (op.operation === "delete" || !op.new_xml) { prepared.push(op) continue } + if (/ 1) { @@ -71,7 +80,12 @@ export function editDiagram( if (check.fixed) { fixes.push(`${op.cell_id}: ${check.fixes.join(", ")}`) } - prepared.push({ ...op, new_xml: check.fixed ?? op.new_xml }) + // The defaults the model is told not to write (style-classes.ts), + // added after the fixes so a repaired cell gets them too + prepared.push({ + ...op, + new_xml: addDefaultStyles(check.fixed ?? op.new_xml, edges), + }) } if (errors.length > 0) return { ok: false, errors, pageError: false } diff --git a/packages/mcp-server/src/index.ts b/packages/mcp-server/src/index.ts index fb8939b9..4e408dae 100644 --- a/packages/mcp-server/src/index.ts +++ b/packages/mcp-server/src/index.ts @@ -48,7 +48,11 @@ import { } from "./http-server.ts" import { parseDrawioFileContent } from "./load-diagram.ts" import { log } from "./logger.ts" -import { prepareNewDiagram, reservedIdError } from "./new-diagram.ts" +import { + prepareNewDiagram, + reservedIdError, + takeStyleDefinitions, +} from "./new-diagram.ts" import { addPageToDoc, deletePageFromDoc, @@ -66,6 +70,7 @@ import { } from "./pages.ts" import { Autosaver, defaultDataDir, expandHome } from "./persistence.ts" import { getShapeLibrary, SHAPE_LIBRARY_LIST } from "./shape-library.ts" +import { addDefaultStyles, applyStyleClasses } from "./style-classes.ts" import { validateAndFixXml } from "./xml-validation.ts" // DOMParser/XMLSerializer globals for the XML helpers (Node has neither) @@ -336,11 +341,11 @@ Before using icon shapes (AWS, Azure, GCP, Kubernetes, Cisco...), call get_shape Accepted xml: 1) Only the mxCell elements of one page (recommended). The server adds , , and the root cells "0" and "1": - + 2) A bare with (one page). 3) A full with one or more pages. Every page's must start with . -Rules: cells are siblings (never nested), ids are unique per page and start from "2", parent="1" for top-level shapes, no XML comments, and shapes stay within x 0 to 800 and y 0 to 600.`, +Rules: cells are siblings (never nested), ids are unique per page and start from "2", parent="1" for top-level shapes, no XML comments, and shapes stay within x 0 to 800 and y 0 to 600. A style used by several cells is defined once with before the cells and used by name (see the drawing guide); html=1 and whiteSpace=wrap are added automatically.`, inputSchema: { xml: z .string() @@ -592,7 +597,7 @@ registerWriteTool( "For add/update, new_xml must be a complete mxCell element including mxGeometry. No XML comments. " + 'Every " inside new_xml must be escaped as \\" in the JSON.\n\n' + "Example - Add a rectangle on the default (first) page:\n" + - '{"operations": [{"operation": "add", "cell_id": "rect-1", "new_xml": ""}]}\n\n' + + '{"operations": [{"operation": "add", "cell_id": "rect-1", "new_xml": ""}]}\n\n' + "Example - Delete a cell on the default page:\n" + '{"operations": [{"operation": "delete", "cell_id": "rect-1"}]}', inputSchema: { @@ -1587,14 +1592,28 @@ registerWriteTool( // If caller provided XML, validate it before splicing it in so we // never get a half-broken mxfile written to the session. - const reserved = xml && reservedIdError(xml) + // Named style definitions come out first and are expanded on the + // validated XML, like prepareNewDiagram + const { + classes, + xml: startXml, + error: styleError, + } = takeStyleDefinitions(xml ?? "") + if (styleError) { + return { + content: [{ type: "text", text: `Error: ${styleError}` }], + isError: true, + } + } + const reserved = startXml && reservedIdError(startXml) if (reserved) { return { content: [{ type: "text", text: `Error: ${reserved}` }], isError: true, } } - let cleanXml: string | undefined = xml && wrapCellsInModel(xml) + let cleanXml: string | undefined = + startXml && wrapCellsInModel(startXml) if (cleanXml) { const { valid, error, fixed, fixes } = validateAndFixXml(cleanXml) @@ -1615,6 +1634,9 @@ registerWriteTool( isError: true, } } + cleanXml = addDefaultStyles( + applyStyleClasses(cleanXml, classes), + ) } let info diff --git a/packages/mcp-server/src/new-diagram.ts b/packages/mcp-server/src/new-diagram.ts index 2cb7b82b..f65cd8d6 100644 --- a/packages/mcp-server/src/new-diagram.ts +++ b/packages/mcp-server/src/new-diagram.ts @@ -2,7 +2,13 @@ * A whole new diagram written by the model, for the create_new_diagram tool * and the web app's display_diagram tool. */ -import { normalizeToMxfile, wrapCellsInModel } from "./pages.ts" +import { hasCells, normalizeToMxfile, wrapCellsInModel } from "./pages.ts" +import { + addDefaultStyles, + applyStyleClasses, + readStyleClasses, + type StyleClasses, +} from "./style-classes.ts" import { readAttributes } from "./xml-attributes.ts" import { validateAndFixXml } from "./xml-validation.ts" @@ -42,23 +48,54 @@ export function reservedIdError(input: string): string | null { } /** - * Bare cells get the wrapper and root cells first, since the strict parser - * rejects several top-level elements. Then the XML is validated and - * auto-fixed while it is still a bare model, where duplicate ids are - * renamed, and finally turned into an . + * The named style definitions are taken out first (style-classes.ts). Bare + * cells then get the wrapper and root cells, since the strict parser rejects + * several top-level elements. Then the XML is validated and auto-fixed while + * it is still a bare model, where duplicate ids are renamed. The names are + * expanded and the default styles added on the fixed XML, so repaired cells + * get them too, and finally it is turned into an . */ +/** + * Take the named style definitions out of the model's XML (style-classes.ts). + * Returns the error for the model when a definition never closed or when + * nothing but definitions was sent. + */ +export function takeStyleDefinitions(input: string): { + classes: StyleClasses + xml: string + error: string | null +} { + const { classes, xml } = readStyleClasses(input) + let error: string | null = null + if (/ before the cells.' + } else if (classes.size > 0 && !hasCells(xml)) { + error = + "Only named style definitions were sent, no cells. Send the mxCell elements after the definitions." + } + return { classes, xml, error } +} + export function prepareNewDiagram( input: string, page: { pageId?: string; pageName?: string } = {}, ): NewDiagram { - const reserved = reservedIdError(input) + const { + classes, + xml: cells, + error: styleError, + } = takeStyleDefinitions(input) + if (styleError) return { ok: false, error: styleError } + const reserved = reservedIdError(cells) if (reserved) return { ok: false, error: reserved } - let xml = wrapCellsInModel(input) + let xml = wrapCellsInModel(cells) const { valid, error, fixed, fixes } = validateAndFixXml(xml) if (fixed) xml = fixed if (!valid) { return { ok: false, error: `XML validation failed - ${error}` } } + xml = addDefaultStyles(applyStyleClasses(xml, classes)) const normalized = normalizeToMxfile(xml, page) if (!normalized) { return { diff --git a/packages/mcp-server/src/style-classes.ts b/packages/mcp-server/src/style-classes.ts new file mode 100644 index 00000000..33681d43 --- /dev/null +++ b/packages/mcp-server/src/style-classes.ts @@ -0,0 +1,232 @@ +/** + * Named styles, written once and used by name like CSS classes. + * + * The model defines a style before the cells and refers to it from any + * number of cells, so a style string is not repeated in every cell: + * + * + * ... + * ... + * + * draw.io resolves a style token without "=" through its stylesheet, which + * is not saved in the file, so the names are expanded here and the saved + * XML is plain draw.io XML. A name without a definition is left as it is: + * draw.io ignores it, unless its own stylesheet knows the name (text, + * ellipse, label, blue, green, ...), which then keeps its meaning. + * + * html=1 and whiteSpace=wrap, which draw.io itself puts on every new shape, + * are added here too, so the model never has to write them. + * + * The callers validate and auto-fix the XML first and rewrite the styles + * after, so a repaired cell (quotes written as entities, a lowercase tag) + * gets its defaults too, and the rewrite only ever sees proper attributes. + */ +import { readAttributes } from "./xml-attributes.ts" + +export type StyleClasses = Map + +// with the whitespace after it, also when +// written as or ; quoted values may hold ">" +const STYLE_DEFINITION = + /"']|"[^"]*"|'[^']*')*?)\s*\/?>(?:\s*<\/mxStyle>)?\s*/gi + +// An mxCell opening tag; the attributes stop before a closing "/" so a +// self-closing cell keeps its "/>" +const CELL_TAG = /"']|"[^"]*"|'[^']*')*?)\s*(\/?)>/g + +// draw.io reads keys and tokens as written, " html" is not "html", so the +// checks are exact too +/** Whether the style sets the key (as `key=` at the start or after a ";") */ +const hasKey = (style: string, key: string) => + new RegExp(`(^|;)${key}=`).test(style) + +/** Whether the style holds this exact token */ +const hasToken = (style: string, token: string) => + style.split(";").includes(token) + +/** + * Replace the names among the tokens with their definitions. The tokens are + * kept exactly as written, empty ones included: a leading ";" is draw.io's + * "no default style" marker, and an entity such as " ends in ";" too. + * A name whose definition is empty is dropped, so it does not leave such a + * marker behind. + */ +function applyToStyle(style: string, classes: StyleClasses): string { + const out: string[] = [] + for (const token of style.split(";")) { + const definition = token.includes("=") + ? undefined + : classes.get(token.trim()) + if (definition === undefined) out.push(token) + else if (definition !== "") out.push(definition) + } + return out.join(";") +} + +/** + * A definition's value goes into a style attribute unchecked by the XML + * validator (the definitions are taken out before it runs), so the two + * characters that would break the attribute are escaped here, the way the + * validator repairs them in cells. + */ +function escapeForAttribute(value: string): string { + return value + .replace( + /&(?!(?:lt|gt|amp|quot|apos|#[0-9]+|#x[0-9a-fA-F]+);)/g, + "&", + ) + .replace(/ { + const attrs = new Map( + readAttributes(attrText).map((a) => [a.name, a.value]), + ) + // "name" may end in the ";" the model is used to; "style" is the + // attribute the model may reach for instead of "value" + const name = attrs.get("name")?.trim().replace(/;+$/, "") + const value = attrs.get("value") ?? attrs.get("style") + if (name && value !== undefined) { + // Without the trailing ";", so a name followed by ";" in a cell + // expands to "...;" and not to "...;;" + classes.set(name, escapeForAttribute(value).replace(/;+$/, "")) + } + return "" + }) + // Two passes resolve names used by other definitions, in either order + for (let pass = 0; pass < 2; pass++) { + for (const [name, value] of classes) { + const others = new Map(classes) + others.delete(name) + classes.set(name, applyToStyle(value, others)) + } + } + return { classes, xml: rest } +} + +/** + * Rewrite the style of every mxCell. The callback returns the new style, or + * undefined to leave the cell alone. A cell without a style attribute gets + * one when the callback returns a style for it. + */ +function rewriteCellStyles( + xml: string, + rewrite: ( + style: string | undefined, + attrs: Map, + ) => string | undefined, +): string { + return xml.replace(CELL_TAG, (tag, attrText: string, selfClose: string) => { + const attributes = readAttributes(attrText) + const attrs = new Map(attributes.map((a) => [a.name, a.value])) + const style = attributes.find((a) => a.name === "style") + const next = rewrite(style?.value, attrs) + if (next === undefined || next === style?.value) return tag + // The value is copied from attributes as written, so only a quote + // from a single-quoted source needs escaping + const written = `style="${next.replace(/"/g, """)}"` + if (!style) { + return `` + } + // Keep the attribute where it was, with the whitespace before it + const space = + attrText.slice(style.start, style.end).match(/^\s*/)?.[0] ?? " " + return `` + }) +} + +/** Expand the names in every cell's style; the overrides after a name still win */ +export function applyStyleClasses(xml: string, classes: StyleClasses): string { + if (classes.size === 0) return xml + return rewriteCellStyles(xml, (style) => + style === undefined ? undefined : applyToStyle(style, classes), + ) +} + +/** The ids of the edges in the XML, whose child vertices are edge labels */ +export function edgeIdsOf(xml: string): Set { + const ids = new Set() + for (const [, attrText] of xml.matchAll(CELL_TAG)) { + const attrs = new Map( + readAttributes(attrText).map((a) => [a.name, a.value]), + ) + const id = attrs.get("id") + if (id && attrs.get("edge") === "1") ids.add(id) + } + return ids +} + +/** + * Whether the label sits outside the shape, as with library icons (their + * styles position the label below) and images, so wrapping it to the + * shape's width would be wrong. + */ +function hasLabelOutside(style: string): boolean { + return ( + hasKey(style, "verticalLabelPosition") || + hasKey(style, "labelPosition") || + hasKey(style, "image") || + hasToken(style, "shape=image") + ) +} + +/** + * draw.io's own defaults for a new shape: html=1 (so labels may hold
+ * and ) and whiteSpace=wrap. Wrapping is not added to shapes whose label + * sits outside (icons, images), to text that sizes itself (autosize=1), to + * edge labels (their label box has no width of its own) or next to html=0 + * (draw.io renders a wrapping label as HTML). Edges get html=1. Root cells + * and cells that already set a key are left alone. + * + * An edit adds cells to a page the XML does not show, so the caller passes + * the ids of the edges already on it (edgeIdsOf), or new edge labels would + * be wrapped like shapes. + */ +export function addDefaultStyles( + xml: string, + knownEdgeIds: Iterable = [], +): string { + const edges = new Set([...knownEdgeIds, ...edgeIdsOf(xml)]) + return rewriteCellStyles(xml, (style, attrs) => { + // A cell marked as both is drawn as an edge + const isEdge = attrs.get("edge") === "1" + const isVertex = attrs.get("vertex") === "1" && !isEdge + if (!isVertex && !isEdge) return undefined + const current = style ?? "" + const isEdgeLabel = + hasToken(current, "edgeLabel") || + attrs.get("connectable") === "0" || + edges.has(attrs.get("parent") ?? "") + const added: string[] = [] + if ( + isVertex && + !hasKey(current, "whiteSpace") && + !hasLabelOutside(current) && + !isEdgeLabel && + !hasToken(current, "autosize=1") && + !hasToken(current, "html=0") + ) { + added.push("whiteSpace=wrap") + } + if (!hasKey(current, "html")) added.push("html=1") + if (added.length === 0) return undefined + const separator = current === "" || current.endsWith(";") ? "" : ";" + return `${current}${separator}${added.join(";")};` + }) +} + +/** Expand the named styles of a model's XML and add the default styles */ +export function expandStyles(xml: string): string { + const { classes, xml: cells } = readStyleClasses(xml) + return addDefaultStyles(applyStyleClasses(cells, classes)) +} diff --git a/packages/mcp-server/src/xml-examples.ts b/packages/mcp-server/src/xml-examples.ts index f1758693..29ef0a94 100644 --- a/packages/mcp-server/src/xml-examples.ts +++ b/packages/mcp-server/src/xml-examples.ts @@ -15,18 +15,18 @@ export const SWIMLANE_EXAMPLE = ` - + ` -export const TWO_EDGES_EXAMPLE = ` +export const TWO_EDGES_EXAMPLE = ` - + ` -export const WAYPOINT_EXAMPLE = ` +export const WAYPOINT_EXAMPLE = ` @@ -35,6 +35,15 @@ export const WAYPOINT_EXAMPLE = ` + + + + + +` + /** Indent every line, for an indented code block in Markdown */ export const indent = (text: string, prefix = " ") => text diff --git a/tests/unit/style-classes.test.ts b/tests/unit/style-classes.test.ts new file mode 100644 index 00000000..d71d244f --- /dev/null +++ b/tests/unit/style-classes.test.ts @@ -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 = + '' +const cell = (id: string, style: string, extra = "") => + `${geometry}` +const edge = (id: string, style: string) => + `` + +const BLUE = + '' + +describe("readStyleClasses", () => { + it("reads the definitions and removes them from the XML", () => { + const xml = `${BLUE}\n\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}` + 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} { + 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;")}${geometry}` + expect(applyStyleClasses(xml, classes)).toBe(xml) + }) + + it("works inside a wrapped model and keeps the attribute order", () => { + const xml = `${geometry}` + 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 = `${geometry}` + expect(addDefaultStyles(xml)).toBe( + `${geometry}`, + ) + }) + + 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 = '' + 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\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 = + '' + expect(addDefaultStyles(edge)).toBe( + '', + ) + const shape = '' + expect(addDefaultStyles(shape)).toBe( + '', + ) + }) + + it("expands a name in a cell wrapped in a UserObject", () => { + const xml = `${BLUE}${geometry}` + expect(expandStyles(xml)).toBe( + `${geometry}`, + ) + }) + + 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( + `\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 = `${cell("2", "rounded=1;")}` + const prepared = prepareNewDiagram(file) + expect(prepared.ok).toBe(true) + if (!prepared.ok) return + expect(prepared.xml).toContain( + '', + ) + 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} { + it("keeps entities, a leading semicolon and a missing trailing one", () => { + const quoted = cell( + "2", + "fontFamily="Times New Roman";fillColor=red;", + ) + expect(addDefaultStyles(quoted)).toBe( + cell( + "2", + "fontFamily="Times New Roman";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 = `${cell("2", "a;")}` + expect(expandStyles(xml)).toBe( + cell("2", "fontFamily="Arial";whiteSpace=wrap;html=1;"), + ) + }) + + it("resolves references in either order; duplicates and built-in names", () => { + const { classes } = readStyleClasses( + ``, + ) + // 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( + '', + ).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: `${geometry}`, + }, + ], + {}, + ) + 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}${geometry}`, + ) + 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 = `${cell("2", "f;")}` + const prepared = prepareNewDiagram(xml) + expect(prepared.ok).toBe(true) + if (!prepared.ok) return + expect(prepared.xml).toContain( + 'style="fontFamily=A&B <C;whiteSpace=wrap;html=1;"', + ) + }) + + it("accepts a lowercase tag, a style attribute and a name ending in ;", () => { + const { classes, xml } = readStyleClasses( + `${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", "")}` + expect(addDefaultStyles(xml)).toContain('style="align=center;html=1;"') + }) + + it("reports a definition that never closed and definitions without cells", () => { + const open = prepareNewDiagram( + `${BLUE} { + 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 = `` + 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 = `` + expect(addDefaultStyles(both)).toContain('style="rounded=1;html=1;"') + }) + + it("drops a name whose definition is empty instead of leaving a leading ;", () => { + const xml = `${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") + }) +}) diff --git a/tests/unit/use-diagram-streaming.test.tsx b/tests/unit/use-diagram-streaming.test.tsx index e551a242..623031b7 100644 --- a/tests/unit/use-diagram-streaming.test.tsx +++ b/tests/unit/use-diagram-streaming.test.tsx @@ -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 = ` + + + 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;"', + ) + }) +})