From b768e7fc8aef50c8de1c249e4102f6aae470d6c6 Mon Sep 17 00:00:00 2001 From: robertvoy Date: Tue, 30 Dec 2025 12:11:12 +1100 Subject: [PATCH] https handling in host field --- distributed.py | 31 ++++++++++++++++++++++--------- utils/cloudflare.py | 18 +++++++++++++++--- web/main.js | 41 ++++++++++++++++++++++++++++++++++------- 3 files changed, 71 insertions(+), 19 deletions(-) diff --git a/distributed.py b/distributed.py index d171483..175e360 100644 --- a/distributed.py +++ b/distributed.py @@ -4,6 +4,7 @@ from PIL import Image import folder_paths import os import json +import re import asyncio import aiohttp from aiohttp import web @@ -60,6 +61,17 @@ def cleanup(): atexit.register(cleanup) +def normalize_host(value): + if value is None: + return None + if not isinstance(value, str): + return value + host = value.strip() + if not host: + return host + host = re.sub(r"^https?://", "", host, flags=re.IGNORECASE) + return host.split("/")[0] + # --- API Endpoints --- @server.PromptServer.instance.routes.get("/distributed/config") async def get_config_endpoint(request): @@ -367,7 +379,7 @@ async def update_worker_endpoint(request): if data["host"] is None: worker.pop("host", None) else: - worker["host"] = data["host"] + worker["host"] = normalize_host(data["host"]) # Handle cuda_device field - remove it if None if "cuda_device" in data: @@ -396,7 +408,7 @@ async def update_worker_endpoint(request): new_worker = { "id": worker_id, "name": data["name"], - "host": data.get("host", "localhost"), + "host": normalize_host(data.get("host", "localhost")), "port": data["port"], "cuda_device": data["cuda_device"], "enabled": data.get("enabled", False), @@ -673,7 +685,8 @@ async def get_local_worker_status_endpoint(request): for worker in config.get("workers", []): # Only check local workers - if not worker.get("host") or worker.get("host") in ["localhost", "127.0.0.1"]: + host = normalize_host(worker.get("host")) or "" + if not host or host in ["localhost", "127.0.0.1"]: worker_id = worker["id"] port = worker["port"] @@ -1324,7 +1337,7 @@ async def get_client_session(): # Local worker detection functions async def is_local_worker(worker_config): """Check if a worker is running on the same machine as the master.""" - host = worker_config.get('host', 'localhost') + host = normalize_host(worker_config.get('host', 'localhost')) or 'localhost' if host in ['localhost', '127.0.0.1', '0.0.0.0', ''] or worker_config.get('type') == 'local': return True @@ -1341,7 +1354,7 @@ async def is_same_physical_host(worker_config): master_machine_id = get_machine_id() # Fetch worker's machine ID via API - host = worker_config.get('host', 'localhost') + host = normalize_host(worker_config.get('host', 'localhost')) or 'localhost' port = worker_config.get('port', 8188) session = await get_client_session() @@ -1404,11 +1417,11 @@ async def get_comms_channel(worker_id, worker_config): return f"http://127.0.0.1:{worker_config['port']}" elif worker_config.get('type') == 'cloud' and is_runpod_environment(): # Runpod same-host optimization (if detected) - host = worker_config.get('host', 'localhost') + host = normalize_host(worker_config.get('host', 'localhost')) or 'localhost' return f"http://{host}:{worker_config['port']}" else: # Remote worker: use configured endpoint - host = worker_config.get('host', 'localhost') + host = normalize_host(worker_config.get('host', 'localhost')) or 'localhost' return f"http://{host}:{worker_config['port']}" # Auto-launch workers if enabled @@ -1435,7 +1448,7 @@ def auto_launch_workers(): worker_name = worker.get('name', f'Worker {worker_id}') # Skip remote workers - host = worker.get('host', 'localhost').lower() + host = (normalize_host(worker.get('host', 'localhost')) or 'localhost').lower() if host not in ['localhost', '127.0.0.1', '', None]: debug_log(f"Skipping remote worker {worker_name} (host: {host})") continue @@ -2121,7 +2134,7 @@ class DistributedCollectorNode: if not wrec: debug_log(f"Collector probe: worker {wid} not found in config") continue - host = wrec.get('host') or 'localhost' + host = normalize_host(wrec.get('host') or 'localhost') or 'localhost' port = int(wrec.get('port', 8188)) url = f"http://{host}:{port}/prompt" try: diff --git a/utils/cloudflare.py b/utils/cloudflare.py index b62dfb6..00953dc 100644 --- a/utils/cloudflare.py +++ b/utils/cloudflare.py @@ -30,6 +30,14 @@ PUBLIC_URL_PATTERN = re.compile(r"(https?://[\w.-]+\.(?:trycloudflare\.com|cloud TUNNEL_START_TIMEOUT = 25 +def _normalize_host(value): + if not value or not isinstance(value, str): + return "" + value = value.strip() + value = re.sub(r"^https?://", "", value, flags=re.IGNORECASE) + return value.split("/")[0] + + class CloudflareTunnelManager: def __init__(self): self.process = None @@ -300,13 +308,14 @@ class CloudflareTunnelManager: raise RuntimeError(error_msg) debug_log(f"Cloudflare tunnel ready at {self.public_url}") + master_host = _normalize_host(self.public_url) self._persist_state( status="running", public_url=self.public_url, pid=self.pid, log_file=self.log_file, previous_host=self.previous_master_host or "", - master_host=self.public_url + master_host=master_host ) return { "status": self.status, @@ -339,8 +348,11 @@ class CloudflareTunnelManager: active_url = tunnel_cfg.get("public_url") current_master_host = (config.get("master") or {}).get("host") restore_host = None - if active_url and current_master_host == active_url: - restore_host = self.previous_master_host or "" + if active_url: + active_host = _normalize_host(active_url) + current_host = _normalize_host(current_master_host) + if current_host == active_host: + restore_host = self.previous_master_host or "" self.status = "stopped" self.public_url = None diff --git a/web/main.js b/web/main.js index f8d4abe..485f7ae 100644 --- a/web/main.js +++ b/web/main.js @@ -158,6 +158,24 @@ class DistributedExtension { } } + _parseHostInput(value) { + if (!value) { + return { host: "", port: null }; + } + let cleaned = value.trim().replace(/^https?:\/\//i, ""); + cleaned = cleaned.split("/")[0]; + try { + const url = new URL(`http://${cleaned}`); + const port = url.port ? parseInt(url.port, 10) : null; + return { + host: url.hostname || cleaned, + port: Number.isFinite(port) ? port : null, + }; + } catch (error) { + return { host: cleaned, port: null }; + } + } + updateTunnelUIElements() { const elements = this.tunnelElements || {}; const status = (this.tunnelStatus?.status || "stopped").toLowerCase(); @@ -501,7 +519,9 @@ class DistributedExtension { // Helper to build worker URL getWorkerUrl(worker, endpoint = '') { - const host = worker.host || window.location.hostname; + const parsed = this._parseHostInput(worker.host || window.location.hostname); + const host = parsed.host || window.location.hostname; + const resolvedPort = parsed.port || worker.port; // Cloud workers always use HTTPS const isCloud = worker.type === 'cloud'; @@ -524,13 +544,13 @@ class DistributedExtension { } // Determine protocol: HTTPS for cloud, Runpod proxies, or port 443 - const useHttps = isCloud || isRunpodProxy || worker.port === 443; + const useHttps = isCloud || isRunpodProxy || resolvedPort === 443; const protocol = useHttps ? 'https' : 'http'; // Only add port if non-standard const defaultPort = useHttps ? 443 : 80; - const needsPort = !isRunpodProxy && worker.port !== defaultPort; - const portStr = needsPort ? `:${worker.port}` : ''; + const needsPort = !isRunpodProxy && resolvedPort !== defaultPort; + const portStr = needsPort ? `:${resolvedPort}` : ''; return `${protocol}://${finalHost}${portStr}${endpoint}`; } @@ -974,7 +994,8 @@ class DistributedExtension { return true; } // Otherwise check by host (backward compatibility) - const host = worker.host || window.location.hostname; + const parsed = this._parseHostInput(worker.host || window.location.hostname); + const host = parsed.host || window.location.hostname; return host !== "localhost" && host !== "127.0.0.1" && host !== window.location.hostname; } @@ -1194,10 +1215,16 @@ class DistributedExtension { const workerType = document.getElementById(`worker-type-${workerId}`).value; const isRemote = workerType === 'remote' || workerType === 'cloud'; const isCloud = workerType === 'cloud'; - const host = isRemote ? document.getElementById(`host-${workerId}`).value : window.location.hostname; - const port = parseInt(document.getElementById(`port-${workerId}`).value); + const rawHost = isRemote ? document.getElementById(`host-${workerId}`).value : window.location.hostname; + const parsedHost = isRemote ? this._parseHostInput(rawHost) : { host: window.location.hostname, port: null }; + const host = isRemote ? parsedHost.host : window.location.hostname; + let port = parseInt(document.getElementById(`port-${workerId}`).value); const cudaDevice = isRemote ? undefined : parseInt(document.getElementById(`cuda-${workerId}`).value); const extraArgs = isRemote ? undefined : document.getElementById(`args-${workerId}`).value; + + if (isRemote && Number.isFinite(parsedHost.port)) { + port = parsedHost.port; + } // Validate if (!name.trim()) {