Commit all current refactor changes

.gitignore

@@ -4,3 +4,4 @@ gpu_config.json
 __pycache__/
 AGENTS.md
 docs/REFACTOR_PLAN.md
+docs/refactor-validation-checklist.md

api/job_routes.py

@@ -4,6 +4,7 @@ import io
 import os
 import base64
 import binascii
+import time
 
 from aiohttp import web
 import server
@@ -11,12 +12,13 @@ import server as _server
 import torch
 from PIL import Image
 
-from ..utils.logging import debug_log, log
+from ..utils.logging import debug_log
 from ..utils.image import pil_to_tensor, ensure_contiguous
 from ..utils.network import handle_api_error
-from ..utils.constants import MEMORY_CLEAR_DELAY
+from ..utils.constants import JOB_INIT_GRACE_PERIOD, MEMORY_CLEAR_DELAY
 from ..utils.async_helpers import queue_prompt_payload
 from .queue_orchestration import orchestrate_distributed_execution
+from .queue_request import parse_queue_request_payload
 
 prompt_server = _server.PromptServer.instance
 
@@ -193,54 +195,23 @@ async def clear_memory_endpoint(request):
 async def distributed_queue_endpoint(request):
     """Queue a distributed workflow, mirroring the UI orchestration pipeline."""
     try:
-        data = await request.json()
+        raw_payload = await request.json()
     except Exception as exc:
         return await handle_api_error(request, f"Invalid JSON payload: {exc}", 400)
 
-    auto_prepare = bool(data.get("auto_prepare"))
-    prompt = data.get("prompt")
-    if prompt is None and auto_prepare:
-        # Transitional compatibility: allow payload.workflow.prompt when auto_prepare is enabled.
-        workflow_payload = data.get("workflow")
-        if isinstance(workflow_payload, dict):
-            candidate_prompt = workflow_payload.get("prompt")
-            if isinstance(candidate_prompt, dict):
-                prompt = candidate_prompt
-
-    if not isinstance(prompt, dict):
-        return await handle_api_error(request, "Field 'prompt' must be an object", 400)
-
-    # Accept either enabled_worker_ids (legacy) or workers list (auto-prepare path).
-    workers_field = data.get("workers")
-    if workers_field is not None and data.get("enabled_worker_ids") is None:
-        if isinstance(workers_field, list):
-            enabled_ids = []
-            for entry in workers_field:
-                if isinstance(entry, dict):
-                    worker_id = entry.get("id")
-                else:
-                    worker_id = entry
-                if worker_id is not None:
-                    enabled_ids.append(str(worker_id))
-            data["enabled_worker_ids"] = enabled_ids
-
-    workflow_meta = data.get("workflow")
-    client_id = data.get("client_id")
-    delegate_master = data.get("delegate_master")
-    enabled_ids = data.get("enabled_worker_ids")
-
-    if enabled_ids is not None:
-        if not isinstance(enabled_ids, list):
-            return await handle_api_error(request, "enabled_worker_ids must be a list of worker IDs", 400)
-        enabled_ids = [str(worker_id) for worker_id in enabled_ids]
+    try:
+        payload = parse_queue_request_payload(raw_payload)
+    except ValueError as exc:
+        return await handle_api_error(request, exc, 400)
 
     try:
         prompt_id, worker_count = await orchestrate_distributed_execution(
-            prompt,
-            workflow_meta,
-            client_id,
-            enabled_worker_ids=enabled_ids,
-            delegate_master=delegate_master,
+            payload.prompt,
+            payload.workflow_meta,
+            payload.client_id,
+            enabled_worker_ids=payload.enabled_worker_ids,
+            delegate_master=payload.delegate_master,
+            trace_execution_id=payload.trace_execution_id,
         )
         return web.json_response({
             "prompt_id": prompt_id,
@@ -320,26 +291,34 @@ async def job_complete_endpoint(request):
         multi_job_id = job_id.strip()
         worker_id = worker_id.strip()
 
-        async with prompt_server.distributed_jobs_lock:
-            pending = prompt_server.distributed_pending_jobs.get(multi_job_id)
-            if pending is None:
-                log(f"ERROR: Job {multi_job_id} not found in distributed_pending_jobs")
-                return await handle_api_error(request, "Job not found or already complete", 404)
+        pending = None
+        queue_size = 0
+        deadline = time.monotonic() + float(JOB_INIT_GRACE_PERIOD)
+        while pending is None:
+            async with prompt_server.distributed_jobs_lock:
+                pending = prompt_server.distributed_pending_jobs.get(multi_job_id)
+                if pending is not None:
+                    await pending.put(
+                        {
+                            "tensor": tensor,
+                            "worker_id": worker_id,
+                            "image_index": int(batch_idx),
+                            "is_last": is_last,
+                            "audio": None,
+                        }
+                    )
+                    queue_size = pending.qsize()
+                    break
 
-            await pending.put(
-                {
-                    "tensor": tensor,
-                    "worker_id": worker_id,
-                    "image_index": int(batch_idx),
-                    "is_last": is_last,
-                    "audio": None,
-                }
-            )
-            debug_log(
-                f"job_complete received canonical envelope - job_id: {multi_job_id}, "
-                f"worker: {worker_id}, batch_idx: {batch_idx}, is_last: {is_last}, "
-                f"queue_size: {pending.qsize()}"
-            )
+            if time.monotonic() > deadline:
+                return await handle_api_error(request, "job not initialized", 404)
+            await asyncio.sleep(0.05)
+
+        debug_log(
+            f"job_complete received canonical envelope - job_id: {multi_job_id}, "
+            f"worker: {worker_id}, batch_idx: {batch_idx}, is_last: {is_last}, "
+            f"queue_size: {queue_size}"
+        )
 
         return web.json_response({"status": "success"})
     except Exception as e:

api/orchestration/__init__.py

@@ -0,0 +1,4 @@
+# Orchestration helpers split out from queue_orchestration.py:
+# - prompt_transform.py: graph pruning + hidden input overrides
+# - media_sync.py: remote media/path normalization
+# - dispatch.py: worker probe + prompt dispatch

api/orchestration/dispatch.py

@@ -0,0 +1,184 @@
+import asyncio
+import json
+import uuid
+
+import aiohttp
+
+from ...utils.logging import debug_log, log
+from ...utils.network import build_worker_url, get_client_session, probe_worker
+
+
+def _trace_prefix(trace_execution_id):
+    return f"[Distributed][exec:{trace_execution_id}]"
+
+
+def _trace_debug(trace_execution_id, message):
+    debug_log(f"{_trace_prefix(trace_execution_id)} {message}")
+
+
+def _trace_info(trace_execution_id, message):
+    log(f"{_trace_prefix(trace_execution_id)} {message}")
+
+
+async def worker_is_active(worker):
+    """Ping worker's /prompt endpoint to confirm it's reachable."""
+    url = build_worker_url(worker)
+    return await probe_worker(url, timeout=3.0) is not None
+
+
+async def worker_ws_is_active(worker):
+    """Ping worker's websocket endpoint to confirm it's reachable."""
+    session = await get_client_session()
+    url = build_worker_url(worker, "/distributed/worker_ws")
+    try:
+        ws = await session.ws_connect(url, heartbeat=20, timeout=3)
+        await ws.close()
+        return True
+    except asyncio.TimeoutError:
+        debug_log(f"[Distributed] Worker WS probe timed out: {url}")
+        return False
+    except aiohttp.ClientConnectorError:
+        debug_log(f"[Distributed] Worker WS unreachable: {url}")
+        return False
+    except Exception as e:
+        debug_log(f"[Distributed] Worker WS probe unexpected error: {e}")
+        return False
+
+
+def _coerce_positive_int(value, default):
+    try:
+        parsed = int(value)
+    except (TypeError, ValueError):
+        return max(1, int(default))
+    return max(1, parsed)
+
+
+async def _probe_worker_active(worker, use_websocket, semaphore):
+    async with semaphore:
+        is_active = await (worker_ws_is_active(worker) if use_websocket else worker_is_active(worker))
+        return worker, is_active
+
+
+async def _dispatch_via_websocket(worker_url, payload, client_id, timeout=60.0):
+    """Open a fresh worker websocket, dispatch one prompt, wait for ack, then close."""
+    request_id = uuid.uuid4().hex
+    ws_payload = {
+        "type": "dispatch_prompt",
+        "request_id": request_id,
+        "prompt": payload.get("prompt"),
+        "workflow": payload.get("workflow"),
+        "client_id": client_id,
+    }
+    ws_url = worker_url.replace("http://", "ws://").replace("https://", "wss://")
+    ws_url = f"{ws_url}/distributed/worker_ws"
+    session = await get_client_session()
+
+    async with session.ws_connect(ws_url, heartbeat=20, timeout=timeout) as ws:
+        await ws.send_json(ws_payload)
+        async for msg in ws:
+            if msg.type == aiohttp.WSMsgType.TEXT:
+                data = json.loads(msg.data or "{}")
+                if data.get("type") == "dispatch_ack" and data.get("request_id") == request_id:
+                    if data.get("ok"):
+                        return
+                    raise RuntimeError(data.get("error") or "Worker rejected websocket dispatch.")
+            elif msg.type in (aiohttp.WSMsgType.ERROR, aiohttp.WSMsgType.CLOSED):
+                raise RuntimeError(f"Worker websocket closed unexpectedly: {msg.type}")
+
+    raise RuntimeError("Worker websocket closed before dispatch_ack was received.")
+
+
+async def dispatch_worker_prompt(
+    worker,
+    prompt_obj,
+    workflow_meta,
+    client_id=None,
+    use_websocket=False,
+    trace_execution_id=None,
+):
+    """Send the prepared prompt to a worker ComfyUI instance."""
+    worker_url = build_worker_url(worker)
+    url = build_worker_url(worker, "/prompt")
+    payload = {"prompt": prompt_obj}
+    extra_data = {}
+    if workflow_meta:
+        extra_data.setdefault("extra_pnginfo", {})["workflow"] = workflow_meta
+    if extra_data:
+        payload["extra_data"] = extra_data
+
+    if use_websocket:
+        try:
+            await _dispatch_via_websocket(
+                worker_url,
+                {
+                    "prompt": prompt_obj,
+                    "workflow": workflow_meta,
+                },
+                client_id,
+            )
+            return
+        except Exception as exc:
+            worker_id = worker.get("id")
+            if trace_execution_id:
+                _trace_info(trace_execution_id, f"Websocket dispatch failed for worker {worker_id}: {exc}")
+            else:
+                log(f"[Distributed] Websocket dispatch failed for worker {worker_id}: {exc}")
+            raise
+
+    session = await get_client_session()
+    async with session.post(
+        url,
+        json=payload,
+        timeout=aiohttp.ClientTimeout(total=60),
+    ) as resp:
+        resp.raise_for_status()
+
+
+async def select_active_workers(
+    workers,
+    use_websocket,
+    delegate_master,
+    trace_execution_id=None,
+    probe_concurrency=8,
+):
+    """Probe workers and return (active_workers, updated_delegate_master)."""
+    probe_limit = _coerce_positive_int(probe_concurrency, 8)
+    probe_semaphore = asyncio.Semaphore(probe_limit)
+
+    if trace_execution_id and workers:
+        _trace_debug(
+            trace_execution_id,
+            f"Probing {len(workers)} workers with probe_concurrency={probe_limit}",
+        )
+
+    probe_results = await asyncio.gather(
+        *[
+            _probe_worker_active(worker, use_websocket, probe_semaphore)
+            for worker in workers
+        ]
+    )
+
+    active_workers = []
+    for worker, is_active in probe_results:
+        if is_active:
+            active_workers.append(worker)
+        else:
+            if trace_execution_id:
+                _trace_info(trace_execution_id, f"Worker {worker['name']} ({worker['id']}) is offline, skipping.")
+            else:
+                log(f"[Distributed] Worker {worker['name']} ({worker['id']}) is offline, skipping.")
+
+    if trace_execution_id and workers:
+        _trace_debug(
+            trace_execution_id,
+            f"Worker probe complete: active={len(active_workers)}/{len(workers)}",
+        )
+
+    if not active_workers and delegate_master:
+        if trace_execution_id:
+            _trace_debug(trace_execution_id, "All workers offline while delegate-only requested; enabling master participation.")
+        else:
+            debug_log("All workers offline while delegate-only requested; enabling master participation.")
+        delegate_master = False
+
+    return active_workers, delegate_master

api/orchestration/media_sync.py

@@ -0,0 +1,220 @@
+import asyncio
+import hashlib
+import mimetypes
+import os
+import re
+
+import aiohttp
+
+from ...utils.logging import debug_log, log
+from ...utils.network import build_worker_url, get_client_session
+
+
+LIKELY_FILENAME_RE = re.compile(
+    r"\.(ckpt|safetensors|pt|pth|bin|yaml|json|png|jpg|jpeg|webp|gif|bmp|latent|txt|vae|lora|embedding)"
+    r"(\s*\[\w+\])?$",
+    re.IGNORECASE,
+)
+IMAGE_OR_VIDEO_RE = re.compile(
+    r"\.(png|jpg|jpeg|webp|gif|bmp|mp4|avi|mov|mkv|webm)(\s*\[\w+\])?$",
+    re.IGNORECASE,
+)
+
+
+def _trace_prefix(trace_execution_id):
+    return f"[Distributed][exec:{trace_execution_id}]"
+
+
+def _trace_debug(trace_execution_id, message):
+    debug_log(f"{_trace_prefix(trace_execution_id)} {message}")
+
+
+def _trace_info(trace_execution_id, message):
+    log(f"{_trace_prefix(trace_execution_id)} {message}")
+
+
+def convert_paths_for_platform(obj, target_separator):
+    """Recursively normalize likely file paths for the worker platform separator."""
+    if target_separator not in ("/", "\\"):
+        return obj
+
+    def _convert(value):
+        if isinstance(value, str):
+            if ("/" in value or "\\" in value) and LIKELY_FILENAME_RE.search(value):
+                trimmed = value.strip()
+                has_drive = bool(re.match(r"^[A-Za-z]:(\\\\|/)", trimmed))
+                is_absolute = trimmed.startswith("/") or trimmed.startswith("\\\\")
+                has_protocol = bool(re.match(r"^\w+://", trimmed))
+
+                # Keep relative media-style paths in forward-slash form (Comfy-style annotated paths).
+                if not has_drive and not is_absolute and not has_protocol and IMAGE_OR_VIDEO_RE.search(trimmed):
+                    return re.sub(r"[\\]+", "/", trimmed)
+
+                if target_separator == "\\":
+                    return re.sub(r"[\\/]+", r"\\", trimmed)
+                return re.sub(r"[\\/]+", "/", trimmed)
+            return value
+        if isinstance(value, list):
+            return [_convert(item) for item in value]
+        if isinstance(value, dict):
+            return {key: _convert(item) for key, item in value.items()}
+        return value
+
+    return _convert(obj)
+
+
+def _find_media_references(prompt_obj):
+    """Find media file references in image/video inputs used by worker prompts."""
+    media_refs = set()
+    for node in prompt_obj.values():
+        if not isinstance(node, dict):
+            continue
+        inputs = node.get("inputs", {})
+        for key in ("image", "video"):
+            value = inputs.get(key)
+            if isinstance(value, str):
+                cleaned = re.sub(r"\s*\[\w+\]$", "", value).strip().replace("\\", "/")
+                if IMAGE_OR_VIDEO_RE.search(cleaned):
+                    media_refs.add(cleaned)
+    return sorted(media_refs)
+
+
+def _load_media_file_sync(filename):
+    """Load local media bytes and hash for worker upload sync."""
+    import folder_paths
+
+    full_path = folder_paths.get_annotated_filepath(filename)
+    if not os.path.exists(full_path):
+        raise FileNotFoundError(filename)
+
+    with open(full_path, "rb") as f:
+        file_bytes = f.read()
+
+    file_hash = hashlib.md5(file_bytes).hexdigest()
+    mime_type = mimetypes.guess_type(full_path)[0]
+    if not mime_type:
+        ext = os.path.splitext(full_path)[1].lower()
+        if ext in {".mp4", ".avi", ".mov", ".mkv", ".webm"}:
+            mime_type = "video/mp4"
+        else:
+            mime_type = "image/png"
+    return file_bytes, file_hash, mime_type
+
+
+async def fetch_worker_path_separator(worker, trace_execution_id=None):
+    """Best-effort fetch of a worker's path separator from /distributed/system_info."""
+    url = build_worker_url(worker, "/distributed/system_info")
+    session = await get_client_session()
+    try:
+        async with session.get(url, timeout=aiohttp.ClientTimeout(total=5)) as resp:
+            if resp.status != 200:
+                return None
+            payload = await resp.json()
+            separator = ((payload or {}).get("platform") or {}).get("path_separator")
+            return separator if separator in ("/", "\\") else None
+    except Exception as exc:
+        if trace_execution_id:
+            _trace_debug(trace_execution_id, f"Failed to fetch worker system info ({worker.get('id')}): {exc}")
+        else:
+            debug_log(f"[Distributed] Failed to fetch worker system info ({worker.get('id')}): {exc}")
+        return None
+
+
+async def _upload_media_to_worker(worker, filename, file_bytes, file_hash, mime_type, trace_execution_id=None):
+    """Upload one media file to worker iff missing or hash-mismatched."""
+    session = await get_client_session()
+    normalized = filename.replace("\\", "/")
+
+    check_url = build_worker_url(worker, "/distributed/check_file")
+    try:
+        async with session.post(
+            check_url,
+            json={"filename": normalized, "hash": file_hash},
+            timeout=aiohttp.ClientTimeout(total=6),
+        ) as resp:
+            if resp.status == 200:
+                payload = await resp.json()
+                if payload.get("exists") and payload.get("hash_matches"):
+                    return False
+    except Exception as exc:
+        if trace_execution_id:
+            _trace_debug(trace_execution_id, f"Media check failed for '{normalized}' on worker {worker.get('id')}: {exc}")
+        else:
+            debug_log(f"[Distributed] Media check failed for '{normalized}' on worker {worker.get('id')}: {exc}")
+
+    parts = normalized.split("/")
+    clean_name = parts[-1]
+    subfolder = "/".join(parts[:-1])
+
+    form = aiohttp.FormData()
+    form.add_field("image", file_bytes, filename=clean_name, content_type=mime_type)
+    form.add_field("type", "input")
+    form.add_field("subfolder", subfolder)
+    form.add_field("overwrite", "true")
+
+    upload_url = build_worker_url(worker, "/upload/image")
+    async with session.post(
+        upload_url,
+        data=form,
+        timeout=aiohttp.ClientTimeout(total=30),
+    ) as resp:
+        resp.raise_for_status()
+    return True
+
+
+async def sync_worker_media(worker, prompt_obj, trace_execution_id=None):
+    """Sync referenced media files from master to a remote worker before dispatch."""
+    media_refs = _find_media_references(prompt_obj)
+    if not media_refs:
+        return
+
+    loop = asyncio.get_running_loop()
+    uploaded = 0
+    skipped = 0
+    missing = 0
+    for filename in media_refs:
+        try:
+            file_bytes, file_hash, mime_type = await loop.run_in_executor(
+                None, _load_media_file_sync, filename
+            )
+        except FileNotFoundError:
+            missing += 1
+            if trace_execution_id:
+                _trace_info(trace_execution_id, f"Media file '{filename}' not found on master; worker may fail to load it.")
+            else:
+                log(f"[Distributed] Media file '{filename}' not found on master; worker may fail to load it.")
+            continue
+        except Exception as exc:
+            if trace_execution_id:
+                _trace_info(trace_execution_id, f"Failed to load media '{filename}' for worker sync: {exc}")
+            else:
+                log(f"[Distributed] Failed to load media '{filename}' for worker sync: {exc}")
+            continue
+
+        try:
+            changed = await _upload_media_to_worker(
+                worker,
+                filename,
+                file_bytes,
+                file_hash,
+                mime_type,
+                trace_execution_id=trace_execution_id,
+            )
+            if changed:
+                uploaded += 1
+            else:
+                skipped += 1
+        except Exception as exc:
+            if trace_execution_id:
+                _trace_info(trace_execution_id, f"Failed to upload media '{filename}' to worker {worker.get('id')}: {exc}")
+            else:
+                log(f"[Distributed] Failed to upload media '{filename}' to worker {worker.get('id')}: {exc}")
+
+    summary = (
+        f"Media sync for worker {worker.get('id')}: "
+        f"uploaded={uploaded}, skipped={skipped}, missing={missing}, referenced={len(media_refs)}"
+    )
+    if trace_execution_id:
+        _trace_debug(trace_execution_id, summary)
+    else:
+        debug_log(f"[Distributed] {summary}")

api/orchestration/prompt_transform.py

@@ -0,0 +1,335 @@
+import json
+from collections import deque
+
+from ...utils.logging import debug_log
+
+
+class PromptIndex:
+    """Cache prompt metadata for faster worker/master prompt preparation."""
+
+    def __init__(self, prompt_obj):
+        self._prompt_json = json.dumps(prompt_obj)
+        self.nodes_by_class = {}
+        self.class_by_node = {}
+        self.inputs_by_node = {}
+        for node_id, node in _iter_prompt_nodes(prompt_obj):
+            class_type = node.get("class_type")
+            node_id_str = str(node_id)
+            if class_type:
+                self.nodes_by_class.setdefault(class_type, []).append(node_id_str)
+            self.class_by_node[node_id_str] = class_type
+            self.inputs_by_node[node_id_str] = node.get("inputs", {})
+        self._upstream_cache = {}
+
+    def copy_prompt(self):
+        return json.loads(self._prompt_json)
+
+    def nodes_for_class(self, class_name):
+        return self.nodes_by_class.get(class_name, [])
+
+    def has_upstream(self, start_node_id, target_class):
+        cache_key = (str(start_node_id), target_class)
+        if cache_key in self._upstream_cache:
+            return self._upstream_cache[cache_key]
+
+        visited = set()
+        stack = [str(start_node_id)]
+        while stack:
+            node_id = stack.pop()
+            if node_id in visited:
+                continue
+            visited.add(node_id)
+            inputs = self.inputs_by_node.get(node_id, {})
+            for value in inputs.values():
+                if isinstance(value, list) and len(value) == 2:
+                    upstream_id = str(value[0])
+                    if self.class_by_node.get(upstream_id) == target_class:
+                        self._upstream_cache[cache_key] = True
+                        return True
+                    if upstream_id in self.inputs_by_node:
+                        stack.append(upstream_id)
+
+        self._upstream_cache[cache_key] = False
+        return False
+
+
+def _iter_prompt_nodes(prompt_obj):
+    for node_id, node in prompt_obj.items():
+        if isinstance(node, dict):
+            yield str(node_id), node
+
+
+def find_nodes_by_class(prompt_obj, class_name):
+    nodes = []
+    for node_id, node in _iter_prompt_nodes(prompt_obj):
+        if node.get("class_type") == class_name:
+            nodes.append(node_id)
+    return nodes
+
+
+def _find_downstream_nodes(prompt_obj, start_ids):
+    """Return all nodes reachable downstream from the provided IDs."""
+    adjacency = {}
+    for node_id, node in _iter_prompt_nodes(prompt_obj):
+        inputs = node.get("inputs", {})
+        for value in inputs.values():
+            if isinstance(value, list) and len(value) == 2:
+                source_id = str(value[0])
+                adjacency.setdefault(source_id, set()).add(str(node_id))
+
+    connected = set(start_ids)
+    queue = deque(start_ids)
+    while queue:
+        current = queue.popleft()
+        for dependent in adjacency.get(current, ()):  # pragma: no branch - simple iteration
+            if dependent not in connected:
+                connected.add(dependent)
+                queue.append(dependent)
+    return connected
+
+
+def _create_numeric_id_generator(prompt_obj):
+    """Return a closure that yields new numeric string IDs."""
+    max_id = 0
+    for node_id in prompt_obj.keys():
+        try:
+            numeric = int(node_id)
+        except (TypeError, ValueError):
+            continue
+        max_id = max(max_id, numeric)
+
+    counter = max_id
+
+    def _next_id():
+        nonlocal counter
+        counter += 1
+        return str(counter)
+
+    return _next_id
+
+
+def _find_upstream_nodes(prompt_obj, start_ids):
+    """Return all nodes reachable upstream from start_ids, including start nodes."""
+    connected = set(str(node_id) for node_id in start_ids)
+    queue = deque(connected)
+    while queue:
+        node_id = queue.popleft()
+        node = prompt_obj.get(node_id) or {}
+        inputs = node.get("inputs", {})
+        for value in inputs.values():
+            if isinstance(value, list) and len(value) == 2:
+                source_id = str(value[0])
+                if source_id in prompt_obj and source_id not in connected:
+                    connected.add(source_id)
+                    queue.append(source_id)
+    return connected
+
+
+def prune_prompt_for_worker(prompt_obj):
+    """Prune worker prompt to distributed nodes and their upstream dependencies."""
+    collector_ids = find_nodes_by_class(prompt_obj, "DistributedCollector")
+    upscale_ids = find_nodes_by_class(prompt_obj, "UltimateSDUpscaleDistributed")
+    distributed_ids = collector_ids + upscale_ids
+    if not distributed_ids:
+        return prompt_obj
+
+    connected = _find_upstream_nodes(prompt_obj, distributed_ids)
+    pruned_prompt = {}
+    for node_id in connected:
+        node = prompt_obj.get(node_id)
+        if node is not None:
+            pruned_prompt[node_id] = json.loads(json.dumps(node))
+
+    next_id = _create_numeric_id_generator(pruned_prompt)
+    for dist_id in distributed_ids:
+        if dist_id not in pruned_prompt:
+            continue
+        downstream = _find_downstream_nodes(prompt_obj, [dist_id])
+        has_removed_downstream = any(node_id != dist_id for node_id in downstream)
+        if has_removed_downstream:
+            preview_id = next_id()
+            pruned_prompt[preview_id] = {
+                "inputs": {
+                    "images": [dist_id, 0],
+                },
+                "class_type": "PreviewImage",
+                "_meta": {
+                    "title": "Preview Image (auto-added)",
+                },
+            }
+
+    return pruned_prompt
+
+
+def prepare_delegate_master_prompt(prompt_obj, collector_ids):
+    """Prune master prompt so it only executes post-collector nodes in delegate mode."""
+    downstream = _find_downstream_nodes(prompt_obj, collector_ids)
+    nodes_to_keep = set(collector_ids)
+    nodes_to_keep.update(downstream)
+
+    pruned_prompt = {}
+    for node_id in nodes_to_keep:
+        node = prompt_obj.get(node_id)
+        if node is not None:
+            pruned_prompt[node_id] = json.loads(json.dumps(node))
+
+    pruned_ids = set(pruned_prompt.keys())
+    for node_id, node in pruned_prompt.items():
+        inputs = node.get("inputs")
+        if not inputs:
+            continue
+        for input_name, input_value in list(inputs.items()):
+            if isinstance(input_value, list) and len(input_value) == 2:
+                source_id = str(input_value[0])
+                if source_id not in pruned_ids:
+                    inputs.pop(input_name, None)
+                    debug_log(
+                        f"Removed upstream reference '{input_name}' from node {node_id} for delegate-only master prompt."
+                    )
+
+    next_id = _create_numeric_id_generator(pruned_prompt)
+    for collector_id in collector_ids:
+        collector_entry = pruned_prompt.get(collector_id)
+        if not collector_entry:
+            continue
+        placeholder_id = next_id()
+        pruned_prompt[placeholder_id] = {
+            "class_type": "DistributedEmptyImage",
+            "inputs": {
+                "height": 64,
+                "width": 64,
+                "channels": 3,
+            },
+            "_meta": {
+                "title": "Distributed Empty Image (auto-added)",
+            },
+        }
+        collector_entry.setdefault("inputs", {})["images"] = [placeholder_id, 0]
+        debug_log(
+            f"Inserted placeholder node {placeholder_id} for collector {collector_id} in delegate-only master prompt."
+        )
+
+    return pruned_prompt
+
+
+def generate_job_id_map(prompt_index, prefix):
+    """Create stable per-node job IDs for distributed nodes."""
+    job_map = {}
+    distributed_nodes = prompt_index.nodes_for_class("DistributedCollector") + prompt_index.nodes_for_class(
+        "UltimateSDUpscaleDistributed"
+    )
+    for node_id in distributed_nodes:
+        job_map[node_id] = f"{prefix}_{node_id}"
+    return job_map
+
+
+def _override_seed_nodes(prompt_copy, prompt_index, is_master, participant_id, worker_index_map):
+    """Configure DistributedSeed nodes for master or worker role."""
+    for node_id in prompt_index.nodes_for_class("DistributedSeed"):
+        node = prompt_copy.get(node_id)
+        if not isinstance(node, dict):
+            continue
+        inputs = node.setdefault("inputs", {})
+        inputs["is_worker"] = not is_master
+        if is_master:
+            inputs["worker_id"] = ""
+        else:
+            inputs["worker_id"] = f"worker_{worker_index_map.get(participant_id, 0)}"
+
+
+def _override_collector_nodes(
+    prompt_copy,
+    prompt_index,
+    is_master,
+    participant_id,
+    job_id_map,
+    master_url,
+    enabled_json,
+    delegate_master,
+):
+    """Configure DistributedCollector nodes for master or worker role."""
+    for node_id in prompt_index.nodes_for_class("DistributedCollector"):
+        node = prompt_copy.get(node_id)
+        if not isinstance(node, dict):
+            continue
+
+        if prompt_index.has_upstream(node_id, "UltimateSDUpscaleDistributed"):
+            node.setdefault("inputs", {})["pass_through"] = True
+            continue
+
+        inputs = node.setdefault("inputs", {})
+        inputs["multi_job_id"] = job_id_map.get(node_id, node_id)
+        inputs["is_worker"] = not is_master
+        inputs["enabled_worker_ids"] = enabled_json
+        if is_master:
+            inputs["delegate_only"] = bool(delegate_master)
+            inputs.pop("master_url", None)
+            inputs.pop("worker_id", None)
+        else:
+            inputs["master_url"] = master_url
+            inputs["worker_id"] = participant_id
+            inputs["delegate_only"] = False
+
+
+def _override_upscale_nodes(
+    prompt_copy,
+    prompt_index,
+    is_master,
+    participant_id,
+    job_id_map,
+    master_url,
+    enabled_json,
+):
+    """Configure UltimateSDUpscaleDistributed nodes for master or worker role."""
+    for node_id in prompt_index.nodes_for_class("UltimateSDUpscaleDistributed"):
+        node = prompt_copy.get(node_id)
+        if not isinstance(node, dict):
+            continue
+        inputs = node.setdefault("inputs", {})
+        inputs["multi_job_id"] = job_id_map.get(node_id, node_id)
+        inputs["is_worker"] = not is_master
+        inputs["enabled_worker_ids"] = enabled_json
+        if is_master:
+            inputs.pop("master_url", None)
+            inputs.pop("worker_id", None)
+        else:
+            inputs["master_url"] = master_url
+            inputs["worker_id"] = participant_id
+
+
+def apply_participant_overrides(
+    prompt_copy,
+    participant_id,
+    enabled_worker_ids,
+    job_id_map,
+    master_url,
+    delegate_master,
+    prompt_index,
+):
+    """Return a prompt copy with hidden inputs configured for master/worker."""
+    is_master = participant_id == "master"
+    worker_index_map = {wid: idx for idx, wid in enumerate(enabled_worker_ids)}
+    enabled_json = json.dumps(enabled_worker_ids)
+
+    _override_seed_nodes(prompt_copy, prompt_index, is_master, participant_id, worker_index_map)
+    _override_collector_nodes(
+        prompt_copy,
+        prompt_index,
+        is_master,
+        participant_id,
+        job_id_map,
+        master_url,
+        enabled_json,
+        delegate_master,
+    )
+    _override_upscale_nodes(
+        prompt_copy,
+        prompt_index,
+        is_master,
+        participant_id,
+        job_id_map,
+        master_url,
+        enabled_json,
+    )
+
+    return prompt_copy

api/queue_request.py

@@ -0,0 +1,79 @@
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional
+
+
+@dataclass(frozen=True)
+class QueueRequestPayload:
+    prompt: Dict[str, Any]
+    workflow_meta: Any
+    client_id: str
+    delegate_master: Optional[bool]
+    enabled_worker_ids: List[str]
+    auto_prepare: bool
+    trace_execution_id: Optional[str]
+
+
+def parse_queue_request_payload(data: Any) -> QueueRequestPayload:
+    """Parse and validate /distributed/queue payload into a normalized shape."""
+    if not isinstance(data, dict):
+        raise ValueError("Expected a JSON object body")
+
+    auto_prepare_raw = data.get("auto_prepare", True)
+    if not isinstance(auto_prepare_raw, bool):
+        raise ValueError("auto_prepare must be a boolean when provided")
+    auto_prepare = auto_prepare_raw
+
+    prompt = data.get("prompt")
+    # Auto-prepare is always on server-side; keep the field for wire compatibility.
+    if prompt is None:
+        workflow_payload = data.get("workflow")
+        if isinstance(workflow_payload, dict):
+            candidate_prompt = workflow_payload.get("prompt")
+            if isinstance(candidate_prompt, dict):
+                prompt = candidate_prompt
+
+    if not isinstance(prompt, dict):
+        raise ValueError("Field 'prompt' must be an object")
+
+    enabled_ids_raw = data.get("enabled_worker_ids")
+    workers_field = data.get("workers")
+    if enabled_ids_raw is None and workers_field is not None:
+        if not isinstance(workers_field, list):
+            raise ValueError("Field 'workers' must be a list when provided")
+        enabled_ids_raw = []
+        for entry in workers_field:
+            worker_id = entry.get("id") if isinstance(entry, dict) else entry
+            if worker_id is not None:
+                enabled_ids_raw.append(str(worker_id))
+
+    if enabled_ids_raw is None:
+        raise ValueError("enabled_worker_ids required")
+    else:
+        if not isinstance(enabled_ids_raw, list):
+            raise ValueError("enabled_worker_ids must be a list of worker IDs")
+        enabled_ids = [str(worker_id).strip() for worker_id in enabled_ids_raw if str(worker_id).strip()]
+
+    delegate_master = data.get("delegate_master")
+    if delegate_master is not None and not isinstance(delegate_master, bool):
+        raise ValueError("delegate_master must be a boolean when provided")
+
+    client_id = data.get("client_id")
+    if not isinstance(client_id, str) or not client_id.strip():
+        raise ValueError("client_id required")
+    client_id = client_id.strip()
+
+    trace_execution_id = data.get("trace_execution_id")
+    if trace_execution_id is not None:
+        if not isinstance(trace_execution_id, str):
+            raise ValueError("trace_execution_id must be a string when provided")
+        trace_execution_id = trace_execution_id.strip() or None
+
+    return QueueRequestPayload(
+        prompt=prompt,
+        workflow_meta=data.get("workflow"),
+        client_id=client_id,
+        delegate_master=delegate_master,
+        enabled_worker_ids=enabled_ids,
+        auto_prepare=auto_prepare,
+        trace_execution_id=trace_execution_id,
+    )

api/schemas.py

@@ -0,0 +1,36 @@
+def require_fields(data: dict, *fields) -> list[str]:
+    """Return field names that are missing or empty in a JSON object."""
+    if not isinstance(data, dict):
+        return list(fields)
+
+    missing = []
+    for field in fields:
+        if field not in data:
+            missing.append(field)
+            continue
+        value = data.get(field)
+        if value is None:
+            missing.append(field)
+            continue
+        if isinstance(value, str) and not value.strip():
+            missing.append(field)
+
+    return missing
+
+
+def validate_worker_id(worker_id: str, config: dict) -> bool:
+    """Return True when worker_id exists in config['workers']."""
+    worker_id_str = str(worker_id)
+    workers = (config or {}).get("workers", [])
+    return any(str(worker.get("id")) == worker_id_str for worker in workers)
+
+
+def validate_positive_int(value, field_name: str) -> str | None:
+    """Validate positive integers and return an error string when invalid."""
+    try:
+        parsed = int(value)
+    except (TypeError, ValueError):
+        return f"Field '{field_name}' must be a positive integer."
+    if parsed <= 0:
+        return f"Field '{field_name}' must be a positive integer."
+    return None

api/worker_routes.py

@@ -13,9 +13,10 @@ import server
 
 from ..utils.config import load_config
 from ..utils.logging import debug_log, log
-from ..utils.network import handle_api_error, normalize_host, get_client_session
+from ..utils.network import build_worker_url, handle_api_error, normalize_host, probe_worker
 from ..utils.constants import CHUNK_SIZE
 from ..workers import get_worker_manager
+from .schemas import require_fields, validate_worker_id
 from ..workers.detection import (
     is_local_worker,
     is_same_physical_host,
@@ -96,10 +97,14 @@ async def clear_launching_state(request):
     try:
         wm = get_worker_manager()
         data = await request.json()
-        worker_id = str(data.get('worker_id'))
-        
-        if not worker_id:
-            return await handle_api_error(request, "worker_id is required", 400)
+        missing = require_fields(data, "worker_id")
+        if missing:
+            return await handle_api_error(request, f"Missing required field(s): {', '.join(missing)}", 400)
+
+        worker_id = str(data.get("worker_id")).strip()
+        config = load_config()
+        if not validate_worker_id(worker_id, config):
+            return await handle_api_error(request, f"Worker {worker_id} not found", 404)
         
         # Clear launching flag in managed processes
         if worker_id in wm.processes:
@@ -353,19 +358,22 @@ async def launch_worker_endpoint(request):
     try:
         wm = get_worker_manager()
         data = await request.json()
-        worker_id = data.get("worker_id")
-        
-        if not worker_id:
-            return await handle_api_error(request, "Missing worker_id", 400)
+        missing = require_fields(data, "worker_id")
+        if missing:
+            return await handle_api_error(request, f"Missing required field(s): {', '.join(missing)}", 400)
+
+        worker_id = str(data.get("worker_id")).strip()
         
         # Find worker config
         config = load_config()
-        worker = next((w for w in config.get("workers", []) if w["id"] == worker_id), None)
+        if not validate_worker_id(worker_id, config):
+            return await handle_api_error(request, f"Worker {worker_id} not found", 404)
+        worker = next((w for w in config.get("workers", []) if str(w.get("id")) == worker_id), None)
         if not worker:
             return await handle_api_error(request, f"Worker {worker_id} not found", 404)
         
         # Ensure consistent string ID
-        worker_id_str = str(worker_id)
+        worker_id_str = worker_id
         
         # Check if already running (managed by this instance)
         if worker_id_str in wm.processes:
@@ -416,11 +424,15 @@ async def stop_worker_endpoint(request):
     try:
         wm = get_worker_manager()
         data = await request.json()
-        worker_id = data.get("worker_id")
-        
-        if not worker_id:
-            return await handle_api_error(request, "Missing worker_id", 400)
-        
+        missing = require_fields(data, "worker_id")
+        if missing:
+            return await handle_api_error(request, f"Missing required field(s): {', '.join(missing)}", 400)
+
+        worker_id = str(data.get("worker_id")).strip()
+        config = load_config()
+        if not validate_worker_id(worker_id, config):
+            return await handle_api_error(request, f"Worker {worker_id} not found", 404)
+
         success, message = wm.stop_worker(worker_id)
         
         if success:
@@ -472,28 +484,24 @@ async def get_local_worker_status_endpoint(request):
                 
                 # Try to connect to worker
                 try:
-                    session = await get_client_session()
-                    async with session.get(
-                        f"http://localhost:{port}/prompt",
-                        timeout=aiohttp.ClientTimeout(total=2)
-                    ) as resp:
-                        if resp.status == 200:
-                            data = await resp.json()
-                            queue_remaining = data.get("exec_info", {}).get("queue_remaining", 0)
-                            worker_statuses[worker_id] = {
-                                "online": True,
-                                "enabled": True,
-                                "processing": queue_remaining > 0,
-                                "queue_count": queue_remaining
-                            }
-                        else:
-                            worker_statuses[worker_id] = {
-                                "online": False,
-                                "enabled": True,
-                                "processing": False,
-                                "queue_count": 0,
-                                "error": f"HTTP {resp.status}"
-                            }
+                    worker_url = build_worker_url(worker)
+                    data = await probe_worker(worker_url, timeout=2.0)
+                    if data is None:
+                        worker_statuses[worker_id] = {
+                            "online": False,
+                            "enabled": True,
+                            "processing": False,
+                            "queue_count": 0,
+                            "error": "Unavailable",
+                        }
+                        continue
+                    queue_remaining = data.get("exec_info", {}).get("queue_remaining", 0)
+                    worker_statuses[worker_id] = {
+                        "online": True,
+                        "enabled": True,
+                        "processing": queue_remaining > 0,
+                        "queue_count": queue_remaining
+                    }
                 except asyncio.TimeoutError:
                     worker_statuses[worker_id] = {
                         "online": False,

docs/comfyui-distributed-api.md

@@ -48,34 +48,38 @@ Queue a workflow for distributed execution.
 {
   "prompt": { "<node_id>": { "class_type": "...", "inputs": { } } },
   "workflow": { },
-  "client_id": "optional",
+  "client_id": "external-client",
   "delegate_master": false,
   "enabled_worker_ids": ["1", "2"],
   "workers": ["1", "2"],
-  "auto_prepare": true
+  "auto_prepare": true,
+  "trace_execution_id": "exec_1700000000_ab12cd"
 }
 ```
 
 #### Fields
 
-- `prompt` (required, object)
+- `prompt` (required unless `workflow.prompt` is present, object)
   - The ComfyUI prompt/workflow graph, same shape as used by `POST /prompt`.
 - `workflow` (optional, object)
   - Workflow metadata that ComfyUI normally stores in `extra_pnginfo.workflow`.
   - If you don’t care about UI metadata, you can omit it.
-- `client_id` (optional, string)
+- `client_id` (required, string)
   - Passed through as `extra_data.client_id` (useful if you consume ComfyUI websocket events).
 - `delegate_master` (optional, boolean)
   - If `true`, attempts “workers-only” execution for workflows based on `DistributedCollector`.
   - Current limitation: delegate-only mode **does not support** `UltimateSDUpscaleDistributed` and will fall back to running the full prompt on master.
-- `enabled_worker_ids` (optional, array of strings)
-  - If provided, only these worker IDs will be considered.
-  - If omitted, the plugin uses workers marked as enabled in the UI config.
+- `enabled_worker_ids` (required, array of strings)
+  - The explicit worker IDs to consider for this run.
 - `workers` (optional, array of strings or objects with `id`)
-  - Transitional alias for `enabled_worker_ids` used by the frontend auto-prepare path.
+  - Transitional alias for `enabled_worker_ids` used by older clients.
 - `auto_prepare` (optional, boolean)
-  - Feature flag used by newer clients. Current backend still requires `prompt`, but can also read
-    `workflow.prompt` when this flag is true.
+  - Kept for wire compatibility.
+  - Backend orchestration always runs with auto-prepare semantics.
+  - If top-level `prompt` is omitted, backend will attempt `workflow.prompt`.
+- `trace_execution_id` (optional, string)
+  - Passed through to orchestration logs.
+  - Server log lines include the marker as `[exec:<trace_execution_id>]`.
 
 ##### How to get `enabled_worker_ids`
 
@@ -184,7 +188,9 @@ Where `payload.json` contains at least:
 {
   "prompt": {
     "1": {"class_type": "KSampler", "inputs": {} }
-  }
+  },
+  "enabled_worker_ids": [],
+  "client_id": "external-client"
 }
 ```
 
@@ -197,6 +203,7 @@ url = "http://127.0.0.1:8188/distributed/queue"
 payload = {
     "prompt": {...},
     "workflow": {...},
+    "client_id": "external-client",
     "delegate_master": False,
     "enabled_worker_ids": ["1", "2"],
 }
@@ -214,6 +221,7 @@ const url = "http://127.0.0.1:8188/distributed/queue";
 const payload = {
   prompt: {/* ... */},
   workflow: {/* ... */},
+  client_id: "external-client",
   delegate_master: false,
   enabled_worker_ids: ["1", "2"],
 };

nodes/collector.py

@@ -13,8 +13,9 @@ from comfy.utils import ProgressBar
 
 from ..utils.logging import debug_log, log
 from ..utils.config import get_worker_timeout_seconds, load_config, is_master_delegate_only
+from ..utils.constants import HEARTBEAT_INTERVAL
 from ..utils.image import tensor_to_pil, pil_to_tensor, ensure_contiguous
-from ..utils.network import get_client_session, normalize_host
+from ..utils.network import build_worker_url, get_client_session, probe_worker
 from ..utils.async_helpers import run_async_in_server_loop
 from ..workers.detection import is_local_worker, get_comms_channel
 
@@ -196,6 +197,15 @@ class DistributedCollectorNode:
             if num_workers == 0:
                 return (images, audio if audio is not None else self.EMPTY_AUDIO)
 
+            # Create the queue before any expensive local work to avoid job_complete race.
+            async with prompt_server.distributed_jobs_lock:
+                if multi_job_id not in prompt_server.distributed_pending_jobs:
+                    prompt_server.distributed_pending_jobs[multi_job_id] = asyncio.Queue()
+                    debug_log(f"Master - Initialized queue early for job {multi_job_id}")
+                else:
+                    existing_size = prompt_server.distributed_pending_jobs[multi_job_id].qsize()
+                    debug_log(f"Master - Using existing queue for job {multi_job_id} (current size: {existing_size})")
+
             if delegate_mode:
                 master_batch_size = 0
                 images_on_cpu = None
@@ -215,22 +225,13 @@ class DistributedCollectorNode:
             worker_images = {}  # Dict to store images by worker_id and index
             worker_audio = {}   # Dict to store audio by worker_id
             
-            # Get the existing queue - it should already exist from prepare_job
-            async with prompt_server.distributed_jobs_lock:
-                if multi_job_id not in prompt_server.distributed_pending_jobs:
-                    log(f"Master - WARNING: Queue doesn't exist for job {multi_job_id}, creating one")
-                    prompt_server.distributed_pending_jobs[multi_job_id] = asyncio.Queue()
-                else:
-                    existing_size = prompt_server.distributed_pending_jobs[multi_job_id].qsize()
-                    debug_log(f"Master - Using existing queue for job {multi_job_id} (current size: {existing_size})")
-            
             # Collect images until all workers report they're done
             collected_count = 0
             workers_done = set()
             
             # Use unified worker timeout from config/UI with simple sliced waits
             base_timeout = float(get_worker_timeout_seconds())
-            slice_timeout = min(0.5, base_timeout)  # small per-wait slice to recheck interrupt
+            slice_timeout = min(max(0.1, HEARTBEAT_INTERVAL / 20.0), base_timeout)
             last_activity = time.time()
             
             
@@ -284,37 +285,44 @@ class DistributedCollectorNode:
                         # Re-check for user interruption after timeout expiry
                         comfy.model_management.throw_exception_if_processing_interrupted()
                         missing_workers = set(str(w) for w in enabled_workers) - workers_done
-                        log(f"Master - Timeout. Still waiting for workers: {list(missing_workers)}")
+                        elapsed = time.time() - last_activity
+                        for missing_worker_id in sorted(missing_workers):
+                            log(
+                                "Master - Heartbeat timeout: "
+                                f"worker={missing_worker_id}, elapsed={elapsed:.1f}s"
+                            )
+                        log(
+                            f"Master - Heartbeat timeout. Still waiting for workers: {list(missing_workers)} "
+                            f"(elapsed={elapsed:.1f}s)"
+                        )
 
                         # Probe missing workers' /prompt endpoints to check if they are actively processing
                         any_busy = False
                         try:
                             cfg = load_config()
                             cfg_workers = cfg.get('workers', [])
-                            session = await get_client_session()
                             for wid in list(missing_workers):
                                 wrec = next((w for w in cfg_workers if str(w.get('id')) == str(wid)), None)
                                 if not wrec:
                                     debug_log(f"Collector probe: worker {wid} not found in config")
                                     continue
-                                host = normalize_host(wrec.get('host') or 'localhost') or 'localhost'
-                                port = int(wrec.get('port', 8188))
-                                url = f"http://{host}:{port}/prompt"
+                                worker_url = build_worker_url(wrec)
                                 try:
-                                    async with session.get(url, timeout=aiohttp.ClientTimeout(total=2.0)) as resp:
-                                        status = resp.status
-                                        q = None
-                                        if status == 200:
-                                            try:
-                                                payload = await resp.json()
-                                                q = int(payload.get('exec_info', {}).get('queue_remaining', 0))
-                                            except Exception:
-                                                q = 0
-                                        debug_log(f"Collector probe: worker {wid} status={status} queue_remaining={q}")
-                                        if status == 200 and q and q > 0:
-                                            any_busy = True
-                                            log(f"Master - Probe grace: worker {wid} appears busy (queue_remaining={q}). Continuing to wait.")
-                                            break
+                                    payload = await probe_worker(worker_url, timeout=2.0)
+                                    queue_remaining = None
+                                    if payload is not None:
+                                        queue_remaining = int(payload.get('exec_info', {}).get('queue_remaining', 0))
+                                    debug_log(
+                                        "Collector probe: worker "
+                                        f"{wid} online={payload is not None} queue_remaining={queue_remaining}"
+                                    )
+                                    if payload is not None and queue_remaining and queue_remaining > 0:
+                                        any_busy = True
+                                        log(
+                                            f"Master - Probe grace: worker {wid} appears busy "
+                                            f"(queue_remaining={queue_remaining}). Continuing to wait."
+                                        )
+                                        break
                                 except Exception as e:
                                     debug_log(f"Collector probe failed for worker {wid}: {e}")
                         except Exception as e:

tests/api/__init__.py

@@ -0,0 +1 @@
+

tests/api/test_config_routes.py

@@ -0,0 +1,148 @@
+import copy
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+from unittest.mock import patch
+
+
+class _FakeResponse:
+    def __init__(self, payload, status=200):
+        self.payload = payload
+        self.status = status
+
+
+class _FakeRequest:
+    def __init__(self, payload=None):
+        self._payload = payload
+
+    async def json(self):
+        return self._payload
+
+
+def _load_config_routes_module():
+    module_path = Path(__file__).resolve().parents[2] / "api" / "config_routes.py"
+    package_name = "dist_api_config_testpkg"
+
+    for mod_name in list(sys.modules):
+        if mod_name == package_name or mod_name.startswith(f"{package_name}."):
+            del sys.modules[mod_name]
+
+    root_pkg = types.ModuleType(package_name)
+    root_pkg.__path__ = []
+    sys.modules[package_name] = root_pkg
+
+    api_pkg = types.ModuleType(f"{package_name}.api")
+    api_pkg.__path__ = []
+    sys.modules[f"{package_name}.api"] = api_pkg
+
+    utils_pkg = types.ModuleType(f"{package_name}.utils")
+    utils_pkg.__path__ = []
+    sys.modules[f"{package_name}.utils"] = utils_pkg
+
+    created_aiohttp_stub = False
+    if "aiohttp" not in sys.modules:
+        created_aiohttp_stub = True
+        aiohttp_module = types.ModuleType("aiohttp")
+        aiohttp_module.web = types.SimpleNamespace(
+            json_response=lambda payload, status=200: _FakeResponse(payload, status=status)
+        )
+        sys.modules["aiohttp"] = aiohttp_module
+
+    class _Routes:
+        def get(self, _path):
+            def _decorator(fn):
+                return fn
+            return _decorator
+
+        def post(self, _path):
+            def _decorator(fn):
+                return fn
+            return _decorator
+
+    server_module = types.ModuleType("server")
+    server_module.PromptServer = types.SimpleNamespace(instance=types.SimpleNamespace(routes=_Routes()))
+    sys.modules["server"] = server_module
+
+    logging_module = types.ModuleType(f"{package_name}.utils.logging")
+    logging_module.debug_log = lambda *_args, **_kwargs: None
+    logging_module.log = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.utils.logging"] = logging_module
+
+    network_module = types.ModuleType(f"{package_name}.utils.network")
+
+    async def _handle_api_error(_request, error, status=500):
+        return _FakeResponse({"status": "error", "message": str(error)}, status=status)
+
+    network_module.handle_api_error = _handle_api_error
+    network_module.normalize_host = lambda value: value
+    sys.modules[f"{package_name}.utils.network"] = network_module
+
+    default_config = {
+        "workers": [],
+        "master": {"host": ""},
+        "settings": {"debug": False},
+        "tunnel": {},
+    }
+
+    config_module = types.ModuleType(f"{package_name}.utils.config")
+    config_module.load_config = lambda: copy.deepcopy(default_config)
+    config_module.save_config = lambda _cfg: True
+    sys.modules[f"{package_name}.utils.config"] = config_module
+
+    spec = importlib.util.spec_from_file_location(f"{package_name}.api.config_routes", module_path)
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    spec.loader.exec_module(module)
+
+    if created_aiohttp_stub:
+        sys.modules.pop("aiohttp", None)
+
+    return module
+
+
+config_routes = _load_config_routes_module()
+
+
+class ConfigRoutesTests(unittest.IsolatedAsyncioTestCase):
+    async def test_get_config_returns_core_sections(self):
+        cfg = {"workers": [], "master": {}, "settings": {}, "tunnel": {}}
+        with patch.object(config_routes, "load_config", return_value=cfg):
+            response = await config_routes.get_config_endpoint(_FakeRequest())
+
+        self.assertEqual(response.status, 200)
+        self.assertIn("workers", response.payload)
+        self.assertIn("master", response.payload)
+        self.assertIn("settings", response.payload)
+
+    async def test_update_config_valid_field_persists(self):
+        cfg = {"workers": [], "master": {}, "settings": {"debug": False}, "tunnel": {}}
+        with patch.object(config_routes, "load_config", return_value=cfg), patch.object(
+            config_routes, "save_config", return_value=True
+        ):
+            response = await config_routes.update_config_endpoint(_FakeRequest({"debug": True}))
+
+        self.assertEqual(response.status, 200)
+        self.assertEqual(response.payload["status"], "success")
+        self.assertTrue(response.payload["config"]["settings"]["debug"])
+
+    async def test_update_config_unknown_field_returns_400(self):
+        cfg = {"workers": [], "master": {}, "settings": {"debug": False}, "tunnel": {}}
+        with patch.object(config_routes, "load_config", return_value=cfg):
+            response = await config_routes.update_config_endpoint(_FakeRequest({"unknown_field": 1}))
+
+        self.assertEqual(response.status, 400)
+        self.assertIn("unknown_field", " ".join(response.payload.get("error", [])).lower())
+
+    async def test_update_config_wrong_type_returns_400(self):
+        cfg = {"workers": [], "master": {}, "settings": {"debug": False}, "tunnel": {}}
+        with patch.object(config_routes, "load_config", return_value=cfg):
+            response = await config_routes.update_config_endpoint(_FakeRequest({"debug": "true"}))
+
+        self.assertEqual(response.status, 400)
+        self.assertIn("debug", " ".join(response.payload.get("error", [])).lower())
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/api/test_distributed_queue.py

@@ -0,0 +1,220 @@
+import importlib.util
+import sys
+import types
+import unittest
+from dataclasses import dataclass
+from pathlib import Path
+from unittest.mock import AsyncMock, patch
+
+
+class _FakeResponse:
+    def __init__(self, payload, status=200):
+        self.payload = payload
+        self.status = status
+
+
+class _FakeRequest:
+    def __init__(self, payload):
+        self._payload = payload
+
+    async def json(self):
+        return self._payload
+
+
+def _load_job_routes_module():
+    module_path = Path(__file__).resolve().parents[2] / "api" / "job_routes.py"
+    package_name = "dist_api_queue_testpkg"
+
+    # Reset package namespace to avoid stale module state across test runs.
+    for mod_name in list(sys.modules):
+        if mod_name == package_name or mod_name.startswith(f"{package_name}."):
+            del sys.modules[mod_name]
+
+    root_pkg = types.ModuleType(package_name)
+    root_pkg.__path__ = []
+    sys.modules[package_name] = root_pkg
+
+    api_pkg = types.ModuleType(f"{package_name}.api")
+    api_pkg.__path__ = []
+    sys.modules[f"{package_name}.api"] = api_pkg
+
+    utils_pkg = types.ModuleType(f"{package_name}.utils")
+    utils_pkg.__path__ = []
+    sys.modules[f"{package_name}.utils"] = utils_pkg
+
+    # aiohttp.web stub
+    created_aiohttp_stub = False
+    if "aiohttp" not in sys.modules:
+        created_aiohttp_stub = True
+        aiohttp_module = types.ModuleType("aiohttp")
+        aiohttp_module.web = types.SimpleNamespace(
+            json_response=lambda payload, status=200: _FakeResponse(payload, status=status)
+        )
+        sys.modules["aiohttp"] = aiohttp_module
+
+    # server module stub with route decorators
+    class _Routes:
+        def get(self, _path):
+            def _decorator(fn):
+                return fn
+            return _decorator
+
+        def post(self, _path):
+            def _decorator(fn):
+                return fn
+            return _decorator
+
+    prompt_server_instance = types.SimpleNamespace(
+        routes=_Routes(),
+        distributed_jobs_lock=None,
+        distributed_pending_jobs={},
+    )
+    server_module = types.ModuleType("server")
+    server_module.PromptServer = types.SimpleNamespace(instance=prompt_server_instance)
+    sys.modules["server"] = server_module
+
+    # torch stub (only needed to satisfy import)
+    if "torch" not in sys.modules:
+        torch_module = types.ModuleType("torch")
+        torch_module.cuda = types.SimpleNamespace(
+            is_available=lambda: False,
+            empty_cache=lambda: None,
+            ipc_collect=lambda: None,
+        )
+        sys.modules["torch"] = torch_module
+
+    # PIL stub (only needed to satisfy import)
+    if "PIL" not in sys.modules:
+        pil_module = types.ModuleType("PIL")
+        image_module = types.ModuleType("PIL.Image")
+        pil_module.Image = image_module
+        sys.modules["PIL"] = pil_module
+        sys.modules["PIL.Image"] = image_module
+
+    logging_module = types.ModuleType(f"{package_name}.utils.logging")
+    logging_module.debug_log = lambda *_args, **_kwargs: None
+    logging_module.log = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.utils.logging"] = logging_module
+
+    image_module = types.ModuleType(f"{package_name}.utils.image")
+    image_module.pil_to_tensor = lambda *_args, **_kwargs: None
+    image_module.ensure_contiguous = lambda tensor: tensor
+    sys.modules[f"{package_name}.utils.image"] = image_module
+
+    network_module = types.ModuleType(f"{package_name}.utils.network")
+
+    async def _handle_api_error(_request, error, status=500):
+        return _FakeResponse({"status": "error", "message": str(error)}, status=status)
+
+    network_module.handle_api_error = _handle_api_error
+    sys.modules[f"{package_name}.utils.network"] = network_module
+
+    constants_module = types.ModuleType(f"{package_name}.utils.constants")
+    constants_module.MEMORY_CLEAR_DELAY = 0.0
+    constants_module.JOB_INIT_GRACE_PERIOD = 10.0
+    sys.modules[f"{package_name}.utils.constants"] = constants_module
+
+    async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
+    async_helpers_module.queue_prompt_payload = AsyncMock(return_value="prompt_local")
+    sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
+
+    queue_orchestration_module = types.ModuleType(f"{package_name}.api.queue_orchestration")
+    queue_orchestration_module.orchestrate_distributed_execution = AsyncMock(return_value=("prompt_dist", 1))
+    sys.modules[f"{package_name}.api.queue_orchestration"] = queue_orchestration_module
+
+    @dataclass(frozen=True)
+    class _QueuePayload:
+        prompt: dict
+        workflow_meta: object
+        client_id: str
+        delegate_master: object
+        enabled_worker_ids: list
+        auto_prepare: bool
+        trace_execution_id: object
+
+    def _parse_queue_request_payload(data):
+        if not isinstance(data, dict):
+            raise ValueError("Expected a JSON object body")
+        prompt = data.get("prompt")
+        if not isinstance(prompt, dict):
+            raise ValueError("Field 'prompt' must be an object")
+        enabled = data.get("enabled_worker_ids")
+        if not isinstance(enabled, list):
+            raise ValueError("enabled_worker_ids required")
+        client_id = data.get("client_id")
+        if not isinstance(client_id, str) or not client_id.strip():
+            raise ValueError("client_id required")
+        return _QueuePayload(
+            prompt=prompt,
+            workflow_meta=data.get("workflow"),
+            client_id=client_id,
+            delegate_master=data.get("delegate_master"),
+            enabled_worker_ids=enabled,
+            auto_prepare=bool(data.get("auto_prepare", True)),
+            trace_execution_id=data.get("trace_execution_id"),
+        )
+
+    queue_request_module = types.ModuleType(f"{package_name}.api.queue_request")
+    queue_request_module.parse_queue_request_payload = _parse_queue_request_payload
+    sys.modules[f"{package_name}.api.queue_request"] = queue_request_module
+
+    spec = importlib.util.spec_from_file_location(f"{package_name}.api.job_routes", module_path)
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    spec.loader.exec_module(module)
+
+    if created_aiohttp_stub:
+        sys.modules.pop("aiohttp", None)
+
+    return module
+
+
+job_routes = _load_job_routes_module()
+
+
+class DistributedQueueEndpointTests(unittest.IsolatedAsyncioTestCase):
+    async def test_distributed_queue_happy_path_returns_prompt_id(self):
+        request = _FakeRequest(
+            {
+                "prompt": {"1": {"class_type": "Node"}},
+                "enabled_worker_ids": ["w1"],
+                "client_id": "client-1",
+                "auto_prepare": True,
+            }
+        )
+        with patch.object(
+            job_routes,
+            "orchestrate_distributed_execution",
+            new=AsyncMock(return_value=("prompt_123", 2)),
+        ):
+            response = await job_routes.distributed_queue_endpoint(request)
+
+        self.assertEqual(response.status, 200)
+        self.assertEqual(response.payload.get("prompt_id"), "prompt_123")
+        self.assertTrue(response.payload.get("auto_prepare_supported"))
+
+    async def test_distributed_queue_missing_prompt_returns_400(self):
+        request = _FakeRequest(
+            {
+                "enabled_worker_ids": ["w1"],
+                "client_id": "client-1",
+            }
+        )
+        response = await job_routes.distributed_queue_endpoint(request)
+        self.assertEqual(response.status, 400)
+        self.assertIn("prompt", response.payload.get("message", "").lower())
+
+    async def test_distributed_queue_missing_enabled_worker_ids_returns_400(self):
+        request = _FakeRequest(
+            {
+                "prompt": {"1": {"class_type": "Node"}},
+                "client_id": "client-1",
+            }
+        )
+        response = await job_routes.distributed_queue_endpoint(request)
+        self.assertEqual(response.status, 400)
+        self.assertIn("enabled_worker_ids", response.payload.get("message", "").lower())
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/api/test_worker_routes.py

@@ -0,0 +1,288 @@
+import importlib.util
+import os
+import sys
+import tempfile
+import types
+import unittest
+from pathlib import Path
+from unittest.mock import patch
+
+
+class _FakeResponse:
+    def __init__(self, payload, status=200):
+        self.payload = payload
+        self.status = status
+
+
+class _FakeRequest:
+    def __init__(self, payload=None, match_info=None, query=None):
+        self._payload = payload
+        self.match_info = match_info or {}
+        self.query = query or {}
+
+    async def json(self):
+        return self._payload
+
+
+class _DummyWorkerManager:
+    def __init__(self):
+        self.processes = {}
+
+    def launch_worker(self, worker):
+        worker_id = str(worker["id"])
+        self.processes[worker_id] = {
+            "pid": 12345,
+            "log_file": f"/tmp/distributed_worker_{worker_id}.log",
+            "process": None,
+        }
+        return 12345
+
+    def _is_process_running(self, _pid):
+        return False
+
+    def save_processes(self):
+        return None
+
+    def stop_worker(self, _worker_id):
+        return True, "Stopped"
+
+    def get_managed_workers(self):
+        return []
+
+
+class _ImmediateLoop:
+    async def run_in_executor(self, _executor, func, *args):
+        return func(*args)
+
+
+def _load_worker_routes_module():
+    module_path = Path(__file__).resolve().parents[2] / "api" / "worker_routes.py"
+    package_name = "dist_api_worker_testpkg"
+
+    for mod_name in list(sys.modules):
+        if mod_name == package_name or mod_name.startswith(f"{package_name}."):
+            del sys.modules[mod_name]
+
+    root_pkg = types.ModuleType(package_name)
+    root_pkg.__path__ = []
+    sys.modules[package_name] = root_pkg
+
+    api_pkg = types.ModuleType(f"{package_name}.api")
+    api_pkg.__path__ = []
+    sys.modules[f"{package_name}.api"] = api_pkg
+
+    utils_pkg = types.ModuleType(f"{package_name}.utils")
+    utils_pkg.__path__ = []
+    sys.modules[f"{package_name}.utils"] = utils_pkg
+
+    workers_pkg = types.ModuleType(f"{package_name}.workers")
+    workers_pkg.__path__ = []
+    workers_pkg.get_worker_manager = lambda: _DummyWorkerManager()
+    sys.modules[f"{package_name}.workers"] = workers_pkg
+
+    detection_module = types.ModuleType(f"{package_name}.workers.detection")
+    detection_module.is_local_worker = lambda *_args, **_kwargs: True
+    detection_module.is_same_physical_host = lambda *_args, **_kwargs: True
+    detection_module.get_machine_id = lambda: "machine-id"
+    detection_module.is_docker_environment = lambda: False
+    detection_module.is_runpod_environment = lambda: False
+    detection_module.get_comms_channel = lambda *_args, **_kwargs: "lan"
+    sys.modules[f"{package_name}.workers.detection"] = detection_module
+
+    created_aiohttp_stub = False
+    if "aiohttp" not in sys.modules:
+        created_aiohttp_stub = True
+        aiohttp_module = types.ModuleType("aiohttp")
+
+        class _ClientTimeout:
+            def __init__(self, total=None):
+                self.total = total
+
+        class _WSMsgType:
+            TEXT = "TEXT"
+            ERROR = "ERROR"
+            CLOSED = "CLOSED"
+
+        class _WebSocketResponse:
+            def __init__(self, *args, **kwargs):
+                self.args = args
+                self.kwargs = kwargs
+
+            async def prepare(self, _request):
+                return None
+
+            async def send_json(self, _payload):
+                return None
+
+            def __aiter__(self):
+                async def _empty():
+                    if False:
+                        yield None
+                return _empty()
+
+        aiohttp_module.ClientTimeout = _ClientTimeout
+        aiohttp_module.WSMsgType = _WSMsgType
+        aiohttp_module.web = types.SimpleNamespace(
+            json_response=lambda payload, status=200: _FakeResponse(payload, status=status),
+            WebSocketResponse=_WebSocketResponse,
+        )
+        sys.modules["aiohttp"] = aiohttp_module
+
+    class _Routes:
+        def get(self, _path):
+            def _decorator(fn):
+                return fn
+            return _decorator
+
+        def post(self, _path):
+            def _decorator(fn):
+                return fn
+            return _decorator
+
+    server_module = types.ModuleType("server")
+    server_module.PromptServer = types.SimpleNamespace(instance=types.SimpleNamespace(routes=_Routes()))
+    sys.modules["server"] = server_module
+
+    if "torch" not in sys.modules:
+        torch_module = types.ModuleType("torch")
+        torch_module.cuda = types.SimpleNamespace(
+            is_available=lambda: False,
+            empty_cache=lambda: None,
+            ipc_collect=lambda: None,
+            current_device=lambda: 0,
+            device_count=lambda: 0,
+        )
+        sys.modules["torch"] = torch_module
+
+    logging_module = types.ModuleType(f"{package_name}.utils.logging")
+    logging_module.debug_log = lambda *_args, **_kwargs: None
+    logging_module.log = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.utils.logging"] = logging_module
+
+    config_module = types.ModuleType(f"{package_name}.utils.config")
+    config_module.load_config = lambda: {"workers": []}
+    sys.modules[f"{package_name}.utils.config"] = config_module
+
+    network_module = types.ModuleType(f"{package_name}.utils.network")
+
+    async def _handle_api_error(_request, error, status=500):
+        return _FakeResponse({"status": "error", "message": str(error)}, status=status)
+
+    network_module.handle_api_error = _handle_api_error
+    network_module.normalize_host = lambda value: value
+    network_module.build_worker_url = lambda worker, endpoint="": f"http://localhost:{worker.get('port', 8188)}{endpoint}"
+
+    async def _probe_worker(*_args, **_kwargs):
+        return None
+
+    network_module.probe_worker = _probe_worker
+
+    async def _get_client_session():
+        raise RuntimeError("not used in these tests")
+
+    network_module.get_client_session = _get_client_session
+    sys.modules[f"{package_name}.utils.network"] = network_module
+
+    constants_module = types.ModuleType(f"{package_name}.utils.constants")
+    constants_module.CHUNK_SIZE = 8192
+    sys.modules[f"{package_name}.utils.constants"] = constants_module
+
+    async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
+
+    async def _queue_prompt_payload(*_args, **_kwargs):
+        return "prompt-id"
+
+    async_helpers_module.queue_prompt_payload = _queue_prompt_payload
+    sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
+
+    schemas_module = types.ModuleType(f"{package_name}.api.schemas")
+
+    def _require_fields(data, *fields):
+        missing = []
+        for field in fields:
+            value = data.get(field) if isinstance(data, dict) else None
+            if value is None or (isinstance(value, str) and not value.strip()):
+                missing.append(field)
+        return missing
+
+    def _validate_worker_id(worker_id, config):
+        return any(str(worker.get("id")) == str(worker_id) for worker in config.get("workers", []))
+
+    schemas_module.require_fields = _require_fields
+    schemas_module.validate_worker_id = _validate_worker_id
+    sys.modules[f"{package_name}.api.schemas"] = schemas_module
+
+    spec = importlib.util.spec_from_file_location(f"{package_name}.api.worker_routes", module_path)
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    spec.loader.exec_module(module)
+
+    if created_aiohttp_stub:
+        sys.modules.pop("aiohttp", None)
+
+    return module
+
+
+worker_routes = _load_worker_routes_module()
+
+
+class WorkerRoutesTests(unittest.IsolatedAsyncioTestCase):
+    async def test_launch_worker_valid_id_returns_200(self):
+        manager = _DummyWorkerManager()
+        config = {"workers": [{"id": "worker-a", "name": "Worker A", "port": 8188}]}
+        request = _FakeRequest({"worker_id": "worker-a"})
+
+        with patch.object(worker_routes, "get_worker_manager", return_value=manager), patch.object(
+            worker_routes, "load_config", return_value=config
+        ), patch.object(
+            worker_routes.asyncio, "get_running_loop", return_value=_ImmediateLoop()
+        ):
+            response = await worker_routes.launch_worker_endpoint(request)
+
+        self.assertEqual(response.status, 200)
+        self.assertEqual(response.payload.get("status"), "success")
+        self.assertEqual(response.payload.get("pid"), 12345)
+
+    async def test_launch_worker_unknown_id_returns_404(self):
+        manager = _DummyWorkerManager()
+        config = {"workers": [{"id": "worker-a", "name": "Worker A", "port": 8188}]}
+        request = _FakeRequest({"worker_id": "missing-worker"})
+
+        with patch.object(worker_routes, "get_worker_manager", return_value=manager), patch.object(
+            worker_routes, "load_config", return_value=config
+        ):
+            response = await worker_routes.launch_worker_endpoint(request)
+
+        self.assertEqual(response.status, 404)
+        self.assertIn("not found", response.payload.get("message", "").lower())
+
+    async def test_worker_log_returns_content_json(self):
+        manager = _DummyWorkerManager()
+        with tempfile.NamedTemporaryFile("w", delete=False, encoding="utf-8") as handle:
+            handle.write("line-1\nline-2\nline-3\n")
+            log_path = handle.name
+
+        manager.processes["worker-a"] = {
+            "pid": 9999,
+            "log_file": log_path,
+            "process": None,
+        }
+
+        request = _FakeRequest(match_info={"worker_id": "worker-a"}, query={"lines": "2"})
+        try:
+            with patch.object(worker_routes, "get_worker_manager", return_value=manager), patch.object(
+                worker_routes.asyncio, "get_running_loop", return_value=_ImmediateLoop()
+            ):
+                response = await worker_routes.get_worker_log_endpoint(request)
+        finally:
+            if os.path.exists(log_path):
+                os.remove(log_path)
+
+        self.assertEqual(response.status, 200)
+        self.assertEqual(response.payload.get("status"), "success")
+        self.assertIn("content", response.payload)
+        self.assertIn("line-3", response.payload["content"])
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_dispatch_selection.py

@@ -0,0 +1,192 @@
+import asyncio
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+from unittest.mock import patch
+
+
+def _load_dispatch_module():
+    module_path = Path(__file__).resolve().parents[1] / "api" / "orchestration" / "dispatch.py"
+
+    package_name = "dist_dispatch_testpkg"
+    root_pkg = types.ModuleType(package_name)
+    root_pkg.__path__ = []
+    sys.modules[package_name] = root_pkg
+
+    api_pkg = types.ModuleType(f"{package_name}.api")
+    api_pkg.__path__ = []
+    sys.modules[f"{package_name}.api"] = api_pkg
+
+    orch_pkg = types.ModuleType(f"{package_name}.api.orchestration")
+    orch_pkg.__path__ = []
+    sys.modules[f"{package_name}.api.orchestration"] = orch_pkg
+
+    utils_pkg = types.ModuleType(f"{package_name}.utils")
+    utils_pkg.__path__ = []
+    sys.modules[f"{package_name}.utils"] = utils_pkg
+
+    logging_module = types.ModuleType(f"{package_name}.utils.logging")
+    logging_module.debug_log = lambda *_args, **_kwargs: None
+    logging_module.log = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.utils.logging"] = logging_module
+
+    network_module = types.ModuleType(f"{package_name}.utils.network")
+    network_module.build_worker_url = lambda *_args, **_kwargs: "http://example.invalid"
+
+    async def _probe_worker(*_args, **_kwargs):
+        return None
+
+    network_module.probe_worker = _probe_worker
+
+    async def _fake_session():
+        raise RuntimeError("get_client_session should be mocked in these tests")
+
+    network_module.get_client_session = _fake_session
+    sys.modules[f"{package_name}.utils.network"] = network_module
+
+    created_aiohttp_stub = False
+    if "aiohttp" not in sys.modules:
+        created_aiohttp_stub = True
+        aiohttp_module = types.ModuleType("aiohttp")
+
+        class _ClientTimeout:
+            def __init__(self, total=None):
+                self.total = total
+
+        class _ClientConnectorError(Exception):
+            pass
+
+        class _WSMsgType:
+            TEXT = "TEXT"
+            ERROR = "ERROR"
+            CLOSED = "CLOSED"
+
+        class _TCPConnector:
+            def __init__(self, *args, **kwargs):
+                self.args = args
+                self.kwargs = kwargs
+
+        class _ClientSession:
+            def __init__(self, *args, **kwargs):
+                self.closed = False
+
+            async def close(self):
+                self.closed = True
+
+        aiohttp_module.ClientTimeout = _ClientTimeout
+        aiohttp_module.ClientConnectorError = _ClientConnectorError
+        aiohttp_module.WSMsgType = _WSMsgType
+        aiohttp_module.TCPConnector = _TCPConnector
+        aiohttp_module.ClientSession = _ClientSession
+        aiohttp_module.web = types.SimpleNamespace(
+            json_response=lambda payload, status=200: {"payload": payload, "status": status}
+        )
+        sys.modules["aiohttp"] = aiohttp_module
+
+    spec = importlib.util.spec_from_file_location(
+        f"{package_name}.api.orchestration.dispatch",
+        module_path,
+    )
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    spec.loader.exec_module(module)
+    if created_aiohttp_stub:
+        sys.modules.pop("aiohttp", None)
+    return module
+
+
+dispatch = _load_dispatch_module()
+
+
+class DispatchSelectionTests(unittest.IsolatedAsyncioTestCase):
+    async def test_select_active_workers_filters_offline(self):
+        workers = [
+            {"id": "w1", "name": "Worker 1"},
+            {"id": "w2", "name": "Worker 2"},
+            {"id": "w3", "name": "Worker 3"},
+        ]
+
+        async def fake_probe(worker):
+            return worker["id"] != "w2"
+
+        with patch.object(dispatch, "worker_is_active", side_effect=fake_probe):
+            active_workers, delegate_master = await dispatch.select_active_workers(
+                workers,
+                use_websocket=False,
+                delegate_master=False,
+                probe_concurrency=3,
+            )
+
+        self.assertEqual([w["id"] for w in active_workers], ["w1", "w3"])
+        self.assertFalse(delegate_master)
+
+    async def test_select_active_workers_disables_delegate_when_all_offline(self):
+        workers = [{"id": "w1", "name": "Worker 1"}]
+
+        async def fake_probe(_worker):
+            return False
+
+        with patch.object(dispatch, "worker_is_active", side_effect=fake_probe):
+            active_workers, delegate_master = await dispatch.select_active_workers(
+                workers,
+                use_websocket=False,
+                delegate_master=True,
+                probe_concurrency=1,
+            )
+
+        self.assertEqual(active_workers, [])
+        self.assertFalse(delegate_master)
+
+    async def test_select_active_workers_uses_websocket_probe_when_enabled(self):
+        workers = [{"id": "w1", "name": "Worker 1"}, {"id": "w2", "name": "Worker 2"}]
+
+        async def fake_http_probe(_worker):
+            return False
+
+        async def fake_ws_probe(_worker):
+            return True
+
+        with patch.object(dispatch, "worker_is_active", side_effect=fake_http_probe) as http_probe, patch.object(
+            dispatch,
+            "worker_ws_is_active",
+            side_effect=fake_ws_probe,
+        ) as ws_probe:
+            active_workers, _ = await dispatch.select_active_workers(
+                workers,
+                use_websocket=True,
+                delegate_master=False,
+                probe_concurrency=4,
+            )
+
+        self.assertEqual([w["id"] for w in active_workers], ["w1", "w2"])
+        self.assertEqual(ws_probe.call_count, 2)
+        self.assertEqual(http_probe.call_count, 0)
+
+    async def test_probe_concurrency_is_bounded(self):
+        workers = [{"id": f"w{i}", "name": f"Worker {i}"} for i in range(6)]
+        state = {"in_flight": 0, "max_in_flight": 0}
+
+        async def fake_probe(_worker):
+            state["in_flight"] += 1
+            state["max_in_flight"] = max(state["max_in_flight"], state["in_flight"])
+            await asyncio.sleep(0.01)
+            state["in_flight"] -= 1
+            return True
+
+        with patch.object(dispatch, "worker_is_active", side_effect=fake_probe):
+            active_workers, _ = await dispatch.select_active_workers(
+                workers,
+                use_websocket=False,
+                delegate_master=False,
+                probe_concurrency=2,
+            )
+
+        self.assertEqual(len(active_workers), len(workers))
+        self.assertLessEqual(state["max_in_flight"], 2)
+        self.assertGreaterEqual(state["max_in_flight"], 2)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_network_helpers.py

@@ -0,0 +1,103 @@
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+
+
+def _load_network_module():
+    module_path = Path(__file__).resolve().parents[1] / "utils" / "network.py"
+
+    package_name = "dist_utils_testpkg"
+    package_module = types.ModuleType(package_name)
+    package_module.__path__ = []  # mark as package
+    sys.modules[package_name] = package_module
+
+    logging_module = types.ModuleType(f"{package_name}.logging")
+    logging_module.debug_log = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.logging"] = logging_module
+
+    server_module = types.ModuleType("server")
+    server_module.PromptServer = types.SimpleNamespace(
+        instance=types.SimpleNamespace(address="127.0.0.1", port=8188, loop=None)
+    )
+    sys.modules["server"] = server_module
+
+    if "aiohttp" not in sys.modules:
+        aiohttp_module = types.ModuleType("aiohttp")
+
+        class _TCPConnector:
+            def __init__(self, *args, **kwargs):
+                self.args = args
+                self.kwargs = kwargs
+
+        class _ClientSession:
+            def __init__(self, *args, **kwargs):
+                self.closed = False
+
+            async def close(self):
+                self.closed = True
+
+        aiohttp_module.TCPConnector = _TCPConnector
+        aiohttp_module.ClientSession = _ClientSession
+        aiohttp_module.web = types.SimpleNamespace(
+            json_response=lambda payload, status=200: {"payload": payload, "status": status}
+        )
+        sys.modules["aiohttp"] = aiohttp_module
+
+    spec = importlib.util.spec_from_file_location(f"{package_name}.network", module_path)
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    spec.loader.exec_module(module)
+    return module
+
+
+network = _load_network_module()
+
+
+class NetworkHelpersTests(unittest.TestCase):
+    def test_normalize_host_strips_protocol_and_path(self):
+        self.assertEqual(network.normalize_host("  https://example.com/a/b  "), "example.com")
+
+    def test_normalize_host_keeps_none(self):
+        self.assertIsNone(network.normalize_host(None))
+
+    def test_build_worker_url_defaults_to_server_address(self):
+        worker = {"id": "w1", "port": 8189}
+        self.assertEqual(network.build_worker_url(worker, "/prompt"), "http://127.0.0.1:8189/prompt")
+
+    def test_build_worker_url_cloud_defaults_to_https(self):
+        worker = {"id": "w2", "host": "foo.proxy.runpod.net", "port": 443}
+        self.assertEqual(network.build_worker_url(worker), "https://foo.proxy.runpod.net")
+
+    def test_build_worker_url_keeps_explicit_scheme(self):
+        worker = {"id": "w3", "host": "https://worker.example.com", "port": 1234}
+        self.assertEqual(network.build_worker_url(worker, "/prompt"), "https://worker.example.com/prompt")
+
+    def test_build_master_url_uses_https_for_cloud_host(self):
+        cfg = {"master": {"host": "demo.proxy.runpod.net"}}
+        prompt_server = types.SimpleNamespace(address="127.0.0.1", port=8188)
+        self.assertEqual(
+            network.build_master_url(config=cfg, prompt_server_instance=prompt_server),
+            "https://demo.proxy.runpod.net",
+        )
+
+    def test_build_master_url_keeps_explicit_scheme(self):
+        cfg = {"master": {"host": "https://master.example.com/"}}
+        prompt_server = types.SimpleNamespace(address="127.0.0.1", port=8188)
+        self.assertEqual(
+            network.build_master_url(config=cfg, prompt_server_instance=prompt_server),
+            "https://master.example.com",
+        )
+
+    def test_build_master_url_falls_back_to_server_address(self):
+        cfg = {"master": {"host": ""}}
+        prompt_server = types.SimpleNamespace(address="0.0.0.0", port=8190)
+        self.assertEqual(
+            network.build_master_url(config=cfg, prompt_server_instance=prompt_server),
+            "http://127.0.0.1:8190",
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_queue_request.py

@@ -0,0 +1,151 @@
+import importlib.util
+import unittest
+from pathlib import Path
+
+
+def _load_queue_request_module():
+    module_path = Path(__file__).resolve().parents[1] / "api" / "queue_request.py"
+    spec = importlib.util.spec_from_file_location("queue_request", module_path)
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    spec.loader.exec_module(module)
+    return module
+
+
+queue_request = _load_queue_request_module()
+parse_queue_request_payload = queue_request.parse_queue_request_payload
+
+
+class QueueRequestPayloadTests(unittest.TestCase):
+    def _base_payload(self):
+        return {
+            "prompt": {"1": {"class_type": "Anything"}},
+            "enabled_worker_ids": ["worker-1"],
+            "client_id": "client-1",
+        }
+
+    def test_normalizes_enabled_worker_ids(self):
+        payload_data = self._base_payload()
+        payload_data["enabled_worker_ids"] = ["a", 2, 3]
+        payload_data["delegate_master"] = True
+        payload = parse_queue_request_payload(
+            payload_data
+        )
+        self.assertEqual(payload.enabled_worker_ids, ["a", "2", "3"])
+        self.assertTrue(payload.delegate_master)
+
+    def test_supports_legacy_workers_field(self):
+        payload_data = self._base_payload()
+        payload_data.pop("enabled_worker_ids", None)
+        payload_data["workers"] = [{"id": "w1"}, "w2", {"id": 3}, {"name": "no-id"}]
+        payload = parse_queue_request_payload(
+            payload_data
+        )
+        self.assertEqual(payload.enabled_worker_ids, ["w1", "w2", "3"])
+
+    def test_supports_auto_prepare_prompt_fallback(self):
+        payload_data = self._base_payload()
+        payload_data.pop("prompt", None)
+        payload_data["auto_prepare"] = True
+        payload_data["workflow"] = {
+            "prompt": {
+                "10": {"class_type": "DistributedCollector"},
+            }
+        }
+        payload = parse_queue_request_payload(
+            payload_data
+        )
+        self.assertIn("10", payload.prompt)
+        self.assertTrue(payload.auto_prepare)
+
+    def test_normalizes_trace_execution_id(self):
+        payload_data = self._base_payload()
+        payload_data["trace_execution_id"] = "  exec_123  "
+        payload = parse_queue_request_payload(
+            payload_data
+        )
+        self.assertEqual(payload.trace_execution_id, "exec_123")
+
+    def test_blank_trace_execution_id_normalizes_to_none(self):
+        payload_data = self._base_payload()
+        payload_data["trace_execution_id"] = "   "
+        payload = parse_queue_request_payload(
+            payload_data
+        )
+        self.assertIsNone(payload.trace_execution_id)
+
+    def test_auto_prepare_defaults_true(self):
+        payload = parse_queue_request_payload(self._base_payload())
+        self.assertTrue(payload.auto_prepare)
+
+    def test_workers_field_must_be_list(self):
+        payload_data = self._base_payload()
+        payload_data.pop("enabled_worker_ids", None)
+        payload_data["workers"] = "worker-a"
+        with self.assertRaisesRegex(ValueError, "Field 'workers' must be a list"):
+            parse_queue_request_payload(payload_data)
+
+    def test_trace_execution_id_must_be_string(self):
+        payload_data = self._base_payload()
+        payload_data["trace_execution_id"] = 123
+        with self.assertRaisesRegex(ValueError, "trace_execution_id must be a string"):
+            parse_queue_request_payload(payload_data)
+
+    def test_auto_prepare_false_still_falls_back_to_workflow_prompt(self):
+        payload_data = self._base_payload()
+        payload_data.pop("prompt", None)
+        payload_data["auto_prepare"] = False
+        payload_data["workflow"] = {
+            "prompt": {"10": {"class_type": "DistributedCollector"}},
+        }
+        payload = parse_queue_request_payload(payload_data)
+        self.assertIn("10", payload.prompt)
+        self.assertFalse(payload.auto_prepare)
+
+    def test_auto_prepare_must_be_boolean(self):
+        payload_data = self._base_payload()
+        payload_data["auto_prepare"] = "true"
+        with self.assertRaisesRegex(ValueError, "auto_prepare must be a boolean"):
+            parse_queue_request_payload(payload_data)
+
+    def test_invalid_delegate_master_type_raises(self):
+        payload_data = self._base_payload()
+        payload_data["delegate_master"] = "yes"
+        with self.assertRaisesRegex(ValueError, "delegate_master must be a boolean"):
+            parse_queue_request_payload(payload_data)
+
+    def test_invalid_enabled_worker_ids_type_raises(self):
+        payload_data = self._base_payload()
+        payload_data["enabled_worker_ids"] = "worker-a"
+        with self.assertRaisesRegex(ValueError, "enabled_worker_ids must be a list"):
+            parse_queue_request_payload(payload_data)
+
+    def test_invalid_top_level_payload_raises(self):
+        with self.assertRaisesRegex(ValueError, "Expected a JSON object body"):
+            parse_queue_request_payload(["not", "an", "object"])
+
+    def test_missing_prompt_raises(self):
+        with self.assertRaisesRegex(ValueError, "Field 'prompt' must be an object"):
+            parse_queue_request_payload(
+                {
+                    "workflow": {},
+                    "enabled_worker_ids": ["worker-1"],
+                    "client_id": "client-1",
+                }
+            )
+
+    def test_missing_enabled_worker_ids_raises(self):
+        payload_data = self._base_payload()
+        payload_data.pop("enabled_worker_ids", None)
+        with self.assertRaisesRegex(ValueError, "enabled_worker_ids required"):
+            parse_queue_request_payload(payload_data)
+
+    def test_missing_client_id_raises(self):
+        payload_data = self._base_payload()
+        payload_data.pop("client_id", None)
+        with self.assertRaisesRegex(ValueError, "client_id required"):
+            parse_queue_request_payload(payload_data)
+
+
+if __name__ == "__main__":
+    unittest.main()

upscale/worker_comms.py

@@ -10,9 +10,11 @@ from ..utils.image import tensor_to_pil
 
 class WorkerCommsMixin:
     async def send_tiles_batch_to_master(self, processed_tiles, multi_job_id, master_url, 
-                                       padding, worker_id):
+                                       padding, worker_id, is_final_flush=False):
         """Send all processed tiles to master, chunked if large."""
         if not processed_tiles:
+            if is_final_flush:
+                await self._send_tiles_completion_signal(multi_job_id, master_url, worker_id)
             return  # Early exit if empty
 
         total_tiles = len(processed_tiles)
@@ -75,7 +77,7 @@ class WorkerCommsMixin:
 
             chunk_size = j - i
             is_chunk_last = (j >= total_tiles)
-            data.add_field('is_last', str(is_chunk_last))
+            data.add_field('is_last', str(bool(is_final_flush and is_chunk_last)))
             data.add_field('batch_size', str(chunk_size))
             data.add_field('tiles_metadata', json.dumps(metadata), content_type='application/json')
 
@@ -101,100 +103,85 @@ class WorkerCommsMixin:
             chunk_index += 1
             i = j
 
-    async def _request_image_from_master(self, multi_job_id, master_url, worker_id):
-        """Request an image index to process from master in dynamic mode."""
-        # Enhanced retries with 404-specific delay, backoff cap, and total timeout to handle init races
+    async def _send_tiles_completion_signal(self, multi_job_id, master_url, worker_id):
+        """Send completion signal to master in static mode when no tiles are left."""
+        data = aiohttp.FormData()
+        data.add_field('multi_job_id', multi_job_id)
+        data.add_field('worker_id', str(worker_id))
+        data.add_field('is_last', 'true')
+        data.add_field('batch_size', '0')
+
+        session = await get_client_session()
+        url = f"{master_url}/distributed/submit_tiles"
+        async with session.post(url, data=data) as response:
+            response.raise_for_status()
+            debug_log(f"Worker {worker_id} sent static completion signal")
+
+    async def _request_work_item_from_master(
+        self,
+        multi_job_id,
+        master_url,
+        worker_id,
+        endpoint="/distributed/request_image",
+    ):
+        """Request one work item from master with retry/backoff and total timeout."""
         max_retries = 10
         retry_delay = 0.5
-        start_time = time.time()
-        
+        start_time = time.monotonic()
+        url = f"{master_url}{endpoint}"
+
         for attempt in range(max_retries):
-            # Check total timeout
-            if time.time() - start_time > 30:
+            if time.monotonic() - start_time > 30:
                 log(f"Total request timeout after 30s for worker {worker_id}")
-                return None, 0
-                
+                return None
+
             try:
                 session = await get_client_session()
-                url = f"{master_url}/distributed/request_image"
-                
                 async with session.post(url, json={
                     'worker_id': str(worker_id),
                     'multi_job_id': multi_job_id
                 }) as response:
                     if response.status == 200:
-                        data = await response.json()
-                        image_idx = data.get('image_idx')
-                        estimated_remaining = data.get('estimated_remaining', 0)
-                        return image_idx, estimated_remaining
-                    elif response.status == 404:
-                        # Special handling for 404 - job not found yet
+                        return await response.json()
+                    if response.status == 404:
                         text = await response.text()
                         debug_log(f"Job not found (404), will retry: {text}")
                         await asyncio.sleep(1.0)
                     else:
                         text = await response.text()
-                        debug_log(f"Request image failed: {response.status} - {text}")
-                        
-            except Exception as e:
+                        debug_log(
+                            f"Request work item failed ({response.status}) for worker {worker_id}: {text}"
+                        )
+
+            except Exception as exc:
                 if attempt < max_retries - 1:
-                    debug_log(f"Retry {attempt + 1}/{max_retries} after error: {e}")
+                    debug_log(f"Retry {attempt + 1}/{max_retries} after error: {exc}")
                     await asyncio.sleep(retry_delay)
-                    retry_delay *= 2
-                    retry_delay = min(retry_delay, 5.0)  # Cap backoff at 5s
+                    retry_delay = min(retry_delay * 2, 5.0)
                 else:
-                    log(f"Failed to request image after {max_retries} attempts: {e}")
+                    log(f"Failed to request work item after {max_retries} attempts: {exc}")
                     raise
-        
-        return None, 0
+
+        return None
+
+    async def _request_image_from_master(self, multi_job_id, master_url, worker_id):
+        """Request an image index to process from master in dynamic mode."""
+        data = await self._request_work_item_from_master(multi_job_id, master_url, worker_id)
+        if not data:
+            return None, 0
+        image_idx = data.get('image_idx')
+        estimated_remaining = data.get('estimated_remaining', 0)
+        return image_idx, estimated_remaining
 
     async def _request_tile_from_master(self, multi_job_id, master_url, worker_id):
         """Request a tile index to process from master in static mode (reusing dynamic infrastructure)."""
-        # Reuse the same retry logic as dynamic mode
-        max_retries = 10
-        retry_delay = 0.5
-        start_time = time.time()
-        
-        for attempt in range(max_retries):
-            # Check total timeout
-            if time.time() - start_time > 30:
-                log(f"Total request timeout after 30s for worker {worker_id}")
-                return None, 0
-                
-            try:
-                session = await get_client_session()
-                url = f"{master_url}/distributed/request_image"  # Same endpoint
-                
-                async with session.post(url, json={
-                    'worker_id': str(worker_id),
-                    'multi_job_id': multi_job_id
-                }) as response:
-                    if response.status == 200:
-                        data = await response.json()
-                        tile_idx = data.get('tile_idx')
-                        estimated_remaining = data.get('estimated_remaining', 0)
-                        batched_static = data.get('batched_static', False)
-                        return tile_idx, estimated_remaining, batched_static
-                    elif response.status == 404:
-                        # Special handling for 404 - job not found yet
-                        text = await response.text()
-                        debug_log(f"Job not found (404), will retry: {text}")
-                        await asyncio.sleep(1.0)
-                    else:
-                        text = await response.text()
-                        debug_log(f"Request tile failed: {response.status} - {text}")
-                        
-            except Exception as e:
-                if attempt < max_retries - 1:
-                    debug_log(f"Retry {attempt + 1}/{max_retries} after error: {e}")
-                    await asyncio.sleep(retry_delay)
-                    retry_delay *= 2
-                    retry_delay = min(retry_delay, 5.0)  # Cap backoff at 5s
-                else:
-                    log(f"Failed to request tile after {max_retries} attempts: {e}")
-                    raise
-        
-        return None, 0, False
+        data = await self._request_work_item_from_master(multi_job_id, master_url, worker_id)
+        if not data:
+            return None, 0, False
+        tile_idx = data.get('tile_idx')
+        estimated_remaining = data.get('estimated_remaining', 0)
+        batched_static = data.get('batched_static', False)
+        return tile_idx, estimated_remaining, batched_static
 
     async def _send_full_image_to_master(self, image_pil, image_idx, multi_job_id, 
                                         master_url, worker_id, is_last):

utils/config.py

@@ -26,7 +26,11 @@ def get_default_config():
             "auto_launch_workers": False,
             "stop_workers_on_master_exit": True,
             "master_delegate_only": False,
-            "websocket_orchestration": True
+            "websocket_orchestration": True,
+            "worker_probe_concurrency": 8,
+            "worker_prep_concurrency": 4,
+            "media_sync_concurrency": 2,
+            "media_sync_timeout_seconds": 120
         },
         "tunnel": {
             "status": "stopped",

utils/constants.py

@@ -30,6 +30,7 @@ TILE_TRANSFER_TIMEOUT = 30.0
 PROCESS_WAIT_TIMEOUT = 3.0
 QUEUE_INIT_TIMEOUT = 5.0
 TILE_SEND_TIMEOUT = 60.0
+JOB_INIT_GRACE_PERIOD = 10.0
 
 # Memory operations  
 MEMORY_CLEAR_DELAY = 0.5
@@ -38,4 +39,19 @@ MEMORY_CLEAR_DELAY = 0.5
 MAX_BATCH = int(os.environ.get('COMFYUI_MAX_BATCH', '20'))  # Maximum items per batch to prevent timeouts/OOM (~100MB chunks for 512x512 PNGs)
 
 # Heartbeat monitoring
+HEARTBEAT_INTERVAL = float(os.environ.get('COMFYUI_HEARTBEAT_INTERVAL', '10'))  # Heartbeat/check interval in seconds
 HEARTBEAT_TIMEOUT = int(os.environ.get('COMFYUI_HEARTBEAT_TIMEOUT', '60'))  # Worker heartbeat timeout in seconds (default 60s)
+
+# Orchestration pipeline
+ORCHESTRATION_WORKER_PROBE_CONCURRENCY = int(
+    os.environ.get('COMFYUI_ORCHESTRATION_WORKER_PROBE_CONCURRENCY', '8')
+)
+ORCHESTRATION_WORKER_PREP_CONCURRENCY = int(
+    os.environ.get('COMFYUI_ORCHESTRATION_WORKER_PREP_CONCURRENCY', '4')
+)
+ORCHESTRATION_MEDIA_SYNC_CONCURRENCY = int(
+    os.environ.get('COMFYUI_ORCHESTRATION_MEDIA_SYNC_CONCURRENCY', '2')
+)
+ORCHESTRATION_MEDIA_SYNC_TIMEOUT = float(
+    os.environ.get('COMFYUI_ORCHESTRATION_MEDIA_SYNC_TIMEOUT', '120')
+)

utils/network.py

@@ -78,3 +78,83 @@ def build_worker_url(worker, endpoint=""):
         base = f"{scheme}://{host}{port_part}"
 
     return f"{base}{endpoint}"
+
+
+async def probe_worker(worker_url: str, timeout: float = 5.0) -> dict | None:
+    """GET {worker_url}/prompt. Returns parsed JSON or None on any failure."""
+    base_url = (worker_url or "").strip().rstrip("/")
+    if not base_url:
+        return None
+    probe_url = base_url if base_url.endswith("/prompt") else f"{base_url}/prompt"
+    session = await get_client_session()
+    try:
+        async with session.get(
+            probe_url,
+            timeout=aiohttp.ClientTimeout(total=float(timeout)),
+        ) as response:
+            if response.status != 200:
+                debug_log(f"[Distributed] Worker probe non-200 status: {response.status} ({probe_url})")
+                return None
+            payload = await response.json()
+            if isinstance(payload, dict):
+                return payload
+            debug_log(f"[Distributed] Worker probe returned non-object JSON: {probe_url}")
+            return None
+    except asyncio.TimeoutError:
+        debug_log(f"[Distributed] Worker probe timed out: {probe_url}")
+        return None
+    except aiohttp.ClientConnectorError:
+        debug_log(f"[Distributed] Worker unreachable: {probe_url}")
+        return None
+    except Exception as exc:
+        debug_log(f"[Distributed] Worker probe error ({probe_url}): {exc}")
+        return None
+
+
+def build_master_url(config=None, prompt_server_instance=None):
+    """Build the best public URL workers should use to reach the master."""
+    if config is None:
+        from .config import load_config
+        config = load_config()
+
+    prompt_server_instance = prompt_server_instance or server.PromptServer.instance
+    master_cfg = (config or {}).get("master", {}) or {}
+    configured_host = (master_cfg.get("host") or "").strip()
+    configured_port = master_cfg.get("port")
+    default_port = getattr(prompt_server_instance, "port", 8188) or 8188
+    try:
+        port = int(configured_port or default_port)
+    except (TypeError, ValueError):
+        port = int(default_port)
+
+    def _needs_https(hostname):
+        hostname = hostname.lower()
+        https_domains = (
+            ".proxy.runpod.net",
+            ".ngrok-free.app",
+            ".ngrok-free.dev",
+            ".ngrok.io",
+            ".trycloudflare.com",
+            ".cloudflare.dev",
+        )
+        return any(hostname.endswith(suffix) for suffix in https_domains)
+
+    if configured_host:
+        if configured_host.startswith(("http://", "https://")):
+            return configured_host.rstrip("/")
+
+        host = configured_host
+        scheme = "https" if _needs_https(host) or port == 443 else "http"
+        default_port_for_scheme = 443 if scheme == "https" else 80
+        if configured_port is None and scheme == "https" and _needs_https(host):
+            port = default_port_for_scheme
+        port_part = "" if port == default_port_for_scheme else f":{port}"
+        return f"{scheme}://{host}{port_part}"
+
+    address = getattr(prompt_server_instance, "address", "127.0.0.1") or "127.0.0.1"
+    if address in ("0.0.0.0", "::"):
+        address = "127.0.0.1"
+    scheme = "https" if port == 443 else "http"
+    default_port_for_scheme = 443 if scheme == "https" else 80
+    port_part = "" if port == default_port_for_scheme else f":{port}"
+    return f"{scheme}://{address}{port_part}"

web/apiClient.js

@@ -97,27 +97,48 @@ export function createApiClient(baseUrl) {
             });
         },
         
-        // Job preparation
-        async prepareJob(multiJobId) {
-            return request('/distributed/prepare_job', {
-                method: 'POST',
-                body: JSON.stringify({ multi_job_id: multiJobId })
-            });
-        },
-
         async queueDistributed(payload) {
             return request('/distributed/queue', {
                 method: 'POST',
                 body: JSON.stringify(payload)
             });
         },
-        
-        // Image loading
-        async loadImage(imagePath) {
-            return request('/distributed/load_image', {
-                method: 'POST',
-                body: JSON.stringify({ image_path: imagePath })
+
+        async probeWorker(workerUrl, timeoutMs = TIMEOUTS.STATUS_CHECK) {
+            const normalizedWorkerUrl = normalizeWorkerUrl(workerUrl);
+            const response = await fetch(`${normalizedWorkerUrl}/prompt`, {
+                method: 'GET',
+                mode: 'cors',
+                cache: 'no-store',
+                signal: AbortSignal.timeout(timeoutMs)
             });
+
+            if (!response.ok) {
+                return { ok: false, status: response.status, queueRemaining: null };
+            }
+
+            const data = await response.json().catch(() => ({}));
+            return {
+                ok: true,
+                status: response.status,
+                queueRemaining: data.exec_info?.queue_remaining || 0,
+            };
+        },
+
+        async dispatchToWorker(workerUrl, promptPayload) {
+            const normalizedWorkerUrl = normalizeWorkerUrl(workerUrl);
+            const response = await fetch(`${normalizedWorkerUrl}/prompt`, {
+                method: 'POST',
+                headers: { 'Content-Type': 'application/json' },
+                mode: 'cors',
+                body: JSON.stringify(promptPayload),
+            });
+
+            if (!response.ok) {
+                throw new Error(`Worker returned ${response.status}`);
+            }
+
+            return response.json();
         },
         
         // Network info

web/constants.js

@@ -17,6 +17,7 @@ export const BUTTON_STYLES = {
     launch: "background-color: #4a7c4a;",
     stop: "background-color: #7c4a4a;",
     log: "background-color: #685434;",
+    working: "background-color: #666;",
     clearMemory: "background-color: #555; padding: 6px 14px;",
     interrupt: "background-color: #555; padding: 6px 14px;",
 };

web/distributed.css

@@ -0,0 +1,32 @@
+:root {
+    --btn-stop: #7c4a4a;
+    --btn-launch: #4a7c4a;
+    --btn-log: #685434;
+    --btn-working: #666;
+    --btn-success: #3a6a3a;
+    --btn-error: #6a3a3a;
+}
+
+.btn--stop {
+    background-color: var(--btn-stop) !important;
+}
+
+.btn--launch {
+    background-color: var(--btn-launch) !important;
+}
+
+.btn--log {
+    background-color: var(--btn-log) !important;
+}
+
+.btn--working {
+    background-color: var(--btn-working) !important;
+}
+
+.btn--success {
+    background-color: var(--btn-success) !important;
+}
+
+.btn--error {
+    background-color: var(--btn-error) !important;
+}

web/executionDecisionUtils.js

@@ -0,0 +1,30 @@
+import { NODE_CLASSES } from "./constants.js";
+
+
+export function selectLeastBusyWorker(extension, statuses) {
+    // Find idle workers (queue_remaining == 0)
+    const idleWorkers = statuses.filter((s) => s.queueRemaining === 0);
+
+    if (idleWorkers.length > 0) {
+        // Round-robin among idle workers
+        if (!extension._distributedQueueRRIndex) {
+            extension._distributedQueueRRIndex = 0;
+        }
+        const index = extension._distributedQueueRRIndex % idleWorkers.length;
+        extension._distributedQueueRRIndex++;
+        return idleWorkers[index];
+    }
+
+    // No idle workers - pick the one with the shortest queue
+    return statuses.reduce((min, s) => (s.queueRemaining < min.queueRemaining ? s : min));
+}
+
+
+export function markSkipDispatch(promptObj) {
+    for (const node of Object.values(promptObj)) {
+        if (node && typeof node === "object" && node.class_type === NODE_CLASSES.DISTRIBUTED_QUEUE) {
+            node.inputs = node.inputs || {};
+            node.inputs.skip_dispatch = true;
+        }
+    }
+}

web/executionUtils.js

@@ -1,59 +1,10 @@
 import { api } from "../../scripts/api.js";
-import { findNodesByClass, findImageReferences, hasUpstreamNode, pruneWorkflowForWorker, getCachedWorkerSystemInfo, findCollectorDownstreamNodes } from './workerUtils.js';
-import { TIMEOUTS, NODE_CLASSES, generateUUID, STATUS_COLORS } from './constants.js';
+import { findNodesByClass } from './workerUtils.js';
+import { TIMEOUTS, NODE_CLASSES, generateUUID } from './constants.js';
+import { markSkipDispatch, selectLeastBusyWorker } from './executionDecisionUtils.js';
+import { checkAllWorkerStatuses, getWorkerUrl } from './workerLifecycle.js';
 
-/**
- * Convert paths in the API prompt to match the target platform's separator
- * @param {Object} apiPrompt - The workflow API prompt
- * @param {string} targetSeparator - The target path separator ('\\' or '/')
- * @returns {Object} The converted API prompt
- */
-function convertPathsForPlatform(apiPrompt, targetSeparator) {
-    // Validate target separator
-    if (!targetSeparator || (targetSeparator !== '/' && targetSeparator !== '\\')) {
-        console.warn('[Distributed] Invalid target separator:', targetSeparator, '- skipping path conversion');
-        return apiPrompt;
-    }
-    
-    // Regex to identify likely file paths with extensions
-    const isLikelyFilename = (value) => {
-        return value.match(/\.(ckpt|safetensors|pt|pth|bin|yaml|json|png|jpg|jpeg|webp|gif|bmp|latent|txt|vae|lora|embedding)(\s*\[\w+\])?$/i);
-    };
-    const isImageOrVideo = (value) => {
-        return value.match(/\.(png|jpg|jpeg|webp|gif|bmp|mp4|avi|mov|mkv|webm)(\s*\[\w+\])?$/i);
-    };
-    
-    function convert(obj) {
-        if (typeof obj === 'string') {
-            // Only convert strings that look like file paths
-            if ((obj.includes('\\') || obj.includes('/')) && isLikelyFilename(obj)) {
-                const trimmed = obj.trim();
-                const hasDrive = /^[A-Za-z]:\\\\|^[A-Za-z]:\//.test(trimmed);
-                const isAbsolute = trimmed.startsWith('/') || trimmed.startsWith('\\\\');
-                const hasProtocol = /^\w+:\/\//.test(trimmed);
-
-                // For annotated relative image/video paths, keep forward slashes
-                if (!hasDrive && !isAbsolute && !hasProtocol && isImageOrVideo(trimmed)) {
-                    return trimmed.replace(/[\\\\]/g, '/');
-                }
-                // Otherwise replace any path separator with the worker's target separator
-                return trimmed.replace(/[\\\\\/]/g, targetSeparator);
-            }
-            return obj;
-        } else if (Array.isArray(obj)) {
-            return obj.map(convert);
-        } else if (typeof obj === 'object' && obj !== null) {
-            const newObj = {};
-            for (const [key, value] of Object.entries(obj)) {
-                newObj[key] = convert(value);
-            }
-            return newObj;
-        }
-        return obj;
-    }
-    
-    return convert(apiPrompt);
-}
+export { markSkipDispatch, selectLeastBusyWorker };
 
 export function setupInterceptor(extension) {
     api.queuePrompt = async (number, prompt) => {
@@ -65,17 +16,17 @@ export function setupInterceptor(extension) {
             if (hasCollector || hasDistUpscale) {
                 const result = await executeParallelDistributed(extension, prompt);
                 // Immediate status check for instant feedback
-                extension.checkAllWorkerStatuses();
+                checkAllWorkerStatuses(extension);
                 // Another check after a short delay to catch state changes
-                setTimeout(() => extension.checkAllWorkerStatuses(), TIMEOUTS.POST_ACTION_DELAY);
+                setTimeout(() => checkAllWorkerStatuses(extension), TIMEOUTS.POST_ACTION_DELAY);
                 return result;
             }
 
             // DistributedQueue: route entire workflow to least-busy worker
             if (hasDistQueue) {
                 const result = await executeQueueDistributed(extension, prompt);
-                extension.checkAllWorkerStatuses();
-                setTimeout(() => extension.checkAllWorkerStatuses(), TIMEOUTS.POST_ACTION_DELAY);
+                checkAllWorkerStatuses(extension);
+                setTimeout(() => checkAllWorkerStatuses(extension), TIMEOUTS.POST_ACTION_DELAY);
                 return result;
             }
         }
@@ -84,9 +35,10 @@ export function setupInterceptor(extension) {
 }
 
 export async function executeParallelDistributed(extension, promptWrapper) {
+    const traceExecutionId = `exec_${Date.now()}_${generateUUID().slice(0, 6)}`;
     try {
-        const executionPrefix = "exec_" + Date.now(); // Unique ID for this specific execution
         const enabledWorkers = extension.enabledWorkers;
+        extension.log(`[exec:${traceExecutionId}] Starting distributed execution`, "debug");
         
         // Pre-flight health check on all enabled workers
         const activeWorkers = await performPreflightCheck(extension, enabledWorkers);
@@ -136,98 +88,32 @@ export async function executeParallelDistributed(extension, promptWrapper) {
             }
         }
 
-        try {
-            const queueResponse = await extension.api.queueDistributed({
-                prompt: promptWrapper.output,
-                workflow: promptWrapper.workflow,
-                enabled_worker_ids: activeWorkers.map((worker) => worker.id),
-                workers: activeWorkers.map((worker) => ({ id: worker.id })),
-                client_id: api.clientId,
-                delegate_master: Boolean(extension.config?.settings?.master_delegate_only),
-                auto_prepare: true,
-            });
-            if (queueResponse?.prompt_id && queueResponse?.auto_prepare_supported !== false) {
-                extension.log(
-                    `Distributed queue accepted by backend (prompt_id=${queueResponse.prompt_id}, workers=${queueResponse.worker_count ?? activeWorkers.length})`,
-                    "debug"
-                );
-                return queueResponse;
-            }
-            extension.log("Backend responded without auto-prepare support; using legacy frontend orchestration.", "warn");
-        } catch (error) {
-            extension.log(`Backend auto-prepare dispatch failed, using legacy frontend path: ${error.message}`, "warn");
-        }
-
-        // Find all distributed nodes in the workflow
-        const collectorNodes = findNodesByClass(promptWrapper.output, NODE_CLASSES.DISTRIBUTED_COLLECTOR);
-        const upscaleNodes = findNodesByClass(promptWrapper.output, NODE_CLASSES.UPSCALE_DISTRIBUTED);
-        const allDistributedNodes = [...collectorNodes, ...upscaleNodes];
-
-        const wantsDelegate = Boolean(extension.config?.settings?.master_delegate_only);
-        let masterDelegateActive = false;
-        if (wantsDelegate) {
-            if (activeWorkers.length === 0) {
-                extension.log("Master delegate-only mode disabled: no active workers detected. Falling back to master execution.", "debug");
-            } else if (upscaleNodes.length > 0) {
-                extension.log("Master delegate-only mode is not yet supported for UltimateSDUpscaleDistributed nodes. Falling back to normal execution.", "warn");
-            } else if (collectorNodes.length === 0) {
-                extension.log("Master delegate-only mode requested but no DistributedCollector nodes found. Falling back to normal execution.", "debug");
-            } else {
-                masterDelegateActive = true;
-                extension.log("Master delegate-only mode enabled: master will orchestrate without running upstream nodes.", "debug");
-            }
-        }
-        
-        // Map original node IDs to truly unique job IDs for this specific run
-        const job_id_map = new Map(allDistributedNodes.map(node => [node.id, `${executionPrefix}_${node.id}`]));
-        
-        // Prepare a separate job queue on the backend for each unique job ID
-        const preparePromises = Array.from(job_id_map.values()).map(uniqueId => prepareDistributedJob(extension, uniqueId));
-        await Promise.all(preparePromises);
-
-        const jobs = [];
-        // Use only active workers
-        const participants = ['master', ...activeWorkers.map(w => w.id)];
-
-        for (const participantId of participants) {
-            const options = { 
-                enabled_worker_ids: activeWorkers.map(w => w.id), 
-                workflow: promptWrapper.workflow,
-                job_id_map: job_id_map, // Pass the map of unique IDs
-                delegate_master: masterDelegateActive && participantId === 'master'
-            };
-            
-            const jobApiPrompt = await prepareApiPromptForParticipant(
-                extension, promptWrapper.output, participantId, options
+        const queueResponse = await extension.api.queueDistributed({
+            prompt: promptWrapper.output,
+            workflow: promptWrapper.workflow,
+            enabled_worker_ids: activeWorkers.map((worker) => worker.id),
+            workers: activeWorkers.map((worker) => ({ id: worker.id })),
+            client_id: api.clientId,
+            delegate_master: Boolean(extension.config?.settings?.master_delegate_only),
+            auto_prepare: true,
+            trace_execution_id: traceExecutionId,
+        });
+        if (queueResponse?.prompt_id) {
+            extension.log(
+                `[exec:${traceExecutionId}] Distributed queue accepted by backend (prompt_id=${queueResponse.prompt_id}, workers=${queueResponse.worker_count ?? activeWorkers.length})`,
+                "debug"
             );
-            
-            if (participantId === 'master') {
-                jobs.push({ type: 'master', promptWrapper: { ...promptWrapper, output: jobApiPrompt } });
-            } else {
-                const worker = activeWorkers.find(w => w.id === participantId);
-                if (worker) {
-                    const job = {
-                        type: 'worker',
-                        worker,
-                        prompt: jobApiPrompt,
-                        workflow: promptWrapper.workflow
-                    };
-                    
-                    // Add image references if found for remote workers
-                    if (options._imageReferences) {
-                        job.imageReferences = options._imageReferences;
-                    }
-                    
-                    jobs.push(job);
-                }
-            }
+            return queueResponse;
         }
-        
-        const result = await executeJobs(extension, jobs);
-        return result;
+        throw new Error(
+            `[exec:${traceExecutionId}] Backend did not return a prompt_id for distributed queue.`
+        );
     } catch (error) {
-        extension.log("Parallel execution failed: " + error.message, "error");
-        throw error;
+        extension.log(`[exec:${traceExecutionId}] Distributed execution failed: ${error.message}`, "error");
+        if (extension.ui?.showToast) {
+            extension.ui.showToast(extension.app, "error", "Distributed Failed", error.message, 5000);
+        }
+        return null;
     }
 }
 
@@ -264,24 +150,13 @@ export async function executeQueueDistributed(extension, promptWrapper) {
         markSkipDispatch(promptToSend);
 
         // Dispatch to the selected worker
-        const workerUrl = extension.getWorkerUrl(worker);
+        const workerUrl = getWorkerUrl(extension, worker);
         try {
-            const response = await fetch(`${workerUrl}/prompt`, {
-                method: 'POST',
-                headers: { 'Content-Type': 'application/json' },
-                mode: 'cors',
-                body: JSON.stringify({
-                    prompt: promptToSend,
-                    extra_data: { extra_pnginfo: { workflow: promptWrapper.workflow } },
-                    client_id: api.clientId
-                })
+            const result = await extension.api.dispatchToWorker(workerUrl, {
+                prompt: promptToSend,
+                extra_data: { extra_pnginfo: { workflow: promptWrapper.workflow } },
+                client_id: api.clientId
             });
-
-            if (!response.ok) {
-                throw new Error(`Worker returned ${response.status}`);
-            }
-
-            const result = await response.json();
             extension.log(`DistributedQueue: Dispatched to ${worker.name}, prompt_id=${result.prompt_id}`, "info");
 
             return { prompt_id: result.prompt_id, worker_id: worker.id };
@@ -300,20 +175,14 @@ async function fetchWorkerQueueStatuses(extension, workers) {
     const statuses = [];
 
     const checkPromises = workers.map(async (worker) => {
-        const url = extension.getWorkerUrl(worker, '/prompt');
+        const workerUrl = getWorkerUrl(extension, worker);
         try {
-            const response = await fetch(url, {
-                method: 'GET',
-                mode: 'cors',
-                cache: 'no-store',
-                signal: AbortSignal.timeout(TIMEOUTS.STATUS_CHECK)
-            });
+            const probeResult = await extension.api.probeWorker(workerUrl, TIMEOUTS.STATUS_CHECK);
 
-            if (response.ok) {
-                const data = await response.json();
-                const queueRemaining = data.exec_info?.queue_remaining || 0;
-                return { worker, queueRemaining, online: true };
+            if (probeResult.ok) {
+                return { worker, queueRemaining: probeResult.queueRemaining || 0, online: true };
             }
+            extension.log(`DistributedQueue: Worker ${worker.name} returned ${probeResult.status}`, "debug");
         } catch (error) {
             extension.log(`DistributedQueue: Worker ${worker.name} unreachable: ${error.message}`, "debug");
         }
@@ -324,551 +193,6 @@ async function fetchWorkerQueueStatuses(extension, workers) {
     return results.filter(r => r !== null);
 }
 
-function selectLeastBusyWorker(extension, statuses) {
-    // Find idle workers (queue_remaining == 0)
-    const idleWorkers = statuses.filter(s => s.queueRemaining === 0);
-
-    if (idleWorkers.length > 0) {
-        // Round-robin among idle workers
-        if (!extension._distributedQueueRRIndex) {
-            extension._distributedQueueRRIndex = 0;
-        }
-        const index = extension._distributedQueueRRIndex % idleWorkers.length;
-        extension._distributedQueueRRIndex++;
-        return idleWorkers[index];
-    }
-
-    // No idle workers - pick the one with the shortest queue
-    return statuses.reduce((min, s) => s.queueRemaining < min.queueRemaining ? s : min);
-}
-
-function markSkipDispatch(promptObj) {
-    for (const node of Object.values(promptObj)) {
-        if (node && typeof node === 'object' && node.class_type === NODE_CLASSES.DISTRIBUTED_QUEUE) {
-            node.inputs = node.inputs || {};
-            node.inputs.skip_dispatch = true;
-        }
-    }
-}
-
-export async function prepareApiPromptForParticipant(extension, baseApiPrompt, participantId, options = {}) {
-    let jobApiPrompt = JSON.parse(JSON.stringify(baseApiPrompt));
-    const isMaster = participantId === 'master';
-    const delegateMaster = Boolean(options.delegate_master);
-    
-    // Find all distributed nodes once (before pruning)
-    let collectorNodes = findNodesByClass(jobApiPrompt, NODE_CLASSES.DISTRIBUTED_COLLECTOR);
-    const upscaleNodes = findNodesByClass(jobApiPrompt, NODE_CLASSES.UPSCALE_DISTRIBUTED);
-    let allDistributedNodes = [...collectorNodes, ...upscaleNodes];
-
-    if (isMaster && delegateMaster) {
-        if (upscaleNodes.length > 0) {
-            extension.log("Delegate-only master mode does not support UltimateSDUpscaleDistributed nodes yet. Using full prompt.", "warn");
-        } else if (collectorNodes.length === 0) {
-            extension.log("Delegate-only master mode requested but no collectors found in master prompt. Using full prompt.", "debug");
-        } else {
-            jobApiPrompt = prepareMasterDelegatePrompt(extension, jobApiPrompt, collectorNodes);
-            collectorNodes = findNodesByClass(jobApiPrompt, NODE_CLASSES.DISTRIBUTED_COLLECTOR);
-            allDistributedNodes = [...collectorNodes, ...upscaleNodes];
-        }
-    }
-    
-    // For workers, handle platform-specific path conversion
-    if (!isMaster) {
-        const workerInfo = extension.config.workers.find(w => w.id === participantId);
-        
-        if (workerInfo && workerInfo.host) {
-            // Remote or cloud worker - needs path translation
-            try {
-                const workerUrl = extension.getWorkerUrl(workerInfo);
-                const systemInfo = await getCachedWorkerSystemInfo(workerUrl);
-                const targetSeparator = systemInfo?.platform?.path_separator;
-                
-                if (targetSeparator) {
-                    // Convert paths to match worker's platform
-                    jobApiPrompt = convertPathsForPlatform(jobApiPrompt, targetSeparator);
-                    extension.log(`Converted paths for ${systemInfo.platform.system} worker ${participantId} (separator: '${targetSeparator}')`, "debug");
-                } else {
-                    extension.log(`No path separator found for worker ${participantId}, skipping path conversion`, "debug");
-                }
-            } catch (e) {
-                extension.log(`Failed to get system info for worker ${participantId}: ${e.message}`, "warn");
-                // Continue without path conversion
-            }
-        }
-        
-        // Prune the workflow to only include distributed node dependencies
-        if (allDistributedNodes.length > 0) {
-            jobApiPrompt = pruneWorkflowForWorker(extension, jobApiPrompt, allDistributedNodes);
-        }
-    }
-    
-    // Handle image references for remote workers
-    if (!isMaster && options.enabled_worker_ids) {
-        // Check if this is a remote worker
-        const workerId = participantId;
-        const workerInfo = extension.config.workers.find(w => w.id === workerId);
-        const isRemote = workerInfo && workerInfo.host;
-        
-        if (isRemote) {
-            // Find all image/video references in the pruned workflow
-            const imageReferences = findImageReferences(extension, jobApiPrompt);
-            if (imageReferences.size > 0) {
-                extension.log(`Found ${imageReferences.size} media references (images/videos) for remote worker ${workerId}`, "debug");
-                // Store image references for later processing
-                options._imageReferences = imageReferences;
-            }
-        }
-    }
-    
-    // Handle Distributed seed nodes
-    const distributorNodes = findNodesByClass(jobApiPrompt, NODE_CLASSES.DISTRIBUTED_SEED);
-    if (distributorNodes.length > 0) {
-        extension.log(`Found ${distributorNodes.length} seed node(s)`, "debug");
-    }
-    
-    for (const seedNode of distributorNodes) {
-        const { inputs } = jobApiPrompt[seedNode.id];
-        inputs.is_worker = !isMaster;
-        if (!isMaster) {
-            const workerIndex = options.enabled_worker_ids.indexOf(participantId);
-            inputs.worker_id = `worker_${workerIndex}`;
-            extension.log(`Set seed node ${seedNode.id} for worker ${workerIndex}`, "debug");
-        }
-    }
-    
-    // Handle Distributed collector nodes (already found above)
-    for (const collector of collectorNodes) {
-        const { inputs } = jobApiPrompt[collector.id];
-        
-        // Check if this collector is downstream from a distributed upscaler
-        const hasUpstreamDistributedUpscaler = hasUpstreamNode(
-            jobApiPrompt, 
-            collector.id, 
-            NODE_CLASSES.UPSCALE_DISTRIBUTED
-        );
-        
-        if (hasUpstreamDistributedUpscaler) {
-            // Set pass_through mode for this collector
-            inputs.pass_through = true;
-            extension.log(`Collector ${collector.id} set to pass-through mode (downstream from distributed upscaler)`, "debug");
-        } else {
-            // Normal collector behavior
-            // Get the unique job ID from the map created for this execution
-            const uniqueJobId = options.job_id_map ? options.job_id_map.get(collector.id) : collector.id;
-            
-            // Use the truly unique ID for this execution
-            inputs.multi_job_id = uniqueJobId;
-            inputs.is_worker = !isMaster;
-            if (isMaster) {
-                inputs.enabled_worker_ids = JSON.stringify(options.enabled_worker_ids || []);
-                if (delegateMaster) {
-                    inputs.delegate_only = true;
-                }
-            } else {
-                inputs.master_url = extension.getMasterUrl();
-                // Also make the worker_job_id unique to prevent potential caching issues
-                inputs.worker_job_id = `${uniqueJobId}_worker_${participantId}`;
-                inputs.worker_id = participantId;
-            }
-        }
-    }
-    
-    // Handle Ultimate SD Upscale Distributed nodes
-    for (const upscaleNode of upscaleNodes) {
-        const { inputs } = jobApiPrompt[upscaleNode.id];
-        
-        // Get the unique job ID from the map
-        const uniqueJobId = options.job_id_map ? options.job_id_map.get(upscaleNode.id) : upscaleNode.id;
-        
-        inputs.multi_job_id = uniqueJobId;
-        inputs.is_worker = !isMaster;
-        
-        if (isMaster) {
-            inputs.enabled_worker_ids = JSON.stringify(options.enabled_worker_ids || []);
-        } else {
-            inputs.master_url = extension.getMasterUrl();
-            inputs.worker_id = participantId;
-            // Workers also need the enabled_worker_ids to calculate tile distribution
-            inputs.enabled_worker_ids = JSON.stringify(options.enabled_worker_ids || []);
-        }
-    }
-    
-    return jobApiPrompt;
-}
-
-export async function prepareDistributedJob(extension, multi_job_id) {
-    try {
-        await extension.api.prepareJob(multi_job_id);
-    } catch (error) {
-        extension.log("Error preparing job: " + error.message, "error");
-        throw error;
-    }
-}
-
-function createNumericIdGenerator(promptObj) {
-    let maxId = 0;
-    for (const key of Object.keys(promptObj)) {
-        const numeric = parseInt(key, 10);
-        if (!Number.isNaN(numeric)) {
-            maxId = Math.max(maxId, numeric);
-        }
-    }
-    return () => {
-        maxId += 1;
-        return String(maxId);
-    };
-}
-
-function prepareMasterDelegatePrompt(extension, apiPrompt, collectorNodes) {
-    const collectorIds = collectorNodes.map((node) => node.id);
-    const nodesToKeep = findCollectorDownstreamNodes(apiPrompt, collectorIds);
-
-    // Ensure collectors themselves are present
-    collectorIds.forEach((id) => nodesToKeep.add(id));
-
-    const prunedPrompt = {};
-    nodesToKeep.forEach((nodeId) => {
-        const nodeData = apiPrompt[nodeId];
-        if (nodeData) {
-            prunedPrompt[nodeId] = JSON.parse(JSON.stringify(nodeData));
-        }
-    });
-
-    const prunedIds = new Set(Object.keys(prunedPrompt));
-
-    // Remove dangling references to trimmed nodes
-    for (const [nodeId, node] of Object.entries(prunedPrompt)) {
-        if (!node.inputs) continue;
-        for (const [inputName, inputValue] of Object.entries(node.inputs)) {
-            if (Array.isArray(inputValue) && inputValue.length === 2) {
-                const sourceId = String(inputValue[0]);
-                if (!prunedIds.has(sourceId)) {
-                    delete node.inputs[inputName];
-                    extension.log(`Removed upstream reference '${inputName}' from node ${nodeId} for delegate-only master prompt.`, "debug");
-                }
-            }
-        }
-    }
-
-    const nextId = createNumericIdGenerator(prunedPrompt);
-    collectorIds.forEach((collectorId) => {
-        const collectorEntry = prunedPrompt[collectorId];
-        if (!collectorEntry) return;
-
-        const placeholderId = nextId();
-        prunedPrompt[placeholderId] = {
-            class_type: NODE_CLASSES.DISTRIBUTED_EMPTY_IMAGE,
-            inputs: {
-                height: 64,
-                width: 64,
-                channels: 3
-            }
-        };
-        prunedIds.add(placeholderId);
-
-        collectorEntry.inputs = collectorEntry.inputs || {};
-        collectorEntry.inputs.images = [placeholderId, 0];
-    });
-
-    extension.log(`Prepared delegate-only master prompt with ${Object.keys(prunedPrompt).length} nodes (kept ${nodesToKeep.size}).`, "debug");
-    return prunedPrompt;
-}
-
-export async function executeJobs(extension, jobs) {
-    let masterPromptId = null;
-    
-    // Pre-load all unique images before dispatching to workers
-    const allImageReferences = new Map();
-    for (const job of jobs) {
-        if (job.type === 'worker' && job.imageReferences) {
-            for (const [filename, info] of job.imageReferences) {
-                allImageReferences.set(filename, info);
-            }
-        }
-    }
-    
-    if (allImageReferences.size > 0) {
-        extension.log(`Pre-loading ${allImageReferences.size} unique media file(s) for all workers`, "debug");
-        await loadImagesForWorker(extension, allImageReferences);
-    }
-    
-    // Now dispatch jobs in parallel
-    const promises = jobs.map(job => {
-        if (job.type === 'master') {
-            return extension.originalQueuePrompt(0, job.promptWrapper).then(result => {
-                masterPromptId = result;
-                return result;
-            });
-        } else {
-            return dispatchToWorker(extension, job.worker, job.prompt, job.workflow, job.imageReferences);
-        }
-    });
-    await Promise.all(promises);
-    
-    // Trigger immediate status check for instant feedback
-    extension.checkAllWorkerStatuses();
-    
-    return masterPromptId || { "prompt_id": "distributed-job-dispatched" };
-}
-
-async function dispatchToWorker(extension, worker, prompt, workflow, imageReferences) {
-    const workerUrl = extension.getWorkerUrl(worker);
-    
-    // Debug logging - always log to console for debugging
-    extension.log(`[Distributed] === Dispatching to ${worker.name} (${worker.id}) ===`, "debug");
-    extension.log('[Distributed] Worker URL: ' + workerUrl, "debug");
-    
-    // Handle image uploads for remote workers
-    if (imageReferences && imageReferences.size > 0) {
-        // Check if this is a local worker (same host as master)
-        const isLocalWorker = workerUrl.includes('127.0.0.1') || workerUrl.includes('localhost');
-        
-        if (isLocalWorker) {
-            extension.log(`[Distributed] Skipping image processing for local worker ${worker.name} (shares filesystem with master)`, "debug");
-        } else {
-            extension.log(`[Distributed] Processing ${imageReferences.size} image(s) for remote worker`, "debug");
-            
-            try {
-                // Load images from master
-                const images = await loadImagesForWorker(extension, imageReferences);
-                
-                // Upload images to worker
-                if (images.length > 0) {
-                    await uploadImagesToWorker(extension, workerUrl, images);
-                    extension.log(`[Distributed] Successfully uploaded ${images.length} image(s) to worker`, "debug");
-                }
-            } catch (error) {
-                extension.log(`Failed to process images for worker ${worker.name}: ${error.message}`, "error");
-                // Continue with workflow execution even if image upload fails
-            }
-        }
-    }
-    
-    const promptToSend = {
-        prompt,
-        workflow,
-        client_id: api.clientId
-    };
-
-    extension.log('[Distributed] Prompt data: ' + JSON.stringify(promptToSend), "debug");
-
-    try {
-        if (extension.config?.settings?.websocket_orchestration) {
-            const wsUrl = buildWorkerWebSocketUrl(workerUrl);
-            await dispatchPromptViaWebSocket(extension, wsUrl, promptToSend);
-            return;
-        }
-        await fetch(`${workerUrl}/prompt`, {
-            method: 'POST',
-            headers: { 'Content-Type': 'application/json' },
-            mode: 'cors',
-            body: JSON.stringify({
-                prompt,
-                extra_data: { extra_pnginfo: { workflow } },
-                client_id: api.clientId
-            })
-        });
-    } catch (e) {
-        extension.log(`Failed to connect to worker ${worker.name} at ${workerUrl}: ${e.message}`, "error");
-    }
-}
-
-function buildWorkerWebSocketUrl(workerUrl) {
-    const url = new URL(workerUrl);
-    url.protocol = url.protocol === 'https:' ? 'wss:' : 'ws:';
-    url.pathname = '/distributed/worker_ws';
-    url.search = '';
-    return url.toString();
-}
-
-async function dispatchPromptViaWebSocket(extension, wsUrl, payload) {
-    const requestId = generateUUID();
-    const message = {
-        type: 'dispatch_prompt',
-        request_id: requestId,
-        prompt: payload.prompt,
-        workflow: payload.workflow,
-        client_id: payload.client_id
-    };
-
-    await new Promise((resolve, reject) => {
-        const ws = new WebSocket(wsUrl);
-        const timeoutId = setTimeout(() => {
-            ws.close();
-            reject(new Error('Websocket dispatch timed out.'));
-        }, 60000);
-
-        ws.addEventListener('open', () => {
-            ws.send(JSON.stringify(message));
-        });
-
-        ws.addEventListener('message', (event) => {
-            try {
-                const data = JSON.parse(event.data || '{}');
-                if (data.type !== 'dispatch_ack' || data.request_id !== requestId) {
-                    return;
-                }
-                clearTimeout(timeoutId);
-                ws.close();
-                if (data.ok) {
-                    resolve();
-                } else {
-                    reject(new Error(data.error || 'Worker rejected websocket dispatch.'));
-                }
-            } catch (error) {
-                clearTimeout(timeoutId);
-                ws.close();
-                reject(error);
-            }
-        });
-
-        ws.addEventListener('error', () => {
-            clearTimeout(timeoutId);
-            ws.close();
-            reject(new Error('Websocket connection failed.'));
-        });
-    });
-    extension.log(`[Distributed] Websocket dispatch succeeded for ${wsUrl}`, "debug");
-}
-
-export async function loadImagesForWorker(extension, imageReferences) {
-    const images = [];
-    
-    // Use a cache to avoid loading the same image multiple times
-    if (!extension._imageCache) {
-        extension._imageCache = new Map();
-    }
-    
-    for (const [filename, info] of imageReferences) {
-        try {
-            // Check cache first
-            if (extension._imageCache.has(filename)) {
-                images.push(extension._imageCache.get(filename));
-                extension.log(`Using cached image: ${filename}`, "debug");
-                continue;
-            }
-            
-            // Limit cache size
-            if (extension._imageCache.size >= 10) {
-                const oldestKey = extension._imageCache.keys().next().value;
-                extension._imageCache.delete(oldestKey);
-                extension.log(`Evicted oldest cache entry: ${oldestKey} (cache limit reached)`, "debug");
-            }
-            
-            // Load image from master's filesystem via API
-            try {
-                const data = await extension.api.loadImage(filename);
-                const imageData = {
-                    name: filename,
-                    image: data.image_data,
-                    hash: data.hash  // Include hash from the response
-                };
-                images.push(imageData);
-                
-                // Cache the image for future use
-                extension._imageCache.set(filename, imageData);
-                extension.log(`Loaded and cached image: ${filename}`, "debug");
-            } catch (loadError) {
-                extension.log(`Failed to load image ${filename}: ${loadError.message}`, "error");
-                throw loadError;
-            }
-        } catch (error) {
-            extension.log(`Error loading image ${filename}: ${error.message}`, "error");
-        }
-    }
-    
-    // Clear cache after a reasonable time to avoid memory issues
-    setTimeout(() => {
-        if (extension._imageCache && extension._imageCache.size > 0) {
-            extension.log(`Clearing image cache (${extension._imageCache.size} images)`, "debug");
-            extension._imageCache.clear();
-        }
-    }, TIMEOUTS.IMAGE_CACHE_CLEAR); // Clear after 30 seconds
-    
-    return images;
-}
-
-export async function uploadImagesToWorker(extension, workerUrl, images) {
-    // Upload images to worker's ComfyUI instance
-    for (const imageData of images) {
-        // Check if file already exists with matching hash
-        if (imageData.hash) {
-            try {
-                const checkResponse = await fetch(`${workerUrl}/distributed/check_file`, {
-                    method: 'POST',
-                    headers: { 'Content-Type': 'application/json' },
-                    mode: 'cors',
-                    body: JSON.stringify({ 
-                        filename: imageData.name, 
-                        hash: imageData.hash 
-                    })
-                });
-                
-                if (checkResponse.ok) {
-                    const result = await checkResponse.json();
-                    if (result.exists && result.hash_matches) {
-                        extension.log(`File ${imageData.name} already exists on worker with matching hash, skipping upload`, "debug");
-                        continue;
-                    }
-                }
-            } catch (error) {
-                // If check fails, proceed with upload
-                extension.log(`Failed to check file existence for ${imageData.name}: ${error.message}`, "debug");
-            }
-        }
-        
-        const formData = new FormData();
-        
-        // Detect MIME type from base64 header (supports both image and video)
-        const mimeMatch = imageData.image.match(/^data:((?:image|video)\/\w+);base64,/);
-        const mimeType = mimeMatch ? mimeMatch[1] : 'image/png'; // Default to PNG if not detected
-        
-        // Convert base64 to blob
-        const base64Data = imageData.image.replace(/^data:(?:image|video)\/\w+;base64,/, '');
-        const byteCharacters = atob(base64Data);
-        const byteNumbers = new Array(byteCharacters.length);
-        for (let i = 0; i < byteCharacters.length; i++) {
-            byteNumbers[i] = byteCharacters.charCodeAt(i);
-        }
-        const byteArray = new Uint8Array(byteNumbers);
-        const blob = new Blob([byteArray], { type: mimeType });
-        
-        // Use original filename without heavy cleaning
-        let cleanName = imageData.name;
-        let subfolder = '';
-        
-        // Extract subfolder if present (handle both slash styles)
-        if (cleanName.includes('/') || cleanName.includes('\\')) {
-            const parts = cleanName.replace(/\\/g, '/').split('/');
-            subfolder = parts.slice(0, -1).join('/');
-            cleanName = parts[parts.length - 1];
-        }
-        
-        formData.append('image', blob, cleanName);
-        formData.append('type', 'input');
-        formData.append('subfolder', subfolder);
-        formData.append('overwrite', 'true');
-        
-        try {
-            const response = await fetch(`${workerUrl}/upload/image`, {
-                method: 'POST',
-                mode: 'cors',
-                body: formData
-            });
-            
-            if (!response.ok) {
-                throw new Error(`Upload failed: ${response.statusText}`);
-            }
-            
-            extension.log(`Uploaded image to worker: ${imageData.name} -> ${subfolder}/${cleanName}`, "debug");
-        } catch (error) {
-            extension.log(`Failed to upload ${imageData.name}: ${error.message}`, "error");
-            // Continue with other images
-        }
-    }
-}
-
 export async function performPreflightCheck(extension, workers) {
     if (workers.length === 0) return [];
     
@@ -876,23 +200,18 @@ export async function performPreflightCheck(extension, workers) {
     const startTime = Date.now();
     
     const checkPromises = workers.map(async (worker) => {
-        const url = extension.getWorkerUrl(worker, '/prompt');
+        const workerUrl = getWorkerUrl(extension, worker);
         
-        extension.log(`Pre-flight checking ${worker.name} at: ${url}`, "debug");
+        extension.log(`Pre-flight checking ${worker.name} at: ${workerUrl}`, "debug");
         
         try {
-            const response = await fetch(url, {
-                method: 'GET',
-                mode: 'cors',
-                cache: 'no-store',
-                signal: AbortSignal.timeout(TIMEOUTS.STATUS_CHECK)
-            });
+            const probeResult = await extension.api.probeWorker(workerUrl, TIMEOUTS.STATUS_CHECK);
 
-            if (response.ok) {
+            if (probeResult.ok) {
                 extension.log(`Worker ${worker.name} is active`, "debug");
                 return { worker, active: true };
             } else {
-                extension.log(`Worker ${worker.name} returned ${response.status}`, "debug");
+                extension.log(`Worker ${worker.name} returned ${probeResult.status}`, "debug");
                 return { worker, active: false };
             }
         } catch (error) {
@@ -916,8 +235,8 @@ export async function performPreflightCheck(extension, workers) {
         const statusDot = document.getElementById(`status-${r.worker.id}`);
         if (statusDot) {
             // Remove pulsing animation once status is determined
-            statusDot.classList.remove('status-pulsing');
-            statusDot.style.backgroundColor = STATUS_COLORS.OFFLINE_RED;
+            statusDot.classList.remove('status-pulsing', 'worker-status--online', 'worker-status--unknown');
+            statusDot.classList.add('worker-status--offline');
             statusDot.title = "Offline - Cannot connect";
         }
     });

web/main.js

@@ -7,17 +7,16 @@ import { createApiClient } from './apiClient.js';
 import { renderSidebarContent, updateWorkerCard } from './sidebarRenderer.js';
 import { handleInterruptWorkers, handleClearMemory } from './workerUtils.js';
 import { setupInterceptor } from './executionUtils.js';
-import { BUTTON_STYLES, PULSE_ANIMATION_CSS, TIMEOUTS, STATUS_COLORS } from './constants.js';
+import { PULSE_ANIMATION_CSS, TIMEOUTS, STATUS_COLORS } from './constants.js';
 import { updateTunnelUIElements, refreshTunnelStatus, handleTunnelToggle } from './tunnelManager.js';
-import { checkAllWorkerStatuses, getWorkerUrl, checkWorkerStatus, launchWorker, stopWorker, loadManagedWorkers, updateWorkerControls, viewWorkerLog, startLogAutoRefresh, stopLogAutoRefresh, toggleWorkerExpanded } from './workerLifecycle.js';
-import { isRemoteWorker, saveWorkerSettings, cancelWorkerSettings, deleteWorker, addNewWorker } from './workerSettings.js';
+import { checkAllWorkerStatuses, checkWorkerStatus, loadManagedWorkers } from './workerLifecycle.js';
 import { detectMasterIP } from './masterDetection.js';
+import { parseHostInput, getMasterUrl as buildMasterUrl } from './urlUtils.js';
 
 class DistributedExtension {
     constructor() {
         this.config = null;
         this.originalQueuePrompt = api.queuePrompt.bind(api);
-        this.statusCheckInterval = null;
         this.logAutoRefreshInterval = null;
         this.masterSettingsExpanded = false;
         this.app = app; // Store app reference for toast notifications
@@ -47,7 +46,7 @@ class DistributedExtension {
             this.setupInterceptor();
             // Don't start polling until panel opens
             // this.startStatusChecking();
-            this.loadManagedWorkers();
+            loadManagedWorkers(this);
             // Detect master IP after everything is set up
             this.detectMasterIP();
         });
@@ -71,6 +70,21 @@ class DistributedExtension {
             style.textContent = PULSE_ANIMATION_CSS;
             document.head.appendChild(style);
         }
+
+        const fileStyleId = 'distributed-file-styles';
+        if (!document.getElementById(fileStyleId)) {
+            const style = document.createElement('style');
+            style.id = fileStyleId;
+            fetch(new URL('./distributed.css', import.meta.url))
+                .then((response) => response.text())
+                .then((cssText) => {
+                    style.textContent = cssText;
+                })
+                .catch((error) => {
+                    this.log(`Failed to load distributed.css: ${error.message}`, "error");
+                });
+            document.head.appendChild(style);
+        }
     }
 
     // --- State & Config Management (Single Source of Truth) ---
@@ -150,21 +164,7 @@ 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 };
-        }
+        return parseHostInput(value);
     }
 
     updateTunnelUIElements(isRunning, isStarting) {
@@ -271,7 +271,7 @@ class DistributedExtension {
     onPanelOpen() {
         this.log("Panel opened - starting status polling", "debug");
         if (!this.statusCheckTimeout) {
-            this.checkAllWorkerStatuses();
+            checkAllWorkerStatuses(this);
         }
     }
     
@@ -299,63 +299,14 @@ class DistributedExtension {
         setupInterceptor(this);
     }
 
-    // TODO: Gradually migrate external call sites to direct module functions, then remove these wrappers.
-    async checkAllWorkerStatuses() {
-        return checkAllWorkerStatuses(this);
-    }
-
-    // Helper to build worker URL
-    getWorkerUrl(worker, endpoint = '') {
-        return getWorkerUrl(this, worker, endpoint);
-    }
-
     updateWorkerCard(workerId, newStatus) {
         return updateWorkerCard(this, workerId, newStatus);
     }
 
-    async launchWorker(workerId) {
-        return launchWorker(this, workerId);
-    }
-
-    async stopWorker(workerId) {
-        return stopWorker(this, workerId);
-    }
-
-    // Generic async button action handler
-    async handleAsyncButtonAction(button, action, successText, errorText, resetDelay = TIMEOUTS.BUTTON_RESET) {
-        const originalText = button.textContent;
-        const originalStyle = button.style.cssText;
-        button.disabled = true;
-        
-        try {
-            await action();
-            button.textContent = successText;
-            button.style.cssText = originalStyle;
-            button.style.backgroundColor = BUTTON_STYLES.success;
-            return true;
-        } catch (error) {
-            button.textContent = errorText || `Error: ${error.message}`;
-            button.style.cssText = originalStyle;
-            button.style.backgroundColor = BUTTON_STYLES.error;
-            throw error;
-        } finally {
-            setTimeout(() => {
-                button.textContent = originalText;
-                button.style.cssText = originalStyle;
-                button.disabled = false;
-            }, resetDelay);
-        }
-    }
-
     /**
      * Cleanup method to stop intervals and listeners
      */
     cleanup() {
-        if (this.statusCheckInterval) {
-            clearInterval(this.statusCheckInterval);
-            this.statusCheckInterval = null;
-        }
-        
         if (this.logAutoRefreshInterval) {
             clearInterval(this.logAutoRefreshInterval);
             this.logAutoRefreshInterval = null;
@@ -369,94 +320,14 @@ class DistributedExtension {
         this.log("Cleaned up intervals", "debug");
     }
 
-    async loadManagedWorkers() {
-        return loadManagedWorkers(this);
-    }
-
-    updateWorkerControls(workerId) {
-        return updateWorkerControls(this, workerId);
-    }
-
-    async viewWorkerLog(workerId) {
-        return viewWorkerLog(this, workerId);
-    }
-
-    isRemoteWorker(worker) {
-        return isRemoteWorker(this, worker);
-    }
-
     getMasterUrl() {
-        // Always use the detected/configured master IP for consistency
-        if (this.config?.master?.host) {
-            const configuredHost = this.config.master.host;
-            
-            // If the configured host already includes protocol, use as-is
-            if (configuredHost.startsWith('http://') || configuredHost.startsWith('https://')) {
-                return configuredHost;
-            }
-            
-            // For domain names (not IPs), default to HTTPS
-            const isIP = /^(\d{1,3}\.){3}\d{1,3}$/.test(configuredHost);
-            const isLocalhost = configuredHost === 'localhost' || configuredHost === '127.0.0.1';
-            
-            if (!isIP && !isLocalhost && configuredHost.includes('.')) {
-                // It's a domain name, use HTTPS
-                return `https://${configuredHost}`;
-            } else {
-                // For IPs and localhost, use current access method
-                const port = window.location.port || (window.location.protocol === 'https:' ? '443' : '80');
-                if ((window.location.protocol === 'https:' && port === '443') || 
-                    (window.location.protocol === 'http:' && port === '80')) {
-                    return `${window.location.protocol}//${configuredHost}`;
-                }
-                return `${window.location.protocol}//${configuredHost}:${port}`;
-            }
-        }
-        
-        // If no master IP is set but we're on a network address, use it
-        const hostname = window.location.hostname;
-        if (hostname !== 'localhost' && hostname !== '127.0.0.1') {
-            return window.location.origin;
-        }
-        
-        // Fallback warning - this won't work for remote workers
-        this.log("No master host configured - remote workers won't be able to connect. " +
-                     "Master host should be auto-detected on startup.", "debug");
-        return window.location.origin;
+        return buildMasterUrl(this.config, window.location, (message, level) => this.log(message, level));
     }
 
     async detectMasterIP() {
         return detectMasterIP(this);
     }
 
-    async saveWorkerSettings(workerId) {
-        return saveWorkerSettings(this, workerId);
-    }
-
-    cancelWorkerSettings(workerId) {
-        return cancelWorkerSettings(this, workerId);
-    }
-
-    async deleteWorker(workerId) {
-        return deleteWorker(this, workerId);
-    }
-
-    async addNewWorker() {
-        return addNewWorker(this);
-    }
-
-    startLogAutoRefresh(workerId) {
-        return startLogAutoRefresh(this, workerId);
-    }
-
-    stopLogAutoRefresh() {
-        return stopLogAutoRefresh(this);
-    }
-
-    toggleWorkerExpanded(workerId) {
-        return toggleWorkerExpanded(this, workerId);
-    }
-
     _handleInterruptWorkers(button) {
         return handleInterruptWorkers(this, button);
     }

web/sidebar/actionsSection.js

@@ -0,0 +1,30 @@
+import { BUTTON_STYLES } from "../constants.js";
+
+export function renderActionsSection(extension) {
+    const actionsSection = document.createElement("div");
+    actionsSection.style.cssText = "padding-top: 10px; margin-bottom: 15px; border-top: 1px solid #444;";
+
+    const buttonRow = document.createElement("div");
+    buttonRow.style.cssText = "display: flex; gap: 8px;";
+
+    const clearMemButton = extension.ui.createButtonHelper(
+        "Clear Worker VRAM",
+        (event) => extension._handleClearMemory(event.target),
+        BUTTON_STYLES.clearMemory
+    );
+    clearMemButton.title = "Clear VRAM on all enabled worker GPUs (not master)";
+    clearMemButton.style.cssText = BUTTON_STYLES.base + " flex: 1;" + BUTTON_STYLES.clearMemory;
+
+    const interruptButton = extension.ui.createButtonHelper(
+        "Interrupt Workers",
+        (event) => extension._handleInterruptWorkers(event.target),
+        BUTTON_STYLES.interrupt
+    );
+    interruptButton.title = "Cancel/interrupt execution on all enabled worker GPUs";
+    interruptButton.style.cssText = BUTTON_STYLES.base + " flex: 1;" + BUTTON_STYLES.interrupt;
+
+    buttonRow.appendChild(clearMemButton);
+    buttonRow.appendChild(interruptButton);
+    actionsSection.appendChild(buttonRow);
+    return actionsSection;
+}

web/sidebar/settingsSection.js

@@ -0,0 +1,121 @@
+import { createCheckboxSetting, createNumberSetting } from "../ui/buttonHelpers.js";
+
+export function renderSettingsSection(extension) {
+    const settingsSection = document.createElement("div");
+    settingsSection.style.cssText = "border-top: 1px solid #444; margin-bottom: 10px;";
+
+    const settingsToggleArea = document.createElement("div");
+    settingsToggleArea.style.cssText = "padding: 16.5px 0; cursor: pointer; user-select: none;";
+
+    const settingsHeader = document.createElement("div");
+    settingsHeader.style.cssText = "display: flex; align-items: center; justify-content: space-between;";
+
+    const workerSettingsTitle = document.createElement("h4");
+    workerSettingsTitle.textContent = "Settings";
+    workerSettingsTitle.style.cssText = "margin: 0; font-size: 14px;";
+
+    const workerSettingsToggle = document.createElement("span");
+    workerSettingsToggle.textContent = "▶";
+    workerSettingsToggle.style.cssText = "font-size: 12px; color: #888; transition: all 0.2s ease;";
+
+    settingsHeader.appendChild(workerSettingsTitle);
+    settingsHeader.appendChild(workerSettingsToggle);
+    settingsToggleArea.appendChild(settingsHeader);
+
+    settingsToggleArea.onmouseover = () => {
+        workerSettingsToggle.style.color = "#fff";
+    };
+    settingsToggleArea.onmouseout = () => {
+        workerSettingsToggle.style.color = "#888";
+    };
+
+    const settingsSeparator = document.createElement("div");
+    settingsSeparator.style.cssText = "border-bottom: 1px solid #444; margin: 0;";
+
+    const settingsContent = document.createElement("div");
+    settingsContent.style.cssText =
+        "max-height: 0; overflow: hidden; opacity: 0; transition: max-height 0.3s ease, opacity 0.3s ease;";
+
+    const settingsDiv = document.createElement("div");
+    settingsDiv.style.cssText =
+        "display: grid; grid-template-columns: 1fr auto; row-gap: 10px; column-gap: 10px; padding-top: 10px; align-items: center;";
+
+    let settingsExpanded = false;
+    settingsToggleArea.onclick = () => {
+        settingsExpanded = !settingsExpanded;
+        if (settingsExpanded) {
+            settingsContent.style.maxHeight = "200px";
+            settingsContent.style.opacity = "1";
+            workerSettingsToggle.style.transform = "rotate(90deg)";
+            settingsSeparator.style.display = "none";
+        } else {
+            settingsContent.style.maxHeight = "0";
+            settingsContent.style.opacity = "0";
+            workerSettingsToggle.style.transform = "rotate(0deg)";
+            settingsSeparator.style.display = "block";
+        }
+    };
+
+    const generalLabel = document.createElement("div");
+    generalLabel.textContent = "GENERAL";
+    generalLabel.style.cssText =
+        "grid-column: 1 / span 2; font-size: 11px; color: #888; letter-spacing: 0.06em; padding-top: 2px;";
+
+    const timeoutsLabel = document.createElement("div");
+    timeoutsLabel.textContent = "TIMEOUTS";
+    timeoutsLabel.style.cssText =
+        "grid-column: 1 / span 2; font-size: 11px; color: #888; letter-spacing: 0.06em; padding-top: 4px;";
+
+    settingsDiv.appendChild(generalLabel);
+    settingsDiv.appendChild(
+        createCheckboxSetting(
+            "setting-debug",
+            "Debug Mode",
+            "Enable verbose logging in the browser console.",
+            extension.config?.settings?.debug || false,
+            (event) => extension._updateSetting("debug", event.target.checked)
+        )
+    );
+    settingsDiv.appendChild(
+        createCheckboxSetting(
+            "setting-auto-launch",
+            "Auto-launch Local Workers on Startup",
+            "Start local worker processes automatically when the master starts.",
+            extension.config?.settings?.auto_launch_workers || false,
+            (event) => extension._updateSetting("auto_launch_workers", event.target.checked)
+        )
+    );
+    settingsDiv.appendChild(
+        createCheckboxSetting(
+            "setting-stop-on-exit",
+            "Stop Local Workers on Master Exit",
+            "Stop local worker processes automatically when the master exits.",
+            extension.config?.settings?.stop_workers_on_master_exit !== false,
+            (event) => extension._updateSetting("stop_workers_on_master_exit", event.target.checked)
+        )
+    );
+    settingsDiv.appendChild(timeoutsLabel);
+    settingsDiv.appendChild(
+        createNumberSetting(
+            "setting-worker-timeout",
+            "Worker Timeout",
+            "Seconds without a heartbeat before a worker is considered timed out. Default 60.",
+            extension.config?.settings?.worker_timeout_seconds ?? 60,
+            10,
+            1,
+            (event) => {
+                const value = parseInt(event.target.value, 10);
+                if (!Number.isFinite(value) || value <= 0) {
+                    return;
+                }
+                extension._updateSetting("worker_timeout_seconds", value);
+            }
+        )
+    );
+
+    settingsContent.appendChild(settingsDiv);
+    settingsSection.appendChild(settingsToggleArea);
+    settingsSection.appendChild(settingsSeparator);
+    settingsSection.appendChild(settingsContent);
+    return settingsSection;
+}

web/sidebar/workersSection.js

@@ -0,0 +1,36 @@
+import { addNewWorker } from "../workerSettings.js";
+
+export function renderWorkersSection(extension) {
+    const workersSection = document.createElement("div");
+    workersSection.style.cssText = "flex: 1; overflow-y: auto; margin-bottom: 15px;";
+
+    const workersList = document.createElement("div");
+    const workers = extension.config?.workers || [];
+
+    if (workers.length === 0) {
+        const blueprintDiv = extension.ui.renderEntityCard(
+            "blueprint",
+            { onClick: () => addNewWorker(extension) },
+            extension
+        );
+        workersList.appendChild(blueprintDiv);
+    }
+
+    workers.forEach((worker) => {
+        const workerCard = extension.ui.renderEntityCard("worker", worker, extension);
+        workersList.appendChild(workerCard);
+    });
+
+    workersSection.appendChild(workersList);
+
+    if (workers.length > 0) {
+        const addWorkerDiv = extension.ui.renderEntityCard(
+            "add",
+            { onClick: () => addNewWorker(extension) },
+            extension
+        );
+        workersSection.appendChild(addWorkerDiv);
+    }
+
+    return workersSection;
+}

web/sidebarRenderer.js

@@ -1,5 +1,8 @@
-import { BUTTON_STYLES, STATUS_COLORS } from './constants.js';
-import { createCheckboxSetting, createNumberSetting } from './ui/buttonHelpers.js';
+import { STATUS_COLORS } from './constants.js';
+import { checkAllWorkerStatuses, loadManagedWorkers, updateWorkerControls } from './workerLifecycle.js';
+import { renderActionsSection } from './sidebar/actionsSection.js';
+import { renderSettingsSection } from './sidebar/settingsSection.js';
+import { renderWorkersSection } from './sidebar/workersSection.js';
 
 export function updateWorkerCard(extension, workerId, newStatus = {}) {
     const card = document.querySelector(`[data-worker-id="${workerId}"]`);
@@ -21,7 +24,7 @@ export function updateWorkerCard(extension, workerId, newStatus = {}) {
         extension.ui.updateStatusDot(workerId, STATUS_COLORS.OFFLINE_RED, "Offline - Cannot connect", false);
     }
 
-    extension.updateWorkerControls(workerId);
+    updateWorkerControls(extension, workerId);
     return true;
 }
 
@@ -60,11 +63,7 @@ export async function renderSidebarContent(extension, el) {
         document.head.appendChild(style);
         loadingDiv.querySelector('svg').style.animation = 'rotate 1s linear infinite';
         // Preload data outside render
-        await Promise.all([
-            extension.loadConfig(),
-            extension.loadManagedWorkers(),
-            extension.refreshTunnelStatus(),
-        ]);
+        await Promise.all([extension.loadConfig(), loadManagedWorkers(extension), extension.refreshTunnelStatus()]);
         extension.tunnelElements = {};
         el.innerHTML = '';
         // Create toolbar header to match ComfyUI style
@@ -99,145 +98,12 @@ export async function renderSidebarContent(extension, el) {
         // Master Node Section
         const masterDiv = extension.ui.renderEntityCard('master', extension.config?.master, extension);
         container.appendChild(masterDiv);
-        // Workers Section (no heading)
-        const gpuSection = document.createElement("div");
-        gpuSection.style.cssText = "flex: 1; overflow-y: auto; margin-bottom: 15px;";
-        const gpuList = document.createElement("div");
-        const workers = extension.config?.workers || [];
-        // If no workers exist, show a full blueprint placeholder first
-        if (workers.length === 0) {
-            const blueprintDiv = extension.ui.renderEntityCard('blueprint', { onClick: () => extension.addNewWorker() }, extension);
-            gpuList.appendChild(blueprintDiv);
-        }
-        // Show existing workers
-        workers.forEach(worker => {
-            const gpuDiv = extension.ui.renderEntityCard('worker', worker, extension);
-            gpuList.appendChild(gpuDiv);
-        });
-        gpuSection.appendChild(gpuList);
-        // Only show the minimal "Add Worker" box if there are existing workers
-        if (workers.length > 0) {
-            const addWorkerDiv = extension.ui.renderEntityCard('add', { onClick: () => extension.addNewWorker() }, extension);
-            gpuSection.appendChild(addWorkerDiv);
-        }
-        container.appendChild(gpuSection);
-        const actionsSection = document.createElement("div");
-        actionsSection.style.cssText = "padding-top: 10px; margin-bottom: 15px; border-top: 1px solid #444;";
-        // Create a row for both buttons
-        const buttonRow = document.createElement("div");
-        buttonRow.style.cssText = "display: flex; gap: 8px;";
-        const clearMemButton = extension.ui.createButtonHelper(
-            "Clear Worker VRAM",
-            (e) => extension._handleClearMemory(e.target),
-            BUTTON_STYLES.clearMemory
-        );
-        clearMemButton.title = "Clear VRAM on all enabled worker GPUs (not master)";
-        clearMemButton.style.cssText = BUTTON_STYLES.base + " flex: 1;" + BUTTON_STYLES.clearMemory;
-        const interruptButton = extension.ui.createButtonHelper(
-            "Interrupt Workers",
-            (e) => extension._handleInterruptWorkers(e.target),
-            BUTTON_STYLES.interrupt
-        );
-        interruptButton.title = "Cancel/interrupt execution on all enabled worker GPUs";
-        interruptButton.style.cssText = BUTTON_STYLES.base + " flex: 1;" + BUTTON_STYLES.interrupt;
-        buttonRow.appendChild(clearMemButton);
-        buttonRow.appendChild(interruptButton);
-        actionsSection.appendChild(buttonRow);
-        container.appendChild(actionsSection);
-        // Settings section
-        const settingsSection = document.createElement("div");
-        // Top separator only; spacing handled by the clickable toggle area for equal top/bottom spacing
-        settingsSection.style.cssText = "border-top: 1px solid #444; margin-bottom: 10px;";
-        // Settings header with toggle (full-area clickable between separators)
-        const settingsToggleArea = document.createElement("div");
-        // Equal spacing above header (to top separator) and below header (to bottom separator)
-        settingsToggleArea.style.cssText = "padding: 16.5px 0; cursor: pointer; user-select: none;";
-        const settingsHeader = document.createElement("div");
-        settingsHeader.style.cssText = "display: flex; align-items: center; justify-content: space-between;";
-        const workerSettingsTitle = document.createElement("h4");
-        workerSettingsTitle.textContent = "Settings";
-        workerSettingsTitle.style.cssText = "margin: 0; font-size: 14px;";
-        const workerSettingsToggle = document.createElement("span");
-        workerSettingsToggle.textContent = "▶"; // Right arrow when collapsed
-        workerSettingsToggle.style.cssText = "font-size: 12px; color: #888; transition: all 0.2s ease;";
-        settingsHeader.appendChild(workerSettingsTitle);
-        settingsHeader.appendChild(workerSettingsToggle);
-        settingsToggleArea.appendChild(settingsHeader);
-        // Hover effect for toggle area
-        settingsToggleArea.onmouseover = () => { workerSettingsToggle.style.color = "#fff"; };
-        settingsToggleArea.onmouseout = () => { workerSettingsToggle.style.color = "#888"; };
-        // A small separator shown only when collapsed (to make the section boundary obvious)
-        const settingsSeparator = document.createElement("div");
-        // No margin so the bottom spacing is controlled by settingsToggleArea padding-bottom
-        settingsSeparator.style.cssText = "border-bottom: 1px solid #444; margin: 0;";
-        // Collapsible settings content
-        const settingsContent = document.createElement("div");
-        settingsContent.style.cssText = "max-height: 0; overflow: hidden; opacity: 0; transition: max-height 0.3s ease, opacity 0.3s ease;";
-        const settingsDiv = document.createElement("div");
-        settingsDiv.style.cssText = "display: grid; grid-template-columns: 1fr auto; row-gap: 10px; column-gap: 10px; padding-top: 10px; align-items: center;";
-        // Toggle functionality
-        let settingsExpanded = false;
-        settingsToggleArea.onclick = () => {
-            settingsExpanded = !settingsExpanded;
-            if (settingsExpanded) {
-                settingsContent.style.maxHeight = "200px";
-                settingsContent.style.opacity = "1";
-                workerSettingsToggle.style.transform = "rotate(90deg)";
-                settingsSeparator.style.display = "none";
-            } else {
-                settingsContent.style.maxHeight = "0";
-                settingsContent.style.opacity = "0";
-                workerSettingsToggle.style.transform = "rotate(0deg)";
-                settingsSeparator.style.display = "block";
-            }
-        };
-        // Section: General
-        const generalLabel = document.createElement("div");
-        generalLabel.textContent = "GENERAL";
-        generalLabel.style.cssText = "grid-column: 1 / span 2; font-size: 11px; color: #888; letter-spacing: 0.06em; padding-top: 2px;";
-        // Section: Timeouts
-        const timeoutsLabel = document.createElement("div");
-        timeoutsLabel.textContent = "TIMEOUTS";
-        timeoutsLabel.style.cssText = "grid-column: 1 / span 2; font-size: 11px; color: #888; letter-spacing: 0.06em; padding-top: 4px;";
-        settingsDiv.appendChild(generalLabel);
-        settingsDiv.appendChild(createCheckboxSetting(
-            "setting-debug", "Debug Mode",
-            "Enable verbose logging in the browser console.",
-            extension.config?.settings?.debug || false,
-            (e) => extension._updateSetting('debug', e.target.checked)
-        ));
-        settingsDiv.appendChild(createCheckboxSetting(
-            "setting-auto-launch", "Auto-launch Local Workers on Startup",
-            "Start local worker processes automatically when the master starts.",
-            extension.config?.settings?.auto_launch_workers || false,
-            (e) => extension._updateSetting('auto_launch_workers', e.target.checked)
-        ));
-        settingsDiv.appendChild(createCheckboxSetting(
-            "setting-stop-on-exit", "Stop Local Workers on Master Exit",
-            "Stop local worker processes automatically when the master exits.",
-            extension.config?.settings?.stop_workers_on_master_exit !== false,
-            (e) => extension._updateSetting('stop_workers_on_master_exit', e.target.checked)
-        ));
-        settingsDiv.appendChild(timeoutsLabel);
-        settingsDiv.appendChild(createNumberSetting(
-            "setting-worker-timeout", "Worker Timeout",
-            "Seconds without a heartbeat before a worker is considered timed out. Default 60.",
-            extension.config?.settings?.worker_timeout_seconds ?? 60,
-            10, 1,
-            (e) => {
-                const v = parseInt(e.target.value, 10);
-                if (!Number.isFinite(v) || v <= 0) return;
-                extension._updateSetting('worker_timeout_seconds', v);
-            }
-        ));
-        settingsContent.appendChild(settingsDiv);
-        settingsSection.appendChild(settingsToggleArea);
-        settingsSection.appendChild(settingsSeparator);
-        settingsSection.appendChild(settingsContent);
-        container.appendChild(settingsSection);
+        container.appendChild(renderWorkersSection(extension));
+        container.appendChild(renderActionsSection(extension));
+        container.appendChild(renderSettingsSection(extension));
         el.appendChild(container);
         // Start checking worker statuses immediately in parallel
-        setTimeout(() => extension.checkAllWorkerStatuses(), 0);
+        setTimeout(() => checkAllWorkerStatuses(extension), 0);
     } finally {
         // Always reset the rendering flag
         extension._isRendering = false;

web/tests/executionUtils.test.js

@@ -0,0 +1,63 @@
+import { describe, expect, it } from "vitest";
+
+import { NODE_CLASSES } from "../constants.js";
+import { markSkipDispatch, selectLeastBusyWorker } from "../executionDecisionUtils.js";
+import { buildWorkerWebSocketUrl } from "../urlUtils.js";
+
+
+describe("execution decision helpers", () => {
+    it("markSkipDispatch sets skip_dispatch=true on DistributedQueue nodes only", () => {
+        const prompt = {
+            "1": {
+                class_type: NODE_CLASSES.DISTRIBUTED_QUEUE,
+                inputs: { seed: 1 },
+            },
+            "2": {
+                class_type: "KSampler",
+                inputs: { seed: 2 },
+            },
+            "3": {
+                class_type: NODE_CLASSES.DISTRIBUTED_QUEUE,
+                inputs: {},
+            },
+        };
+
+        markSkipDispatch(prompt);
+
+        expect(prompt["1"].inputs.skip_dispatch).toBe(true);
+        expect(prompt["3"].inputs.skip_dispatch).toBe(true);
+        expect(prompt["2"].inputs.skip_dispatch).toBeUndefined();
+    });
+
+    it("selectLeastBusyWorker round-robins among idle workers", () => {
+        const extension = {};
+        const statuses = [
+            { worker: { id: "w1" }, queueRemaining: 0 },
+            { worker: { id: "w2" }, queueRemaining: 0 },
+        ];
+
+        expect(selectLeastBusyWorker(extension, statuses)).toBe(statuses[0]);
+        expect(selectLeastBusyWorker(extension, statuses)).toBe(statuses[1]);
+        expect(selectLeastBusyWorker(extension, statuses)).toBe(statuses[0]);
+    });
+
+    it("selectLeastBusyWorker chooses shortest queue when no idle workers", () => {
+        const extension = {};
+        const statuses = [
+            { worker: { id: "w1" }, queueRemaining: 5 },
+            { worker: { id: "w2" }, queueRemaining: 2 },
+            { worker: { id: "w3" }, queueRemaining: 7 },
+        ];
+
+        expect(selectLeastBusyWorker(extension, statuses)).toBe(statuses[1]);
+    });
+
+    it("buildWorkerWebSocketUrl converts http/https to ws/wss", () => {
+        expect(buildWorkerWebSocketUrl("http://worker.local:8188")).toBe(
+            "ws://worker.local:8188/distributed/worker_ws"
+        );
+        expect(buildWorkerWebSocketUrl("https://worker.example.com")).toBe(
+            "wss://worker.example.com/distributed/worker_ws"
+        );
+    });
+});

web/tests/workerLifecycle.test.js

@@ -0,0 +1,52 @@
+import { afterEach, beforeEach, describe, expect, it } from "vitest";
+
+import { getWorkerUrl } from "../workerLifecycle.js";
+
+
+describe("workerLifecycle URL construction", () => {
+    let originalWindow;
+
+    beforeEach(() => {
+        originalWindow = globalThis.window;
+        globalThis.window = {
+            location: {
+                hostname: "127.0.0.1",
+                protocol: "http:",
+                port: "8190",
+                origin: "http://127.0.0.1:8190",
+            },
+        };
+    });
+
+    afterEach(() => {
+        globalThis.window = originalWindow;
+    });
+
+    it("builds local worker URL with explicit local port", () => {
+        const worker = { id: "w1", port: 8189, type: "local" };
+        expect(getWorkerUrl({}, worker, "/prompt")).toBe("http://127.0.0.1:8189/prompt");
+    });
+
+    it("builds remote worker URL with host:port", () => {
+        const worker = { id: "w2", host: "worker.example.com", port: 9000, type: "remote" };
+        expect(getWorkerUrl({}, worker, "/prompt")).toBe("http://worker.example.com:9000/prompt");
+    });
+
+    it("builds cloud worker URL as https", () => {
+        const worker = { id: "w3", host: "cloud.example.com", port: 443, type: "cloud" };
+        expect(getWorkerUrl({}, worker, "/prompt")).toBe("https://cloud.example.com/prompt");
+    });
+
+    it("rewrites runpod proxy hostname for local worker ports", () => {
+        globalThis.window = {
+            location: {
+                hostname: "podabc.proxy.runpod.net",
+                protocol: "https:",
+                port: "",
+                origin: "https://podabc.proxy.runpod.net",
+            },
+        };
+        const worker = { id: "w4", port: 8189, type: "local" };
+        expect(getWorkerUrl({}, worker, "/prompt")).toBe("https://podabc-8189.proxy.runpod.net/prompt");
+    });
+});

web/ui.js

@@ -3,6 +3,8 @@ import { createButtonHelper as createButtonHelperFn } from './ui/buttonHelpers.j
 import { showCloudflareWarning as showCloudflareWarningFn } from './ui/cloudflareWarning.js';
 import { createWorkerSettingsForm as createWorkerSettingsFormFn } from './ui/settingsForm.js';
 import { renderEntityCard as renderEntityCardFn } from './ui/entityCard.js';
+import { launchWorker, startLogAutoRefresh, stopLogAutoRefresh, stopWorker, updateWorkerControls, viewWorkerLog } from './workerLifecycle.js';
+import { isRemoteWorker } from './workerSettings.js';
 
 const cardConfigs = {
     master: {
@@ -56,7 +58,7 @@ const cardConfigs = {
             id: (data) => `status-${data.id}`
         },
         infoText: (data, extension) => {
-            const isRemote = extension.isRemoteWorker(data);
+            const isRemote = isRemoteWorker(extension, data);
             const isCloud = data.type === 'cloud';
             
             if (isCloud) {
@@ -454,9 +456,9 @@ export class DistributedUI {
         refreshCheckbox.style.cssText = 'cursor: pointer;';
         refreshCheckbox.onchange = (e) => {
             if (e.target.checked) {
-                extension.startLogAutoRefresh(workerId);
+                startLogAutoRefresh(extension, workerId);
             } else {
-                extension.stopLogAutoRefresh();
+                stopLogAutoRefresh(extension);
             }
         };
         
@@ -472,7 +474,7 @@ export class DistributedUI {
         // Close button
         const closeBtn = this.createButton('✕', 
             () => {
-                extension.stopLogAutoRefresh();
+                stopLogAutoRefresh(extension);
                 modal.remove();
             }, 
             'background-color: #c04c4c;');
@@ -528,7 +530,7 @@ export class DistributedUI {
         // Close on background click
         modal.onclick = (e) => {
             if (e.target === modal) {
-                extension.stopLogAutoRefresh();
+                stopLogAutoRefresh(extension);
                 modal.remove();
             }
         };
@@ -536,7 +538,7 @@ export class DistributedUI {
         // Close on Escape key
         const handleEscape = (e) => {
             if (e.key === 'Escape') {
-                extension.stopLogAutoRefresh();
+                stopLogAutoRefresh(extension);
                 modal.remove();
                 document.removeEventListener('keydown', handleEscape);
             }
@@ -546,7 +548,7 @@ export class DistributedUI {
         document.body.appendChild(modal);
         
         // Start auto-refresh
-        extension.startLogAutoRefresh(workerId);
+        startLogAutoRefresh(extension, workerId);
     }
 
     formatFileSize(bytes) {
@@ -736,22 +738,25 @@ export class DistributedUI {
                 controlsWrapper.appendChild(remoteInfo);
             } else {
                 const controls = this.createWorkerControls(data.id, {
-                    launch: () => extension.launchWorker(data.id),
-                    stop: () => extension.stopWorker(data.id),
-                    viewLog: () => extension.viewWorkerLog(data.id)
+                    launch: () => launchWorker(extension, data.id),
+                    stop: () => stopWorker(extension, data.id),
+                    viewLog: () => viewWorkerLog(extension, data.id)
                 });
                 
                 const launchBtn = controls.querySelector(`#launch-${data.id}`);
                 const stopBtn = controls.querySelector(`#stop-${data.id}`);
                 const logBtn = controls.querySelector(`#log-${data.id}`);
                 
-                launchBtn.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.workerControl + BUTTON_STYLES.launch;
+                launchBtn.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.workerControl;
+                launchBtn.classList.add("btn--launch");
                 launchBtn.title = "Launch worker (runs in background with logging)";
                 
-                stopBtn.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.workerControl + BUTTON_STYLES.stop + BUTTON_STYLES.hidden;
+                stopBtn.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.workerControl + BUTTON_STYLES.hidden;
+                stopBtn.classList.add("btn--stop");
                 stopBtn.title = "Stop worker";
                 
-                logBtn.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.workerControl + BUTTON_STYLES.log + BUTTON_STYLES.hidden;
+                logBtn.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.workerControl + BUTTON_STYLES.hidden;
+                logBtn.classList.add("btn--log");
                 
                 while (controls.firstChild) {
                     controlsWrapper.appendChild(controls.firstChild);

web/ui/entityCard.js

@@ -1,9 +1,12 @@
+import { updateWorkerControls, toggleWorkerExpanded } from "../workerLifecycle.js";
+import { isRemoteWorker } from "../workerSettings.js";
+
 export function renderEntityCard(ui, cardConfigs, entityType, data, extension) {
     const config = cardConfigs[entityType] || {};
     const isPlaceholder = entityType === 'blueprint' || entityType === 'add';
     const isWorker = entityType === 'worker';
     const isMaster = entityType === 'master';
-    const isRemote = isWorker && extension.isRemoteWorker(data);
+    const isRemote = isWorker && isRemoteWorker(extension, data);
 
     const cardOptions = {
         onClick: isPlaceholder ? data?.onClick : null,
@@ -52,7 +55,7 @@ export function renderEntityCard(ui, cardConfigs, entityType, data, extension) {
                     arrow.style.transform = "rotate(0deg)";
                 }
             } else {
-                extension.toggleWorkerExpanded(data.id);
+                toggleWorkerExpanded(extension, data.id);
             }
         };
     }
@@ -115,7 +118,7 @@ export function renderEntityCard(ui, cardConfigs, entityType, data, extension) {
     }
 
     if (isWorker && !isRemote) {
-        extension.updateWorkerControls(data.id);
+        updateWorkerControls(extension, data.id);
     }
 
     return card;

web/ui/settingsForm.js

@@ -1,4 +1,5 @@
 import { BUTTON_STYLES } from '../constants.js';
+import { cancelWorkerSettings, deleteWorker, isRemoteWorker, saveWorkerSettings } from '../workerSettings.js';
 
 export function createWorkerSettingsForm(ui, extension, worker) {
     const form = document.createElement("div");
@@ -95,7 +96,7 @@ export function createWorkerSettingsForm(ui, extension, worker) {
 
     const hostGroup = ui.createFormGroup("Host:", worker.host || "", `host-${worker.id}`, "text", "e.g., 192.168.1.100");
     hostGroup.group.id = `host-group-${worker.id}`;
-    hostGroup.group.style.display = (extension.isRemoteWorker(worker) || worker.type === "cloud") ? "flex" : "none";
+    hostGroup.group.style.display = (isRemoteWorker(extension, worker) || worker.type === "cloud") ? "flex" : "none";
     form.appendChild(hostGroup.group);
 
     const portGroup = ui.createFormGroup("Port:", worker.port, `port-${worker.id}`, "number");
@@ -104,21 +105,21 @@ export function createWorkerSettingsForm(ui, extension, worker) {
 
     const cudaGroup = ui.createFormGroup("CUDA Device:", worker.cuda_device || 0, `cuda-${worker.id}`, "number");
     cudaGroup.group.id = `cuda-group-${worker.id}`;
-    cudaGroup.group.style.display = (extension.isRemoteWorker(worker) || worker.type === "cloud") ? "none" : "flex";
+    cudaGroup.group.style.display = (isRemoteWorker(extension, worker) || worker.type === "cloud") ? "none" : "flex";
     form.appendChild(cudaGroup.group);
 
     const argsGroup = ui.createFormGroup("Extra Args:", worker.extra_args || "", `args-${worker.id}`);
     argsGroup.group.id = `args-group-${worker.id}`;
-    argsGroup.group.style.display = (extension.isRemoteWorker(worker) || worker.type === "cloud") ? "none" : "flex";
+    argsGroup.group.style.display = (isRemoteWorker(extension, worker) || worker.type === "cloud") ? "none" : "flex";
     form.appendChild(argsGroup.group);
 
-    const saveBtn = ui.createButton("Save", () => extension.saveWorkerSettings(worker.id), "background-color: #4a7c4a;");
+    const saveBtn = ui.createButton("Save", () => saveWorkerSettings(extension, worker.id), "background-color: #4a7c4a;");
     saveBtn.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.success;
 
-    const cancelBtn = ui.createButton("Cancel", () => extension.cancelWorkerSettings(worker.id), "background-color: #555;");
+    const cancelBtn = ui.createButton("Cancel", () => cancelWorkerSettings(extension, worker.id), "background-color: #555;");
     cancelBtn.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.cancel;
 
-    const deleteBtn = ui.createButton("Delete", () => extension.deleteWorker(worker.id), "background-color: #7c4a4a;");
+    const deleteBtn = ui.createButton("Delete", () => deleteWorker(extension, worker.id), "background-color: #7c4a4a;");
     deleteBtn.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.error + BUTTON_STYLES.marginLeftAuto;
 
     const buttonGroup = ui.createButtonGroup([saveBtn, cancelBtn, deleteBtn], " margin-top: 8px;");
@@ -129,7 +130,7 @@ export function createWorkerSettingsForm(ui, extension, worker) {
     if (worker.type === "cloud") {
         typeSelect.value = "cloud";
         runpodText.style.display = "block";
-    } else if (extension.isRemoteWorker(worker)) {
+    } else if (isRemoteWorker(extension, worker)) {
         typeSelect.value = "remote";
     } else {
         typeSelect.value = "local";

web/urlUtils.js

@@ -20,3 +20,95 @@ export function normalizeWorkerUrl(rawUrl) {
         return withProtocol.replace(/\/$/, "");
     }
 }
+
+export function 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 };
+    }
+}
+
+export function buildWorkerUrl(worker, endpoint = "", windowLocation = window.location) {
+    const parsed = parseHostInput(worker?.host || windowLocation.hostname);
+    const host = parsed.host || windowLocation.hostname;
+    const resolvedPort = parsed.port || worker?.port || 8188;
+
+    const isCloud = worker?.type === "cloud";
+    const isRunpodProxy = host.endsWith(".proxy.runpod.net");
+
+    let finalHost = host;
+    if (!worker?.host && isRunpodProxy) {
+        const match = host.match(/^(.*)\.proxy\.runpod\.net$/);
+        if (match) {
+            finalHost = `${match[1]}-${resolvedPort}.proxy.runpod.net`;
+        } else {
+            console.error(`[Distributed] Failed to parse Runpod proxy host: ${host}`);
+        }
+    }
+
+    const useHttps = isCloud || isRunpodProxy || resolvedPort === 443;
+    const protocol = useHttps ? "https" : "http";
+    const defaultPort = useHttps ? 443 : 80;
+    const needsPort = !isRunpodProxy && resolvedPort !== defaultPort;
+    const portPart = needsPort ? `:${resolvedPort}` : "";
+
+    return normalizeWorkerUrl(`${protocol}://${finalHost}${portPart}${endpoint}`);
+}
+
+export function buildWorkerWebSocketUrl(workerUrl) {
+    const normalized = normalizeWorkerUrl(workerUrl);
+    const wsBase = normalized.replace(/^http:\/\//i, "ws://").replace(/^https:\/\//i, "wss://");
+    return `${wsBase}/distributed/worker_ws`;
+}
+
+export function getMasterUrl(config, windowLocation = window.location, log = null) {
+    const masterHost = config?.master?.host;
+    if (masterHost) {
+        const configuredHost = masterHost;
+
+        // If the configured host already includes protocol, use as-is.
+        if (configuredHost.startsWith("http://") || configuredHost.startsWith("https://")) {
+            return configuredHost;
+        }
+
+        // For domain names (not IPs), default to HTTPS.
+        const isIP = /^(\d{1,3}\.){3}\d{1,3}$/.test(configuredHost);
+        const isLocalhost = configuredHost === "localhost" || configuredHost === "127.0.0.1";
+
+        if (!isIP && !isLocalhost && configuredHost.includes(".")) {
+            return `https://${configuredHost}`;
+        }
+
+        const protocol = windowLocation.protocol || "http:";
+        const port = windowLocation.port || (protocol === "https:" ? "443" : "80");
+        if ((protocol === "https:" && port === "443") || (protocol === "http:" && port === "80")) {
+            return `${protocol}//${configuredHost}`;
+        }
+        return `${protocol}//${configuredHost}:${port}`;
+    }
+
+    const hostname = windowLocation.hostname;
+    if (hostname !== "localhost" && hostname !== "127.0.0.1") {
+        return windowLocation.origin;
+    }
+
+    if (typeof log === "function") {
+        log(
+            "No master host configured - remote workers won't be able to connect. Master host should be auto-detected on startup.",
+            "debug",
+        );
+    }
+    return windowLocation.origin;
+}

web/workerLifecycle.js

@@ -1,8 +1,11 @@
-import { BUTTON_STYLES, TIMEOUTS, STATUS_COLORS, ENDPOINTS } from './constants.js';
-import { normalizeWorkerUrl } from './urlUtils.js';
+import { TIMEOUTS, STATUS_COLORS, ENDPOINTS } from './constants.js';
+import { buildWorkerUrl, normalizeWorkerUrl } from './urlUtils.js';
+import { isRemoteWorker } from './workerSettings.js';
 
 export { normalizeWorkerUrl };
 
+let _statusCheckRunning = false;
+
 function setStatusDotClass(dot, statusClass) {
     if (!dot) {
         return;
@@ -19,46 +22,54 @@ function setStatusDotClass(dot, statusClass) {
     }
 }
 
+function setButtonClass(button, className) {
+    if (!button) {
+        return;
+    }
+    button.classList.remove("btn--stop", "btn--launch", "btn--log", "btn--working", "btn--success", "btn--error");
+    if (className) {
+        button.classList.add(className);
+    }
+}
+
 export async function checkAllWorkerStatuses(extension) {
-    // Don't continue if panel is closed
-    if (!extension.panelElement) {
+    if (_statusCheckRunning || !extension.panelElement) {
         return;
     }
+    _statusCheckRunning = true;
+    let nextInterval = 5000;
 
-    // Create new abort controller for this round of checks
-    extension.statusCheckAbortController = new AbortController();
+    try {
+        // Create a fresh AbortController for this poll cycle.
+        extension.statusCheckAbortController = new AbortController();
 
-    // Check master status
-    checkMasterStatus(extension);
+        await checkMasterStatus(extension);
 
-    if (!extension.config || !extension.config.workers) {
-        return;
-    }
+        if (extension.config?.workers) {
+            await Promise.all(
+                extension.config.workers.map(async (worker) => {
+                    if (worker.enabled || extension.state.isWorkerLaunching(worker.id)) {
+                        await checkWorkerStatus(extension, worker);
+                    }
+                })
+            );
+        }
 
-    for (const worker of extension.config.workers) {
-        // Check status for enabled workers OR workers that are launching
-        if (worker.enabled || extension.state.isWorkerLaunching(worker.id)) {
-            checkWorkerStatus(extension, worker);
+        let isActive = extension.state.getMasterStatus() === "processing";
+        extension.config?.workers?.forEach((worker) => {
+            const workerState = extension.state.getWorker(worker.id);
+            if (workerState.launching || workerState.status?.processing) {
+                isActive = true;
+            }
+        });
+
+        nextInterval = isActive ? 1000 : 5000;
+    } finally {
+        _statusCheckRunning = false;
+        if (extension.panelElement) {
+            extension.statusCheckTimeout = setTimeout(() => checkAllWorkerStatuses(extension), nextInterval);
         }
     }
-
-    // Determine next interval based on current state
-    let isActive = extension.state.getMasterStatus() === "processing"; // Master is busy
-
-    // Check workers for activity
-    extension.config.workers.forEach((worker) => {
-        const ws = extension.state.getWorker(worker.id); // Get worker state
-        if (ws.launching || ws.status?.processing) {
-            // Launching or processing
-            isActive = true;
-        }
-    });
-
-    // Set next delay: 1s if active, 5s if idle
-    const nextInterval = isActive ? 1000 : 5000;
-
-    // Schedule the next check
-    extension.statusCheckTimeout = setTimeout(() => checkAllWorkerStatuses(extension), nextInterval);
 }
 
 export async function checkMasterStatus(extension) {
@@ -111,42 +122,7 @@ export async function checkMasterStatus(extension) {
 
 // Helper to build worker URL
 export function getWorkerUrl(extension, worker, endpoint = "") {
-    const parsed = extension._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";
-
-    // Detect if we're running on Runpod (for local workers on Runpod infrastructure)
-    const isRunpodProxy = host.endsWith(".proxy.runpod.net");
-
-    // For local workers on Runpod, construct the port-specific proxy URL
-    let finalHost = host;
-    if (!worker.host && isRunpodProxy) {
-        const match = host.match(/^(.*)\.proxy\.runpod\.net$/);
-        if (match) {
-            const podId = match[1];
-            const domain = "proxy.runpod.net";
-            finalHost = `${podId}-${worker.port}.${domain}`;
-        } else {
-            // Fallback or log error if no match (shouldn't happen)
-            console.error(`[Distributed] Failed to parse Runpod proxy host: ${host}`);
-        }
-    }
-
-    // Determine protocol: HTTPS for cloud, Runpod proxies, or 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 && resolvedPort !== defaultPort;
-    const portStr = needsPort ? `:${resolvedPort}` : "";
-
-    // NOTE: Runpod host rewriting here intentionally diverges from Python-side heuristics
-    // in utils.network.build_worker_url. Preserve current browser behavior for now.
-    return normalizeWorkerUrl(`${protocol}://${finalHost}${portStr}${endpoint}`);
+    return buildWorkerUrl(worker, endpoint, window.location);
 }
 
 export async function checkWorkerStatus(extension, worker) {
@@ -300,7 +276,7 @@ export async function stopWorker(extension, workerId) {
     if (stopBtn) {
         stopBtn.disabled = true;
         stopBtn.textContent = "Stopping...";
-        stopBtn.style.backgroundColor = "#666";
+        setButtonClass(stopBtn, "btn--working");
     }
 
     try {
@@ -315,7 +291,7 @@ export async function stopWorker(extension, workerId) {
 
             // Flash success feedback
             if (stopBtn) {
-                stopBtn.style.backgroundColor = BUTTON_STYLES.success;
+                setButtonClass(stopBtn, "btn--success");
                 stopBtn.textContent = "Stopped!";
                 setTimeout(() => {
                     updateWorkerControls(extension, workerId);
@@ -329,7 +305,7 @@ export async function stopWorker(extension, workerId) {
 
             // Flash error feedback
             if (stopBtn) {
-                stopBtn.style.backgroundColor = BUTTON_STYLES.error;
+                setButtonClass(stopBtn, "btn--error");
                 stopBtn.textContent = result.message.includes("already stopped") ? "Not Running" : "Failed";
 
                 // If already stopped, update status immediately
@@ -348,7 +324,7 @@ export async function stopWorker(extension, workerId) {
 
         // Reset button on error
         if (stopBtn) {
-            stopBtn.style.backgroundColor = BUTTON_STYLES.error;
+            setButtonClass(stopBtn, "btn--error");
             stopBtn.textContent = "Error";
             setTimeout(() => {
                 updateWorkerControls(extension, workerId);
@@ -403,7 +379,7 @@ export function updateWorkerControls(extension, workerId) {
     }
 
     // Skip button updates for remote workers
-    if (extension.isRemoteWorker(worker)) {
+    if (isRemoteWorker(extension, worker)) {
         return;
     }
 
@@ -419,6 +395,7 @@ export function updateWorkerControls(extension, workerId) {
     // Show log button immediately if we have log file info (even if worker is still starting)
     if (managedInfo?.log_file && logBtn) {
         logBtn.style.display = "";
+        setButtonClass(logBtn, "btn--log");
     } else if (logBtn && !managedInfo) {
         logBtn.style.display = "none";
     }
@@ -432,7 +409,7 @@ export function updateWorkerControls(extension, workerId) {
             stopBtn.style.display = "";
             stopBtn.disabled = false;
             stopBtn.textContent = "Stop";
-            stopBtn.style.backgroundColor = "#7c4a4a"; // Red when enabled
+            setButtonClass(stopBtn, "btn--stop");
         } else {
             // Hide stop button for workers launched outside UI
             stopBtn.style.display = "none";
@@ -442,7 +419,7 @@ export function updateWorkerControls(extension, workerId) {
         launchBtn.style.display = ""; // Show launch button
         launchBtn.disabled = false;
         launchBtn.textContent = "Launch";
-        launchBtn.style.backgroundColor = "#4a7c4a"; // Always green
+        setButtonClass(launchBtn, "btn--launch");
 
         stopBtn.style.display = "none"; // Hide stop button when not running
     }
@@ -460,7 +437,7 @@ export async function viewWorkerLog(extension, workerId) {
     if (logBtn) {
         logBtn.disabled = true;
         logBtn.textContent = "Loading...";
-        logBtn.style.backgroundColor = "#666";
+        setButtonClass(logBtn, "btn--working");
     }
 
     try {
@@ -474,7 +451,7 @@ export async function viewWorkerLog(extension, workerId) {
         if (logBtn) {
             logBtn.disabled = false;
             logBtn.textContent = "View Log";
-            logBtn.style.backgroundColor = "#685434"; // Keep the yellow color
+            setButtonClass(logBtn, "btn--log");
         }
     } catch (error) {
         extension.log("Error viewing log: " + error.message, "error");
@@ -487,12 +464,12 @@ export async function viewWorkerLog(extension, workerId) {
 
         // Flash error and restore button
         if (logBtn) {
-            logBtn.style.backgroundColor = BUTTON_STYLES.error;
+            setButtonClass(logBtn, "btn--error");
             logBtn.textContent = "Error";
             setTimeout(() => {
                 logBtn.disabled = false;
                 logBtn.textContent = "View Log";
-                logBtn.style.backgroundColor = "#685434"; // Keep the yellow color
+                setButtonClass(logBtn, "btn--log");
             }, TIMEOUTS.FLASH_LONG);
         }
     }

web/workerSettings.js

@@ -1,5 +1,7 @@
 import { renderSidebarContent } from './sidebarRenderer.js';
 import { generateUUID } from './constants.js';
+import { parseHostInput } from './urlUtils.js';
+import { toggleWorkerExpanded } from './workerLifecycle.js';
 
 export function isRemoteWorker(extension, worker) {
     // Check if explicitly marked as cloud worker
@@ -7,7 +9,7 @@ export function isRemoteWorker(extension, worker) {
         return true;
     }
     // Otherwise check by host (backward compatibility)
-    const parsed = extension._parseHostInput(worker.host || window.location.hostname);
+    const parsed = 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;
 }
@@ -28,7 +30,7 @@ export async function saveWorkerSettings(extension, workerId) {
     const isRemote = workerType === "remote" || workerType === "cloud";
     const isCloud = workerType === "cloud";
     const rawHost = isRemote ? document.getElementById(`host-${workerId}`).value : window.location.hostname;
-    const parsedHost = isRemote ? extension._parseHostInput(rawHost) : { host: window.location.hostname, port: null };
+    const parsedHost = isRemote ? 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);
@@ -163,7 +165,7 @@ export async function saveWorkerSettings(extension, workerId) {
 
 export function cancelWorkerSettings(extension, workerId) {
     // Collapse the settings panel
-    extension.toggleWorkerExpanded(workerId);
+    toggleWorkerExpanded(extension, workerId);
 
     // Reset form values to original
     const worker = extension.config.workers.find((w) => w.id === workerId);
@@ -177,7 +179,7 @@ export function cancelWorkerSettings(extension, workerId) {
         // Reset remote checkbox
         const remoteCheckbox = document.getElementById(`remote-${workerId}`);
         if (remoteCheckbox) {
-            remoteCheckbox.checked = extension.isRemoteWorker(worker);
+            remoteCheckbox.checked = isRemoteWorker(extension, worker);
         }
     }
 }

web/workerUtils.js

@@ -1,24 +1,37 @@
-import { BUTTON_STYLES, TIMEOUTS, ENDPOINTS, NODE_CLASSES } from './constants.js';
+import { TIMEOUTS, ENDPOINTS } from './constants.js';
+import { checkAllWorkerStatuses, getWorkerUrl } from './workerLifecycle.js';
 
 export async function handleWorkerOperation(extension, button, operation, successText, errorText) {
     const originalText = button.textContent;
     const originalStyle = button.style.cssText;
+    const originalClasses = Array.from(button.classList);
+    const stateClasses = ["btn--working", "btn--success", "btn--error"];
+
+    const setButtonStateClass = (className) => {
+        button.classList.remove(...stateClasses);
+        if (className) {
+            button.classList.add(className);
+        }
+    };
     
     button.textContent = operation.loadingText;
     button.disabled = true;
+    setButtonStateClass("btn--working");
     
     try {
         const urlsToProcess = extension.enabledWorkers.map(w => ({ 
             name: w.name, 
-            url: extension.getWorkerUrl(w)
+            url: getWorkerUrl(extension, w)
         }));
         
         if (urlsToProcess.length === 0) {
             button.textContent = "No Workers";
-            button.style.backgroundColor = "#c04c4c";
+            setButtonStateClass("btn--error");
             setTimeout(() => {
                 button.textContent = originalText;
                 button.style.cssText = originalStyle;
+                button.classList.remove(...stateClasses);
+                button.classList.add(...originalClasses);
                 button.disabled = false;
             }, TIMEOUTS.BUTTON_RESET);
             return;
@@ -38,17 +51,19 @@ export async function handleWorkerOperation(extension, button, operation, succes
         
         if (failures.length === 0) {
             button.textContent = successText;
-            button.style.backgroundColor = BUTTON_STYLES.success.split(':')[1].trim().replace(';', '');
+            setButtonStateClass("btn--success");
             if (operation.onSuccess) operation.onSuccess();
         } else {
             button.textContent = errorText;
-            button.style.backgroundColor = BUTTON_STYLES.error.split(':')[1].trim().replace(';', '');
+            setButtonStateClass("btn--error");
             extension.log(`${operation.name} failed on: ${failures.map(f => f.name).join(", ")}`, "error");
         }
         
         setTimeout(() => {
             button.textContent = originalText;
             button.style.cssText = originalStyle;
+            button.classList.remove(...stateClasses);
+            button.classList.add(...originalClasses);
         }, TIMEOUTS.BUTTON_RESET);
     } finally {
         button.disabled = false;
@@ -60,7 +75,7 @@ export async function handleInterruptWorkers(extension, button) {
         name: "Interrupt",
         endpoint: ENDPOINTS.INTERRUPT,
         loadingText: "Interrupting...",
-        onSuccess: () => setTimeout(() => extension.checkAllWorkerStatuses(), TIMEOUTS.POST_ACTION_DELAY)
+        onSuccess: () => setTimeout(() => checkAllWorkerStatuses(extension), TIMEOUTS.POST_ACTION_DELAY)
     }, "Interrupted!", "Error! See Console");
 }
 
@@ -77,292 +92,3 @@ export function findNodesByClass(apiPrompt, className) {
         .filter(([, nodeData]) => nodeData.class_type === className)
         .map(([nodeId, nodeData]) => ({ id: nodeId, data: nodeData }));
 }
-
-/**
- * Find all image references in the workflow
- * Looks for inputs named "image" that contain filename strings
- */
-export function findImageReferences(extension, apiPrompt) {
-    const images = new Map();
-    // Updated regex to handle:
-    // - Standard files: "image.png"
-    // - Subfolder files: "subfolder/image.png"
-    // - ComfyUI special format: "clipspace/file.png [input]"
-    // - Video files: "video.mp4", "animation.avi", etc.
-    const imageExtensions = /\.(png|jpg|jpeg|gif|webp|bmp|mp4|avi|mov|mkv|webm)(\s*\[\w+\])?$/i;
-    
-    for (const [nodeId, node] of Object.entries(apiPrompt)) {
-        // Check for both 'image' and 'video' inputs
-        const mediaInputs = [];
-        if (node.inputs && node.inputs.image) {
-            mediaInputs.push(node.inputs.image);
-        }
-        if (node.inputs && node.inputs.video) {
-            mediaInputs.push(node.inputs.video);
-        }
-        
-        for (const mediaValue of mediaInputs) {
-            if (typeof mediaValue === 'string') {
-                // Clean special suffixes like [input] or [output]
-                const cleanValue = mediaValue.replace(/\s*\[\w+\]$/, '').trim();
-                // Normalize to forward slashes so subfolder/filename derivation is consistent on Windows
-                const normalizedValue = cleanValue.replace(/\\/g, '/');
-                if (imageExtensions.test(normalizedValue)) {
-                    images.set(normalizedValue, {
-                        nodeId,
-                        nodeType: node.class_type,
-                        inputName: 'image'  // Keep as 'image' for compatibility
-                    });
-                    extension.log(`Found media reference: ${normalizedValue} in node ${nodeId} (${node.class_type})`, "debug");
-                }
-            }
-        }
-    }
-    
-    return images;
-}
-
-/**
- * Find only upstream nodes (inputs) for distributed collector nodes
- * This is used for workers to avoid executing downstream nodes like SaveImage
- * @param {Object} apiPrompt - The API prompt containing the workflow
- * @param {Array<string>} collectorIds - Array of collector node IDs
- * @returns {Set<string>} Set of node IDs that feed into collectors
- */
-export function findCollectorUpstreamNodes(apiPrompt, collectorIds) {
-    const connected = new Set(collectorIds); // Include all collectors
-    const toProcess = [...collectorIds];
-    
-    // Only traverse upstream (inputs)
-    while (toProcess.length > 0) {
-        const nodeId = toProcess.pop();
-        const node = apiPrompt[nodeId];
-        
-        // Traverse upstream (inputs) only
-        if (node && node.inputs) {
-            for (const [inputName, inputValue] of Object.entries(node.inputs)) {
-                if (Array.isArray(inputValue) && inputValue.length === 2) {
-                    const sourceNodeId = String(inputValue[0]);
-                    if (!connected.has(sourceNodeId)) {
-                        connected.add(sourceNodeId);
-                        toProcess.push(sourceNodeId);
-                    }
-                }
-            }
-        }
-    }
-    
-    return connected;
-}
-
-/**
- * Find all downstream nodes from the provided collector IDs.
- * Used to keep only post-collection nodes when master runs in orchestrator mode.
- * @param {Object} apiPrompt
- * @param {Array<string>} collectorIds
- * @returns {Set<string>}
- */
-export function findCollectorDownstreamNodes(apiPrompt, collectorIds) {
-    const adjacency = new Map();
-
-    for (const [nodeId, node] of Object.entries(apiPrompt)) {
-        if (!node.inputs) continue;
-        for (const inputValue of Object.values(node.inputs)) {
-            if (Array.isArray(inputValue) && inputValue.length === 2) {
-                const sourceId = String(inputValue[0]);
-                if (!adjacency.has(sourceId)) {
-                    adjacency.set(sourceId, new Set());
-                }
-                adjacency.get(sourceId).add(String(nodeId));
-            }
-        }
-    }
-
-    const connected = new Set(collectorIds);
-    const queue = [...collectorIds];
-
-    while (queue.length > 0) {
-        const current = queue.shift();
-        const dependents = adjacency.get(current);
-        if (!dependents) continue;
-        dependents.forEach((dependentId) => {
-            if (!connected.has(dependentId)) {
-                connected.add(dependentId);
-                queue.push(dependentId);
-            }
-        });
-    }
-
-    return connected;
-}
-
-/**
- * Prune workflow to only include nodes connected to distributed nodes
- * @param {Object} apiPrompt - The full workflow API prompt
- * @param {Array} distributedNodes - Array of distributed nodes (optional, will find if not provided)
- * @returns {Object} Pruned API prompt with only required nodes
- */
-export function pruneWorkflowForWorker(extension, apiPrompt, distributedNodes = null) {
-    // Find all distributed nodes if not provided
-    if (!distributedNodes) {
-        const collectorNodes = findNodesByClass(apiPrompt, NODE_CLASSES.DISTRIBUTED_COLLECTOR);
-        const upscaleNodes = findNodesByClass(apiPrompt, NODE_CLASSES.UPSCALE_DISTRIBUTED);
-        distributedNodes = [...collectorNodes, ...upscaleNodes];
-    }
-    
-    if (distributedNodes.length === 0) {
-        // No distributed nodes, return full workflow
-        return apiPrompt;
-    }
-    
-    // Get all nodes connected to distributed nodes
-    const distributedIds = distributedNodes.map(node => node.id);
-    
-    // For workers, only include upstream nodes (this removes ALL downstream nodes after collectors)
-    const connectedNodes = findCollectorUpstreamNodes(apiPrompt, distributedIds);
-    
-    extension.log(`Pruning workflow: keeping ${connectedNodes.size} of ${Object.keys(apiPrompt).length} nodes (removed all downstream nodes)`, "debug");
-    
-    // Create pruned prompt with only required nodes
-    const prunedPrompt = {};
-    for (const nodeId of connectedNodes) {
-        prunedPrompt[nodeId] = JSON.parse(JSON.stringify(apiPrompt[nodeId]));
-    }
-    
-    // Check if any distributed node has downstream SaveImage nodes that were removed
-    // If so, add a PreviewImage node after the collector
-    for (const distNode of distributedNodes) {
-        const distNodeId = distNode.id;
-        
-        // Check if this distributed node had any downstream nodes in the original workflow
-        const originalOutputMap = new Map();
-        for (const [nodeId, node] of Object.entries(apiPrompt)) {
-            if (node.inputs) {
-                for (const [inputName, inputValue] of Object.entries(node.inputs)) {
-                    if (Array.isArray(inputValue) && inputValue.length === 2 && String(inputValue[0]) === distNodeId) {
-                        if (!originalOutputMap.has(distNodeId)) {
-                            originalOutputMap.set(distNodeId, []);
-                        }
-                        originalOutputMap.get(distNodeId).push({nodeId, inputName});
-                    }
-                }
-            }
-        }
-        
-        // If this distributed node had downstream nodes that were removed, add a PreviewImage
-        if (originalOutputMap.has(distNodeId) && originalOutputMap.get(distNodeId).length > 0) {
-            // Generate unique numeric ID: max existing numeric key +1
-            const existingIds = Object.keys(prunedPrompt)
-                .filter(k => !isNaN(parseInt(k)))
-                .map(k => parseInt(k));
-            const maxId = existingIds.length > 0 ? Math.max(...existingIds) : 0;
-            const previewNodeId = String(maxId + 1);
-            
-            // Add PreviewImage node connected to the distributed node
-            prunedPrompt[previewNodeId] = {
-                inputs: {
-                    images: [distNodeId, 0]  // Connect to first output of distributed node
-                },
-                class_type: NODE_CLASSES.PREVIEW_IMAGE,
-                _meta: {
-                    title: "Preview Image (auto-added)"
-                }
-            };
-            
-            extension.log(`Added PreviewImage node ${previewNodeId} after distributed node ${distNodeId} for worker`, "debug");
-        }
-    }
-    
-    return prunedPrompt;
-}
-
-/**
- * Check if a node has an upstream node of a specific type
- * @param {Object} apiPrompt - The workflow API prompt
- * @param {string} nodeId - The node to check
- * @param {string} upstreamType - The class_type to look for upstream
- * @returns {boolean} True if an upstream node of the specified type exists
- */
-export function hasUpstreamNode(apiPrompt, nodeId, upstreamType) {
-    const visited = new Set();
-    const toProcess = [nodeId];
-    
-    while (toProcess.length > 0) {
-        const currentId = toProcess.pop();
-        if (visited.has(currentId)) continue;
-        visited.add(currentId);
-        
-        const node = apiPrompt[currentId];
-        if (!node) continue;
-        
-        // Check inputs for upstream connections
-        if (node.inputs) {
-            for (const [inputName, inputValue] of Object.entries(node.inputs)) {
-                if (Array.isArray(inputValue) && inputValue.length === 2) {
-                    const sourceNodeId = String(inputValue[0]);
-                    const sourceNode = apiPrompt[sourceNodeId];
-                    
-                    if (sourceNode && sourceNode.class_type === upstreamType) {
-                        return true;
-                    }
-                    
-                    if (!visited.has(sourceNodeId)) {
-                        toProcess.push(sourceNodeId);
-                    }
-                }
-            }
-        }
-    }
-    
-    return false;
-}
-
-/**
- * Get system information from a worker
- * @param {string} workerUrl - The worker URL
- * @returns {Promise<Object>} System information including platform details
- */
-export async function getWorkerSystemInfo(workerUrl) {
-    try {
-        const response = await fetch(`${workerUrl}/distributed/system_info`);
-        if (!response.ok) {
-            throw new Error(`HTTP ${response.status}`);
-        }
-        return await response.json();
-    } catch (error) {
-        console.warn(`Failed to get system info from ${workerUrl}:`, error);
-        // Return sensible defaults
-        return {
-            platform: {
-                os_name: 'posix',  // Assume Linux
-                path_separator: '/',
-                system: 'Linux'
-            }
-        };
-    }
-}
-
-// Cache system info to avoid repeated calls
-const systemInfoCache = new Map();
-
-/**
- * Get cached system information from a worker
- * @param {string} workerUrl - The worker URL
- * @returns {Promise<Object>} Cached or fresh system information
- */
-export async function getCachedWorkerSystemInfo(workerUrl) {
-    if (systemInfoCache.has(workerUrl)) {
-        return systemInfoCache.get(workerUrl);
-    }
-    
-    const info = await getWorkerSystemInfo(workerUrl);
-    systemInfoCache.set(workerUrl, info);
-    return info;
-}
-
-/**
- * Clear the system info cache
- */
-export function clearSystemInfoCache() {
-    systemInfoCache.clear();
-}

workers/process/__init__.py

@@ -0,0 +1,11 @@
+from .launch_builder import LaunchCommandBuilder
+from .lifecycle import ProcessLifecycle
+from .persistence import ProcessPersistence
+from .root_discovery import ComfyRootDiscovery
+
+__all__ = [
+    "ComfyRootDiscovery",
+    "LaunchCommandBuilder",
+    "ProcessLifecycle",
+    "ProcessPersistence",
+]

workers/process/launch_builder.py

@@ -0,0 +1,79 @@
+import glob
+import os
+import shlex
+import shutil
+
+from ...utils.logging import debug_log
+from ...utils.process import get_python_executable
+
+
+class LaunchCommandBuilder:
+    """Build command-lines for launching worker ComfyUI processes."""
+
+    def _find_windows_terminal(self):
+        """Find Windows Terminal executable."""
+        possible_paths = [
+            os.path.expandvars(r"%LOCALAPPDATA%\Microsoft\WindowsApps\wt.exe"),
+            os.path.expandvars(r"%PROGRAMFILES%\WindowsApps\Microsoft.WindowsTerminal_*\wt.exe"),
+            "wt.exe",
+        ]
+
+        for path in possible_paths:
+            if os.path.exists(path):
+                return path
+            if "*" in path:
+                matches = glob.glob(path)
+                if matches:
+                    return matches[0]
+
+        wt_path = shutil.which("wt")
+        if wt_path:
+            return wt_path
+        return None
+
+    def build_launch_command(self, worker_config, comfy_root):
+        """Build the command to launch a worker."""
+        main_py = os.path.join(comfy_root, "main.py")
+
+        if os.path.exists(main_py):
+            cmd = [
+                get_python_executable(),
+                main_py,
+                "--port",
+                str(worker_config["port"]),
+                "--enable-cors-header",
+            ]
+            debug_log(f"Using main.py: {main_py}")
+        else:
+            error_msg = f"Could not find main.py in {comfy_root}\n"
+            error_msg += f"Searched for: {main_py}\n"
+            error_msg += f"Directory contents of {comfy_root}:\n"
+            try:
+                if os.path.exists(comfy_root):
+                    files = os.listdir(comfy_root)[:20]
+                    error_msg += "  " + "\n  ".join(files)
+                    if len(os.listdir(comfy_root)) > 20:
+                        error_msg += f"\n  ... and {len(os.listdir(comfy_root)) - 20} more files"
+                else:
+                    error_msg += f"  Directory {comfy_root} does not exist!"
+            except Exception as exc:
+                error_msg += f"  Error listing directory: {exc}"
+
+            error_msg += "\n\nPossible solutions:\n"
+            error_msg += "1. Check if ComfyUI is installed in a different location\n"
+            error_msg += "2. For Docker: ComfyUI might be in /ComfyUI or /app\n"
+            error_msg += "3. Ensure the custom node is installed in the correct location\n"
+            raise RuntimeError(error_msg)
+
+        if worker_config.get("extra_args"):
+            raw_args = worker_config["extra_args"].strip()
+            if raw_args:
+                extra_args_list = shlex.split(raw_args)
+                forbidden_chars = set(";|>&<`$()[]{}*!?")
+                for arg in extra_args_list:
+                    if any(char in forbidden_chars for char in arg):
+                        forbidden = "".join(forbidden_chars)
+                        raise ValueError(f"Invalid characters in extra_args: {arg}. Forbidden: {forbidden}")
+                cmd.extend(extra_args_list)
+
+        return cmd

workers/process/lifecycle.py

@@ -0,0 +1,292 @@
+import os
+import platform
+import signal
+import subprocess
+import time
+
+from ...utils.config import load_config, save_config
+from ...utils.constants import PROCESS_TERMINATION_TIMEOUT, PROCESS_WAIT_TIMEOUT, WORKER_CHECK_INTERVAL
+from ...utils.logging import debug_log, log
+from ...utils.process import get_python_executable, is_process_alive, terminate_process
+
+try:
+    import psutil
+
+    PSUTIL_AVAILABLE = True
+except ImportError:
+    log("psutil not available, using fallback process management")
+    PSUTIL_AVAILABLE = False
+
+
+class ProcessLifecycle:
+    """Worker process lifecycle operations operating on manager-owned state."""
+
+    def __init__(self, manager):
+        self._manager = manager
+
+    def launch_worker(self, worker_config, show_window=False):
+        """Launch a worker process with logging."""
+        _ = show_window  # Kept for API compatibility.
+        comfy_root = self._manager.find_comfy_root()
+
+        env = os.environ.copy()
+        env["CUDA_VISIBLE_DEVICES"] = str(worker_config.get("cuda_device", 0))
+        env["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+        env["COMFYUI_MASTER_PID"] = str(os.getpid())
+
+        cmd = self._manager.build_launch_command(worker_config, comfy_root)
+        cwd = comfy_root
+
+        log_dir = os.path.join(comfy_root, "logs", "workers")
+        os.makedirs(log_dir, exist_ok=True)
+
+        date_stamp = time.strftime("%Y%m%d")
+        worker_name = worker_config.get("name", f"Worker{worker_config['id']}")
+        safe_name = "".join(char if char.isalnum() or char in ("-", "_") else "_" for char in worker_name)
+        log_file = os.path.join(log_dir, f"{safe_name}_{date_stamp}.log")
+
+        with open(log_file, "a", encoding="utf-8") as log_handle:
+            log_handle.write(f"\n\n{'=' * 50}\n")
+            log_handle.write("=== ComfyUI Worker Session Started ===\n")
+            log_handle.write(f"Worker: {worker_name}\n")
+            log_handle.write(f"Port: {worker_config['port']}\n")
+            log_handle.write(f"CUDA Device: {worker_config.get('cuda_device', 0)}\n")
+            log_handle.write(f"Started: {time.strftime('%Y-%m-%d %H:%M:%S')}\n")
+            log_handle.write(f"Command: {' '.join(cmd)}\n")
+
+            config = load_config()
+            stop_on_master_exit = config.get("settings", {}).get("stop_workers_on_master_exit", True)
+            if stop_on_master_exit:
+                log_handle.write("Note: Worker will stop when master shuts down\n")
+            else:
+                log_handle.write("Note: Worker will continue running after master shuts down\n")
+            log_handle.write("=" * 30 + "\n\n")
+            log_handle.flush()
+
+            if stop_on_master_exit and env.get("COMFYUI_MASTER_PID"):
+                monitor_script = os.path.join(
+                    os.path.dirname(os.path.dirname(__file__)),
+                    "worker_monitor.py",
+                )
+                monitored_cmd = [get_python_executable(), monitor_script] + cmd
+                log_handle.write(f"[Worker Monitor] Monitoring master PID: {env['COMFYUI_MASTER_PID']}\n")
+                log_handle.flush()
+            else:
+                monitored_cmd = cmd
+
+            if platform.system() == "Windows":
+                create_no_window = 0x08000000
+                process = subprocess.Popen(
+                    monitored_cmd,
+                    env=env,
+                    cwd=cwd,
+                    stdout=log_handle,
+                    stderr=subprocess.STDOUT,
+                    creationflags=create_no_window,
+                )
+            else:
+                process = subprocess.Popen(
+                    monitored_cmd,
+                    env=env,
+                    cwd=cwd,
+                    stdout=log_handle,
+                    stderr=subprocess.STDOUT,
+                    start_new_session=True,
+                )
+
+        worker_id = str(worker_config["id"])
+        self._manager.processes[worker_id] = {
+            "pid": process.pid,
+            "process": process,
+            "started_at": time.time(),
+            "config": worker_config,
+            "log_file": log_file,
+            "is_monitor": stop_on_master_exit and env.get("COMFYUI_MASTER_PID"),
+            "launching": True,
+        }
+
+        self._manager.save_processes()
+
+        if stop_on_master_exit and env.get("COMFYUI_MASTER_PID"):
+            debug_log(f"Launched worker {worker_name} via monitor (Monitor PID: {process.pid})")
+        else:
+            log(f"Launched worker {worker_name} directly (PID: {process.pid})")
+        debug_log(f"Log file: {log_file}")
+        return process.pid
+
+    def stop_worker(self, worker_id):
+        """Stop a worker process."""
+        worker_id = str(worker_id)
+        if worker_id not in self._manager.processes:
+            return False, "Worker not managed by UI"
+
+        proc_info = self._manager.processes[worker_id]
+        process = proc_info.get("process")
+        pid = proc_info["pid"]
+        debug_log(f"Attempting to stop worker {worker_id} (PID: {pid})")
+
+        if not process:
+            try:
+                debug_log("[Distributed] Stopping restored process (no subprocess object)")
+                if self._kill_process_tree(pid):
+                    del self._manager.processes[worker_id]
+                    self._manager.save_processes()
+                    debug_log(f"Successfully stopped worker {worker_id} and all child processes")
+                    return True, "Worker stopped"
+                return False, "Failed to stop worker process"
+            except Exception as exc:
+                log(f"[Distributed] Exception during stop: {exc}")
+                return False, f"Error stopping worker: {str(exc)}"
+
+        if process.poll() is not None:
+            log(f"[Distributed] Worker {worker_id} already stopped")
+            del self._manager.processes[worker_id]
+            self._manager.save_processes()
+            return False, "Worker already stopped"
+
+        try:
+            debug_log(f"Using process tree kill for worker {worker_id}")
+            if self._kill_process_tree(pid):
+                del self._manager.processes[worker_id]
+                self._manager.save_processes()
+                debug_log(f"Successfully stopped worker {worker_id} and all child processes")
+                return True, "Worker stopped"
+
+            log("[Distributed] Process tree kill failed, trying normal termination")
+            terminate_process(process, timeout=PROCESS_TERMINATION_TIMEOUT)
+            del self._manager.processes[worker_id]
+            self._manager.save_processes()
+            return True, "Worker stopped (fallback)"
+        except Exception as exc:
+            log(f"[Distributed] Exception during stop: {exc}")
+            return False, f"Error stopping worker: {str(exc)}"
+
+    def get_managed_workers(self):
+        """Get list of workers managed by this process."""
+        managed = {}
+        for worker_id, proc_info in list(self._manager.processes.items()):
+            is_running, _ = self._check_worker_process(worker_id, proc_info)
+            if is_running:
+                managed[worker_id] = {
+                    "pid": proc_info["pid"],
+                    "started_at": proc_info["started_at"],
+                    "log_file": proc_info.get("log_file"),
+                    "launching": proc_info.get("launching", False),
+                }
+            else:
+                del self._manager.processes[worker_id]
+
+        return managed
+
+    def cleanup_all(self):
+        """Stop all managed workers (called on shutdown)."""
+        for worker_id in list(self._manager.processes.keys()):
+            try:
+                self.stop_worker(worker_id)
+            except Exception as exc:
+                log(f"[Distributed] Error stopping worker {worker_id}: {exc}")
+
+        config = load_config()
+        config["managed_processes"] = {}
+        save_config(config)
+
+    def _is_process_running(self, pid):
+        """Check if a process with given PID is running."""
+        return is_process_alive(pid)
+
+    def _check_worker_process(self, worker_id, proc_info):
+        """Check if a worker process is still running and return status."""
+        _ = worker_id  # Signature retained for compatibility with existing callers.
+        process = proc_info.get("process")
+        pid = proc_info.get("pid")
+
+        if process:
+            return process.poll() is None, True
+        if pid:
+            return self._is_process_running(pid), False
+        return False, False
+
+    def _kill_process_tree(self, pid):
+        """Kill a process and all its children."""
+        if PSUTIL_AVAILABLE:
+            try:
+                parent = psutil.Process(pid)
+                children = parent.children(recursive=True)
+
+                debug_log(f"Killing process tree for PID {pid} ({parent.name()})")
+                for child in children:
+                    debug_log(f"  - Child PID {child.pid} ({child.name()})")
+
+                for child in children:
+                    try:
+                        debug_log(f"Terminating child {child.pid}")
+                        child.terminate()
+                    except psutil.NoSuchProcess:
+                        pass
+
+                _, alive = psutil.wait_procs(children, timeout=PROCESS_WAIT_TIMEOUT)
+                for child in alive:
+                    try:
+                        debug_log(f"Force killing child {child.pid}")
+                        child.kill()
+                    except psutil.NoSuchProcess:
+                        pass
+
+                try:
+                    debug_log(f"Terminating parent {pid}")
+                    parent.terminate()
+                    parent.wait(timeout=PROCESS_WAIT_TIMEOUT)
+                except psutil.TimeoutExpired:
+                    debug_log(f"Force killing parent {pid}")
+                    parent.kill()
+                except psutil.NoSuchProcess:
+                    debug_log(f"Parent process {pid} already gone")
+                return True
+            except psutil.NoSuchProcess:
+                debug_log(f"Process {pid} does not exist")
+                return False
+            except Exception as exc:
+                debug_log(f"Error killing process tree: {exc}")
+
+        debug_log("[Distributed] Using OS commands to kill process tree")
+        if platform.system() == "Windows":
+            try:
+                result = subprocess.run(
+                    ["wmic", "process", "where", f"ParentProcessId={pid}", "get", "ProcessId"],
+                    capture_output=True,
+                    text=True,
+                )
+                if result.returncode == 0:
+                    lines = result.stdout.strip().split("\n")[1:]
+                    child_pids = [line.strip() for line in lines if line.strip().isdigit()]
+                    debug_log(f"[Distributed] Found child processes: {child_pids}")
+                    for child_pid in child_pids:
+                        try:
+                            subprocess.run(
+                                ["taskkill", "/F", "/PID", child_pid],
+                                capture_output=True,
+                                check=False,
+                            )
+                        except (FileNotFoundError, OSError) as exc:
+                            debug_log(f"[Distributed] Warning: taskkill failed for PID {child_pid}: {exc}")
+
+                result = subprocess.run(
+                    ["taskkill", "/F", "/PID", str(pid), "/T"],
+                    capture_output=True,
+                    text=True,
+                )
+                debug_log(f"[Distributed] Taskkill result: {result.stdout.strip()}")
+                return result.returncode == 0
+            except Exception as exc:
+                log(f"[Distributed] Error with taskkill: {exc}")
+                return False
+
+        try:
+            subprocess.run(["pkill", "-TERM", "-P", str(pid)], check=False)
+            time.sleep(WORKER_CHECK_INTERVAL)
+            subprocess.run(["pkill", "-KILL", "-P", str(pid)], check=False)
+            os.kill(pid, signal.SIGKILL)
+            return True
+        except Exception as exc:
+            log(f"[Distributed] Error killing process tree for PID {pid}: {exc}")
+            return False

workers/process/persistence.py

@@ -0,0 +1,48 @@
+from ...utils.config import load_config, save_config
+from ...utils.logging import debug_log
+
+
+class ProcessPersistence:
+    """Persist and restore manager-owned worker process metadata."""
+
+    def __init__(self, manager):
+        self._manager = manager
+
+    def load_processes(self):
+        """Load persisted process information from config."""
+        config = load_config()
+        managed_processes = config.get("managed_processes", {})
+
+        for worker_id, proc_info in managed_processes.items():
+            pid = proc_info.get("pid")
+            if pid and self._manager._is_process_running(pid):
+                self._manager.processes[worker_id] = {
+                    "pid": pid,
+                    "process": None,
+                    "started_at": proc_info.get("started_at"),
+                    "config": proc_info.get("config"),
+                    "log_file": proc_info.get("log_file"),
+                }
+                debug_log(f"[Distributed] Restored worker {worker_id} (PID: {pid})")
+            elif pid:
+                debug_log(f"[Distributed] Worker {worker_id} (PID: {pid}) is no longer running")
+
+    def save_processes(self):
+        """Save process information to config."""
+        config = load_config()
+        managed_processes = {}
+
+        for worker_id, proc_info in self._manager.processes.items():
+            is_running, _ = self._manager._check_worker_process(worker_id, proc_info)
+            if not is_running:
+                continue
+            managed_processes[worker_id] = {
+                "pid": proc_info["pid"],
+                "started_at": proc_info["started_at"],
+                "config": proc_info["config"],
+                "log_file": proc_info.get("log_file"),
+                "launching": proc_info.get("launching", False),
+            }
+
+        config["managed_processes"] = managed_processes
+        save_config(config)

workers/process/root_discovery.py

@@ -0,0 +1,66 @@
+import os
+
+from ...utils.logging import debug_log, log
+
+
+class ComfyRootDiscovery:
+    """Resolve the ComfyUI root directory across local and container layouts."""
+
+    def find_comfy_root(self):
+        # Start from current file location.
+        current_dir = os.path.dirname(os.path.abspath(__file__))
+        potential_root = os.path.dirname(os.path.dirname(current_dir))
+
+        # Method 1: Check for environment variable override.
+        env_root = os.environ.get("COMFYUI_ROOT")
+        if env_root and os.path.exists(os.path.join(env_root, "main.py")):
+            debug_log(f"Found ComfyUI root via COMFYUI_ROOT environment variable: {env_root}")
+            return env_root
+
+        # Method 2: Try going up from custom_nodes directory.
+        if os.path.exists(os.path.join(potential_root, "main.py")):
+            debug_log(f"Found ComfyUI root via directory traversal: {potential_root}")
+            return potential_root
+
+        # Method 3: Look for common Docker paths.
+        docker_paths = [
+            "/basedir",
+            "/ComfyUI",
+            "/app",
+            "/workspace/ComfyUI",
+            "/comfyui",
+            "/opt/ComfyUI",
+            "/workspace",
+        ]
+        for path in docker_paths:
+            if os.path.exists(path) and os.path.exists(os.path.join(path, "main.py")):
+                debug_log(f"Found ComfyUI root in Docker path: {path}")
+                return path
+
+        # Method 4: Search upwards for main.py.
+        search_dir = current_dir
+        for _ in range(5):
+            if os.path.exists(os.path.join(search_dir, "main.py")):
+                debug_log(f"Found ComfyUI root via upward search: {search_dir}")
+                return search_dir
+            parent = os.path.dirname(search_dir)
+            if parent == search_dir:
+                break
+            search_dir = parent
+
+        # Method 5: Try to import and use folder_paths.
+        try:
+            import folder_paths
+
+            if hasattr(folder_paths, "base_path") and os.path.exists(
+                os.path.join(folder_paths.base_path, "main.py")
+            ):
+                debug_log(f"Found ComfyUI root via folder_paths: {folder_paths.base_path}")
+                return folder_paths.base_path
+        except Exception as exc:
+            debug_log(f"folder_paths root detection failed: {exc}")
+
+        log("Warning: Could not reliably determine ComfyUI root directory")
+        log(f"Current directory: {current_dir}")
+        log(f"Initial guess was: {potential_root}")
+        return potential_root

workers/process_manager.py

@@ -1,521 +1,49 @@
-import os
-import json
-import re
-import asyncio
-import subprocess
-import platform
-import time
-import atexit
-import signal
-import sys
-import shlex
-import uuid
-import torch
-
-from ..utils.logging import debug_log, log
-from ..utils.config import load_config, save_config
-from ..utils.process import is_process_alive, terminate_process, get_python_executable
-from ..utils.constants import PROCESS_TERMINATION_TIMEOUT, WORKER_CHECK_INTERVAL, PROCESS_WAIT_TIMEOUT
-
-# Try to import psutil for better process management
-try:
-    import psutil
-    PSUTIL_AVAILABLE = True
-except ImportError:
-    log("psutil not available, using fallback process management")
-    PSUTIL_AVAILABLE = False
+from .process import ComfyRootDiscovery, LaunchCommandBuilder, ProcessLifecycle, ProcessPersistence
 
 
 class WorkerProcessManager:
+    """Thin composition wrapper around worker process subsystems."""
+
     def __init__(self):
-        self.processes = {}  # worker_id -> process info
+        self.processes = {}
+        self._root_discovery = ComfyRootDiscovery()
+        self._launch_builder = LaunchCommandBuilder()
+        self._lifecycle = ProcessLifecycle(self)
+        self._persistence = ProcessPersistence(self)
         self.load_processes()
-        
+
     def find_comfy_root(self):
-        """Find the ComfyUI root directory."""
-        # Start from current file location
-        current_dir = os.path.dirname(os.path.abspath(__file__))
-        
-        # Method 1: Check for environment variable override
-        env_root = os.environ.get('COMFYUI_ROOT')
-        if env_root and os.path.exists(os.path.join(env_root, "main.py")):
-            debug_log(f"Found ComfyUI root via COMFYUI_ROOT environment variable: {env_root}")
-            return env_root
-        
-        # Method 2: Try going up from custom_nodes directory
-        # This file should be in ComfyUI/custom_nodes/ComfyUI-Distributed/
-        potential_root = os.path.dirname(os.path.dirname(current_dir))
-        if os.path.exists(os.path.join(potential_root, "main.py")):
-            debug_log(f"Found ComfyUI root via directory traversal: {potential_root}")
-            return potential_root
-        
-        # Method 3: Look for common Docker paths
-        docker_paths = ["/basedir", "/ComfyUI", "/app", "/workspace/ComfyUI", "/comfyui", "/opt/ComfyUI", "/workspace"]
-        for path in docker_paths:
-            if os.path.exists(path) and os.path.exists(os.path.join(path, "main.py")):
-                debug_log(f"Found ComfyUI root in Docker path: {path}")
-                return path
-        
-        # Method 4: Search upwards for main.py
-        search_dir = current_dir
-        for _ in range(5):  # Limit search depth
-            if os.path.exists(os.path.join(search_dir, "main.py")):
-                debug_log(f"Found ComfyUI root via upward search: {search_dir}")
-                return search_dir
-            parent = os.path.dirname(search_dir)
-            if parent == search_dir:  # Reached root
-                break
-            search_dir = parent
-        
-        # Method 5: Try to import and use folder_paths
-        try:
-            import folder_paths
-            # folder_paths.base_path should point to ComfyUI root
-            if hasattr(folder_paths, 'base_path') and os.path.exists(os.path.join(folder_paths.base_path, "main.py")):
-                debug_log(f"Found ComfyUI root via folder_paths: {folder_paths.base_path}")
-                return folder_paths.base_path
-        except Exception:
-            pass
-        
-        # If all methods fail, log detailed information and return the best guess
-        log(f"Warning: Could not reliably determine ComfyUI root directory")
-        log(f"Current directory: {current_dir}")
-        log(f"Initial guess was: {potential_root}")
-        
-        # Return the initial guess as fallback
-        return potential_root
-        
+        return self._root_discovery.find_comfy_root()
+
     def _find_windows_terminal(self):
-        """Find Windows Terminal executable."""
-        # Common locations for Windows Terminal
-        possible_paths = [
-            os.path.expandvars(r"%LOCALAPPDATA%\Microsoft\WindowsApps\wt.exe"),
-            os.path.expandvars(r"%PROGRAMFILES%\WindowsApps\Microsoft.WindowsTerminal_*\wt.exe"),
-            "wt.exe"  # Try PATH
-        ]
-        
-        for path in possible_paths:
-            if os.path.exists(path):
-                return path
-            # Handle wildcard for WindowsApps
-            if '*' in path:
-                import glob
-                matches = glob.glob(path)
-                if matches:
-                    return matches[0]
-        
-        # Try to find it in PATH
-        import shutil
-        wt_path = shutil.which("wt")
-        if wt_path:
-            return wt_path
-            
-        return None
-        
+        return self._launch_builder._find_windows_terminal()
+
     def build_launch_command(self, worker_config, comfy_root):
-        """Build the command to launch a worker."""
-        # Use main.py directly - it's the most reliable method
-        main_py = os.path.join(comfy_root, "main.py")
-        
-        if os.path.exists(main_py):
-            cmd = [
-                get_python_executable(),
-                main_py,
-                "--port", str(worker_config['port']),
-                "--enable-cors-header"
-            ]
-            debug_log(f"Using main.py: {main_py}")
-        else:
-            # Provide detailed error message
-            error_msg = f"Could not find main.py in {comfy_root}\n"
-            error_msg += f"Searched for: {main_py}\n"
-            error_msg += f"Directory contents of {comfy_root}:\n"
-            try:
-                if os.path.exists(comfy_root):
-                    files = os.listdir(comfy_root)[:20]  # List first 20 files
-                    error_msg += "  " + "\n  ".join(files)
-                    if len(os.listdir(comfy_root)) > 20:
-                        error_msg += f"\n  ... and {len(os.listdir(comfy_root)) - 20} more files"
-                else:
-                    error_msg += f"  Directory {comfy_root} does not exist!"
-            except Exception as e:
-                error_msg += f"  Error listing directory: {e}"
-            
-            # Try to suggest the correct path
-            error_msg += "\n\nPossible solutions:\n"
-            error_msg += "1. Check if ComfyUI is installed in a different location\n"
-            error_msg += "2. For Docker: ComfyUI might be in /ComfyUI or /app\n"
-            error_msg += "3. Ensure the custom node is installed in the correct location\n"
-            
-            raise RuntimeError(error_msg)
-        
-        # Add any extra arguments safely
-        if worker_config.get('extra_args'):
-            raw_args = worker_config['extra_args'].strip()
-            if raw_args:
-                # Safely split using shlex to handle quotes and spaces
-                extra_args_list = shlex.split(raw_args)
-                
-                # Validate: Block dangerous shell meta-characters (e.g., ; | & > < ` $)
-                forbidden_chars = set(';|>&<`$()[]{}*!?')
-                for arg in extra_args_list:
-                    if any(c in forbidden_chars for c in arg):
-                        raise ValueError(f"Invalid characters in extra_args: {arg}. Forbidden: {''.join(forbidden_chars)}")
-                
-                cmd.extend(extra_args_list)
-            
-        return cmd
-        
+        return self._launch_builder.build_launch_command(worker_config, comfy_root)
+
     def launch_worker(self, worker_config, show_window=False):
-        """Launch a worker process with logging."""
-        comfy_root = self.find_comfy_root()
-        
-        # Set up environment
-        env = os.environ.copy()
-        env['CUDA_VISIBLE_DEVICES'] = str(worker_config.get('cuda_device', 0))
-        env['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
-        
-        # Pass master PID to worker so it can monitor if master is still alive
-        env['COMFYUI_MASTER_PID'] = str(os.getpid())
-        
-        cmd = self.build_launch_command(worker_config, comfy_root)
-        
-        # Change to ComfyUI root directory for the process
-        cwd = comfy_root
-        
-        # Create log directory and file
-        log_dir = os.path.join(comfy_root, "logs", "workers")
-        os.makedirs(log_dir, exist_ok=True)
-        
-        # Use daily log files instead of timestamp
-        date_stamp = time.strftime("%Y%m%d")
-        worker_name = worker_config.get('name', f'Worker{worker_config["id"]}')
-        # Clean worker name for filename
-        safe_name = "".join(c if c.isalnum() or c in ('-', '_') else '_' for c in worker_name)
-        log_file = os.path.join(log_dir, f"{safe_name}_{date_stamp}.log")
-        
-        # Launch process with logging (append mode for daily logs)
-        with open(log_file, 'a') as log_handle:
-            # Write startup info to log with timestamp
-            log_handle.write(f"\n\n{'='*50}\n")
-            log_handle.write(f"=== ComfyUI Worker Session Started ===\n")
-            log_handle.write(f"Worker: {worker_name}\n")
-            log_handle.write(f"Port: {worker_config['port']}\n")
-            log_handle.write(f"CUDA Device: {worker_config.get('cuda_device', 0)}\n")
-            log_handle.write(f"Started: {time.strftime('%Y-%m-%d %H:%M:%S')}\n")
-            log_handle.write(f"Command: {' '.join(cmd)}\n")
-            
-            # Note about worker behavior
-            config = load_config()
-            stop_on_master_exit = config.get('settings', {}).get('stop_workers_on_master_exit', True)
-            
-            if stop_on_master_exit:
-                log_handle.write("Note: Worker will stop when master shuts down\n")
-            else:
-                log_handle.write("Note: Worker will continue running after master shuts down\n")
-            
-            log_handle.write("=" * 30 + "\n\n")
-            log_handle.flush()
-            
-            # Wrap command with monitor if needed
-            if stop_on_master_exit and env.get('COMFYUI_MASTER_PID'):
-                # Use the monitor wrapper
-                monitor_script = os.path.join(os.path.dirname(__file__), 'worker_monitor.py')
-                monitored_cmd = [get_python_executable(), monitor_script] + cmd
-                log_handle.write(f"[Worker Monitor] Monitoring master PID: {env['COMFYUI_MASTER_PID']}\n")
-                log_handle.flush()
-            else:
-                monitored_cmd = cmd
-            
-            # Platform-specific process creation - always hidden with logging
-            if platform.system() == "Windows":
-                CREATE_NO_WINDOW = 0x08000000
-                process = subprocess.Popen(
-                    monitored_cmd, env=env, cwd=cwd,
-                    stdout=log_handle, 
-                    stderr=subprocess.STDOUT,
-                    creationflags=CREATE_NO_WINDOW
-                )
-            else:
-                # Unix-like systems
-                process = subprocess.Popen(
-                    monitored_cmd, env=env, cwd=cwd,
-                    stdout=log_handle,
-                    stderr=subprocess.STDOUT,
-                    start_new_session=True  # Detach from parent
-                )
-        
-        # Track the process with log file info - use string ID for consistency
-        worker_id = str(worker_config['id'])
-        self.processes[worker_id] = {
-            'pid': process.pid,
-            'process': process,
-            'started_at': time.time(),
-            'config': worker_config,
-            'log_file': log_file,
-            'is_monitor': stop_on_master_exit and env.get('COMFYUI_MASTER_PID'),  # Track if using monitor
-            'launching': True  # Mark as launching until confirmed running
-        }
-        
-        # Save process info for persistence
-        self.save_processes()
-        
-        if stop_on_master_exit and env.get('COMFYUI_MASTER_PID'):
-            debug_log(f"Launched worker {worker_name} via monitor (Monitor PID: {process.pid})")
-        else:
-            log(f"Launched worker {worker_name} directly (PID: {process.pid})")
-        debug_log(f"Log file: {log_file}")
-        return process.pid
-        
+        return self._lifecycle.launch_worker(worker_config, show_window=show_window)
+
     def stop_worker(self, worker_id):
-        """Stop a worker process."""
-        # Ensure worker_id is string
-        worker_id = str(worker_id)
-        if worker_id not in self.processes:
-            return False, "Worker not managed by UI"
-            
-        proc_info = self.processes[worker_id]
-        process = proc_info.get('process')
-        pid = proc_info['pid']
-        
-        debug_log(f"Attempting to stop worker {worker_id} (PID: {pid})")
-        
-        # For restored processes without subprocess object
-        if not process:
-            try:
-                debug_log(f"[Distributed] Stopping restored process (no subprocess object)")
-                if self._kill_process_tree(pid):
-                    del self.processes[worker_id]
-                    self.save_processes()
-                    debug_log(f"Successfully stopped worker {worker_id} and all child processes")
-                    return True, "Worker stopped"
-                else:
-                    return False, "Failed to stop worker process"
-            except Exception as e:
-                log(f"[Distributed] Exception during stop: {e}")
-                return False, f"Error stopping worker: {str(e)}"
-        
-        # Normal case with subprocess object
-        # Check if still running
-        if process.poll() is not None:
-            # Already stopped
-            log(f"[Distributed] Worker {worker_id} already stopped")
-            del self.processes[worker_id]
-            self.save_processes()
-            return False, "Worker already stopped"
-            
-        # Try to kill the entire process tree
-        try:
-            debug_log(f"Using process tree kill for worker {worker_id}")
-            if self._kill_process_tree(pid):
-                # Clean up tracking
-                del self.processes[worker_id]
-                self.save_processes()
-                debug_log(f"Successfully stopped worker {worker_id} and all child processes")
-                return True, "Worker stopped"
-            else:
-                # Fallback to normal termination
-                log(f"[Distributed] Process tree kill failed, trying normal termination")
-                if process:
-                    terminate_process(process, timeout=PROCESS_TERMINATION_TIMEOUT)
-                
-                del self.processes[worker_id]
-                self.save_processes()
-                return True, "Worker stopped (fallback)"
-                
-        except Exception as e:
-            log(f"[Distributed] Exception during stop: {e}")
-            return False, f"Error stopping worker: {str(e)}"
-            
+        return self._lifecycle.stop_worker(worker_id)
+
     def get_managed_workers(self):
-        """Get list of workers managed by this process."""
-        managed = {}
-        for worker_id, proc_info in list(self.processes.items()):
-            # Check if process is still running
-            is_running, _ = self._check_worker_process(worker_id, proc_info)
-            
-            if is_running:
-                managed[worker_id] = {
-                    'pid': proc_info['pid'],
-                    'started_at': proc_info['started_at'],
-                    'log_file': proc_info.get('log_file'),
-                    'launching': proc_info.get('launching', False)
-                }
-            else:
-                # Process has stopped, remove from tracking
-                del self.processes[worker_id]
-        
-        return managed
-        
+        return self._lifecycle.get_managed_workers()
+
     def cleanup_all(self):
-        """Stop all managed workers (called on shutdown)."""
-        for worker_id in list(self.processes.keys()):
-            try:
-                self.stop_worker(worker_id)
-            except Exception as e:
-                log(f"[Distributed] Error stopping worker {worker_id}: {e}")
-        
-        # Clear all managed processes from config
-        config = load_config()
-        config['managed_processes'] = {}
-        save_config(config)
-    
+        return self._lifecycle.cleanup_all()
+
     def load_processes(self):
-        """Load persisted process information from config."""
-        config = load_config()
-        managed_processes = config.get('managed_processes', {})
-        
-        # Verify each saved process is still running
-        for worker_id, proc_info in managed_processes.items():
-            pid = proc_info.get('pid')
-            if pid and self._is_process_running(pid):
-                # Reconstruct process info
-                self.processes[worker_id] = {
-                    'pid': pid,
-                    'process': None,  # Can't reconstruct subprocess object
-                    'started_at': proc_info.get('started_at'),
-                    'config': proc_info.get('config'),
-                    'log_file': proc_info.get('log_file')
-                }
-                debug_log(f"[Distributed] Restored worker {worker_id} (PID: {pid})")
-            else:
-                if pid:
-                    debug_log(f"[Distributed] Worker {worker_id} (PID: {pid}) is no longer running")
-    
+        return self._persistence.load_processes()
+
     def save_processes(self):
-        """Save process information to config."""
-        config = load_config()
-        
-        # Create serializable version of process info
-        managed_processes = {}
-        for worker_id, proc_info in self.processes.items():
-            # Only save if process is running
-            is_running, _ = self._check_worker_process(worker_id, proc_info)
-            
-            if is_running:
-                managed_processes[worker_id] = {
-                    'pid': proc_info['pid'],
-                    'started_at': proc_info['started_at'],
-                    'config': proc_info['config'],
-                    'log_file': proc_info.get('log_file'),
-                    'launching': proc_info.get('launching', False)
-                }
-        
-        # Update config with managed processes
-        config['managed_processes'] = managed_processes
-        save_config(config)
-    
+        return self._persistence.save_processes()
+
     def _is_process_running(self, pid):
-        """Check if a process with given PID is running."""
-        return is_process_alive(pid)
-    
+        return self._lifecycle._is_process_running(pid)
+
     def _check_worker_process(self, worker_id, proc_info):
-        """Check if a worker process is still running and return status.
-        
-        Returns:
-            tuple: (is_running, has_subprocess_object)
-        """
-        process = proc_info.get('process')
-        pid = proc_info.get('pid')
-        
-        if process:
-            # Normal case with subprocess object
-            return process.poll() is None, True
-        elif pid:
-            # Restored process without subprocess object
-            return self._is_process_running(pid), False
-        else:
-            # No process or PID
-            return False, False
-    
+        return self._lifecycle._check_worker_process(worker_id, proc_info)
+
     def _kill_process_tree(self, pid):
-        """Kill a process and all its children."""
-        if PSUTIL_AVAILABLE:
-            try:
-                parent = psutil.Process(pid)
-                children = parent.children(recursive=True)
-                
-                # Log what we're about to kill
-                debug_log(f"Killing process tree for PID {pid} ({parent.name()})")
-                for child in children:
-                    debug_log(f"  - Child PID {child.pid} ({child.name()})")
-                
-                # Kill children first
-                for child in children:
-                    try:
-                        debug_log(f"Terminating child {child.pid}")
-                        child.terminate()
-                    except psutil.NoSuchProcess:
-                        pass
-                
-                # Wait a bit for graceful termination
-                gone, alive = psutil.wait_procs(children, timeout=PROCESS_WAIT_TIMEOUT)
-                
-                # Force kill any remaining
-                for child in alive:
-                    try:
-                        debug_log(f"Force killing child {child.pid}")
-                        child.kill()
-                    except psutil.NoSuchProcess:
-                        pass
-                
-                # Finally kill the parent
-                try:
-                    debug_log(f"Terminating parent {pid}")
-                    parent.terminate()
-                    parent.wait(timeout=PROCESS_WAIT_TIMEOUT)
-                except psutil.TimeoutExpired:
-                    debug_log(f"Force killing parent {pid}")
-                    parent.kill()
-                except psutil.NoSuchProcess:
-                    debug_log(f"Parent process {pid} already gone")
-                    
-                return True
-                
-            except psutil.NoSuchProcess:
-                debug_log(f"Process {pid} does not exist")
-                return False
-            except Exception as e:
-                debug_log(f"Error killing process tree: {e}")
-                # Fall through to OS commands
-        
-        # Fallback to OS-specific commands
-        debug_log("[Distributed] Using OS commands to kill process tree")
-        if platform.system() == "Windows":
-            try:
-                # Use wmic to find child processes
-                result = subprocess.run(['wmic', 'process', 'where', f'ParentProcessId={pid}', 'get', 'ProcessId'], 
-                                      capture_output=True, text=True)
-                if result.returncode == 0:
-                    lines = result.stdout.strip().split('\n')[1:]  # Skip header
-                    child_pids = [line.strip() for line in lines if line.strip() and line.strip().isdigit()]
-                    
-                    debug_log(f"[Distributed] Found child processes: {child_pids}")
-                    
-                    # Kill each child
-                    for child_pid in child_pids:
-                        try:
-                            subprocess.run(['taskkill', '/F', '/PID', child_pid], 
-                                         capture_output=True, check=False)
-                        except (OSError, FileNotFoundError) as e:
-                            debug_log(f"[Distributed] Warning: taskkill failed for PID {child_pid}: {e}")
-                
-                # Kill the parent with tree flag
-                result = subprocess.run(['taskkill', '/F', '/PID', str(pid), '/T'], 
-                                      capture_output=True, text=True)
-                debug_log(f"[Distributed] Taskkill result: {result.stdout.strip()}")
-                return result.returncode == 0
-            except Exception as e:
-                log(f"[Distributed] Error with taskkill: {e}")
-                return False
-        else:
-            # Unix: use pkill
-            try:
-                subprocess.run(['pkill', '-TERM', '-P', str(pid)], check=False)
-                time.sleep(WORKER_CHECK_INTERVAL)
-                subprocess.run(['pkill', '-KILL', '-P', str(pid)], check=False)
-                os.kill(pid, signal.SIGKILL)
-                return True
-            except Exception as e:
-                log(f"[Distributed] Error killing process tree for PID {pid}: {e}")
-                return False
+        return self._lifecycle._kill_process_tree(pid)