feat: add ComfyUI toggle for Prometheus metrics server

nodes/api/__init__.py

@@ -90,6 +90,7 @@ async def control_server(request):
             # Extract host and port from settings if provided
             host = settings.get("host") if settings else None
             port = settings.get("port") if settings else None
+            enable_metrics = settings.get("enableMetrics") if settings else None
 
             if action == "status":
                 # Simply return the current server status
@@ -98,7 +99,11 @@ async def control_server(request):
                     "status": server_manager.get_status()
                 })
             elif action == "start":
-                success = await server_manager.start(port=port, host=host)
+                success = await server_manager.start(
+                    port=port,
+                    host=host,
+                    enable_metrics=enable_metrics
+                )
                 return web.json_response({
                     "success": success,
                     "status": server_manager.get_status()
@@ -120,7 +125,11 @@ async def control_server(request):
                         "message": "Forced server shutdown due to error"
                     })
             elif action == "restart":
-                success = await server_manager.restart(port=port, host=host)
+                success = await server_manager.restart(
+                    port=port,
+                    host=host,
+                    enable_metrics=enable_metrics
+                )
                 return web.json_response({
                     "success": success,
                     "status": server_manager.get_status()
@@ -178,10 +187,16 @@ async def manage_configuration(request):
                 name = data.get("name")
                 host = data.get("host")
                 port = data.get("port")
+                enable_metrics = data.get("enableMetrics")
                 if not name or not host or not port:
                     return web.json_response({"error": "Missing required parameters"}, status=400)
 
-                success = settings_storage.add_configuration(name, host, port)
+                success = settings_storage.add_configuration(
+                    name, 
+                    host, 
+                    port, 
+                    enable_metrics
+                )
                 return web.json_response({
                     "success": success,
                     "settings": settings_storage.load_settings()
@@ -214,4 +229,3 @@ async def manage_configuration(request):
         except Exception as e:
             logging.error(f"Error managing configuration: {str(e)}")
             return web.json_response({"error": str(e)}, status=500)
-

nodes/server_manager.py

@@ -34,15 +34,18 @@ def __init__(self, host="0.0.0.0", port=None):
         self.host = host
         self.port = port
         self.is_running = False
+        self.enable_metrics = False
         logging.info(f"Initializing {self.__class__.__name__}")
 
     @abstractmethod
-    async def start(self, port=None, host=None) -> bool:
+    async def start(self, port=None, host=None, enable_metrics=None) -> bool:
         """Start the ComfyStream server
 
         Args:
             port: Optional port to use. If None, implementation should choose a port.
             host: Optional host to use. If None, implementation should use the default host.
+            enable_metrics: Optional flag to enable metrics endpoint. If None, use the
+                current setting.
 
         Returns:
             bool: True if server started successfully, False otherwise
@@ -76,22 +79,32 @@ def check_server_health(self) -> bool:
         """
         pass
 
-    async def restart(self, port=None, host=None) -> bool:
-        """Restart the ComfyStream server
+    async def restart(self, port=None, host=None, enable_metrics=None) -> bool:
+        """Restart the ComfyStream server.
 
         Args:
             port: Optional port to use. If None, use the current port.
             host: Optional host to use. If None, use the current host.
+            enable_metrics: Optional flag to enable metrics endpoint. If None, use the
+                current setting.
 
         Returns:
-            bool: True if server restarted successfully, False otherwise
+            bool: True if server restarted successfully, False otherwise.
         """
         logging.info("Restarting ComfyStream server...")
         # Use provided values or current values
         port_to_use = port if port is not None else self.port
         host_to_use = host if host is not None else self.host
+        enable_metrics = (
+            enable_metrics if enable_metrics is not None else self.enable_metrics
+        )
+
         await self.stop()
-        return await self.start(port=port_to_use, host=host_to_use)
+        return await self.start(
+            port=port_to_use,
+            host=host_to_use,
+            enable_metrics=enable_metrics
+        )
 
 class LocalComfyStreamServer(ComfyStreamServerBase):
     """Local ComfyStream server implementation"""
@@ -136,8 +149,18 @@ def log_subprocess_output(self, pipe, level):
         for line in iter(pipe.readline, b''):
             logging.log(level, line.decode().strip())
 
-    async def start(self, port=None, host=None):
-        """Start the ComfyStream server"""
+    async def start(self, port=None, host=None, enable_metrics=None):
+        """Start the ComfyStream server.
+
+        Args:
+            port: Optional port to use. If None, find an available port.
+            host: Optional host to use. If None, use the current host.
+            enable_metrics: Optional flag to enable metrics endpoint. If None, use the
+                current setting.
+
+        Returns:
+            bool: True if server started successfully, False otherwise.
+        """
         if self.is_running:
             logging.info("Server is already running")
             return False
@@ -146,7 +169,9 @@ async def start(self, port=None, host=None):
             self.port = port or self.find_available_port()
             if host is not None:
                 self.host = host
-
+            if enable_metrics is not None:
+                self.enable_metrics = enable_metrics
+
             # Get the path to the ComfyStream server directory and script
             server_dir = Path(__file__).parent.parent / "server"
             server_script = server_dir / "app.py"
@@ -161,6 +186,8 @@ async def start(self, port=None, host=None):
                   "--port", str(self.port),
                   "--host", str(self.host),
                   "--workspace", str(comfyui_workspace)]
+            if self.enable_metrics:
+                cmd.append("--monitor")
 
             logging.info(f"Starting server with command: {' '.join(cmd)}")
 
@@ -247,4 +274,4 @@ def cleanup(self):
             except Exception as e:
                 logging.error(f"Error cleaning up server process: {str(e)}")
             self.process = None
-            self.is_running = False
+            self.is_running = False

nodes/settings_storage.py

@@ -15,6 +15,7 @@
 DEFAULT_SETTINGS = {
     "host": "0.0.0.0",
     "port": 8889,
+    "enableMetrics": False,
     "configurations": [],
     "selectedConfigIndex": -1
 }
@@ -78,12 +79,17 @@ def update_settings(new_settings):
 
     return save_settings(current_settings)
 
-def add_configuration(name, host, port):
+def add_configuration(name, host, port, enable_metrics):
     """Add a new configuration"""
     settings = load_settings()
 
     # Create the new configuration
-    config = {"name": name, "host": host, "port": port}
+    config = {
+        "name": name,
+        "host": host,
+        "port": port,
+        "enableMetrics": enable_metrics
+    }
 
     # Add to configurations list
     settings["configurations"].append(config)
@@ -125,9 +131,10 @@ def select_configuration(index):
             config = settings["configurations"][index]
             settings["host"] = config["host"]
             settings["port"] = config["port"]
+            settings["enableMetrics"] = config["enableMetrics"]
 
         # Save updated settings
         return save_settings(settings)
     else:
         logging.error(f"Invalid configuration index: {index}")
-        return False 
+        return False 

nodes/web/js/launcher.js

@@ -101,7 +101,7 @@ document.addEventListener('comfy-extension-registered', (event) => {
 async function controlServer(action) {
     try {
         // Get settings from the settings manager
-        const settings = settingsManager.getCurrentHostPort();
+        const settings = settingsManager.getCurrentServerSettings();
 
         // Set transitional state based on action
         if (action === 'start') {
@@ -308,4 +308,4 @@ const extension = {
     }
 };
 
-app.registerExtension(extension); 
+app.registerExtension(extension);