bugfix: Implement Graceful WebSocket Server Shutdown Mechanism

- Added shutdown method to BrowserConnector for controlled WebSocket server closure
- Improved process signal handling for SIGINT and SIGTERM
- Updated Chrome extension to handle server shutdown signal
- Bumped package versions to 1.1.0 for browser-tools-mcp and browser-tools-server

browser-tools-mcp/package-lock.json

@@ -1,12 +1,12 @@
 {
   "name": "@agentdeskai/browser-tools-mcp",
-  "version": "1.0.11",
+  "version": "1.1.0",
   "lockfileVersion": 3,
   "requires": true,
   "packages": {
     "": {
       "name": "@agentdeskai/browser-tools-mcp",
-      "version": "1.0.11",
+      "version": "1.1.0",
       "license": "MIT",
       "dependencies": {
         "@modelcontextprotocol/sdk": "^1.4.1",

browser-tools-server/browser-connector.ts

@@ -667,6 +667,47 @@ export class BrowserConnector {
       });
     }
   }
+
+  // Add shutdown method
+  public shutdown() {
+    return new Promise<void>((resolve) => {
+      console.log("Shutting down WebSocket server...");
+
+      // Send close message to client if connection is active
+      if (this.activeConnection && this.activeConnection.readyState === WebSocket.OPEN) {
+        console.log("Notifying client to close connection...");
+        try {
+          this.activeConnection.send(JSON.stringify({ type: "server-shutdown" }));
+        } catch (err) {
+          console.error("Error sending shutdown message to client:", err);
+        }
+      }
+
+      // Set a timeout to force close after 2 seconds
+      const forceCloseTimeout = setTimeout(() => {
+        console.log("Force closing connections after timeout...");
+        if (this.activeConnection) {
+          this.activeConnection.terminate(); // Force close the connection
+          this.activeConnection = null;
+        }
+        this.wss.close();
+        resolve();
+      }, 2000);
+
+      // Close active WebSocket connection if exists
+      if (this.activeConnection) {
+        this.activeConnection.close(1000, "Server shutting down");
+        this.activeConnection = null;
+      }
+
+      // Close WebSocket server
+      this.wss.close(() => {
+        clearTimeout(forceCloseTimeout);
+        console.log("WebSocket server closed gracefully");
+        resolve();
+      });
+    });
+  }
 }
 
 // Move the server creation before BrowserConnector instantiation
@@ -677,10 +718,41 @@ const server = app.listen(PORT, () => {
 // Initialize the browser connector with the existing app AND server
 const browserConnector = new BrowserConnector(app, server);
 
-// Handle shutdown gracefully
-process.on("SIGINT", () => {
-  server.close(() => {
-    console.log("Server shut down");
+// Handle shutdown gracefully with improved error handling
+process.on("SIGINT", async () => {
+  console.log("\nReceived SIGINT signal. Starting graceful shutdown...");
+
+  try {
+    // First shutdown WebSocket connections
+    await browserConnector.shutdown();
+
+    // Then close the HTTP server
+    await new Promise<void>((resolve, reject) => {
+      server.close((err) => {
+        if (err) {
+          console.error("Error closing HTTP server:", err);
+          reject(err);
+        } else {
+          console.log("HTTP server closed successfully");
+          resolve();
+        }
+      });
+    });
+
+    // Clear all logs
+    clearAllLogs();
+
+    console.log("Shutdown completed successfully");
     process.exit(0);
-  });
+  } catch (error) {
+    console.error("Error during shutdown:", error);
+    // Force exit in case of error
+    process.exit(1);
+  }
+});
+
+// Also handle SIGTERM
+process.on("SIGTERM", () => {
+  console.log("\nReceived SIGTERM signal");
+  process.emit("SIGINT");
 });

browser-tools-server/package-lock.json

@@ -1,12 +1,12 @@
 {
   "name": "@agentdeskai/browser-tools-server",
-  "version": "1.0.5",
+  "version": "1.1.0",
   "lockfileVersion": 3,
   "requires": true,
   "packages": {
     "": {
       "name": "@agentdeskai/browser-tools-server",
-      "version": "1.0.5",
+      "version": "1.1.0",
       "license": "MIT",
       "dependencies": {
         "@modelcontextprotocol/sdk": "^1.4.1",

chrome-extension/devtools.js

@@ -536,6 +536,18 @@ function setupWebSocket() {
       const message = JSON.parse(event.data);
       console.log("Chrome Extension: Received WebSocket message:", message);
 
+      if (message.type === "server-shutdown") {
+        console.log("Chrome Extension: Received server shutdown signal");
+        // Clear any reconnection attempts
+        if (wsReconnectTimeout) {
+          clearTimeout(wsReconnectTimeout);
+          wsReconnectTimeout = null;
+        }
+        // Close the connection gracefully
+        ws.close(1000, "Server shutting down");
+        return;
+      }
+
       if (message.type === "take-screenshot") {
         console.log("Chrome Extension: Taking screenshot...");
         // Capture screenshot of the current tab
@@ -574,10 +586,7 @@ function setupWebSocket() {
         });
       }
     } catch (error) {
-      console.error(
-        "Chrome Extension: Error processing WebSocket message:",
-        error
-      );
+      console.error("Chrome Extension: Error processing WebSocket message:", error);
     }
   };
 
@@ -589,11 +598,17 @@ function setupWebSocket() {
     }
   };
 
-  ws.onclose = () => {
+  ws.onclose = (event) => {
     console.log(
-      "Chrome Extension: WebSocket disconnected, attempting to reconnect..."
+      `Chrome Extension: WebSocket disconnected (${event.code}: ${event.reason})`
     );
-    wsReconnectTimeout = setTimeout(setupWebSocket, WS_RECONNECT_DELAY);
+    // Only attempt to reconnect if it wasn't a server shutdown
+    if (event.reason !== "Server shutting down") {
+      console.log("Chrome Extension: Attempting to reconnect...");
+      wsReconnectTimeout = setTimeout(setupWebSocket, WS_RECONNECT_DELAY);
+    } else {
+      console.log("Chrome Extension: Server shutdown detected, not reconnecting");
+    }
   };
 
   ws.onerror = (error) => {