diff --git a/packages/mcp/src/lib/drawio-layout.ts b/packages/mcp/src/lib/drawio-layout.ts index 7adb166f..647bdca8 100644 --- a/packages/mcp/src/lib/drawio-layout.ts +++ b/packages/mcp/src/lib/drawio-layout.ts @@ -10,7 +10,7 @@ // exactly mxGraph's convention for a child of a container, so they map across // directly. Container sizes are computed by ELK; leaf sizes are preserved. -import ELK from "elkjs/lib/elk.bundled.js"; +import { Worker } from "node:worker_threads"; import { JSDOM } from "jsdom"; import { normalizeInput, parseCells, type DrawioCell } from "./drawio-xml.js"; @@ -18,22 +18,33 @@ import { normalizeInput, parseCells, type DrawioCell } from "./drawio-xml.js"; const DEFAULT_W = 140; const DEFAULT_H = 60; -// DoS bounds for the in-process ELK layout. The mxGraph XML is LLM-supplied -// (layout:"elk" in drawioCreate/drawioUpdate) and elkjs runs synchronously on -// the MCP server's event loop, so an unbounded graph would block it for -// seconds-to-minutes. A ~1MB XML (well under the stage-1 16MB cap) can carry -// thousands of nodes. We cap the graph size and race the layout against a -// wall-clock timeout; on either bound we fall back to the ORIGINAL model, the -// same best-effort contract the catch already honours. +// DoS bounds for the ELK layout. The mxGraph XML is LLM-supplied (layout:"elk" +// in drawioCreate/drawioUpdate). elkjs' layout() returns a Promise but runs the +// crossing-minimisation SYNCHRONOUSLY — it blocks whatever thread it runs on for +// the whole pass. A ~1MB XML (well under the stage-1 16MB cap) can carry +// thousands of nodes. We (a) cap the graph size before ever calling ELK and +// (b) run the layout in a WORKER THREAD so the main event loop stays free, with +// the wall-clock timeout enforced by terminating that worker. On either bound we +// fall back to the ORIGINAL model, the same best-effort contract the catch honours. // - 500 nodes lays out in well under a second; beyond that ELK cost climbs // steeply, so refuse and leave the (already-valid) model untouched. // - Edges dominate the layered-crossing cost, so allow a bit more headroom // (1000) than nodes but still bound them. -// - 5s is generous for any graph within the caps yet short enough that a -// pathological input can never wedge the server. +// - The timeout is a HARD kill of the worker thread — the only way to interrupt +// synchronous JS. The in-process setTimeout race we used before was an +// illusion: the timer could never fire while the SAME thread was blocked +// inside elkjs, so it "protected" nothing. Now the timer runs on the main +// thread while ELK runs on the worker, so it can actually fire and terminate. const ELK_MAX_NODES = 500; const ELK_MAX_EDGES = 1000; -const ELK_TIMEOUT_MS = 5000; +// Wall-clock ceiling for a single layout pass. Overridable for tests (a tiny +// value forces the terminate-on-timeout path deterministically); a non-positive +// or unparseable override falls back to the default. +const ELK_TIMEOUT_DEFAULT_MS = 5000; +function resolveElkTimeoutMs(): number { + const raw = Number(process.env.DRAWIO_ELK_TIMEOUT_MS); + return Number.isFinite(raw) && raw > 0 ? Math.floor(raw) : ELK_TIMEOUT_DEFAULT_MS; +} // Spacing is set >=150px on purpose so an ELK layout never trips the linter's // "gap between adjacent shapes < 150px" quality warning (acceptance #3). @@ -78,13 +89,57 @@ interface ElkGraph extends ElkNode { edges?: ElkEdge[]; } +/** + * Run one ELK layered layout on a worker thread and resolve with the laid-out + * graph. The timeout is enforced by `worker.terminate()` — a HARD kill, which is + * the only way to interrupt elkjs' synchronous crossing-minimisation once it has + * started. Rejects on timeout, worker error, or an early exit; the caller treats + * any rejection as "keep the original model" (best-effort layout). + */ +function layoutInWorker(graph: ElkGraph, timeoutMs: number): Promise { + return new Promise((resolve, reject) => { + const worker = new Worker( + new URL("./drawio-layout.worker.js", import.meta.url), + { workerData: { graph } }, + ); + let settled = false; + const finish = (fn: () => void) => { + if (settled) return; + settled = true; + clearTimeout(timer); + // Always tear the worker down: on the happy path so it does not linger, + // on timeout so the blocked synchronous ELK run is actually interrupted. + void worker.terminate(); + fn(); + }; + const timer = setTimeout( + () => finish(() => reject(new Error("ELK layout timed out"))), + timeoutMs, + ); + worker.once("message", (msg: { ok?: boolean; laid?: ElkGraph; error?: string }) => { + finish(() => + msg?.ok + ? resolve(msg.laid as ElkGraph) + : reject(new Error(msg?.error ?? "ELK layout failed")), + ); + }); + worker.once("error", (err) => finish(() => reject(err))); + worker.once("exit", (code) => { + // A clean exit after we already settled is normal (terminate()); only an + // unexpected early exit while still pending is a failure. + if (settled) return; + finish(() => reject(new Error(`ELK worker exited early (code ${code})`))); + }); + }); +} + /** * Apply an ELK layered layout to a drawio input and return a full mxGraphModel * string with rewritten geometry. Accepts the same three input forms as * drawioCreate (a bare model, an , or a list). Async because - * elkjs' layout() is promise-based. On any layout failure the ORIGINAL - * (normalized) model is returned unchanged — layout is best-effort polish, never - * a reason to fail the write. + * the layout runs on a worker thread. On any layout failure (including a + * terminate-on-timeout) the ORIGINAL (normalized) model is returned unchanged — + * layout is best-effort polish, never a reason to fail the write. */ export async function applyElkLayout(inputXml: string): Promise { const modelXml = normalizeInput(inputXml); @@ -150,26 +205,14 @@ export async function applyElkLayout(inputXml: string): Promise { }; let laid: ElkGraph; - let timer: ReturnType | undefined; try { - // elkjs ships a CJS default export whose interop shape varies across - // module systems; resolve the real constructor at runtime, then cast (the - // runtime call is verified — see the layout unit test). - const Ctor: any = (ELK as any).default ?? ELK; - const elk = new Ctor(); - // Race the layout against a wall-clock timeout so a graph that is under the - // node/edge caps but still pathologically slow can never wedge the server. - const timeout = new Promise((_, reject) => { - timer = setTimeout( - () => reject(new Error("ELK layout timed out")), - ELK_TIMEOUT_MS, - ); - }); - laid = (await Promise.race([elk.layout(graph as any), timeout])) as ElkGraph; + // Run the (synchronous-under-the-hood) ELK pass on a worker thread so the + // main event loop is never blocked, and enforce the wall-clock ceiling by + // terminating that worker on timeout. A graph under the node/edge caps but + // still pathologically slow is hard-killed instead of wedging anything. + laid = await layoutInWorker(graph, resolveElkTimeoutMs()); } catch { return modelXml; // best-effort: keep the model as-is on timeout or ELK failure - } finally { - if (timer) clearTimeout(timer); } // Collect computed geometry per node id (coords are parent-relative already). diff --git a/packages/mcp/src/lib/drawio-layout.worker.ts b/packages/mcp/src/lib/drawio-layout.worker.ts new file mode 100644 index 00000000..0313be93 --- /dev/null +++ b/packages/mcp/src/lib/drawio-layout.worker.ts @@ -0,0 +1,36 @@ +// Worker-thread entry for the ELK layered layout (issue #486, commit 1). +// +// elkjs' layout() returns a Promise but runs the actual crossing-minimisation +// SYNCHRONOUSLY — it blocks whatever thread it runs on for the whole pass. On +// the in-app MCP host that thread used to be the main NestJS event loop, so a +// pathological graph at the node/edge cap could wedge ALL HTTP/SSE/loopback +// traffic while it churned. Running it HERE, on a dedicated worker thread, keeps +// the main loop free; the parent enforces the wall-clock timeout by calling +// `worker.terminate()` — the only way to interrupt synchronous JS — since the +// in-process `setTimeout` race the parent used before could never fire while the +// same thread was blocked inside elkjs. +import { parentPort, workerData } from "node:worker_threads"; +import ELK from "elkjs/lib/elk.bundled.js"; + +interface WorkerInput { + graph: unknown; +} + +const { graph } = (workerData ?? {}) as WorkerInput; + +// elkjs ships a CJS default export whose interop shape varies across module +// systems; resolve the real constructor at runtime (same as the parent did). +const Ctor: any = (ELK as any).default ?? ELK; +const elk = new Ctor(); + +elk + .layout(graph as any) + .then((laid: unknown) => { + parentPort?.postMessage({ ok: true, laid }); + }) + .catch((err: unknown) => { + parentPort?.postMessage({ + ok: false, + error: err instanceof Error ? err.message : String(err), + }); + }); diff --git a/packages/mcp/test/unit/drawio-layout.test.mjs b/packages/mcp/test/unit/drawio-layout.test.mjs index a6998d17..ccf5eda3 100644 --- a/packages/mcp/test/unit/drawio-layout.test.mjs +++ b/packages/mcp/test/unit/drawio-layout.test.mjs @@ -101,6 +101,88 @@ test("DoS guard: a graph over the node cap is returned unchanged, quickly", asyn assert.ok(dt < 2000, `cap path should be fast, took ${dt}ms`); }); +/** Build a layered DAG near the caps: `n` vertices, up to ~2 edges each into the + * next layer of `layerSize`. Used as a real worst-case graph for the benchmark. */ +function layeredGraph(n, layerSize) { + let cells = ""; + for (let i = 2; i < 2 + n; i++) { + cells += + `` + + ``; + } + let ei = 0; + for (let i = 2; i < 2 + n; i++) { + for (const off of [layerSize, layerSize + 1]) { + const t = i + off; + if (t < 2 + n) cells += ``; + } + } + return ( + '' + + cells + + "" + ); +} + +test("terminate-on-timeout: a layout that exceeds the wall-clock ceiling is hard-killed and the original model is returned (#486)", async () => { + // A 1ms ceiling fires before the worker can even finish loading elkjs, so the + // parent must terminate() the worker and fall back to the ORIGINAL model. On + // the OLD in-process race this timer could never fire while the SAME thread was + // blocked inside elkjs — the fallback path was unreachable; here it works. + const prev = process.env.DRAWIO_ELK_TIMEOUT_MS; + process.env.DRAWIO_ELK_TIMEOUT_MS = "1"; + try { + const model = layeredGraph(400, 20); + const t0 = Date.now(); + const laid = await applyElkLayout(model); + const dt = Date.now() - t0; + // Original geometry is preserved verbatim: every vertex is still stacked at + // (10,10), proving NO ELK coordinates were applied (the pass was killed). + const verts = parseCells(laid).filter((c) => c.vertex); + assert.equal(verts.length, 400, "all vertices survived the fallback"); + for (const v of verts) { + assert.equal(v.geometry.x, 10, "x untouched -> layout was terminated"); + assert.equal(v.geometry.y, 10, "y untouched -> layout was terminated"); + } + // The kill is prompt: terminate() returns the call well under the natural + // layout time for a 400-node graph. + assert.ok(dt < 2000, `terminate path should be prompt, took ${dt}ms`); + } finally { + if (prev === undefined) delete process.env.DRAWIO_ELK_TIMEOUT_MS; + else process.env.DRAWIO_ELK_TIMEOUT_MS = prev; + } +}); + +test("benchmark guard: a worst-case graph AT the cap lays out without wedging the main event loop (#486)", async () => { + // ~500 nodes / ~1000 edges — a real worst case at the node/edge caps. The + // layout runs on a WORKER thread, so the MAIN event loop must stay responsive + // throughout: a timer scheduled on the main thread keeps firing while ELK + // churns. On the OLD synchronous-on-main-thread code this counter would be + // pinned at 0 for the whole layout (event loop wedged) — exactly the prod fire. + const model = layeredGraph(500, 20); + let mainLoopTicks = 0; + const iv = setInterval(() => { + mainLoopTicks++; + }, 2); + const t0 = Date.now(); + const laid = await applyElkLayout(model); + const dt = Date.now() - t0; + clearInterval(iv); + + assert.ok( + mainLoopTicks > 0, + "main event loop must stay responsive while ELK runs on the worker", + ); + // Benchmark guard: the worst-case graph actually LAYS OUT within the default + // ceiling (it did not fall back). At least one vertex moved off the stack. + const verts = parseCells(laid).filter((c) => c.vertex); + assert.equal(verts.length, 500, "all vertices survived"); + const moved = verts.some((v) => v.geometry.x !== 10 || v.geometry.y !== 10); + assert.ok(moved, "layout was applied (did not time out / fall back)"); + // Sanity ceiling well under the 5s wall-clock timeout. + assert.ok(dt < 5000, `worst-case layout should be under the ceiling, took ${dt}ms`); +}); + test("layout is best-effort: an empty/degenerate model is returned intact", async () => { const model = '';