@mcp-shark/mcp-shark 1.5.4 → 1.5.5

package/ui/server/utils/port.js

@@ -1,28 +1,28 @@
 import { createConnection } from 'node:net';
+import { Defaults } from '#core/constants/Defaults.js';
 
-export function checkPortReady(port, host = 'localhost', timeout = 10000) {
-  return new Promise((resolve, reject) => {
-    const startTime = Date.now();
+function tryConnect(port, host, startTime, timeout, resolve, reject) {
+  const socket = createConnection(port, host);
 
-    const tryConnect = () => {
-      const socket = createConnection(port, host);
-
-      socket.on('connect', () => {
-        socket.destroy();
-        resolve(true);
-      });
+  socket.on('connect', () => {
+    socket.destroy();
+    resolve(true);
+  });
 
-      socket.on('error', (_err) => {
-        socket.destroy();
-        const elapsed = Date.now() - startTime;
-        if (elapsed >= timeout) {
-          reject(new Error(`Port ${port} not ready after ${timeout}ms`));
-        } else {
-          setTimeout(tryConnect, 200);
-        }
-      });
-    };
+  socket.on('error', (_err) => {
+    socket.destroy();
+    const elapsed = Date.now() - startTime;
+    if (elapsed >= timeout) {
+      reject(new Error(`Port ${port} not ready after ${timeout}ms`));
+    } else {
+      setTimeout(() => tryConnect(port, host, startTime, timeout, resolve, reject), 200);
+    }
+  });
+}
 
-    tryConnect();
+export function checkPortReady(port, host = 'localhost', timeout = Defaults.PORT_CHECK_TIMEOUT) {
+  return new Promise((resolve, reject) => {
+    const startTime = Date.now();
+    tryConnect(port, host, startTime, timeout, resolve, reject);
   });
 }

package/ui/server/utils/process.js

@@ -1,13 +1,21 @@
-const MAX_LOG_LINES = 10000;
+import { Defaults } from '#core/constants/Defaults.js';
+
+const MAX_LOG_LINES = Defaults.MAX_LOG_LINES;
+
+function logEntry(mcpSharkLogs, broadcastLogUpdate, type, data) {
+  const timestamp = new Date().toISOString();
+  const line = data.toString();
+  mcpSharkLogs.push({ timestamp, type, line });
+  if (mcpSharkLogs.length > MAX_LOG_LINES) {
+    mcpSharkLogs.shift();
+  }
+  broadcastLogUpdate({ timestamp, type, line });
+}
+
+function createLogEntryWrapper(mcpSharkLogs, broadcastLogUpdate, type, data) {
+  return logEntry(mcpSharkLogs, broadcastLogUpdate, type, data);
+}
 
 export function createLogEntry(mcpSharkLogs, broadcastLogUpdate) {
-  return function logEntry(type, data) {
-    const timestamp = new Date().toISOString();
-    const line = data.toString();
-    mcpSharkLogs.push({ timestamp, type, line });
-    if (mcpSharkLogs.length > MAX_LOG_LINES) {
-      mcpSharkLogs.shift();
-    }
-    broadcastLogUpdate({ timestamp, type, line });
-  };
+  return (type, data) => createLogEntryWrapper(mcpSharkLogs, broadcastLogUpdate, type, data);
 }

package/ui/server/utils/processState.js

@@ -0,0 +1,17 @@
+/**
+ * Set MCP Shark server process in state
+ * @param {object} processState - Process state object
+ * @param {object} server - Server instance
+ */
+export function setMcpSharkProcess(processState, server) {
+  processState.mcpSharkServer = server;
+}
+
+/**
+ * Get MCP Shark server process from state
+ * @param {object} processState - Process state object
+ * @returns {object|null} Server instance or null
+ */
+export function getMcpSharkProcess(processState) {
+  return processState.mcpSharkServer;
+}

package/ui/server/utils/signals.js

@@ -0,0 +1,34 @@
+/**
+ * Handle shutdown signal (SIGTERM, SIGINT)
+ * @param {Function} cleanup - Cleanup function to execute
+ * @param {object} logger - Logger instance
+ */
+export async function shutdown(cleanup, logger) {
+  try {
+    // Set a timeout to force exit if cleanup takes too long
+    const timeout = setTimeout(() => {
+      logger?.warn('Shutdown timeout reached, forcing exit');
+      process.exit(0);
+    }, 2000); // 2 second timeout
+
+    await cleanup();
+    clearTimeout(timeout);
+  } catch (err) {
+    logger?.warn({ error: err.message }, 'Error during shutdown, exiting anyway');
+  } finally {
+    // Always exit, even if cleanup failed
+    process.exit(0);
+  }
+}
+
+/**
+ * Handle process exit event
+ * @param {Function} cleanup - Cleanup function to execute
+ */
+export async function handleExit(cleanup) {
+  try {
+    await cleanup();
+  } catch (_err) {
+    // Ignore errors during exit
+  }
+}

package/ui/server/websocket/broadcast.js

@@ -0,0 +1,33 @@
+import { RequestFilters } from '#core';
+
+/**
+ * Broadcast log update to all connected WebSocket clients
+ * @param {Set} clients - Set of connected WebSocket clients
+ * @param {object} logEntry - Log entry to broadcast
+ */
+export function broadcastLogUpdate(clients, logEntry) {
+  const message = JSON.stringify({ type: 'log', data: logEntry });
+  for (const client of clients) {
+    if (client.readyState === 1) {
+      client.send(message);
+    }
+  }
+}
+
+/**
+ * Notify all connected clients about request updates
+ * @param {Set} clients - Set of connected WebSocket clients
+ * @param {object} requestService - Request service instance
+ * @param {object} serializationLib - Serialization library instance
+ */
+export function notifyClients(clients, requestService, serializationLib) {
+  const filters = new RequestFilters({ limit: 100 });
+  const requests = requestService.getRequests(filters);
+  const serialized = serializationLib.serializeBigInt(requests);
+  const message = JSON.stringify({ type: 'update', data: serialized });
+  for (const client of clients) {
+    if (client.readyState === 1) {
+      client.send(message);
+    }
+  }
+}

package/ui/server/websocket/handler.js

@@ -0,0 +1,52 @@
+import { Server as ServerConstants } from '#core/constants/Server.js';
+
+/**
+ * Handle WebSocket connection with heartbeat and timeout
+ * @param {Set} clients - Set of connected WebSocket clients
+ * @param {WebSocket} ws - WebSocket connection
+ * @param {object} logger - Logger instance
+ */
+export function handleWebSocketConnection(clients, ws, logger) {
+  clients.add(ws);
+
+  // Set up timeout to close stale connections
+  let timeoutId = setTimeout(() => {
+    if (ws.readyState === 1) {
+      logger?.warn('WebSocket connection timeout, closing');
+      ws.close();
+    }
+  }, ServerConstants.WEBSOCKET_TIMEOUT_MS);
+
+  // Set up heartbeat to keep connection alive
+  const heartbeatInterval = setInterval(() => {
+    if (ws.readyState === 1) {
+      ws.ping();
+    } else {
+      clearInterval(heartbeatInterval);
+    }
+  }, ServerConstants.WEBSOCKET_HEARTBEAT_INTERVAL_MS);
+
+  ws.on('pong', () => {
+    // Reset timeout on pong
+    clearTimeout(timeoutId);
+    timeoutId = setTimeout(() => {
+      if (ws.readyState === 1) {
+        logger?.warn('WebSocket connection timeout, closing');
+        ws.close();
+      }
+    }, ServerConstants.WEBSOCKET_TIMEOUT_MS);
+  });
+
+  ws.on('close', () => {
+    clearTimeout(timeoutId);
+    clearInterval(heartbeatInterval);
+    clients.delete(ws);
+  });
+
+  ws.on('error', (error) => {
+    logger?.error({ error: error.message }, 'WebSocket error');
+    clearTimeout(timeoutId);
+    clearInterval(heartbeatInterval);
+    clients.delete(ws);
+  });
+}

package/ui/server.js

@@ -1,245 +1,66 @@
-import { createServer } from 'node:http';
-import * as path from 'node:path';
-import { fileURLToPath, pathToFileURL } from 'node:url';
-import express from 'express';
-import { WebSocketServer } from 'ws';
+import { pathToFileURL } from 'node:url';
 
-import { getDatabaseFile, prepareAppDataSpaces } from '#common/configs';
-import { openDb } from '#common/db/init';
-import { queryRequests } from '#common/db/query';
-import { restoreOriginalConfig } from './server/utils/config.js';
+import { Environment } from '#core/configs/index.js';
+import { bootstrapLogger } from '#core/libraries/index.js';
 
-import { createBackupRoutes } from './server/routes/backups/index.js';
-import { createCompositeRoutes } from './server/routes/composite/index.js';
-import { createConfigRoutes } from './server/routes/config.js';
-import { createConversationsRoutes } from './server/routes/conversations.js';
-import { createHelpRoutes } from './server/routes/help.js';
-import { createLogsRoutes } from './server/routes/logs.js';
-import { createPlaygroundRoutes } from './server/routes/playground.js';
-import { createRequestsRoutes } from './server/routes/requests.js';
-import { createSessionsRoutes } from './server/routes/sessions.js';
-import { createSettingsRoutes } from './server/routes/settings.js';
-import { createSmartScanRoutes } from './server/routes/smartscan.js';
-import { createStatisticsRoutes } from './server/routes/statistics.js';
-import { serializeBigInt } from './server/utils/serialization.js';
-
-const __filename = fileURLToPath(import.meta.url);
-const __dirname = path.dirname(__filename);
-
-const _MAX_LOG_LINES = 10000;
-
-export function createUIServer() {
-  prepareAppDataSpaces();
-
-  const db = openDb(getDatabaseFile());
-  const app = express();
-  const server = createServer(app);
-  const wss = new WebSocketServer({ server });
-
-  app.use(express.json());
-
-  const clients = new Set();
-  const mcpSharkLogs = [];
-  const processState = { mcpSharkServer: null };
-
-  const setMcpSharkProcess = (server) => {
-    processState.mcpSharkServer = server;
-  };
-
-  wss.on('connection', (ws) => {
-    clients.add(ws);
-    ws.on('close', () => clients.delete(ws));
-  });
-
-  const broadcastLogUpdate = (logEntry) => {
-    const message = JSON.stringify({ type: 'log', data: logEntry });
-    clients.forEach((client) => {
-      if (client.readyState === 1) {
-        client.send(message);
-      }
-    });
-  };
-
-  const restoreConfig = () => {
-    return restoreOriginalConfig(mcpSharkLogs, broadcastLogUpdate);
-  };
-
-  const requestsRoutes = createRequestsRoutes(db);
-  const conversationsRoutes = createConversationsRoutes(db);
-  const sessionsRoutes = createSessionsRoutes(db);
-  const statisticsRoutes = createStatisticsRoutes(db);
-  const logsRoutes = createLogsRoutes(mcpSharkLogs, broadcastLogUpdate);
-  const configRoutes = createConfigRoutes();
-  const backupRoutes = createBackupRoutes();
-  const getMcpSharkProcess = () => processState.mcpSharkServer;
-  const compositeRoutes = createCompositeRoutes(
-    getMcpSharkProcess,
-    setMcpSharkProcess,
-    mcpSharkLogs,
-    broadcastLogUpdate
-  );
-  const helpRoutes = createHelpRoutes();
-  const playgroundRoutes = createPlaygroundRoutes();
-  const smartScanRoutes = createSmartScanRoutes();
-  const settingsRoutes = createSettingsRoutes();
-
-  app.get('/api/requests', requestsRoutes.getRequests);
-  app.get('/api/packets', requestsRoutes.getRequests);
-  app.get('/api/requests/:frameNumber', requestsRoutes.getRequest);
-  app.get('/api/packets/:frameNumber', requestsRoutes.getRequest);
-  app.get('/api/requests/export', requestsRoutes.exportRequests);
-  app.post('/api/requests/clear', requestsRoutes.clearRequests);
-
-  app.get('/api/conversations', conversationsRoutes.getConversations);
-
-  app.get('/api/sessions', sessionsRoutes.getSessions);
-  app.get('/api/sessions/:sessionId/requests', sessionsRoutes.getSessionRequests);
-  app.get('/api/sessions/:sessionId/packets', sessionsRoutes.getSessionRequests);
-
-  app.get('/api/statistics', statisticsRoutes.getStatistics);
-
-  app.get('/api/composite/logs', logsRoutes.getLogs);
-  app.post('/api/composite/logs/clear', logsRoutes.clearLogs);
-  app.get('/api/composite/logs/export', logsRoutes.exportLogs);
-
-  app.post('/api/config/services', configRoutes.extractServices);
-  app.get('/api/config/read', configRoutes.readConfig);
-  app.get('/api/config/detect', configRoutes.detectConfig);
-  app.get('/api/config/backups', backupRoutes.listBackups);
-  app.get('/api/config/backup/view', backupRoutes.viewBackup);
-  app.post('/api/config/restore', (req, res) => {
-    backupRoutes.restoreBackup(req, res, mcpSharkLogs, broadcastLogUpdate);
-  });
-  app.post('/api/config/backup/delete', (req, res) => {
-    backupRoutes.deleteBackup(req, res, mcpSharkLogs, broadcastLogUpdate);
-  });
-
-  app.post('/api/composite/setup', compositeRoutes.setup);
-  app.post('/api/composite/stop', (req, res) => {
-    compositeRoutes.stop(req, res, restoreConfig);
-  });
-  app.get('/api/composite/status', compositeRoutes.getStatus);
-  app.get('/api/composite/servers', compositeRoutes.getServers);
-
-  app.get('/api/help/state', helpRoutes.getState);
-  app.post('/api/help/dismiss', helpRoutes.dismiss);
-  app.post('/api/help/reset', helpRoutes.reset);
-
-  app.post('/api/playground/proxy', playgroundRoutes.proxyRequest);
-
-  app.post('/api/smartscan/scans', smartScanRoutes.createScan);
-  app.get('/api/smartscan/scans', smartScanRoutes.listScans);
-  app.get('/api/smartscan/scans/:scanId', smartScanRoutes.getScan);
-  app.get('/api/smartscan/token', smartScanRoutes.getToken);
-  app.post('/api/smartscan/token', smartScanRoutes.saveToken);
-  app.get('/api/smartscan/discover', smartScanRoutes.discoverServers);
-  app.post('/api/smartscan/scans/batch', smartScanRoutes.createBatchScans);
-  app.post('/api/smartscan/cached-results', smartScanRoutes.getCachedResults);
-  app.post('/api/smartscan/cache/clear', smartScanRoutes.clearCache);
-
-  app.get('/api/settings', settingsRoutes.getSettings);
-
-  const staticPath = path.join(__dirname, 'dist');
-  app.use(express.static(staticPath));
-
-  app.get('*', (_req, res) => {
-    res.sendFile(path.join(staticPath, 'index.html'));
-  });
-
-  const notifyClients = async () => {
-    const requests = queryRequests(db, { limit: 100 });
-    const message = JSON.stringify({ type: 'update', data: serializeBigInt(requests) });
-    clients.forEach((client) => {
-      if (client.readyState === 1) {
-        client.send(message);
-      }
-    });
-  };
-
-  const timestampState = { lastTs: 0 };
-  const intervalId = setInterval(() => {
-    const lastCheck = db.prepare('SELECT MAX(timestamp_ns) as max_ts FROM packets').get();
-    if (lastCheck && lastCheck.max_ts > timestampState.lastTs) {
-      timestampState.lastTs = lastCheck.max_ts;
-      notifyClients();
-    }
-  }, 500);
-
-  const cleanup = async () => {
-    console.log('Shutting down UI server...');
-
-    // Clear interval
-    clearInterval(intervalId);
-
-    // Stop MCP Shark server if running
-    if (processState.mcpSharkServer) {
-      try {
-        if (processState.mcpSharkServer.stop) {
-          await processState.mcpSharkServer.stop();
-        }
-        processState.mcpSharkServer = null;
-      } catch (err) {
-        console.error('Error stopping MCP Shark server:', err);
-      }
-    }
-
-    // Close WebSocket connections
-    clients.forEach((client) => {
-      if (client.readyState === 1) {
-        client.close();
-      }
-    });
-    clients.clear();
-
-    // Close WebSocket server
-    wss.close();
-
-    // Restore config
-    restoreConfig();
-
-    // Close HTTP server
-    return new Promise((resolve) => {
-      server.close(() => {
-        console.log('UI server stopped');
-        resolve();
-      });
-    });
-  };
-
-  return { server, cleanup };
-}
+import { createUIServer } from './server/setup.js';
+import { handleExit, shutdown } from './server/utils/signals.js';
 
+/**
+ * Run UI server with signal handlers
+ */
 export async function runUIServer() {
-  const port = Number.parseInt(process.env.UI_PORT) || 9853;
-  const { server, cleanup } = createUIServer();
-
-  const shutdown = async () => {
-    try {
-      await cleanup();
-    } catch (err) {
-      console.error('Error during shutdown:', err);
-    } finally {
-      process.exit(0);
-    }
-  };
+  const port = Environment.getUiPort();
+  const { server, cleanup, logger, wss } = createUIServer();
 
   // Register signal handlers
-  process.on('SIGTERM', shutdown);
-  process.on('SIGINT', shutdown);
-  process.on('exit', async () => {
-    // Final cleanup on exit
-    try {
-      await cleanup();
-    } catch (_err) {
-      // Ignore errors during exit
+  process.on('SIGTERM', () => shutdown(cleanup, logger));
+  process.on('SIGINT', () => shutdown(cleanup, logger));
+  process.on('exit', () => handleExit(cleanup));
+
+  // Handle server errors (e.g., port already in use)
+  server.on('error', (error) => {
+    if (error.code === 'EADDRINUSE') {
+      logger?.error(
+        {
+          port,
+          error: error.message,
+        },
+        `Port ${port} is already in use. Please stop the existing server or use a different port.`
+      );
+      bootstrapLogger.error(
+        `\n❌ Port ${port} is already in use.\n   Please stop the existing server or set MCP_SHARK_PORT environment variable to use a different port.\n`
+      );
+    } else {
+      logger?.error({ error: error.message, stack: error.stack }, 'Server error');
+      bootstrapLogger.error({ error: error.message }, 'Server error');
     }
+    process.exit(1);
   });
 
+  // Handle WebSocket server errors
+  if (wss) {
+    wss.on('error', (error) => {
+      logger?.error({ error: error.message }, 'WebSocket server error');
+      if (error.code === 'EADDRINUSE') {
+        bootstrapLogger.error(
+          `\n❌ WebSocket port conflict. Port ${port} is already in use.\n   Please stop the existing server or use a different port.\n`
+        );
+      }
+    });
+  }
+
   server.listen(port, '0.0.0.0', () => {
-    console.log(`UI server listening on http://localhost:${port}`);
+    logger?.info({ port }, 'UI server listening');
   });
 }
 
 if (import.meta.url === pathToFileURL(process.argv[1]).href) {
-  runUIServer().catch(console.error);
+  runUIServer().catch((error) => {
+    bootstrapLogger.error(
+      { error: error.message, stack: error.stack },
+      'Failed to start UI server'
+    );
+    process.exit(1);
+  });
 }

package/ui/src/App.jsx

@@ -4,6 +4,7 @@ import CompositeSetup from './CompositeSetup';
 import IntroTour from './IntroTour';
 import SmartScan from './SmartScan';
 import TabNavigation from './TabNavigation';
+import ApiDocsButton from './components/App/ApiDocsButton';
 import HelpButton from './components/App/HelpButton';
 import TrafficTab from './components/App/TrafficTab';
 import { useAppState } from './components/App/useAppState';
@@ -83,6 +84,7 @@ function App() {
       <div style={{ position: 'relative' }} data-tour="tabs">
         <TabNavigation activeTab={activeTab} onTabChange={setActiveTab} />
       </div>
+      <ApiDocsButton />
       <HelpButton
         onClick={() => {
           if (showTour) {

package/ui/src/CompositeSetup.jsx

@@ -2,6 +2,7 @@ import { useEffect, useState } from 'react';
 import BackupList from './components/BackupList';
 import ConfigFileSection from './components/ConfigFileSection';
 import ConfigViewerModal from './components/ConfigViewerModal';
+import ConfirmationModal from './components/ConfirmationModal';
 import MessageDisplay from './components/MessageDisplay';
 import ServerControl from './components/ServerControl';
 import SetupHeader from './components/SetupHeader';
@@ -18,6 +19,8 @@ function CompositeSetup() {
   const [loading, setLoading] = useState(false);
   const [message, setMessage] = useState(null);
   const [error, setError] = useState(null);
+  const [showRestoreModal, setShowRestoreModal] = useState(false);
+  const [pendingRestore, setPendingRestore] = useState({ backupPath: null, originalPath: null });
 
   const { services, selectedServices, setSelectedServices } = useServiceExtraction(
     fileContent,
@@ -52,14 +55,14 @@ function CompositeSetup() {
     return () => clearInterval(interval);
   }, []);
 
-  const handleRestore = async (backupPath, originalPath) => {
-    if (
-      !confirm(
-        'Are you sure you want to restore this backup? This will overwrite the current config file.'
-      )
-    ) {
-      return;
-    }
+  const handleRestoreClick = (backupPath, originalPath) => {
+    setPendingRestore({ backupPath, originalPath });
+    setShowRestoreModal(true);
+  };
+
+  const handleRestore = async () => {
+    const { backupPath, originalPath } = pendingRestore;
+    setShowRestoreModal(false);
 
     try {
       const res = await fetch('/api/config/restore', {
@@ -256,7 +259,7 @@ function CompositeSetup() {
           backups={backups}
           loadingBackups={loadingBackups}
           onRefresh={loadBackups}
-          onRestore={handleRestore}
+          onRestore={handleRestoreClick}
           onView={handleViewBackup}
           onDelete={handleDelete}
         />
@@ -264,6 +267,17 @@ function CompositeSetup() {
         <WhatThisDoesSection filePath={filePath} updatePath={updatePath} />
       </div>
 
+      <ConfirmationModal
+        isOpen={showRestoreModal}
+        onClose={() => setShowRestoreModal(false)}
+        onConfirm={handleRestore}
+        title="Restore Backup"
+        message="Are you sure you want to restore this backup? This will overwrite the current config file."
+        confirmText="Restore"
+        cancelText="Cancel"
+        danger={false}
+      />
+
       <ConfigViewerModal
         viewingConfig={viewingConfig}
         configContent={configContent}

package/ui/src/PacketFilters.jsx

@@ -1,6 +1,7 @@
 import { IconSearch, IconTrash } from '@tabler/icons-react';
 import anime from 'animejs';
 import { useEffect, useRef, useState } from 'react';
+import AlertModal from './components/AlertModal';
 import ConfirmationModal from './components/ConfirmationModal';
 import ExportControls from './components/PacketFilters/ExportControls';
 import FilterInput from './components/PacketFilters/FilterInput';
@@ -10,6 +11,8 @@ import { fadeIn } from './utils/animations';
 function RequestFilters({ filters, onFilterChange, stats, onClear }) {
   const filtersRef = useRef(null);
   const [showClearModal, setShowClearModal] = useState(false);
+  const [showAlertModal, setShowAlertModal] = useState(false);
+  const [alertMessage, setAlertMessage] = useState('');
 
   useEffect(() => {
     if (filtersRef.current) {
@@ -56,7 +59,8 @@ function RequestFilters({ filters, onFilterChange, stats, onClear }) {
       document.body.removeChild(a);
     } catch (error) {
       console.error('Failed to export traffic:', error);
-      alert('Failed to export traffic. Please try again.');
+      setAlertMessage('Failed to export traffic. Please try again.');
+      setShowAlertModal(true);
     }
   };
 
@@ -217,11 +221,13 @@ function RequestFilters({ filters, onFilterChange, stats, onClear }) {
               }
             } else {
               const error = await response.json();
-              alert(`Failed to clear traffic: ${error.error || 'Unknown error'}`);
+              setAlertMessage(`Failed to clear traffic: ${error.error || 'Unknown error'}`);
+              setShowAlertModal(true);
             }
           } catch (error) {
             console.error('Failed to clear traffic:', error);
-            alert('Failed to clear traffic. Please try again.');
+            setAlertMessage('Failed to clear traffic. Please try again.');
+            setShowAlertModal(true);
           }
         }}
         title="Clear All Captured Traffic"
@@ -230,6 +236,14 @@ function RequestFilters({ filters, onFilterChange, stats, onClear }) {
         cancelText="Cancel"
         danger={true}
       />
+
+      <AlertModal
+        isOpen={showAlertModal}
+        onClose={() => setShowAlertModal(false)}
+        title="Error"
+        message={alertMessage}
+        type="error"
+      />
     </div>
   );
 }