drawio-mcp-server 1.7.0 → 2.0.3

package/build/index.js

@@ -3,20 +3,24 @@ import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
 import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
 import { WebStandardStreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/webStandardStreamableHttp.js";
 import { serve } from "@hono/node-server";
-import { z } from "zod";
 import { Hono } from "hono";
 import { cors } from "hono/cors";
 import EventEmitter from "node:events";
 import { createServer } from "node:net";
+import { join } from "node:path";
+import { fileURLToPath } from "node:url";
+import { readFileSync, existsSync, statSync, readdirSync, realpathSync, } from "node:fs";
 import { WebSocket, WebSocketServer } from "ws";
-import { buildConfig, shouldShowHelp } from "./config.js";
+const VERSION = process.env.npm_package_version ?? "2.0.0";
+import { buildConfig, shouldShowHelp, getHttpFeatureConfig, } from "./config.js";
 import { bus_reply_stream, bus_request_stream, } from "./types.js";
 import { create_bus } from "./emitter_bus.js";
-import { default_tool } from "./tool.js";
 import { nanoid_id_generator } from "./nanoid_id_generator.js";
 import { create_logger as create_console_logger } from "./mcp_console_logger.js";
 import { create_logger as create_server_logger, validLogLevels, } from "./mcp_server_logger.js";
-const VERSION = "1.7.0";
+import { getLocalPluginPath, isUsingLocalAssets, getAssetRoot, ensureAssets, } from "./assets/index.js";
+import { registerTools } from "./tools/index.js";
+const fatalLog = create_console_logger();
 /**
  * Display help message and exit
  */
@@ -28,12 +32,19 @@ Usage: drawio-mcp-server [options]
 
 Options:
   --extension-port, -p <number>  WebSocket server port for browser extension (default: 3333)
+  --editor, -e                   Enable draw.io editor endpoint
+  --http-port                    HTTP server port for Streamable HTTP transport (default: 3000)
+  --transport                    Transport type: stdio, http (default: stdio)
+  --asset-path <path>           Custom path for downloaded assets
   --help, -h                     Show this help message
 
 Examples:
   drawio-mcp-server                           # Use default extension port 3333
   drawio-mcp-server --extension-port 8080     # Use custom extension port 8080
   drawio-mcp-server -p 8080                   # Short form
+  drawio-mcp-server --editor                  # Enable draw.io editor endpoint
+  drawio-mcp-server -e --http                 # Enable editor and HTTP transport
+  drawio-mcp-server --editor --asset-path /data/assets # Use custom asset path
   `);
     process.exit(0);
 }
@@ -47,65 +58,6 @@ async function checkPortAvailable(port) {
         server.on("error", () => resolve(false));
     });
 }
-const emitter = new EventEmitter();
-const conns = new Set();
-const bus_to_ws_forwarder_listener = (event) => {
-    log.debug(`[bridge] received; forwarding message to #${conns.size} clients`, event);
-    for (const ws of [...conns]) {
-        if (ws.readyState !== WebSocket.OPEN) {
-            conns.delete(ws);
-            continue;
-        }
-        try {
-            ws.send(JSON.stringify(event));
-        }
-        catch (e) {
-            log.debug("[bridge] error forwarding request", e);
-            conns.delete(ws);
-        }
-    }
-};
-emitter.on(bus_request_stream, bus_to_ws_forwarder_listener);
-async function start_websocket_server(extensionPort) {
-    log.debug(`Draw.io MCP Server (${VERSION}) starting (WebSocket extension port: ${extensionPort})`);
-    const isPortAvailable = await checkPortAvailable(extensionPort);
-    if (!isPortAvailable) {
-        console.error(`[start_websocket_server] Error: Port ${extensionPort} is already in use. Please stop the process using this port and try again.`);
-        process.exit(1);
-    }
-    const server = new WebSocketServer({ port: extensionPort });
-    server.on("connection", (ws) => {
-        log.debug(`[ws_handler] A WebSocket client #${conns.size} connected, presumably MCP Extension!`);
-        conns.add(ws);
-        ws.on("message", (data) => {
-            const str = typeof data === "string" ? data : data.toString();
-            try {
-                const json = JSON.parse(str);
-                log.debug(`[ws] received from Extension`, json);
-                emitter.emit(bus_reply_stream, json);
-            }
-            catch (error) {
-                log.debug(`[ws] failed to parse message`, error);
-            }
-        });
-        ws.on("close", (code) => {
-            conns.delete(ws);
-            log.debug(`[ws_handler] WebSocket client closed with code ${code}`);
-        });
-        ws.on("error", (error) => {
-            log.debug(`[ws_handler] WebSocket client error`, error);
-            conns.delete(ws);
-        });
-    });
-    server.on("listening", () => {
-        log.debug(`[start_websocket_server] Listening to port ${extensionPort}`);
-    });
-    server.on("error", (error) => {
-        console.error(`[start_websocket_server] Error: Failed to listen on port ${extensionPort}`, error);
-        process.exit(1);
-    });
-    return server;
-}
 const logger_type = process.env.LOGGER_TYPE;
 let capabilities = {
     resources: {},
@@ -120,267 +72,12 @@ if (logger_type === "mcp_server") {
         },
     };
 }
-// Create server instance
-const server = new McpServer({
-    name: "drawio-mcp-server",
-    version: VERSION,
-}, {
-    capabilities,
-});
-const log = logger_type === "mcp_server"
-    ? create_server_logger(server)
-    : create_console_logger();
-const bus = create_bus(log)(emitter);
-const id_generator = nanoid_id_generator();
-const context = {
-    bus,
-    id_generator,
-    log,
-};
-const TOOL_get_selected_cell = "get-selected-cell";
-server.tool(TOOL_get_selected_cell, "This tool allows you to retrieve selected cell (whether vertex or edge) on the current page of a Draw.io diagram. The response is a JSON containing attributes of the cell.", {}, default_tool(TOOL_get_selected_cell, context));
-const TOOL_add_rectangle = "add-rectangle";
-server.tool(TOOL_add_rectangle, "This tool allows you to add new Rectangle vertex cell (object) on the current page of a Draw.io diagram. It accepts multiple optional input parameter.", {
-    x: z
-        .number()
-        .optional()
-        .describe("X-axis position of the Rectangle vertex cell")
-        .default(100),
-    y: z
-        .number()
-        .optional()
-        .describe("Y-axis position of the Rectangle vertex cell")
-        .default(100),
-    width: z
-        .number()
-        .optional()
-        .describe("Width of the Rectangle vertex cell")
-        .default(200),
-    height: z
-        .number()
-        .optional()
-        .describe("Height of the Rectangle vertex cell")
-        .default(100),
-    text: z
-        .string()
-        .optional()
-        .describe("Text content placed inside of the Rectangle vertex cell")
-        .default("New Cell"),
-    style: z
-        .string()
-        .optional()
-        .describe("Semi-colon separated list of Draw.io visual styles, in the form of `key=value`. Example: `whiteSpace=wrap;html=1;fillColor=#f5f5f5;strokeColor=#666666;`")
-        .default("whiteSpace=wrap;html=1;fillColor=#dae8fc;strokeColor=#6c8ebf;"),
-}, default_tool(TOOL_add_rectangle, context));
-const TOOL_add_edge = "add-edge";
-server.tool(TOOL_add_edge, "This tool creates an edge, sometimes called also a relation, between two vertexes (cells).", {
-    source_id: z
-        .string()
-        .describe("Source ID of a cell. It is represented by `id` attribute."),
-    target_id: z
-        .string()
-        .describe("Target ID of a cell. It is represented by `id` attribute."),
-    text: z
-        .string()
-        .optional()
-        .describe("Text content placed over the edge cell"),
-    style: z
-        .string()
-        .optional()
-        .describe("Semi-colon separated list of Draw.io visual styles, in the form of `key=value`. Example: `edgeStyle=orthogonalEdgeStyle;rounded=0;orthogonalLoop=1;jettySize=auto;html=1;exitX=1;exitY=0.5;exitDx=0;exitDy=0;entryX=0;entryY=0.5;entryDx=0;entryDy=0;`")
-        .default("edgeStyle=orthogonalEdgeStyle;rounded=0;orthogonalLoop=1;jettySize=auto;html=1;exitX=1;exitY=0.5;exitDx=0;exitDy=0;entryX=0;entryY=0.5;entryDx=0;entryDy=0;"),
-}, default_tool(TOOL_add_edge, context));
-const TOOL_delete_cell_by_id = "delete-cell-by-id";
-server.tool(TOOL_delete_cell_by_id, "Deletes a cell, whether it is a vertex or edge.", {
-    cell_id: z
-        .string()
-        .describe("The ID of a cell to delete. The cell can be either vertex or edge. The ID is located in `id` attribute."),
-}, default_tool(TOOL_delete_cell_by_id, context));
-const TOOL_get_shape_categories = "get-shape-categories";
-server.tool(TOOL_get_shape_categories, "Retrieves available shape categories from the diagram's library. Library is split into multiple categories.", {}, default_tool(TOOL_get_shape_categories, context));
-const TOOL_get_shapes_in_category = "get-shapes-in-category";
-server.tool(TOOL_get_shapes_in_category, "Retrieve all shapes in the provided category from the diagram's library. A shape primarily contains `style` based on which you can create new vertex cells.", {
-    category_id: z
-        .string()
-        .describe("Identifier (ID / key) of the category from which all the shapes should be retrieved."),
-}, default_tool(TOOL_get_shapes_in_category, context));
-const TOOL_get_shape_by_name = "get-shape-by-name";
-server.tool(TOOL_get_shape_by_name, "Retrieve a specific shape by its name from all available shapes in the diagram's library. It returns the shape and also the category it belongs.", {
-    shape_name: z
-        .string()
-        .describe("Name of the shape to retrieve from the shape library of the current diagram."),
-}, default_tool(TOOL_get_shape_by_name, context));
-const TOOL_add_cell_of_shape = "add-cell-of-shape";
-server.tool(TOOL_add_cell_of_shape, "This tool allows you to add new vertex cell (object) on the current page of a Draw.io diagram by its shape name. It accepts multiple optional input parameter.", {
-    shape_name: z
-        .string()
-        .describe("Name of the shape to retrieved from the shape library of the current diagram."),
-    x: z
-        .number()
-        .optional()
-        .describe("X-axis position of the vertex cell of the shape")
-        .default(100),
-    y: z
-        .number()
-        .optional()
-        .describe("Y-axis position of the vertex cell of the shape")
-        .default(100),
-    width: z
-        .number()
-        .optional()
-        .describe("Width of the vertex cell of the shape")
-        .default(200),
-    height: z
-        .number()
-        .optional()
-        .describe("Height of the vertex cell of the shape")
-        .default(100),
-    text: z
-        .string()
-        .optional()
-        .describe("Text content placed inside of the vertex cell of the shape"),
-    style: z
-        .string()
-        .optional()
-        .describe("Semi-colon separated list of Draw.io visual styles, in the form of `key=value`. Example: `whiteSpace=wrap;html=1;fillColor=#f5f5f5;strokeColor=#666666;`"),
-}, default_tool(TOOL_add_cell_of_shape, context));
-const TOOL_set_cell_shape = "set-cell-shape";
-server.tool(TOOL_set_cell_shape, "Updates the visual style of an existing vertex cell to match a library shape by name.", {
-    cell_id: z
-        .string()
-        .describe("Identifier (`id` attribute) of the cell whose shape should change."),
-    shape_name: z
-        .string()
-        .describe("Name of the library shape whose style should be applied to the existing cell."),
-}, default_tool(TOOL_set_cell_shape, context));
-const TOOL_set_cell_data = "set-cell-data";
-server.tool(TOOL_set_cell_data, "Sets or updates a custom attribute on an existing cell.", {
-    cell_id: z
-        .string()
-        .describe("Identifier (`id` attribute) of the cell to update with custom data."),
-    key: z.string().describe("Name of the attribute to set on the cell."),
-    value: z
-        .union([z.string(), z.number(), z.boolean()])
-        .describe("Value to store for the attribute. Non-string values are stringified before storage."),
-}, default_tool(TOOL_set_cell_data, context));
-const TOOL_edit_cell = "edit-cell";
-server.tool(TOOL_edit_cell, "Update properties of an existing vertex/shape cell by its ID. Only provided fields are modified; unspecified properties remain unchanged.", {
-    cell_id: z
-        .string()
-        .describe("Identifier (`id` attribute) of the cell to update. Applies to vertex/shape cells."),
-    text: z
-        .string()
-        .optional()
-        .describe("Replace the cell's text/label content."),
-    x: z
-        .number()
-        .optional()
-        .describe("Set a new X-axis position for the cell."),
-    y: z
-        .number()
-        .optional()
-        .describe("Set a new Y-axis position for the cell."),
-    width: z.number().optional().describe("Set a new width for the cell."),
-    height: z.number().optional().describe("Set a new height for the cell."),
-    style: z
-        .string()
-        .optional()
-        .describe("Replace the cell's style string (semi-colon separated `key=value` pairs)."),
-}, default_tool(TOOL_edit_cell, context));
-const TOOL_edit_edge = "edit-edge";
-server.tool(TOOL_edit_edge, "Update properties of an existing edge by its ID. Only provided fields are modified; unspecified properties remain unchanged.", {
-    cell_id: z
-        .string()
-        .describe("Identifier (`id` attribute) of the edge cell to update. The ID must reference an edge."),
-    text: z.string().optional().describe("Replace the edge's label text."),
-    source_id: z
-        .string()
-        .optional()
-        .describe("Reassign the edge's source terminal to a different cell ID."),
-    target_id: z
-        .string()
-        .optional()
-        .describe("Reassign the edge's target terminal to a different cell ID."),
-    style: z
-        .string()
-        .optional()
-        .describe("Replace the edge's style string (semi-colon separated `key=value` pairs)."),
-}, default_tool(TOOL_edit_edge, context));
-const Attributes = z.lazy(() => z
-    .array(z.union([
-    z.string(),
-    Attributes, // recursion: nested arrays
-]))
-    .refine((arr) => arr.length === 0 || typeof arr[0] === "string", {
-    message: "If not empty, the first element must be a string operator",
-})
-    .default([]));
-const TOOL_list_paged_model = "list-paged-model";
-server.tool(TOOL_list_paged_model, "Retrieves a paginated view of all cells (vertices and edges) in the current Draw.io diagram. This tool provides access to the complete model data with essential fields only, sanitized to remove circular dependencies and excessive data. It allows to filter based on multiple criteria and attribute boolean logic. Useful for programmatic inspection of diagram structure without overwhelming response sizes.", {
-    page: z
-        .number()
-        .optional()
-        .describe("Zero-based page number for pagination. Page 0 returns the first batch of cells, page 1 returns the next batch, etc. Default is 0.")
-        .default(0),
-    page_size: z
-        .number()
-        .optional()
-        .describe("Maximum number of cells to return in a single page. Controls response size and performance. Must be between 1 and 1000. Default is 50.")
-        .default(50),
-    filter: z
-        .object({
-        cell_type: z
-            .enum(["edge", "vertex", "object", "layer", "group"])
-            .optional()
-            .describe("Filter by cell type: 'edge' for connection lines, 'vertex' for vertices/shapes, 'object' for any cell type, 'layer' for layer cells, 'group' for grouped cells"),
-        parent_ids: z
-            .array(z.string())
-            .optional()
-            .describe("Filter cells to only those whose parent is one of the specified parent IDs."),
-        layer_ids: z
-            .array(z.string())
-            .optional()
-            .describe("Filter cells to only those whose parent is one of the specified layer IDs. Alias for parent_ids."),
-        ids: z
-            .array(z.string())
-            .optional()
-            .describe("Filter cells to only those whose ID is one of the specified IDs."),
-        attributes: Attributes.optional().describe('Boolean logic array expressions for filtering cell attributes. Format: ["and" | "or", ...expressions] or ["equal", key, value]. Matches against cell attributes and parsed style properties.'),
-    })
-        .optional()
-        .describe("Optional filter criteria to apply to cells before pagination")
-        .default({}),
-}, default_tool(TOOL_list_paged_model, context));
-const TOOL_list_layers = "list-layers";
-server.tool(TOOL_list_layers, "Lists all available layers in the diagram with their IDs and names.", {}, default_tool(TOOL_list_layers, context));
-const TOOL_set_active_layer = "set-active-layer";
-server.tool(TOOL_set_active_layer, "Sets the active layer for creating new elements. All subsequent element creation will happen in this layer.", {
-    layer_id: z.string().describe("ID of the layer to set as active"),
-}, default_tool(TOOL_set_active_layer, context));
-const TOOL_move_cell_to_layer = "move-cell-to-layer";
-server.tool(TOOL_move_cell_to_layer, "Moves a cell from its current layer to a target layer.", {
-    cell_id: z.string().describe("ID of the cell to move"),
-    target_layer_id: z
-        .string()
-        .describe("ID of the target layer where the cell will be moved"),
-}, default_tool(TOOL_move_cell_to_layer, context));
-const TOOL_get_active_layer = "get-active-layer";
-server.tool(TOOL_get_active_layer, "Gets the currently active layer information.", {}, default_tool(TOOL_get_active_layer, context));
-const TOOL_create_layer = "create-layer";
-server.tool(TOOL_create_layer, "Creates a new layer in the diagram.", {
-    name: z.string().describe("Name for the new layer"),
-}, default_tool(TOOL_create_layer, context));
-async function start_stdio_transport() {
+async function start_stdio_transport(server, log) {
     const transport = new StdioServerTransport();
     await server.connect(transport);
     log.debug(`Draw.io MCP Server STDIO transport active`);
 }
-async function start_streamable_http_transport(http_port) {
-    // Create a stateless transport (no options = no session management)
-    const transport = new WebStandardStreamableHTTPServerTransport();
-    // Create the Hono app
-    const app = new Hono();
-    // Enable CORS for all origins
+function setupCors(app) {
     app.use("*", cors({
         origin: "*",
         allowMethods: ["GET", "POST", "DELETE", "OPTIONS"],
@@ -392,16 +89,272 @@ async function start_streamable_http_transport(http_port) {
         ],
         exposeHeaders: ["mcp-session-id", "mcp-protocol-version"],
     }));
-    app.get("/health", (c) => c.json({ status: server.isConnected() ? "ok" : "mcp not ready" }));
+}
+function registerHealthRoute(app) {
+    app.get("/health", (c) => c.json({ status: "ok" }));
+}
+function registerConfigRoute(app, config) {
+    app.get("/api/config", (c) => c.json({
+        websocketPort: config.extensionPort,
+        serverUrl: `http://localhost:${config.httpPort}`,
+    }));
+}
+function registerEditorRoutes(app, config, log) {
+    const assetConfig = {
+        assetPath: config.assetPath,
+    };
+    const isLocal = isUsingLocalAssets(assetConfig);
+    const assetRoot = isLocal ? getAssetRoot(assetConfig) : null;
+    const localPluginPath = getLocalPluginPath();
+    app.get("/js/mcp-plugin.js", (c) => {
+        if (existsSync(localPluginPath)) {
+            const content = readFileSync(localPluginPath);
+            c.header("Content-Type", "application/javascript");
+            return c.body(content);
+        }
+        return c.text("Plugin not found", 404);
+    });
+    app.use("/*", async (c, next) => {
+        let path = c.req.path;
+        if (path === "" || path === "/") {
+            path = "index.html";
+        }
+        if (!isLocal || !assetRoot) {
+            await next();
+            return;
+        }
+        let filePath = join(assetRoot, path);
+        if (existsSync(filePath)) {
+            const fileStats = statSync(filePath);
+            if (fileStats.isDirectory()) {
+                const normalizedPath = path.endsWith("/") ? path.slice(0, -1) : path;
+                const parentDir = join(assetRoot, normalizedPath);
+                if (existsSync(parentDir) && statSync(parentDir).isDirectory()) {
+                    const entries = readdirSync(parentDir);
+                    const firstFile = entries.find((e) => !statSync(join(parentDir, e)).isDirectory());
+                    if (firstFile) {
+                        filePath = join(parentDir, firstFile);
+                    }
+                    else {
+                        await next();
+                        return;
+                    }
+                }
+                else {
+                    await next();
+                    return;
+                }
+            }
+            let content = readFileSync(filePath);
+            const ext = filePath.split(".").pop()?.toLowerCase();
+            if (ext === "html") {
+                const contentStr = content.toString("utf-8");
+                if (contentStr.includes("</body>") &&
+                    !contentStr.includes("mcp-plugin.js")) {
+                    const pluginScript = `<script src="/js/mcp-plugin.js"></script>`;
+                    content = Buffer.from(contentStr.replace("</body>", `${pluginScript}</body>`));
+                }
+            }
+            const contentTypes = {
+                html: "text/html",
+                css: "text/css",
+                js: "application/javascript",
+                json: "application/json",
+                svg: "image/svg+xml",
+                png: "image/png",
+                ico: "image/x-icon",
+                map: "application/json",
+            };
+            c.header("Content-Type", contentTypes[ext || ""] || "text/plain");
+            return c.body(content);
+        }
+        await next();
+    });
+    app.get("/", (c) => {
+        return c.redirect("/index.html?offline=1&local=1");
+    });
+    log.debug(`Draw.io editor enabled at: http://localhost:${config.httpPort}/`);
+}
+function registerMcpRoute(app) {
+    const transport = new WebStandardStreamableHTTPServerTransport();
     app.all("/mcp", (c) => transport.handleRequest(c.req.raw));
-    await server.connect(transport);
-    serve({
+    return transport;
+}
+function createHttpApp(log, config, features) {
+    const app = new Hono();
+    setupCors(app);
+    if (features.enableHealth)
+        registerHealthRoute(app);
+    if (features.enableConfig)
+        registerConfigRoute(app, config);
+    const mcpTransport = features.enableMcp ? registerMcpRoute(app) : undefined;
+    if (features.enableEditor)
+        registerEditorRoutes(app, config, log);
+    return { app, mcpTransport };
+}
+async function startHttpServer(server, log, httpPort, config, features) {
+    const { app, mcpTransport } = createHttpApp(log, config, features);
+    if (mcpTransport) {
+        await server.connect(mcpTransport);
+    }
+    const httpServer = serve({
         fetch: app.fetch,
-        port: http_port,
+        port: httpPort,
+    });
+    const listeningPort = httpPort === 0
+        ? (httpServer.address()?.port ?? httpPort)
+        : httpPort;
+    log.debug(`Draw.io MCP Server HTTP active on port ${listeningPort}`);
+    if (features.enableMcp) {
+        log.debug(`MCP endpoint: http://localhost:${listeningPort}/mcp`);
+    }
+    if (features.enableEditor) {
+        log.debug(`Editor: http://localhost:${listeningPort}/`);
+    }
+    return {
+        server: httpServer,
+        port: listeningPort,
+    };
+}
+export function createDrawioMcpApp(overrides) {
+    const server = new McpServer({
+        name: "drawio-mcp-server",
+        version: VERSION,
+    }, {
+        capabilities,
     });
-    log.debug(`Draw.io MCP Server Streamable HTTP transport active`);
-    log.debug(`Health check: http://localhost:${http_port}/health`);
-    log.debug(`MCP endpoint: http://localhost:${http_port}/mcp`);
+    const log = overrides?.log ??
+        (logger_type === "mcp_server"
+            ? create_server_logger(server)
+            : create_console_logger());
+    const emitter = new EventEmitter();
+    const conns = new Set();
+    const bus = create_bus(log)(emitter);
+    const id_generator = nanoid_id_generator();
+    const context = {
+        bus,
+        id_generator,
+        log,
+    };
+    registerTools(server, context);
+    const bus_to_ws_forwarder_listener = (event) => {
+        log.debug(`[bridge] received; forwarding message to #${conns.size} clients`, event);
+        for (const ws of [...conns]) {
+            if (ws.readyState !== WebSocket.OPEN) {
+                conns.delete(ws);
+                continue;
+            }
+            try {
+                ws.send(JSON.stringify(event));
+            }
+            catch (e) {
+                log.debug("[bridge] error forwarding request", e);
+                conns.delete(ws);
+            }
+        }
+    };
+    emitter.on(bus_request_stream, bus_to_ws_forwarder_listener);
+    let wsServer;
+    let httpServer;
+    async function startWebSocketServer(extensionPort = 3333) {
+        log.debug(`Draw.io MCP Server (${VERSION}) starting (WebSocket extension port: ${extensionPort})`);
+        if (extensionPort !== 0) {
+            const isPortAvailable = await checkPortAvailable(extensionPort);
+            if (!isPortAvailable) {
+                throw new Error(`[start_websocket_server] Error: Port ${extensionPort} is already in use. Please stop the process using this port and try again.`);
+            }
+        }
+        wsServer = new WebSocketServer({ port: extensionPort });
+        wsServer.on("connection", (ws) => {
+            log.debug(`[ws_handler] A WebSocket client #${conns.size} connected, presumably MCP Extension!`);
+            conns.add(ws);
+            ws.on("message", (data) => {
+                const str = typeof data === "string" ? data : data.toString();
+                try {
+                    const json = JSON.parse(str);
+                    log.debug(`[ws] received from Extension`, json);
+                    emitter.emit(bus_reply_stream, json);
+                }
+                catch (error) {
+                    log.debug(`[ws] failed to parse message`, error);
+                }
+            });
+            ws.on("close", (code) => {
+                conns.delete(ws);
+                log.debug(`[ws_handler] WebSocket client closed with code ${code}`);
+            });
+            ws.on("error", (error) => {
+                log.debug(`[ws_handler] WebSocket client error`, error);
+                conns.delete(ws);
+            });
+        });
+        await new Promise((resolve, reject) => {
+            const onListening = () => {
+                wsServer?.off("error", onError);
+                const address = wsServer?.address();
+                log.debug(`[start_websocket_server] Listening to port ${address?.port ?? extensionPort}`);
+                resolve();
+            };
+            const onError = (error) => {
+                wsServer?.off("listening", onListening);
+                reject(error);
+            };
+            wsServer?.once("listening", onListening);
+            wsServer?.once("error", onError);
+        });
+        return wsServer;
+    }
+    async function close() {
+        emitter.off(bus_request_stream, bus_to_ws_forwarder_listener);
+        for (const ws of [...conns]) {
+            try {
+                ws.close();
+            }
+            catch {
+                // ignore
+            }
+        }
+        conns.clear();
+        await server.close();
+        if (wsServer) {
+            await new Promise((resolve, reject) => {
+                wsServer?.close((error) => {
+                    if (error) {
+                        reject(error);
+                        return;
+                    }
+                    resolve();
+                });
+            });
+            wsServer = undefined;
+        }
+        if (httpServer) {
+            await new Promise((resolve, reject) => {
+                httpServer?.close((error) => {
+                    if (error) {
+                        reject(error);
+                        return;
+                    }
+                    resolve();
+                });
+            });
+            httpServer = undefined;
+        }
+    }
+    return {
+        server,
+        log,
+        context,
+        emitter,
+        close,
+        startWebSocketServer,
+        startStdioTransport: () => start_stdio_transport(server, log),
+        startHttpServer: async (httpPort, config, features) => {
+            const started = await startHttpServer(server, log, httpPort, config, features);
+            httpServer = started.server;
+            return started;
+        },
+    };
 }
 async function main() {
     // Check if help was requested (before parsing config)
@@ -417,16 +370,31 @@ async function main() {
         process.exit(1);
     }
     const config = configResult;
-    await start_websocket_server(config.extensionPort);
-    if (config.transports.indexOf("stdio") > -1) {
-        await start_stdio_transport();
+    const features = getHttpFeatureConfig(config);
+    // Initialize assets if needed
+    if (features.enableEditor) {
+        console.log("Initializing draw.io assets...");
+        const assetConfig = {
+            assetPath: config.assetPath,
+        };
+        await ensureAssets(assetConfig, (msg) => console.log(msg));
+        console.log("Assets ready!");
     }
-    if (config.transports.indexOf("http") > -1) {
-        start_streamable_http_transport(config.httpPort);
+    const app = createDrawioMcpApp();
+    await app.startWebSocketServer(config.extensionPort);
+    if (config.transports.indexOf("stdio") > -1) {
+        await app.startStdioTransport();
     }
-    log.debug(`Draw.io MCP Server running on ${config.transports}`);
+    await app.startHttpServer(config.httpPort, config, features);
+    app.log.debug(`Draw.io MCP Server running on ${config.transports}`);
+}
+const isMainModule = process.argv[1]
+    ? realpathSync(fileURLToPath(import.meta.url)) ===
+        realpathSync(process.argv[1])
+    : false;
+if (isMainModule) {
+    main().catch((error) => {
+        fatalLog.debug("Fatal error in main():", error);
+        process.exit(1);
+    });
 }
-main().catch((error) => {
-    log.debug("Fatal error in main():", error);
-    process.exit(1);
-});