From e56e3288e1adde7ff5b1a7fca74ac7ec9db303a4 Mon Sep 17 00:00:00 2001
From: shanks-20million <shanks.20million@gmail.com>
Date: Tue, 29 Jul 2025 11:09:17 +0530
Subject: [PATCH] feat: add authentication token support for ComfyUI API calls

- Add ComfyUIToken setting to frontend configuration panel
- Implement token-based authentication for all ComfyUI API endpoints:
  - WebSocket connections (/ws)
  - Prompt execution (/prompt)
  - Image retrieval (/view)
  - History access (/history)
- Support RunPod-style authentication with query parameter tokens
- Maintain backward compatibility when no token is configured
- Enable secure access to protected ComfyUI instances
---
 config.py                   |  5 +++++
 js/index.js                 |  5 +++++
 services/comfyui_service.py | 32 +++++++++++++++++++++++---------
 3 files changed, 33 insertions(+), 9 deletions(-)

diff --git a/config.py b/config.py
index e24d9a2..d7c51c2 100644
--- a/config.py
+++ b/config.py
@@ -81,6 +81,11 @@ class Config:
             port = self._get_port_from_args()
         
         return f"{host}:{port}"
+    
+    @property
+    def comfy_token(self):
+        """Get ComfyUI authentication token from settings"""
+        return self.user_settings.get("Connect.ComfyUIToken", "")
 
 
 config = Config()
diff --git a/js/index.js b/js/index.js
index 2205b07..bb1ea1e 100644
--- a/js/index.js
+++ b/js/index.js
@@ -15,6 +15,11 @@ app.registerExtension({
       name: "ComfyUI Port",
       type: "text",
     },
+    {
+      id: "Connect.ComfyUIToken",
+      name: "ComfyUI Authentication Token",
+      type: "text",
+    },
     {
       id: "Connect.GatewayEndpoint",
       name: "ComfyUI Gateway Endpoint",
diff --git a/services/comfyui_service.py b/services/comfyui_service.py
index 0b1b642..bdbaeba 100644
--- a/services/comfyui_service.py
+++ b/services/comfyui_service.py
@@ -43,9 +43,13 @@ class ComfyUIService:
             return
 
         self.session = aiohttp.ClientSession()
-        self.ws = await self.session.ws_connect(
-            f"ws://{config.comfy_endpoint}/ws?clientId={self.CLIENT_ID}"
-        )
+        
+        # Build WebSocket URL with token if available
+        ws_url = f"ws://{config.comfy_endpoint}/ws?clientId={self.CLIENT_ID}"
+        if config.comfy_token:
+            ws_url += f"&token={config.comfy_token}"
+            
+        self.ws = await self.session.ws_connect(ws_url)
         # Start the global websocket listener
         self._listener_task = asyncio.create_task(self._listen_websocket())
         self._connected = True
@@ -93,15 +97,21 @@ class ComfyUIService:
         await self._ensure_connected()
         payload = {"prompt": prompt, "client_id": self.CLIENT_ID}
         data = json.dumps(payload).encode("utf-8")
-        async with self.session.post(
-            f"http://{config.comfy_endpoint}/prompt", data=data
-        ) as response:
+        
+        # Build URL with token if available
+        url = f"http://{config.comfy_endpoint}/prompt"
+        if config.comfy_token:
+            url += f"?token={config.comfy_token}"
+            
+        async with self.session.post(url, data=data) as response:
             return await response.json()
 
     async def get_image(self, filename, subfolder, folder_type):
         """Retrieve an image from ComfyUI"""
         await self._ensure_connected()
         params = {"filename": filename, "subfolder": subfolder, "type": folder_type}
+        if config.comfy_token:
+            params["token"] = config.comfy_token
         url_values = urllib.parse.urlencode(params)
         async with self.session.get(
             f"http://{config.comfy_endpoint}/view?{url_values}"
@@ -113,9 +123,13 @@ class ComfyUIService:
     async def get_history(self, prompt_id):
         """Get execution history for a prompt"""
         await self._ensure_connected()
-        async with self.session.get(
-            f"http://{config.comfy_endpoint}/history/{prompt_id}"
-        ) as response:
+        
+        # Build URL with token if available, works for comfyui-login plugin
+        url = f"http://{config.comfy_endpoint}/history/{prompt_id}"
+        if config.comfy_token:
+            url += f"?token={config.comfy_token}"
+            
+        async with self.session.get(url) as response:
             return await response.json()
 
     async def run_workflow(self, workflow: dict) -> dict: