drawio-mcp-server 1.8.0 → 2.0.3

package/README.md

@@ -3,16 +3,20 @@
 Let's do some Vibe Diagramming with the most wide-spread diagramming tool called Draw.io (Diagrams.net).
 
 [![Discord channel](https://shields.io/static/v1?logo=discord&message=draw.io%20mcp&label=chat&color=5865F2&logoColor=white)](https://discord.gg/dM4PWdf42q)
-[![Build project](https://github.com/lgazo/drawio-mcp-server/actions/workflows/ci.yml/badge.svg)](https://github.com/lgazo/drawio-mcp-server/actions/workflows/ci.yml)
+[![Build project](https://github.com/lgazo/drawio-mcp-server/actions/workflows/server-ci.yml/badge.svg)](https://github.com/lgazo/drawio-mcp-server/actions/workflows/server-ci.yml)
 [![Verified on MseeP](https://mseep.ai/badge.svg)](https://mseep.ai/app/5fc2b7fe-8ceb-4683-97bd-6d31e07b5888)
 [![Version](https://img.shields.io/npm/v/drawio-mcp-server)](https://npmjs.com/package/drawio-mcp-server)
 
 ## Key Highlights
 
+- Import and export diagrams from/to XML, SVG (with embedded XML), or PNG (with embedded XML) ![v2.0.0](https://img.shields.io/badge/v2.0.0-blue)
+- Edge geometry control with waypoints and automatic self-connector routing ![v2.0.0](https://img.shields.io/badge/v2.0.0-blue)
+- Parent-child relationships for nested shapes and grouping ![v2.0.0](https://img.shields.io/badge/v2.0.0-blue)
+- Unified Server and Extension in the same mono-repo ![v2.0.0](https://img.shields.io/badge/v2.0.0-blue)
 - Built-in Draw.io editor - no browser extension required ![v1.8.0](https://img.shields.io/badge/v1.8.0-blue)
 - MCP server that lets AI agents control Draw.io diagrams
 - Programmatic diagram creation, inspection, and modification via MCP tools
-- Layer management for complex diagrams ![v1.7.0](https://img.shields.io/badge/v1.7.0-blue)
+- Layer management for complex diagrams
 - Works with any MCP client (Claude Desktop, Claude Code, Zed, Codex, etc.)
 
 ## Introduction
@@ -32,7 +36,7 @@ Two ways to use:
 No additional requirements - runs out of the box with `--editor` flag.
 
 ### For Browser Extension
-- **Browser extension** - [drawio-mcp-extension](https://github.com/lgazo/drawio-mcp-extension)
+- **Browser extension** - [drawio-mcp-extension](./packages/drawio-mcp-extension/README.md)
 - Draw.io open in your browser
 
 ### Optional
@@ -152,7 +156,7 @@ The server runs as part of your MCP host. Detailed configuration for all support
 
 ## Alternative: Browser Extension
 
-Instead of the built-in editor, you can use the [browser extension](https://github.com/lgazo/drawio-mcp-extension) to connect to Draw.io running in your browser. This works with or without the `--editor` flag.
+Instead of the built-in editor, you can use the [browser extension](./packages/drawio-mcp-extension/README.md) to connect to Draw.io running in your browser. This works with or without the `--editor` flag.
 
 1. Open [Draw.io in your browser](https://app.diagrams.net/)
 2. Install the Draw.io MCP Browser Extension:
@@ -173,7 +177,7 @@ Configuration without `--editor`:
 }
 ```
 
-See the [extension documentation](https://github.com/lgazo/drawio-mcp-extension) for more details.
+See the [extension documentation](./packages/drawio-mcp-extension/README.md) for more details.
 
 ## Related Resources
 

package/build/emitter_bus.js

@@ -1,7 +1,6 @@
-import { bus_reply_stream, bus_request_stream, } from "./types.js";
+import { bus_reply_stream, bus_request_stream } from "./types.js";
 export function create_bus(log) {
     return function (emitter) {
-        const listeners = [];
         const bus = {
             send_to_extension: (request) => {
                 log.debug(`[bus] sending to Extension`, request);
@@ -11,11 +10,11 @@ export function create_bus(log) {
                 const listener = (emitter_data) => {
                     log.debug(`[bus] received from Extension`, emitter_data);
                     if (emitter_data && emitter_data.__event === event_name) {
+                        emitter.off(bus_reply_stream, listener);
                         reply(emitter_data);
                     }
                 };
                 emitter.on(bus_reply_stream, listener);
-                listeners.push(reply);
             },
         };
         return bus;

package/build/index.js

@@ -3,23 +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 { readFileSync, existsSync, statSync, readdirSync } from "node:fs";
+import { fileURLToPath } from "node:url";
+import { readFileSync, existsSync, statSync, readdirSync, realpathSync, } from "node:fs";
 import { WebSocket, WebSocketServer } from "ws";
+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";
 import { getLocalPluginPath, isUsingLocalAssets, getAssetRoot, ensureAssets, } from "./assets/index.js";
-const VERSION = "1.8.0";
+import { registerTools } from "./tools/index.js";
+const fatalLog = create_console_logger();
 /**
  * Display help message and exit
  */
@@ -57,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: {},
@@ -130,257 +72,7 @@ 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`);
@@ -399,7 +91,7 @@ function setupCors(app) {
     }));
 }
 function registerHealthRoute(app) {
-    app.get("/health", (c) => c.json({ status: server.isConnected() ? "ok" : "mcp not ready" }));
+    app.get("/health", (c) => c.json({ status: "ok" }));
 }
 function registerConfigRoute(app, config) {
     app.get("/api/config", (c) => c.json({
@@ -407,7 +99,7 @@ function registerConfigRoute(app, config) {
         serverUrl: `http://localhost:${config.httpPort}`,
     }));
 }
-function registerEditorRoutes(app, config) {
+function registerEditorRoutes(app, config, log) {
     const assetConfig = {
         assetPath: config.assetPath,
     };
@@ -488,7 +180,7 @@ function registerMcpRoute(app) {
     app.all("/mcp", (c) => transport.handleRequest(c.req.raw));
     return transport;
 }
-function createHttpApp(config, features) {
+function createHttpApp(log, config, features) {
     const app = new Hono();
     setupCors(app);
     if (features.enableHealth)
@@ -497,25 +189,172 @@ function createHttpApp(config, features) {
         registerConfigRoute(app, config);
     const mcpTransport = features.enableMcp ? registerMcpRoute(app) : undefined;
     if (features.enableEditor)
-        registerEditorRoutes(app, config);
+        registerEditorRoutes(app, config, log);
     return { app, mcpTransport };
 }
-async function startHttpServer(httpPort, config, features) {
-    const { app, mcpTransport } = createHttpApp(config, features);
+async function startHttpServer(server, log, httpPort, config, features) {
+    const { app, mcpTransport } = createHttpApp(log, config, features);
     if (mcpTransport) {
         await server.connect(mcpTransport);
     }
-    serve({
+    const httpServer = serve({
         fetch: app.fetch,
         port: httpPort,
     });
-    log.debug(`Draw.io MCP Server HTTP active on 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:${httpPort}/mcp`);
+        log.debug(`MCP endpoint: http://localhost:${listeningPort}/mcp`);
     }
     if (features.enableEditor) {
-        log.debug(`Editor: http://localhost:${httpPort}/`);
+        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,
+    });
+    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)
@@ -541,14 +380,21 @@ async function main() {
         await ensureAssets(assetConfig, (msg) => console.log(msg));
         console.log("Assets ready!");
     }
-    await start_websocket_server(config.extensionPort);
+    const app = createDrawioMcpApp();
+    await app.startWebSocketServer(config.extensionPort);
     if (config.transports.indexOf("stdio") > -1) {
-        await start_stdio_transport();
+        await app.startStdioTransport();
     }
-    startHttpServer(config.httpPort, config, features);
-    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);
-});