mirror of
https://github.com/DayuanJiang/next-ai-draw-io.git
synced 2026-10-11 03:59:58 +08:00
refactor(web): edit and wrap diagrams with the MCP server's code
- edit_diagram runs the MCP server's editDiagram: every new_xml is checked first, one cell per operation, and after the edit only the target page is checked, rejecting only errors this edit introduced. An unrelated problem elsewhere in the document no longer blocks every edit. The error lists each failed operation - The streaming edit preview uses the MCP applyDiagramOperations - Delete applyDiagramOperations (292 lines) and wrapWithMxFile from lib/utils.ts, and the unused hand-copied scripts/test-diagram-operations.mjs - One blank document (BLANK_MXFILE) for the web app and the MCP preview, replacing four copies - Saving a .drawio wraps a bare model with normalizeToMxfile - The empty-diagram check uses hasCells, which also counts cells wrapped in a UserObject/object - DiagramOperation is the MCP type - The wrapped-cell and empty-diagram tests now run against the MCP code - New e2e test: edit_diagram changes the canvas, and a failing edit leaves it as it was
This commit is contained in:
dayuan.jiang
@@ -1,34 +1,45 @@
|
||||
import { expect, test } from "@playwright/test"
|
||||
import { expect, type Page, test } from "@playwright/test"
|
||||
import { getIframe, sendMessage, waitForCompleteCount } from "./lib/fixtures"
|
||||
|
||||
/**
|
||||
* Checks what draw.io actually shows after display_diagram, not only the
|
||||
* Checks what draw.io actually shows after the diagram tools, not only the
|
||||
* tool card. The tool input is streamed in chunks like a real model, and
|
||||
* the browser tool handler (not the server) completes the tool call.
|
||||
*/
|
||||
function streamedToolCall(xml: string) {
|
||||
function streamedToolCall(toolName: string, input: unknown) {
|
||||
const toolCallId = `call_${Math.random().toString(36).slice(2)}`
|
||||
const input = JSON.stringify({ xml })
|
||||
const chunks = input.match(/[\s\S]{1,40}/g) ?? []
|
||||
const chunks = JSON.stringify(input).match(/[\s\S]{1,40}/g) ?? []
|
||||
const events = [
|
||||
{ type: "start", messageId: `msg_${toolCallId}` },
|
||||
{ type: "tool-input-start", toolCallId, toolName: "display_diagram" },
|
||||
{ type: "tool-input-start", toolCallId, toolName },
|
||||
...chunks.map((inputTextDelta) => ({
|
||||
type: "tool-input-delta",
|
||||
toolCallId,
|
||||
inputTextDelta,
|
||||
})),
|
||||
{
|
||||
type: "tool-input-available",
|
||||
toolCallId,
|
||||
toolName: "display_diagram",
|
||||
input: { xml },
|
||||
},
|
||||
{ type: "tool-input-available", toolCallId, toolName, input },
|
||||
{ type: "finish" },
|
||||
]
|
||||
return `${events.map((e) => `data: ${JSON.stringify(e)}\n\n`).join("")}data: [DONE]\n\n`
|
||||
}
|
||||
|
||||
const END_TURN =
|
||||
'data: {"type":"start"}\n\ndata: {"type":"finish"}\n\ndata: [DONE]\n\n'
|
||||
|
||||
/** Answer each chat request with the next reply, then end the turn */
|
||||
async function mockReplies(p: Page, replies: string[]) {
|
||||
await p.route("**/api/chat", async (route) => {
|
||||
await route.fulfill({
|
||||
status: 200,
|
||||
contentType: "text/event-stream",
|
||||
body: replies.shift() ?? END_TURN,
|
||||
})
|
||||
})
|
||||
await p.goto("/", { waitUntil: "networkidle" })
|
||||
await getIframe(p).waitFor({ state: "visible", timeout: 30000 })
|
||||
return p.frameLocator("iframe")
|
||||
}
|
||||
|
||||
const cell = (id: string, label: string, x: number) =>
|
||||
`<mxCell id="${id}" value="${label}" style="rounded=1;" vertex="1" parent="1"><mxGeometry x="${x}" y="40" width="120" height="60" as="geometry"/></mxCell>`
|
||||
const page = (id: string, cells: string) =>
|
||||
@@ -46,20 +57,10 @@ const NEW_CELLS =
|
||||
test("display_diagram replaces the document with the fixed diagram", async ({
|
||||
page: p,
|
||||
}) => {
|
||||
const replies = [TWO_PAGES, NEW_CELLS]
|
||||
await p.route("**/api/chat", async (route) => {
|
||||
const xml = replies.shift()
|
||||
await route.fulfill({
|
||||
status: 200,
|
||||
contentType: "text/event-stream",
|
||||
body: xml
|
||||
? streamedToolCall(xml)
|
||||
: 'data: {"type":"start"}\n\ndata: {"type":"finish"}\n\ndata: [DONE]\n\n',
|
||||
})
|
||||
})
|
||||
await p.goto("/", { waitUntil: "networkidle" })
|
||||
await getIframe(p).waitFor({ state: "visible", timeout: 30000 })
|
||||
const canvas = p.frameLocator("iframe")
|
||||
const canvas = await mockReplies(p, [
|
||||
streamedToolCall("display_diagram", { xml: TWO_PAGES }),
|
||||
streamedToolCall("display_diagram", { xml: NEW_CELLS }),
|
||||
])
|
||||
|
||||
await sendMessage(p, "Draw two pages")
|
||||
await waitForCompleteCount(p, 1)
|
||||
@@ -79,3 +80,59 @@ test("display_diagram replaces the document with the fixed diagram", async ({
|
||||
await expect(canvas.getByText("Old A")).toHaveCount(0)
|
||||
await expect(canvas.getByText("Second", { exact: true })).toHaveCount(0)
|
||||
})
|
||||
|
||||
test("edit_diagram applies all operations or none", async ({ page: p }) => {
|
||||
const canvas = await mockReplies(p, [
|
||||
streamedToolCall("display_diagram", {
|
||||
xml: cell("a", "Alpha", 40) + cell("b", "Beta", 220),
|
||||
}),
|
||||
streamedToolCall("edit_diagram", {
|
||||
operations: [
|
||||
{ operation: "delete", cell_id: "a" },
|
||||
{
|
||||
operation: "update",
|
||||
cell_id: "b",
|
||||
new_xml: cell("b", "Beta two", 220),
|
||||
},
|
||||
{
|
||||
operation: "add",
|
||||
cell_id: "c",
|
||||
new_xml: cell("c", "Gamma", 400),
|
||||
},
|
||||
],
|
||||
}),
|
||||
// The first operation is fine, the second fails: nothing is kept
|
||||
streamedToolCall("edit_diagram", {
|
||||
operations: [
|
||||
{
|
||||
operation: "update",
|
||||
cell_id: "c",
|
||||
new_xml: cell("c", "Broken", 400),
|
||||
},
|
||||
{ operation: "delete", cell_id: "missing" },
|
||||
],
|
||||
}),
|
||||
])
|
||||
|
||||
await sendMessage(p, "Draw two boxes")
|
||||
await waitForCompleteCount(p, 1)
|
||||
await expect(canvas.getByText("Alpha", { exact: true })).toBeVisible({
|
||||
timeout: 15000,
|
||||
})
|
||||
|
||||
await sendMessage(p, "Change them")
|
||||
await waitForCompleteCount(p, 2)
|
||||
await expect(canvas.getByText("Gamma", { exact: true })).toBeVisible({
|
||||
timeout: 15000,
|
||||
})
|
||||
await expect(canvas.getByText("Beta two", { exact: true })).toBeVisible()
|
||||
await expect(canvas.getByText("Alpha", { exact: true })).toHaveCount(0)
|
||||
|
||||
await sendMessage(p, "Change again")
|
||||
await expect(p.getByText(/No changes were made/).first()).toBeAttached({
|
||||
timeout: 15000,
|
||||
})
|
||||
await p.waitForTimeout(1000)
|
||||
await expect(canvas.getByText("Gamma", { exact: true })).toBeVisible()
|
||||
await expect(canvas.getByText("Broken", { exact: true })).toHaveCount(0)
|
||||
})
|
||||
|
||||
@@ -6,7 +6,6 @@ import { describe, expect, it } from "vitest"
|
||||
import {
|
||||
dropInvalidToolCalls,
|
||||
fixToolInputJson,
|
||||
isMinimalDiagram,
|
||||
replaceHistoricalToolInputs,
|
||||
validateFileParts,
|
||||
} from "@/lib/chat-helpers"
|
||||
@@ -95,36 +94,6 @@ describe("validateFileParts", () => {
|
||||
})
|
||||
})
|
||||
|
||||
describe("isMinimalDiagram", () => {
|
||||
it("returns true for empty diagram", () => {
|
||||
const xml = '<mxCell id="0"/><mxCell id="1" parent="0"/>'
|
||||
expect(isMinimalDiagram(xml)).toBe(true)
|
||||
})
|
||||
|
||||
it("returns false for diagram with content", () => {
|
||||
const xml =
|
||||
'<mxCell id="0"/><mxCell id="1" parent="0"/><mxCell id="2" value="Hello"/>'
|
||||
expect(isMinimalDiagram(xml)).toBe(false)
|
||||
})
|
||||
|
||||
it("handles whitespace correctly", () => {
|
||||
const xml = ' <mxCell id="0"/> <mxCell id="1" parent="0"/> '
|
||||
expect(isMinimalDiagram(xml)).toBe(true)
|
||||
})
|
||||
|
||||
it("returns false for a shape drawn in draw.io with a random id", () => {
|
||||
const xml =
|
||||
'<mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/><mxCell id="xY3kQ9-1" value="" style="rounded=0;" vertex="1" parent="1"><mxGeometry x="10" y="10" width="120" height="60" as="geometry"/></mxCell></root></mxGraphModel>'
|
||||
expect(isMinimalDiagram(xml)).toBe(false)
|
||||
})
|
||||
|
||||
it("does not mistake ids that start with 0 or 1 for root cells", () => {
|
||||
const xml =
|
||||
'<mxCell id="0"/><mxCell id="1" parent="0"/><mxCell id="10"/>'
|
||||
expect(isMinimalDiagram(xml)).toBe(false)
|
||||
})
|
||||
})
|
||||
|
||||
describe("replaceHistoricalToolInputs", () => {
|
||||
it("replaces display_diagram tool inputs with placeholder", () => {
|
||||
const messages = [
|
||||
|
||||
@@ -180,3 +180,86 @@ describe("autoFixXml", () => {
|
||||
expect(result.fixed).toContain('<object id="3"')
|
||||
})
|
||||
})
|
||||
|
||||
describe("hasCells (was isMinimalDiagram in the web app)", () => {
|
||||
it("returns true for empty diagram", () => {
|
||||
const xml = '<mxCell id="0"/><mxCell id="1" parent="0"/>'
|
||||
expect(hasCells(xml)).toBe(false)
|
||||
})
|
||||
|
||||
it("returns false for diagram with content", () => {
|
||||
const xml =
|
||||
'<mxCell id="0"/><mxCell id="1" parent="0"/><mxCell id="2" value="Hello"/>'
|
||||
expect(hasCells(xml)).toBe(true)
|
||||
})
|
||||
|
||||
it("handles whitespace correctly", () => {
|
||||
const xml = ' <mxCell id="0"/> <mxCell id="1" parent="0"/> '
|
||||
expect(hasCells(xml)).toBe(false)
|
||||
})
|
||||
|
||||
it("returns false for a shape drawn in draw.io with a random id", () => {
|
||||
const xml =
|
||||
'<mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/><mxCell id="xY3kQ9-1" value="" style="rounded=0;" vertex="1" parent="1"><mxGeometry x="10" y="10" width="120" height="60" as="geometry"/></mxCell></root></mxGraphModel>'
|
||||
expect(hasCells(xml)).toBe(true)
|
||||
})
|
||||
|
||||
it("does not mistake ids that start with 0 or 1 for root cells", () => {
|
||||
const xml =
|
||||
'<mxCell id="0"/><mxCell id="1" parent="0"/><mxCell id="10"/>'
|
||||
expect(hasCells(xml)).toBe(true)
|
||||
})
|
||||
|
||||
it("counts a cell wrapped in a UserObject", () => {
|
||||
const xml =
|
||||
'<mxCell id="0"/><mxCell id="1" parent="0"/><UserObject id="u" link="x"><mxCell vertex="1" parent="1"/></UserObject>'
|
||||
expect(hasCells(xml)).toBe(true)
|
||||
})
|
||||
})
|
||||
|
||||
describe("applyDiagramOperations with wrapped cells", () => {
|
||||
const xml = `<mxfile><diagram id="p1"><mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/><UserObject id="5" label="Docs" link="https://example.com"><mxCell vertex="1" parent="1"><mxGeometry as="geometry"/></mxCell></UserObject><mxCell id="6" value="B" vertex="1" parent="1"><mxGeometry as="geometry"/></mxCell><mxCell id="e1" edge="1" parent="1" source="5" target="6"><mxGeometry relative="1" as="geometry"/></mxCell></root></mxGraphModel></diagram></mxfile>`
|
||||
|
||||
it("deletes a wrapped cell and its edges", () => {
|
||||
const { result, errors } = applyDiagramOperations(xml, [
|
||||
{ operation: "delete", cell_id: "5" },
|
||||
{ operation: "delete", cell_id: "e1" },
|
||||
])
|
||||
expect(errors).toEqual([])
|
||||
expect(result).not.toContain("UserObject")
|
||||
expect(result).not.toContain('id="e1"')
|
||||
expect(result).toContain('id="6"')
|
||||
})
|
||||
|
||||
it("rejects adding a cell with the id of a wrapped cell", () => {
|
||||
const { errors } = applyDiagramOperations(xml, [
|
||||
{
|
||||
operation: "add",
|
||||
cell_id: "5",
|
||||
new_xml: '<mxCell id="5" vertex="1" parent="1"/>',
|
||||
},
|
||||
])
|
||||
expect(errors[0]?.message).toContain("already exists")
|
||||
})
|
||||
|
||||
it("updates a wrapped cell", () => {
|
||||
const { result, errors } = applyDiagramOperations(xml, [
|
||||
{
|
||||
operation: "update",
|
||||
cell_id: "5",
|
||||
new_xml:
|
||||
'<UserObject id="5" label="New" link="https://example.org"><mxCell vertex="1" parent="1"><mxGeometry as="geometry"/></mxCell></UserObject>',
|
||||
},
|
||||
])
|
||||
expect(errors).toEqual([])
|
||||
expect(result).toContain('label="New"')
|
||||
expect(result).not.toContain('label="Docs"')
|
||||
})
|
||||
|
||||
it("reports deleting a cell that does not exist", () => {
|
||||
const { errors } = applyDiagramOperations(xml, [
|
||||
{ operation: "delete", cell_id: "missing" },
|
||||
])
|
||||
expect(errors[0]?.message).toContain("not found")
|
||||
})
|
||||
})
|
||||
|
||||
@@ -1,11 +1,5 @@
|
||||
import { describe, expect, it } from "vitest"
|
||||
import {
|
||||
applyDiagramOperations,
|
||||
cn,
|
||||
extractCompleteMxCells,
|
||||
isMxCellXmlComplete,
|
||||
wrapWithMxFile,
|
||||
} from "@/lib/utils"
|
||||
import { cn, extractCompleteMxCells, isMxCellXmlComplete } from "@/lib/utils"
|
||||
|
||||
describe("isMxCellXmlComplete", () => {
|
||||
it("returns false for empty/null input", () => {
|
||||
@@ -71,36 +65,6 @@ describe("isMxCellXmlComplete", () => {
|
||||
})
|
||||
})
|
||||
|
||||
describe("wrapWithMxFile", () => {
|
||||
it("wraps empty string with default structure", () => {
|
||||
const result = wrapWithMxFile("")
|
||||
expect(result).toContain("<mxfile>")
|
||||
expect(result).toContain("<mxGraphModel>")
|
||||
expect(result).toContain('<mxCell id="0"/>')
|
||||
expect(result).toContain('<mxCell id="1" parent="0"/>')
|
||||
})
|
||||
|
||||
it("wraps raw mxCell content", () => {
|
||||
const xml = '<mxCell id="2" value="Hello"/>'
|
||||
const result = wrapWithMxFile(xml)
|
||||
expect(result).toContain("<mxfile>")
|
||||
expect(result).toContain(xml)
|
||||
expect(result).toContain("</mxfile>")
|
||||
})
|
||||
|
||||
it("returns full mxfile unchanged", () => {
|
||||
const fullXml =
|
||||
'<mxfile><diagram name="Page-1"><mxGraphModel></mxGraphModel></diagram></mxfile>'
|
||||
const result = wrapWithMxFile(fullXml)
|
||||
expect(result).toBe(fullXml)
|
||||
})
|
||||
|
||||
it("handles whitespace in input", () => {
|
||||
const result = wrapWithMxFile(" ")
|
||||
expect(result).toContain("<mxfile>")
|
||||
})
|
||||
})
|
||||
|
||||
describe("cn (class name utility)", () => {
|
||||
it("merges class names", () => {
|
||||
expect(cn("foo", "bar")).toBe("foo bar")
|
||||
@@ -132,50 +96,3 @@ describe("extractCompleteMxCells", () => {
|
||||
)
|
||||
})
|
||||
})
|
||||
|
||||
describe("applyDiagramOperations with wrapped cells", () => {
|
||||
const xml = `<mxfile><diagram id="p1"><mxGraphModel><root><mxCell id="0"/><mxCell id="1" parent="0"/><UserObject id="5" label="Docs" link="https://example.com"><mxCell vertex="1" parent="1"><mxGeometry as="geometry"/></mxCell></UserObject><mxCell id="6" value="B" vertex="1" parent="1"><mxGeometry as="geometry"/></mxCell><mxCell id="e1" edge="1" parent="1" source="5" target="6"><mxGeometry relative="1" as="geometry"/></mxCell></root></mxGraphModel></diagram></mxfile>`
|
||||
|
||||
it("deletes a wrapped cell and its edges", () => {
|
||||
const { result, errors } = applyDiagramOperations(xml, [
|
||||
{ operation: "delete", cell_id: "5" },
|
||||
{ operation: "delete", cell_id: "e1" },
|
||||
])
|
||||
expect(errors).toEqual([])
|
||||
expect(result).not.toContain("UserObject")
|
||||
expect(result).not.toContain('id="e1"')
|
||||
expect(result).toContain('id="6"')
|
||||
})
|
||||
|
||||
it("rejects adding a cell with the id of a wrapped cell", () => {
|
||||
const { errors } = applyDiagramOperations(xml, [
|
||||
{
|
||||
operation: "add",
|
||||
cell_id: "5",
|
||||
new_xml: '<mxCell id="5" vertex="1" parent="1"/>',
|
||||
},
|
||||
])
|
||||
expect(errors[0]?.message).toContain("already exists")
|
||||
})
|
||||
|
||||
it("updates a wrapped cell", () => {
|
||||
const { result, errors } = applyDiagramOperations(xml, [
|
||||
{
|
||||
operation: "update",
|
||||
cell_id: "5",
|
||||
new_xml:
|
||||
'<UserObject id="5" label="New" link="https://example.org"><mxCell vertex="1" parent="1"><mxGeometry as="geometry"/></mxCell></UserObject>',
|
||||
},
|
||||
])
|
||||
expect(errors).toEqual([])
|
||||
expect(result).toContain('label="New"')
|
||||
expect(result).not.toContain('label="Docs"')
|
||||
})
|
||||
|
||||
it("reports deleting a cell that does not exist", () => {
|
||||
const { errors } = applyDiagramOperations(xml, [
|
||||
{ operation: "delete", cell_id: "missing" },
|
||||
])
|
||||
expect(errors[0]?.message).toContain("not found")
|
||||
})
|
||||
})
|
||||
|
||||
Reference in New Issue
Block a user