refactor: implement REFACTOR_PLAN_2 across distributed stack

api/__init__.py

@@ -2,3 +2,4 @@ from . import config_routes  # noqa: F401
 from . import tunnel_routes  # noqa: F401
 from . import worker_routes  # noqa: F401
 from . import job_routes  # noqa: F401
+from . import usdu_routes  # noqa: F401

api/config_routes.py

@@ -3,7 +3,7 @@ import json
 from aiohttp import web
 import server
 
-from ..utils.config import load_config, save_config
+from ..utils.config import config_transaction, load_config
 from ..utils.logging import debug_log, log
 from ..utils.network import handle_api_error, normalize_host
 
@@ -37,6 +37,36 @@ _SETTINGS_FIELDS = {
     "has_auto_populated_workers",
 }
 
+_WORKER_FIELDS = [
+    ("enabled", None, False),
+    ("name", None, False),
+    ("port", None, False),
+    ("host", normalize_host, True),
+    ("cuda_device", None, True),
+    ("extra_args", None, True),
+    ("type", None, False),
+]
+
+_MASTER_FIELDS = [
+    ("name", None, False),
+    ("host", normalize_host, True),
+    ("port", None, False),
+    ("cuda_device", None, False),
+    ("extra_args", None, False),
+]
+
+
+def _apply_field_patch(target: dict, data: dict, field_rules: list) -> None:
+    """Apply a partial update to a target dict based on field rules."""
+    for key, normalizer, remove_on_none in field_rules:
+        if key not in data:
+            continue
+        value = data[key]
+        if value is None and remove_on_none:
+            target.pop(key, None)
+        else:
+            target[key] = normalizer(value) if (normalizer and value is not None) else value
+
 
 @server.PromptServer.instance.routes.get("/distributed/config")
 async def get_config_endpoint(request):
@@ -55,8 +85,8 @@ async def update_config_endpoint(request):
     if not isinstance(data, dict):
         return await handle_api_error(request, "Config payload must be an object", 400)
 
-    config = load_config()
-    settings = config.setdefault("settings", {})
+    validated_settings = {}
+    validated_root = {}
     errors = []
 
     for key, value in data.items():
@@ -74,16 +104,22 @@ async def update_config_endpoint(request):
             continue
 
         if key in _SETTINGS_FIELDS:
-            settings[key] = value
+            validated_settings[key] = value
         else:
-            config[key] = value
+            validated_root[key] = value
 
     if errors:
-        return web.json_response({"error": errors}, status=400)
+        return await handle_api_error(request, errors, 400)
 
-    if save_config(config):
-        return web.json_response({"status": "success", "config": config})
-    return await handle_api_error(request, "Failed to save config")
+    try:
+        async with config_transaction() as config:
+            settings = config.setdefault("settings", {})
+            settings.update(validated_settings)
+            for key, value in validated_root.items():
+                config[key] = value
+            return web.json_response({"status": "success", "config": config})
+    except Exception as e:
+        return await handle_api_error(request, e)
 
 
 @server.PromptServer.instance.routes.get("/distributed/queue_status/{job_id}")
@@ -93,7 +129,7 @@ async def queue_status_endpoint(request):
         job_id = request.match_info['job_id']
         
         # Import to ensure initialization
-        from ..nodes.distributed_upscale import ensure_tile_jobs_initialized
+        from ..upscale.job_store import ensure_tile_jobs_initialized
         prompt_server = ensure_tile_jobs_initialized()
         
         async with prompt_server.distributed_tile_jobs_lock:
@@ -112,72 +148,39 @@ async def update_worker_endpoint(request):
         
         if worker_id is None:
             return await handle_api_error(request, "Missing worker_id", 400)
-            
-        config = load_config()
-        worker_found = False
-        
-        for worker in config.get("workers", []):
-            if worker["id"] == worker_id:
-                # Update all provided fields
-                if "enabled" in data:
-                    worker["enabled"] = data["enabled"]
-                if "name" in data:
-                    worker["name"] = data["name"]
-                if "port" in data:
-                    worker["port"] = data["port"]
-                    
-                # Handle host field - remove it if None
-                if "host" in data:
-                    if data["host"] is None:
-                        worker.pop("host", None)
-                    else:
-                        worker["host"] = normalize_host(data["host"])
-                        
-                # Handle cuda_device field - remove it if None
-                if "cuda_device" in data:
-                    if data["cuda_device"] is None:
-                        worker.pop("cuda_device", None)
-                    else:
-                        worker["cuda_device"] = data["cuda_device"]
-                        
-                # Handle extra_args field - remove it if None
-                if "extra_args" in data:
-                    if data["extra_args"] is None:
-                        worker.pop("extra_args", None)
-                    else:
-                        worker["extra_args"] = data["extra_args"]
-                        
-                # Handle type field
-                if "type" in data:
-                    worker["type"] = data["type"]
-                        
-                worker_found = True
-                break
-                
-        if not worker_found:
-            # If worker not found and all required fields are provided, create new worker
-            if all(key in data for key in ["name", "port", "cuda_device"]):
-                new_worker = {
-                    "id": worker_id,
-                    "name": data["name"],
-                    "host": normalize_host(data.get("host", "localhost")),
-                    "port": data["port"],
-                    "cuda_device": data["cuda_device"],
-                    "enabled": data.get("enabled", False),
-                    "extra_args": data.get("extra_args", ""),
-                    "type": data.get("type", "local")
-                }
-                if "workers" not in config:
-                    config["workers"] = []
-                config["workers"].append(new_worker)
-                worker_found = True
-            else:
-                return await handle_api_error(request, f"Worker {worker_id} not found and missing required fields for creation", 404)
-            
-        if save_config(config):
+
+        async with config_transaction() as config:
+            worker_found = False
+            workers = config.setdefault("workers", [])
+
+            for worker in workers:
+                if worker["id"] == worker_id:
+                    _apply_field_patch(worker, data, _WORKER_FIELDS)
+                    worker_found = True
+                    break
+
+            if not worker_found:
+                # If worker not found and all required fields are provided, create new worker
+                if all(key in data for key in ["name", "port", "cuda_device"]):
+                    new_worker = {
+                        "id": worker_id,
+                        "name": data["name"],
+                        "host": normalize_host(data.get("host", "localhost")),
+                        "port": data["port"],
+                        "cuda_device": data["cuda_device"],
+                        "enabled": data.get("enabled", False),
+                        "extra_args": data.get("extra_args", ""),
+                        "type": data.get("type", "local")
+                    }
+                    workers.append(new_worker)
+                else:
+                    return await handle_api_error(
+                        request,
+                        f"Worker {worker_id} not found and missing required fields for creation",
+                        404,
+                    )
+
             return web.json_response({"status": "success"})
-        else:
-            return await handle_api_error(request, "Failed to save config")
     except Exception as e:
         return await handle_api_error(request, e, 400)
 
@@ -190,29 +193,26 @@ async def delete_worker_endpoint(request):
         if worker_id is None:
             return await handle_api_error(request, "Missing worker_id", 400)
             
-        config = load_config()
-        workers = config.get("workers", [])
-        
-        # Find and remove the worker
-        worker_index = -1
-        for i, worker in enumerate(workers):
-            if worker["id"] == worker_id:
-                worker_index = i
-                break
-                
-        if worker_index == -1:
-            return await handle_api_error(request, f"Worker {worker_id} not found", 404)
-            
-        # Remove the worker
-        removed_worker = workers.pop(worker_index)
-        
-        if save_config(config):
+        async with config_transaction() as config:
+            workers = config.get("workers", [])
+
+            # Find and remove the worker
+            worker_index = -1
+            for i, worker in enumerate(workers):
+                if worker["id"] == worker_id:
+                    worker_index = i
+                    break
+
+            if worker_index == -1:
+                return await handle_api_error(request, f"Worker {worker_id} not found", 404)
+
+            # Remove the worker
+            removed_worker = workers.pop(worker_index)
+
             return web.json_response({
                 "status": "success",
                 "message": f"Worker {removed_worker.get('name', worker_id)} deleted"
             })
-        else:
-            return await handle_api_error(request, "Failed to save config")
     except Exception as e:
         return await handle_api_error(request, e, 400)
 
@@ -229,16 +229,13 @@ async def update_setting_endpoint(request):
         if key not in _SETTINGS_FIELDS:
             return await handle_api_error(request, f"Unknown setting: {key}", 400)
 
-        config = load_config()
-        if 'settings' not in config:
-            config['settings'] = {}
-        
-        config['settings'][key] = value
+        async with config_transaction() as config:
+            if 'settings' not in config:
+                config['settings'] = {}
+
+            config['settings'][key] = value
 
-        if save_config(config):
             return web.json_response({"status": "success", "message": f"Setting '{key}' updated."})
-        else:
-            return await handle_api_error(request, "Failed to save config")
     except Exception as e:
         return await handle_api_error(request, e, 400)
 
@@ -248,28 +245,11 @@ async def update_master_endpoint(request):
     try:
         data = await request.json()
         
-        config = load_config()
-        if 'master' not in config:
-            config['master'] = {}
-        
-        # Update all provided fields
-        if "name" in data:
-            config['master']['name'] = data['name']
-        if "host" in data:
-            if data["host"] is None:
-                config['master'].pop('host', None)
-            else:
-                config['master']['host'] = normalize_host(data["host"])
-        if "port" in data:
-            config['master']['port'] = data['port']
-        if "cuda_device" in data:
-            config['master']['cuda_device'] = data['cuda_device']
-        if "extra_args" in data:
-            config['master']['extra_args'] = data['extra_args']
-            
-        if save_config(config):
+        async with config_transaction() as config:
+            if 'master' not in config:
+                config['master'] = {}
+            _apply_field_patch(config['master'], data, _MASTER_FIELDS)
+
             return web.json_response({"status": "success", "message": "Master configuration updated."})
-        else:
-            return await handle_api_error(request, "Failed to save config")
     except Exception as e:
         return await handle_api_error(request, e, 400)

api/job_routes.py

@@ -8,7 +8,6 @@ import time
 
 from aiohttp import web
 import server
-import server as _server
 import torch
 from PIL import Image
 
@@ -17,10 +16,10 @@ from ..utils.image import pil_to_tensor, ensure_contiguous
 from ..utils.network import handle_api_error
 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_orchestration import ensure_distributed_state, orchestrate_distributed_execution
 from .queue_request import parse_queue_request_payload
 
-prompt_server = _server.PromptServer.instance
+prompt_server = server.PromptServer.instance
 
 # Canonical worker result envelope accepted by POST /distributed/job_complete:
 # { "job_id": str, "worker_id": str, "batch_idx": int, "image": <base64 PNG>, "is_last": bool }
@@ -136,6 +135,7 @@ async def prepare_job_endpoint(request):
         if not multi_job_id:
             return await handle_api_error(request, "Missing multi_job_id", 400)
 
+        ensure_distributed_state()
         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()
@@ -285,7 +285,7 @@ async def job_complete_endpoint(request):
         if not isinstance(is_last, bool):
             errors.append("is_last: expected boolean")
         if errors:
-            return web.json_response({"error": errors}, status=400)
+            return await handle_api_error(request, errors, 400)
 
         tensor = _decode_canonical_png_tensor(image_payload)
         multi_job_id = job_id.strip()

api/orchestration/dispatch.py

@@ -6,18 +6,8 @@ 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}")
+from ...utils.trace_logger import trace_debug, trace_info
+from ..schemas import parse_positive_int
 
 
 async def worker_is_active(worker):
@@ -45,14 +35,6 @@ async def worker_ws_is_active(worker):
         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))
@@ -120,7 +102,7 @@ async def dispatch_worker_prompt(
         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}")
+                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
@@ -142,11 +124,11 @@ async def select_active_workers(
     probe_concurrency=8,
 ):
     """Probe workers and return (active_workers, updated_delegate_master)."""
-    probe_limit = _coerce_positive_int(probe_concurrency, 8)
+    probe_limit = parse_positive_int(probe_concurrency, 8)
     probe_semaphore = asyncio.Semaphore(probe_limit)
 
     if trace_execution_id and workers:
-        _trace_debug(
+        trace_debug(
             trace_execution_id,
             f"Probing {len(workers)} workers with probe_concurrency={probe_limit}",
         )
@@ -164,19 +146,19 @@ async def select_active_workers(
             active_workers.append(worker)
         else:
             if trace_execution_id:
-                _trace_info(trace_execution_id, f"Worker {worker['name']} ({worker['id']}) is offline, skipping.")
+                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_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.")
+            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

api/orchestration/media_sync.py

@@ -8,6 +8,7 @@ import aiohttp
 
 from ...utils.logging import debug_log, log
 from ...utils.network import build_worker_url, get_client_session
+from ...utils.trace_logger import trace_debug, trace_info
 
 
 LIKELY_FILENAME_RE = re.compile(
@@ -20,19 +21,6 @@ IMAGE_OR_VIDEO_RE = re.compile(
     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 ("/", "\\"):
@@ -114,7 +102,7 @@ async def fetch_worker_path_separator(worker, trace_execution_id=None):
             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}")
+            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
@@ -138,7 +126,7 @@ async def _upload_media_to_worker(worker, filename, file_bytes, file_hash, mime_
                     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}")
+            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}")
 
@@ -180,13 +168,13 @@ async def sync_worker_media(worker, prompt_obj, trace_execution_id=None):
         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.")
+                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}")
+                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
@@ -206,7 +194,7 @@ async def sync_worker_media(worker, prompt_obj, trace_execution_id=None):
                 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}")
+                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}")
 
@@ -215,6 +203,6 @@ async def sync_worker_media(worker, prompt_obj, trace_execution_id=None):
         f"uploaded={uploaded}, skipped={skipped}, missing={missing}, referenced={len(media_refs)}"
     )
     if trace_execution_id:
-        _trace_debug(trace_execution_id, summary)
+        trace_debug(trace_execution_id, summary)
     else:
         debug_log(f"[Distributed] {summary}")

api/queue_orchestration.py

@@ -14,6 +14,8 @@ from ..utils.constants import (
 )
 from ..utils.logging import debug_log, log
 from ..utils.network import build_master_url
+from ..utils.trace_logger import trace_debug
+from .schemas import parse_positive_float, parse_positive_int
 from .orchestration.dispatch import dispatch_worker_prompt, select_active_workers
 from .orchestration.media_sync import convert_paths_for_platform, fetch_worker_path_separator, sync_worker_media
 from .orchestration.prompt_transform import (
@@ -33,20 +35,17 @@ def _generate_execution_trace_id():
     return f"exec_{int(time.time() * 1000)}_{uuid.uuid4().hex[:6]}"
 
 
-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 ensure_distributed_state():
+def ensure_distributed_state(server_instance=None):
     """Ensure prompt_server has the state used by distributed queue orchestration."""
-    if not hasattr(prompt_server, "distributed_pending_jobs"):
-        prompt_server.distributed_pending_jobs = {}
-    if not hasattr(prompt_server, "distributed_jobs_lock"):
-        prompt_server.distributed_jobs_lock = asyncio.Lock()
+    ps = server_instance or prompt_server
+    if not hasattr(ps, "distributed_pending_jobs"):
+        ps.distributed_pending_jobs = {}
+    if not hasattr(ps, "distributed_jobs_lock"):
+        ps.distributed_jobs_lock = asyncio.Lock()
+
+
+# Initialize top-level distributed queue state at module import time.
+ensure_distributed_state()
 
 
 async def _ensure_distributed_queue(job_id):
@@ -90,40 +89,22 @@ def _resolve_enabled_workers(config, requested_ids=None):
     return workers
 
 
-def _parse_positive_int(value, default):
-    """Parse positive integer setting values with fallback."""
-    try:
-        parsed = int(value)
-    except (TypeError, ValueError):
-        return default
-    return max(1, parsed)
-
-
-def _parse_positive_float(value, default):
-    """Parse positive float setting values with fallback."""
-    try:
-        parsed = float(value)
-    except (TypeError, ValueError):
-        return default
-    return max(1.0, parsed)
-
-
 def _resolve_orchestration_limits(config):
     """Resolve bounded concurrency/timeouts for worker preparation pipeline."""
     settings = (config or {}).get("settings", {}) or {}
-    worker_probe_concurrency = _parse_positive_int(
+    worker_probe_concurrency = parse_positive_int(
         settings.get("worker_probe_concurrency"),
         ORCHESTRATION_WORKER_PROBE_CONCURRENCY,
     )
-    worker_prep_concurrency = _parse_positive_int(
+    worker_prep_concurrency = parse_positive_int(
         settings.get("worker_prep_concurrency"),
         ORCHESTRATION_WORKER_PREP_CONCURRENCY,
     )
-    media_sync_concurrency = _parse_positive_int(
+    media_sync_concurrency = parse_positive_int(
         settings.get("media_sync_concurrency"),
         ORCHESTRATION_MEDIA_SYNC_CONCURRENCY,
     )
-    media_sync_timeout_seconds = _parse_positive_float(
+    media_sync_timeout_seconds = parse_positive_float(
         settings.get("media_sync_timeout_seconds"),
         ORCHESTRATION_MEDIA_SYNC_TIMEOUT,
     )
@@ -176,7 +157,7 @@ async def _prepare_worker_payload(
                         timeout=media_sync_timeout_seconds,
                     )
                 except asyncio.TimeoutError:
-                    _trace_debug(
+                    trace_debug(
                         trace_execution_id,
                         (
                             f"Media sync timed out after {media_sync_timeout_seconds:.1f}s "
@@ -214,7 +195,7 @@ async def orchestrate_distributed_execution(
     requested_ids = enabled_worker_ids if enabled_worker_ids is not None else None
     workers = _resolve_enabled_workers(config, requested_ids)
     prompt_index = PromptIndex(prompt_obj)
-    _trace_debug(
+    trace_debug(
         execution_trace_id,
         (
             f"Orchestration start: requested_workers={len(workers)}, "
@@ -232,7 +213,7 @@ async def orchestrate_distributed_execution(
         delegate_master = bool(config.get("settings", {}).get("master_delegate_only", False))
 
     if not workers and delegate_master:
-        _trace_debug(
+        trace_debug(
             execution_trace_id,
             "Delegate-only requested but no workers are enabled. Falling back to master execution.",
         )
@@ -252,7 +233,7 @@ async def orchestrate_distributed_execution(
     job_id_map = generate_job_id_map(prompt_index, discovery_prefix)
 
     if not job_id_map:
-        _trace_debug(execution_trace_id, "No distributed nodes detected; queueing prompt on master only.")
+        trace_debug(execution_trace_id, "No distributed nodes detected; queueing prompt on master only.")
         prompt_id = await queue_prompt_payload(prompt_obj, workflow_meta, client_id)
         return prompt_id, 0
 
@@ -286,7 +267,7 @@ async def orchestrate_distributed_execution(
             master_prompt = prepare_delegate_master_prompt(master_prompt, collector_ids)
 
     if active_workers:
-        _trace_debug(
+        trace_debug(
             execution_trace_id,
             "Active distributed workers: "
             + ", ".join(f"{worker['name']} ({worker['id']})" for worker in active_workers),
@@ -329,7 +310,7 @@ async def orchestrate_distributed_execution(
         )
 
     prompt_id = await queue_prompt_payload(master_prompt, workflow_meta, client_id)
-    _trace_debug(
+    trace_debug(
         execution_trace_id,
         f"Orchestration complete: prompt_id={prompt_id}, dispatched_workers={len(worker_payloads)}, delegate_master={delegate_master}",
     )

api/schemas.py

@@ -34,3 +34,21 @@ def validate_positive_int(value, field_name: str) -> str | None:
     if parsed <= 0:
         return f"Field '{field_name}' must be a positive integer."
     return None
+
+
+def parse_positive_int(value, default: int) -> int:
+    """Parse value as positive int, returning default on failure."""
+    try:
+        parsed = int(value)
+    except (TypeError, ValueError):
+        return max(1, int(default))
+    return max(1, parsed)
+
+
+def parse_positive_float(value, default: float) -> float:
+    """Parse value as positive float, returning default on failure."""
+    try:
+        parsed = float(value)
+    except (TypeError, ValueError):
+        return max(0.0, float(default))
+    return max(0.0, parsed)

api/usdu_routes.py

@@ -0,0 +1,228 @@
+import asyncio
+import io
+import time
+
+from aiohttp import web
+from PIL import Image
+import server
+
+from ..upscale.job_models import BaseJobState, ImageJobState, TileJobState
+from ..upscale.job_store import MAX_PAYLOAD_SIZE, ensure_tile_jobs_initialized
+from ..upscale.payload_parsers import _parse_tiles_from_form
+from ..utils.logging import debug_log
+from ..utils.network import handle_api_error
+
+
+@server.PromptServer.instance.routes.post("/distributed/heartbeat")
+async def heartbeat_endpoint(request):
+    try:
+        data = await request.json()
+        worker_id = data.get('worker_id')
+        multi_job_id = data.get('multi_job_id')
+
+        if not worker_id or not multi_job_id:
+            return await handle_api_error(request, "Missing worker_id or multi_job_id", 400)
+
+        prompt_server = ensure_tile_jobs_initialized()
+        async with prompt_server.distributed_tile_jobs_lock:
+            if multi_job_id in prompt_server.distributed_pending_tile_jobs:
+                job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
+                if isinstance(job_data, BaseJobState):
+                    job_data.worker_status[worker_id] = time.time()
+                    debug_log(f"Heartbeat from worker {worker_id}")
+                    return web.json_response({"status": "success"})
+                return await handle_api_error(request, "Worker status tracking not available", 400)
+            return await handle_api_error(request, "Job not found", 404)
+    except Exception as e:
+        return await handle_api_error(request, e, 500)
+
+
+@server.PromptServer.instance.routes.post("/distributed/submit_tiles")
+async def submit_tiles_endpoint(request):
+    """Endpoint for workers to submit processed tiles in static mode."""
+    try:
+        content_length = request.headers.get('content-length')
+        if content_length and int(content_length) > MAX_PAYLOAD_SIZE:
+            return await handle_api_error(request, f"Payload too large: {content_length} bytes", 413)
+
+        data = await request.post()
+        multi_job_id = data.get('multi_job_id')
+        worker_id = data.get('worker_id')
+        is_last = data.get('is_last', 'False').lower() == 'true'
+
+        if multi_job_id is None or worker_id is None:
+            return await handle_api_error(request, "Missing multi_job_id or worker_id", 400)
+
+        prompt_server = ensure_tile_jobs_initialized()
+
+        batch_size = int(data.get('batch_size', 0))
+
+        # Handle completion signal
+        if batch_size == 0 and is_last:
+            async with prompt_server.distributed_tile_jobs_lock:
+                if multi_job_id in prompt_server.distributed_pending_tile_jobs:
+                    job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
+                    if not isinstance(job_data, TileJobState):
+                        return await handle_api_error(request, "Job not configured for tile submissions", 400)
+                    await job_data.queue.put({
+                        'worker_id': worker_id,
+                        'is_last': True,
+                        'tiles': [],
+                    })
+                    debug_log(f"Received completion signal from worker {worker_id}")
+                    return web.json_response({"status": "success"})
+
+        try:
+            tiles = _parse_tiles_from_form(data)
+        except ValueError as e:
+            return await handle_api_error(request, str(e), 400)
+
+        # Submit tiles to queue
+        async with prompt_server.distributed_tile_jobs_lock:
+            if multi_job_id in prompt_server.distributed_pending_tile_jobs:
+                job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
+                if not isinstance(job_data, TileJobState):
+                    return await handle_api_error(request, "Job not configured for tile submissions", 400)
+
+                q = job_data.queue
+                if batch_size > 0 or len(tiles) > 0:
+                    await q.put({
+                        'worker_id': worker_id,
+                        'tiles': tiles,
+                        'is_last': is_last,
+                    })
+                    debug_log(f"Received {len(tiles)} tiles from worker {worker_id} (is_last={is_last})")
+                else:
+                    await q.put({
+                        'worker_id': worker_id,
+                        'is_last': True,
+                        'tiles': [],
+                    })
+
+                return web.json_response({"status": "success"})
+            return await handle_api_error(request, "Job not found", 404)
+    except Exception as e:
+        return await handle_api_error(request, e, 500)
+
+
+@server.PromptServer.instance.routes.post("/distributed/submit_image")
+async def submit_image_endpoint(request):
+    """Endpoint for workers to submit processed images in dynamic mode."""
+    try:
+        content_length = request.headers.get('content-length')
+        if content_length and int(content_length) > MAX_PAYLOAD_SIZE:
+            return await handle_api_error(request, f"Payload too large: {content_length} bytes", 413)
+
+        data = await request.post()
+        multi_job_id = data.get('multi_job_id')
+        worker_id = data.get('worker_id')
+        is_last = data.get('is_last', 'False').lower() == 'true'
+
+        if multi_job_id is None or worker_id is None:
+            return await handle_api_error(request, "Missing multi_job_id or worker_id", 400)
+
+        prompt_server = ensure_tile_jobs_initialized()
+
+        # Handle image submission
+        if 'full_image' in data and 'image_idx' in data:
+            image_idx = int(data.get('image_idx'))
+            img_data = data['full_image'].file.read()
+            img = Image.open(io.BytesIO(img_data)).convert("RGB")
+
+            debug_log(f"Received full image {image_idx} from worker {worker_id}")
+
+            async with prompt_server.distributed_tile_jobs_lock:
+                if multi_job_id in prompt_server.distributed_pending_tile_jobs:
+                    job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
+                    if not isinstance(job_data, ImageJobState):
+                        return await handle_api_error(request, "Job not configured for image submissions", 400)
+                    await job_data.queue.put({
+                        'worker_id': worker_id,
+                        'image_idx': image_idx,
+                        'image': img,
+                        'is_last': is_last,
+                    })
+                    return web.json_response({"status": "success"})
+
+        # Handle completion signal (no image data)
+        elif is_last:
+            async with prompt_server.distributed_tile_jobs_lock:
+                if multi_job_id in prompt_server.distributed_pending_tile_jobs:
+                    job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
+                    if not isinstance(job_data, ImageJobState):
+                        return await handle_api_error(request, "Job not configured for image submissions", 400)
+                    await job_data.queue.put({
+                        'worker_id': worker_id,
+                        'is_last': True,
+                    })
+                    debug_log(f"Received completion signal from worker {worker_id}")
+                    return web.json_response({"status": "success"})
+        else:
+            return await handle_api_error(request, "Missing image data or invalid request", 400)
+
+        return await handle_api_error(request, "Job not found", 404)
+    except Exception as e:
+        return await handle_api_error(request, e, 500)
+
+
+@server.PromptServer.instance.routes.post("/distributed/request_image")
+async def request_image_endpoint(request):
+    """Endpoint for workers to request tasks (images in dynamic mode, tiles in static mode)."""
+    try:
+        data = await request.json()
+        worker_id = data.get('worker_id')
+        multi_job_id = data.get('multi_job_id')
+
+        if not worker_id or not multi_job_id:
+            return await handle_api_error(request, "Missing worker_id or multi_job_id", 400)
+
+        prompt_server = ensure_tile_jobs_initialized()
+        async with prompt_server.distributed_tile_jobs_lock:
+            if multi_job_id in prompt_server.distributed_pending_tile_jobs:
+                job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
+                if not isinstance(job_data, BaseJobState):
+                    return await handle_api_error(request, "Invalid job data structure", 500)
+
+                mode = job_data.mode
+                if isinstance(job_data, ImageJobState):
+                    pending_queue = job_data.pending_images
+                elif isinstance(job_data, TileJobState):
+                    pending_queue = job_data.pending_tasks
+                else:
+                    return await handle_api_error(request, "Invalid job configuration", 400)
+
+                try:
+                    task_idx = await asyncio.wait_for(pending_queue.get(), timeout=0.1)
+                    job_data.assigned_to_workers.setdefault(worker_id, []).append(task_idx)
+                    job_data.worker_status[worker_id] = time.time()
+                    remaining = pending_queue.qsize()
+
+                    if mode == 'dynamic':
+                        debug_log(f"UltimateSDUpscale API - Assigned image {task_idx} to worker {worker_id}")
+                        return web.json_response({"image_idx": task_idx, "estimated_remaining": remaining})
+                    debug_log(f"UltimateSDUpscale API - Assigned tile {task_idx} to worker {worker_id}")
+                    return web.json_response({
+                        "tile_idx": task_idx,
+                        "estimated_remaining": remaining,
+                        "batched_static": job_data.batched_static,
+                    })
+                except asyncio.TimeoutError:
+                    if mode == 'dynamic':
+                        return web.json_response({"image_idx": None})
+                    return web.json_response({"tile_idx": None})
+            return await handle_api_error(request, "Job not found", 404)
+    except Exception as e:
+        return await handle_api_error(request, e, 500)
+
+
+@server.PromptServer.instance.routes.get("/distributed/job_status")
+async def job_status_endpoint(request):
+    """Endpoint to check if a job is ready."""
+    multi_job_id = request.query.get('multi_job_id')
+    if not multi_job_id:
+        return web.json_response({"ready": False})
+    prompt_server = ensure_tile_jobs_initialized()
+    async with prompt_server.distributed_tile_jobs_lock:
+        job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
+        ready = bool(isinstance(job_data, BaseJobState) and job_data.queue is not None)
+        return web.json_response({"ready": ready})

api/worker_routes.py

@@ -18,12 +18,9 @@ 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,
     get_machine_id,
     is_docker_environment,
     is_runpod_environment,
-    get_comms_channel,
 )
 from ..utils.async_helpers import queue_prompt_payload
 
@@ -317,13 +314,7 @@ async def get_network_info_endpoint(request):
             "message": "Auto-detected network configuration"
         })
     except Exception as e:
-        return web.json_response({
-            "status": "error",
-            "message": str(e),
-            "hostname": "unknown",
-            "all_ips": [],
-            "recommended_ip": None
-        })
+        return await handle_api_error(request, e, 500)
 
 @server.PromptServer.instance.routes.get("/distributed/system_info")
 async def get_system_info_endpoint(request):
@@ -347,10 +338,7 @@ async def get_system_info_endpoint(request):
             "runpod_pod_id": os.environ.get('RUNPOD_POD_ID')
         })
     except Exception as e:
-        return web.json_response({
-            "status": "error",
-            "message": str(e)
-        }, status=500)
+        return await handle_api_error(request, e, 500)
 
 @server.PromptServer.instance.routes.post("/distributed/launch_worker")
 async def launch_worker_endpoint(request):
@@ -389,12 +377,7 @@ async def launch_worker_endpoint(request):
                 is_running = wm._is_process_running(proc_info['pid'])
             
             if is_running:
-                return web.json_response({
-                    "status": "error",
-                    "message": "Worker already running (managed by UI)",
-                    "pid": proc_info['pid'],
-                    "log_file": proc_info.get('log_file')
-                }, status=409)
+                return await handle_api_error(request, "Worker already running (managed by UI)", 409)
             else:
                 # Process is dead, remove it
                 del wm.processes[worker_id_str]
@@ -438,8 +421,11 @@ async def stop_worker_endpoint(request):
         if success:
             return web.json_response({"status": "success", "message": message})
         else:
-            return web.json_response({"status": "error", "message": message}, 
-                                   status=404 if "not managed" in message else 409)
+            return await handle_api_error(
+                request,
+                message,
+                404 if "not managed" in message else 409,
+            )
             
     except Exception as e:
         return await handle_api_error(request, e, 400)

distributed.py

@@ -2,7 +2,6 @@
 ComfyUI-Distributed: thin entry point.
 All implementation lives in workers/, nodes/, api/.
 """
-import asyncio
 import atexit
 import os
 
@@ -28,6 +27,7 @@ from .nodes import (
     any_type,
 )
 from . import api  # noqa: F401 - triggers all @routes.* registrations
+from .api.queue_orchestration import ensure_distributed_state
 
 ensure_config_exists()
 
@@ -40,9 +40,7 @@ atexit.register(lambda: None)  # placeholder; real cleanup in sync_cleanup
 
 # Initialize distributed job state on prompt_server
 prompt_server = server.PromptServer.instance
-if not hasattr(prompt_server, 'distributed_pending_jobs'):
-    prompt_server.distributed_pending_jobs = {}
-    prompt_server.distributed_jobs_lock = asyncio.Lock()
+ensure_distributed_state(prompt_server)
 
 if not hasattr(prompt_server, 'distributed_pending_tile_jobs'):
     prompt_server.distributed_pending_tile_jobs = {}

nodes/collector.py

@@ -17,7 +17,6 @@ from ..utils.constants import HEARTBEAT_INTERVAL
 from ..utils.image import tensor_to_pil, pil_to_tensor, ensure_contiguous
 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
 
 prompt_server = _server.PromptServer.instance
 

nodes/distributed_upscale.py

@@ -6,7 +6,7 @@ import comfy.samplers
 
 from ..utils.logging import debug_log, log
 from ..utils.async_helpers import run_async_in_server_loop
-from ..utils.usdu_managment import ensure_tile_jobs_initialized
+from ..upscale.job_store import ensure_tile_jobs_initialized
 
 from ..upscale.tile_ops import TileOpsMixin
 from ..upscale.result_collector import ResultCollectorMixin

upscale/conditioning.py

@@ -0,0 +1,32 @@
+import copy
+
+
+def clone_control_chain(control, clone_hint=True):
+    """Shallow copy the ControlNet chain, optionally cloning hints but sharing models."""
+    if control is None:
+        return None
+    new_control = copy.copy(control)
+    if clone_hint and hasattr(control, 'cond_hint_original'):
+        hint = getattr(control, 'cond_hint_original', None)
+        new_control.cond_hint_original = hint.clone() if hint is not None else None
+    if hasattr(control, 'previous_controlnet'):
+        new_control.previous_controlnet = clone_control_chain(control.previous_controlnet, clone_hint)
+    return new_control
+
+
+def clone_conditioning(cond_list, clone_hints=True):
+    """Clone conditioning without duplicating ControlNet models."""
+    new_cond = []
+    for emb, cond_dict in cond_list:
+        new_emb = emb.clone() if emb is not None else None
+        new_dict = cond_dict.copy()
+        if 'control' in new_dict:
+            new_dict['control'] = clone_control_chain(new_dict['control'], clone_hints)
+        if 'mask' in new_dict and new_dict['mask'] is not None:
+            new_dict['mask'] = new_dict['mask'].clone()
+        if 'pooled_output' in new_dict and new_dict['pooled_output'] is not None:
+            new_dict['pooled_output'] = new_dict['pooled_output'].clone()
+        if 'area' in new_dict:
+            new_dict['area'] = new_dict['area'][:]
+        new_cond.append([new_emb, new_dict])
+    return new_cond

upscale/job_state.py

@@ -1,10 +1,8 @@
 import asyncio
 
 from ..utils.logging import debug_log
-from ..utils.usdu_managment import (
-    _check_and_requeue_timed_out_workers as _requeue_usdu,
-    ensure_tile_jobs_initialized,
-)
+from .job_store import ensure_tile_jobs_initialized
+from .job_timeout import _check_and_requeue_timed_out_workers as _requeue_usdu
 from .job_models import ImageJobState, TileJobState
 
 

upscale/job_store.py

@@ -0,0 +1,180 @@
+import asyncio
+import os
+import time
+from typing import List, Optional
+
+import server
+
+from ..utils.logging import debug_log
+from .job_models import BaseJobState, ImageJobState, TileJobState
+
+# Configure maximum payload size (50MB default, configurable via environment variable)
+MAX_PAYLOAD_SIZE = int(os.environ.get('COMFYUI_MAX_PAYLOAD_SIZE', str(50 * 1024 * 1024)))
+
+
+def ensure_tile_jobs_initialized():
+    """Ensure tile job storage is initialized on the server instance."""
+    prompt_server = server.PromptServer.instance
+    if not hasattr(prompt_server, 'distributed_pending_tile_jobs'):
+        debug_log("Initializing persistent tile job queue on server instance.")
+        prompt_server.distributed_pending_tile_jobs = {}
+        prompt_server.distributed_tile_jobs_lock = asyncio.Lock()
+    else:
+        invalid_job_ids = [
+            job_id
+            for job_id, job_data in prompt_server.distributed_pending_tile_jobs.items()
+            if not isinstance(job_data, BaseJobState)
+        ]
+        for job_id in invalid_job_ids:
+            debug_log(f"Removing invalid job state for {job_id}")
+            del prompt_server.distributed_pending_tile_jobs[job_id]
+    return prompt_server
+
+
+async def _init_job_queue(
+    multi_job_id,
+    mode,
+    batch_size=None,
+    num_tiles_per_image=None,
+    all_indices=None,
+    enabled_workers=None,
+    batched_static: bool = False,
+):
+    """Unified initialization for job queues in static and dynamic modes."""
+    prompt_server = ensure_tile_jobs_initialized()
+    async with prompt_server.distributed_tile_jobs_lock:
+        if multi_job_id in prompt_server.distributed_pending_tile_jobs:
+            debug_log(f"Queue already exists for {multi_job_id}")
+            return
+
+        if mode == 'dynamic':
+            job_data = ImageJobState(multi_job_id=multi_job_id)
+        elif mode == 'static':
+            job_data = TileJobState(multi_job_id=multi_job_id)
+        else:
+            raise ValueError(f"Unknown mode: {mode}")
+
+        job_data.worker_status = {w: time.time() for w in enabled_workers or []}
+        job_data.assigned_to_workers = {w: [] for w in enabled_workers or []}
+
+        if mode == 'dynamic':
+            job_data.batch_size = int(batch_size or 0)
+            pending_queue = job_data.pending_images
+            for i in (all_indices or range(int(batch_size or 0))):
+                await pending_queue.put(i)
+            debug_log(f"Initialized image queue with {batch_size} pending items")
+        elif mode == 'static':
+            job_data.num_tiles_per_image = int(num_tiles_per_image or 0)
+            job_data.batch_size = int(batch_size or 0)
+            job_data.batched_static = bool(batched_static)
+            # For batched static distribution, populate only tile ids [0..num_tiles_per_image-1]
+            pending_queue = job_data.pending_tasks
+            if batched_static and num_tiles_per_image is not None:
+                for i in range(num_tiles_per_image):
+                    await pending_queue.put(i)
+            else:
+                total_tiles = int(batch_size or 0) * int(num_tiles_per_image or 0)
+                for i in range(total_tiles):
+                    await pending_queue.put(i)
+
+        prompt_server.distributed_pending_tile_jobs[multi_job_id] = job_data
+
+
+async def init_dynamic_job(
+    multi_job_id: str,
+    batch_size: int,
+    enabled_workers: List[str],
+    all_indices: Optional[List[int]] = None,
+):
+    """Initialize queue for dynamic mode (per-image), with collector fields."""
+    await _init_job_queue(
+        multi_job_id,
+        'dynamic',
+        batch_size=batch_size,
+        all_indices=all_indices or list(range(batch_size)),
+        enabled_workers=enabled_workers,
+    )
+    debug_log(f"Job {multi_job_id} initialized with {batch_size} images")
+
+
+async def init_static_job_batched(
+    multi_job_id: str,
+    batch_size: int,
+    num_tiles_per_image: int,
+    enabled_workers: List[str],
+):
+    """Initialize queue for static mode (batched-per-tile)."""
+    await _init_job_queue(
+        multi_job_id,
+        'static',
+        batch_size=batch_size,
+        num_tiles_per_image=num_tiles_per_image,
+        enabled_workers=enabled_workers,
+        batched_static=True,
+    )
+
+
+async def _drain_results_queue(multi_job_id):
+    """Drain pending results from queue and update completed_tasks. Returns count drained."""
+    prompt_server = ensure_tile_jobs_initialized()
+    async with prompt_server.distributed_tile_jobs_lock:
+        job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
+        if not isinstance(job_data, BaseJobState):
+            return 0
+        q = job_data.queue
+        completed_tasks = job_data.completed_tasks
+
+        collected = 0
+        while True:
+            try:
+                result = q.get_nowait()
+            except asyncio.QueueEmpty:
+                break
+
+            worker_id = result['worker_id']
+            is_last = result.get('is_last', False)
+
+            if 'image_idx' in result and 'image' in result:
+                task_id = result['image_idx']
+                if task_id not in completed_tasks:
+                    completed_tasks[task_id] = result['image']
+                    collected += 1
+            elif 'tiles' in result:
+                for tile_data in result['tiles']:
+                    task_id = tile_data.get('global_idx', tile_data['tile_idx'])
+                    if task_id not in completed_tasks:
+                        completed_tasks[task_id] = tile_data
+                        collected += 1
+            if is_last:
+                if worker_id in job_data.worker_status:
+                    del job_data.worker_status[worker_id]
+
+        return collected
+
+
+async def _get_completed_count(multi_job_id):
+    """Get count of completed tasks."""
+    prompt_server = ensure_tile_jobs_initialized()
+    async with prompt_server.distributed_tile_jobs_lock:
+        job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
+        if isinstance(job_data, BaseJobState):
+            return len(job_data.completed_tasks)
+        return 0
+
+
+async def _mark_task_completed(multi_job_id, task_id, result):
+    """Mark a task as completed."""
+    prompt_server = ensure_tile_jobs_initialized()
+    async with prompt_server.distributed_tile_jobs_lock:
+        job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
+        if isinstance(job_data, BaseJobState):
+            job_data.completed_tasks[task_id] = result
+
+
+async def _cleanup_job(multi_job_id):
+    """Cleanup the job data."""
+    prompt_server = ensure_tile_jobs_initialized()
+    async with prompt_server.distributed_tile_jobs_lock:
+        if multi_job_id in prompt_server.distributed_pending_tile_jobs:
+            del prompt_server.distributed_pending_tile_jobs[multi_job_id]
+            debug_log(f"Cleaned up job {multi_job_id}")

upscale/job_timeout.py

@@ -0,0 +1,150 @@
+import time
+
+from ..utils.config import load_config
+from ..utils.constants import HEARTBEAT_TIMEOUT
+from ..utils.logging import debug_log, log
+from ..utils.network import build_worker_url, probe_worker
+from .job_models import BaseJobState
+from .job_store import ensure_tile_jobs_initialized
+
+
+def _find_worker_record(worker_id):
+    """Return worker config entry by id, or None when missing."""
+    workers = load_config().get("workers", [])
+    return next((w for w in workers if str(w.get("id")) == str(worker_id)), None)
+
+
+async def _check_and_requeue_timed_out_workers(multi_job_id, total_tasks):
+    """Check timed out workers and requeue their tasks. Returns requeued count."""
+    prompt_server = ensure_tile_jobs_initialized()
+    current_time = time.time()
+
+    # Allow override via config setting 'worker_timeout_seconds'
+    cfg = load_config()
+    hb_timeout = int(cfg.get("settings", {}).get("worker_timeout_seconds", HEARTBEAT_TIMEOUT))
+
+    # Snapshot timed-out workers and job details under lock.
+    async with prompt_server.distributed_tile_jobs_lock:
+        job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
+        if not isinstance(job_data, BaseJobState):
+            return 0
+
+        completed_tasks_snapshot = set(job_data.completed_tasks.keys())
+        batched_static_snapshot = bool(job_data.batched_static)
+        num_tiles_per_image_snapshot = int(job_data.num_tiles_per_image or 1)
+        batch_size_snapshot = int(job_data.batch_size or 1)
+
+        timed_out_workers = []
+        for worker, last_heartbeat in list(job_data.worker_status.items()):
+            age = current_time - float(last_heartbeat)
+            debug_log(f"Timeout check: worker={worker} age={age:.1f}s threshold={hb_timeout}s")
+            if age > hb_timeout:
+                timed_out_workers.append(
+                    {
+                        "worker_id": worker,
+                        "last_heartbeat": float(last_heartbeat),
+                        "assigned_tasks": list(job_data.assigned_to_workers.get(worker, [])),
+                    }
+                )
+
+    if not timed_out_workers:
+        return 0
+
+    # Probe outside lock to avoid lock contention on network latency.
+    workers_to_requeue = []
+    workers_graced = []
+    for worker_info in timed_out_workers:
+        worker = worker_info["worker_id"]
+        assigned = worker_info["assigned_tasks"]
+        age = current_time - worker_info["last_heartbeat"]
+
+        incomplete_assigned = 0
+        try:
+            if assigned:
+                if batched_static_snapshot:
+                    for task_id in assigned:
+                        for b in range(batch_size_snapshot):
+                            gidx = b * num_tiles_per_image_snapshot + task_id
+                            if gidx not in completed_tasks_snapshot:
+                                incomplete_assigned += 1
+                                break
+                else:
+                    for task_id in assigned:
+                        if task_id not in completed_tasks_snapshot:
+                            incomplete_assigned += 1
+            debug_log(
+                f"Assigned diagnostics: total_assigned={len(assigned)} "
+                f"incomplete_assigned={incomplete_assigned}"
+            )
+        except Exception as e:
+            debug_log(f"Assigned diagnostics failed for worker {worker}: {e}")
+
+        busy = False
+        probe_queue = None
+        try:
+            worker_record = _find_worker_record(worker)
+            if worker_record:
+                worker_url = build_worker_url(worker_record)
+                debug_log(f"Probing worker {worker} at {worker_url}/prompt")
+                payload = await probe_worker(worker_url, timeout=2.0)
+                if payload is not None:
+                    probe_queue = int(payload.get("exec_info", {}).get("queue_remaining", 0))
+                    busy = probe_queue is not None and probe_queue > 0
+            else:
+                debug_log(f"Probe skipped; worker {worker} not found in config")
+        except Exception as e:
+            debug_log(f"Probe failed for worker {worker}: {e}")
+        finally:
+            debug_log(
+                f"Probe diagnostics: online={probe_queue is not None} queue_remaining={probe_queue}"
+            )
+
+        if busy:
+            workers_graced.append(worker)
+            debug_log(f"Heartbeat grace: worker {worker} busy via probe; skipping requeue")
+            continue
+
+        log(f"Worker {worker} heartbeat timed out after {age:.1f}s")
+        workers_to_requeue.append((worker, assigned))
+
+    # Re-acquire lock and apply requeue/cleanup decisions.
+    async with prompt_server.distributed_tile_jobs_lock:
+        job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
+        if not isinstance(job_data, BaseJobState):
+            return 0
+
+        # Refresh heartbeat for workers that we proved are still busy.
+        for worker in workers_graced:
+            if worker in job_data.worker_status:
+                job_data.worker_status[worker] = current_time
+
+        requeued_count = 0
+        completed_tasks = job_data.completed_tasks
+        batched_static = bool(job_data.batched_static)
+        num_tiles_per_image = int(job_data.num_tiles_per_image or 1)
+        batch_size = int(job_data.batch_size or 1)
+        for worker, assigned_snapshot in workers_to_requeue:
+            # Use current assignments if present, falling back to the snapshot.
+            assigned_tasks = list(job_data.assigned_to_workers.get(worker, assigned_snapshot))
+            for task_id in assigned_tasks:
+                # If batched_static, task_id is a tile_idx; consider it complete only if
+                # all corresponding global_idx entries are present in completed_tasks.
+                if batched_static:
+                    all_done = True
+                    for b in range(batch_size):
+                        gidx = b * num_tiles_per_image + task_id
+                        if gidx not in completed_tasks:
+                            all_done = False
+                            break
+                    if not all_done:
+                        await job_data.pending_tasks.put(task_id)
+                        requeued_count += 1
+                else:
+                    if task_id not in completed_tasks:
+                        await job_data.pending_tasks.put(task_id)
+                        requeued_count += 1
+            job_data.worker_status.pop(worker, None)
+            if worker in job_data.assigned_to_workers:
+                job_data.assigned_to_workers[worker] = []
+
+        return requeued_count

upscale/modes/dynamic.py

@@ -6,13 +6,19 @@ from ...utils.image import tensor_to_pil, pil_to_tensor
 from ...utils.async_helpers import run_async_in_server_loop
 from ...utils.config import get_worker_timeout_seconds
 from ...utils.constants import TILE_WAIT_TIMEOUT, TILE_SEND_TIMEOUT
-from ...utils.usdu_managment import (
-    ensure_tile_jobs_initialized, init_dynamic_job,
-    clone_conditioning, _send_heartbeat_to_master,
-)
+from ..job_store import ensure_tile_jobs_initialized, init_dynamic_job
 
 
 class DynamicModeMixin:
+    """
+    Dynamic (per-image queue) USDU mode behaviors for master and worker roles.
+
+    Expected co-mixins on `self`:
+    - TileOpsMixin (`calculate_tiles`, `_slice_conditioning`, `_process_and_blend_tile`).
+    - JobStateMixin (image queue/task completion helpers).
+    - WorkerCommsMixin (`_request_image_from_master`, `_send_full_image_to_master`, `_send_heartbeat_to_master`).
+    """
+
     def process_master_dynamic(self, upscaled_image, model, positive, negative, vae,
                               seed, steps, cfg, sampler_name, scheduler, denoise,
                               tile_width, tile_height, padding, mask_blur,
@@ -71,22 +77,7 @@ class DynamicModeMixin:
                 image_seed = seed + image_idx * 1000
                 
                 # Pre-slice conditioning once per image (not per tile)
-                positive_sliced = clone_conditioning(positive)
-                negative_sliced = clone_conditioning(negative)
-                for cond_list in [positive_sliced, negative_sliced]:
-                    for i in range(len(cond_list)):
-                        emb, cond_dict = cond_list[i]
-                        if emb.shape[0] > 1:
-                            cond_list[i][0] = emb[image_idx:image_idx+1]
-                        if 'control' in cond_dict:
-                            control = cond_dict['control']
-                            while control is not None:
-                                hint = control.cond_hint_original
-                                if hint.shape[0] > 1:
-                                    control.cond_hint_original = hint[image_idx:image_idx+1]
-                                control = control.previous_controlnet
-                        if 'mask' in cond_dict and cond_dict['mask'].shape[0] > 1:
-                            cond_dict['mask'] = cond_dict['mask'][image_idx:image_idx+1]
+                positive_sliced, negative_sliced = self._slice_conditioning(positive, negative, image_idx)
                 
                 for tile_idx, pos in enumerate(all_tiles):
                     local_image = self._process_and_blend_tile(
@@ -268,55 +259,40 @@ class DynamicModeMixin:
             image_seed = seed + image_idx * 1000
 
             # Pre-slice conditioning once per image (not per tile)
-            positive_sliced = clone_conditioning(positive)
-            negative_sliced = clone_conditioning(negative)
-            for cond_list in [positive_sliced, negative_sliced]:
-                for i in range(len(cond_list)):
-                        emb, cond_dict = cond_list[i]
-                        if emb.shape[0] > 1:
-                            cond_list[i][0] = emb[image_idx:image_idx+1]
-                        if 'control' in cond_dict:
-                            control = cond_dict['control']
-                            while control is not None:
-                                hint = control.cond_hint_original
-                                if hint.shape[0] > 1:
-                                    control.cond_hint_original = hint[image_idx:image_idx+1]
-                                control = control.previous_controlnet
-                        if 'mask' in cond_dict and cond_dict['mask'].shape[0] > 1:
-                            cond_dict['mask'] = cond_dict['mask'][image_idx:image_idx+1]
+            positive_sliced, negative_sliced = self._slice_conditioning(positive, negative, image_idx)
 
-                for tile_idx, pos in enumerate(all_tiles):
-                    local_image = self._process_and_blend_tile(
-                        tile_idx, pos, single_tensor, local_image,
-                        model, positive_sliced, negative_sliced, vae, image_seed, steps, cfg,
-                        sampler_name, scheduler, denoise, tile_width, tile_height,
-                        padding, mask_blur, width, height, force_uniform_tiles,
-                        tiled_decode, batch_idx=image_idx
-                    )
-                    run_async_in_server_loop(
-                        _send_heartbeat_to_master(multi_job_id, master_url, worker_id),
-                        timeout=5.0
-                    )
+            for tile_idx, pos in enumerate(all_tiles):
+                local_image = self._process_and_blend_tile(
+                    tile_idx, pos, single_tensor, local_image,
+                    model, positive_sliced, negative_sliced, vae, image_seed, steps, cfg,
+                    sampler_name, scheduler, denoise, tile_width, tile_height,
+                    padding, mask_blur, width, height, force_uniform_tiles,
+                    tiled_decode, batch_idx=image_idx
+                )
+                run_async_in_server_loop(
+                    self._send_heartbeat_to_master(multi_job_id, master_url, worker_id),
+                    timeout=5.0
+                )
 
-                # Send processed image back to master
-                try:
-                    # Use the estimated remaining to determine if this is the last image
-                    is_last = is_last_for_worker
-                    run_async_in_server_loop(
-                        self._send_full_image_to_master(local_image, image_idx, multi_job_id,
-                                                        master_url, worker_id, is_last),
-                        timeout=TILE_SEND_TIMEOUT
-                    )
-                    # Send heartbeat after processing
-                    run_async_in_server_loop(
-                        _send_heartbeat_to_master(multi_job_id, master_url, worker_id),
-                        timeout=5.0
-                    )
-                    if is_last:
-                        break
-                except Exception as e:
-                    log(f"USDU Dist Worker[{worker_id[:8]}] - Error sending image {image_idx}: {e}")
-                    # Continue processing other images
+            # Send processed image back to master
+            try:
+                # Use the estimated remaining to determine if this is the last image
+                is_last = is_last_for_worker
+                run_async_in_server_loop(
+                    self._send_full_image_to_master(local_image, image_idx, multi_job_id,
+                                                    master_url, worker_id, is_last),
+                    timeout=TILE_SEND_TIMEOUT
+                )
+                # Send heartbeat after processing
+                run_async_in_server_loop(
+                    self._send_heartbeat_to_master(multi_job_id, master_url, worker_id),
+                    timeout=5.0
+                )
+                if is_last:
+                    break
+            except Exception as e:
+                log(f"USDU Dist Worker[{worker_id[:8]}] - Error sending image {image_idx}: {e}")
+                # Continue processing other images
 
         # Send final is_last signal
         debug_log(f"Worker[{worker_id[:8]}] processed {processed_count} images, sending completion signal")

upscale/modes/static.py

@@ -5,17 +5,32 @@ from ...utils.logging import debug_log, log
 from ...utils.image import tensor_to_pil, pil_to_tensor
 from ...utils.async_helpers import run_async_in_server_loop
 from ...utils.config import get_worker_timeout_seconds
-from ...utils.constants import HEARTBEAT_INTERVAL, MAX_BATCH, TILE_SEND_TIMEOUT, TILE_WAIT_TIMEOUT
-from ...utils.usdu_managment import (
+from ...utils.constants import (
+    HEARTBEAT_INTERVAL,
+    JOB_POLL_INTERVAL,
+    JOB_POLL_MAX_ATTEMPTS,
+    MAX_BATCH,
+    TILE_SEND_TIMEOUT,
+    TILE_WAIT_TIMEOUT,
+)
+from ..job_store import (
     ensure_tile_jobs_initialized, init_static_job_batched,
-    _mark_task_completed, _cleanup_job, _drain_results_queue,
-    _send_heartbeat_to_master, _get_completed_count,
+    _mark_task_completed, _cleanup_job, _drain_results_queue, _get_completed_count,
 )
 from ..job_models import TileJobState
 
 
 class StaticModeMixin:
-    def _poll_job_ready(self, multi_job_id, master_url, worker_id=None, max_attempts=20):
+    """
+    Static (tile-queue) USDU mode behaviors for master and worker roles.
+
+    Expected co-mixins on `self`:
+    - TileOpsMixin (`calculate_tiles`, tile extract/blend helpers).
+    - JobStateMixin (`_get_next_tile_index`, `_get_all_completed_tasks`, requeue checks).
+    - WorkerCommsMixin (`send_tiles_batch_to_master`, `_request_tile_from_master`, `_send_heartbeat_to_master`).
+    """
+
+    def _poll_job_ready(self, multi_job_id, master_url, worker_id=None, max_attempts=JOB_POLL_MAX_ATTEMPTS):
         """Poll master for job readiness to avoid worker/master initialization race."""
         for attempt in range(max_attempts):
             ready = run_async_in_server_loop(
@@ -28,7 +43,7 @@ class StaticModeMixin:
                 else:
                     debug_log(f"Job {multi_job_id} ready after {attempt} attempts")
                 return True
-            time.sleep(1.0)
+            time.sleep(JOB_POLL_INTERVAL)
         return False
 
     def _extract_and_process_tile(
@@ -194,7 +209,7 @@ class StaticModeMixin:
         log(f"USDU Dist Worker[{worker_id[:8]}]: Canvas {width}x{height} | Tile {tile_width}x{tile_height} | Tiles/image {num_tiles_per_image} | Batch {batch_size}")
         processed_count = 0
 
-        max_poll_attempts = 20
+        max_poll_attempts = JOB_POLL_MAX_ATTEMPTS
         if not self._poll_job_ready(multi_job_id, master_url, worker_id=worker_id, max_attempts=max_poll_attempts):
             log(f"Job {multi_job_id} not ready after {max_poll_attempts} attempts, aborting")
             return (upscaled_image,)
@@ -254,7 +269,7 @@ class StaticModeMixin:
             # Send heartbeat
             try:
                 run_async_in_server_loop(
-                    _send_heartbeat_to_master(multi_job_id, master_url, worker_id),
+                    self._send_heartbeat_to_master(multi_job_id, master_url, worker_id),
                     timeout=5.0
                 )
             except Exception as e:

upscale/payload_parsers.py

@@ -0,0 +1,64 @@
+import io
+import json
+
+from PIL import Image
+
+
+def _parse_tiles_from_form(data):
+    """Parse tiles submitted via multipart/form-data into a list of tile dicts."""
+    try:
+        padding = int(data.get('padding', 0)) if data.get('padding') is not None else 0
+    except Exception:
+        padding = 0
+
+    meta_raw = data.get('tiles_metadata')
+    if meta_raw is None:
+        raise ValueError("Missing tiles_metadata")
+
+    try:
+        metadata = json.loads(meta_raw)
+    except Exception as e:
+        raise ValueError(f"Invalid tiles_metadata JSON: {e}")
+
+    if not isinstance(metadata, list):
+        raise ValueError("tiles_metadata must be a list")
+
+    tiles = []
+    for i, meta in enumerate(metadata):
+        file_field = data.get(f'tile_{i}')
+        if file_field is None or not hasattr(file_field, 'file'):
+            raise ValueError(f"Missing tile data for index {i}")
+
+        raw = file_field.file.read()
+        try:
+            img = Image.open(io.BytesIO(raw)).convert("RGB")
+        except Exception as e:
+            raise ValueError(f"Invalid image data for tile {i}: {e}")
+
+        try:
+            tile_info = {
+                'image': img,
+                'tile_idx': int(meta.get('tile_idx', i)),
+                'x': int(meta.get('x', 0)),
+                'y': int(meta.get('y', 0)),
+                'extracted_width': int(meta.get('extracted_width', img.width)),
+                'extracted_height': int(meta.get('extracted_height', img.height)),
+                'padding': int(padding),
+            }
+        except Exception as e:
+            raise ValueError(f"Invalid metadata values for tile {i}: {e}")
+
+        if 'batch_idx' in meta:
+            try:
+                tile_info['batch_idx'] = int(meta['batch_idx'])
+            except Exception:
+                pass
+        if 'global_idx' in meta:
+            try:
+                tile_info['global_idx'] = int(meta['global_idx'])
+            except Exception:
+                pass
+
+        tiles.append(tile_info)
+
+    return tiles

upscale/result_collector.py

@@ -1,17 +1,38 @@
 import asyncio, time
 import comfy.model_management
 import server
-from ..utils.constants import HEARTBEAT_INTERVAL
+from ..utils.constants import DYNAMIC_MODE_MAX_POLL_TIMEOUT, HEARTBEAT_INTERVAL
 from ..utils.logging import debug_log, log
 from ..utils.config import get_worker_timeout_seconds
-from ..utils.usdu_managment import (
-    ensure_tile_jobs_initialized, _mark_task_completed,
-    _check_and_requeue_timed_out_workers,
-)
+from .job_store import ensure_tile_jobs_initialized, _mark_task_completed
+from .job_timeout import _check_and_requeue_timed_out_workers
 from .job_models import BaseJobState, ImageJobState, TileJobState
 
 
 class ResultCollectorMixin:
+    """
+    Mixin for master-side result collection in USDU distributed jobs.
+
+    Expected co-mixins/attributes:
+    - JobStateMixin methods for queue/task access.
+    - `self._check_and_requeue_timed_out_workers(...)` coroutine.
+    - `self._async_yield(...)` optional helper from WorkerCommsMixin.
+    """
+
+    def _log_worker_timeout_status(self, job_data, current_time: float, multi_job_id: str) -> list[str]:
+        """Log timeout elapsed seconds for each tracked worker and return worker ids."""
+        if not isinstance(job_data, BaseJobState):
+            return []
+
+        worker_status = dict(job_data.worker_status)
+        for worker_id, last_seen in worker_status.items():
+            elapsed = max(0.0, current_time - float(last_seen))
+            log(
+                "UltimateSDUpscale Master - Heartbeat timeout: "
+                f"job={multi_job_id}, worker={worker_id}, elapsed={elapsed:.1f}s"
+            )
+        return list(worker_status.keys())
+
     async def _async_collect_results(self, multi_job_id, num_workers, mode='static', 
                                    remaining_to_collect=None, batch_size=None):
         """Unified async helper to collect results from workers (tiles or images)."""
@@ -59,10 +80,16 @@ class ResultCollectorMixin:
             # For dynamic mode with remaining_to_collect, check if we've collected enough
             if mode == 'dynamic' and remaining_to_collect and collected_count >= remaining_to_collect:
                 break
-                
+
+            job_data_snapshot = None
+            async with prompt_server.distributed_tile_jobs_lock:
+                current_job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
+                if isinstance(current_job_data, BaseJobState):
+                    job_data_snapshot = current_job_data
+
             try:
                 # Shorter poll for dynamic mode, but never exceed the configured timeout
-                wait_timeout = (min(10.0, timeout) if mode == 'dynamic' else timeout)
+                wait_timeout = (min(DYNAMIC_MODE_MAX_POLL_TIMEOUT, timeout) if mode == 'dynamic' else timeout)
                 result = await asyncio.wait_for(q.get(), timeout=wait_timeout)
                 worker_id = result['worker_id']
                 is_last = result.get('is_last', False)
@@ -115,6 +142,7 @@ class ResultCollectorMixin:
                     
             except asyncio.TimeoutError:
                 current_time = time.time()
+                waiting_workers = self._log_worker_timeout_status(job_data_snapshot, current_time, multi_job_id)
                 if mode == 'dynamic':
                     # Check for worker timeouts periodically
                     if current_time - last_heartbeat_check >= HEARTBEAT_INTERVAL:
@@ -126,16 +154,6 @@ class ResultCollectorMixin:
                     
                     # Check if we've been waiting too long overall
                     if current_time - wait_started_at > timeout:
-                        async with prompt_server.distributed_tile_jobs_lock:
-                            job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
-                            worker_status = dict(job_data.worker_status) if isinstance(job_data, BaseJobState) else {}
-                            waiting_workers = list(worker_status.keys())
-                        for worker_id, last_seen in worker_status.items():
-                            elapsed_for_worker = max(0.0, current_time - float(last_seen))
-                            log(
-                                "UltimateSDUpscale Master - Heartbeat timeout: "
-                                f"worker={worker_id}, elapsed={elapsed_for_worker:.1f}s"
-                            )
                         elapsed = current_time - wait_started_at
                         log(
                             "UltimateSDUpscale Master - Heartbeat timeout while waiting for images; "
@@ -143,16 +161,6 @@ class ResultCollectorMixin:
                         )
                         break
                 else:
-                    async with prompt_server.distributed_tile_jobs_lock:
-                        job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
-                        worker_status = dict(job_data.worker_status) if isinstance(job_data, BaseJobState) else {}
-                        waiting_workers = list(worker_status.keys())
-                    for worker_id, last_seen in worker_status.items():
-                        elapsed_for_worker = max(0.0, current_time - float(last_seen))
-                        log(
-                            "UltimateSDUpscale Master - Heartbeat timeout: "
-                            f"worker={worker_id}, elapsed={elapsed_for_worker:.1f}s"
-                        )
                     elapsed = current_time - wait_started_at
                     log(
                         f"UltimateSDUpscale Master - Heartbeat timeout waiting for {item_type}; "

upscale/tile_ops.py

@@ -5,7 +5,7 @@ import comfy.samplers, comfy.model_management
 from ..utils.logging import debug_log, log
 from ..utils.image import tensor_to_pil, pil_to_tensor
 from ..utils.usdu_utils import crop_cond, get_crop_region, expand_crop
-from ..utils.usdu_managment import clone_conditioning
+from .conditioning import clone_conditioning
 
 
 class TileOpsMixin:

upscale/worker_comms.py

@@ -9,6 +9,10 @@ from ..utils.image import tensor_to_pil
 
 
 class WorkerCommsMixin:
+    async def _send_heartbeat_to_master(self, multi_job_id, master_url, worker_id):
+        """Proxy heartbeat helper used by worker processing mixins."""
+        await _send_heartbeat_to_master(multi_job_id, master_url, worker_id)
+
     async def send_tiles_batch_to_master(self, processed_tiles, multi_job_id, master_url, 
                                        padding, worker_id, is_final_flush=False):
         """Send all processed tiles to master, chunked if large."""

utils/cloudflare.py

@@ -1,388 +0,0 @@
-"""
-Lightweight Cloudflare Tunnel manager.
-
-Responsibilities:
-- Locate or download the cloudflared binary for the current platform.
-- Start/stop a quick tunnel pointing to the local ComfyUI server.
-- Track status and persist minimal state in the shared config.
-"""
-import asyncio
-import os
-import platform
-import re
-import shutil
-import signal
-import stat
-import subprocess
-import threading
-import time
-from urllib import request, error as urlerror
-
-from .config import load_config, save_config
-from .logging import debug_log, log
-from .network import get_server_port, normalize_host
-from .process import is_process_alive, terminate_process
-
-# Regex to capture the generated public URL from cloudflared output
-PUBLIC_URL_PATTERN = re.compile(r"(https?://[\w.-]+\.(?:trycloudflare\.com|cloudflare\.dev))", re.IGNORECASE)
-
-# Default timeout (seconds) to wait for a tunnel URL before giving up
-TUNNEL_START_TIMEOUT = 25
-
-class CloudflareTunnelManager:
-    def __init__(self):
-        self.process = None
-        self.pid = None
-        self.public_url = None
-        self.last_error = None
-        self.log_file = None
-        self.status = "stopped"
-        self.previous_master_host = None
-
-        self._lock = asyncio.Lock()
-        self._url_event = None
-        self._reader_thread = None
-        self._loop = None
-        self._recent_logs = []
-
-        self.binary_path = None
-        self._restore_state()
-
-    # --- Binary discovery & download ---
-    @property
-    def base_dir(self):
-        return os.path.dirname(os.path.dirname(__file__))
-
-    @property
-    def bin_dir(self):
-        return os.path.join(self.base_dir, "bin")
-
-    def _get_asset_name(self):
-        system = platform.system().lower()
-        machine = platform.machine().lower()
-
-        if system == "windows":
-            if "arm" in machine:
-                return "cloudflared-windows-arm64.exe"
-            return "cloudflared-windows-amd64.exe"
-        if system == "darwin":
-            if machine in ("arm64", "aarch64"):
-                return "cloudflared-darwin-arm64"
-            return "cloudflared-darwin-amd64"
-        if system == "linux":
-            if machine in ("arm64", "aarch64"):
-                return "cloudflared-linux-arm64"
-            return "cloudflared-linux-amd64"
-
-        raise RuntimeError(f"Unsupported platform for cloudflared: {system}/{machine}")
-
-    def _download_cloudflared(self):
-        asset = self._get_asset_name()
-        url = f"https://github.com/cloudflare/cloudflared/releases/latest/download/{asset}"
-        os.makedirs(self.bin_dir, exist_ok=True)
-        target_path = os.path.join(self.bin_dir, "cloudflared.exe" if asset.endswith(".exe") else "cloudflared")
-
-        debug_log(f"Downloading cloudflared from {url}")
-        try:
-            with request.urlopen(url, timeout=30) as resp:
-                with open(target_path, "wb") as f:
-                    shutil.copyfileobj(resp, f)
-        except urlerror.URLError as exc:
-            raise RuntimeError(f"Failed to download cloudflared: {exc}") from exc
-
-        # Make executable
-        st = os.stat(target_path)
-        os.chmod(target_path, st.st_mode | stat.S_IEXEC)
-        debug_log(f"Downloaded cloudflared to {target_path}")
-        return target_path
-
-    def ensure_binary(self):
-        # Allow explicit override
-        env_path = os.environ.get("CLOUDFLARED_PATH")
-        if env_path and os.path.exists(env_path):
-            self.binary_path = env_path
-            return env_path
-
-        # Check cached path
-        cached = self.binary_path
-        if cached and os.path.exists(cached):
-            return cached
-
-        # Local bin directory
-        candidate = os.path.join(self.bin_dir, "cloudflared.exe" if platform.system().lower() == "windows" else "cloudflared")
-        if os.path.exists(candidate):
-            self.binary_path = candidate
-            return candidate
-
-        # System PATH
-        path_binary = shutil.which("cloudflared")
-        if path_binary:
-            self.binary_path = path_binary
-            return path_binary
-
-        # Download on demand
-        self.binary_path = self._download_cloudflared()
-        return self.binary_path
-
-    # --- State helpers ---
-    def _persist_state(self, status=None, public_url=None, pid=None, log_file=None, previous_host=None, master_host=None):
-        cfg = load_config()
-        tunnel_cfg = cfg.get("tunnel", {}) if isinstance(cfg.get("tunnel", {}), dict) else {}
-
-        if status is not None:
-            tunnel_cfg["status"] = status
-        if public_url is not None:
-            tunnel_cfg["public_url"] = public_url
-        if pid is not None:
-            tunnel_cfg["pid"] = pid
-        if log_file is not None:
-            tunnel_cfg["log_file"] = log_file
-        if previous_host is not None:
-            tunnel_cfg["previous_master_host"] = previous_host
-        if master_host is not None:
-            cfg.setdefault("master", {})["host"] = master_host
-
-        cfg["tunnel"] = tunnel_cfg
-        save_config(cfg)
-
-    def _restore_state(self):
-        cfg = load_config()
-        tunnel_cfg = cfg.get("tunnel", {}) if isinstance(cfg.get("tunnel", {}), dict) else {}
-
-        self.public_url = tunnel_cfg.get("public_url") or None
-        self.previous_master_host = tunnel_cfg.get("previous_master_host")
-        self.log_file = tunnel_cfg.get("log_file")
-        pid = tunnel_cfg.get("pid")
-
-        if pid and is_process_alive(pid):
-            self.pid = pid
-            self.status = tunnel_cfg.get("status", "running")
-            debug_log(f"Detected existing cloudflared process (pid={pid})")
-        else:
-            # Clear stale info
-            self._persist_state(status="stopped", public_url="", pid=None, log_file=None)
-            self.status = "stopped"
-            self.pid = None
-
-    # --- Process output handling ---
-    def _append_log(self, line):
-        if self.log_file:
-            try:
-                with open(self.log_file, "a", encoding="utf-8", errors="replace") as f:
-                    f.write(line + "\n")
-            except Exception as exc:  # pragma: no cover - best effort logging
-                debug_log(f"Failed to write tunnel log: {exc}")
-
-        self._recent_logs.append(line)
-        # Keep last ~200 lines to avoid unbounded growth
-        if len(self._recent_logs) > 200:
-            self._recent_logs = self._recent_logs[-200:]
-
-    def _reader(self):
-        assert self.process is not None
-        loop = self._loop
-        for raw_line in iter(self.process.stdout.readline, ""):
-            line = raw_line.strip()
-            if not line:
-                continue
-
-            self._append_log(line)
-            match = PUBLIC_URL_PATTERN.search(line)
-            if match and not self.public_url:
-                self.public_url = match.group(1).rstrip("/")
-                self.status = "running"
-                if self._url_event and loop:
-                    loop.call_soon_threadsafe(self._url_event.set)
-
-            # Capture obvious error strings
-            if "error" in line.lower() and not self.last_error:
-                self.last_error = line
-
-        # Process exited
-        if self.status == "starting" and not self.public_url:
-            self.status = "error"
-            if not self.last_error:
-                self.last_error = "Cloudflare tunnel exited before becoming ready"
-            if self._url_event and loop:
-                loop.call_soon_threadsafe(self._url_event.set)
-        elif self.status == "running":
-            # Treat unexpected exit as stopped so UI can retry
-            self.status = "stopped"
-        elif self.status == "stopped":
-            # Normal stop path, nothing to do
-            pass
-
-    # --- Public API ---
-    async def start_tunnel(self):
-        async with self._lock:
-            if self.process and self.process.poll() is None:
-                return {
-                    "status": self.status,
-                    "public_url": self.public_url,
-                    "pid": self.process.pid,
-                    "log_file": self.log_file,
-                }
-
-            # If we had a stale pid, ensure it's not running
-            if self.pid and is_process_alive(self.pid):
-                debug_log(f"Stopping stale cloudflared pid {self.pid} before starting a new one")
-                await self.stop_tunnel()
-
-            binary = await asyncio.to_thread(self.ensure_binary)
-            port = get_server_port()
-            self.status = "starting"
-            self.last_error = None
-            self.public_url = None
-            self._recent_logs = []
-
-            # Remember current master host so we can restore it if the tunnel stops
-            config = load_config()
-            master_host = (config.get("master") or {}).get("host") or ""
-            tunnel_cfg = config.get("tunnel") or {}
-            if tunnel_cfg.get("previous_master_host"):
-                self.previous_master_host = tunnel_cfg.get("previous_master_host")
-            else:
-                self.previous_master_host = master_host
-
-            os.makedirs(os.path.join(self.base_dir, "logs"), exist_ok=True)
-            timestamp = time.strftime("%Y%m%d-%H%M%S")
-            self.log_file = os.path.join(self.base_dir, "logs", f"cloudflare-{timestamp}.log")
-
-            cmd = [
-                binary,
-                "tunnel",
-                "--no-autoupdate",
-                "--url",
-                f"http://127.0.0.1:{port}",
-            ]
-
-            debug_log(f"Starting cloudflared: {' '.join(cmd)}")
-            try:
-                self.process = subprocess.Popen(
-                    cmd,
-                    stdout=subprocess.PIPE,
-                    stderr=subprocess.STDOUT,
-                    text=True,
-                    bufsize=1,
-                )
-            except FileNotFoundError:
-                self.status = "error"
-                raise RuntimeError("cloudflared binary not found")
-            except Exception as exc:
-                self.status = "error"
-                raise RuntimeError(f"Failed to start cloudflared: {exc}") from exc
-
-            self.pid = self.process.pid
-            self._persist_state(
-                status="starting",
-                pid=self.pid,
-                log_file=self.log_file,
-                previous_host=self.previous_master_host
-            )
-
-            # Kick off background reader
-            self._loop = asyncio.get_running_loop()
-            self._url_event = asyncio.Event()
-            self._reader_thread = threading.Thread(target=self._reader, daemon=True)
-            self._reader_thread.start()
-
-            try:
-                await asyncio.wait_for(self._url_event.wait(), timeout=TUNNEL_START_TIMEOUT)
-            except asyncio.TimeoutError:
-                self.last_error = "Timed out waiting for Cloudflare to assign a URL"
-                await self.stop_tunnel()
-                raise RuntimeError(self.last_error)
-
-            if self.status != "running" or not self.public_url:
-                error_msg = self.last_error or "Cloudflare tunnel failed to start"
-                await self.stop_tunnel()
-                raise RuntimeError(error_msg)
-
-            debug_log(f"Cloudflare tunnel ready at {self.public_url}")
-            master_host = normalize_host(self.public_url)
-            self._persist_state(
-                status="running",
-                public_url=self.public_url,
-                pid=self.pid,
-                log_file=self.log_file,
-                previous_host=self.previous_master_host or "",
-                master_host=master_host
-            )
-            return {
-                "status": self.status,
-                "public_url": self.public_url,
-                "pid": self.pid,
-                "log_file": self.log_file,
-            }
-
-    async def stop_tunnel(self):
-        async with self._lock:
-            pid = self.process.pid if self.process else self.pid
-            if not pid:
-                self._persist_state(status="stopped", public_url=None, pid=None, log_file=None)
-                self.status = "stopped"
-                return {"status": "stopped"}
-
-            debug_log(f"Stopping cloudflared (pid={pid})")
-            if self.process:
-                terminate_process(self.process, timeout=5)
-            else:
-                try:
-                    os.kill(pid, signal.SIGTERM)
-                    time.sleep(0.5)
-                except Exception as exc:  # pragma: no cover - best effort
-                    debug_log(f"Error stopping cloudflared pid {pid}: {exc}")
-
-            config = load_config()
-            tunnel_cfg = config.get("tunnel", {}) if isinstance(config.get("tunnel", {}), dict) else {}
-            self.previous_master_host = tunnel_cfg.get("previous_master_host", self.previous_master_host)
-            active_url = tunnel_cfg.get("public_url")
-            current_master_host = (config.get("master") or {}).get("host")
-            restore_host = None
-            if active_url:
-                active_host = normalize_host(active_url)
-                current_host = normalize_host(current_master_host)
-                if current_host == active_host:
-                    restore_host = self.previous_master_host or ""
-
-            self.status = "stopped"
-            self.public_url = None
-            self.pid = None
-            self.process = None
-            self.last_error = None
-            self._url_event = None
-            if self._reader_thread and self._reader_thread.is_alive():
-                self._reader_thread.join(timeout=1)
-            self._reader_thread = None
-            self._persist_state(
-                status="stopped",
-                public_url="",
-                pid=None,
-                log_file=self.log_file,
-                previous_host=self.previous_master_host,
-                master_host=restore_host
-            )
-            return {"status": "stopped"}
-
-    def get_status(self):
-        alive = False
-        pid = self.process.pid if self.process else self.pid
-        if pid:
-            alive = is_process_alive(pid)
-            if not alive and self.status == "running":
-                self.status = "stopped"
-
-        return {
-            "status": self.status,
-            "public_url": self.public_url,
-            "pid": pid,
-            "log_file": self.log_file,
-            "last_error": self.last_error,
-            "binary_path": self.binary_path or shutil.which("cloudflared"),
-            "recent_logs": self._recent_logs[-20:],
-            "previous_master_host": self.previous_master_host,
-        }
-
-
-# Singleton tunnel manager
-cloudflare_tunnel_manager = CloudflareTunnelManager()

utils/cloudflare/__init__.py

@@ -0,0 +1,5 @@
+from .tunnel import CloudflareTunnelManager
+
+cloudflare_tunnel_manager = CloudflareTunnelManager()
+
+__all__ = ["CloudflareTunnelManager", "cloudflare_tunnel_manager"]

utils/cloudflare/binary.py

@@ -0,0 +1,83 @@
+"""Cloudflared binary discovery and download helpers."""
+
+import os
+import platform
+import shutil
+import stat
+from urllib import error as urlerror
+from urllib import request
+
+from ..logging import debug_log
+
+
+def _get_project_root():
+    return os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
+
+
+def _get_cloudflared_dir():
+    return os.path.join(_get_project_root(), "bin")
+
+
+def _get_platform_binary_name():
+    system = platform.system().lower()
+    machine = platform.machine().lower()
+
+    if system == "windows":
+        if "arm" in machine:
+            return "cloudflared-windows-arm64.exe"
+        return "cloudflared-windows-amd64.exe"
+    if system == "darwin":
+        if machine in ("arm64", "aarch64"):
+            return "cloudflared-darwin-arm64"
+        return "cloudflared-darwin-amd64"
+    if system == "linux":
+        if machine in ("arm64", "aarch64"):
+            return "cloudflared-linux-arm64"
+        return "cloudflared-linux-amd64"
+
+    raise RuntimeError(f"Unsupported platform for cloudflared: {system}/{machine}")
+
+
+def _get_binary_path(bin_dir=None):
+    bin_dir = bin_dir or _get_cloudflared_dir()
+    binary_name = "cloudflared.exe" if platform.system().lower() == "windows" else "cloudflared"
+    return os.path.join(bin_dir, binary_name)
+
+
+def _download_cloudflared():
+    asset = _get_platform_binary_name()
+    url = f"https://github.com/cloudflare/cloudflared/releases/latest/download/{asset}"
+
+    bin_dir = _get_cloudflared_dir()
+    os.makedirs(bin_dir, exist_ok=True)
+    target_path = _get_binary_path(bin_dir)
+
+    debug_log(f"Downloading cloudflared from {url}")
+    try:
+        with request.urlopen(url, timeout=30) as resp:
+            with open(target_path, "wb") as f:
+                shutil.copyfileobj(resp, f)
+    except urlerror.URLError as exc:
+        raise RuntimeError(f"Failed to download cloudflared: {exc}") from exc
+
+    st = os.stat(target_path)
+    os.chmod(target_path, st.st_mode | stat.S_IEXEC)
+    debug_log(f"Downloaded cloudflared to {target_path}")
+    return target_path
+
+
+def ensure_binary() -> str:
+    """Return a usable cloudflared binary path, downloading if necessary."""
+    env_path = os.environ.get("CLOUDFLARED_PATH")
+    if env_path and os.path.exists(env_path):
+        return env_path
+
+    local_candidate = _get_binary_path()
+    if os.path.exists(local_candidate):
+        return local_candidate
+
+    path_binary = shutil.which("cloudflared")
+    if path_binary:
+        return path_binary
+
+    return _download_cloudflared()

utils/cloudflare/process_reader.py

@@ -0,0 +1,97 @@
+"""Background cloudflared process output reader."""
+
+import asyncio
+import re
+import threading
+
+from ..constants import CLOUDFLARE_LOG_BUFFER_SIZE
+from ..logging import debug_log
+
+PUBLIC_URL_PATTERN = re.compile(
+    r"(https?://[\w.-]+\.(?:trycloudflare\.com|cloudflare\.dev))",
+    re.IGNORECASE,
+)
+class ProcessReader:
+    def __init__(self, log_file=None):
+        self._process = None
+        self._thread = None
+        self._loop = None
+        self._url_event = None
+        self._public_url = None
+        self._last_error = None
+        self._recent_logs = []
+        self._log_file = log_file
+
+    def set_log_file(self, log_file):
+        self._log_file = log_file
+
+    def _append_log(self, line):
+        if self._log_file:
+            try:
+                with open(self._log_file, "a", encoding="utf-8", errors="replace") as f:
+                    f.write(line + "\n")
+            except Exception as exc:  # pragma: no cover
+                debug_log(f"Failed to write tunnel log: {exc}")
+
+        self._recent_logs.append(line)
+        if len(self._recent_logs) > CLOUDFLARE_LOG_BUFFER_SIZE:
+            self._recent_logs = self._recent_logs[-CLOUDFLARE_LOG_BUFFER_SIZE:]
+
+    def _reader(self):
+        process = self._process
+        if process is None:
+            return
+
+        loop = self._loop
+        for raw_line in iter(process.stdout.readline, ""):
+            line = raw_line.strip()
+            if not line:
+                continue
+
+            self._append_log(line)
+            match = PUBLIC_URL_PATTERN.search(line)
+            if match and not self._public_url:
+                self._public_url = match.group(1).rstrip("/")
+                if self._url_event and loop:
+                    loop.call_soon_threadsafe(self._url_event.set)
+
+            if "error" in line.lower() and not self._last_error:
+                self._last_error = line
+
+        if self._url_event and loop:
+            if not self._last_error and not self._public_url:
+                self._last_error = "Cloudflare tunnel exited before becoming ready"
+            loop.call_soon_threadsafe(self._url_event.set)
+
+    def start(self, process, loop):
+        self._process = process
+        self._loop = loop
+        self._url_event = asyncio.Event()
+        self._public_url = None
+        self._last_error = None
+        self._recent_logs = []
+        self._thread = threading.Thread(target=self._reader, daemon=True)
+        self._thread.start()
+
+    async def wait_for_url(self, timeout):
+        if not self._url_event:
+            return None
+        await asyncio.wait_for(self._url_event.wait(), timeout=timeout)
+        return self._public_url
+
+    def stop(self):
+        if self._thread and self._thread.is_alive():
+            self._thread.join(timeout=1)
+        self._thread = None
+        self._process = None
+        self._loop = None
+        self._url_event = None
+
+    def get_url(self):
+        return self._public_url
+
+    def get_last_error(self):
+        return self._last_error
+
+    def get_recent_logs(self):
+        return list(self._recent_logs)

utils/cloudflare/state.py

@@ -0,0 +1,81 @@
+"""Cloudflare tunnel state persistence helpers."""
+
+from ..config import load_config, save_config
+from ..network import normalize_host
+
+
+def _get_tunnel_config(cfg):
+    tunnel_cfg = cfg.get("tunnel", {})
+    if isinstance(tunnel_cfg, dict):
+        return tunnel_cfg
+    return {}
+
+
+def load_tunnel_state():
+    cfg = load_config()
+    tunnel_cfg = _get_tunnel_config(cfg)
+    master_cfg = cfg.get("master", {}) if isinstance(cfg.get("master", {}), dict) else {}
+    return {
+        "status": tunnel_cfg.get("status", "stopped"),
+        "public_url": tunnel_cfg.get("public_url") or None,
+        "pid": tunnel_cfg.get("pid"),
+        "log_file": tunnel_cfg.get("log_file"),
+        "previous_master_host": tunnel_cfg.get("previous_master_host"),
+        "master_host": master_cfg.get("host"),
+    }
+
+
+def persist_tunnel_state(
+    status=None,
+    public_url=None,
+    pid=None,
+    log_file=None,
+    previous_host=None,
+    master_host=None,
+):
+    cfg = load_config()
+    tunnel_cfg = _get_tunnel_config(cfg)
+
+    if status is not None:
+        tunnel_cfg["status"] = status
+    if public_url is not None:
+        tunnel_cfg["public_url"] = public_url
+    if pid is not None:
+        tunnel_cfg["pid"] = pid
+    if log_file is not None:
+        tunnel_cfg["log_file"] = log_file
+    if previous_host is not None:
+        tunnel_cfg["previous_master_host"] = previous_host
+    if master_host is not None:
+        cfg.setdefault("master", {})["host"] = master_host
+
+    cfg["tunnel"] = tunnel_cfg
+    save_config(cfg)
+
+
+def clear_tunnel_state(log_file=None, previous_host=None, master_host=None):
+    persist_tunnel_state(
+        status="stopped",
+        public_url="",
+        pid=None,
+        log_file=log_file,
+        previous_host=previous_host,
+        master_host=master_host,
+    )
+
+
+def resolve_restore_master_host(previous_master_host):
+    """Determine whether master host should be restored after tunnel stop."""
+    cfg = load_config()
+    tunnel_cfg = _get_tunnel_config(cfg)
+    active_url = tunnel_cfg.get("public_url")
+    current_master_host = (cfg.get("master") or {}).get("host")
+
+    if not active_url:
+        return None
+
+    active_host = normalize_host(active_url)
+    current_host = normalize_host(current_master_host)
+    if current_host == active_host:
+        return previous_master_host or ""
+    return None

utils/cloudflare/tunnel.py

@@ -0,0 +1,207 @@
+"""Cloudflare tunnel lifecycle manager."""
+
+import asyncio
+import os
+import shutil
+import signal
+import subprocess
+import time
+
+from ..constants import TUNNEL_START_TIMEOUT
+from ..logging import debug_log
+from ..network import get_server_port, normalize_host
+from ..process import is_process_alive, terminate_process
+from .binary import ensure_binary
+from .process_reader import ProcessReader
+from .state import clear_tunnel_state, load_tunnel_state, persist_tunnel_state, resolve_restore_master_host
+
+
+class CloudflareTunnelManager:
+    def __init__(self):
+        self.process = None
+        self.pid = None
+        self.public_url = None
+        self.last_error = None
+        self.log_file = None
+        self.status = "stopped"
+        self.previous_master_host = None
+
+        self._lock = asyncio.Lock()
+        self._reader = ProcessReader()
+        self.binary_path = None
+
+        self._restore_state()
+
+    @property
+    def base_dir(self):
+        return os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
+
+    def _restore_state(self):
+        state = load_tunnel_state()
+
+        self.public_url = state.get("public_url") or None
+        self.previous_master_host = state.get("previous_master_host")
+        self.log_file = state.get("log_file")
+        pid = state.get("pid")
+
+        if pid and is_process_alive(pid):
+            self.pid = pid
+            self.status = state.get("status") or "running"
+            debug_log(f"Detected existing cloudflared process (pid={pid})")
+        else:
+            clear_tunnel_state(log_file=self.log_file, previous_host=self.previous_master_host)
+            self.status = "stopped"
+            self.pid = None
+
+    async def start_tunnel(self):
+        async with self._lock:
+            if self.process and self.process.poll() is None:
+                return {
+                    "status": self.status,
+                    "public_url": self.public_url,
+                    "pid": self.process.pid,
+                    "log_file": self.log_file,
+                }
+
+            if self.pid and is_process_alive(self.pid):
+                debug_log(f"Stopping stale cloudflared pid {self.pid} before starting a new one")
+                await self.stop_tunnel()
+
+            binary = await asyncio.to_thread(ensure_binary)
+            self.binary_path = binary
+            port = get_server_port()
+            self.status = "starting"
+            self.last_error = None
+            self.public_url = None
+
+            state = load_tunnel_state()
+            master_host = state.get("master_host") or ""
+            if state.get("previous_master_host"):
+                self.previous_master_host = state.get("previous_master_host")
+            else:
+                self.previous_master_host = master_host
+
+            os.makedirs(os.path.join(self.base_dir, "logs"), exist_ok=True)
+            timestamp = time.strftime("%Y%m%d-%H%M%S")
+            self.log_file = os.path.join(self.base_dir, "logs", f"cloudflare-{timestamp}.log")
+
+            cmd = [
+                binary,
+                "tunnel",
+                "--no-autoupdate",
+                "--url",
+                f"http://127.0.0.1:{port}",
+            ]
+
+            debug_log(f"Starting cloudflared: {' '.join(cmd)}")
+            try:
+                self.process = subprocess.Popen(
+                    cmd,
+                    stdout=subprocess.PIPE,
+                    stderr=subprocess.STDOUT,
+                    text=True,
+                    bufsize=1,
+                )
+            except FileNotFoundError:
+                self.status = "error"
+                raise RuntimeError("cloudflared binary not found")
+            except Exception as exc:
+                self.status = "error"
+                raise RuntimeError(f"Failed to start cloudflared: {exc}") from exc
+
+            self.pid = self.process.pid
+            persist_tunnel_state(
+                status="starting",
+                pid=self.pid,
+                log_file=self.log_file,
+                previous_host=self.previous_master_host,
+            )
+
+            loop = asyncio.get_running_loop()
+            self._reader.set_log_file(self.log_file)
+            self._reader.start(self.process, loop)
+
+            try:
+                await self._reader.wait_for_url(timeout=TUNNEL_START_TIMEOUT)
+            except asyncio.TimeoutError:
+                self.last_error = "Timed out waiting for Cloudflare to assign a URL"
+                await self.stop_tunnel()
+                raise RuntimeError(self.last_error)
+
+            public_url = self._reader.get_url()
+            if not public_url:
+                self.last_error = self._reader.get_last_error() or "Cloudflare tunnel failed to start"
+                await self.stop_tunnel()
+                raise RuntimeError(self.last_error)
+
+            self.public_url = public_url
+            self.status = "running"
+            debug_log(f"Cloudflare tunnel ready at {self.public_url}")
+
+            persist_tunnel_state(
+                status="running",
+                public_url=self.public_url,
+                pid=self.pid,
+                log_file=self.log_file,
+                previous_host=self.previous_master_host or "",
+                master_host=normalize_host(self.public_url),
+            )
+            return {
+                "status": self.status,
+                "public_url": self.public_url,
+                "pid": self.pid,
+                "log_file": self.log_file,
+            }
+
+    async def stop_tunnel(self):
+        async with self._lock:
+            pid = self.process.pid if self.process else self.pid
+            if not pid:
+                clear_tunnel_state(log_file=self.log_file, previous_host=self.previous_master_host)
+                self.status = "stopped"
+                return {"status": "stopped"}
+
+            debug_log(f"Stopping cloudflared (pid={pid})")
+            if self.process:
+                terminate_process(self.process, timeout=5)
+            else:
+                try:
+                    os.kill(pid, signal.SIGTERM)
+                    time.sleep(0.5)
+                except Exception as exc:  # pragma: no cover
+                    debug_log(f"Error stopping cloudflared pid {pid}: {exc}")
+
+            restore_host = resolve_restore_master_host(self.previous_master_host)
+
+            self.status = "stopped"
+            self.public_url = None
+            self.pid = None
+            self.process = None
+            self.last_error = None
+            self._reader.stop()
+
+            clear_tunnel_state(
+                log_file=self.log_file,
+                previous_host=self.previous_master_host,
+                master_host=restore_host,
+            )
+            return {"status": "stopped"}
+
+    def get_status(self):
+        alive = False
+        pid = self.process.pid if self.process else self.pid
+        if pid:
+            alive = is_process_alive(pid)
+            if not alive and self.status == "running":
+                self.status = "stopped"
+
+        return {
+            "status": self.status,
+            "public_url": self.public_url,
+            "pid": pid,
+            "log_file": self.log_file,
+            "last_error": self.last_error or self._reader.get_last_error(),
+            "binary_path": self.binary_path or shutil.which("cloudflared"),
+            "recent_logs": self._reader.get_recent_logs()[-20:],
+            "previous_master_host": self.previous_master_host,
+        }

utils/config.py

@@ -1,8 +1,10 @@
 """
 Configuration management for ComfyUI-Distributed.
 """
+import asyncio
 import os
 import json
+from contextlib import asynccontextmanager
 from .logging import log
 
 # Import defaults for timeout fallbacks
@@ -11,6 +13,7 @@ from .constants import HEARTBEAT_TIMEOUT
 CONFIG_FILE = os.path.join(os.path.dirname(os.path.dirname(__file__)), "gpu_config.json")
 _config_cache = None
 _config_mtime = 0.0
+_config_lock = asyncio.Lock()
 
 
 def _config_path():
@@ -95,15 +98,36 @@ def load_config():
 
 def save_config(config):
     """Saves the configuration to file."""
+    tmp_path = f"{_config_path()}.tmp"
     try:
-        with open(_config_path(), 'w', encoding='utf-8') as f:
+        with open(tmp_path, 'w', encoding='utf-8') as f:
             json.dump(config, f, indent=2)
+            f.flush()
+            os.fsync(f.fileno())
+        os.replace(tmp_path, _config_path())
         invalidate_config_cache()
         return True
     except Exception as e:
+        try:
+            os.unlink(tmp_path)
+        except OSError:
+            pass
         log(f"Error saving config: {e}")
         return False
 
+
+@asynccontextmanager
+async def config_transaction():
+    """Acquire config lock, yield loaded config, and save if changed."""
+    async with _config_lock:
+        config = load_config()
+        original_snapshot = json.dumps(config, sort_keys=True)
+        yield config
+        updated_snapshot = json.dumps(config, sort_keys=True)
+        if updated_snapshot != original_snapshot:
+            if not save_config(config):
+                raise RuntimeError("Failed to save config")
+
 def ensure_config_exists():
     """Creates default config file if it doesn't exist. Used by __init__.py"""
     if not os.path.exists(_config_path()):

utils/constants.py

@@ -13,6 +13,10 @@ PROCESS_TERMINATION_TIMEOUT = 5.0
 WORKER_CHECK_INTERVAL = 2.0
 STATUS_CHECK_INTERVAL = 5.0
 
+# Cloudflare tunnel
+TUNNEL_START_TIMEOUT = float(os.environ.get("TUNNEL_START_TIMEOUT", "25"))
+CLOUDFLARE_LOG_BUFFER_SIZE = 200
+
 # Network
 CHUNK_SIZE = 8192
 LOG_TAIL_BYTES = 65536  # 64KB
@@ -42,6 +46,13 @@ MAX_BATCH = int(os.environ.get('COMFYUI_MAX_BATCH', '20'))  # Maximum items per
 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)
 
+# USDU result collection
+DYNAMIC_MODE_MAX_POLL_TIMEOUT = 10.0
+
+# Static mode job poll loop
+JOB_POLL_INTERVAL = 1.0
+JOB_POLL_MAX_ATTEMPTS = 20
+
 # Orchestration pipeline
 ORCHESTRATION_WORKER_PROBE_CONCURRENCY = int(
     os.environ.get('COMFYUI_ORCHESTRATION_WORKER_PROBE_CONCURRENCY', '8')

utils/network.py

@@ -34,8 +34,14 @@ async def cleanup_client_session():
 
 async def handle_api_error(request, error, status=500):
     """Standardized error response handler."""
-    debug_log(f"API Error: {error}")
-    return web.json_response({"status": "error", "message": str(error)}, status=status)
+    if isinstance(error, list):
+        messages = [str(item) for item in error]
+        debug_log(f"API Error [{status}]: {messages}")
+        return web.json_response({"errors": messages}, status=status)
+
+    message = str(error)
+    debug_log(f"API Error [{status}]: {message}")
+    return web.json_response({"error": message}, status=status)
 
 def get_server_port():
     """Get the ComfyUI server port."""

utils/trace_logger.py

@@ -0,0 +1,13 @@
+from .logging import debug_log, log
+
+
+def trace_prefix(trace_execution_id: str) -> str:
+    return f"[Distributed][exec:{trace_execution_id}]"
+
+
+def trace_debug(trace_execution_id: str, message: str) -> None:
+    debug_log(f"{trace_prefix(trace_execution_id)} {message}")
+
+
+def trace_info(trace_execution_id: str, message: str) -> None:
+    log(f"{trace_prefix(trace_execution_id)} {message}")

utils/usdu_managment.py

@@ -1,343 +1,32 @@
-import asyncio
-import time
-import json
-import copy
-import os
-import io
-from aiohttp import web
-import server
-from PIL import Image
+"""Backward-compatibility shim for USDU helpers.
 
-# Import from other utilities
-from .logging import debug_log, log
-from .network import build_worker_url, handle_api_error, get_client_session, probe_worker
-from ..upscale.job_models import BaseJobState, ImageJobState, TileJobState
-# We avoid converting to tensors on the master for tiles; blending uses PIL
-from typing import List, Optional
+Route handlers and job logic now live in:
+- upscale.job_store
+- upscale.job_timeout
+- upscale.payload_parsers
+- api.usdu_routes
+"""
 
-# Configure maximum payload size (50MB default, configurable via environment variable)
-MAX_PAYLOAD_SIZE = int(os.environ.get('COMFYUI_MAX_PAYLOAD_SIZE', str(50 * 1024 * 1024)))
+from ..upscale.conditioning import clone_conditioning, clone_control_chain
+from ..upscale.job_store import (
+    MAX_PAYLOAD_SIZE,
+    _cleanup_job,
+    _drain_results_queue,
+    _get_completed_count,
+    _init_job_queue,
+    _mark_task_completed,
+    ensure_tile_jobs_initialized,
+    init_dynamic_job,
+    init_static_job_batched,
+)
+from ..upscale.job_timeout import _check_and_requeue_timed_out_workers
+from ..upscale.payload_parsers import _parse_tiles_from_form
+from .logging import debug_log
+from .network import get_client_session
 
-# Import HEARTBEAT_TIMEOUT from constants
-from .constants import HEARTBEAT_TIMEOUT
-from .config import load_config
-
-
-def _parse_tiles_from_form(data):
-    """Parse tiles submitted via multipart/form-data into a list of tile dicts.
-
-    Expects the following fields in the aiohttp form data:
-    - 'tiles_metadata': JSON list with per-tile metadata items containing at least
-      'tile_idx', 'x', 'y', 'extracted_width', 'extracted_height'. Optional
-      'batch_idx' and 'global_idx' are included when available.
-    - 'tile_{i}': image bytes for each tile described in tiles_metadata (PNG).
-    - 'padding': integer padding used during extraction (optional; defaults 0).
-
-    Returns: list of dicts with keys: 'image', 'tile_idx', 'x', 'y',
-    'extracted_width', 'extracted_height', and optional 'batch_idx', 'global_idx',
-    plus 'padding'.
-    """
-    try:
-        # Parse padding if present
-        padding = int(data.get('padding', 0)) if data.get('padding') is not None else 0
-    except Exception:
-        padding = 0
-
-    # Parse tiles metadata (JSON list)
-    meta_raw = data.get('tiles_metadata')
-    if meta_raw is None:
-        raise ValueError("Missing tiles_metadata")
-
-    try:
-        metadata = json.loads(meta_raw)
-    except Exception as e:
-        raise ValueError(f"Invalid tiles_metadata JSON: {e}")
-
-    if not isinstance(metadata, list):
-        raise ValueError("tiles_metadata must be a list")
-
-    tiles = []
-    # Iterate over metadata items and corresponding uploaded files tile_0, tile_1, ...
-    for i, meta in enumerate(metadata):
-        file_field = data.get(f'tile_{i}')
-        if file_field is None or not hasattr(file_field, 'file'):
-            raise ValueError(f"Missing tile data for index {i}")
-
-        # Read image bytes and decode to PIL
-        raw = file_field.file.read()
-        try:
-            img = Image.open(io.BytesIO(raw)).convert("RGB")
-        except Exception as e:
-            raise ValueError(f"Invalid image data for tile {i}: {e}")
-
-        # Build tile dictionary (store PIL only; master blends via PIL)
-        try:
-            tile_info = {
-                'image': img,
-                'tile_idx': int(meta.get('tile_idx', i)),
-                'x': int(meta.get('x', 0)),
-                'y': int(meta.get('y', 0)),
-                'extracted_width': int(meta.get('extracted_width', img.width)),
-                'extracted_height': int(meta.get('extracted_height', img.height)),
-                'padding': int(padding),
-            }
-        except Exception as e:
-            raise ValueError(f"Invalid metadata values for tile {i}: {e}")
-
-        # Optional fields
-        if 'batch_idx' in meta:
-            try:
-                tile_info['batch_idx'] = int(meta['batch_idx'])
-            except Exception:
-                pass
-        if 'global_idx' in meta:
-            try:
-                tile_info['global_idx'] = int(meta['global_idx'])
-            except Exception:
-                pass
-
-        tiles.append(tile_info)
-
-    return tiles
-
-async def init_dynamic_job(multi_job_id: str, batch_size: int, enabled_workers: List[str], all_indices: Optional[List[int]] = None):
-    """Initialize queue for dynamic mode (per-image), with collector fields.
-
-    - Creates pending image queue entries
-    - Initializes dynamic completion/status tracking
-    """
-    await _init_job_queue(
-        multi_job_id,
-        'dynamic',
-        batch_size=batch_size,
-        all_indices=all_indices or list(range(batch_size)),
-        enabled_workers=enabled_workers,
-    )
-    debug_log(f"Job {multi_job_id} initialized with {batch_size} images")
-
-
-async def init_static_job_batched(multi_job_id: str, batch_size: int, num_tiles_per_image: int, enabled_workers: List[str]):
-    """Initialize queue for static mode (batched-per-tile).
-
-    - Populates pending tile queue with tile ids [0..num_tiles_per_image-1]
-    """
-    await _init_job_queue(
-        multi_job_id,
-        'static',
-        batch_size=batch_size,
-        num_tiles_per_image=num_tiles_per_image,
-        enabled_workers=enabled_workers,
-        batched_static=True,
-    )
-    # Initialization handled by master; avoid duplicate init logs here
-
-async def _init_job_queue(
-    multi_job_id,
-    mode,
-    batch_size=None,
-    num_tiles_per_image=None,
-    all_indices=None,
-    enabled_workers=None,
-    batched_static: bool = False,
-):
-    """Unified initialization for job queues in static and dynamic modes."""
-    prompt_server = ensure_tile_jobs_initialized()
-    async with prompt_server.distributed_tile_jobs_lock:
-        if multi_job_id in prompt_server.distributed_pending_tile_jobs:
-            debug_log(f"Queue already exists for {multi_job_id}")
-            return
-
-        if mode == 'dynamic':
-            job_data = ImageJobState(multi_job_id=multi_job_id)
-        elif mode == 'static':
-            job_data = TileJobState(multi_job_id=multi_job_id)
-        else:
-            raise ValueError(f"Unknown mode: {mode}")
-
-        job_data.worker_status = {w: time.time() for w in enabled_workers or []}
-        job_data.assigned_to_workers = {w: [] for w in enabled_workers or []}
-
-        if mode == 'dynamic':
-            job_data.batch_size = int(batch_size or 0)
-            pending_queue = job_data.pending_images
-            for i in (all_indices or range(int(batch_size or 0))):
-                await pending_queue.put(i)
-            debug_log(f"Initialized image queue with {batch_size} pending items")
-        elif mode == 'static':
-            job_data.num_tiles_per_image = int(num_tiles_per_image or 0)
-            job_data.batch_size = int(batch_size or 0)
-            job_data.batched_static = bool(batched_static)
-            # For batched static distribution, populate only tile ids [0..num_tiles_per_image-1]
-            pending_queue = job_data.pending_tasks
-            if batched_static and num_tiles_per_image is not None:
-                for i in range(num_tiles_per_image):
-                    await pending_queue.put(i)
-            else:
-                total_tiles = int(batch_size or 0) * int(num_tiles_per_image or 0)
-                for i in range(total_tiles):
-                    await pending_queue.put(i)
-            
-        prompt_server.distributed_pending_tile_jobs[multi_job_id] = job_data
-
-async def _drain_results_queue(multi_job_id):
-    """Drain pending results from queue and update completed_tasks. Returns count drained.
-
-    Uses non-blocking get_nowait to avoid await timeouts and reduce latency.
-    """
-    prompt_server = ensure_tile_jobs_initialized()
-    async with prompt_server.distributed_tile_jobs_lock:
-        job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
-        if not isinstance(job_data, BaseJobState):
-            return 0
-        q = job_data.queue
-        completed_tasks = job_data.completed_tasks
-
-        collected = 0
-        while True:
-            try:
-                result = q.get_nowait()
-            except asyncio.QueueEmpty:
-                break
-
-            worker_id = result['worker_id']
-            is_last = result.get('is_last', False)
-
-            if 'image_idx' in result and 'image' in result:
-                task_id = result['image_idx']
-                if task_id not in completed_tasks:
-                    completed_tasks[task_id] = result['image']
-                    collected += 1
-            elif 'tiles' in result:
-                for tile_data in result['tiles']:
-                    task_id = tile_data.get('global_idx', tile_data['tile_idx'])
-                    if task_id not in completed_tasks:
-                        completed_tasks[task_id] = tile_data
-                        collected += 1
-            if is_last:
-                # Track worker completion
-                if worker_id in job_data.worker_status:
-                    del job_data.worker_status[worker_id]
-
-        return collected
-
-async def _check_and_requeue_timed_out_workers(multi_job_id, total_tasks):
-    """Check timed out workers and requeue their tasks. Returns requeued count."""
-    prompt_server = ensure_tile_jobs_initialized()
-    async with prompt_server.distributed_tile_jobs_lock:
-        job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
-        if not job_data:
-            return 0
-
-        current_time = time.time()
-        if not isinstance(job_data, BaseJobState):
-            return 0
-
-        requeued_count = 0
-        completed_tasks = job_data.completed_tasks
-
-        # Allow override via config setting 'worker_timeout_seconds'
-        cfg = load_config()
-        hb_timeout = int(cfg.get('settings', {}).get('worker_timeout_seconds', HEARTBEAT_TIMEOUT))
-
-        for worker, last_heartbeat in list(job_data.worker_status.items()):
-            age = current_time - last_heartbeat
-            debug_log(f"Timeout check: worker={worker} age={age:.1f}s threshold={hb_timeout}s")
-            if age > hb_timeout:
-                # Busy-only grace policy: require positive signal from worker (/prompt)
-                # We also log assignment state for diagnostics but do not grace on it alone.
-                assigned = job_data.assigned_to_workers.get(worker, [])
-                incomplete_assigned = 0
-                try:
-                    if assigned:
-                        batched_static = bool(job_data.batched_static)
-                        if batched_static:
-                            num_tiles_per_image = int(job_data.num_tiles_per_image or 1)
-                            batch_size = int(job_data.batch_size or 1)
-                            for task_id in assigned:
-                                for b in range(batch_size):
-                                    gidx = b * num_tiles_per_image + task_id
-                                    if gidx not in completed_tasks:
-                                        incomplete_assigned += 1
-                                        break
-                        else:
-                            for task_id in assigned:
-                                if task_id not in completed_tasks:
-                                    incomplete_assigned += 1
-                    debug_log(f"Assigned diagnostics: total_assigned={len(assigned)} incomplete_assigned={incomplete_assigned}")
-                except Exception as e:
-                    debug_log(f"Assigned diagnostics failed for worker {worker}: {e}")
-
-                busy = False
-                probe_queue = None
-                try:
-                    cfg_workers = load_config().get('workers', [])
-                    wrec = next((w for w in cfg_workers if str(w.get('id')) == str(worker)), None)
-                    if wrec:
-                        worker_url = build_worker_url(wrec)
-                        debug_log(f"Probing worker {worker} at {worker_url}/prompt")
-                        payload = await probe_worker(worker_url, timeout=2.0)
-                        if payload is not None:
-                            probe_queue = int(payload.get('exec_info', {}).get('queue_remaining', 0))
-                            busy = probe_queue is not None and probe_queue > 0
-                except Exception as e:
-                    debug_log(f"Probe failed for worker {worker}: {e}")
-                finally:
-                    debug_log(
-                        f"Probe diagnostics: online={probe_queue is not None} queue_remaining={probe_queue}"
-                    )
-
-                if busy:
-                    job_data.worker_status[worker] = current_time
-                    debug_log(f"Heartbeat grace: worker {worker} busy via probe; skipping requeue")
-                    continue
-
-                log(f"Worker {worker} heartbeat timed out after {age:.1f}s")
-                for task_id in job_data.assigned_to_workers.get(worker, []):
-                    # If batched_static, task_id is a tile_idx; consider it complete only if
-                    # all corresponding global_idx entries are present in completed_tasks.
-                    batched_static = bool(job_data.batched_static)
-                    if batched_static:
-                        num_tiles_per_image = int(job_data.num_tiles_per_image or 1)
-                        batch_size = int(job_data.batch_size or 1)
-                        # Check all global indices for this tile across the batch
-                        all_done = True
-                        for b in range(batch_size):
-                            gidx = b * num_tiles_per_image + task_id
-                            if gidx not in completed_tasks:
-                                all_done = False
-                                break
-                        if not all_done:
-                            await job_data.pending_tasks.put(task_id)
-                            requeued_count += 1
-                    else:
-                        if task_id not in completed_tasks:
-                            await job_data.pending_tasks.put(task_id)
-                            requeued_count += 1
-                job_data.worker_status.pop(worker, None)
-                if worker in job_data.assigned_to_workers:
-                    job_data.assigned_to_workers[worker] = []
-
-        return requeued_count
-
-async def _get_completed_count(multi_job_id):
-    """Get count of completed tasks."""
-    prompt_server = ensure_tile_jobs_initialized()
-    async with prompt_server.distributed_tile_jobs_lock:
-        job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
-        if isinstance(job_data, BaseJobState):
-            return len(job_data.completed_tasks)
-        return 0
-
-async def _mark_task_completed(multi_job_id, task_id, result):
-    """Mark a task as completed."""
-    prompt_server = ensure_tile_jobs_initialized()
-    async with prompt_server.distributed_tile_jobs_lock:
-        job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
-        if isinstance(job_data, BaseJobState):
-            job_data.completed_tasks[task_id] = result
 
 async def _send_heartbeat_to_master(multi_job_id, master_url, worker_id):
-    """Send heartbeat to master."""
+    """Send heartbeat to master from worker-side processing loops."""
     try:
         data = {'multi_job_id': multi_job_id, 'worker_id': str(worker_id)}
         session = await get_client_session()
@@ -347,282 +36,20 @@ async def _send_heartbeat_to_master(multi_job_id, master_url, worker_id):
     except Exception as e:
         debug_log(f"Heartbeat failed: {e}")
 
-async def _cleanup_job(multi_job_id):
-    """Cleanup the job data."""
-    prompt_server = ensure_tile_jobs_initialized()
-    async with prompt_server.distributed_tile_jobs_lock:
-        if multi_job_id in prompt_server.distributed_pending_tile_jobs:
-            del prompt_server.distributed_pending_tile_jobs[multi_job_id]
-            debug_log(f"Cleaned up job {multi_job_id}")
 
-# API Endpoints (generalized)
-
-@server.PromptServer.instance.routes.post("/distributed/heartbeat")
-async def heartbeat_endpoint(request):
-    try:
-        data = await request.json()
-        worker_id = data.get('worker_id')
-        multi_job_id = data.get('multi_job_id')
-        
-        if not worker_id or not multi_job_id:
-            return await handle_api_error(request, "Missing worker_id or multi_job_id", 400)
-        
-        prompt_server = ensure_tile_jobs_initialized()
-        async with prompt_server.distributed_tile_jobs_lock:
-            if multi_job_id in prompt_server.distributed_pending_tile_jobs:
-                job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
-                if isinstance(job_data, BaseJobState):
-                    job_data.worker_status[worker_id] = time.time()
-                    debug_log(f"Heartbeat from worker {worker_id}")
-                    return web.json_response({"status": "success"})
-                else:
-                    return await handle_api_error(request, "Worker status tracking not available", 400)
-            else:
-                return await handle_api_error(request, "Job not found", 404)
-    except Exception as e:
-        return await handle_api_error(request, e, 500)
-
-@server.PromptServer.instance.routes.post("/distributed/submit_tiles")
-async def submit_tiles_endpoint(request):
-    """Endpoint for workers to submit processed tiles in static mode."""
-    try:
-        content_length = request.headers.get('content-length')
-        if content_length and int(content_length) > MAX_PAYLOAD_SIZE:
-            return await handle_api_error(request, f"Payload too large: {content_length} bytes", 413)
-        
-        data = await request.post()
-        multi_job_id = data.get('multi_job_id')
-        worker_id = data.get('worker_id')
-        is_last = data.get('is_last', 'False').lower() == 'true'
-        
-        if multi_job_id is None or worker_id is None:
-            return await handle_api_error(request, "Missing multi_job_id or worker_id", 400)
-
-        prompt_server = ensure_tile_jobs_initialized()
-        
-        batch_size = int(data.get('batch_size', 0))
-        tiles = []
-        
-        # Handle completion signal
-        if batch_size == 0 and is_last:
-            async with prompt_server.distributed_tile_jobs_lock:
-                if multi_job_id in prompt_server.distributed_pending_tile_jobs:
-                    job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
-                    if not isinstance(job_data, TileJobState):
-                        return await handle_api_error(request, "Job not configured for tile submissions", 400)
-                    await job_data.queue.put({
-                        'worker_id': worker_id,
-                        'is_last': True,
-                        'tiles': []
-                    })
-                    debug_log(f"Received completion signal from worker {worker_id}")
-                    return web.json_response({"status": "success"})
-        
-        try:
-            tiles = _parse_tiles_from_form(data)
-        except ValueError as e:
-            return await handle_api_error(request, str(e), 400)
-
-        # Submit tiles to queue
-        async with prompt_server.distributed_tile_jobs_lock:
-            if multi_job_id in prompt_server.distributed_pending_tile_jobs:
-                job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
-                if not isinstance(job_data, TileJobState):
-                    return await handle_api_error(request, "Job not configured for tile submissions", 400)
-
-                q = job_data.queue
-                if batch_size > 0 or len(tiles) > 0:
-                    await q.put({
-                        'worker_id': worker_id,
-                        'tiles': tiles,
-                        'is_last': is_last
-                    })
-                    debug_log(f"Received {len(tiles)} tiles from worker {worker_id} (is_last={is_last})")
-                else:
-                    await q.put({
-                        'worker_id': worker_id,
-                        'is_last': True,
-                        'tiles': []
-                    })
-                    
-                return web.json_response({"status": "success"})
-            else:
-                return await handle_api_error(request, "Job not found", 404)
-    except Exception as e:
-        return await handle_api_error(request, e, 500)
-
-@server.PromptServer.instance.routes.post("/distributed/submit_image")
-async def submit_image_endpoint(request):
-    """Endpoint for workers to submit processed images in dynamic mode."""
-    try:
-        content_length = request.headers.get('content-length')
-        if content_length and int(content_length) > MAX_PAYLOAD_SIZE:
-            return await handle_api_error(request, f"Payload too large: {content_length} bytes", 413)
-        
-        data = await request.post()
-        multi_job_id = data.get('multi_job_id')
-        worker_id = data.get('worker_id')
-        is_last = data.get('is_last', 'False').lower() == 'true'
-        
-        if multi_job_id is None or worker_id is None:
-            return await handle_api_error(request, "Missing multi_job_id or worker_id", 400)
-
-        prompt_server = ensure_tile_jobs_initialized()
-        
-        # Handle image submission
-        if 'full_image' in data and 'image_idx' in data:
-            image_idx = int(data.get('image_idx'))
-            img_data = data['full_image'].file.read()
-            img = Image.open(io.BytesIO(img_data)).convert("RGB")
-            
-            debug_log(f"Received full image {image_idx} from worker {worker_id}")
-            
-            async with prompt_server.distributed_tile_jobs_lock:
-                if multi_job_id in prompt_server.distributed_pending_tile_jobs:
-                    job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
-                    if not isinstance(job_data, ImageJobState):
-                        return await handle_api_error(request, "Job not configured for image submissions", 400)
-                    await job_data.queue.put({
-                        'worker_id': worker_id,
-                        'image_idx': image_idx,
-                        'image': img,
-                        'is_last': is_last
-                    })
-                    return web.json_response({"status": "success"})
-        
-        # Handle completion signal (no image data)
-        elif is_last:
-            async with prompt_server.distributed_tile_jobs_lock:
-                if multi_job_id in prompt_server.distributed_pending_tile_jobs:
-                    job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
-                    if not isinstance(job_data, ImageJobState):
-                        return await handle_api_error(request, "Job not configured for image submissions", 400)
-                    await job_data.queue.put({
-                        'worker_id': worker_id,
-                        'is_last': True,
-                    })
-                    debug_log(f"Received completion signal from worker {worker_id}")
-                    return web.json_response({"status": "success"})
-        else:
-            return await handle_api_error(request, "Missing image data or invalid request", 400)
-            
-        return await handle_api_error(request, "Job not found", 404)
-    except Exception as e:
-        return await handle_api_error(request, e, 500)
-
-# Helper functions for shallow copying conditioning without duplicating models
-def clone_control_chain(control, clone_hint=True):
-    """Shallow copy the ControlNet chain, optionally cloning hints but sharing models."""
-    if control is None:
-        return None
-    new_control = copy.copy(control)  # Shallow copy (shares model)
-    if clone_hint and hasattr(control, 'cond_hint_original'):
-        hint = getattr(control, 'cond_hint_original', None)
-        new_control.cond_hint_original = hint.clone() if hint is not None else None
-    if hasattr(control, 'previous_controlnet'):
-        new_control.previous_controlnet = clone_control_chain(control.previous_controlnet, clone_hint)
-    return new_control
-
-def clone_conditioning(cond_list, clone_hints=True):
-    """Clone conditioning without duplicating ControlNet models."""
-    new_cond = []
-    for emb, cond_dict in cond_list:
-        new_emb = emb.clone() if emb is not None else None
-        new_dict = cond_dict.copy()
-        if 'control' in new_dict:
-            new_dict['control'] = clone_control_chain(new_dict['control'], clone_hints)
-        if 'mask' in new_dict:
-            if new_dict['mask'] is not None:
-                new_dict['mask'] = new_dict['mask'].clone()
-        # Handle other potential fields if needed
-        if 'pooled_output' in new_dict:
-            if new_dict['pooled_output'] is not None:
-                new_dict['pooled_output'] = new_dict['pooled_output'].clone()
-        if 'area' in new_dict:
-            new_dict['area'] = new_dict['area'][:]  # Shallow copy list/tuple
-        new_cond.append([new_emb, new_dict])
-    return new_cond
-
-# Helper function to ensure persistent state is initialized
-def ensure_tile_jobs_initialized():
-    """Ensure tile job storage is initialized on the server instance."""
-    prompt_server = server.PromptServer.instance
-    if not hasattr(prompt_server, 'distributed_pending_tile_jobs'):
-        debug_log("Initializing persistent tile job queue on server instance.")
-        prompt_server.distributed_pending_tile_jobs = {}
-        prompt_server.distributed_tile_jobs_lock = asyncio.Lock()
-    else:
-        invalid_job_ids = [
-            job_id
-            for job_id, job_data in prompt_server.distributed_pending_tile_jobs.items()
-            if not isinstance(job_data, BaseJobState)
-        ]
-        for job_id in invalid_job_ids:
-            debug_log(f"Removing invalid job state for {job_id}")
-            del prompt_server.distributed_pending_tile_jobs[job_id]
-    return prompt_server
-
-# API Endpoint for tile completion
-@server.PromptServer.instance.routes.post("/distributed/request_image")
-async def request_image_endpoint(request):
-    """Endpoint for workers to request tasks (images in dynamic mode, tiles in static mode)."""
-    try:
-        data = await request.json()
-        worker_id = data.get('worker_id')
-        multi_job_id = data.get('multi_job_id')
-        
-        if not worker_id or not multi_job_id:
-            return await handle_api_error(request, "Missing worker_id or multi_job_id", 400)
-        
-        prompt_server = ensure_tile_jobs_initialized()
-        async with prompt_server.distributed_tile_jobs_lock:
-            if multi_job_id in prompt_server.distributed_pending_tile_jobs:
-                job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
-                if not isinstance(job_data, BaseJobState):
-                    return await handle_api_error(request, "Invalid job data structure", 500)
-
-                mode = job_data.mode
-                if isinstance(job_data, ImageJobState):
-                    pending_queue = job_data.pending_images
-                elif isinstance(job_data, TileJobState):
-                    pending_queue = job_data.pending_tasks
-                else:
-                    return await handle_api_error(request, "Invalid job configuration", 400)
-                
-                try:
-                    task_idx = await asyncio.wait_for(pending_queue.get(), timeout=0.1)
-                    # Track assigned task
-                    job_data.assigned_to_workers.setdefault(worker_id, []).append(task_idx)
-                    # Update worker heartbeat
-                    job_data.worker_status[worker_id] = time.time()
-                    # Get estimated remaining count
-                    remaining = pending_queue.qsize()  # Approximate
-                    
-                    # Return appropriate response based on mode
-                    if mode == 'dynamic':
-                        debug_log(f"UltimateSDUpscale API - Assigned image {task_idx} to worker {worker_id}")
-                        return web.json_response({"image_idx": task_idx, "estimated_remaining": remaining})
-                    else:  # static
-                        debug_log(f"UltimateSDUpscale API - Assigned tile {task_idx} to worker {worker_id}")
-                        return web.json_response({"tile_idx": task_idx, "estimated_remaining": remaining, "batched_static": job_data.batched_static})
-                except asyncio.TimeoutError:
-                    if mode == 'dynamic':
-                        return web.json_response({"image_idx": None})
-                    else:
-                        return web.json_response({"tile_idx": None})
-            else:
-                return await handle_api_error(request, "Job not found", 404)
-    except Exception as e:
-        return await handle_api_error(request, e, 500)
-
-@server.PromptServer.instance.routes.get("/distributed/job_status")
-async def job_status_endpoint(request):
-    """Endpoint to check if a job is ready."""
-    multi_job_id = request.query.get('multi_job_id')
-    if not multi_job_id:
-        return web.json_response({"ready": False})
-    prompt_server = ensure_tile_jobs_initialized()
-    async with prompt_server.distributed_tile_jobs_lock:
-        job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
-        ready = bool(isinstance(job_data, BaseJobState) and job_data.queue is not None)
-        return web.json_response({"ready": ready})
+__all__ = [
+    "MAX_PAYLOAD_SIZE",
+    "_check_and_requeue_timed_out_workers",
+    "_cleanup_job",
+    "_drain_results_queue",
+    "_get_completed_count",
+    "_init_job_queue",
+    "_mark_task_completed",
+    "_parse_tiles_from_form",
+    "_send_heartbeat_to_master",
+    "clone_conditioning",
+    "clone_control_chain",
+    "ensure_tile_jobs_initialized",
+    "init_dynamic_job",
+    "init_static_job_batched",
+]

web/apiClient.js

@@ -4,27 +4,40 @@ import { normalizeWorkerUrl } from './urlUtils.js';
 export function createApiClient(baseUrl) {
     const normalizedBaseUrl = normalizeWorkerUrl(baseUrl);
 
-    const request = async (endpoint, options = {}, retries = TIMEOUTS.MAX_RETRIES) => {
+    const request = async (
+        endpoint,
+        options = {},
+        { retries = TIMEOUTS.MAX_RETRIES, retry = true } = {},
+    ) => {
+        const maxAttempts = retry ? retries : 1;
         let lastError;
         let delay = TIMEOUTS.RETRY_DELAY; // Initial delay for exponential backoff
 
-        for (let attempt = 0; attempt < retries; attempt++) {
+        for (let attempt = 0; attempt < maxAttempts; attempt++) {
             try {
+                const headers = {
+                    'Content-Type': 'application/json',
+                    ...(options.headers || {}),
+                };
                 const response = await fetch(`${normalizedBaseUrl}${endpoint}`, {
-                    headers: { 'Content-Type': 'application/json' },
-                    ...options
+                    ...options,
+                    headers,
                 });
                 
                 if (!response.ok) {
-                    const error = await response.json().catch(() => ({ message: 'Request failed' }));
-                    throw new Error(error.message || `HTTP ${response.status}`);
+                    const error = await response.json().catch(() => ({}));
+                    const message = error.message
+                        || error.error
+                        || (Array.isArray(error.errors) ? error.errors.join('; ') : null)
+                        || `HTTP ${response.status}`;
+                    throw new Error(message);
                 }
                 
                 return await response.json();
             } catch (error) {
                 lastError = error;
-                console.log(`API Error (attempt ${attempt + 1}/${retries}): ${endpoint} - ${error.message}`);
-                if (attempt < retries - 1) {
+                console.log(`API Error (attempt ${attempt + 1}/${maxAttempts}): ${endpoint} - ${error.message}`);
+                if (attempt < maxAttempts - 1) {
                     await new Promise(resolve => setTimeout(resolve, delay));
                     delay *= 2; // Exponential backoff
                 }
@@ -32,6 +45,39 @@ export function createApiClient(baseUrl) {
         }
         throw lastError;
     };
+
+    const requestUrl = async (
+        url,
+        options = {},
+        { retries = TIMEOUTS.MAX_RETRIES, retry = true } = {},
+    ) => {
+        const maxAttempts = retry ? retries : 1;
+        let lastError;
+        let delay = TIMEOUTS.RETRY_DELAY;
+
+        for (let attempt = 0; attempt < maxAttempts; attempt++) {
+            try {
+                const response = await fetch(url, options);
+                if (!response.ok) {
+                    const error = await response.json().catch(() => ({}));
+                    const message = error.message
+                        || error.error
+                        || (Array.isArray(error.errors) ? error.errors.join('; ') : null)
+                        || `HTTP ${response.status}`;
+                    throw new Error(message);
+                }
+                return await response.json();
+            } catch (error) {
+                lastError = error;
+                console.log(`API Error (attempt ${attempt + 1}/${maxAttempts}): ${url} - ${error.message}`);
+                if (attempt < maxAttempts - 1) {
+                    await new Promise(resolve => setTimeout(resolve, delay));
+                    delay *= 2;
+                }
+            }
+        }
+        throw lastError;
+    };
     
     return {
         // Config endpoints
@@ -43,28 +89,28 @@ export function createApiClient(baseUrl) {
             return request('/distributed/config/update_worker', {
                 method: 'POST',
                 body: JSON.stringify({ worker_id: workerId, ...data })
-            });
+            }, { retry: false });
         },
         
         async deleteWorker(workerId) {
             return request('/distributed/config/delete_worker', {
                 method: 'POST',
                 body: JSON.stringify({ worker_id: workerId })
-            });
+            }, { retry: false });
         },
         
         async updateSetting(key, value) {
             return request('/distributed/config/update_setting', {
                 method: 'POST',
                 body: JSON.stringify({ key, value })
-            });
+            }, { retry: false });
         },
         
         async updateMaster(data) {
             return request('/distributed/config/update_master', {
                 method: 'POST',
                 body: JSON.stringify(data)
-            });
+            }, { retry: false });
         },
         
         // Worker management endpoints
@@ -72,14 +118,14 @@ export function createApiClient(baseUrl) {
             return request('/distributed/launch_worker', {
                 method: 'POST',
                 body: JSON.stringify({ worker_id: workerId })
-            });
+            }, { retry: false });
         },
         
         async stopWorker(workerId) {
             return request('/distributed/stop_worker', {
                 method: 'POST',
                 body: JSON.stringify({ worker_id: workerId })
-            });
+            }, { retry: false });
         },
         
         async getManagedWorkers() {
@@ -94,51 +140,59 @@ export function createApiClient(baseUrl) {
             return request('/distributed/worker/clear_launching', {
                 method: 'POST',
                 body: JSON.stringify({ worker_id: workerId })
-            });
+            }, { retry: false });
         },
         
         async queueDistributed(payload) {
             return request('/distributed/queue', {
                 method: 'POST',
+                headers: {
+                    ...(payload?.trace_execution_id
+                        ? { 'X-Idempotency-Key': payload.trace_execution_id }
+                        : {}),
+                },
                 body: JSON.stringify(payload)
-            });
+            }, { retry: false });
         },
 
-        async probeWorker(workerUrl, timeoutMs = TIMEOUTS.STATUS_CHECK) {
+        async probeWorker(workerUrl, timeoutMs = TIMEOUTS.STATUS_CHECK, signal = null) {
             const normalizedWorkerUrl = normalizeWorkerUrl(workerUrl);
-            const response = await fetch(`${normalizedWorkerUrl}/prompt`, {
-                method: 'GET',
-                mode: 'cors',
-                cache: 'no-store',
-                signal: AbortSignal.timeout(timeoutMs)
-            });
+            const controller = new AbortController();
+            const timeoutId = setTimeout(() => controller.abort(), timeoutMs);
+            const effectiveSignal = signal
+                ? AbortSignal.any([controller.signal, signal])
+                : controller.signal;
+            try {
+                const response = await fetch(`${normalizedWorkerUrl}/prompt`, {
+                    method: 'GET',
+                    mode: 'cors',
+                    cache: 'no-store',
+                    signal: effectiveSignal,
+                });
 
-            if (!response.ok) {
-                return { ok: false, status: response.status, queueRemaining: null };
+                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,
+                };
+            } finally {
+                clearTimeout(timeoutId);
             }
-
-            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`, {
+            return requestUrl(`${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();
+            }, { retry: false });
         },
         
         // Network info
@@ -179,14 +233,14 @@ export function createApiClient(baseUrl) {
             return request('/distributed/tunnel/start', {
                 method: 'POST',
                 body: JSON.stringify({})
-            });
+            }, { retry: false });
         },
 
         async stopTunnel() {
             return request('/distributed/tunnel/stop', {
                 method: 'POST',
                 body: JSON.stringify({})
-            });
+            }, { retry: false });
         },
 
         async getTunnelStatus() {

web/distributed.css

@@ -5,6 +5,13 @@
     --btn-working: #666;
     --btn-success: #3a6a3a;
     --btn-error: #6a3a3a;
+    --tunnel-enable: #665533;
+    --tunnel-disable: #7c4a4a;
+    --master-badge-fallback-bg: #243024;
+    --master-badge-fallback-text: #6bd06b;
+    --master-badge-fallback-border: #335533;
+    --master-badge-delegate-bg: #3a3a3a;
+    --master-badge-delegate-text: #ffcc66;
 }
 
 .btn--stop {
@@ -30,3 +37,167 @@
 .btn--error {
     background-color: var(--btn-error) !important;
 }
+
+.master-info-badge--fallback {
+    background-color: var(--master-badge-fallback-bg) !important;
+    color: var(--master-badge-fallback-text) !important;
+    border: 1px solid var(--master-badge-fallback-border) !important;
+}
+
+.master-info-badge--delegate {
+    background-color: var(--master-badge-delegate-bg) !important;
+    color: var(--master-badge-delegate-text) !important;
+}
+
+.entity-card-content--hovered {
+    background-color: #333 !important;
+}
+
+.placeholder-card--blueprint {
+    border-color: #555 !important;
+    background-color: rgba(255, 255, 255, 0.02) !important;
+}
+
+.placeholder-card--blueprint.is-hovered {
+    border-color: #777 !important;
+    background-color: rgba(255, 255, 255, 0.05) !important;
+}
+
+.placeholder-card--add {
+    border-color: #444 !important;
+    background-color: transparent !important;
+}
+
+.placeholder-card--add.is-hovered {
+    border-color: #666 !important;
+    background-color: rgba(255, 255, 255, 0.02) !important;
+}
+
+.placeholder-column--blueprint {
+    color: #777 !important;
+}
+
+.placeholder-column--blueprint.is-hovered {
+    color: #999 !important;
+}
+
+.placeholder-column--add {
+    color: #555 !important;
+    border-color: #444 !important;
+}
+
+.placeholder-column--add.is-hovered {
+    color: #888 !important;
+    border-color: #666 !important;
+}
+
+.tunnel-button--enable {
+    background-color: var(--tunnel-enable) !important;
+}
+
+.tunnel-button--disable {
+    background-color: var(--tunnel-disable) !important;
+}
+
+.tunnel-status--enable {
+    background-color: var(--tunnel-enable) !important;
+}
+
+.tunnel-status--disable {
+    background-color: var(--tunnel-disable) !important;
+}
+
+.is-hidden {
+    display: none !important;
+}
+
+.settings-arrow--expanded {
+    transform: rotate(90deg) !important;
+}
+
+.log-modal {
+    position: fixed;
+    top: 0;
+    left: 0;
+    width: 100%;
+    height: 100%;
+    background: rgba(0, 0, 0, 0.8);
+    display: flex;
+    align-items: center;
+    justify-content: center;
+    z-index: 10000;
+}
+
+.log-modal__content {
+    background: #1e1e1e;
+    border-radius: 8px;
+    width: 90%;
+    max-width: 1200px;
+    height: 80%;
+    display: flex;
+    flex-direction: column;
+    border: 1px solid #444;
+}
+
+.log-modal__header {
+    padding: 15px 20px;
+    border-bottom: 1px solid #444;
+    display: flex;
+    justify-content: space-between;
+    align-items: center;
+}
+
+.log-modal__title {
+    margin: 0;
+    color: #fff;
+}
+
+.log-modal__header-buttons {
+    display: flex;
+    gap: 20px;
+    align-items: center;
+}
+
+.log-modal__refresh {
+    display: flex;
+    align-items: center;
+    gap: 4px;
+}
+
+.log-modal__refresh-input {
+    cursor: pointer;
+}
+
+.log-modal__refresh-label {
+    font-size: 12px;
+    color: #ccc;
+    cursor: pointer;
+    white-space: nowrap;
+}
+
+.log-modal__close {
+    background-color: #c04c4c !important;
+    padding: 5px 10px !important;
+    font-size: 14px !important;
+    font-weight: bold !important;
+}
+
+.log-modal__body {
+    flex: 1;
+    overflow: auto;
+    padding: 15px;
+    font-family: 'Consolas', 'Monaco', 'Courier New', monospace;
+    font-size: 12px;
+    line-height: 1.4;
+    color: #ddd;
+    background: #0d0d0d;
+    white-space: pre-wrap;
+    word-wrap: break-word;
+}
+
+.log-modal__status {
+    padding: 10px 20px;
+    border-top: 1px solid #444;
+    font-size: 11px;
+    color: #888;
+}

web/executionUtils.js

@@ -1,5 +1,5 @@
 import { api } from "../../scripts/api.js";
-import { findNodesByClass } from './workerUtils.js';
+import { applyProbeResultToWorkerDot, findNodesByClass } from './workerUtils.js';
 import { TIMEOUTS, NODE_CLASSES, generateUUID } from './constants.js';
 import { markSkipDispatch, selectLeastBusyWorker } from './executionDecisionUtils.js';
 import { checkAllWorkerStatuses, getWorkerUrl } from './workerLifecycle.js';
@@ -232,13 +232,7 @@ export async function performPreflightCheck(extension, workers) {
     
     // Update UI status indicators for inactive workers
     results.filter(r => !r.active).forEach(r => {
-        const statusDot = document.getElementById(`status-${r.worker.id}`);
-        if (statusDot) {
-            // Remove pulsing animation once status is determined
-            statusDot.classList.remove('status-pulsing', 'worker-status--online', 'worker-status--unknown');
-            statusDot.classList.add('worker-status--offline');
-            statusDot.title = "Offline - Cannot connect";
-        }
+        applyProbeResultToWorkerDot(r.worker.id, { ok: false });
     });
     
     return activeWorkers;

web/ui.js

@@ -3,7 +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 { createLogModal } from './ui/logModal.js';
+import { launchWorker, stopWorker, updateWorkerControls, viewWorkerLog } from './workerLifecycle.js';
 import { isRemoteWorker } from './workerSettings.js';
 
 const cardConfigs = {
@@ -377,184 +378,34 @@ export class DistributedUI {
             [STATUS_COLORS.PROCESSING_YELLOW]: "worker-status--processing",
         };
 
-        const statusClass = colorClassMap[color];
-        if (statusClass) {
-            statusDot.classList.add(statusClass);
-            statusDot.style.backgroundColor = "";
-        } else {
-            statusDot.style.backgroundColor = color;
-        }
+        const statusClass = colorClassMap[color] || "worker-status--unknown";
+        statusDot.classList.add(statusClass);
+        statusDot.style.backgroundColor = "";
         statusDot.title = title;
         statusDot.classList.toggle('status-pulsing', pulsing);
     }
 
     showLogModal(extension, workerId, logData) {
-        // Remove any existing modal
-        const existingModal = document.getElementById('distributed-log-modal');
-        if (existingModal) {
-            existingModal.remove();
+        if (this._logModal) {
+            this._logModal.unmount();
+            this._logModal = null;
         }
-        
+
         const worker = extension.config.workers.find(w => w.id === workerId);
         const workerName = worker?.name || `Worker ${workerId}`;
-        
-        // Create modal container
-        const modal = document.createElement('div');
-        modal.id = 'distributed-log-modal';
-        modal.style.cssText = `
-            position: fixed;
-            top: 0;
-            left: 0;
-            width: 100%;
-            height: 100%;
-            background: rgba(0, 0, 0, 0.8);
-            display: flex;
-            align-items: center;
-            justify-content: center;
-            z-index: 10000;
-        `;
-        
-        // Create modal content
-        const content = document.createElement('div');
-        content.style.cssText = `
-            background: #1e1e1e;
-            border-radius: 8px;
-            width: 90%;
-            max-width: 1200px;
-            height: 80%;
-            display: flex;
-            flex-direction: column;
-            border: 1px solid #444;
-        `;
-        
-        // Header
-        const header = document.createElement('div');
-        header.style.cssText = `
-            padding: 15px 20px;
-            border-bottom: 1px solid #444;
-            display: flex;
-            justify-content: space-between;
-            align-items: center;
-        `;
-        
-        const title = document.createElement('h3');
-        title.textContent = `${workerName} - Log Viewer`;
-        title.style.cssText = 'margin: 0; color: #fff;';
-        
-        const headerButtons = document.createElement('div');
-        headerButtons.style.cssText = 'display: flex; gap: 20px; align-items: center;';
-        
-        // Auto-refresh container
-        const refreshContainer = document.createElement('div');
-        refreshContainer.style.cssText = 'display: flex; align-items: center; gap: 4px;';
-        
-        // Auto-refresh checkbox
-        const refreshCheckbox = document.createElement('input');
-        refreshCheckbox.type = 'checkbox';
-        refreshCheckbox.id = 'log-auto-refresh';
-        refreshCheckbox.checked = true; // Enabled by default
-        refreshCheckbox.style.cssText = 'cursor: pointer;';
-        refreshCheckbox.onchange = (e) => {
-            if (e.target.checked) {
-                startLogAutoRefresh(extension, workerId);
-            } else {
-                stopLogAutoRefresh(extension);
-            }
-        };
-        
-        const refreshLabel = document.createElement('label');
-        refreshLabel.htmlFor = 'log-auto-refresh';
-        refreshLabel.style.cssText = 'font-size: 12px; color: #ccc; cursor: pointer; white-space: nowrap;';
-        refreshLabel.textContent = 'Auto-refresh';
-        
-        // Add checkbox and label to container
-        refreshContainer.appendChild(refreshCheckbox);
-        refreshContainer.appendChild(refreshLabel);
-        
-        // Close button
-        const closeBtn = this.createButton('✕', 
-            () => {
-                stopLogAutoRefresh(extension);
-                modal.remove();
-            }, 
-            'background-color: #c04c4c;');
-        closeBtn.style.cssText += ' padding: 5px 10px; font-size: 14px; font-weight: bold;';
-        
-        headerButtons.appendChild(refreshContainer);
-        headerButtons.appendChild(closeBtn);
-        
-        header.appendChild(title);
-        header.appendChild(headerButtons);
-        
-        // Log content area
-        const logContainer = document.createElement('div');
-        logContainer.style.cssText = `
-            flex: 1;
-            overflow: auto;
-            padding: 15px;
-            font-family: 'Consolas', 'Monaco', 'Courier New', monospace;
-            font-size: 12px;
-            line-height: 1.4;
-            color: #ddd;
-            background: #0d0d0d;
-            white-space: pre-wrap;
-            word-wrap: break-word;
-        `;
-        logContainer.id = 'distributed-log-content';
-        logContainer.textContent = logData.content;
-        
-        // Auto-scroll to bottom
-        setTimeout(() => {
-            logContainer.scrollTop = logContainer.scrollHeight;
-        }, 0);
-        
-        // Status bar
-        const statusBar = document.createElement('div');
-        statusBar.style.cssText = `
-            padding: 10px 20px;
-            border-top: 1px solid #444;
-            font-size: 11px;
-            color: #888;
-        `;
-        statusBar.textContent = `Log file: ${logData.log_file}`;
-        if (logData.truncated) {
-            statusBar.textContent += ` (showing last ${logData.lines_shown} lines of ${this.formatFileSize(logData.file_size)})`;
-        }
-        
-        // Assemble modal
-        content.appendChild(header);
-        content.appendChild(logContainer);
-        content.appendChild(statusBar);
-        modal.appendChild(content);
-        
-        // Close on background click
-        modal.onclick = (e) => {
-            if (e.target === modal) {
-                stopLogAutoRefresh(extension);
-                modal.remove();
-            }
-        };
-        
-        // Close on Escape key
-        const handleEscape = (e) => {
-            if (e.key === 'Escape') {
-                stopLogAutoRefresh(extension);
-                modal.remove();
-                document.removeEventListener('keydown', handleEscape);
-            }
-        };
-        document.addEventListener('keydown', handleEscape);
-        
-        document.body.appendChild(modal);
-        
-        // Start auto-refresh
-        startLogAutoRefresh(extension, workerId);
-    }
 
-    formatFileSize(bytes) {
-        if (bytes < 1024) return bytes + ' B';
-        if (bytes < 1024 * 1024) return (bytes / 1024).toFixed(1) + ' KB';
-        return (bytes / (1024 * 1024)).toFixed(1) + ' MB';
+        const modal = createLogModal();
+        this._logModal = modal;
+        modal.mount(document.body, {
+            workerName,
+            logData,
+            fetchLog: async () => extension.api.getWorkerLog(workerId, 1000),
+            onClose: () => {
+                if (this._logModal === modal) {
+                    this._logModal = null;
+                }
+            },
+        });
     }
 
     createWorkerSettingsForm(extension, worker) {
@@ -712,17 +563,15 @@ export class DistributedUI {
             let message;
             const badge = document.createElement("div");
             badge.style.cssText = this.styles.infoBox;
+            badge.classList.add("master-info-badge");
             if (fallbackActive) {
                 message = "No workers selected. Master fallback execution active.";
                 badge.textContent = message;
-                badge.style.backgroundColor = "#243024";
-                badge.style.color = "#6bd06b";
-                badge.style.border = "1px solid #335533";
+                badge.classList.add("master-info-badge--fallback");
             } else if (!participationEnabled) {
                 message = "Master disabled: running as orchestrator only.";
                 badge.textContent = message;
-                badge.style.backgroundColor = "#3a3a3a";
-                badge.style.color = "#ffcc66";
+                badge.classList.add("master-info-badge--delegate");
             } else {
                 message = "Master participating in workflows.";
                 badge.textContent = message;
@@ -822,9 +671,10 @@ export class DistributedUI {
         settingsForm.appendChild(hostResult.group);
 
         // Cloudflare tunnel toggle (simple button inside master settings)
-        const tunnelBtn = this.createButton("Enable Cloudflare Tunnel", (e) => extension.handleTunnelToggle(e.target), "background-color: #665533;");
+        const tunnelBtn = this.createButton("Enable Cloudflare Tunnel", (e) => extension.handleTunnelToggle(e.target));
         tunnelBtn.id = "cloudflare-tunnel-button";
-        tunnelBtn.style.cssText = BUTTON_STYLES.base + " background-color: #665533; margin: 4px 0 -5px 0;";
+        tunnelBtn.style.cssText = BUTTON_STYLES.base + " margin: 4px 0 -5px 0;";
+        tunnelBtn.classList.add("tunnel-button", "tunnel-button--enable");
         settingsForm.appendChild(tunnelBtn);
         extension.tunnelElements = { button: tunnelBtn };
         extension.updateTunnelUIElements();
@@ -890,30 +740,19 @@ export class DistributedUI {
     }
 
     addPlaceholderHover(card, leftColumn, entityType) {
+        const cardTypeClass = entityType === 'blueprint' ? 'placeholder-card--blueprint' : 'placeholder-card--add';
+        const columnTypeClass = entityType === 'blueprint' ? 'placeholder-column--blueprint' : 'placeholder-column--add';
+        card.classList.add('placeholder-card', cardTypeClass);
+        leftColumn.classList.add('placeholder-column', columnTypeClass);
+
         card.onmouseover = () => {
-            if (entityType === 'blueprint') {
-                card.style.borderColor = "#777";
-                card.style.backgroundColor = "rgba(255, 255, 255, 0.05)";
-                leftColumn.style.color = "#999";
-            } else {
-                card.style.borderColor = "#666";
-                card.style.backgroundColor = "rgba(255, 255, 255, 0.02)";
-                leftColumn.style.color = "#888";
-                leftColumn.style.borderColor = "#666";
-            }
+            card.classList.add('is-hovered');
+            leftColumn.classList.add('is-hovered');
         };
         
         card.onmouseout = () => {
-            if (entityType === 'blueprint') {
-                card.style.borderColor = "#555";
-                card.style.backgroundColor = "rgba(255, 255, 255, 0.02)";
-                leftColumn.style.color = "#777";
-            } else {
-                card.style.borderColor = "#444";
-                card.style.backgroundColor = "transparent";
-                leftColumn.style.color = "#555";
-                leftColumn.style.borderColor = "#444";
-            }
+            card.classList.remove('is-hovered');
+            leftColumn.classList.remove('is-hovered');
         };
     }
 

web/ui/logModal.js

@@ -0,0 +1,192 @@
+import { TIMEOUTS } from '../constants.js';
+
+function formatFileSize(bytes) {
+    if (bytes < 1024) return bytes + ' B';
+    if (bytes < 1024 * 1024) return (bytes / 1024).toFixed(1) + ' KB';
+    return (bytes / (1024 * 1024)).toFixed(1) + ' MB';
+}
+
+export function createLogModal() {
+    let _modalEl = null;
+    let _keydownHandler = null;
+    let _refreshTimer = null;
+    let _fetching = false;
+    let _fetchLog = null;
+    let _onClose = null;
+    let _logContentEl = null;
+    let _statusBarEl = null;
+    let _refreshCheckbox = null;
+
+    const updateLogView = (logData) => {
+        if (!_logContentEl || !_statusBarEl || !logData) {
+            return;
+        }
+
+        const shouldAutoScroll =
+            _logContentEl.scrollTop + _logContentEl.clientHeight >= _logContentEl.scrollHeight - 50;
+        _logContentEl.textContent = logData.content || '';
+        if (shouldAutoScroll) {
+            _logContentEl.scrollTop = _logContentEl.scrollHeight;
+        }
+
+        let statusText = `Log file: ${logData.log_file || 'unknown'}`;
+        if (logData.truncated) {
+            statusText += ` (showing last ${logData.lines_shown} lines of ${formatFileSize(logData.file_size || 0)})`;
+        }
+        _statusBarEl.textContent = statusText;
+    };
+
+    const refreshLog = async () => {
+        if (_fetching || !_fetchLog || !_refreshCheckbox?.checked) {
+            return;
+        }
+        _fetching = true;
+        try {
+            const data = await _fetchLog();
+            if (data) {
+                updateLogView(data);
+            }
+        } catch (_error) {
+            // Keep the modal open and continue retrying on next interval.
+        } finally {
+            _fetching = false;
+        }
+    };
+
+    const stopRefresh = () => {
+        if (_refreshTimer) {
+            clearInterval(_refreshTimer);
+            _refreshTimer = null;
+        }
+    };
+
+    const startRefresh = () => {
+        stopRefresh();
+        _refreshTimer = setInterval(() => {
+            refreshLog();
+        }, TIMEOUTS.LOG_REFRESH);
+    };
+
+    const unmount = () => {
+        stopRefresh();
+
+        if (_keydownHandler) {
+            document.removeEventListener('keydown', _keydownHandler);
+            _keydownHandler = null;
+        }
+
+        if (_modalEl) {
+            _modalEl.remove();
+            _modalEl = null;
+        }
+
+        const onClose = _onClose;
+        _onClose = null;
+        if (onClose) {
+            onClose();
+        }
+    };
+
+    const mount = (container, { workerName, logData, onClose, fetchLog }) => {
+        _onClose = onClose;
+        _fetchLog = fetchLog;
+
+        const modal = document.createElement('div');
+        modal.id = 'distributed-log-modal';
+        modal.className = 'log-modal';
+
+        const content = document.createElement('div');
+        content.className = 'log-modal__content';
+
+        const header = document.createElement('div');
+        header.className = 'log-modal__header';
+
+        const title = document.createElement('h3');
+        title.className = 'log-modal__title';
+        title.textContent = `${workerName} - Log Viewer`;
+
+        const headerButtons = document.createElement('div');
+        headerButtons.className = 'log-modal__header-buttons';
+
+        const refreshContainer = document.createElement('div');
+        refreshContainer.className = 'log-modal__refresh';
+
+        const refreshCheckbox = document.createElement('input');
+        refreshCheckbox.type = 'checkbox';
+        refreshCheckbox.id = 'log-auto-refresh';
+        refreshCheckbox.className = 'log-modal__refresh-input';
+        refreshCheckbox.checked = true;
+
+        const refreshLabel = document.createElement('label');
+        refreshLabel.htmlFor = 'log-auto-refresh';
+        refreshLabel.className = 'log-modal__refresh-label';
+        refreshLabel.textContent = 'Auto-refresh';
+
+        refreshContainer.appendChild(refreshCheckbox);
+        refreshContainer.appendChild(refreshLabel);
+
+        const closeBtn = document.createElement('button');
+        closeBtn.className = 'distributed-button log-modal__close';
+        closeBtn.textContent = '✕';
+
+        headerButtons.appendChild(refreshContainer);
+        headerButtons.appendChild(closeBtn);
+
+        header.appendChild(title);
+        header.appendChild(headerButtons);
+
+        const logContainer = document.createElement('div');
+        logContainer.className = 'log-modal__body';
+        logContainer.id = 'distributed-log-content';
+
+        const statusBar = document.createElement('div');
+        statusBar.className = 'log-modal__status';
+
+        content.appendChild(header);
+        content.appendChild(logContainer);
+        content.appendChild(statusBar);
+        modal.appendChild(content);
+
+        closeBtn.addEventListener('click', unmount);
+        modal.addEventListener('click', (event) => {
+            if (event.target === modal) {
+                unmount();
+            }
+        });
+
+        _keydownHandler = (event) => {
+            if (event.key === 'Escape') {
+                unmount();
+            }
+        };
+        document.addEventListener('keydown', _keydownHandler);
+
+        _modalEl = modal;
+        _logContentEl = logContainer;
+        _statusBarEl = statusBar;
+        _refreshCheckbox = refreshCheckbox;
+
+        refreshCheckbox.addEventListener('change', () => {
+            if (refreshCheckbox.checked) {
+                refreshLog();
+            }
+        });
+
+        container.appendChild(modal);
+        updateLogView(logData);
+
+        requestAnimationFrame(() => {
+            if (_logContentEl) {
+                _logContentEl.scrollTop = _logContentEl.scrollHeight;
+            }
+        });
+
+        startRefresh();
+    };
+
+    return {
+        mount,
+        unmount,
+        update: updateLogView,
+    };
+}

web/workerLifecycle.js

@@ -1,6 +1,7 @@
-import { TIMEOUTS, STATUS_COLORS, ENDPOINTS } from './constants.js';
+import { TIMEOUTS, STATUS_COLORS } from './constants.js';
 import { buildWorkerUrl, normalizeWorkerUrl } from './urlUtils.js';
 import { isRemoteWorker } from './workerSettings.js';
+import { applyProbeResultToWorkerDot } from './workerUtils.js';
 
 export { normalizeWorkerUrl };
 
@@ -74,40 +75,45 @@ export async function checkAllWorkerStatuses(extension) {
 
 export async function checkMasterStatus(extension) {
     try {
-        const response = await fetch(`${window.location.origin}${ENDPOINTS.PROMPT}`, {
-            method: "GET",
-            signal: AbortSignal.timeout(TIMEOUTS.STATUS_CHECK),
-        });
+        const signal = extension.statusCheckAbortController?.signal || null;
+        const probeResult = await extension.api.probeWorker(
+            window.location.origin,
+            TIMEOUTS.STATUS_CHECK,
+            signal,
+        );
+        if (!probeResult.ok) {
+            throw new Error(`HTTP ${probeResult.status}`);
+        }
 
-        if (response.ok) {
-            const data = await response.json();
-            const queueRemaining = data.exec_info?.queue_remaining || 0;
-            const isProcessing = queueRemaining > 0;
+        const queueRemaining = probeResult.queueRemaining || 0;
+        const isProcessing = queueRemaining > 0;
 
-            // Update master status in state
-            extension.state.setMasterStatus(isProcessing ? "processing" : "online");
+        // Update master status in state
+        extension.state.setMasterStatus(isProcessing ? "processing" : "online");
 
-            // Update master status dot
-            const statusDot = document.getElementById("master-status");
-            if (statusDot) {
-                if (!extension.isMasterParticipating()) {
-                    if (isProcessing) {
-                        setStatusDotClass(statusDot, "worker-status--processing");
-                        statusDot.title = `Orchestrating (${queueRemaining} in queue)`;
-                    } else {
-                        setStatusDotClass(statusDot, "worker-status--unknown");
-                        statusDot.title = "Master orchestrator only";
-                    }
-                } else if (isProcessing) {
+        // Update master status dot
+        const statusDot = document.getElementById("master-status");
+        if (statusDot) {
+            if (!extension.isMasterParticipating()) {
+                if (isProcessing) {
                     setStatusDotClass(statusDot, "worker-status--processing");
-                    statusDot.title = `Processing (${queueRemaining} in queue)`;
+                    statusDot.title = `Orchestrating (${queueRemaining} in queue)`;
                 } else {
-                    setStatusDotClass(statusDot, "worker-status--online");
-                    statusDot.title = "Online";
+                    setStatusDotClass(statusDot, "worker-status--unknown");
+                    statusDot.title = "Master orchestrator only";
                 }
+            } else if (isProcessing) {
+                setStatusDotClass(statusDot, "worker-status--processing");
+                statusDot.title = `Processing (${queueRemaining} in queue)`;
+            } else {
+                setStatusDotClass(statusDot, "worker-status--online");
+                statusDot.title = "Online";
             }
         }
     } catch (error) {
+        if (error?.name === "AbortError") {
+            return;
+        }
         // Master is always online (we're running on it), so keep it green
         const statusDot = document.getElementById("master-status");
         if (statusDot) {
@@ -127,47 +133,36 @@ export function getWorkerUrl(extension, worker, endpoint = "") {
 
 export async function checkWorkerStatus(extension, worker) {
     // Assume caller ensured enabled; proceed with check
-    const url = getWorkerUrl(extension, worker, ENDPOINTS.PROMPT);
+    const workerUrl = getWorkerUrl(extension, worker);
 
     try {
-        // Combine timeout with abort controller signal
-        const timeoutSignal = AbortSignal.timeout(TIMEOUTS.STATUS_CHECK);
-        const signal = extension.statusCheckAbortController
-            ? AbortSignal.any([timeoutSignal, extension.statusCheckAbortController.signal])
-            : timeoutSignal;
-
-        const response = await fetch(url, {
-            method: "GET",
-            mode: "cors",
+        const signal = extension.statusCheckAbortController?.signal || null;
+        const probeResult = await extension.api.probeWorker(
+            workerUrl,
+            TIMEOUTS.STATUS_CHECK,
             signal,
+        );
+        if (!probeResult.ok) {
+            throw new Error(`HTTP ${probeResult.status}`);
+        }
+
+        const queueRemaining = probeResult.queueRemaining || 0;
+        const isProcessing = queueRemaining > 0;
+
+        // Update status
+        extension.state.setWorkerStatus(worker.id, {
+            online: true,
+            processing: isProcessing,
+            queueCount: queueRemaining,
         });
 
-        if (response.ok) {
-            const data = await response.json();
-            const queueRemaining = data.exec_info?.queue_remaining || 0;
-            const isProcessing = queueRemaining > 0;
+        // Update status dot based on probe result
+        applyProbeResultToWorkerDot(worker.id, probeResult);
 
-            // Update status
-            extension.state.setWorkerStatus(worker.id, {
-                online: true,
-                processing: isProcessing,
-                queueCount: queueRemaining,
-            });
-
-            // Update status dot based on processing state
-            if (isProcessing) {
-                extension.ui.updateStatusDot(worker.id, STATUS_COLORS.PROCESSING_YELLOW, `Online - Processing (${queueRemaining} in queue)`, false);
-            } else {
-                extension.ui.updateStatusDot(worker.id, STATUS_COLORS.ONLINE_GREEN, "Online - Idle", false);
-            }
-
-            // Clear launching state since worker is now online
-            if (extension.state.isWorkerLaunching(worker.id)) {
-                extension.state.setWorkerLaunching(worker.id, false);
-                clearLaunchingFlag(extension, worker.id);
-            }
-        } else {
-            throw new Error(`HTTP ${response.status}`);
+        // Clear launching state since worker is now online
+        if (extension.state.isWorkerLaunching(worker.id)) {
+            extension.state.setWorkerLaunching(worker.id, false);
+            clearLaunchingFlag(extension, worker.id);
         }
     } catch (error) {
         // Don't process aborted requests
@@ -186,8 +181,8 @@ export async function checkWorkerStatus(extension, worker) {
         if (extension.state.isWorkerLaunching(worker.id)) {
             extension.ui.updateStatusDot(worker.id, STATUS_COLORS.PROCESSING_YELLOW, "Launching...", true);
         } else if (worker.enabled) {
-            // Only update to red if not currently launching AND still enabled
-            extension.ui.updateStatusDot(worker.id, STATUS_COLORS.OFFLINE_RED, "Offline - Cannot connect", false);
+            // Only update to red if not currently launching AND still enabled.
+            applyProbeResultToWorkerDot(worker.id, { ok: false });
         }
         // If disabled, don't update the dot (leave it gray)
 
@@ -394,34 +389,34 @@ 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 = "";
+        logBtn.classList.remove("is-hidden");
         setButtonClass(logBtn, "btn--log");
     } else if (logBtn && !managedInfo) {
-        logBtn.style.display = "none";
+        logBtn.classList.add("is-hidden");
     }
 
     if (status?.online || managedInfo) {
         // Worker is running or we just launched it
-        launchBtn.style.display = "none"; // Hide launch button when running
+        launchBtn.classList.add("is-hidden");
 
         if (managedInfo) {
             // Only show stop button if we manage this worker
-            stopBtn.style.display = "";
+            stopBtn.classList.remove("is-hidden");
             stopBtn.disabled = false;
             stopBtn.textContent = "Stop";
             setButtonClass(stopBtn, "btn--stop");
         } else {
             // Hide stop button for workers launched outside UI
-            stopBtn.style.display = "none";
+            stopBtn.classList.add("is-hidden");
         }
     } else {
         // Worker is not running
-        launchBtn.style.display = ""; // Show launch button
+        launchBtn.classList.remove("is-hidden");
         launchBtn.disabled = false;
         launchBtn.textContent = "Launch";
         setButtonClass(launchBtn, "btn--launch");
 
-        stopBtn.style.display = "none"; // Hide stop button when not running
+        stopBtn.classList.add("is-hidden");
     }
 }
 
@@ -545,21 +540,13 @@ export function toggleWorkerExpanded(extension, workerId) {
         extension.state.setWorkerExpanded(workerId, false);
         settingsDiv.classList.remove("expanded");
         if (settingsArrow) {
-            settingsArrow.style.transform = "rotate(0deg)";
+            settingsArrow.classList.remove("settings-arrow--expanded");
         }
-        // Animate padding to 0
-        settingsDiv.style.padding = "0 12px";
-        settingsDiv.style.marginTop = "0";
-        settingsDiv.style.marginBottom = "0";
     } else {
         extension.state.setWorkerExpanded(workerId, true);
         settingsDiv.classList.add("expanded");
         if (settingsArrow) {
-            settingsArrow.style.transform = "rotate(90deg)";
+            settingsArrow.classList.add("settings-arrow--expanded");
         }
-        // Animate padding to full
-        settingsDiv.style.padding = "12px";
-        settingsDiv.style.marginTop = "8px";
-        settingsDiv.style.marginBottom = "8px";
     }
 }

web/workerUtils.js

@@ -92,3 +92,34 @@ export function findNodesByClass(apiPrompt, className) {
         .filter(([, nodeData]) => nodeData.class_type === className)
         .map(([nodeId, nodeData]) => ({ id: nodeId, data: nodeData }));
 }
+
+
+export function applyProbeResultToWorkerDot(workerId, probeResult) {
+    const dot = document.getElementById(`status-${workerId}`);
+    if (!dot) {
+        return;
+    }
+
+    dot.classList.remove(
+        'worker-status--online',
+        'worker-status--offline',
+        'worker-status--processing',
+        'worker-status--unknown',
+        'status-pulsing',
+    );
+
+    if (!probeResult || !probeResult.ok) {
+        dot.classList.add('worker-status--offline');
+        dot.title = 'Offline - Cannot connect';
+        return;
+    }
+
+    if ((probeResult.queueRemaining || 0) > 0) {
+        dot.classList.add('worker-status--processing');
+        dot.title = `Processing (${probeResult.queueRemaining} queued)`;
+        return;
+    }
+
+    dot.classList.add('worker-status--online');
+    dot.title = 'Online - Idle';
+}

workers/detection.py

@@ -1,16 +1,11 @@
-import asyncio
 import os
-import socket
-import json
 import platform
 import uuid
-from multiprocessing import Queue
 
 import aiohttp
 
 from ..utils.network import normalize_host, get_client_session
-from ..utils.logging import debug_log, log
-from ..utils.config import load_config
+from ..utils.logging import debug_log
 
 
 async def is_local_worker(worker_config):
@@ -76,30 +71,3 @@ def is_runpod_environment():
     """Check if running in Runpod environment."""
     return (os.environ.get('RUNPOD_POD_ID') is not None or
             os.environ.get('RUNPOD_API_KEY') is not None)
-
-async def get_comms_channel(worker_id, worker_config):
-    """Get communication channel for a worker (HTTP URL or IPC Queue)."""
-    if await is_local_worker(worker_config):
-        comm_mode = worker_config.get('communicationMode', 'http')
-        
-        if comm_mode == 'ipc':
-            # Use multiprocessing Queue for IPC
-            from . import get_worker_manager
-            wm = get_worker_manager()
-            if worker_id not in wm.queues:
-                wm.queues[worker_id] = Queue()
-            return wm.queues[worker_id]
-        elif is_docker_environment():
-            # Docker: use host.docker.internal
-            return f"http://host.docker.internal:{worker_config['port']}"
-        else:
-            # Local: use loopback
-            return f"http://127.0.0.1:{worker_config['port']}"
-    elif worker_config.get('type') == 'cloud' and is_runpod_environment():
-        # Runpod same-host optimization (if detected)
-        host = normalize_host(worker_config.get('host', 'localhost')) or 'localhost'
-        return f"http://{host}:{worker_config['port']}"
-    else:
-        # Remote worker: use configured endpoint
-        host = normalize_host(worker_config.get('host', 'localhost')) or 'localhost'
-        return f"http://{host}:{worker_config['port']}"