feat: implement centralized shutdown manager and logging

src/server.ts

@@ -12,6 +12,9 @@ import {
 } from "./services";
 import { schema, tool } from "./tools";
 import { renderMermaid } from "./utils";
+import { createServerCleanup } from "./utils/cleanupHelpers";
+import { Logger } from "./utils/logger";
+import { shutdownManager } from "./utils/shutdownManager";
 
 /**
  * Creates and configures an MCP server for mermaid generation.
@@ -31,13 +34,10 @@ export function createServer(): Server {
 
   setupToolHandlers(server);
 
-  server.onerror = (error) => console.error("[MCP Error]", error);
-  // Graceful shutdown on Ctrl-C, remove listener afterwards to prevent leaks
-  const sigintHandler = async () => {
-    await server.close();
-    process.exit(0);
-  };
-  process.once("SIGINT", sigintHandler);
+  server.onerror = (error) => Logger.error("MCP Error", error);
+
+  // Register server cleanup with shutdown manager
+  shutdownManager.registerCleanup(createServerCleanup(server));
 
   return server;
 }

src/services/sse.ts

@@ -2,6 +2,7 @@ import type { IncomingMessage, ServerResponse } from "node:http";
 import type { Server } from "@modelcontextprotocol/sdk/server/index.js";
 import { SSEServerTransport } from "@modelcontextprotocol/sdk/server/sse.js";
 import { type RequestHandlers, createBaseHttpServer } from "../utils";
+import { Logger } from "../utils/logger";
 
 export const startSSEMcpServer = async (
   server: Server,
@@ -36,7 +37,7 @@ export const startSSEMcpServer = async (
         try {
           await server.close();
         } catch (error) {
-          console.error("Error closing server:", error);
+          Logger.error("Error closing server", error);
         }
 
         delete activeTransports[transport.sessionId];
@@ -52,8 +53,7 @@ export const startSSEMcpServer = async (
         });
       } catch (error) {
         if (!closed) {
-          console.error("Error connecting to server:", error);
-
+          Logger.error("Error connecting to server", error);
           res.writeHead(500).end("Error connecting to server");
         }
       }
@@ -93,9 +93,12 @@ export const startSSEMcpServer = async (
   const cleanup = () => {
     // Close all active transports
     for (const transport of Object.values(activeTransports)) {
-      transport.close();
+      try {
+        transport.close();
+      } catch (error) {
+        Logger.error("Error closing SSE transport", error);
+      }
     }
-    server.close();
   };
 
   // Create the HTTP server using our factory

src/services/stdio.ts

@@ -1,7 +1,13 @@
 import type { Server } from "@modelcontextprotocol/sdk/server/index.js";
 import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
+import { createTransportCleanup } from "../utils/cleanupHelpers";
+import { shutdownManager } from "../utils/shutdownManager";
 
 export async function startStdioMcpServer(server: Server): Promise<void> {
   const transport = new StdioServerTransport();
+
+  // Register transport cleanup
+  shutdownManager.registerCleanup(createTransportCleanup(transport));
+
   await server.connect(transport);
 }

src/services/streamable.ts

@@ -12,6 +12,7 @@ import {
   createBaseHttpServer,
   getBody,
 } from "../utils";
+import { Logger } from "../utils/logger";
 
 export const startHTTPStreamableServer = async (
   createServer: () => Server,
@@ -82,7 +83,7 @@ export const startHTTPStreamableServer = async (
               try {
                 await server?.close();
               } catch (error) {
-                console.error("Error closing server:", error);
+                Logger.error("Error closing server", error);
               }
 
               // delete used transport and server to avoid memory leak.
@@ -126,7 +127,7 @@ export const startHTTPStreamableServer = async (
         // Handle the request if the server is already created.
         await transport.handleRequest(req, res, body);
       } catch (error) {
-        console.error("Error handling request:", error);
+        Logger.error("Error handling request", error);
         res.setHeader("Content-Type", "application/json");
         res.writeHead(500).end(
           JSON.stringify({
@@ -204,8 +205,11 @@ export const startHTTPStreamableServer = async (
   // Custom cleanup for streamable server
   const cleanup = () => {
     for (const { server, transport } of Object.values(activeTransports)) {
-      transport.close();
-      server.close();
+      try {
+        transport.close();
+      } catch (error) {
+        Logger.error("Error closing streamable transport", error);
+      }
     }
   };
 

src/utils/cleanupHelpers.ts

@@ -0,0 +1,47 @@
+import type { Server } from "@modelcontextprotocol/sdk/server/index.js";
+import { Logger } from "./logger";
+
+/**
+ * Create a server cleanup handler
+ */
+export function createServerCleanup(server: Server): () => void {
+  return () => {
+    try {
+      server.close();
+    } catch (error) {
+      Logger.error("Error closing server", error);
+    }
+  };
+}
+
+/**
+ * Create a transport cleanup handler
+ */
+export function createTransportCleanup(transport: {
+  close: () => void;
+}): () => void {
+  return () => {
+    try {
+      transport.close();
+    } catch (error) {
+      Logger.error("Error closing transport", error);
+    }
+  };
+}
+
+/**
+ * Create a combined cleanup handler for server and transport
+ */
+export function createCombinedCleanup(
+  server: Server,
+  transport: { close: () => void },
+): () => void {
+  return () => {
+    try {
+      transport.close();
+      server.close();
+    } catch (error) {
+      Logger.error("Error during combined cleanup", error);
+    }
+  };
+}

src/utils/httpServer.ts

@@ -1,5 +1,7 @@
 import http from "node:http";
 import type { IncomingMessage, ServerResponse } from "node:http";
+import { Logger } from "./logger";
+import { shutdownManager } from "./shutdownManager";
 
 /**
  * Interface for request handlers that will be passed to the server factory
@@ -33,7 +35,7 @@ function handleCORS(req: IncomingMessage, res: ServerResponse): void {
       res.setHeader("Access-Control-Allow-Methods", "GET, POST, OPTIONS");
       res.setHeader("Access-Control-Allow-Headers", "*");
     } catch (error) {
-      console.error("Error parsing origin:", error);
+      Logger.error("Error parsing origin", error);
     }
   }
 }
@@ -65,31 +67,33 @@ function handleCommonEndpoints(
 }
 
 /**
- * Sets up signal handlers for graceful shutdown
+ * Sets up cleanup handlers for the HTTP server
  */
 function setupCleanupHandlers(
   httpServer: http.Server,
   customCleanup?: () => void,
 ): void {
-  const cleanup = () => {
-    console.log("\nClosing server...");
+  // Register HTTP server cleanup
+  shutdownManager.registerCleanup(() => {
+    return new Promise<void>((resolve) => {
+      httpServer.close(() => {
+        Logger.cleanup("HTTP server closed");
+        resolve();
+      });
+    });
+  });
 
-    // Execute custom cleanup if provided
-    if (customCleanup) customCleanup();
+  // Register custom cleanup if provided
+  if (customCleanup) {
+    shutdownManager.registerCleanup(customCleanup);
+  }
 
-    httpServer.close(() => {
-      console.log("Server closed");
-      // exit after cleanup
-      process.exit(0);
-    });
-  };
+  // Setup signal handlers only once
+  shutdownManager.setupSignalHandlers();
 
-  // Attach signal handlers and remove them on server close to avoid leaks
-  process.on("SIGINT", cleanup);
-  process.on("SIGTERM", cleanup);
   httpServer.once("close", () => {
-    process.removeListener("SIGINT", cleanup);
-    process.removeListener("SIGTERM", cleanup);
+    process.removeAllListeners("SIGINT");
+    process.removeAllListeners("SIGTERM");
   });
 }
 
@@ -101,16 +105,7 @@ function logServerStartup(
   port: number,
   endpoint: string,
 ): void {
-  const serverUrl = `http://localhost:${port}${endpoint}`;
-  const healthUrl = `http://localhost:${port}/health`;
-  const pingUrl = `http://localhost:${port}/ping`;
-
-  console.log(
-    `${serverType} running on: \x1b[32m\u001B[4m${serverUrl}\u001B[0m\x1b[0m`,
-  );
-  console.log("\nTest endpoints:");
-  console.log(`• Health check: \u001B[4m${healthUrl}\u001B[0m`);
-  console.log(`• Ping test: \u001B[4m${pingUrl}\u001B[0m`);
+  Logger.serverStartup(serverType, port, endpoint);
 }
 
 /**
@@ -138,7 +133,7 @@ export function createBaseHttpServer(
     try {
       await handlers.handleRequest(req, res);
     } catch (error) {
-      console.error(`Error in ${handlers.serverType} request handler:`, error);
+      Logger.error(`Error in ${handlers.serverType} request handler`, error);
       res.writeHead(500).end("Internal Server Error");
     }
   });

src/utils/index.ts

@@ -3,3 +3,10 @@ export { InMemoryEventStore } from "./InMemoryEventStore";
 export { getBody } from "./getBody";
 export { createBaseHttpServer, type RequestHandlers } from "./httpServer";
 export { renderMermaid } from "./render";
+export { shutdownManager, ShutdownManager } from "./shutdownManager";
+export { Logger } from "./logger";
+export {
+  createServerCleanup,
+  createTransportCleanup,
+  createCombinedCleanup,
+} from "./cleanupHelpers";

src/utils/logger.ts

@@ -0,0 +1,71 @@
+/**
+ * Unified logger for consistent logging across the application
+ */
+const prefix = "[MCP-Mermaid]";
+
+/**
+ * Log info message
+ */
+export function info(message: string, ...args: unknown[]): void {
+  console.log(`${prefix} ℹ️  ${message}`, ...args);
+}
+
+/**
+ * Log warning message
+ */
+export function warn(message: string, ...args: unknown[]): void {
+  console.warn(`${prefix} ⚠️  ${message}`, ...args);
+}
+
+/**
+ * Log error message
+ */
+export function error(message: string, error?: unknown): void {
+  console.error(`${prefix} ❌ ${message}`, error || "");
+}
+
+/**
+ * Log success message
+ */
+export function success(message: string, ...args: unknown[]): void {
+  console.log(`${prefix} ✅ ${message}`, ...args);
+}
+
+/**
+ * Log server startup information
+ */
+export function serverStartup(
+  serverType: string,
+  port: number,
+  endpoint: string,
+): void {
+  const serverUrl = `http://localhost:${port}${endpoint}`;
+  const healthUrl = `http://localhost:${port}/health`;
+  const pingUrl = `http://localhost:${port}/ping`;
+
+  console.log(
+    `${prefix} 🚀 ${serverType} running on: \x1b[32m\u001B[4m${serverUrl}\u001B[0m\x1b[0m`,
+  );
+  console.log("\nTest endpoints:");
+  console.log(`• Health check: \u001B[4m${healthUrl}\u001B[0m`);
+  console.log(`• Ping test: \u001B[4m${pingUrl}\u001B[0m`);
+}
+
+/**
+ * Log cleanup information
+ */
+export function cleanup(message: string): void {
+  console.log(`${prefix} 🧹 ${message}`);
+}
+
+/**
+ * Logger object for backward compatibility
+ */
+export const Logger = {
+  info,
+  warn,
+  error,
+  success,
+  serverStartup,
+  cleanup,
+};

src/utils/render.ts

@@ -1,6 +1,6 @@
-import * as path from "path";
-import * as fs from "fs";
-import * as os from "os";
+import * as fs from "node:fs";
+import * as os from "node:os";
+import * as path from "node:path";
 import {
   type MermaidRenderer,
   type RenderResult,
@@ -23,7 +23,7 @@ export async function renderMermaid(
 ): Promise<RenderResult> {
   if (!renderer) renderer = createMermaidRenderer();
   const cssContent = `svg { background: ${backgroundColor}; }`;
-  const cssTmpPath = path.join(os.tmpdir(), 'mermaid-tmp-css.css');
+  const cssTmpPath = path.join(os.tmpdir(), "mermaid-tmp-css.css");
   fs.writeFileSync(cssTmpPath, cssContent);
 
   const r = await renderer([mermaid], {