Remove list splitter/collector feature and refactor distributed modules

- Remove DistributedListSplitter and DistributedListCollector nodes
- Remove /distributed/init_list_queue and /distributed/request_list_item endpoints
- Clean up all references from prompt_transform, queue_orchestration, constants, and tests
- Refactor and reorganize distributed modules for quality improvements

.claude/skills/desloppify/SKILL.md

@@ -0,0 +1,151 @@
+<!-- desloppify-begin -->
+<!-- desloppify-skill-version: 2 -->
+---
+name: desloppify
+description: >
+  Codebase health scanner and technical debt tracker. Use when the user asks
+  about code quality, technical debt, dead code, large files, god classes,
+  duplicate functions, code smells, naming issues, import cycles, or coupling
+  problems. Also use when asked for a health score, what to fix next, or to
+  create a cleanup plan. Supports 28 languages.
+allowed-tools: Bash(desloppify *)
+---
+
+# Desloppify
+
+## 1. Your Job
+
+Improve code quality by maximising the **strict score** honestly.
+
+**The main thing you do is run `desloppify next`** — it tells you exactly what to fix and how. Fix it, resolve it, run `next` again. Keep going.
+
+Follow the scan output's **INSTRUCTIONS FOR AGENTS** — don't substitute your own analysis.
+
+## 2. The Workflow
+
+Two loops. The **outer loop** rescans periodically to measure progress.
+The **inner loop** is where you spend most of your time: fixing issues one by one.
+
+### Outer loop — scan and check
+
+```bash
+desloppify scan --path .       # analyse the codebase
+desloppify status              # check scores — are we at target?
+```
+If not at target, work the inner loop. Rescan periodically — especially after clearing a cluster or batch of related fixes. Issues cascade-resolve and new ones may surface.
+
+### Inner loop — fix issues
+
+Repeat until the queue is clear:
+
+```
+1. desloppify next              ← tells you exactly what to fix next
+2. Fix the issue in code
+3. Resolve it (next shows you the exact command including required attestation)
+```
+
+Score may temporarily drop after fixes — cascade effects are normal, keep going.
+If `next` suggests an auto-fixer, run `desloppify fix <fixer> --dry-run` to preview, then apply.
+
+**To be strategic**, use `plan` to shape what `next` gives you:
+```bash
+desloppify plan                        # see the full ordered queue
+desloppify plan move <pat> top         # reorder — what unblocks the most?
+desloppify plan cluster create <name>  # group related issues to batch-fix
+desloppify plan focus <cluster>        # scope next to one cluster
+desloppify plan defer <pat>            # push low-value items aside
+desloppify plan skip <pat>             # hide from next
+desloppify plan done <pat>             # mark complete
+desloppify plan reopen <pat>           # reopen
+```
+
+### Subjective reviews
+
+The scan will prompt you when a subjective review is needed — just follow its instructions.
+If you need to trigger one manually:
+```bash
+desloppify review --run-batches --runner codex --parallel --scan-after-import
+```
+
+### Other useful commands
+
+```bash
+desloppify next --count 5                         # top 5 priorities
+desloppify next --cluster <name>                  # drill into a cluster
+desloppify show <pattern>                         # filter by file/detector/ID
+desloppify show --status open                     # all open findings
+desloppify plan skip --permanent "<id>" --note "reason" # accept debt (lowers strict score)
+desloppify scan --path . --reset-subjective       # reset subjective baseline to 0
+```
+
+## 3. Reference
+
+### How scoring works
+
+Overall score = **40% mechanical** + **60% subjective**.
+
+- **Mechanical (40%)**: auto-detected issues — duplication, dead code, smells, unused imports, security. Fixed by changing code and rescanning.
+- **Subjective (60%)**: design quality review — naming, error handling, abstractions, clarity. Starts at **0%** until reviewed. The scan will prompt you when a review is needed.
+- **Strict score** is the north star: wontfix items count as open. The gap between overall and strict is your wontfix debt.
+- **Score types**: overall (lenient), strict (wontfix counts), objective (mechanical only), verified (confirmed fixes only).
+
+### Subjective reviews in detail
+
+- **Preferred**: `desloppify review --run-batches --runner codex --parallel --scan-after-import` — does everything in one command.
+- **Manual path**: `desloppify review --prepare` → review per dimension → `desloppify review --import file.json`.
+- Import first, fix after — import creates tracked state entries for correlation.
+- Integrity: reviewers score from evidence only. Scores hitting exact targets trigger auto-reset.
+- Even moderate scores (60-80) dramatically improve overall health.
+- Stale dimensions auto-surface in `next` — just follow the queue.
+
+### Key concepts
+
+- **Tiers**: T1 auto-fix → T2 quick manual → T3 judgment call → T4 major refactor.
+- **Auto-clusters**: related findings are auto-grouped in `next`. Drill in with `next --cluster <name>`.
+- **Zones**: production/script (scored), test/config/generated/vendor (not scored). Fix with `zone set`.
+- **Wontfix cost**: widens the lenient↔strict gap. Challenge past decisions when the gap grows.
+- Score can temporarily drop after fixes (cascade effects are normal).
+
+## 4. Escalate Tool Issues Upstream
+
+When desloppify itself appears wrong or inconsistent:
+
+1. Capture a minimal repro (`command`, `path`, `expected`, `actual`).
+2. Open a GitHub issue in `peteromallet/desloppify`.
+3. If you can fix it safely, open a PR linked to that issue.
+4. If unsure whether it is tool bug vs user workflow, issue first, PR second.
+
+## Prerequisite
+
+`command -v desloppify >/dev/null 2>&1 && echo "desloppify: installed" || echo "NOT INSTALLED — run: pip install --upgrade git+https://github.com/peteromallet/desloppify.git"`
+
+<!-- desloppify-end -->
+
+## Claude Code Overlay
+
+Use Claude subagents for subjective scoring work that should be context-isolated.
+
+### Subjective review
+
+1. **Preferred**: `desloppify review --run-batches --runner codex --parallel --scan-after-import` — does everything in one command.
+2. **Claude cloud path**: `desloppify review --external-start --external-runner claude` → use generated `claude_launch_prompt.md` + `review_result.template.json` → run printed `desloppify review --external-submit --session-id <id> --import <file>`.
+3. **Manual path**: split dimensions across N subagents (one message, multiple Task calls), merge outputs, then `desloppify review --import findings.json`.
+
+For the manual path:
+- Read `dimension_prompts` from `query.json` for dimension definitions.
+- Give each agent the codebase path, dimensions, and output format. Let agents decide what to read.
+- Each agent writes output to a separate file. Merge assessments (average overlaps) and findings.
+- Import first, fix after — import creates tracked state for correlation.
+
+### Subagent rules
+
+1. Prefer delegating review tasks to a project subagent in `.claude/agents/`.
+2. Set `context: fork` so prior chat context does not leak into scoring.
+3. For blind reviews, consume `.desloppify/review_packet_blind.json` instead of full `query.json`.
+4. Score from evidence only; do not anchor to target thresholds. When mixed, score lower.
+5. Return machine-readable JSON matching the format in the base skill doc. For `--external-submit`, include `session` from the generated template.
+6. `findings` MUST match `query.system_prompt` exactly. Use `"findings": []` when no defects found.
+7. Import is fail-closed: invalid findings abort unless `--allow-partial` is passed.
+
+<!-- desloppify-overlay: claude -->
+<!-- desloppify-end -->

README.md

@@ -157,9 +157,9 @@ Accelerate Ultimate SD Upscaler by distributing video tiles across multiple work
 
 Control your distributed cluster programmatically without opening the browser.
 
-* **Endpoint:** `POST /distributed/queue`
-* **Functionality:** Accepts a standard ComfyUI workflow JSON, automatically distributes it to available workers, and returns the execution ID.
-* **Documentation:** [See API Examples & Scripts](https://github.com/robertvoy/ComfyUI-Distributed/blob/main/docs/comfyui-distributed-api.md)
+* **Endpoint:** `POST /distributed/queue`
+* **Functionality:** Accepts a distributed queue payload (`prompt`, `client_id`, `enabled_worker_ids`, optional `workflow`/`delegate_master`/`trace_execution_id`), dispatches to healthy workers, and returns `{prompt_id, worker_count}`.
+* **Documentation:** [See API Examples & Scripts](https://github.com/robertvoy/ComfyUI-Distributed/blob/main/docs/comfyui-distributed-api.md)
 
 > **⚠️ Security Warning:** Do not expose your ComfyUI port to the public internet. If you need remote access, run ComfyUI behind a secure proxy (like Cloudflare or a VPN).
 

__init__.py

@@ -1,29 +1,11 @@
-# Import everything needed from the main module
-from .distributed import (
-    NODE_CLASS_MAPPINGS as DISTRIBUTED_CLASS_MAPPINGS, 
-    NODE_DISPLAY_NAME_MAPPINGS as DISTRIBUTED_DISPLAY_NAME_MAPPINGS
-)
+"""ComfyUI-Distributed package entrypoint."""
+from __future__ import annotations
 
-# Import utilities
-from .utils.config import ensure_config_exists, CONFIG_FILE
-from .utils.logging import debug_log
-
-# Import distributed upscale nodes
-from .nodes.distributed_upscale import (
-    NODE_CLASS_MAPPINGS as UPSCALE_CLASS_MAPPINGS,
-    NODE_DISPLAY_NAME_MAPPINGS as UPSCALE_DISPLAY_NAME_MAPPINGS
-)
+from .bootstrap.entrypoint import build_node_mappings, initialize_runtime
 
 WEB_DIRECTORY = "./web"
 
-ensure_config_exists()
+NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS = build_node_mappings()
+initialize_runtime()
 
-# Merge node mappings
-NODE_CLASS_MAPPINGS = {**DISTRIBUTED_CLASS_MAPPINGS, **UPSCALE_CLASS_MAPPINGS}
-NODE_DISPLAY_NAME_MAPPINGS = {**DISTRIBUTED_DISPLAY_NAME_MAPPINGS, **UPSCALE_DISPLAY_NAME_MAPPINGS}
-
-__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
-
-debug_log("Loaded Distributed nodes.")
-debug_log(f"Config file: {CONFIG_FILE}")
-debug_log(f"Available nodes: {list(NODE_CLASS_MAPPINGS.keys())}")
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]

api/__init__.py

@@ -1,5 +1,21 @@
-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
+"""Distributed API package with explicit route bootstrap."""
+from __future__ import annotations
+
+from importlib import import_module
+
+_ROUTE_MODULES = (
+    "config_routes",
+    "tunnel_routes",
+    "worker_routes",
+    "job_routes",
+    "usdu_routes",
+)
+
+
+def bootstrap_routes() -> None:
+    """Import route modules once so aiohttp decorators register endpoints."""
+    for module_name in _ROUTE_MODULES:
+        import_module(f"{__name__}.{module_name}")
+
+
+__all__ = ["bootstrap_routes"]

api/config_routes.py

@@ -1,31 +1,12 @@
 from __future__ import annotations
 
-import json
-from contextlib import asynccontextmanager
-
 from aiohttp import web
 import server
 
-try:
-    from ..utils.config import config_transaction, load_config, save_config
-except ImportError:
-    from ..utils.config import load_config
-
-    try:
-        from ..utils.config import save_config
-    except ImportError:
-        def save_config(_config):
-            return True
-
-    @asynccontextmanager
-    async def config_transaction():
-        config = load_config()
-        original_snapshot = json.dumps(config, sort_keys=True)
-        yield config
-        if json.dumps(config, sort_keys=True) != original_snapshot:
-            save_config(config)
-from ..utils.logging import debug_log, log
+from ..utils.config import config_transaction, load_config
+from ..utils.logging import debug_log
 from ..utils.network import handle_api_error, normalize_host
+from .endpoint_policy import run_authorized_endpoint
 
 
 def _positive_int(value: int) -> bool:
@@ -88,68 +69,103 @@ def _apply_field_patch(target: dict, data: dict, field_rules: list) -> None:
             target[key] = normalizer(value) if (normalizer and value is not None) else value
 
 
+def _validate_workers_payload(workers_value: object) -> list[str]:
+    """Validate top-level workers payload shape before persisting config."""
+    if not isinstance(workers_value, list):
+        return ["workers: expected list"]
+
+    errors: list[str] = []
+    seen_ids: set[str] = set()
+    for index, worker in enumerate(workers_value):
+        if not isinstance(worker, dict):
+            errors.append(f"workers[{index}]: expected object")
+            continue
+
+        worker_id = str(worker.get("id", "")).strip()
+        if not worker_id:
+            errors.append(f"workers[{index}].id: expected non-empty string")
+        elif worker_id in seen_ids:
+            errors.append(f"workers[{index}].id: duplicate id '{worker_id}'")
+        else:
+            seen_ids.add(worker_id)
+
+        if "port" in worker and worker.get("port") is not None:
+            try:
+                port = int(worker.get("port"))
+            except (TypeError, ValueError):
+                errors.append(f"workers[{index}].port: expected integer")
+            else:
+                if port <= 0:
+                    errors.append(f"workers[{index}].port: expected positive integer")
+
+    return errors
+
+
 @server.PromptServer.instance.routes.get("/distributed/config")
 async def get_config_endpoint(request: web.Request) -> web.StreamResponse:
-    config = load_config()
-    return web.json_response(config)
+    async def _operation() -> web.StreamResponse:
+        return web.json_response(load_config())
+
+    return await run_authorized_endpoint(request, _operation)
 
 
 @server.PromptServer.instance.routes.post("/distributed/config")
 async def update_config_endpoint(request: web.Request) -> web.StreamResponse:
     """Bulk config update with schema validation."""
-    try:
-        data = await request.json()
-    except Exception as e:
-        return await handle_api_error(request, f"Invalid JSON payload: {e}", 400)
+    async def _operation() -> web.StreamResponse:
+        try:
+            data = await request.json()
+        except Exception as exc:
+            return await handle_api_error(request, f"Invalid JSON payload: {exc}", 400)
 
-    if not isinstance(data, dict):
-        return await handle_api_error(request, "Config payload must be an object", 400)
+        if not isinstance(data, dict):
+            return await handle_api_error(request, "Config payload must be an object", 400)
 
-    validated_settings = {}
-    validated_root = {}
-    errors = []
+        validated_settings = {}
+        validated_root = {}
+        errors = []
 
-    for key, value in data.items():
-        if key not in CONFIG_SCHEMA:
-            errors.append(f"Unknown field: {key}")
-            continue
+        for key, value in data.items():
+            if key not in CONFIG_SCHEMA:
+                errors.append(f"Unknown field: {key}")
+                continue
 
-        expected_type, validator = CONFIG_SCHEMA[key]
-        if not isinstance(value, expected_type):
-            errors.append(f"{key}: expected {expected_type.__name__}")
-            continue
+            expected_type, validator = CONFIG_SCHEMA[key]
+            if not isinstance(value, expected_type):
+                errors.append(f"{key}: expected {expected_type.__name__}")
+                continue
 
-        if validator and not validator(value):
-            errors.append(f"{key}: value {value!r} failed validation")
-            continue
+            if validator and not validator(value):
+                errors.append(f"{key}: value {value!r} failed validation")
+                continue
 
-        if key in _SETTINGS_FIELDS:
-            validated_settings[key] = value
-        else:
-            validated_root[key] = value
+            if key == "workers":
+                errors.extend(_validate_workers_payload(value))
+                if errors:
+                    continue
 
-    if errors:
-        return web.json_response({
-            "status": "error",
-            "error": errors,
-            "message": "; ".join(errors),
-        }, status=400)
+            if key in _SETTINGS_FIELDS:
+                validated_settings[key] = value
+            else:
+                validated_root[key] = value
+
+        if errors:
+            return await handle_api_error(request, errors, 400)
 
-    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)
+
+    return await run_authorized_endpoint(request, _operation)
 
 
 @server.PromptServer.instance.routes.get("/distributed/queue_status/{job_id}")
 async def queue_status_endpoint(request: web.Request) -> web.StreamResponse:
     """Check if a job queue is initialized."""
-    try:
+    async def _operation() -> web.StreamResponse:
         job_id = request.match_info['job_id']
         
         # Import to ensure initialization
@@ -161,12 +177,12 @@ async def queue_status_endpoint(request: web.Request) -> web.StreamResponse:
         
         debug_log(f"Queue status check for job {job_id}: {'exists' if exists else 'not found'}")
         return web.json_response({"exists": exists, "job_id": job_id})
-    except Exception as e:
-        return await handle_api_error(request, e, 500)
+
+    return await run_authorized_endpoint(request, _operation, unexpected_status=500)
 
 @server.PromptServer.instance.routes.post("/distributed/config/update_worker")
 async def update_worker_endpoint(request: web.Request) -> web.StreamResponse:
-    try:
+    async def _operation() -> web.StreamResponse:
         data = await request.json()
         worker_id = data.get("worker_id")
         
@@ -205,12 +221,12 @@ async def update_worker_endpoint(request: web.Request) -> web.StreamResponse:
                     )
 
             return web.json_response({"status": "success"})
-    except Exception as e:
-        return await handle_api_error(request, e, 400)
+
+    return await run_authorized_endpoint(request, _operation, unexpected_status=500)
 
 @server.PromptServer.instance.routes.post("/distributed/config/delete_worker")
 async def delete_worker_endpoint(request: web.Request) -> web.StreamResponse:
-    try:
+    async def _operation() -> web.StreamResponse:
         data = await request.json()
         worker_id = data.get("worker_id")
         
@@ -237,13 +253,13 @@ async def delete_worker_endpoint(request: web.Request) -> web.StreamResponse:
                 "status": "success",
                 "message": f"Worker {removed_worker.get('name', worker_id)} deleted"
             })
-    except Exception as e:
-        return await handle_api_error(request, e, 400)
+
+    return await run_authorized_endpoint(request, _operation, unexpected_status=500)
 
 @server.PromptServer.instance.routes.post("/distributed/config/update_setting")
 async def update_setting_endpoint(request: web.Request) -> web.StreamResponse:
     """Updates a specific key in the settings object."""
-    try:
+    async def _operation() -> web.StreamResponse:
         data = await request.json()
         key = data.get("key")
         value = data.get("value")
@@ -260,13 +276,13 @@ async def update_setting_endpoint(request: web.Request) -> web.StreamResponse:
             config['settings'][key] = value
 
             return web.json_response({"status": "success", "message": f"Setting '{key}' updated."})
-    except Exception as e:
-        return await handle_api_error(request, e, 400)
+
+    return await run_authorized_endpoint(request, _operation, unexpected_status=500)
 
 @server.PromptServer.instance.routes.post("/distributed/config/update_master")
 async def update_master_endpoint(request: web.Request) -> web.StreamResponse:
     """Updates master configuration."""
-    try:
+    async def _operation() -> web.StreamResponse:
         data = await request.json()
         
         async with config_transaction() as config:
@@ -275,5 +291,5 @@ async def update_master_endpoint(request: web.Request) -> web.StreamResponse:
             _apply_field_patch(config['master'], data, _MASTER_FIELDS)
 
             return web.json_response({"status": "success", "message": "Master configuration updated."})
-    except Exception as e:
-        return await handle_api_error(request, e, 400)
+
+    return await run_authorized_endpoint(request, _operation, unexpected_status=500)

api/endpoint_policy.py

@@ -0,0 +1,25 @@
+"""Shared endpoint policy helpers for distributed API routes."""
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable
+
+from aiohttp import web
+
+from ..utils.network import handle_api_error
+from .request_guards import authorization_error_or_none
+
+
+async def run_authorized_endpoint(
+    request: web.Request,
+    operation: Callable[[], Awaitable[web.StreamResponse]],
+    *,
+    unexpected_status: int = 500,
+) -> web.StreamResponse:
+    """Run endpoint operation with shared auth guard and fallback error mapping."""
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
+    try:
+        return await operation()
+    except Exception as exc:
+        return await handle_api_error(request, exc, unexpected_status)

api/job_routes.py

@@ -18,25 +18,20 @@ from ..utils.logging import debug_log
 from ..utils.image import pil_to_tensor, ensure_contiguous
 from ..utils.network import handle_api_error
 from ..utils.constants import JOB_INIT_GRACE_PERIOD, MEMORY_CLEAR_DELAY
-try:
-    from .queue_orchestration import ensure_distributed_state, orchestrate_distributed_execution
-except ImportError:
-    from .queue_orchestration import orchestrate_distributed_execution
-
-    def ensure_distributed_state() -> None:
-        prompt_server_instance = server.PromptServer.instance
-        if not hasattr(prompt_server_instance, "distributed_jobs_lock"):
-            prompt_server_instance.distributed_jobs_lock = asyncio.Lock()
-        if not hasattr(prompt_server_instance, "distributed_pending_jobs"):
-            prompt_server_instance.distributed_pending_jobs = {}
+from ..utils.runtime_state import ensure_distributed_runtime_state
+from .request_guards import authorization_error_or_none
+from .schemas import require_bool_literal
+from .queue_orchestration import orchestrate_distributed_execution
 from .queue_request import parse_queue_request_payload
 
-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 }
 
 
+def _runtime_state():
+    return ensure_distributed_runtime_state()
+
+
 def _decode_image_sync(image_path: str) -> dict[str, Any]:
     """Decode image/video file and compute hash in a threadpool worker."""
     import base64
@@ -148,16 +143,19 @@ def _decode_audio_payload(audio_payload: dict[str, Any]) -> dict[str, Any]:
 
 @server.PromptServer.instance.routes.post("/distributed/prepare_job")
 async def prepare_job_endpoint(request: web.Request) -> web.StreamResponse:
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
     try:
         data = await request.json()
         multi_job_id = data.get('multi_job_id')
         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()
+        runtime_state = _runtime_state()
+        async with runtime_state.distributed_jobs_lock:
+            if multi_job_id not in runtime_state.distributed_pending_jobs:
+                runtime_state.distributed_pending_jobs[multi_job_id] = asyncio.Queue()
         
         debug_log(f"Prepared queue for job {multi_job_id}")
         return web.json_response({"status": "success"})
@@ -167,7 +165,11 @@ async def prepare_job_endpoint(request: web.Request) -> web.StreamResponse:
 @server.PromptServer.instance.routes.post("/distributed/clear_memory")
 async def clear_memory_endpoint(request: web.Request) -> web.StreamResponse:
     debug_log("Received request to clear VRAM.")
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
     try:
+        warnings: list[str] = []
         # Use ComfyUI's prompt server queue system like the /free endpoint does
         if hasattr(server.PromptServer.instance, 'prompt_queue'):
             server.PromptServer.instance.prompt_queue.set_flag("unload_models", True)
@@ -184,12 +186,16 @@ async def clear_memory_endpoint(request: web.Request) -> web.StreamResponse:
         try:
             mm.unload_all_models()
         except AttributeError as e:
-            debug_log(f"Warning during model unload: {e}")
+            warning = f"Model unload warning: {e}"
+            warnings.append(warning)
+            debug_log(warning)
         
         try:
             mm.soft_empty_cache()
         except Exception as e:
-            debug_log(f"Warning during cache clear: {e}")
+            warning = f"Cache clear warning: {e}"
+            warnings.append(warning)
+            debug_log(warning)
         
         for _ in range(3):
             gc.collect()
@@ -198,6 +204,16 @@ async def clear_memory_endpoint(request: web.Request) -> web.StreamResponse:
             torch.cuda.empty_cache()
             torch.cuda.ipc_collect()
         
+        if warnings:
+            debug_log("VRAM cleared with warnings.")
+            return web.json_response(
+                {
+                    "status": "partial",
+                    "message": "GPU memory cleared with warnings.",
+                    "warnings": warnings,
+                }
+            )
+
         debug_log("VRAM cleared successfully.")
         return web.json_response({"status": "success", "message": "GPU memory cleared."})
     except Exception as e:
@@ -206,13 +222,16 @@ async def clear_memory_endpoint(request: web.Request) -> web.StreamResponse:
         gc.collect()
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
-        debug_log(f"Partial VRAM clear completed with warning: {e}")
-        return web.json_response({"status": "success", "message": "GPU memory cleared (with warnings)"})
+        debug_log(f"VRAM clear failed: {e}")
+        return await handle_api_error(request, f"GPU memory clear failed: {e}", 500)
 
 
 @server.PromptServer.instance.routes.post("/distributed/queue")
 async def distributed_queue_endpoint(request: web.Request) -> web.StreamResponse:
     """Queue a distributed workflow, mirroring the UI orchestration pipeline."""
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
     try:
         raw_payload = await request.json()
     except Exception as exc:
@@ -235,7 +254,6 @@ async def distributed_queue_endpoint(request: web.Request) -> web.StreamResponse
         return web.json_response({
             "prompt_id": prompt_id,
             "worker_count": worker_count,
-            "auto_prepare_supported": True,
         })
     except Exception as exc:
         return await handle_api_error(request, exc, 500)
@@ -243,6 +261,9 @@ async def distributed_queue_endpoint(request: web.Request) -> web.StreamResponse
 @server.PromptServer.instance.routes.post("/distributed/load_image")
 async def load_image_endpoint(request: web.Request) -> web.StreamResponse:
     """Load an image or video file and return it as base64 data with hash."""
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
     try:
         data = await request.json()
         image_path = data.get("image_path")
@@ -260,6 +281,9 @@ async def load_image_endpoint(request: web.Request) -> web.StreamResponse:
 @server.PromptServer.instance.routes.post("/distributed/check_file")
 async def check_file_endpoint(request: web.Request) -> web.StreamResponse:
     """Check if a file exists and matches the given hash."""
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
     try:
         data = await request.json()
         filename = data.get("filename")
@@ -277,6 +301,9 @@ async def check_file_endpoint(request: web.Request) -> web.StreamResponse:
 
 @server.PromptServer.instance.routes.post("/distributed/job_complete")
 async def job_complete_endpoint(request: web.Request) -> web.StreamResponse:
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
     try:
         data = await request.json()
     except Exception as exc:
@@ -291,7 +318,7 @@ async def job_complete_endpoint(request: web.Request) -> web.StreamResponse:
         batch_idx = data.get("batch_idx")
         image_payload = data.get("image")
         audio_payload = data.get("audio")
-        is_last = data.get("is_last")
+        is_last_raw = data.get("is_last")
 
         errors = []
         if not isinstance(job_id, str) or not job_id.strip():
@@ -304,8 +331,12 @@ async def job_complete_endpoint(request: web.Request) -> web.StreamResponse:
             errors.append("image: expected non-empty base64 PNG string")
         if audio_payload is not None and not isinstance(audio_payload, dict):
             errors.append("audio: expected object when provided")
-        if not isinstance(is_last, bool):
-            errors.append("is_last: expected boolean")
+
+        try:
+            is_last = require_bool_literal(is_last_raw, field_name="is_last")
+        except ValueError as exc:
+            errors.append(str(exc))
+            is_last = False
         if errors:
             return await handle_api_error(request, errors, 400)
 
@@ -314,12 +345,21 @@ async def job_complete_endpoint(request: web.Request) -> web.StreamResponse:
         multi_job_id = job_id.strip()
         worker_id = worker_id.strip()
 
+        runtime_state = _runtime_state()
+        allowed_workers = runtime_state.distributed_job_allowed_workers.get(multi_job_id)
+        if allowed_workers is not None and worker_id not in allowed_workers:
+            return await handle_api_error(
+                request,
+                f"Unauthorized worker_id for job {multi_job_id}",
+                403,
+            )
+
         pending = None
         queue_size = 0
         deadline = time.monotonic() + float(JOB_INIT_GRACE_PERIOD)
         while pending is None:
-            async with prompt_server.distributed_jobs_lock:
-                pending = prompt_server.distributed_pending_jobs.get(multi_job_id)
+            async with runtime_state.distributed_jobs_lock:
+                pending = runtime_state.distributed_pending_jobs.get(multi_job_id)
                 if pending is not None:
                     queue_item = {
                         "tensor": tensor,

api/orchestration/__init__.py

@@ -1,4 +1,7 @@
-# Orchestration helpers split out from queue_orchestration.py:
-# - prompt_transform.py: graph pruning + hidden input overrides
-# - media_sync.py: remote media/path normalization
-# - dispatch.py: worker probe + prompt dispatch
+"""Orchestration helpers used by distributed queue execution."""
+
+__all__ = [
+    "dispatch",
+    "media_sync",
+    "prompt_transform",
+]

api/orchestration/dispatch.py

@@ -5,26 +5,12 @@ from typing import Any
 
 import aiohttp
 
+from ...utils.auth import distributed_auth_headers
+from ...utils.config import load_config
 from ...utils.logging import debug_log, log
 from ...utils.network import build_worker_url, get_client_session, probe_worker
-try:
-    from ...utils.trace_logger import trace_debug, trace_info
-except ImportError:
-    def _trace_fallback(_execution_id: str, message: str) -> None:
-        debug_log(message)
-
-    trace_debug = _trace_fallback
-    trace_info = _trace_fallback
-
-try:
-    from ..schemas import parse_positive_int
-except ImportError:
-    def parse_positive_int(value, default):
-        try:
-            parsed = int(value)
-            return parsed if parsed > 0 else default
-        except (TypeError, ValueError):
-            return default
+from ...utils.trace_logger import trace_debug, trace_info
+from ..schemas import coerce_positive_int
 
 _least_busy_rr_index = 0
 
@@ -39,8 +25,9 @@ async def worker_ws_is_active(worker: dict[str, Any]) -> bool:
     """Ping worker's websocket endpoint to confirm it's reachable."""
     session = await get_client_session()
     url = build_worker_url(worker, "/distributed/worker_ws")
+    headers = distributed_auth_headers(load_config())
     try:
-        ws = await session.ws_connect(url, heartbeat=20, timeout=3)
+        ws = await session.ws_connect(url, heartbeat=20, timeout=3, headers=headers)
         await ws.close()
         return True
     except asyncio.TimeoutError:
@@ -60,7 +47,13 @@ async def _probe_worker_active(worker, use_websocket, semaphore):
         return worker, is_active
 
 
-async def _dispatch_via_websocket(worker_url, payload, client_id, timeout=60.0):
+async def _dispatch_via_websocket(
+    worker_url,
+    payload,
+    client_id,
+    headers: dict[str, str],
+    timeout=60.0,
+):
     """Open a fresh worker websocket, dispatch one prompt, wait for ack, then close."""
     request_id = uuid.uuid4().hex
     ws_payload = {
@@ -74,7 +67,12 @@ async def _dispatch_via_websocket(worker_url, payload, client_id, timeout=60.0):
     ws_url = f"{ws_url}/distributed/worker_ws"
     session = await get_client_session()
 
-    async with session.ws_connect(ws_url, heartbeat=20, timeout=timeout) as ws:
+    async with session.ws_connect(
+        ws_url,
+        heartbeat=20,
+        timeout=timeout,
+        headers=headers,
+    ) as ws:
         await ws.send_json(ws_payload)
         async for msg in ws:
             if msg.type == aiohttp.WSMsgType.TEXT:
@@ -115,6 +113,7 @@ async def dispatch_worker_prompt(
         payload["extra_data"] = extra_data
 
     if use_websocket:
+        ws_headers = distributed_auth_headers(load_config())
         try:
             await _dispatch_via_websocket(
                 worker_url,
@@ -123,6 +122,7 @@ async def dispatch_worker_prompt(
                     "workflow": workflow_meta,
                 },
                 client_id,
+                ws_headers,
             )
             return
         except Exception as exc:
@@ -150,7 +150,7 @@ async def select_active_workers(
     probe_concurrency: int = 8,
 ) -> tuple[list[dict[str, Any]], bool]:
     """Probe workers and return (active_workers, updated_delegate_master)."""
-    probe_limit = parse_positive_int(probe_concurrency, 8)
+    probe_limit = coerce_positive_int(probe_concurrency, 8)
     probe_semaphore = asyncio.Semaphore(probe_limit)
 
     if trace_execution_id and workers:
@@ -233,7 +233,7 @@ async def select_least_busy_worker(
     if not workers:
         return None
 
-    probe_limit = parse_positive_int(probe_concurrency, 8)
+    probe_limit = coerce_positive_int(probe_concurrency, 8)
     probe_semaphore = asyncio.Semaphore(probe_limit)
     statuses = await asyncio.gather(
         *[
@@ -279,7 +279,7 @@ async def rank_workers_by_load(
     if not workers:
         return []
 
-    probe_limit = parse_positive_int(probe_concurrency, 8)
+    probe_limit = coerce_positive_int(probe_concurrency, 8)
     probe_semaphore = asyncio.Semaphore(probe_limit)
     statuses = await asyncio.gather(
         *[

api/orchestration/media_sync.py

@@ -7,6 +7,8 @@ from typing import Any
 
 import aiohttp
 
+from ...utils.auth import distributed_auth_headers
+from ...utils.config import load_config
 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
@@ -132,8 +134,9 @@ async def fetch_worker_path_separator(
     """Best-effort fetch of a worker's path separator from /distributed/system_info."""
     url = build_worker_url(worker, "/distributed/system_info")
     session = await get_client_session()
+    headers = distributed_auth_headers(load_config())
     try:
-        async with session.get(url, timeout=aiohttp.ClientTimeout(total=5)) as resp:
+        async with session.get(url, headers=headers, timeout=aiohttp.ClientTimeout(total=5)) as resp:
             if resp.status != 200:
                 return None
             payload = await resp.json()
@@ -150,6 +153,7 @@ async def fetch_worker_path_separator(
 async def _upload_media_to_worker(worker, filename, file_bytes, file_hash, mime_type, trace_execution_id=None):
     """Upload one media file to worker iff missing or hash-mismatched."""
     session = await get_client_session()
+    headers = distributed_auth_headers(load_config())
     normalized = filename.replace("\\", "/")
 
     check_url = build_worker_url(worker, "/distributed/check_file")
@@ -157,6 +161,7 @@ async def _upload_media_to_worker(worker, filename, file_bytes, file_hash, mime_
         async with session.post(
             check_url,
             json={"filename": normalized, "hash": file_hash},
+            headers=headers,
             timeout=aiohttp.ClientTimeout(total=6),
         ) as resp:
             if resp.status == 200:

api/orchestration/prompt_transform.py

@@ -332,11 +332,10 @@ def prune_prompt_for_branch_worker(
 def prune_prompt_for_worker(prompt_obj: dict[str, Any]) -> dict[str, Any]:
     """Prune worker prompt to distributed nodes and their upstream dependencies."""
     collector_ids = find_nodes_by_class(prompt_obj, "DistributedCollector")
-    list_collector_ids = find_nodes_by_class(prompt_obj, "DistributedListCollector")
     branch_collector_ids = find_nodes_by_class(prompt_obj, "DistributedBranchCollector")
     upscale_ids = find_nodes_by_class(prompt_obj, "UltimateSDUpscaleDistributed")
     branch_ids = find_nodes_by_class(prompt_obj, "DistributedBranch")
-    distributed_ids = collector_ids + list_collector_ids + branch_collector_ids + upscale_ids + branch_ids
+    distributed_ids = collector_ids + branch_collector_ids + upscale_ids + branch_ids
     if not distributed_ids:
         return prompt_obj
 
@@ -439,8 +438,6 @@ def generate_job_id_map(prompt_index: PromptIndex, prefix: str) -> dict[str, str
     job_map = {}
     distributed_nodes = (
         prompt_index.nodes_for_class("DistributedCollector")
-        + prompt_index.nodes_for_class("DistributedListSplitter")
-        + prompt_index.nodes_for_class("DistributedListCollector")
         + prompt_index.nodes_for_class("DistributedBranchCollector")
         + prompt_index.nodes_for_class("UltimateSDUpscaleDistributed")
         + prompt_index.nodes_for_class("DistributedBranch")
@@ -450,7 +447,7 @@ def generate_job_id_map(prompt_index: PromptIndex, prefix: str) -> dict[str, str
     return job_map
 
 
-def _override_seed_nodes(prompt_copy, prompt_index, is_master, participant_id, worker_index_map):
+def _override_seed_nodes(prompt_copy, prompt_index, is_master, participant_id):
     """Configure DistributedSeed nodes for master or worker role."""
     for node_id in prompt_index.nodes_for_class("DistributedSeed"):
         node = prompt_copy.get(node_id)
@@ -461,7 +458,7 @@ def _override_seed_nodes(prompt_copy, prompt_index, is_master, participant_id, w
         if is_master:
             inputs["worker_id"] = ""
         else:
-            inputs["worker_id"] = f"worker_{worker_index_map.get(participant_id, 0)}"
+            inputs["worker_id"] = str(participant_id)
 
 
 def _override_collector_nodes(
@@ -524,7 +521,7 @@ def _override_upscale_nodes(
             inputs["worker_id"] = participant_id
 
 
-def _override_value_nodes(prompt_copy, prompt_index, is_master, participant_id, worker_index_map):
+def _override_value_nodes(prompt_copy, prompt_index, is_master, participant_id):
     """Configure DistributedValue nodes for master or worker role."""
     for node_id in prompt_index.nodes_for_class("DistributedValue"):
         node = prompt_copy.get(node_id)
@@ -535,74 +532,7 @@ def _override_value_nodes(prompt_copy, prompt_index, is_master, participant_id,
         if is_master:
             inputs["worker_id"] = ""
         else:
-            inputs["worker_id"] = f"worker_{worker_index_map.get(participant_id, 0)}"
-
-
-def _override_list_splitter_nodes(
-    prompt_copy,
-    prompt_index,
-    is_master,
-    participant_id,
-    enabled_worker_ids,
-    job_id_map,
-    master_url,
-    delegate_master,
-):
-    """Configure participant chunk mapping for DistributedListSplitter nodes."""
-    participants = _resolve_participants(enabled_worker_ids, delegate_master)
-    total_participants = max(1, len(participants))
-    participant_id = str(participant_id)
-
-    if participant_id in participants:
-        participant_index = participants.index(participant_id)
-    else:
-        participant_index = 0
-
-    for node_id in prompt_index.nodes_for_class("DistributedListSplitter"):
-        node = prompt_copy.get(node_id)
-        if not isinstance(node, dict):
-            continue
-        inputs = node.setdefault("inputs", {})
-        inputs["participant_index"] = participant_index
-        inputs["total_participants"] = total_participants
-        inputs["multi_job_id"] = job_id_map.get(node_id, node_id)
-        inputs["is_worker"] = not is_master
-        if is_master:
-            inputs.pop("master_url", None)
-            inputs["worker_id"] = "master"
-        else:
-            inputs["master_url"] = master_url
-            inputs["worker_id"] = participant_id
-
-
-def _override_list_collector_nodes(
-    prompt_copy,
-    prompt_index,
-    is_master,
-    participant_id,
-    job_id_map,
-    master_url,
-    enabled_json,
-    delegate_master,
-):
-    """Configure DistributedListCollector nodes for master/worker role."""
-    for node_id in prompt_index.nodes_for_class("DistributedListCollector"):
-        node = prompt_copy.get(node_id)
-        if not isinstance(node, dict):
-            continue
-
-        inputs = node.setdefault("inputs", {})
-        inputs["multi_job_id"] = job_id_map.get(node_id, node_id)
-        inputs["is_worker"] = not is_master
-        inputs["enabled_worker_ids"] = enabled_json
-        if is_master:
-            inputs["delegate_only"] = bool(delegate_master)
-            inputs.pop("master_url", None)
-            inputs.pop("worker_id", None)
-        else:
-            inputs["master_url"] = master_url
             inputs["worker_id"] = str(participant_id)
-            inputs["delegate_only"] = False
 
 
 def _override_branch_nodes(
@@ -740,21 +670,10 @@ def apply_participant_overrides(
 ) -> dict[str, Any]:
     """Return a prompt copy with hidden inputs configured for master/worker."""
     is_master = participant_id == "master"
-    worker_index_map = {wid: idx for idx, wid in enumerate(enabled_worker_ids)}
     enabled_json = json.dumps(enabled_worker_ids)
 
-    _override_seed_nodes(prompt_copy, prompt_index, is_master, participant_id, worker_index_map)
-    _override_value_nodes(prompt_copy, prompt_index, is_master, participant_id, worker_index_map)
-    _override_list_splitter_nodes(
-        prompt_copy,
-        prompt_index,
-        is_master,
-        participant_id,
-        enabled_worker_ids,
-        job_id_map,
-        master_url,
-        delegate_master,
-    )
+    _override_seed_nodes(prompt_copy, prompt_index, is_master, participant_id)
+    _override_value_nodes(prompt_copy, prompt_index, is_master, participant_id)
     _override_collector_nodes(
         prompt_copy,
         prompt_index,
@@ -765,16 +684,6 @@ def apply_participant_overrides(
         enabled_json,
         delegate_master,
     )
-    _override_list_collector_nodes(
-        prompt_copy,
-        prompt_index,
-        is_master,
-        participant_id,
-        job_id_map,
-        master_url,
-        enabled_json,
-        delegate_master,
-    )
     _override_upscale_nodes(
         prompt_copy,
         prompt_index,

api/queue_orchestration.py

@@ -3,8 +3,6 @@ import time
 import uuid
 from typing import Any
 
-import server
-
 from ..utils.async_helpers import queue_prompt_payload
 from ..utils.config import load_config
 from ..utils.constants import (
@@ -15,8 +13,9 @@ from ..utils.constants import (
 )
 from ..utils.logging import debug_log, log
 from ..utils.network import build_master_url
+from ..utils.runtime_state import ensure_distributed_runtime_state, get_prompt_server_instance
 from ..utils.trace_logger import trace_debug
-from .schemas import parse_positive_float, parse_positive_int
+from .schemas import coerce_positive_float, coerce_positive_int
 from .orchestration.dispatch import (
     dispatch_worker_prompt,
     rank_workers_by_load,
@@ -33,33 +32,21 @@ from .orchestration.prompt_transform import (
     prune_prompt_for_worker,
 )
 
-
-prompt_server = server.PromptServer.instance
-
-
 def _generate_execution_trace_id():
     return f"exec_{int(time.time() * 1000)}_{uuid.uuid4().hex[:6]}"
 
 
 def ensure_distributed_state(server_instance=None):
     """Ensure prompt_server has the state used by distributed queue orchestration."""
-    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()
+    ensure_distributed_runtime_state(server_instance)
 
 
 async def _ensure_distributed_queue(job_id):
     """Ensure a queue exists for the given distributed job ID."""
-    ensure_distributed_state()
-    async with prompt_server.distributed_jobs_lock:
-        if job_id not in prompt_server.distributed_pending_jobs:
-            prompt_server.distributed_pending_jobs[job_id] = asyncio.Queue()
+    runtime_state = ensure_distributed_runtime_state()
+    async with runtime_state.distributed_jobs_lock:
+        if job_id not in runtime_state.distributed_pending_jobs:
+            runtime_state.distributed_pending_jobs[job_id] = asyncio.Queue()
 
 
 def _resolve_enabled_workers(config, requested_ids=None):
@@ -98,19 +85,19 @@ def _resolve_enabled_workers(config, requested_ids=None):
 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 = coerce_positive_int(
         settings.get("worker_probe_concurrency"),
         ORCHESTRATION_WORKER_PROBE_CONCURRENCY,
     )
-    worker_prep_concurrency = parse_positive_int(
+    worker_prep_concurrency = coerce_positive_int(
         settings.get("worker_prep_concurrency"),
         ORCHESTRATION_WORKER_PREP_CONCURRENCY,
     )
-    media_sync_concurrency = parse_positive_int(
+    media_sync_concurrency = coerce_positive_int(
         settings.get("media_sync_concurrency"),
         ORCHESTRATION_MEDIA_SYNC_CONCURRENCY,
     )
-    media_sync_timeout_seconds = parse_positive_float(
+    media_sync_timeout_seconds = coerce_positive_float(
         settings.get("media_sync_timeout_seconds"),
         ORCHESTRATION_MEDIA_SYNC_TIMEOUT,
     )
@@ -272,13 +259,20 @@ async def _select_execution_workers(
 def _distributed_queue_nodes(prompt_index):
     return (
         prompt_index.nodes_for_class("DistributedCollector")
-        + prompt_index.nodes_for_class("DistributedListCollector")
         + prompt_index.nodes_for_class("DistributedBranchCollector")
         + prompt_index.nodes_for_class("DistributedBranch")
         + prompt_index.nodes_for_class("UltimateSDUpscaleDistributed")
     )
 
 
+def _register_job_allowed_workers(job_id_map, enabled_ids):
+    runtime_state = ensure_distributed_runtime_state()
+    allowed = {str(worker_id) for worker_id in (enabled_ids or []) if str(worker_id).strip()}
+    for job_id in job_id_map.values():
+        if job_id:
+            runtime_state.distributed_job_allowed_workers[str(job_id)] = set(allowed)
+
+
 async def _ensure_job_queues(prompt_index, job_id_map):
     for node_id in _distributed_queue_nodes(prompt_index):
         job_id = job_id_map.get(node_id)
@@ -309,7 +303,6 @@ def _build_master_prompt(
 
     collector_ids = (
         find_nodes_by_class(master_prompt, "DistributedCollector")
-        + find_nodes_by_class(master_prompt, "DistributedListCollector")
         + find_nodes_by_class(master_prompt, "DistributedBranchCollector")
     )
     upscale_nodes = find_nodes_by_class(master_prompt, "UltimateSDUpscaleDistributed")
@@ -405,7 +398,7 @@ async def orchestrate_distributed_execution(
 
     config = load_config()
     use_websocket = bool(config.get("settings", {}).get("websocket_orchestration", False))
-    master_url = build_master_url(config=config, prompt_server_instance=prompt_server)
+    master_url = build_master_url(config=config, prompt_server_instance=get_prompt_server_instance())
     (
         worker_probe_concurrency,
         worker_prep_concurrency,
@@ -463,6 +456,7 @@ async def orchestrate_distributed_execution(
         prompt_id = await queue_prompt_payload(prompt_obj, workflow_meta, client_id)
         return prompt_id, 0
 
+    _register_job_allowed_workers(job_id_map, enabled_ids)
     await _ensure_job_queues(prompt_index, job_id_map)
 
     master_prompt = _build_master_prompt(

api/queue_request.py

@@ -1,16 +1,20 @@
 from dataclasses import dataclass
-from typing import Any, Dict, List, Optional
+from typing import Any
+
+try:
+    from ..utils.worker_ids import require_enabled_worker_ids
+except ImportError:  # pragma: no cover - supports direct module loading in isolated tests.
+    from utils.worker_ids import require_enabled_worker_ids
 
 
 @dataclass(frozen=True)
 class QueueRequestPayload:
-    prompt: Dict[str, Any]
-    workflow_meta: Any
+    prompt: dict[str, Any]
+    workflow_meta: dict[str, Any] | None
     client_id: str
-    delegate_master: Optional[bool]
-    enabled_worker_ids: List[str]
-    auto_prepare: bool
-    trace_execution_id: Optional[str]
+    delegate_master: bool | None
+    enabled_worker_ids: list[str]
+    trace_execution_id: str | None
 
 
 def parse_queue_request_payload(data: Any) -> QueueRequestPayload:
@@ -18,11 +22,6 @@ def parse_queue_request_payload(data: Any) -> QueueRequestPayload:
     if not isinstance(data, dict):
         raise ValueError("Expected a JSON object body")
 
-    auto_prepare_raw = data.get("auto_prepare", True)
-    if not isinstance(auto_prepare_raw, bool):
-        raise ValueError("auto_prepare must be a boolean when provided")
-    auto_prepare = auto_prepare_raw
-
     prompt = data.get("prompt")
     # Auto-prepare is always on server-side; keep the field for wire compatibility.
     if prompt is None:
@@ -35,6 +34,10 @@ def parse_queue_request_payload(data: Any) -> QueueRequestPayload:
     if not isinstance(prompt, dict):
         raise ValueError("Field 'prompt' must be an object")
 
+    workflow_meta = data.get("workflow")
+    if workflow_meta is not None and not isinstance(workflow_meta, dict):
+        raise ValueError("Field 'workflow' must be an object when provided")
+
     enabled_ids_raw = data.get("enabled_worker_ids")
     workers_field = data.get("workers")
     if enabled_ids_raw is None and workers_field is not None:
@@ -51,7 +54,10 @@ def parse_queue_request_payload(data: Any) -> QueueRequestPayload:
     else:
         if not isinstance(enabled_ids_raw, list):
             raise ValueError("enabled_worker_ids must be a list of worker IDs")
-        enabled_ids = [str(worker_id).strip() for worker_id in enabled_ids_raw if str(worker_id).strip()]
+        try:
+            enabled_ids = require_enabled_worker_ids(enabled_ids_raw)
+        except ValueError as exc:
+            raise ValueError(f"enabled_worker_ids invalid: {exc}") from exc
 
     delegate_master = data.get("delegate_master")
     if delegate_master is not None and not isinstance(delegate_master, bool):
@@ -70,10 +76,9 @@ def parse_queue_request_payload(data: Any) -> QueueRequestPayload:
 
     return QueueRequestPayload(
         prompt=prompt,
-        workflow_meta=data.get("workflow"),
+        workflow_meta=workflow_meta,
         client_id=client_id,
         delegate_master=delegate_master,
         enabled_worker_ids=enabled_ids,
-        auto_prepare=auto_prepare,
         trace_execution_id=trace_execution_id,
     )

api/request_guards.py

@@ -0,0 +1,15 @@
+"""Shared request guards for distributed API endpoints."""
+from __future__ import annotations
+
+from aiohttp import web
+
+from ..utils.config import load_config
+from ..utils.network import handle_api_error
+from .schemas import is_authorized_request
+
+
+async def authorization_error_or_none(request: web.Request) -> web.StreamResponse | None:
+    """Return a 403 response when distributed API auth fails, else None."""
+    if is_authorized_request(request, load_config()):
+        return None
+    return await handle_api_error(request, "Unauthorized", 403)

api/schemas.py

@@ -1,3 +1,29 @@
+from __future__ import annotations
+
+from typing import Any
+
+from ..utils.auth import (
+    AUTH_HEADER_NAME,
+    distributed_auth_headers,
+    is_authorized_request,
+)
+from ..utils.parsing import coerce_bool, require_bool_literal
+
+__all__ = [
+    "AUTH_HEADER_NAME",
+    "coerce_bool",
+    "coerce_positive_float",
+    "coerce_positive_int",
+    "distributed_auth_headers",
+    "is_authorized_request",
+    "require_bool_literal",
+    "require_fields",
+    "require_worker_id",
+    "require_positive_int",
+    "validate_worker_id",
+]
+
+
 def require_fields(data: dict, *fields) -> list[str]:
     """Return field names that are missing or empty in a JSON object."""
     if not isinstance(data, dict):
@@ -18,26 +44,43 @@ def require_fields(data: dict, *fields) -> list[str]:
     return missing
 
 
-def validate_worker_id(worker_id: str, config: dict) -> bool:
+def validate_worker_id(worker_id: str, config: dict[str, Any]) -> bool:
     """Return True when worker_id exists in config['workers']."""
-    worker_id_str = str(worker_id)
+    try:
+        require_worker_id(worker_id, config)
+        return True
+    except ValueError:
+        return False
+
+
+def require_worker_id(worker_id: str, config: dict[str, Any], field_name: str = "worker_id") -> str:
+    """Strict worker-id validator for request boundaries."""
+    worker_id_str = str(worker_id).strip()
+    if not worker_id_str:
+        raise ValueError(f"Field '{field_name}' must be a non-empty worker id.")
+
     workers = (config or {}).get("workers", [])
-    return any(str(worker.get("id")) == worker_id_str for worker in workers)
+    for worker in workers:
+        if not isinstance(worker, dict):
+            continue
+        if str(worker.get("id")).strip() == worker_id_str:
+            return worker_id_str
+    raise ValueError(f"Worker {worker_id_str} not found.")
 
 
-def validate_positive_int(value, field_name: str) -> str | None:
-    """Validate positive integers and return an error string when invalid."""
+def require_positive_int(value, field_name: str = "value") -> int:
+    """Strict positive-int parser for request boundary validation."""
     try:
         parsed = int(value)
-    except (TypeError, ValueError):
-        return f"Field '{field_name}' must be a positive integer."
+    except (TypeError, ValueError) as exc:
+        raise ValueError(f"Field '{field_name}' must be a positive integer.") from exc
     if parsed <= 0:
-        return f"Field '{field_name}' must be a positive integer."
-    return None
+        raise ValueError(f"Field '{field_name}' must be a positive integer.")
+    return parsed
 
 
-def parse_positive_int(value, default: int) -> int:
-    """Parse value as positive int, returning default on failure."""
+def coerce_positive_int(value, default: int) -> int:
+    """Permissive positive-int coercion used for non-boundary defaults."""
     try:
         parsed = int(value)
     except (TypeError, ValueError):
@@ -45,10 +88,13 @@ def parse_positive_int(value, default: int) -> int:
     return max(1, parsed)
 
 
-def parse_positive_float(value, default: float) -> float:
-    """Parse value as positive float, returning default on failure."""
+def coerce_positive_float(value, default: float) -> float:
+    """Permissive positive-float coercion used for non-boundary defaults."""
+    fallback = float(default)
+    if fallback <= 0.0:
+        fallback = 0.1
     try:
         parsed = float(value)
     except (TypeError, ValueError):
-        return max(0.0, float(default))
-    return max(0.0, parsed)
+        return fallback
+    return parsed if parsed > 0.0 else fallback

api/tunnel_routes.py

@@ -3,51 +3,61 @@ from __future__ import annotations
 from aiohttp import web
 import server
 
+from .endpoint_policy import run_authorized_endpoint
 from ..utils.cloudflare import cloudflare_tunnel_manager
 from ..utils.config import load_config
-from ..utils.logging import debug_log, log
-from ..utils.network import handle_api_error
 
 
 @server.PromptServer.instance.routes.get("/distributed/tunnel/status")
 async def tunnel_status_endpoint(request: web.Request) -> web.StreamResponse:
     """Return Cloudflare tunnel status and last known details."""
-    try:
+
+    async def _operation() -> web.StreamResponse:
         status = cloudflare_tunnel_manager.get_status()
         config = load_config()
         master_host = (config.get("master") or {}).get("host")
-        return web.json_response({
-            "status": "success",
-            "tunnel": status,
-            "master_host": master_host
-        })
-    except Exception as e:
-        return await handle_api_error(request, e, 500)
+        return web.json_response(
+            {
+                "status": "success",
+                "tunnel": status,
+                "master_host": master_host,
+            }
+        )
+
+    return await run_authorized_endpoint(request, _operation)
+
 
 @server.PromptServer.instance.routes.post("/distributed/tunnel/start")
 async def tunnel_start_endpoint(request: web.Request) -> web.StreamResponse:
     """Start a Cloudflare tunnel pointing at the current ComfyUI server."""
-    try:
+
+    async def _operation() -> web.StreamResponse:
         result = await cloudflare_tunnel_manager.start_tunnel()
         config = load_config()
-        return web.json_response({
-            "status": "success",
-            "tunnel": result,
-            "master_host": (config.get("master") or {}).get("host")
-        })
-    except Exception as e:
-        return await handle_api_error(request, e, 500)
+        return web.json_response(
+            {
+                "status": "success",
+                "tunnel": result,
+                "master_host": (config.get("master") or {}).get("host"),
+            }
+        )
+
+    return await run_authorized_endpoint(request, _operation)
+
 
 @server.PromptServer.instance.routes.post("/distributed/tunnel/stop")
 async def tunnel_stop_endpoint(request: web.Request) -> web.StreamResponse:
     """Stop the managed Cloudflare tunnel if running."""
-    try:
+
+    async def _operation() -> web.StreamResponse:
         result = await cloudflare_tunnel_manager.stop_tunnel()
         config = load_config()
-        return web.json_response({
-            "status": "success",
-            "tunnel": result,
-            "master_host": (config.get("master") or {}).get("host")
-        })
-    except Exception as e:
-        return await handle_api_error(request, e, 500)
+        return web.json_response(
+            {
+                "status": "success",
+                "tunnel": result,
+                "master_host": (config.get("master") or {}).get("host"),
+            }
+        )
+
+    return await run_authorized_endpoint(request, _operation)

api/usdu_routes.py

@@ -8,15 +8,50 @@ from aiohttp import web
 from PIL import Image
 import server
 
+from .request_guards import authorization_error_or_none
+from .schemas import require_bool_literal
 from ..upscale.job_models import BaseJobState, ImageJobState, TileJobState
-from ..upscale.job_store import MAX_PAYLOAD_SIZE, ensure_tile_jobs_initialized, init_dynamic_job
+from ..upscale.job_store import ensure_tile_jobs_initialized, init_dynamic_job
 from ..upscale.payload_parsers import parse_tiles_from_form
 from ..utils.logging import debug_log
 from ..utils.network import handle_api_error
+from ..utils.usdu_management import MAX_PAYLOAD_SIZE
+
+
+def _parse_int_field(value, field_name: str, *, minimum: int | None = None) -> int:
+    try:
+        parsed = int(value)
+    except (TypeError, ValueError) as exc:
+        raise ValueError(f"{field_name} must be an integer") from exc
+    if minimum is not None and parsed < minimum:
+        raise ValueError(f"{field_name} must be >= {minimum}")
+    return parsed
+
+
+def _parse_bool_field(
+    value,
+    field_name: str,
+    *,
+    default: bool = False,
+) -> bool:
+    if value is None:
+        return default
+    return require_bool_literal(value, field_name=field_name)
+
+
+def _is_worker_allowed(job_data: BaseJobState, worker_id: str) -> bool:
+    known_workers = {str(wid) for wid in getattr(job_data, "worker_status", {}).keys()}
+    known_workers.update(str(wid) for wid in getattr(job_data, "assigned_to_workers", {}).keys())
+    if not known_workers:
+        return True
+    return str(worker_id) in known_workers
 
 
 @server.PromptServer.instance.routes.post("/distributed/heartbeat")
 async def heartbeat_endpoint(request: web.Request) -> web.StreamResponse:
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
     try:
         data = await request.json()
         worker_id = data.get('worker_id')
@@ -30,6 +65,8 @@ async def heartbeat_endpoint(request: web.Request) -> web.StreamResponse:
             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):
+                    if not _is_worker_allowed(job_data, worker_id):
+                        return await handle_api_error(request, "Unauthorized worker_id", 403)
                     job_data.worker_status[worker_id] = time.time()
                     debug_log(f"Heartbeat from worker {worker_id}")
                     return web.json_response({"status": "success"})
@@ -42,22 +79,36 @@ async def heartbeat_endpoint(request: web.Request) -> web.StreamResponse:
 @server.PromptServer.instance.routes.post("/distributed/submit_tiles")
 async def submit_tiles_endpoint(request: web.Request) -> web.StreamResponse:
     """Endpoint for workers to submit processed tiles in static mode."""
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
     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)
+        if content_length:
+            try:
+                payload_size = int(content_length)
+            except ValueError:
+                return await handle_api_error(request, "Invalid content-length header", 400)
+            if payload_size > 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'
+        try:
+            is_last = _parse_bool_field(data.get('is_last', False), "is_last", default=False)
+        except ValueError as e:
+            return await handle_api_error(request, str(e), 400)
 
         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))
+        try:
+            batch_size = _parse_int_field(data.get('batch_size', 0), "batch_size", minimum=0)
+        except ValueError as e:
+            return await handle_api_error(request, str(e), 400)
 
         # Handle completion signal
         if batch_size == 0 and is_last:
@@ -66,6 +117,8 @@ async def submit_tiles_endpoint(request: web.Request) -> web.StreamResponse:
                     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)
+                    if not _is_worker_allowed(job_data, worker_id):
+                        return await handle_api_error(request, "Unauthorized worker_id", 403)
                     await job_data.queue.put({
                         'worker_id': worker_id,
                         'is_last': True,
@@ -85,6 +138,8 @@ async def submit_tiles_endpoint(request: web.Request) -> web.StreamResponse:
                 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)
+                if not _is_worker_allowed(job_data, worker_id):
+                    return await handle_api_error(request, "Unauthorized worker_id", 403)
 
                 q = job_data.queue
                 if batch_size > 0 or len(tiles) > 0:
@@ -110,15 +165,26 @@ async def submit_tiles_endpoint(request: web.Request) -> web.StreamResponse:
 @server.PromptServer.instance.routes.post("/distributed/submit_image")
 async def submit_image_endpoint(request: web.Request) -> web.StreamResponse:
     """Endpoint for workers to submit processed images in dynamic mode."""
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
     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)
+        if content_length:
+            try:
+                payload_size = int(content_length)
+            except ValueError:
+                return await handle_api_error(request, "Invalid content-length header", 400)
+            if payload_size > 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'
+        try:
+            is_last = _parse_bool_field(data.get('is_last', False), "is_last", default=False)
+        except ValueError as e:
+            return await handle_api_error(request, str(e), 400)
 
         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)
@@ -127,7 +193,10 @@ async def submit_image_endpoint(request: web.Request) -> web.StreamResponse:
 
         # Handle image submission
         if 'full_image' in data and 'image_idx' in data:
-            image_idx = int(data.get('image_idx'))
+            try:
+                image_idx = _parse_int_field(data.get('image_idx'), "image_idx", minimum=0)
+            except ValueError as e:
+                return await handle_api_error(request, str(e), 400)
             img_data = data['full_image'].file.read()
             img = Image.open(io.BytesIO(img_data)).convert("RGB")
 
@@ -138,6 +207,8 @@ async def submit_image_endpoint(request: web.Request) -> web.StreamResponse:
                     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)
+                    if not _is_worker_allowed(job_data, worker_id):
+                        return await handle_api_error(request, "Unauthorized worker_id", 403)
                     await job_data.queue.put({
                         'worker_id': worker_id,
                         'image_idx': image_idx,
@@ -153,6 +224,8 @@ async def submit_image_endpoint(request: web.Request) -> web.StreamResponse:
                     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)
+                    if not _is_worker_allowed(job_data, worker_id):
+                        return await handle_api_error(request, "Unauthorized worker_id", 403)
                     await job_data.queue.put({
                         'worker_id': worker_id,
                         'is_last': True,
@@ -170,6 +243,9 @@ async def submit_image_endpoint(request: web.Request) -> web.StreamResponse:
 @server.PromptServer.instance.routes.post("/distributed/request_image")
 async def request_image_endpoint(request: web.Request) -> web.StreamResponse:
     """Endpoint for workers to request tasks (images in dynamic mode, tiles in static mode)."""
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
     try:
         data = await request.json()
         worker_id = data.get('worker_id')
@@ -192,6 +268,8 @@ async def request_image_endpoint(request: web.Request) -> web.StreamResponse:
                     pending_queue = job_data.pending_tasks
                 else:
                     return await handle_api_error(request, "Invalid job configuration", 400)
+                if not _is_worker_allowed(job_data, worker_id):
+                    return await handle_api_error(request, "Unauthorized worker_id", 403)
 
                 try:
                     task_idx = await asyncio.wait_for(pending_queue.get(), timeout=0.1)
@@ -201,99 +279,33 @@ async def request_image_endpoint(request: web.Request) -> web.StreamResponse:
 
                     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})
+                        return web.json_response(
+                            {
+                                "kind": "image",
+                                "task_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,
+                        "kind": "tile",
+                        "task_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 web.json_response({"kind": "none", "task_idx": None, "estimated_remaining": 0})
             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/init_list_queue")
-async def init_list_queue_endpoint(request: web.Request) -> web.StreamResponse:
-    """Initialize a shared queue of list indices for dynamic DistributedListSplitter mode."""
-    try:
-        data = await request.json()
-        multi_job_id = data.get("multi_job_id")
-        list_size = data.get("list_size", 0)
-        enabled_workers = data.get("enabled_workers", [])
-
-        if not multi_job_id:
-            return await handle_api_error(request, "Missing multi_job_id", 400)
-
-        try:
-            list_size = max(0, int(list_size))
-        except (TypeError, ValueError):
-            return await handle_api_error(request, "list_size must be an integer", 400)
-
-        worker_ids = [str(worker_id) for worker_id in (enabled_workers or []) if str(worker_id).strip()]
-        await init_dynamic_job(
-            multi_job_id=str(multi_job_id),
-            batch_size=list_size,
-            enabled_workers=worker_ids,
-            all_indices=list(range(list_size)),
-        )
-
-        prompt_server = ensure_tile_jobs_initialized()
-        async with prompt_server.distributed_tile_jobs_lock:
-            job_data = prompt_server.distributed_pending_tile_jobs.get(str(multi_job_id))
-            remaining = (
-                job_data.pending_images.qsize()
-                if isinstance(job_data, ImageJobState)
-                else 0
-            )
-
-        return web.json_response({"status": "success", "remaining": remaining})
-    except Exception as e:
-        return await handle_api_error(request, e, 500)
-
-
-@server.PromptServer.instance.routes.post("/distributed/request_list_item")
-async def request_list_item_endpoint(request: web.Request) -> web.StreamResponse:
-    """Allow workers to pull one list item index for dynamic list distribution."""
-    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:
-            job_data = prompt_server.distributed_pending_tile_jobs.get(str(multi_job_id))
-            if not isinstance(job_data, BaseJobState):
-                return await handle_api_error(request, "Job not found", 404)
-
-            pending_queue = getattr(job_data, "pending_images", None) or getattr(job_data, "pending_tasks", None)
-            if pending_queue is None:
-                return await handle_api_error(request, "Job not configured for list item requests", 400)
-
-            try:
-                item_idx = await asyncio.wait_for(pending_queue.get(), timeout=0.1)
-            except asyncio.TimeoutError:
-                return web.json_response({"item_idx": None})
-
-            worker_id = str(worker_id)
-            job_data.assigned_to_workers.setdefault(worker_id, []).append(item_idx)
-            job_data.worker_status[worker_id] = time.time()
-            remaining = pending_queue.qsize()
-            return web.json_response({"item_idx": int(item_idx), "estimated_remaining": remaining})
-    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: web.Request) -> web.StreamResponse:
     """Endpoint to check if a job is ready."""
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
     multi_job_id = request.query.get('multi_job_id')
     if not multi_job_id:
         return web.json_response({"ready": False})

api/worker_routes.py

@@ -25,27 +25,27 @@ from ..utils.network import (
 )
 from ..utils.constants import CHUNK_SIZE
 from ..workers import get_worker_manager
-from .schemas import require_fields, validate_worker_id
+from .request_guards import authorization_error_or_none
+from .schemas import (
+    distributed_auth_headers,
+    require_fields,
+    require_worker_id,
+)
 from ..workers.detection import (
     get_machine_id,
     is_docker_environment,
     is_runpod_environment,
 )
-try:
-    from ..utils.async_helpers import PromptValidationError, queue_prompt_payload
-except ImportError:
-    from ..utils.async_helpers import queue_prompt_payload
-
-    class PromptValidationError(RuntimeError):
-        def __init__(self, message, validation_error=None, node_errors=None):
-            super().__init__(str(message))
-            self.validation_error = validation_error if isinstance(validation_error, dict) else {}
-            self.node_errors = node_errors if isinstance(node_errors, dict) else {}
+from ..utils.async_helpers import PromptValidationError, queue_prompt_payload
 
 
 @server.PromptServer.instance.routes.get("/distributed/worker_ws")
-async def worker_ws_endpoint(request: web.Request) -> web.WebSocketResponse:
+async def worker_ws_endpoint(request: web.Request) -> web.StreamResponse:
     """WebSocket endpoint for worker prompt dispatch."""
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
+
     ws = web.WebSocketResponse(heartbeat=30)
     await ws.prepare(request)
 
@@ -118,6 +118,9 @@ async def worker_ws_endpoint(request: web.Request) -> web.WebSocketResponse:
 @server.PromptServer.instance.routes.post("/distributed/worker/clear_launching")
 async def clear_launching_state(request: web.Request) -> web.StreamResponse:
     """Clear the launching flag when worker is confirmed running."""
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
     try:
         wm = get_worker_manager()
         data = await request.json()
@@ -127,8 +130,10 @@ async def clear_launching_state(request: web.Request) -> web.StreamResponse:
 
         worker_id = str(data.get("worker_id")).strip()
         config = load_config()
-        if not validate_worker_id(worker_id, config):
-            return await handle_api_error(request, f"Worker {worker_id} not found", 404)
+        try:
+            worker_id = require_worker_id(worker_id, config)
+        except ValueError as exc:
+            return await handle_api_error(request, exc, 404)
         
         # Clear launching flag in managed processes
         if worker_id in wm.processes:
@@ -372,6 +377,9 @@ def _find_worker_by_id(config: dict[str, Any], worker_id: str) -> dict[str, Any]
 @server.PromptServer.instance.routes.get("/distributed/local_log")
 async def get_local_log_endpoint(request: web.Request) -> web.StreamResponse:
     """Return this instance's in-memory ComfyUI log buffer."""
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
     try:
         from app.logger import get_logs
     except Exception as e:
@@ -417,6 +425,9 @@ async def get_local_log_endpoint(request: web.Request) -> web.StreamResponse:
 @server.PromptServer.instance.routes.get("/distributed/network_info")
 async def get_network_info_endpoint(request: web.Request) -> web.StreamResponse:
     """Get network interfaces and recommend best IP for master."""
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
     try:
         loop = asyncio.get_running_loop()
         info = await loop.run_in_executor(None, _collect_network_info_sync)
@@ -432,6 +443,9 @@ async def get_network_info_endpoint(request: web.Request) -> web.StreamResponse:
 @server.PromptServer.instance.routes.get("/distributed/system_info")
 async def get_system_info_endpoint(request: web.Request) -> web.StreamResponse:
     """Get system information including machine ID for local worker detection."""
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
     try:
         import socket
         
@@ -456,6 +470,9 @@ async def get_system_info_endpoint(request: web.Request) -> web.StreamResponse:
 @server.PromptServer.instance.routes.post("/distributed/launch_worker")
 async def launch_worker_endpoint(request: web.Request) -> web.StreamResponse:
     """Launch a worker process from the UI."""
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
     try:
         wm = get_worker_manager()
         data = await request.json()
@@ -467,8 +484,10 @@ async def launch_worker_endpoint(request: web.Request) -> web.StreamResponse:
         
         # Find worker config
         config = load_config()
-        if not validate_worker_id(worker_id, config):
-            return await handle_api_error(request, f"Worker {worker_id} not found", 404)
+        try:
+            worker_id = require_worker_id(worker_id, config)
+        except ValueError as exc:
+            return await handle_api_error(request, exc, 404)
         worker = next((w for w in config.get("workers", []) if str(w.get("id")) == worker_id), None)
         if not worker:
             return await handle_api_error(request, f"Worker {worker_id} not found", 404)
@@ -487,7 +506,7 @@ async def launch_worker_endpoint(request: web.Request) -> web.StreamResponse:
                 is_running = process.poll() is None
             else:
                 # Restored process without subprocess object
-                is_running = wm._is_process_running(proc_info['pid'])
+                is_running = wm.is_process_running(proc_info['pid'])
             
             if is_running:
                 return await handle_api_error(request, "Worker already running (managed by UI)", 409)
@@ -517,6 +536,9 @@ async def launch_worker_endpoint(request: web.Request) -> web.StreamResponse:
 @server.PromptServer.instance.routes.post("/distributed/stop_worker")
 async def stop_worker_endpoint(request: web.Request) -> web.StreamResponse:
     """Stop a worker process that was launched from the UI."""
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
     try:
         wm = get_worker_manager()
         data = await request.json()
@@ -526,8 +548,10 @@ async def stop_worker_endpoint(request: web.Request) -> web.StreamResponse:
 
         worker_id = str(data.get("worker_id")).strip()
         config = load_config()
-        if not validate_worker_id(worker_id, config):
-            return await handle_api_error(request, f"Worker {worker_id} not found", 404)
+        try:
+            worker_id = require_worker_id(worker_id, config)
+        except ValueError as exc:
+            return await handle_api_error(request, exc, 404)
 
         success, message = wm.stop_worker(worker_id)
         
@@ -547,6 +571,9 @@ async def stop_worker_endpoint(request: web.Request) -> web.StreamResponse:
 @server.PromptServer.instance.routes.get("/distributed/managed_workers")
 async def get_managed_workers_endpoint(request: web.Request) -> web.StreamResponse:
     """Get list of workers managed by this UI instance."""
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
     try:
         managed = get_worker_manager().get_managed_workers()
         return web.json_response({
@@ -560,6 +587,9 @@ async def get_managed_workers_endpoint(request: web.Request) -> web.StreamRespon
 @server.PromptServer.instance.routes.get("/distributed/local-worker-status")
 async def get_local_worker_status_endpoint(request: web.Request) -> web.StreamResponse:
     """Check status of all local workers (localhost/no host specified)."""
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
     try:
         config = load_config()
         worker_statuses = {}
@@ -630,6 +660,9 @@ async def get_local_worker_status_endpoint(request: web.Request) -> web.StreamRe
 @server.PromptServer.instance.routes.get("/distributed/worker_log/{worker_id}")
 async def get_worker_log_endpoint(request: web.Request) -> web.StreamResponse:
     """Get log content for a specific worker."""
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
     try:
         wm = get_worker_manager()
         worker_id = request.match_info['worker_id']
@@ -673,6 +706,9 @@ async def get_worker_log_endpoint(request: web.Request) -> web.StreamResponse:
 @server.PromptServer.instance.routes.get("/distributed/remote_worker_log/{worker_id}")
 async def get_remote_worker_log_endpoint(request: web.Request) -> web.StreamResponse:
     """Proxy a remote worker log request to the worker's local in-memory log endpoint."""
+    auth_error = await authorization_error_or_none(request)
+    if auth_error is not None:
+        return auth_error
     try:
         worker_id = str(request.match_info["worker_id"]).strip()
         config = load_config()
@@ -695,6 +731,7 @@ async def get_remote_worker_log_endpoint(request: web.Request) -> web.StreamResp
         async with session.get(
             worker_url,
             params={"lines": str(lines_to_read)},
+            headers=distributed_auth_headers(config),
             timeout=aiohttp.ClientTimeout(total=5),
         ) as resp:
             if resp.status >= 400:

bootstrap/__init__.py

@@ -0,0 +1,5 @@
+"""Bootstrap entrypoints for ComfyUI-Distributed."""
+
+from .entrypoint import build_node_mappings, initialize_runtime
+
+__all__ = ["build_node_mappings", "initialize_runtime"]

bootstrap/entrypoint.py

@@ -0,0 +1,66 @@
+"""Single startup owner for node surface and runtime initialization."""
+from __future__ import annotations
+
+import atexit
+import os
+from functools import lru_cache
+from typing import Any
+
+
+@lru_cache(maxsize=1)
+def build_node_mappings() -> tuple[dict[str, Any], dict[str, str]]:
+    """Build and cache the full exported node mapping surface."""
+    from ..nodes import (
+        NODE_CLASS_MAPPINGS as distributed_class_mappings,
+        NODE_DISPLAY_NAME_MAPPINGS as distributed_display_name_mappings,
+    )
+    from ..nodes.distributed_upscale import (
+        NODE_CLASS_MAPPINGS as upscale_class_mappings,
+        NODE_DISPLAY_NAME_MAPPINGS as upscale_display_name_mappings,
+    )
+
+    return (
+        {**distributed_class_mappings, **upscale_class_mappings},
+        {**distributed_display_name_mappings, **upscale_display_name_mappings},
+    )
+
+
+@lru_cache(maxsize=1)
+def _runtime_initialized() -> dict[str, bool]:
+    return {"done": False}
+
+
+def initialize_runtime(prompt_server: Any | None = None) -> None:
+    """Initialize distributed runtime side effects once."""
+    state = _runtime_initialized()
+    if state["done"]:
+        return
+
+    import server
+
+    from ..api import bootstrap_routes
+    from ..api.queue_orchestration import ensure_distributed_state
+    from ..upscale.job_store import ensure_tile_jobs_initialized
+    from ..utils.config import CONFIG_FILE, ensure_config_exists
+    from ..utils.logging import debug_log
+    from ..workers.startup import delayed_auto_launch, register_async_signals, sync_cleanup
+
+    ensure_config_exists()
+    bootstrap_routes()
+
+    server_instance = prompt_server if prompt_server is not None else server.PromptServer.instance
+    ensure_distributed_state(server_instance)
+    ensure_tile_jobs_initialized()
+
+    atexit.register(lambda: None)  # placeholder; real cleanup in sync_cleanup
+    if not os.environ.get("COMFYUI_IS_WORKER"):
+        atexit.register(sync_cleanup)
+        delayed_auto_launch()
+        register_async_signals()
+
+    state["done"] = True
+
+    node_mappings, _ = build_node_mappings()
+    debug_log("Loaded Distributed nodes.")
+    debug_log(f"Config file: {CONFIG_FILE}")
+    debug_log(f"Available nodes: {list(node_mappings.keys())}")

conftest.py

@@ -1,28 +1,4 @@
-# conftest.py — project-level pytest configuration.
-#
-# Problem: custom_nodes/ComfyUI-Distributed/__init__.py uses relative imports
-# (from .distributed import ...) that fail when pytest tries to import it as a
-# standalone module during Package.setup() for the root package node.
-#
-# Fix: patch Package.setup() to skip the root-package's __init__.py import.
-# All actual package context is provided by each test module via
-# importlib.util.spec_from_file_location with synthetic stub packages.
-
-from _pytest.python import Package
-
-_orig_pkg_setup = Package.setup
-
-
-def _patched_pkg_setup(self) -> None:
-    # Skip the root package setup — its __init__.py uses relative imports
-    # that require a parent package (ComfyUI's plugin loader) which is not
-    # available in the test environment.
-    if self.path == self.config.rootpath:
-        return
-    _orig_pkg_setup(self)
-
-
-Package.setup = _patched_pkg_setup
+"""Project-level pytest collection rules for plugin-style package layout."""
 
 collect_ignore = [
     "__init__.py",

distributed.py

@@ -1,55 +1,15 @@
-"""
-ComfyUI-Distributed: thin entry point.
-All implementation lives in workers/, nodes/, api/.
-"""
-import atexit
-import os
+"""Compatibility facade for legacy imports from `distributed`."""
+from __future__ import annotations
 
-import server
-
-from .utils.config import ensure_config_exists
-from .utils.logging import debug_log
-from .utils.network import cleanup_client_session
-from .workers import get_worker_manager
-from .workers.startup import delayed_auto_launch, register_async_signals, sync_cleanup
-from .upscale.job_store import ensure_tile_jobs_initialized
-from .nodes import (
-    NODE_CLASS_MAPPINGS,
-    NODE_DISPLAY_NAME_MAPPINGS,
-    ImageBatchDivider,
-    DistributedCollectorNode,
-    DistributedSeed,
-    DistributedModelName,
-    DistributedValue,
-    AudioBatchDivider,
-    DistributedEmptyImage,
-    DistributedListSplitter,
-    DistributedListCollector,
-    DistributedBranch,
-    DistributedBranchCollector,
-    AnyType,
-    ByPassTypeTuple,
-    any_type,
+from .bootstrap.entrypoint import (
+    build_node_mappings,
+    initialize_runtime,
 )
-from . import api  # noqa: F401 - triggers all @routes.* registrations
-from .api.queue_orchestration import ensure_distributed_state
 
-ensure_config_exists()
+NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS = build_node_mappings()
 
-# Aiohttp session cleanup
-async def _cleanup_session():
-    await cleanup_client_session()
-
-
-atexit.register(lambda: None)  # placeholder; real cleanup in sync_cleanup
-
-# Initialize distributed job state on prompt_server
-prompt_server = server.PromptServer.instance
-ensure_distributed_state(prompt_server)
-ensure_tile_jobs_initialized()
-
-# Worker startup
-if not os.environ.get('COMFYUI_IS_WORKER'):
-    atexit.register(sync_cleanup)
-    delayed_auto_launch()
-    register_async_signals()
+__all__ = [
+    "NODE_CLASS_MAPPINGS",
+    "NODE_DISPLAY_NAME_MAPPINGS",
+    "initialize_runtime",
+]

docs/comfyui-distributed-api.md

@@ -53,7 +53,6 @@ Queue a workflow for distributed execution.
   "delegate_master": false,
   "enabled_worker_ids": ["1", "2"],
   "workers": ["1", "2"],
-  "auto_prepare": true,
   "trace_execution_id": "exec_1700000000_ab12cd"
 }
 ```
@@ -74,10 +73,6 @@ Queue a workflow for distributed execution.
   - The explicit worker IDs to consider for this run.
 - `workers` (optional, array of strings or objects with `id`)
   - Transitional alias for `enabled_worker_ids` used by older clients.
-- `auto_prepare` (optional, boolean)
-  - Kept for wire compatibility.
-  - Backend orchestration always runs with auto-prepare semantics.
-  - If top-level `prompt` is omitted, backend will attempt `workflow.prompt`.
 - `trace_execution_id` (optional, string)
   - Passed through to orchestration logs.
   - Server log lines include the marker as `[exec:<trace_execution_id>]`.
@@ -105,8 +100,7 @@ $cfg.workers | Select-Object id,name,enabled,host,port,type | Format-Table -Auto
 ```json
 {
   "prompt_id": "<uuid>",
-  "worker_count": 2,
-  "auto_prepare_supported": true
+  "worker_count": 2
 }
 ```
 

nodes/__init__.py

@@ -10,15 +10,11 @@ from .utilities import (
     any_type,
 )
 from .collector import DistributedCollectorNode
-from .list_splitter import DistributedListSplitter
-from .list_collector import DistributedListCollector
 from .branch import DistributedBranch
 from .branch_collector import DistributedBranchCollector
 
 NODE_CLASS_MAPPINGS = {
     "DistributedCollector": DistributedCollectorNode,
-    "DistributedListSplitter": DistributedListSplitter,
-    "DistributedListCollector": DistributedListCollector,
     "DistributedBranch": DistributedBranch,
     "DistributedBranchCollector": DistributedBranchCollector,
     "DistributedSeed": DistributedSeed,
@@ -30,8 +26,6 @@ NODE_CLASS_MAPPINGS = {
 }
 NODE_DISPLAY_NAME_MAPPINGS = {
     "DistributedCollector": "Distributed Collector",
-    "DistributedListSplitter": "Distributed List Splitter",
-    "DistributedListCollector": "Distributed List Collector",
     "DistributedBranch": "Distributed Branch",
     "DistributedBranchCollector": "Distributed Branch Collector",
     "DistributedSeed": "Distributed Seed",

nodes/branch_collector.py

@@ -1,16 +1,17 @@
 import asyncio
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Any
 
 import aiohttp
 import torch
 
+from ..utils.auth import distributed_auth_headers
 from ..utils.async_helpers import run_async_in_server_loop
-from ..utils.config import is_master_delegate_only
+from ..utils.config import is_master_delegate_only, load_config
 from ..utils.image import encode_tensor_png_data_url, ensure_contiguous
 from ..utils.logging import log
 from ..utils.network import get_client_session
-from ..utils.worker_ids import parse_enabled_worker_ids
+from ..utils.worker_ids import coerce_enabled_worker_ids
 from .context_kwargs import parse_distributed_hidden_context
 from .hidden_inputs import build_distributed_hidden_inputs
 from .queue_wait import collect_worker_queue_results
@@ -29,7 +30,7 @@ class BranchRunContext:
     multi_job_id: str = ""
     is_worker: bool = False
     master_url: str = ""
-    enabled_worker_ids: str = "[]"
+    enabled_worker_ids: list[str] = field(default_factory=list)
     worker_id: str = ""
     assigned_branch: int = -1
     delegate_only: bool = False
@@ -190,6 +191,7 @@ class DistributedBranchCollector:
         async with session.post(
             url,
             json=payload,
+            headers=distributed_auth_headers(load_config()),
             timeout=aiohttp.ClientTimeout(total=60),
         ) as response:
             response.raise_for_status()
@@ -236,7 +238,7 @@ class DistributedBranchCollector:
             return tuple(outputs)
 
         delegate_mode = bool(run_context.delegate_only or is_master_delegate_only())
-        enabled_workers = parse_enabled_worker_ids(run_context.enabled_worker_ids)
+        enabled_workers = coerce_enabled_worker_ids(run_context.enabled_worker_ids)
         slots_by_participant = self._participant_branch_slots(enabled_workers, delegate_mode, branch_count)
         expected_worker_ids = self._expected_worker_ids_for_branches(enabled_workers, delegate_mode, branch_count)
         expected_workers = set(expected_worker_ids)
@@ -282,6 +284,8 @@ class DistributedBranchCollector:
         finally:
             async with prompt_server.distributed_jobs_lock:
                 prompt_server.distributed_pending_jobs.pop(run_context.multi_job_id, None)
+                if hasattr(prompt_server, "distributed_job_allowed_workers"):
+                    prompt_server.distributed_job_allowed_workers.pop(run_context.multi_job_id, None)
 
         missing_workers = sorted(expected_workers - workers_done)
         if missing_workers:

nodes/collector.py

@@ -3,7 +3,7 @@ import base64
 import io
 import time
 from collections.abc import Callable
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Any
 
 import aiohttp
@@ -12,12 +12,13 @@ import comfy.model_management
 import torch
 from comfy.utils import ProgressBar
 
+from ..utils.auth import distributed_auth_headers
 from ..utils.logging import debug_log, log
 from ..utils.config import get_worker_timeout_seconds, load_config, is_master_delegate_only
 from ..utils.constants import HEARTBEAT_INTERVAL
 from ..utils.image import tensor_to_pil, pil_to_tensor, ensure_contiguous
 from ..utils.network import build_worker_url, get_client_session, probe_worker
-from ..utils.worker_ids import parse_enabled_worker_ids
+from ..utils.worker_ids import coerce_enabled_worker_ids
 from ..utils.audio_payload import encode_audio_payload
 from ..utils.async_helpers import run_async_in_server_loop
 from .context_kwargs import parse_distributed_hidden_context
@@ -31,7 +32,7 @@ class CollectorRunContext:
     multi_job_id: str = ""
     is_worker: bool = False
     master_url: str = ""
-    enabled_worker_ids: str = "[]"
+    enabled_worker_ids: list[str] = field(default_factory=list)
     worker_batch_size: int = 1
     worker_id: str = ""
     pass_through: bool = False
@@ -142,6 +143,7 @@ class DistributedCollectorNode:
                 async with session.post(
                     url,
                     json=payload,
+                    headers=distributed_auth_headers(load_config()),
                     timeout=aiohttp.ClientTimeout(total=60),
                 ) as response:
                     response.raise_for_status()
@@ -280,6 +282,8 @@ class DistributedCollectorNode:
         async with prompt_server.distributed_jobs_lock:
             if multi_job_id in prompt_server.distributed_pending_jobs:
                 del prompt_server.distributed_pending_jobs[multi_job_id]
+            if hasattr(prompt_server, "distributed_job_allowed_workers"):
+                prompt_server.distributed_job_allowed_workers.pop(multi_job_id, None)
 
     def _build_master_inputs(self, images, audio, delegate_mode, multi_job_id, num_workers):
         if delegate_mode:
@@ -452,7 +456,7 @@ class DistributedCollectorNode:
         delegate_only,
     ):
         delegate_mode = delegate_only or is_master_delegate_only()
-        enabled_workers = parse_enabled_worker_ids(enabled_worker_ids)
+        enabled_workers = coerce_enabled_worker_ids(enabled_worker_ids)
         expected_workers = set(enabled_workers)
         num_workers = len(expected_workers)
         if num_workers == 0:

nodes/context_kwargs.py

@@ -1,13 +1,26 @@
 from typing import Any
 
+from ..utils.parsing import coerce_bool
+from ..utils.worker_ids import coerce_enabled_worker_ids
+
+
+def _parse_bool(value: Any, default: bool = False) -> bool:
+    return coerce_bool(value, default=default)
+
+
+def _parse_enabled_workers(value: Any) -> list[str]:
+    if value is None:
+        return []
+    return coerce_enabled_worker_ids(value)
+
 
 def parse_distributed_hidden_context(kwargs: dict[str, Any]) -> dict[str, Any]:
     """Parse common distributed hidden inputs from node kwargs."""
     return {
-        "multi_job_id": str(kwargs.get("multi_job_id", "")),
-        "is_worker": bool(kwargs.get("is_worker", False)),
-        "master_url": str(kwargs.get("master_url", "")),
-        "enabled_worker_ids": str(kwargs.get("enabled_worker_ids", "[]")),
-        "worker_id": str(kwargs.get("worker_id", "")),
-        "delegate_only": bool(kwargs.get("delegate_only", False)),
+        "multi_job_id": str(kwargs.get("multi_job_id", "") or ""),
+        "is_worker": _parse_bool(kwargs.get("is_worker", False)),
+        "master_url": str(kwargs.get("master_url", "") or ""),
+        "enabled_worker_ids": _parse_enabled_workers(kwargs.get("enabled_worker_ids", "[]")),
+        "worker_id": str(kwargs.get("worker_id", "") or ""),
+        "delegate_only": _parse_bool(kwargs.get("delegate_only", False)),
     }

nodes/distributed_upscale.py

@@ -1,4 +1,3 @@
-import json
 import math
 from functools import wraps
 from typing import Any, Callable
@@ -18,6 +17,7 @@ from ..upscale.processing_args import UpscaleCoreArgs
 from ..upscale.modes.single_gpu import SingleGpuModeMixin
 from ..upscale.modes.static import StaticModeMixin
 from ..upscale.modes.dynamic import DynamicModeMixin
+from ..utils.worker_ids import coerce_enabled_worker_ids
 
 def sync_wrapper(async_func: Callable[..., Any]) -> Callable[..., Any]:
     """Decorator to wrap async methods for synchronous execution."""
@@ -30,21 +30,10 @@ def sync_wrapper(async_func: Callable[..., Any]) -> Callable[..., Any]:
         )
     return sync_func
 
-def _parse_enabled_worker_ids(enabled_worker_ids):
-    """Parse enabled worker IDs from either JSON or list input."""
-    if isinstance(enabled_worker_ids, list):
-        return [str(worker_id) for worker_id in enabled_worker_ids]
-    if not enabled_worker_ids:
-        return []
-    if isinstance(enabled_worker_ids, str):
-        try:
-            parsed = json.loads(enabled_worker_ids)
-        except json.JSONDecodeError:
-            log("USDU Dist: Invalid enabled_worker_ids JSON; defaulting to no workers.")
-            return []
-        if isinstance(parsed, list):
-            return [str(wid) for wid in parsed]
-    return []
+
+def _parse_enabled_worker_ids(enabled_worker_ids: str | list[str] | None) -> list[str]:
+    """Backward-compatible alias for enabled-worker normalization."""
+    return coerce_enabled_worker_ids(enabled_worker_ids)
 
 class UltimateSDUpscaleDistributed(
     DynamicModeMixin,
@@ -255,7 +244,7 @@ class UltimateSDUpscaleDistributed(
         
         # Ensure mode consistency across master/workers via shared threshold
         # Determine mode (must match master's logic)
-        enabled_workers = _parse_enabled_worker_ids(enabled_worker_ids)
+        enabled_workers = coerce_enabled_worker_ids(enabled_worker_ids)
         num_workers = len(enabled_workers)
         # Compute number of tiles for this image to decide if tile distribution makes sense
         _, height, width, _ = upscaled_image.shape
@@ -264,8 +253,8 @@ class UltimateSDUpscaleDistributed(
 
         mode = self._determine_processing_mode(batch_size, num_workers, dynamic_threshold)
         # For USDU-style processing, we want tile distribution whenever workers are available
-        # and there is more than one tile to process, even if batch == 1.
-        if num_workers > 0 and num_tiles_per_image > 1:
+        # and there is more than one tile to process for single-image runs.
+        if num_workers > 0 and batch_size <= 1 and num_tiles_per_image > 1:
             mode = "static"
             
         debug_log(f"USDU Dist Worker - Batch size {batch_size}")
@@ -324,14 +313,14 @@ class UltimateSDUpscaleDistributed(
         )
         
         # Parse enabled workers
-        enabled_workers = _parse_enabled_worker_ids(enabled_worker_ids)
+        enabled_workers = coerce_enabled_worker_ids(enabled_worker_ids)
         num_workers = len(enabled_workers)
         
         # Determine processing mode
         mode = self._determine_processing_mode(batch_size, num_workers, dynamic_threshold)
         # Prefer tile-based static distribution when workers are available and there are multiple tiles,
-        # even for batch == 1, to spread tiles across GPUs like the legacy dynamic tile queue.
-        if num_workers > 0 and num_tiles_per_image > 1:
+        # for single-image jobs to spread tiles across GPUs like the legacy tile queue.
+        if num_workers > 0 and batch_size <= 1 and num_tiles_per_image > 1:
             mode = "static"
         
         log(f"USDU Dist: Workers {num_workers} | Mode {mode} | Threshold {dynamic_threshold}")
@@ -370,18 +359,14 @@ class UltimateSDUpscaleDistributed(
                                                all_tiles, num_tiles_per_image)
 
     def _determine_processing_mode(self, batch_size: int, num_workers: int, dynamic_threshold: int) -> str:
-        """Determines processing mode per requested policy:
-        - any workers     => prefer static (tile-based) for USDU
-        - no workers      => single_gpu
-        """
+        """Determine mode from worker availability and configured threshold."""
         if num_workers == 0:
             return "single_gpu"
-        # Default to static when distributed; master/worker may still override if special cases arise
+        threshold = max(1, int(dynamic_threshold))
+        if int(batch_size) >= threshold:
+            return "dynamic"
         return "static"
 
-# Ensure initialization before registering routes
-ensure_tile_jobs_initialized()
-
 # Node registration
 NODE_CLASS_MAPPINGS = {
     "UltimateSDUpscaleDistributed": UltimateSDUpscaleDistributed,

nodes/hidden_inputs.py

@@ -1,6 +1,15 @@
 from typing import Any
 
 
+def build_worker_identity_hidden_inputs() -> dict[str, tuple[str, dict[str, Any]]]:
+    """Hidden-input contract for worker identity context."""
+    return {
+        "is_worker": ("BOOLEAN", {"default": False}),
+        "enabled_worker_ids": ("STRING", {"default": "[]"}),
+        "worker_id": ("STRING", {"default": ""}),
+    }
+
+
 def build_distributed_hidden_inputs(
     *,
     include_assigned_branch: bool = False,
@@ -10,12 +19,10 @@ def build_distributed_hidden_inputs(
 ) -> dict[str, tuple[str, dict[str, Any]]]:
     hidden_inputs: dict[str, tuple[str, dict[str, Any]]] = {
         "multi_job_id": ("STRING", {"default": ""}),
-        "is_worker": ("BOOLEAN", {"default": False}),
         "master_url": ("STRING", {"default": ""}),
-        "enabled_worker_ids": ("STRING", {"default": "[]"}),
-        "worker_id": ("STRING", {"default": ""}),
         "delegate_only": ("BOOLEAN", {"default": False}),
     }
+    hidden_inputs.update(build_worker_identity_hidden_inputs())
     if include_assigned_branch:
         hidden_inputs["assigned_branch"] = (
             "INT",

nodes/list_collector.py

@@ -1,238 +0,0 @@
-import asyncio
-from dataclasses import dataclass
-from typing import Any
-
-import aiohttp
-import torch
-
-from ..utils.async_helpers import run_async_in_server_loop
-from ..utils.config import is_master_delegate_only
-from ..utils.image import encode_tensor_png_data_url, ensure_contiguous
-from ..utils.logging import debug_log, log
-from ..utils.network import get_client_session
-from ..utils.worker_ids import parse_enabled_worker_ids
-from .context_kwargs import parse_distributed_hidden_context
-from .hidden_inputs import build_distributed_hidden_inputs
-from .queue_wait import collect_worker_queue_results
-from .runtime_helpers import (
-    get_prompt_server_instance as _get_prompt_server_instance,
-    throw_if_processing_interrupted as _throw_if_processing_interrupted,
-)
-
-
-@dataclass(frozen=True)
-class ListCollectorRunContext:
-    multi_job_id: str = ""
-    is_worker: bool = False
-    master_url: str = ""
-    enabled_worker_ids: str = "[]"
-    worker_id: str = ""
-    delegate_only: bool = False
-
-
-class DistributedListCollector:
-    INPUT_IS_LIST = True
-    OUTPUT_IS_LIST = (True,)
-
-    @classmethod
-    def INPUT_TYPES(cls: type["DistributedListCollector"]) -> dict[str, Any]:
-        return {
-            "required": {
-                "images": ("IMAGE",),
-            },
-            "hidden": build_distributed_hidden_inputs(),
-        }
-
-    RETURN_TYPES = ("IMAGE",)
-    RETURN_NAMES = ("images",)
-    FUNCTION = "run"
-    CATEGORY = "image"
-
-    def _build_run_context(self, **kwargs: Any) -> ListCollectorRunContext:
-        return ListCollectorRunContext(**parse_distributed_hidden_context(kwargs))
-
-    def run(
-        self,
-        images: list[Any] | torch.Tensor | None,
-        **kwargs: Any,
-    ) -> tuple[list[Any]]:
-        context = self._build_run_context(**kwargs)
-        image_list = self._normalize_image_list(images)
-        if not context.multi_job_id:
-            return (image_list,)
-
-        return run_async_in_server_loop(
-            self.execute(
-                image_list,
-                context=context,
-            )
-        )
-
-    def _normalize_image_list(self, images: list[Any] | torch.Tensor | None) -> list[Any]:
-        if images is None:
-            return []
-        if isinstance(images, list):
-            return list(images)
-        return [images]
-
-    async def send_list_to_master(
-        self,
-        image_list: list[Any],
-        multi_job_id: str,
-        master_url: str,
-        worker_id: str,
-    ) -> None:
-        if not image_list:
-            return
-
-        sendable = []
-        for idx, image in enumerate(image_list):
-            if not isinstance(image, torch.Tensor):
-                continue
-            batch = image
-            if batch.ndim == 3:
-                batch = batch.unsqueeze(0)
-            if batch.ndim != 4 or batch.shape[0] <= 0:
-                continue
-            sendable.append((idx, batch))
-
-        if not sendable:
-            return
-
-        session = await get_client_session()
-        url = f"{master_url}/distributed/job_complete"
-
-        for send_pos, (image_index, image_batch) in enumerate(sendable):
-            payload = {
-                "job_id": str(multi_job_id),
-                "worker_id": str(worker_id),
-                "batch_idx": int(image_index),
-                "image": encode_tensor_png_data_url(image_batch, 0),
-                "is_last": bool(send_pos == len(sendable) - 1),
-            }
-            try:
-                async with session.post(
-                    url,
-                    json=payload,
-                    timeout=aiohttp.ClientTimeout(total=60),
-                ) as response:
-                    response.raise_for_status()
-            except Exception as exc:
-                log(f"Worker - Failed to send list item to master: {exc}")
-                debug_log(f"Worker - Full error details: URL={url}")
-                raise
-
-    def _store_worker_result(self, worker_images, item):
-        worker_id = str(item.get("worker_id", ""))
-        tensor = item.get("tensor")
-        image_index = item.get("image_index")
-        if not worker_id or tensor is None or image_index is None:
-            return 0
-
-        worker_images.setdefault(worker_id, {})
-        worker_images[worker_id][int(image_index)] = tensor
-        return 1
-
-    def _reorder_and_combine_list(self, worker_images, worker_order, master_images, delegate_mode):
-        ordered = []
-        if not delegate_mode:
-            ordered.extend(master_images)
-
-        ordered_worker_ids = [str(worker_id) for worker_id in (worker_order or [])]
-        seen = set()
-
-        for worker_id in ordered_worker_ids:
-            seen.add(worker_id)
-            worker_entries = worker_images.get(worker_id, {})
-            for image_index in sorted(worker_entries.keys()):
-                ordered.append(worker_entries[image_index])
-
-        for worker_id in sorted(worker_images.keys()):
-            if worker_id in seen:
-                continue
-            worker_entries = worker_images.get(worker_id, {})
-            for image_index in sorted(worker_entries.keys()):
-                ordered.append(worker_entries[image_index])
-
-        combined = []
-        for image in ordered:
-            if isinstance(image, torch.Tensor):
-                if image.is_cuda:
-                    image = image.cpu()
-                image = ensure_contiguous(image)
-            combined.append(image)
-
-        return combined
-
-    async def execute(self, *args: Any, **kwargs: Any) -> tuple[list[Any]]:
-        """Compatibility wrapper with a uniform execute() signature across collectors."""
-        return await self._execute_list(*args, **kwargs)
-
-    async def _execute_list(
-        self,
-        images: list[Any] | torch.Tensor | None,
-        context: ListCollectorRunContext | None = None,
-    ) -> tuple[list[Any]]:
-        run_context = context or ListCollectorRunContext()
-        image_list = self._normalize_image_list(images)
-
-        if run_context.is_worker:
-            debug_log(
-                "Worker - Job "
-                f"{run_context.multi_job_id} complete. Sending {len(image_list)} image list item(s) to master"
-            )
-            await self.send_list_to_master(
-                image_list,
-                run_context.multi_job_id,
-                run_context.master_url,
-                run_context.worker_id,
-            )
-            return (image_list,)
-
-        delegate_mode = bool(run_context.delegate_only or is_master_delegate_only())
-        enabled_workers = parse_enabled_worker_ids(run_context.enabled_worker_ids)
-        expected_workers = set(enabled_workers)
-        if not expected_workers:
-            return (image_list,)
-
-        prompt_server = _get_prompt_server_instance()
-
-        async with prompt_server.distributed_jobs_lock:
-            if run_context.multi_job_id not in prompt_server.distributed_pending_jobs:
-                prompt_server.distributed_pending_jobs[run_context.multi_job_id] = asyncio.Queue()
-                debug_log(
-                    "Master - Initialized queue early for list collector job "
-                    f"{run_context.multi_job_id}"
-                )
-
-        master_images = [] if delegate_mode else image_list
-
-        worker_images = {}
-        workers_done: set[str] = set()
-
-        def _handle_queue_result(result: dict[str, Any]) -> None:
-            self._store_worker_result(worker_images, result)
-
-        try:
-            workers_done = await collect_worker_queue_results(
-                prompt_server=prompt_server,
-                multi_job_id=run_context.multi_job_id,
-                expected_workers=expected_workers,
-                on_result=_handle_queue_result,
-                timeout_log_prefix=(
-                    "Master - List collector heartbeat timeout. "
-                    "Still waiting for workers: "
-                ),
-                throw_if_interrupted=_throw_if_processing_interrupted,
-            )
-
-        finally:
-            async with prompt_server.distributed_jobs_lock:
-                prompt_server.distributed_pending_jobs.pop(run_context.multi_job_id, None)
-
-        combined = self._reorder_and_combine_list(worker_images, enabled_workers, master_images, delegate_mode)
-        debug_log(
-            "Master - List collector job "
-            f"{run_context.multi_job_id} complete. Combined {len(combined)} image list item(s)."
-        )
-        return (combined,)

nodes/list_splitter.py

@@ -1,220 +0,0 @@
-import asyncio
-from dataclasses import dataclass
-from typing import Any
-
-import aiohttp
-
-from ..upscale.job_models import BaseJobState, ImageJobState
-from ..upscale.job_store import ensure_tile_jobs_initialized, init_dynamic_job
-from ..utils.async_helpers import run_async_in_server_loop
-from ..utils.logging import debug_log
-from ..utils.network import get_client_session
-from .utilities import _chunk_bounds
-
-
-@dataclass(frozen=True)
-class ListSplitterRunContext:
-    participant_index: int = 0
-    total_participants: int = 1
-    multi_job_id: str = ""
-    is_worker: bool = False
-    master_url: str = ""
-    worker_id: str = ""
-
-
-class DistributedListSplitter:
-    INPUT_IS_LIST = True
-    OUTPUT_IS_LIST = (True,)
-
-    @classmethod
-    def INPUT_TYPES(cls: type["DistributedListSplitter"]) -> dict[str, Any]:
-        return {
-            "required": {
-                "images": ("IMAGE",),
-                "mode": (["static", "dynamic"], {"default": "static"}),
-            },
-            "hidden": {
-                "participant_index": ("INT", {"default": 0, "min": 0, "max": 1024}),
-                "total_participants": ("INT", {"default": 1, "min": 1, "max": 1024}),
-                "multi_job_id": ("STRING", {"default": ""}),
-                "is_worker": ("BOOLEAN", {"default": False}),
-                "master_url": ("STRING", {"default": ""}),
-                "worker_id": ("STRING", {"default": ""}),
-            },
-        }
-
-    RETURN_TYPES = ("IMAGE",)
-    RETURN_NAMES = ("images",)
-    FUNCTION = "split"
-    CATEGORY = "image"
-
-    def _build_run_context(self, **kwargs: Any) -> ListSplitterRunContext:
-        try:
-            participant_index = int(kwargs.get("participant_index", 0))
-        except (TypeError, ValueError):
-            participant_index = 0
-        try:
-            total_participants = int(kwargs.get("total_participants", 1))
-        except (TypeError, ValueError):
-            total_participants = 1
-
-        return ListSplitterRunContext(
-            participant_index=participant_index,
-            total_participants=max(total_participants, 1),
-            multi_job_id=str(kwargs.get("multi_job_id", "")),
-            is_worker=bool(kwargs.get("is_worker", False)),
-            master_url=str(kwargs.get("master_url", "")),
-            worker_id=str(kwargs.get("worker_id", "")),
-        )
-
-    def split(
-        self,
-        images: list[Any] | Any | None,
-        mode: str = "static",
-        **kwargs: Any,
-    ) -> tuple[list[Any]]:
-        context = self._build_run_context(**kwargs)
-        if images is None:
-            return ([],)
-
-        image_list = images if isinstance(images, list) else [images]
-        if mode != "dynamic" or not context.multi_job_id:
-            return self._split_static(
-                image_list,
-                context.participant_index,
-                context.total_participants,
-            )
-
-        return run_async_in_server_loop(
-            self._split_dynamic(
-                image_list,
-                context.multi_job_id,
-                context.is_worker,
-                context.master_url,
-                context.worker_id,
-            )
-        )
-
-    def _split_static(self, image_list, participant_index, total_participants):
-        split_count = max(1, int(total_participants or 1))
-        if split_count <= 1:
-            return (list(image_list),)
-
-        bounds = _chunk_bounds(len(image_list), split_count)
-        index = int(participant_index or 0)
-        if index < 0:
-            index = 0
-        if index >= split_count:
-            index = split_count - 1
-
-        start, end = bounds[index]
-        return (list(image_list[start:end]),)
-
-    async def _ensure_local_dynamic_queue(self, multi_job_id, total_items):
-        prompt_server = ensure_tile_jobs_initialized()
-        async with prompt_server.distributed_tile_jobs_lock:
-            existing = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
-            if isinstance(existing, ImageJobState):
-                return
-        await init_dynamic_job(
-            multi_job_id,
-            batch_size=total_items,
-            enabled_workers=[],
-            all_indices=list(range(total_items)),
-        )
-
-    async def _pull_items_from_local_queue(self, multi_job_id, worker_id):
-        prompt_server = ensure_tile_jobs_initialized()
-        pulled = []
-        worker_key = str(worker_id or "master")
-
-        while True:
-            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):
-                    break
-                pending_queue = getattr(job_data, "pending_images", None) or getattr(job_data, "pending_tasks", None)
-                if pending_queue is None:
-                    break
-
-                try:
-                    item_idx = pending_queue.get_nowait()
-                except asyncio.QueueEmpty:
-                    break
-
-                job_data.assigned_to_workers.setdefault(worker_key, []).append(item_idx)
-                pulled.append(int(item_idx))
-
-        return pulled
-
-    async def _request_item_from_master(self, multi_job_id, master_url, worker_id):
-        session = await get_client_session()
-        url = f"{master_url}/distributed/request_list_item"
-        try:
-            async with session.post(
-                url,
-                json={"multi_job_id": str(multi_job_id), "worker_id": str(worker_id or "")},
-                timeout=aiohttp.ClientTimeout(total=10),
-            ) as response:
-                if response.status == 404:
-                    return "not_ready", None
-                if response.status != 200:
-                    return "error", None
-                payload = await response.json()
-                return "ok", payload.get("item_idx")
-        except Exception:
-            return "error", None
-
-    async def _pull_items_from_master(self, multi_job_id, master_url, worker_id):
-        pulled = []
-        consecutive_not_ready = 0
-        max_not_ready_polls = 40
-
-        while True:
-            status, item_idx = await self._request_item_from_master(multi_job_id, master_url, worker_id)
-
-            if status == "not_ready":
-                if consecutive_not_ready >= max_not_ready_polls:
-                    break
-                consecutive_not_ready += 1
-                await asyncio.sleep(0.25)
-                continue
-
-            if status == "error":
-                if pulled:
-                    break
-                await asyncio.sleep(0.25)
-                continue
-
-            consecutive_not_ready = 0
-            if item_idx is None:
-                break
-
-            try:
-                pulled.append(int(item_idx))
-            except (TypeError, ValueError) as exc:
-                debug_log(f"DistributedListSplitter dynamic mode: invalid item index '{item_idx}' ignored: {exc}")
-                continue
-
-        return pulled
-
-    async def _split_dynamic(self, image_list, multi_job_id, is_worker, master_url, worker_id):
-        if not image_list:
-            return ([],)
-
-        if is_worker:
-            pulled_indices = await self._pull_items_from_master(multi_job_id, master_url, worker_id)
-        else:
-            await self._ensure_local_dynamic_queue(multi_job_id, len(image_list))
-            pulled_indices = await self._pull_items_from_local_queue(multi_job_id, worker_id)
-
-        selected_images = [
-            image_list[index]
-            for index in pulled_indices
-            if isinstance(index, int) and 0 <= index < len(image_list)
-        ]
-        debug_log(
-            f"DistributedListSplitter dynamic mode: role={'worker' if is_worker else 'master'} "
-            f"pulled={len(selected_images)} items for job {multi_job_id}"
-        )
-        return (selected_images,)

nodes/utilities.py

@@ -3,6 +3,8 @@ import json
 from typing import Any
 
 from ..utils.logging import debug_log, log
+from ..utils.worker_ids import coerce_enabled_worker_ids, parse_worker_index, worker_value_key
+from .hidden_inputs import build_worker_identity_hidden_inputs
 
 
 def _chunk_bounds(total_items: int, n_splits: int) -> list[tuple[int, int]]:
@@ -40,8 +42,7 @@ class DistributedSeed:
                 }),
             },
             "hidden": {
-                "is_worker": ("BOOLEAN", {"default": False}),
-                "worker_id": ("STRING", {"default": ""}),
+                **build_worker_identity_hidden_inputs(),
             },
         }
     
@@ -55,30 +56,24 @@ class DistributedSeed:
         seed: int,
         is_worker: bool = False,
         worker_id: str = "",
+        enabled_worker_ids: str = "[]",
     ) -> tuple[int]:
         if not is_worker:
             # Master node: pass through original values
             debug_log(f"Distributor - Master: seed={seed}")
             return (seed,)
         else:
-            # Worker node: apply offset based on worker index
-            # Find worker index from enabled_worker_ids
-            try:
-                # Worker IDs are passed as "worker_0", "worker_1", etc.
-                if worker_id.startswith("worker_"):
-                    worker_index = int(worker_id.split("_")[1])
-                else:
-                    # Fallback: try to parse as direct index
-                    worker_index = int(worker_id)
-                
+            enabled_workers = coerce_enabled_worker_ids(enabled_worker_ids)
+            worker_index = parse_worker_index(worker_id, enabled_workers)
+            if worker_index is not None:
                 offset = worker_index + 1
                 new_seed = seed + offset
                 debug_log(f"Distributor - Worker {worker_index}: seed={seed} → {new_seed}")
                 return (new_seed,)
-            except (ValueError, IndexError) as e:
-                debug_log(f"Distributor - Error parsing worker_id '{worker_id}': {e}")
-                # Fallback: return original seed
-                return (seed,)
+
+            debug_log(f"Distributor - Error parsing worker_id '{worker_id}': no worker index resolved")
+            # Fallback: return original seed
+            return (seed,)
 
 
 # Define ByPassTypeTuple for flexible return types
@@ -105,8 +100,7 @@ class DistributedValue:
                 "worker_values": ("STRING", {"default": "{}"}),
             },
             "hidden": {
-                "is_worker": ("BOOLEAN", {"default": False}),
-                "worker_id": ("STRING", {"default": ""}),
+                **build_worker_identity_hidden_inputs(),
             },
         }
 
@@ -149,12 +143,14 @@ class DistributedValue:
         worker_values: str | dict[str, Any] = "{}",
         is_worker: bool = False,
         worker_id: str = "",
+        enabled_worker_ids: str = "[]",
     ) -> tuple[Any]:
         values = {}
         value_type = "STRING"
 
         try:
-            values = json.loads(worker_values) if isinstance(worker_values, str) else worker_values
+            raw_values = json.loads(worker_values) if isinstance(worker_values, str) else worker_values
+            values = dict(raw_values) if isinstance(raw_values, dict) else {}
             if not isinstance(values, dict):
                 values = {}
         except json.JSONDecodeError as e:
@@ -165,26 +161,23 @@ class DistributedValue:
         value_type = values.get("_type", inferred_type)
         if value_type not in {"STRING", "COMBO", "INT", "FLOAT"}:
             value_type = inferred_type
-        values["_type"] = value_type
         coerced_default = self._coerce_safe(default_value, value_type)
 
         if not is_worker:
             debug_log(f"DistributedValue - Master: returning default '{coerced_default}'")
             return (coerced_default,)
 
-        try:
-            if worker_id.startswith("worker_"):
-                idx = int(worker_id.split("_")[1])
-            else:
-                idx = int(worker_id)
-            key = str(idx + 1)  # worker_0 → key "1" (1-indexed)
-            raw = values.get(key, "")
-            if raw:
+        enabled_workers = coerce_enabled_worker_ids(enabled_worker_ids)
+        direct_key = str(worker_id).strip()
+        lookup_key = direct_key if direct_key in values else worker_value_key(worker_id, enabled_workers)
+        raw = values.get(lookup_key)
+        if raw is not None and raw != "":
+            try:
                 coerced = self._coerce(raw, value_type)
-                debug_log(f"DistributedValue - Worker {idx}: returning '{coerced}'")
+                debug_log(f"DistributedValue - Worker key {lookup_key}: returning '{coerced}'")
                 return (coerced,)
-        except (ValueError, IndexError) as e:
-            debug_log(f"DistributedValue - Error: {e}")
+            except (TypeError, ValueError) as e:
+                debug_log(f"DistributedValue - Error coercing worker value for key {lookup_key}: {e}")
         debug_log(f"DistributedValue - Worker fallback: returning default '{coerced_default}'")
         return (coerced_default,)
 

pyproject.toml

@@ -3,7 +3,10 @@ name = "ComfyUI-Distributed"
 description = "ComfyUI extension that enables multi-GPU processing locally, remotely and in the cloud"
 version = "1.4.1"
 license = {file = "LICENSE"}
-dependencies  = []
+dependencies = [
+    "aiohttp>=3.9,<4",
+    "Pillow>=10,<12",
+]
 
 [project.urls]
 Repository = "https://github.com/robertvoy/ComfyUI-Distributed"

tests/api/harness.py

@@ -0,0 +1,163 @@
+"""Shared stubs/utilities for API route module unit tests."""
+from __future__ import annotations
+
+import sys
+import types
+from typing import Any, Callable
+
+
+class RoutesStub:
+    """Simple decorator-style route registrar used in unit tests."""
+
+    def get(self, _path: str):
+        return lambda fn: fn
+
+    def post(self, _path: str):
+        return lambda fn: fn
+
+
+class ClientTimeoutStub:
+    def __init__(self, total: float | None = None):
+        self.total = total
+
+
+class WSMsgTypeStub:
+    TEXT = "TEXT"
+    ERROR = "ERROR"
+    CLOSED = "CLOSED"
+
+
+class WebSocketResponseStub:
+    def __init__(self, *args: Any, **kwargs: Any):
+        self.args = args
+        self.kwargs = kwargs
+
+    async def prepare(self, _request: Any) -> None:
+        return None
+
+    async def send_json(self, _payload: Any) -> None:
+        return None
+
+    def __aiter__(self):
+        async def _empty():
+            if False:
+                yield None
+
+        return _empty()
+
+
+class FormDataStub:
+    def add_field(self, *_args: Any, **_kwargs: Any) -> None:
+        return None
+
+
+def reset_package_namespace(package_name: str) -> None:
+    """Clear test package modules from sys.modules."""
+    prefix = f"{package_name}."
+    for mod_name in list(sys.modules):
+        if mod_name == package_name or mod_name.startswith(prefix):
+            del sys.modules[mod_name]
+
+
+def ensure_namespace_package(module_name: str) -> types.ModuleType:
+    """Create a namespace-like package module in sys.modules."""
+    module = types.ModuleType(module_name)
+    module.__path__ = []
+    sys.modules[module_name] = module
+    return module
+
+
+def bootstrap_test_package(
+    package_name: str,
+    *,
+    with_api: bool = True,
+    with_utils: bool = True,
+    with_workers: bool = False,
+    with_upscale: bool = False,
+    with_orchestration: bool = False,
+) -> None:
+    """Create a clean package scaffold for module-level import tests."""
+    reset_package_namespace(package_name)
+    ensure_namespace_package(package_name)
+
+    if with_api:
+        ensure_namespace_package(f"{package_name}.api")
+    if with_utils:
+        ensure_namespace_package(f"{package_name}.utils")
+    if with_workers:
+        ensure_namespace_package(f"{package_name}.workers")
+    if with_upscale:
+        ensure_namespace_package(f"{package_name}.upscale")
+    if with_orchestration:
+        ensure_namespace_package(f"{package_name}.api.orchestration")
+
+
+def install_server_stub(prompt_server_instance: Any | None = None) -> Any:
+    """Install a minimal `server` module exposing PromptServer.instance.routes."""
+    if prompt_server_instance is None:
+        prompt_server_instance = types.SimpleNamespace()
+    if not hasattr(prompt_server_instance, "routes"):
+        prompt_server_instance.routes = RoutesStub()
+
+    server_module = types.ModuleType("server")
+    server_module.PromptServer = types.SimpleNamespace(instance=prompt_server_instance)
+    sys.modules["server"] = server_module
+    return prompt_server_instance
+
+
+def install_aiohttp_stub(
+    json_response_factory: Callable[..., Any],
+) -> bool:
+    """Install/augment `aiohttp` with common members used by API tests."""
+    created = False
+    if "aiohttp" not in sys.modules:
+        created = True
+        aiohttp_module = types.ModuleType("aiohttp")
+        sys.modules["aiohttp"] = aiohttp_module
+
+    aiohttp_module = sys.modules["aiohttp"]
+    if not hasattr(aiohttp_module, "ClientTimeout"):
+        aiohttp_module.ClientTimeout = ClientTimeoutStub
+    if not hasattr(aiohttp_module, "WSMsgType"):
+        aiohttp_module.WSMsgType = WSMsgTypeStub
+    if not hasattr(aiohttp_module, "FormData"):
+        aiohttp_module.FormData = FormDataStub
+
+    web_obj = getattr(aiohttp_module, "web", None)
+    if web_obj is None:
+        web_obj = types.SimpleNamespace()
+        aiohttp_module.web = web_obj
+    web_obj.json_response = json_response_factory
+    if not hasattr(web_obj, "WebSocketResponse"):
+        web_obj.WebSocketResponse = WebSocketResponseStub
+
+    return created
+
+
+def cleanup_optional_module(module_name: str, created: bool) -> None:
+    """Remove stubbed module only when this test created it."""
+    if created:
+        sys.modules.pop(module_name, None)
+
+
+def install_request_guards_stub(package_name: str) -> None:
+    """Install permissive request guard for isolated route tests."""
+    request_guards_module = types.ModuleType(f"{package_name}.api.request_guards")
+
+    async def _authorization_error_or_none(_request):
+        return None
+
+    request_guards_module.authorization_error_or_none = _authorization_error_or_none
+    sys.modules[f"{package_name}.api.request_guards"] = request_guards_module
+
+
+def install_endpoint_policy_passthrough_stub(package_name: str) -> None:
+    """Install endpoint policy stub that executes operation directly."""
+    endpoint_policy_module = types.ModuleType(f"{package_name}.api.endpoint_policy")
+
+    async def _run_authorized_endpoint(_request, operation, unexpected_status=500):
+        _ = unexpected_status
+        return await operation()
+
+    endpoint_policy_module.run_authorized_endpoint = _run_authorized_endpoint
+    sys.modules[f"{package_name}.api.endpoint_policy"] = endpoint_policy_module

tests/api/test_auth_guard_integration.py

@@ -0,0 +1,423 @@
+import asyncio
+import copy
+import importlib.util
+import sys
+import types
+import unittest
+from contextlib import asynccontextmanager
+from pathlib import Path
+
+from tests.api.harness import (
+    bootstrap_test_package,
+    cleanup_optional_module,
+    install_aiohttp_stub,
+    install_server_stub,
+)
+
+
+AUTH_TOKEN = "integration-secret"
+
+
+class _FakeResponse:
+    def __init__(self, payload, status=200):
+        self.payload = payload
+        self.status = status
+
+
+class _FakeRequest:
+    def __init__(self, *, headers=None, json_payload=None, post_payload=None, query=None):
+        self.headers = headers or {}
+        self._json_payload = json_payload or {}
+        self._post_payload = post_payload or {}
+        self.query = query or {}
+
+    async def json(self):
+        return self._json_payload
+
+    async def post(self):
+        return self._post_payload
+
+
+def _load_module(package_name: str, rel_path: str, module_name: str):
+    module_path = Path(__file__).resolve().parents[2] / rel_path
+    spec = importlib.util.spec_from_file_location(f"{package_name}.{module_name}", module_path)
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+def _install_common_utils(package_name: str, config_data: dict):
+    logging_module = types.ModuleType(f"{package_name}.utils.logging")
+    logging_module.debug_log = lambda *_args, **_kwargs: None
+    logging_module.log = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.utils.logging"] = logging_module
+
+    config_module = types.ModuleType(f"{package_name}.utils.config")
+    config_module.load_config = lambda: copy.deepcopy(config_data)
+    config_module.save_config = lambda _cfg: True
+
+    @asynccontextmanager
+    async def _config_transaction():
+        yield config_module.load_config()
+
+    config_module.config_transaction = _config_transaction
+    sys.modules[f"{package_name}.utils.config"] = config_module
+
+    network_module = types.ModuleType(f"{package_name}.utils.network")
+
+    async def _handle_api_error(_request, error, status=500):
+        return _FakeResponse({"status": "error", "message": str(error)}, status=status)
+
+    network_module.handle_api_error = _handle_api_error
+    network_module.normalize_host = lambda value: value
+    network_module.build_worker_url = lambda worker, path="": f"http://{worker.get('host', '127.0.0.1')}{path}"
+    network_module.get_client_session = lambda: None
+    network_module.probe_worker = lambda *_args, **_kwargs: {"ok": True}
+    sys.modules[f"{package_name}.utils.network"] = network_module
+
+
+def _load_real_guard_stack(package_name: str):
+    _load_module(package_name, "utils/auth.py", "utils.auth")
+    _load_module(package_name, "utils/parsing.py", "utils.parsing")
+    _load_module(package_name, "api/schemas.py", "api.schemas")
+    _load_module(package_name, "api/request_guards.py", "api.request_guards")
+    _load_module(package_name, "api/endpoint_policy.py", "api.endpoint_policy")
+
+
+def _load_request_guards_only_module():
+    package_name = "dist_auth_guard_testpkg"
+    config_data = {"settings": {"distributed_api_token": AUTH_TOKEN}}
+    bootstrap_test_package(package_name, with_api=True, with_utils=True)
+    install_server_stub()
+    created_aiohttp_stub = install_aiohttp_stub(
+        lambda payload, status=200: _FakeResponse(payload, status=status)
+    )
+    _install_common_utils(package_name, config_data)
+    _load_real_guard_stack(package_name)
+    request_guards = sys.modules[f"{package_name}.api.request_guards"]
+    schemas = sys.modules[f"{package_name}.api.schemas"]
+    cleanup_optional_module("aiohttp", created_aiohttp_stub)
+    return request_guards, schemas
+
+
+def _load_config_routes_real_guard_module():
+    package_name = "dist_config_auth_integration_testpkg"
+    config_data = {
+        "workers": [],
+        "master": {"host": "127.0.0.1"},
+        "settings": {"distributed_api_token": AUTH_TOKEN},
+        "tunnel": {},
+    }
+    bootstrap_test_package(package_name, with_api=True, with_utils=True)
+    install_server_stub()
+    created_aiohttp_stub = install_aiohttp_stub(
+        lambda payload, status=200: _FakeResponse(payload, status=status)
+    )
+    _install_common_utils(package_name, config_data)
+    _load_real_guard_stack(package_name)
+    module = _load_module(package_name, "api/config_routes.py", "api.config_routes")
+    cleanup_optional_module("aiohttp", created_aiohttp_stub)
+    return module
+
+
+def _load_tunnel_routes_real_guard_module():
+    package_name = "dist_tunnel_auth_integration_testpkg"
+    config_data = {
+        "workers": [],
+        "master": {"host": "127.0.0.1"},
+        "settings": {"distributed_api_token": AUTH_TOKEN},
+        "tunnel": {},
+    }
+    bootstrap_test_package(package_name, with_api=True, with_utils=True)
+    install_server_stub()
+    created_aiohttp_stub = install_aiohttp_stub(
+        lambda payload, status=200: _FakeResponse(payload, status=status)
+    )
+    _install_common_utils(package_name, config_data)
+    _load_real_guard_stack(package_name)
+
+    cloudflare_module = types.ModuleType(f"{package_name}.utils.cloudflare")
+    cloudflare_module.cloudflare_tunnel_manager = types.SimpleNamespace(
+        get_status=lambda: {"active": False},
+        start_tunnel=lambda: {"active": True},
+        stop_tunnel=lambda: {"active": False},
+    )
+    sys.modules[f"{package_name}.utils.cloudflare"] = cloudflare_module
+
+    module = _load_module(package_name, "api/tunnel_routes.py", "api.tunnel_routes")
+    cleanup_optional_module("aiohttp", created_aiohttp_stub)
+    return module
+
+
+def _load_usdu_routes_real_guard_module():
+    package_name = "dist_usdu_auth_integration_testpkg"
+    config_data = {"settings": {"distributed_api_token": AUTH_TOKEN}}
+    bootstrap_test_package(package_name, with_api=True, with_utils=True, with_upscale=True)
+    install_server_stub()
+    created_aiohttp_stub = install_aiohttp_stub(
+        lambda payload, status=200: _FakeResponse(payload, status=status)
+    )
+    _install_common_utils(package_name, config_data)
+    _load_real_guard_stack(package_name)
+
+    usdu_management_module = types.ModuleType(f"{package_name}.utils.usdu_management")
+    usdu_management_module.MAX_PAYLOAD_SIZE = 1024 * 1024
+    sys.modules[f"{package_name}.utils.usdu_management"] = usdu_management_module
+
+    prompt_server_holder = {
+        "value": types.SimpleNamespace(
+            distributed_tile_jobs_lock=asyncio.Lock(),
+            distributed_pending_tile_jobs={},
+        )
+    }
+
+    job_store_module = types.ModuleType(f"{package_name}.upscale.job_store")
+    job_store_module.ensure_tile_jobs_initialized = lambda: prompt_server_holder["value"]
+    job_store_module.init_dynamic_job = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.upscale.job_store"] = job_store_module
+
+    job_models_module = types.ModuleType(f"{package_name}.upscale.job_models")
+
+    class _BaseJobState:
+        pass
+
+    class _ImageJobState(_BaseJobState):
+        def __init__(self):
+            self.mode = "dynamic"
+            self.queue = asyncio.Queue()
+            self.pending_images = asyncio.Queue()
+            self.completed_images = {}
+            self.worker_status = {}
+            self.assigned_to_workers = {}
+
+    class _TileJobState(_BaseJobState):
+        def __init__(self):
+            self.mode = "static"
+            self.queue = asyncio.Queue()
+            self.pending_tasks = asyncio.Queue()
+            self.completed_tasks = {}
+            self.worker_status = {}
+            self.assigned_to_workers = {}
+            self.batched_static = False
+
+    job_models_module.BaseJobState = _BaseJobState
+    job_models_module.ImageJobState = _ImageJobState
+    job_models_module.TileJobState = _TileJobState
+    sys.modules[f"{package_name}.upscale.job_models"] = job_models_module
+
+    parsers_module = types.ModuleType(f"{package_name}.upscale.payload_parsers")
+    parsers_module.parse_tiles_from_form = lambda _data: []
+    sys.modules[f"{package_name}.upscale.payload_parsers"] = parsers_module
+
+    module = _load_module(package_name, "api/usdu_routes.py", "api.usdu_routes")
+    module._prompt_server_holder = prompt_server_holder
+    module._ImageJobState = _ImageJobState
+    cleanup_optional_module("aiohttp", created_aiohttp_stub)
+    return module
+
+
+def _install_torch_stub_if_missing():
+    if "torch" in sys.modules:
+        return False
+
+    torch_module = types.ModuleType("torch")
+    torch_module.cuda = types.SimpleNamespace(
+        is_available=lambda: False,
+        empty_cache=lambda: None,
+        ipc_collect=lambda: None,
+    )
+    sys.modules["torch"] = torch_module
+    return True
+
+
+def _load_worker_routes_real_guard_module():
+    package_name = "dist_worker_auth_integration_testpkg"
+    config_data = {"settings": {"distributed_api_token": AUTH_TOKEN}, "workers": []}
+    bootstrap_test_package(package_name, with_api=True, with_utils=True, with_workers=True)
+    install_server_stub()
+    created_aiohttp_stub = install_aiohttp_stub(
+        lambda payload, status=200: _FakeResponse(payload, status=status)
+    )
+    created_torch_stub = _install_torch_stub_if_missing()
+
+    _install_common_utils(package_name, config_data)
+    _load_real_guard_stack(package_name)
+
+    constants_module = types.ModuleType(f"{package_name}.utils.constants")
+    constants_module.CHUNK_SIZE = 65536
+    sys.modules[f"{package_name}.utils.constants"] = constants_module
+
+    workers_module = types.ModuleType(f"{package_name}.workers")
+    workers_module.get_worker_manager = lambda: types.SimpleNamespace(
+        processes={},
+        save_processes=lambda: None,
+        is_process_running=lambda _pid: False,
+    )
+    sys.modules[f"{package_name}.workers"] = workers_module
+
+    detection_module = types.ModuleType(f"{package_name}.workers.detection")
+    detection_module.get_machine_id = lambda: "machine-id"
+    detection_module.is_docker_environment = lambda: False
+    detection_module.is_runpod_environment = lambda: False
+    sys.modules[f"{package_name}.workers.detection"] = detection_module
+
+    async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
+
+    class _PromptValidationError(Exception):
+        def __init__(self, message="validation error"):
+            super().__init__(message)
+            self.validation_error = {}
+            self.node_errors = {}
+
+    async def _queue_prompt_payload(*_args, **_kwargs):
+        return "prompt-id"
+
+    async_helpers_module.PromptValidationError = _PromptValidationError
+    async_helpers_module.queue_prompt_payload = _queue_prompt_payload
+    sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
+
+    module = _load_module(package_name, "api/worker_routes.py", "api.worker_routes")
+    cleanup_optional_module("aiohttp", created_aiohttp_stub)
+    if created_torch_stub:
+        sys.modules.pop("torch", None)
+    return module
+
+
+def _load_job_routes_real_guard_module():
+    package_name = "dist_job_auth_integration_testpkg"
+    config_data = {"settings": {"distributed_api_token": AUTH_TOKEN}}
+    bootstrap_test_package(package_name, with_api=True, with_utils=True)
+    install_server_stub()
+    created_aiohttp_stub = install_aiohttp_stub(
+        lambda payload, status=200: _FakeResponse(payload, status=status)
+    )
+    created_torch_stub = _install_torch_stub_if_missing()
+
+    _install_common_utils(package_name, config_data)
+    _load_real_guard_stack(package_name)
+
+    image_module = types.ModuleType(f"{package_name}.utils.image")
+    image_module.pil_to_tensor = lambda _img: None
+    image_module.ensure_contiguous = lambda tensor: tensor
+    sys.modules[f"{package_name}.utils.image"] = image_module
+
+    constants_module = types.ModuleType(f"{package_name}.utils.constants")
+    constants_module.JOB_INIT_GRACE_PERIOD = 0.1
+    constants_module.MEMORY_CLEAR_DELAY = 0.0
+    sys.modules[f"{package_name}.utils.constants"] = constants_module
+
+    runtime_state_module = types.ModuleType(f"{package_name}.utils.runtime_state")
+    runtime_state = types.SimpleNamespace(
+        distributed_jobs_lock=asyncio.Lock(),
+        distributed_pending_jobs={},
+    )
+    runtime_state_module.ensure_distributed_runtime_state = lambda: runtime_state
+    sys.modules[f"{package_name}.utils.runtime_state"] = runtime_state_module
+
+    orchestration_module = types.ModuleType(f"{package_name}.api.queue_orchestration")
+    orchestration_module.orchestrate_distributed_execution = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.api.queue_orchestration"] = orchestration_module
+
+    queue_request_module = types.ModuleType(f"{package_name}.api.queue_request")
+    queue_request_module.parse_queue_request_payload = lambda data: data
+    sys.modules[f"{package_name}.api.queue_request"] = queue_request_module
+
+    module = _load_module(package_name, "api/job_routes.py", "api.job_routes")
+    cleanup_optional_module("aiohttp", created_aiohttp_stub)
+    if created_torch_stub:
+        sys.modules.pop("torch", None)
+    return module
+
+
+request_guards_module, schemas_module = _load_request_guards_only_module()
+config_routes_module = _load_config_routes_real_guard_module()
+tunnel_routes_module = _load_tunnel_routes_real_guard_module()
+usdu_routes_module = _load_usdu_routes_real_guard_module()
+worker_routes_module = _load_worker_routes_real_guard_module()
+job_routes_module = _load_job_routes_real_guard_module()
+
+
+class AuthGuardContractTests(unittest.IsolatedAsyncioTestCase):
+    async def test_is_authorized_request_missing_token_fails(self):
+        request = _FakeRequest(headers={})
+        config = {"settings": {"distributed_api_token": AUTH_TOKEN}}
+        self.assertFalse(schemas_module.is_authorized_request(request, config))
+
+    async def test_is_authorized_request_matching_token_succeeds(self):
+        request = _FakeRequest(headers={schemas_module.AUTH_HEADER_NAME: AUTH_TOKEN})
+        config = {"settings": {"distributed_api_token": AUTH_TOKEN}}
+        self.assertTrue(schemas_module.is_authorized_request(request, config))
+
+    async def test_authorization_error_or_none_returns_403_for_invalid_token(self):
+        request = _FakeRequest(headers={schemas_module.AUTH_HEADER_NAME: "wrong-token"})
+        response = await request_guards_module.authorization_error_or_none(request)
+        self.assertIsNotNone(response)
+        self.assertEqual(response.status, 403)
+
+    async def test_authorization_error_or_none_returns_none_for_valid_token(self):
+        request = _FakeRequest(headers={schemas_module.AUTH_HEADER_NAME: AUTH_TOKEN})
+        response = await request_guards_module.authorization_error_or_none(request)
+        self.assertIsNone(response)
+
+
+class RouteAuthIntegrationTests(unittest.IsolatedAsyncioTestCase):
+    def _auth_headers(self):
+        return {schemas_module.AUTH_HEADER_NAME: AUTH_TOKEN}
+
+    async def test_config_endpoint_enforces_real_guard(self):
+        unauthorized = await config_routes_module.get_config_endpoint(_FakeRequest(headers={}))
+        self.assertEqual(unauthorized.status, 403)
+
+        authorized = await config_routes_module.get_config_endpoint(_FakeRequest(headers=self._auth_headers()))
+        self.assertEqual(authorized.status, 200)
+        self.assertIn("settings", authorized.payload)
+
+    async def test_tunnel_endpoint_enforces_real_guard(self):
+        unauthorized = await tunnel_routes_module.tunnel_status_endpoint(_FakeRequest(headers={}))
+        self.assertEqual(unauthorized.status, 403)
+
+        authorized = await tunnel_routes_module.tunnel_status_endpoint(_FakeRequest(headers=self._auth_headers()))
+        self.assertEqual(authorized.status, 200)
+        self.assertEqual(authorized.payload.get("status"), "success")
+
+    async def test_usdu_endpoint_enforces_real_guard(self):
+        unauthorized = await usdu_routes_module.job_status_endpoint(_FakeRequest(headers={}))
+        self.assertEqual(unauthorized.status, 403)
+
+        authorized = await usdu_routes_module.job_status_endpoint(
+            _FakeRequest(headers=self._auth_headers(), query={"multi_job_id": "job-1"})
+        )
+        self.assertEqual(authorized.status, 200)
+        self.assertIn("ready", authorized.payload)
+
+    async def test_worker_endpoint_enforces_real_guard(self):
+        worker_routes_module._collect_network_info_sync = lambda: {
+            "interfaces": [],
+            "recommended_host": "127.0.0.1",
+        }
+        unauthorized = await worker_routes_module.get_network_info_endpoint(_FakeRequest(headers={}))
+        self.assertEqual(unauthorized.status, 403)
+
+        authorized = await worker_routes_module.get_network_info_endpoint(
+            _FakeRequest(headers=self._auth_headers())
+        )
+        self.assertEqual(authorized.status, 200)
+        self.assertEqual(authorized.payload.get("status"), "success")
+
+    async def test_job_endpoint_enforces_real_guard(self):
+        unauthorized = await job_routes_module.prepare_job_endpoint(
+            _FakeRequest(headers={}, json_payload={"multi_job_id": "job-1"})
+        )
+        self.assertEqual(unauthorized.status, 403)
+
+        authorized = await job_routes_module.prepare_job_endpoint(
+            _FakeRequest(headers=self._auth_headers(), json_payload={"multi_job_id": "job-1"})
+        )
+        self.assertEqual(authorized.status, 200)
+        self.assertEqual(authorized.payload.get("status"), "success")
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/api/test_config_routes.py

@@ -3,9 +3,19 @@ import importlib.util
 import sys
 import types
 import unittest
+from contextlib import asynccontextmanager
 from pathlib import Path
 from unittest.mock import patch
 
+from tests.api.harness import (
+    bootstrap_test_package,
+    cleanup_optional_module,
+    install_aiohttp_stub,
+    install_endpoint_policy_passthrough_stub,
+    install_request_guards_stub,
+    install_server_stub,
+)
+
 
 class _FakeResponse:
     def __init__(self, payload, status=200):
@@ -25,45 +35,18 @@ def _load_config_routes_module():
     module_path = Path(__file__).resolve().parents[2] / "api" / "config_routes.py"
     package_name = "dist_api_config_testpkg"
 
-    for mod_name in list(sys.modules):
-        if mod_name == package_name or mod_name.startswith(f"{package_name}."):
-            del sys.modules[mod_name]
+    bootstrap_test_package(package_name, with_api=True, with_utils=True)
 
-    root_pkg = types.ModuleType(package_name)
-    root_pkg.__path__ = []
-    sys.modules[package_name] = root_pkg
+    schemas_module = types.ModuleType(f"{package_name}.api.schemas")
+    schemas_module.is_authorized_request = lambda _request, _config: True
+    sys.modules[f"{package_name}.api.schemas"] = schemas_module
 
-    api_pkg = types.ModuleType(f"{package_name}.api")
-    api_pkg.__path__ = []
-    sys.modules[f"{package_name}.api"] = api_pkg
-
-    utils_pkg = types.ModuleType(f"{package_name}.utils")
-    utils_pkg.__path__ = []
-    sys.modules[f"{package_name}.utils"] = utils_pkg
-
-    created_aiohttp_stub = False
-    if "aiohttp" not in sys.modules:
-        created_aiohttp_stub = True
-        aiohttp_module = types.ModuleType("aiohttp")
-        aiohttp_module.web = types.SimpleNamespace(
-            json_response=lambda payload, status=200: _FakeResponse(payload, status=status)
-        )
-        sys.modules["aiohttp"] = aiohttp_module
-
-    class _Routes:
-        def get(self, _path):
-            def _decorator(fn):
-                return fn
-            return _decorator
-
-        def post(self, _path):
-            def _decorator(fn):
-                return fn
-            return _decorator
-
-    server_module = types.ModuleType("server")
-    server_module.PromptServer = types.SimpleNamespace(instance=types.SimpleNamespace(routes=_Routes()))
-    sys.modules["server"] = server_module
+    install_request_guards_stub(package_name)
+    install_endpoint_policy_passthrough_stub(package_name)
+    created_aiohttp_stub = install_aiohttp_stub(
+        lambda payload, status=200: _FakeResponse(payload, status=status)
+    )
+    install_server_stub()
 
     logging_module = types.ModuleType(f"{package_name}.utils.logging")
     logging_module.debug_log = lambda *_args, **_kwargs: None
@@ -73,7 +56,16 @@ def _load_config_routes_module():
     network_module = types.ModuleType(f"{package_name}.utils.network")
 
     async def _handle_api_error(_request, error, status=500):
-        return _FakeResponse({"status": "error", "message": str(error)}, status=status)
+        if isinstance(error, list):
+            message = "; ".join(str(item) for item in error)
+            error_payload = [str(item) for item in error]
+        else:
+            message = str(error)
+            error_payload = str(error)
+        return _FakeResponse(
+            {"status": "error", "error": error_payload, "message": message},
+            status=status,
+        )
 
     network_module.handle_api_error = _handle_api_error
     network_module.normalize_host = lambda value: value
@@ -89,6 +81,12 @@ def _load_config_routes_module():
     config_module = types.ModuleType(f"{package_name}.utils.config")
     config_module.load_config = lambda: copy.deepcopy(default_config)
     config_module.save_config = lambda _cfg: True
+
+    @asynccontextmanager
+    async def _config_transaction():
+        yield config_module.load_config()
+
+    config_module.config_transaction = _config_transaction
     sys.modules[f"{package_name}.utils.config"] = config_module
 
     spec = importlib.util.spec_from_file_location(f"{package_name}.api.config_routes", module_path)
@@ -96,8 +94,7 @@ def _load_config_routes_module():
     assert spec is not None and spec.loader is not None
     spec.loader.exec_module(module)
 
-    if created_aiohttp_stub:
-        sys.modules.pop("aiohttp", None)
+    cleanup_optional_module("aiohttp", created_aiohttp_stub)
 
     return module
 
@@ -118,9 +115,7 @@ class ConfigRoutesTests(unittest.IsolatedAsyncioTestCase):
 
     async def test_update_config_valid_field_persists(self):
         cfg = {"workers": [], "master": {}, "settings": {"debug": False}, "tunnel": {}}
-        with patch.object(config_routes, "load_config", return_value=cfg), patch.object(
-            config_routes, "save_config", return_value=True
-        ):
+        with patch.object(config_routes, "load_config", return_value=cfg):
             response = await config_routes.update_config_endpoint(_FakeRequest({"debug": True}))
 
         self.assertEqual(response.status, 200)

tests/api/test_job_routes.py

@@ -11,6 +11,14 @@ from unittest.mock import AsyncMock, patch
 import numpy as np
 import torch
 
+from tests.api.harness import (
+    bootstrap_test_package,
+    cleanup_optional_module,
+    install_aiohttp_stub,
+    install_request_guards_stub,
+    install_server_stub,
+)
+
 
 class _FakeResponse:
     def __init__(self, payload, status=200):
@@ -19,8 +27,9 @@ class _FakeResponse:
 
 
 class _FakeRequest:
-    def __init__(self, payload):
+    def __init__(self, payload, headers=None):
         self._payload = payload
+        self.headers = headers or {}
 
     async def json(self):
         return self._payload
@@ -30,53 +39,19 @@ def _load_job_routes_module():
     module_path = Path(__file__).resolve().parents[2] / "api" / "job_routes.py"
     package_name = "dist_api_queue_testpkg"
 
-    # Reset package namespace to avoid stale module state across test runs.
-    for mod_name in list(sys.modules):
-        if mod_name == package_name or mod_name.startswith(f"{package_name}."):
-            del sys.modules[mod_name]
-
-    root_pkg = types.ModuleType(package_name)
-    root_pkg.__path__ = []
-    sys.modules[package_name] = root_pkg
-
-    api_pkg = types.ModuleType(f"{package_name}.api")
-    api_pkg.__path__ = []
-    sys.modules[f"{package_name}.api"] = api_pkg
-
-    utils_pkg = types.ModuleType(f"{package_name}.utils")
-    utils_pkg.__path__ = []
-    sys.modules[f"{package_name}.utils"] = utils_pkg
+    bootstrap_test_package(package_name, with_api=True, with_utils=True)
 
     # aiohttp.web stub
-    created_aiohttp_stub = False
-    if "aiohttp" not in sys.modules:
-        created_aiohttp_stub = True
-        aiohttp_module = types.ModuleType("aiohttp")
-        aiohttp_module.web = types.SimpleNamespace(
-            json_response=lambda payload, status=200: _FakeResponse(payload, status=status)
-        )
-        sys.modules["aiohttp"] = aiohttp_module
+    created_aiohttp_stub = install_aiohttp_stub(
+        lambda payload, status=200: _FakeResponse(payload, status=status)
+    )
 
     # server module stub with route decorators
-    class _Routes:
-        def get(self, _path):
-            def _decorator(fn):
-                return fn
-            return _decorator
-
-        def post(self, _path):
-            def _decorator(fn):
-                return fn
-            return _decorator
-
     prompt_server_instance = types.SimpleNamespace(
-        routes=_Routes(),
         distributed_jobs_lock=None,
         distributed_pending_jobs={},
     )
-    server_module = types.ModuleType("server")
-    server_module.PromptServer = types.SimpleNamespace(instance=prompt_server_instance)
-    sys.modules["server"] = server_module
+    install_server_stub(prompt_server_instance)
 
     # torch stub (only needed to satisfy import)
     created_torch_stub = False
@@ -157,6 +132,47 @@ def _load_job_routes_module():
     constants_module.JOB_INIT_GRACE_PERIOD = 10.0
     sys.modules[f"{package_name}.utils.constants"] = constants_module
 
+    config_module = types.ModuleType(f"{package_name}.utils.config")
+    config_module.load_config = lambda: {"settings": {}}
+    sys.modules[f"{package_name}.utils.config"] = config_module
+
+    runtime_state_module = types.ModuleType(f"{package_name}.utils.runtime_state")
+
+    def _ensure_distributed_runtime_state(server_instance=None):
+        ps = server_instance or prompt_server_instance
+        if not hasattr(ps, "distributed_pending_jobs"):
+            ps.distributed_pending_jobs = {}
+        if not hasattr(ps, "distributed_jobs_lock") or ps.distributed_jobs_lock is None:
+            ps.distributed_jobs_lock = asyncio.Lock()
+        if not hasattr(ps, "distributed_job_allowed_workers"):
+            ps.distributed_job_allowed_workers = {}
+        if not hasattr(ps, "distributed_pending_tile_jobs"):
+            ps.distributed_pending_tile_jobs = {}
+        if not hasattr(ps, "distributed_tile_jobs_lock"):
+            ps.distributed_tile_jobs_lock = asyncio.Lock()
+        return types.SimpleNamespace(
+            distributed_pending_jobs=ps.distributed_pending_jobs,
+            distributed_jobs_lock=ps.distributed_jobs_lock,
+            distributed_job_allowed_workers=ps.distributed_job_allowed_workers,
+            distributed_pending_tile_jobs=ps.distributed_pending_tile_jobs,
+            distributed_tile_jobs_lock=ps.distributed_tile_jobs_lock,
+        )
+
+    runtime_state_module.ensure_distributed_runtime_state = _ensure_distributed_runtime_state
+    runtime_state_module.get_prompt_server_instance = lambda: prompt_server_instance
+    sys.modules[f"{package_name}.utils.runtime_state"] = runtime_state_module
+
+    schemas_module = types.ModuleType(f"{package_name}.api.schemas")
+    schemas_module.is_authorized_request = lambda _request, _config: True
+    schemas_module.require_bool_literal = (
+        lambda value, field_name="value": value
+        if isinstance(value, bool)
+        else (_ for _ in ()).throw(ValueError(f"Field '{field_name}' must be a boolean literal."))
+    )
+    sys.modules[f"{package_name}.api.schemas"] = schemas_module
+
+    install_request_guards_stub(package_name)
+
     async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
     async_helpers_module.queue_prompt_payload = AsyncMock(return_value="prompt_local")
     sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
@@ -172,7 +188,6 @@ def _load_job_routes_module():
         client_id: str
         delegate_master: object
         enabled_worker_ids: list
-        auto_prepare: bool
         trace_execution_id: object
 
     def _parse_queue_request_payload(data):
@@ -193,7 +208,6 @@ def _load_job_routes_module():
             client_id=client_id,
             delegate_master=data.get("delegate_master"),
             enabled_worker_ids=enabled,
-            auto_prepare=bool(data.get("auto_prepare", True)),
             trace_execution_id=data.get("trace_execution_id"),
         )
 
@@ -206,13 +220,11 @@ def _load_job_routes_module():
     assert spec is not None and spec.loader is not None
     spec.loader.exec_module(module)
 
-    if created_aiohttp_stub:
-        sys.modules.pop("aiohttp", None)
-    if created_torch_stub:
-        sys.modules.pop("torch", None)
+    cleanup_optional_module("aiohttp", created_aiohttp_stub)
+    cleanup_optional_module("torch", created_torch_stub)
     if created_pil_stub:
-        sys.modules.pop("PIL.Image", None)
-        sys.modules.pop("PIL", None)
+        cleanup_optional_module("PIL.Image", True)
+        cleanup_optional_module("PIL", True)
 
     return module
 
@@ -239,7 +251,6 @@ class DistributedQueueEndpointTests(unittest.IsolatedAsyncioTestCase):
 
         self.assertEqual(response.status, 200)
         self.assertEqual(response.payload.get("prompt_id"), "prompt_123")
-        self.assertTrue(response.payload.get("auto_prepare_supported"))
 
     async def test_distributed_queue_missing_prompt_returns_400(self):
         request = _FakeRequest(
@@ -276,8 +287,10 @@ class JobCompleteAudioPayloadTests(unittest.IsolatedAsyncioTestCase):
 
     async def test_job_complete_accepts_audio_payload(self):
         queue = asyncio.Queue()
-        job_routes.prompt_server.distributed_jobs_lock = asyncio.Lock()
-        job_routes.prompt_server.distributed_pending_jobs = {"job-1": queue}
+        prompt_server = job_routes.server.PromptServer.instance
+        prompt_server.distributed_jobs_lock = asyncio.Lock()
+        prompt_server.distributed_pending_jobs = {"job-1": queue}
+        prompt_server.distributed_job_allowed_workers = {"job-1": {"worker-1"}}
         request = _FakeRequest(
             {
                 "job_id": "job-1",
@@ -300,6 +313,28 @@ class JobCompleteAudioPayloadTests(unittest.IsolatedAsyncioTestCase):
         self.assertEqual(queued["audio"]["sample_rate"], 44100)
         self.assertEqual(tuple(queued["audio"]["waveform"].shape), (1, 2, 4))
 
+    async def test_job_complete_rejects_worker_not_in_job_allowlist(self):
+        queue = asyncio.Queue()
+        prompt_server = job_routes.server.PromptServer.instance
+        prompt_server.distributed_jobs_lock = asyncio.Lock()
+        prompt_server.distributed_pending_jobs = {"job-allow": queue}
+        prompt_server.distributed_job_allowed_workers = {"job-allow": {"worker-expected"}}
+        request = _FakeRequest(
+            {
+                "job_id": "job-allow",
+                "worker_id": "worker-unexpected",
+                "batch_idx": 0,
+                "image": "data:image/png;base64,AAAA",
+                "is_last": True,
+            }
+        )
+
+        with patch.object(job_routes, "_decode_canonical_png_tensor", return_value="tensor-data"):
+            response = await job_routes.job_complete_endpoint(request)
+
+        self.assertEqual(response.status, 403)
+        self.assertIn("unauthorized", response.payload.get("message", "").lower())
+
     def test_decode_audio_payload_rejects_bad_shape(self):
         bad = {
             "sample_rate": 44100,

tests/api/test_media_sync.py

@@ -4,30 +4,29 @@ import types
 import unittest
 from pathlib import Path
 
+from tests.api.harness import (
+    bootstrap_test_package,
+    cleanup_optional_module,
+    install_aiohttp_stub,
+)
+
+
+class _AiohttpResponse:
+    def __init__(self, payload, status=200):
+        self.payload = payload
+        self.status = status
+
 
 def _load_media_sync_module():
     module_path = Path(__file__).resolve().parents[2] / "api" / "orchestration" / "media_sync.py"
     package_name = "dist_ms_testpkg"
 
-    for mod_name in list(sys.modules):
-        if mod_name == package_name or mod_name.startswith(f"{package_name}."):
-            del sys.modules[mod_name]
-
-    root_pkg = types.ModuleType(package_name)
-    root_pkg.__path__ = []
-    sys.modules[package_name] = root_pkg
-
-    api_pkg = types.ModuleType(f"{package_name}.api")
-    api_pkg.__path__ = []
-    sys.modules[f"{package_name}.api"] = api_pkg
-
-    orch_pkg = types.ModuleType(f"{package_name}.api.orchestration")
-    orch_pkg.__path__ = []
-    sys.modules[f"{package_name}.api.orchestration"] = orch_pkg
-
-    utils_pkg = types.ModuleType(f"{package_name}.utils")
-    utils_pkg.__path__ = []
-    sys.modules[f"{package_name}.utils"] = utils_pkg
+    bootstrap_test_package(
+        package_name,
+        with_api=True,
+        with_utils=True,
+        with_orchestration=True,
+    )
 
     logging_module = types.ModuleType(f"{package_name}.utils.logging")
     logging_module.debug_log = lambda *_args, **_kwargs: None
@@ -48,22 +47,17 @@ def _load_media_sync_module():
     trace_module.trace_info = lambda *_args, **_kwargs: None
     sys.modules[f"{package_name}.utils.trace_logger"] = trace_module
 
-    created_aiohttp_stub = False
-    if "aiohttp" not in sys.modules:
-        created_aiohttp_stub = True
-        aiohttp_module = types.ModuleType("aiohttp")
+    auth_module = types.ModuleType(f"{package_name}.utils.auth")
+    auth_module.distributed_auth_headers = lambda _config=None: {}
+    sys.modules[f"{package_name}.utils.auth"] = auth_module
 
-        class _ClientTimeout:
-            def __init__(self, total=None):
-                pass
+    config_module = types.ModuleType(f"{package_name}.utils.config")
+    config_module.load_config = lambda: {"settings": {}}
+    sys.modules[f"{package_name}.utils.config"] = config_module
 
-        class _FormData:
-            def add_field(self, *args, **kwargs):
-                pass
-
-        aiohttp_module.ClientTimeout = _ClientTimeout
-        aiohttp_module.FormData = _FormData
-        sys.modules["aiohttp"] = aiohttp_module
+    created_aiohttp_stub = install_aiohttp_stub(
+        lambda payload, status=200: _AiohttpResponse(payload, status=status)
+    )
 
     spec = importlib.util.spec_from_file_location(
         f"{package_name}.api.orchestration.media_sync",
@@ -73,8 +67,7 @@ def _load_media_sync_module():
     assert spec is not None and spec.loader is not None
     spec.loader.exec_module(module)
 
-    if created_aiohttp_stub:
-        sys.modules.pop("aiohttp", None)
+    cleanup_optional_module("aiohttp", created_aiohttp_stub)
 
     return module
 
@@ -260,3 +253,7 @@ class RewritePromptMediaInputsTests(unittest.TestCase):
 
 if __name__ == "__main__":
     unittest.main()
+    class _AiohttpResponse:
+        def __init__(self, payload, status=200):
+            self.payload = payload
+            self.status = status

tests/api/test_queue_orchestration.py

@@ -5,32 +5,22 @@ import types
 import unittest
 from pathlib import Path
 
+from tests.api.harness import (
+    bootstrap_test_package,
+    install_server_stub,
+)
+
 
 def _bootstrap_package(package_name):
-    for mod_name in list(sys.modules):
-        if mod_name == package_name or mod_name.startswith(f"{package_name}."):
-            del sys.modules[mod_name]
-
-    root_pkg = types.ModuleType(package_name)
-    root_pkg.__path__ = []
-    sys.modules[package_name] = root_pkg
-
-    api_pkg = types.ModuleType(f"{package_name}.api")
-    api_pkg.__path__ = []
-    sys.modules[f"{package_name}.api"] = api_pkg
-
-    orchestration_pkg = types.ModuleType(f"{package_name}.api.orchestration")
-    orchestration_pkg.__path__ = []
-    sys.modules[f"{package_name}.api.orchestration"] = orchestration_pkg
-
-    utils_pkg = types.ModuleType(f"{package_name}.utils")
-    utils_pkg.__path__ = []
-    sys.modules[f"{package_name}.utils"] = utils_pkg
+    bootstrap_test_package(
+        package_name,
+        with_api=True,
+        with_utils=True,
+        with_orchestration=True,
+    )
 
     prompt_server_instance = types.SimpleNamespace(address="127.0.0.1", port=8188)
-    server_module = types.ModuleType("server")
-    server_module.PromptServer = types.SimpleNamespace(instance=prompt_server_instance)
-    sys.modules["server"] = server_module
+    install_server_stub(prompt_server_instance)
 
     async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
 
@@ -65,9 +55,37 @@ def _bootstrap_package(package_name):
     trace_logger_module.trace_debug = lambda *_args, **_kwargs: None
     sys.modules[f"{package_name}.utils.trace_logger"] = trace_logger_module
 
+    runtime_state_module = types.ModuleType(f"{package_name}.utils.runtime_state")
+
+    def _ensure_distributed_runtime_state(server_instance=None):
+        ps = server_instance or prompt_server_instance
+        if not hasattr(ps, "distributed_pending_jobs"):
+            ps.distributed_pending_jobs = {}
+        if not hasattr(ps, "distributed_jobs_lock"):
+            ps.distributed_jobs_lock = asyncio.Lock()
+        if not hasattr(ps, "distributed_job_allowed_workers"):
+            ps.distributed_job_allowed_workers = {}
+        if not hasattr(ps, "distributed_pending_tile_jobs"):
+            ps.distributed_pending_tile_jobs = {}
+        if not hasattr(ps, "distributed_tile_jobs_lock"):
+            ps.distributed_tile_jobs_lock = asyncio.Lock()
+        return types.SimpleNamespace(
+            distributed_pending_jobs=ps.distributed_pending_jobs,
+            distributed_jobs_lock=ps.distributed_jobs_lock,
+            distributed_job_allowed_workers=ps.distributed_job_allowed_workers,
+            distributed_pending_tile_jobs=ps.distributed_pending_tile_jobs,
+            distributed_tile_jobs_lock=ps.distributed_tile_jobs_lock,
+        )
+
+    runtime_state_module.ensure_distributed_runtime_state = _ensure_distributed_runtime_state
+    runtime_state_module.get_prompt_server_instance = lambda: prompt_server_instance
+    sys.modules[f"{package_name}.utils.runtime_state"] = runtime_state_module
+
     schemas_module = types.ModuleType(f"{package_name}.api.schemas")
-    schemas_module.parse_positive_int = lambda value, default: int(value) if str(value).isdigit() and int(value) > 0 else default
-    schemas_module.parse_positive_float = (
+    schemas_module.coerce_positive_int = (
+        lambda value, default: int(value) if str(value).isdigit() and int(value) > 0 else default
+    )
+    schemas_module.coerce_positive_float = (
         lambda value, default: float(value) if isinstance(value, (int, float, str)) and float(value) > 0 else default
     )
     sys.modules[f"{package_name}.api.schemas"] = schemas_module

tests/api/test_tunnel_routes.py

@@ -4,6 +4,14 @@ import types
 import unittest
 from pathlib import Path
 
+from tests.api.harness import (
+    bootstrap_test_package,
+    install_aiohttp_stub,
+    install_endpoint_policy_passthrough_stub,
+    install_request_guards_stub,
+    install_server_stub,
+)
+
 
 class _FakeResponse:
     def __init__(self, payload, status=200):
@@ -15,38 +23,9 @@ def _load_tunnel_routes_module():
     module_path = Path(__file__).resolve().parents[2] / "api" / "tunnel_routes.py"
     package_name = "dist_tunnel_routes_testpkg"
 
-    for mod_name in list(sys.modules):
-        if mod_name == package_name or mod_name.startswith(f"{package_name}."):
-            del sys.modules[mod_name]
-
-    root_pkg = types.ModuleType(package_name)
-    root_pkg.__path__ = []
-    sys.modules[package_name] = root_pkg
-
-    api_pkg = types.ModuleType(f"{package_name}.api")
-    api_pkg.__path__ = []
-    sys.modules[f"{package_name}.api"] = api_pkg
-
-    utils_pkg = types.ModuleType(f"{package_name}.utils")
-    utils_pkg.__path__ = []
-    sys.modules[f"{package_name}.utils"] = utils_pkg
-
-    class _Routes:
-        def get(self, _path):
-            return lambda fn: fn
-
-        def post(self, _path):
-            return lambda fn: fn
-
-    server_module = types.ModuleType("server")
-    server_module.PromptServer = types.SimpleNamespace(instance=types.SimpleNamespace(routes=_Routes()))
-    sys.modules["server"] = server_module
-
-    aiohttp_module = types.ModuleType("aiohttp")
-    aiohttp_module.web = types.SimpleNamespace(
-        json_response=lambda payload, status=200: _FakeResponse(payload, status=status)
-    )
-    sys.modules["aiohttp"] = aiohttp_module
+    bootstrap_test_package(package_name, with_api=True, with_utils=True)
+    install_server_stub()
+    install_aiohttp_stub(lambda payload, status=200: _FakeResponse(payload, status=status))
 
     cloudflare_module = types.ModuleType(f"{package_name}.utils.cloudflare")
 
@@ -80,6 +59,9 @@ def _load_tunnel_routes_module():
     network_module.handle_api_error = _handle_api_error
     sys.modules[f"{package_name}.utils.network"] = network_module
 
+    install_request_guards_stub(package_name)
+    install_endpoint_policy_passthrough_stub(package_name)
+
     spec = importlib.util.spec_from_file_location(f"{package_name}.api.tunnel_routes", module_path)
     module = importlib.util.module_from_spec(spec)
     assert spec is not None and spec.loader is not None

tests/api/test_usdu_routes.py

@@ -9,6 +9,14 @@ from pathlib import Path
 
 from PIL import Image
 
+from tests.api.harness import (
+    bootstrap_test_package,
+    cleanup_optional_module,
+    install_aiohttp_stub,
+    install_request_guards_stub,
+    install_server_stub,
+)
+
 
 class _FakeResponse:
     def __init__(self, payload, status=200):
@@ -30,43 +38,30 @@ class _FakeRequest:
         return self._post_payload
 
 
-class _Routes:
-    def post(self, _path):
-        def _decorator(fn):
-            return fn
-
-        return _decorator
-
-    def get(self, _path):
-        def _decorator(fn):
-            return fn
-
-        return _decorator
-
-
 def _load_usdu_routes_module():
     module_path = Path(__file__).resolve().parents[2] / "api" / "usdu_routes.py"
     package_name = "dist_api_usdu_testpkg"
 
-    for mod_name in list(sys.modules):
-        if mod_name == package_name or mod_name.startswith(f"{package_name}."):
-            del sys.modules[mod_name]
+    bootstrap_test_package(package_name, with_api=True, with_utils=True, with_upscale=True)
 
-    root_pkg = types.ModuleType(package_name)
-    root_pkg.__path__ = []
-    sys.modules[package_name] = root_pkg
+    schemas_module = types.ModuleType(f"{package_name}.api.schemas")
 
-    api_pkg = types.ModuleType(f"{package_name}.api")
-    api_pkg.__path__ = []
-    sys.modules[f"{package_name}.api"] = api_pkg
+    def _require_bool_literal(value, field_name: str = "value"):
+        if isinstance(value, bool):
+            return value
+        if isinstance(value, str):
+            normalized = value.strip().lower()
+            if normalized == "true":
+                return True
+            if normalized == "false":
+                return False
+        raise ValueError(f"Field '{field_name}' must be a boolean literal.")
 
-    upscale_pkg = types.ModuleType(f"{package_name}.upscale")
-    upscale_pkg.__path__ = []
-    sys.modules[f"{package_name}.upscale"] = upscale_pkg
+    schemas_module.require_bool_literal = _require_bool_literal
+    schemas_module.is_authorized_request = lambda _request, _config: True
+    sys.modules[f"{package_name}.api.schemas"] = schemas_module
 
-    utils_pkg = types.ModuleType(f"{package_name}.utils")
-    utils_pkg.__path__ = []
-    sys.modules[f"{package_name}.utils"] = utils_pkg
+    install_request_guards_stub(package_name)
 
     prompt_server_holder = {
         "value": types.SimpleNamespace(
@@ -75,23 +70,23 @@ def _load_usdu_routes_module():
         )
     }
 
-    created_aiohttp_stub = False
-    if "aiohttp" not in sys.modules:
-        created_aiohttp_stub = True
-        aiohttp_module = types.ModuleType("aiohttp")
-        aiohttp_module.web = types.SimpleNamespace(
-            json_response=lambda payload, status=200: _FakeResponse(payload, status=status)
-        )
-        sys.modules["aiohttp"] = aiohttp_module
-
-    server_module = types.ModuleType("server")
-    server_module.PromptServer = types.SimpleNamespace(instance=types.SimpleNamespace(routes=_Routes()))
-    sys.modules["server"] = server_module
+    created_aiohttp_stub = install_aiohttp_stub(
+        lambda payload, status=200: _FakeResponse(payload, status=status)
+    )
+    install_server_stub()
 
     logging_module = types.ModuleType(f"{package_name}.utils.logging")
     logging_module.debug_log = lambda *_args, **_kwargs: None
     sys.modules[f"{package_name}.utils.logging"] = logging_module
 
+    config_module = types.ModuleType(f"{package_name}.utils.config")
+    config_module.load_config = lambda: {"settings": {}}
+    sys.modules[f"{package_name}.utils.config"] = config_module
+
+    usdu_management_module = types.ModuleType(f"{package_name}.utils.usdu_management")
+    usdu_management_module.MAX_PAYLOAD_SIZE = 1024
+    sys.modules[f"{package_name}.utils.usdu_management"] = usdu_management_module
+
     network_module = types.ModuleType(f"{package_name}.utils.network")
 
     async def _handle_api_error(_request, error, status=500):
@@ -174,8 +169,7 @@ def _load_usdu_routes_module():
     assert spec is not None and spec.loader is not None
     spec.loader.exec_module(module)
 
-    if created_aiohttp_stub:
-        sys.modules.pop("aiohttp", None)
+    cleanup_optional_module("aiohttp", created_aiohttp_stub)
 
     module.web = types.SimpleNamespace(
         json_response=lambda payload, status=200: _FakeResponse(payload, status=status)
@@ -238,7 +232,8 @@ class USDURoutesTests(unittest.IsolatedAsyncioTestCase):
         response = await usdu_routes.request_image_endpoint(request)
 
         self.assertEqual(response.status, 200)
-        self.assertEqual(response.payload.get("image_idx"), 7)
+        self.assertEqual(response.payload.get("kind"), "image")
+        self.assertEqual(response.payload.get("task_idx"), 7)
         self.assertEqual(response.payload.get("estimated_remaining"), 0)
         self.assertEqual(job_data.assigned_to_workers["worker-a"], [7])
         self.assertIn("worker-a", job_data.worker_status)
@@ -254,38 +249,11 @@ class USDURoutesTests(unittest.IsolatedAsyncioTestCase):
         response = await usdu_routes.request_image_endpoint(request)
 
         self.assertEqual(response.status, 200)
-        self.assertEqual(response.payload.get("tile_idx"), 4)
+        self.assertEqual(response.payload.get("kind"), "tile")
+        self.assertEqual(response.payload.get("task_idx"), 4)
         self.assertTrue(response.payload.get("batched_static"))
         self.assertEqual(job_data.assigned_to_workers["worker-a"], [4])
 
-    async def test_init_list_queue_initializes_pending_items(self):
-        request = _FakeRequest(
-            json_payload={
-                "multi_job_id": "list-job-1",
-                "list_size": 3,
-                "enabled_workers": ["worker-a", "worker-b"],
-            }
-        )
-        response = await usdu_routes.init_list_queue_endpoint(request)
-
-        self.assertEqual(response.status, 200)
-        self.assertEqual(response.payload.get("status"), "success")
-        self.assertEqual(response.payload.get("remaining"), 3)
-
-    async def test_request_list_item_assigns_next_index(self):
-        prompt_server = usdu_routes._prompt_server_holder["value"]
-        job_data = usdu_routes._ImageJobState("list-job-2")
-        await job_data.pending_images.put(11)
-        prompt_server.distributed_pending_tile_jobs["list-job-2"] = job_data
-
-        request = _FakeRequest(json_payload={"worker_id": "worker-a", "multi_job_id": "list-job-2"})
-        response = await usdu_routes.request_list_item_endpoint(request)
-
-        self.assertEqual(response.status, 200)
-        self.assertEqual(response.payload.get("item_idx"), 11)
-        self.assertEqual(response.payload.get("estimated_remaining"), 0)
-        self.assertEqual(job_data.assigned_to_workers["worker-a"], [11])
-
     async def test_submit_tiles_completion_signal_enqueues_last_marker(self):
         prompt_server = usdu_routes._prompt_server_holder["value"]
         job_data = usdu_routes._TileJobState("job-4")

tests/api/test_worker_routes.py

@@ -8,6 +8,14 @@ from collections import deque
 from pathlib import Path
 from unittest.mock import patch
 
+from tests.api.harness import (
+    bootstrap_test_package,
+    cleanup_optional_module,
+    install_aiohttp_stub,
+    install_request_guards_stub,
+    install_server_stub,
+)
+
 
 class _FakeResponse:
     def __init__(self, payload, status=200):
@@ -16,10 +24,11 @@ class _FakeResponse:
 
 
 class _FakeRequest:
-    def __init__(self, payload=None, match_info=None, query=None):
+    def __init__(self, payload=None, match_info=None, query=None, headers=None):
         self._payload = payload
         self.match_info = match_info or {}
         self.query = query or {}
+        self.headers = headers or {}
 
     async def json(self):
         return self._payload
@@ -49,8 +58,8 @@ class _FakeHTTPClientSession:
         self._status = status
         self.calls = []
 
-    def get(self, url, params=None, timeout=None):
-        self.calls.append({"url": url, "params": params, "timeout": timeout})
+    def get(self, url, params=None, headers=None, timeout=None):
+        self.calls.append({"url": url, "params": params, "headers": headers, "timeout": timeout})
         return _FakeHTTPClientResponse(self._payload, status=self._status)
 
 
@@ -67,7 +76,7 @@ class _DummyWorkerManager:
         }
         return 12345
 
-    def _is_process_running(self, _pid):
+    def is_process_running(self, _pid):
         return False
 
     def save_processes(self):
@@ -89,21 +98,7 @@ def _load_worker_routes_module():
     module_path = Path(__file__).resolve().parents[2] / "api" / "worker_routes.py"
     package_name = "dist_api_worker_testpkg"
 
-    for mod_name in list(sys.modules):
-        if mod_name == package_name or mod_name.startswith(f"{package_name}."):
-            del sys.modules[mod_name]
-
-    root_pkg = types.ModuleType(package_name)
-    root_pkg.__path__ = []
-    sys.modules[package_name] = root_pkg
-
-    api_pkg = types.ModuleType(f"{package_name}.api")
-    api_pkg.__path__ = []
-    sys.modules[f"{package_name}.api"] = api_pkg
-
-    utils_pkg = types.ModuleType(f"{package_name}.utils")
-    utils_pkg.__path__ = []
-    sys.modules[f"{package_name}.utils"] = utils_pkg
+    bootstrap_test_package(package_name, with_api=True, with_utils=True, with_workers=True)
 
     workers_pkg = types.ModuleType(f"{package_name}.workers")
     workers_pkg.__path__ = []
@@ -119,59 +114,10 @@ def _load_worker_routes_module():
     detection_module.get_comms_channel = lambda *_args, **_kwargs: "lan"
     sys.modules[f"{package_name}.workers.detection"] = detection_module
 
-    created_aiohttp_stub = False
-    if "aiohttp" not in sys.modules:
-        created_aiohttp_stub = True
-        aiohttp_module = types.ModuleType("aiohttp")
-
-        class _ClientTimeout:
-            def __init__(self, total=None):
-                self.total = total
-
-        class _WSMsgType:
-            TEXT = "TEXT"
-            ERROR = "ERROR"
-            CLOSED = "CLOSED"
-
-        class _WebSocketResponse:
-            def __init__(self, *args, **kwargs):
-                self.args = args
-                self.kwargs = kwargs
-
-            async def prepare(self, _request):
-                return None
-
-            async def send_json(self, _payload):
-                return None
-
-            def __aiter__(self):
-                async def _empty():
-                    if False:
-                        yield None
-                return _empty()
-
-        aiohttp_module.ClientTimeout = _ClientTimeout
-        aiohttp_module.WSMsgType = _WSMsgType
-        aiohttp_module.web = types.SimpleNamespace(
-            json_response=lambda payload, status=200: _FakeResponse(payload, status=status),
-            WebSocketResponse=_WebSocketResponse,
-        )
-        sys.modules["aiohttp"] = aiohttp_module
-
-    class _Routes:
-        def get(self, _path):
-            def _decorator(fn):
-                return fn
-            return _decorator
-
-        def post(self, _path):
-            def _decorator(fn):
-                return fn
-            return _decorator
-
-    server_module = types.ModuleType("server")
-    server_module.PromptServer = types.SimpleNamespace(instance=types.SimpleNamespace(routes=_Routes()))
-    sys.modules["server"] = server_module
+    created_aiohttp_stub = install_aiohttp_stub(
+        lambda payload, status=200: _FakeResponse(payload, status=status)
+    )
+    install_server_stub()
 
     created_torch_stub = False
     if "torch" not in sys.modules:
@@ -248,19 +194,29 @@ def _load_worker_routes_module():
     def _validate_worker_id(worker_id, config):
         return any(str(worker.get("id")) == str(worker_id) for worker in config.get("workers", []))
 
+    def _require_worker_id(worker_id, config, field_name="worker_id"):
+        _ = field_name
+        worker_id_str = str(worker_id).strip()
+        if any(str(worker.get("id")) == worker_id_str for worker in config.get("workers", [])):
+            return worker_id_str
+        raise ValueError(f"Worker {worker_id_str} not found")
+
     schemas_module.require_fields = _require_fields
     schemas_module.validate_worker_id = _validate_worker_id
+    schemas_module.require_worker_id = _require_worker_id
+    schemas_module.is_authorized_request = lambda _request, _config: True
+    schemas_module.distributed_auth_headers = lambda _config: {}
     sys.modules[f"{package_name}.api.schemas"] = schemas_module
 
+    install_request_guards_stub(package_name)
+
     spec = importlib.util.spec_from_file_location(f"{package_name}.api.worker_routes", module_path)
     module = importlib.util.module_from_spec(spec)
     assert spec is not None and spec.loader is not None
     spec.loader.exec_module(module)
 
-    if created_aiohttp_stub:
-        sys.modules.pop("aiohttp", None)
-    if created_torch_stub:
-        sys.modules.pop("torch", None)
+    cleanup_optional_module("aiohttp", created_aiohttp_stub)
+    cleanup_optional_module("torch", created_torch_stub)
 
     return module
 

tests/conftest.py

@@ -1,2 +0,0 @@
-# This conftest.py marks the tests/ directory as the pytest collection root,
-# preventing pytest from traversing into the parent package's __init__.py.

tests/unit/nodes/test_distributed_upscale.py

@@ -37,6 +37,14 @@ def _bootstrap_package(package_name):
     logging_module.log = lambda *_args, **_kwargs: None
     sys.modules[f"{package_name}.utils.logging"] = logging_module
 
+    worker_ids_module = types.ModuleType(f"{package_name}.utils.worker_ids")
+    worker_ids_module.coerce_enabled_worker_ids = lambda value: (
+        [str(v) for v in value]
+        if isinstance(value, list)
+        else []
+    )
+    sys.modules[f"{package_name}.utils.worker_ids"] = worker_ids_module
+
     async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
     async_helpers_module.run_async_in_server_loop = lambda coro, timeout=None: (_ for _ in ()).throw(
         RuntimeError(f"run_async_in_server_loop should not be used in these tests (timeout={timeout})")
@@ -139,6 +147,7 @@ def _load_distributed_upscale_module():
     package_name = "dist_distributed_upscale_testpkg"
     _bootstrap_package(package_name)
     _load_module(package_name, "nodes/hidden_inputs.py", "nodes.hidden_inputs")
+    _load_module(package_name, "upscale/mode_contexts.py", "upscale.mode_contexts")
     _load_module(package_name, "upscale/processing_args.py", "upscale.processing_args")
     return _load_module(package_name, "nodes/distributed_upscale.py", "nodes.distributed_upscale")
 

tests/unit/nodes/test_distributed_value.py

@@ -28,6 +28,18 @@ class DistributedValueTests(unittest.TestCase):
         root_pkg.__path__ = []
         sys.modules[pkg_name] = root_pkg
 
+        nodes_pkg = types.ModuleType(f"{pkg_name}.nodes")
+        nodes_pkg.__path__ = []
+        sys.modules[f"{pkg_name}.nodes"] = nodes_pkg
+
+        hidden_inputs_mod = types.ModuleType(f"{pkg_name}.nodes.hidden_inputs")
+        hidden_inputs_mod.build_worker_identity_hidden_inputs = lambda: {
+            "is_worker": ("BOOLEAN", {"default": False}),
+            "enabled_worker_ids": ("STRING", {"default": "[]"}),
+            "worker_id": ("STRING", {"default": ""}),
+        }
+        sys.modules[f"{pkg_name}.nodes.hidden_inputs"] = hidden_inputs_mod
+
         utils_pkg = types.ModuleType(f"{pkg_name}.utils")
         utils_pkg.__path__ = []
         sys.modules[f"{pkg_name}.utils"] = utils_pkg
@@ -37,6 +49,16 @@ class DistributedValueTests(unittest.TestCase):
         logging_mod.log = lambda *_a, **_k: None
         sys.modules[f"{pkg_name}.utils.logging"] = logging_mod
 
+        worker_ids_path = Path(__file__).resolve().parents[3] / "utils" / "worker_ids.py"
+        worker_ids_spec = importlib.util.spec_from_file_location(
+            f"{pkg_name}.utils.worker_ids",
+            worker_ids_path,
+        )
+        worker_ids_mod = importlib.util.module_from_spec(worker_ids_spec)
+        assert worker_ids_spec is not None and worker_ids_spec.loader is not None
+        worker_ids_spec.loader.exec_module(worker_ids_mod)
+        sys.modules[f"{pkg_name}.utils.worker_ids"] = worker_ids_mod
+
         spec = importlib.util.spec_from_file_location(
             f"{pkg_name}.nodes.utilities", module_path
         )

tests/unit/nodes/test_list_collector.py

@@ -1,237 +0,0 @@
-import asyncio
-import base64
-import importlib.util
-import io
-import json
-import sys
-import types
-import unittest
-from pathlib import Path
-
-import torch
-from PIL import Image
-
-
-def _bootstrap_package(package_name):
-    for mod_name in list(sys.modules):
-        if mod_name == package_name or mod_name.startswith(f"{package_name}."):
-            del sys.modules[mod_name]
-
-    root_pkg = types.ModuleType(package_name)
-    root_pkg.__path__ = []
-    sys.modules[package_name] = root_pkg
-
-    nodes_pkg = types.ModuleType(f"{package_name}.nodes")
-    nodes_pkg.__path__ = []
-    sys.modules[f"{package_name}.nodes"] = nodes_pkg
-
-    utils_pkg = types.ModuleType(f"{package_name}.utils")
-    utils_pkg.__path__ = []
-    sys.modules[f"{package_name}.utils"] = utils_pkg
-
-    logging_module = types.ModuleType(f"{package_name}.utils.logging")
-    logging_module.debug_log = lambda *_args, **_kwargs: None
-    logging_module.log = lambda *_args, **_kwargs: None
-    sys.modules[f"{package_name}.utils.logging"] = logging_module
-
-    config_module = types.ModuleType(f"{package_name}.utils.config")
-    config_module.get_worker_timeout_seconds = lambda: 0.2
-    config_module.is_master_delegate_only = lambda: False
-    sys.modules[f"{package_name}.utils.config"] = config_module
-
-    constants_module = types.ModuleType(f"{package_name}.utils.constants")
-    constants_module.HEARTBEAT_INTERVAL = 1.0
-    sys.modules[f"{package_name}.utils.constants"] = constants_module
-
-    image_module = types.ModuleType(f"{package_name}.utils.image")
-
-    def ensure_contiguous(tensor):
-        return tensor.contiguous()
-
-    def tensor_to_pil(image_batch, index):
-        tensor = image_batch[index].detach().cpu().clamp(0, 1)
-        array = (tensor.numpy() * 255.0).astype("uint8")
-        return Image.fromarray(array)
-
-    def encode_tensor_png_data_url(image_batch, batch_index=0):
-        image = tensor_to_pil(image_batch, batch_index)
-        byte_io = io.BytesIO()
-        image.save(byte_io, format="PNG", compress_level=0)
-        encoded = base64.b64encode(byte_io.getvalue()).decode("utf-8")
-        return f"data:image/png;base64,{encoded}"
-
-    image_module.ensure_contiguous = ensure_contiguous
-    image_module.tensor_to_pil = tensor_to_pil
-    image_module.encode_tensor_png_data_url = encode_tensor_png_data_url
-    sys.modules[f"{package_name}.utils.image"] = image_module
-
-    network_module = types.ModuleType(f"{package_name}.utils.network")
-
-    async def get_client_session():
-        raise RuntimeError("test should patch get_client_session")
-
-    network_module.get_client_session = get_client_session
-    sys.modules[f"{package_name}.utils.network"] = network_module
-
-    async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
-    async_helpers_module.run_async_in_server_loop = lambda coro: asyncio.run(coro)
-    sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
-
-
-def _load_module(package_name, module_rel_path, module_name):
-    module_path = Path(__file__).resolve().parents[3] / module_rel_path
-    spec = importlib.util.spec_from_file_location(
-        f"{package_name}.{module_name}",
-        module_path,
-    )
-    module = importlib.util.module_from_spec(spec)
-    assert spec is not None and spec.loader is not None
-    sys.modules[spec.name] = module
-    spec.loader.exec_module(module)
-    return module
-
-
-def _load_list_collector_module():
-    package_name = "dist_list_collector_testpkg"
-    _bootstrap_package(package_name)
-    _load_module(package_name, "nodes/hidden_inputs.py", "nodes.hidden_inputs")
-    _load_module(package_name, "nodes/runtime_helpers.py", "nodes.runtime_helpers")
-    _load_module(package_name, "nodes/queue_wait.py", "nodes.queue_wait")
-    _load_module(package_name, "utils/worker_ids.py", "utils.worker_ids")
-    return _load_module(package_name, "nodes/list_collector.py", "nodes.list_collector")
-
-
-collector_module = _load_list_collector_module()
-
-
-def _context(**overrides):
-    defaults = {
-        "multi_job_id": "",
-        "is_worker": False,
-        "master_url": "",
-        "enabled_worker_ids": "[]",
-        "worker_id": "",
-        "delegate_only": False,
-    }
-    defaults.update(overrides)
-    return collector_module.ListCollectorRunContext(**defaults)
-
-
-class _FakeResponseCtx:
-    def __init__(self, recorder, payload, url):
-        self._recorder = recorder
-        self._payload = payload
-        self._url = url
-
-    async def __aenter__(self):
-        self._recorder.append({"url": self._url, "payload": self._payload})
-        return self
-
-    async def __aexit__(self, exc_type, exc, tb):
-        return False
-
-    def raise_for_status(self):
-        return None
-
-
-class _FakeSession:
-    def __init__(self):
-        self.calls = []
-
-    def post(self, url, json=None, timeout=None):
-        return _FakeResponseCtx(self.calls, json, url)
-
-
-class _FakePromptServer:
-    def __init__(self):
-        self.distributed_pending_jobs = {}
-        self.distributed_jobs_lock = asyncio.Lock()
-
-
-class DistributedListCollectorTests(unittest.IsolatedAsyncioTestCase):
-    def setUp(self):
-        self.node = collector_module.DistributedListCollector()
-
-    async def test_worker_mode_sends_images_to_master(self):
-        fake_session = _FakeSession()
-
-        async def _get_session():
-            return fake_session
-
-        collector_module.get_client_session = _get_session
-
-        images = [
-            torch.zeros((1, 4, 4, 3), dtype=torch.float32),
-            torch.ones((1, 4, 4, 3), dtype=torch.float32),
-        ]
-
-        result = await self.node.execute(
-            images,
-            context=_context(
-                multi_job_id="job-list",
-                is_worker=True,
-                master_url="http://master.local:8188",
-                worker_id="worker-a",
-            ),
-        )
-
-        self.assertEqual(len(result[0]), 2)
-        self.assertEqual(len(fake_session.calls), 2)
-        self.assertEqual(fake_session.calls[0]["url"], "http://master.local:8188/distributed/job_complete")
-        self.assertEqual(fake_session.calls[0]["payload"]["batch_idx"], 0)
-        self.assertFalse(fake_session.calls[0]["payload"]["is_last"])
-        self.assertEqual(fake_session.calls[1]["payload"]["batch_idx"], 1)
-        self.assertTrue(fake_session.calls[1]["payload"]["is_last"])
-
-    async def test_master_mode_collects_results_as_list(self):
-        fake_server = _FakePromptServer()
-        collector_module._get_prompt_server_instance = lambda: fake_server
-
-        queue = asyncio.Queue()
-        worker_a = torch.full((1, 4, 4, 3), 0.25, dtype=torch.float32)
-        worker_b = torch.full((1, 4, 4, 3), 0.75, dtype=torch.float32)
-        await queue.put({"tensor": worker_a, "worker_id": "worker-a", "image_index": 0, "is_last": True})
-        await queue.put({"tensor": worker_b, "worker_id": "worker-b", "image_index": 0, "is_last": True})
-        fake_server.distributed_pending_jobs["job-list"] = queue
-
-        master_images = [torch.zeros((1, 4, 4, 3), dtype=torch.float32)]
-
-        result = await self.node.execute(
-            master_images,
-            context=_context(
-                multi_job_id="job-list",
-                enabled_worker_ids=json.dumps(["worker-a", "worker-b"]),
-            ),
-        )
-
-        self.assertEqual(len(result[0]), 3)
-        self.assertTrue(torch.equal(result[0][0], master_images[0]))
-        self.assertTrue(torch.equal(result[0][1], worker_a))
-        self.assertTrue(torch.equal(result[0][2], worker_b))
-
-    async def test_delegate_only_mode_returns_worker_images_only(self):
-        fake_server = _FakePromptServer()
-        collector_module._get_prompt_server_instance = lambda: fake_server
-
-        queue = asyncio.Queue()
-        worker_img = torch.full((1, 4, 4, 3), 0.5, dtype=torch.float32)
-        await queue.put({"tensor": worker_img, "worker_id": "worker-a", "image_index": 0, "is_last": True})
-        fake_server.distributed_pending_jobs["job-list"] = queue
-
-        master_images = [torch.zeros((1, 4, 4, 3), dtype=torch.float32)]
-
-        result = await self.node.execute(
-            master_images,
-            context=_context(
-                multi_job_id="job-list",
-                enabled_worker_ids=json.dumps(["worker-a"]),
-                delegate_only=True,
-            ),
-        )
-
-        self.assertEqual(len(result[0]), 1)
-        self.assertTrue(torch.equal(result[0][0], worker_img))
-
-
-if __name__ == "__main__":
-    unittest.main()

tests/unit/nodes/test_list_splitter.py

@@ -1,261 +0,0 @@
-import asyncio
-import importlib.util
-import sys
-import types
-import unittest
-from dataclasses import dataclass, field
-from pathlib import Path
-
-import torch
-
-
-def _bootstrap_package(package_name):
-    for mod_name in list(sys.modules):
-        if mod_name == package_name or mod_name.startswith(f"{package_name}."):
-            del sys.modules[mod_name]
-
-    root_pkg = types.ModuleType(package_name)
-    root_pkg.__path__ = []
-    sys.modules[package_name] = root_pkg
-
-    nodes_pkg = types.ModuleType(f"{package_name}.nodes")
-    nodes_pkg.__path__ = []
-    sys.modules[f"{package_name}.nodes"] = nodes_pkg
-
-    utils_pkg = types.ModuleType(f"{package_name}.utils")
-    utils_pkg.__path__ = []
-    sys.modules[f"{package_name}.utils"] = utils_pkg
-
-    upscale_pkg = types.ModuleType(f"{package_name}.upscale")
-    upscale_pkg.__path__ = []
-    sys.modules[f"{package_name}.upscale"] = upscale_pkg
-
-    logging_module = types.ModuleType(f"{package_name}.utils.logging")
-    logging_module.debug_log = lambda *_args, **_kwargs: None
-    logging_module.log = lambda *_args, **_kwargs: None
-    sys.modules[f"{package_name}.utils.logging"] = logging_module
-
-    async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
-    async_helpers_module.run_async_in_server_loop = lambda coro: asyncio.run(coro)
-    sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
-
-    network_module = types.ModuleType(f"{package_name}.utils.network")
-
-    async def _get_client_session():
-        raise RuntimeError("test should patch get_client_session")
-
-    network_module.get_client_session = _get_client_session
-    sys.modules[f"{package_name}.utils.network"] = network_module
-
-    job_models_module = types.ModuleType(f"{package_name}.upscale.job_models")
-
-    class BaseJobState:
-        pass
-
-    @dataclass
-    class ImageJobState(BaseJobState):
-        multi_job_id: str
-        mode: str = field(default="dynamic", init=False)
-        queue: asyncio.Queue = field(default_factory=asyncio.Queue)
-        pending_images: asyncio.Queue = field(default_factory=asyncio.Queue)
-        completed_images: dict = field(default_factory=dict)
-        worker_status: dict = field(default_factory=dict)
-        assigned_to_workers: dict = field(default_factory=dict)
-
-        @property
-        def pending_tasks(self):
-            return self.pending_images
-
-    job_models_module.BaseJobState = BaseJobState
-    job_models_module.ImageJobState = ImageJobState
-    sys.modules[f"{package_name}.upscale.job_models"] = job_models_module
-
-    prompt_server_holder = {
-        "value": types.SimpleNamespace(
-            distributed_tile_jobs_lock=asyncio.Lock(),
-            distributed_pending_tile_jobs={},
-        )
-    }
-
-    job_store_module = types.ModuleType(f"{package_name}.upscale.job_store")
-
-    def ensure_tile_jobs_initialized():
-        return prompt_server_holder["value"]
-
-    async def init_dynamic_job(multi_job_id, batch_size, enabled_workers, all_indices=None):
-        prompt_server = prompt_server_holder["value"]
-        async with prompt_server.distributed_tile_jobs_lock:
-            if multi_job_id in prompt_server.distributed_pending_tile_jobs:
-                return
-            job_data = ImageJobState(multi_job_id=multi_job_id)
-            for idx in (all_indices if all_indices is not None else list(range(batch_size))):
-                await job_data.pending_images.put(idx)
-            prompt_server.distributed_pending_tile_jobs[multi_job_id] = job_data
-
-    job_store_module.ensure_tile_jobs_initialized = ensure_tile_jobs_initialized
-    job_store_module.init_dynamic_job = init_dynamic_job
-    job_store_module._prompt_server_holder = prompt_server_holder
-    sys.modules[f"{package_name}.upscale.job_store"] = job_store_module
-
-
-def _load_module(package_name, module_rel_path, module_name):
-    module_path = Path(__file__).resolve().parents[3] / module_rel_path
-    spec = importlib.util.spec_from_file_location(
-        f"{package_name}.{module_name}",
-        module_path,
-    )
-    module = importlib.util.module_from_spec(spec)
-    assert spec is not None and spec.loader is not None
-    sys.modules[spec.name] = module
-    spec.loader.exec_module(module)
-    return module
-
-
-def _load_list_splitter_module():
-    package_name = "dist_list_splitter_testpkg"
-    _bootstrap_package(package_name)
-    _load_module(package_name, "nodes/utilities.py", "nodes.utilities")
-    module = _load_module(package_name, "nodes/list_splitter.py", "nodes.list_splitter")
-    return module
-
-
-splitter_module = _load_list_splitter_module()
-
-
-class _FakeResponse:
-    def __init__(self, payload, status=200):
-        self.payload = payload
-        self.status = status
-
-    async def __aenter__(self):
-        return self
-
-    async def __aexit__(self, exc_type, exc, tb):
-        return False
-
-    async def json(self):
-        return self.payload
-
-
-class _FakeSession:
-    def __init__(self, responses):
-        self.responses = list(responses)
-        self.calls = []
-
-    def post(self, url, json=None, timeout=None):
-        self.calls.append({"url": url, "json": json})
-        if not self.responses:
-            return _FakeResponse({"item_idx": None}, status=200)
-        payload = self.responses.pop(0)
-        return _FakeResponse(payload, status=200)
-
-
-class DistributedListSplitterTests(unittest.TestCase):
-    def setUp(self):
-        self.node = splitter_module.DistributedListSplitter()
-        prompt_server = splitter_module.ensure_tile_jobs_initialized()
-        prompt_server.distributed_pending_tile_jobs = {}
-
-    def _images(self, count):
-        return [torch.full((1, 4, 4, 3), float(idx), dtype=torch.float32) for idx in range(count)]
-
-    def test_single_participant_returns_full_list(self):
-        images = self._images(5)
-        out = self.node.split(images, mode="static", participant_index=0, total_participants=1)[0]
-        self.assertEqual(len(out), 5)
-        self.assertEqual([int(t[0, 0, 0, 0].item()) for t in out], [0, 1, 2, 3, 4])
-
-    def test_two_participants_split_six_images_evenly(self):
-        images = self._images(6)
-        p0 = self.node.split(images, mode="static", participant_index=0, total_participants=2)[0]
-        p1 = self.node.split(images, mode="static", participant_index=1, total_participants=2)[0]
-        self.assertEqual(len(p0), 3)
-        self.assertEqual(len(p1), 3)
-        self.assertEqual([int(t[0, 0, 0, 0].item()) for t in p0], [0, 1, 2])
-        self.assertEqual([int(t[0, 0, 0, 0].item()) for t in p1], [3, 4, 5])
-
-    def test_three_participants_split_seven_images(self):
-        images = self._images(7)
-        p0 = self.node.split(images, mode="static", participant_index=0, total_participants=3)[0]
-        p1 = self.node.split(images, mode="static", participant_index=1, total_participants=3)[0]
-        p2 = self.node.split(images, mode="static", participant_index=2, total_participants=3)[0]
-        self.assertEqual([len(p0), len(p1), len(p2)], [3, 2, 2])
-
-    def test_more_participants_than_images_yields_empty_slices(self):
-        images = self._images(2)
-        outputs = [self.node.split(images, mode="static", participant_index=i, total_participants=5)[0] for i in range(5)]
-        self.assertEqual([len(chunk) for chunk in outputs], [1, 1, 0, 0, 0])
-
-    def test_mixed_dimensions_are_preserved(self):
-        images = [
-            torch.zeros((1, 4, 4, 3), dtype=torch.float32),
-            torch.zeros((1, 8, 6, 3), dtype=torch.float32),
-            torch.zeros((1, 3, 9, 3), dtype=torch.float32),
-        ]
-        p0 = self.node.split(images, mode="static", participant_index=0, total_participants=2)[0]
-        p1 = self.node.split(images, mode="static", participant_index=1, total_participants=2)[0]
-        self.assertEqual([tuple(t.shape) for t in p0], [(1, 4, 4, 3), (1, 8, 6, 3)])
-        self.assertEqual([tuple(t.shape) for t in p1], [(1, 3, 9, 3)])
-
-    def test_dynamic_mode_master_pulls_from_local_queue(self):
-        images = self._images(4)
-
-        pulled = self.node.split(
-            images,
-            mode="dynamic",
-            multi_job_id="job-dynamic",
-            is_worker=False,
-            worker_id="master",
-        )[0]
-
-        self.assertEqual([int(t[0, 0, 0, 0].item()) for t in pulled], [0, 1, 2, 3])
-
-    def test_dynamic_mode_worker_pulls_via_http(self):
-        images = self._images(5)
-        fake_session = _FakeSession([
-            {"item_idx": 1, "estimated_remaining": 3},
-            {"item_idx": 3, "estimated_remaining": 2},
-            {"item_idx": None},
-        ])
-
-        async def _get_session():
-            return fake_session
-
-        splitter_module.get_client_session = _get_session
-
-        pulled = self.node.split(
-            images,
-            mode="dynamic",
-            multi_job_id="job-dynamic",
-            is_worker=True,
-            master_url="http://master.local:8188",
-            worker_id="worker-a",
-        )[0]
-
-        self.assertEqual([int(t[0, 0, 0, 0].item()) for t in pulled], [1, 3])
-        self.assertGreaterEqual(len(fake_session.calls), 3)
-
-    def test_dynamic_mode_exhausted_queue_returns_empty(self):
-        images = self._images(3)
-
-        first_pull = self.node.split(
-            images,
-            mode="dynamic",
-            multi_job_id="job-dynamic",
-            is_worker=False,
-            worker_id="master",
-        )[0]
-        second_pull = self.node.split(
-            images,
-            mode="dynamic",
-            multi_job_id="job-dynamic",
-            is_worker=False,
-            worker_id="master",
-        )[0]
-
-        self.assertEqual(len(first_pull), 3)
-        self.assertEqual(second_pull, [])
-
-
-if __name__ == "__main__":
-    unittest.main()

tests/unit/system/test_config.py

@@ -21,14 +21,13 @@ def _load_config_module():
     root_pkg.__path__ = []
     sys.modules[package_name] = root_pkg
 
-    logging_module = types.ModuleType(f"{package_name}.logging")
-    logging_module.log = lambda *_args, **_kwargs: None
-    logging_module.debug_log = lambda *_args, **_kwargs: None
-    sys.modules[f"{package_name}.logging"] = logging_module
-
     constants_module = types.ModuleType(f"{package_name}.constants")
     constants_module.GPU_CONFIG_FILE = "gpu_config.json"
     constants_module.HEARTBEAT_TIMEOUT = 30
+    constants_module.ORCHESTRATION_WORKER_PROBE_CONCURRENCY = 8
+    constants_module.ORCHESTRATION_WORKER_PREP_CONCURRENCY = 4
+    constants_module.ORCHESTRATION_MEDIA_SYNC_CONCURRENCY = 2
+    constants_module.ORCHESTRATION_MEDIA_SYNC_TIMEOUT = 120
     sys.modules[f"{package_name}.constants"] = constants_module
 
     spec = importlib.util.spec_from_file_location(f"{package_name}.config", module_path)

tests/unit/system/test_distributed.py

@@ -1,21 +1,17 @@
 import importlib.util
-import os
 import sys
 import types
 import unittest
 from pathlib import Path
 
 
-def _load_distributed_module(is_worker):
+def _load_distributed_module():
     module_path = Path(__file__).resolve().parents[3] / "distributed.py"
     package_name = "dist_distributed_entry_testpkg"
     calls = {
-        "ensure_config_exists": 0,
-        "ensure_distributed_state": 0,
-        "ensure_tile_jobs_initialized": 0,
-        "delayed_auto_launch": 0,
-        "register_async_signals": 0,
-        "sync_cleanup_registered": 0,
+        "build_node_mappings": 0,
+        "initialize_runtime": 0,
+        "initialize_runtime_args": [],
     }
 
     for mod_name in list(sys.modules):
@@ -26,126 +22,49 @@ def _load_distributed_module(is_worker):
     root_pkg.__path__ = []
     sys.modules[package_name] = root_pkg
 
-    # Absolute dependency used by module.
-    server_module = types.ModuleType("server")
-    server_module.PromptServer = types.SimpleNamespace(instance=types.SimpleNamespace())
-    sys.modules["server"] = server_module
+    bootstrap_pkg = types.ModuleType(f"{package_name}.bootstrap")
+    bootstrap_pkg.__path__ = []
+    sys.modules[f"{package_name}.bootstrap"] = bootstrap_pkg
 
-    utils_pkg = types.ModuleType(f"{package_name}.utils")
-    utils_pkg.__path__ = []
-    sys.modules[f"{package_name}.utils"] = utils_pkg
+    entrypoint_module = types.ModuleType(f"{package_name}.bootstrap.entrypoint")
 
-    config_module = types.ModuleType(f"{package_name}.utils.config")
-    config_module.ensure_config_exists = lambda: calls.__setitem__("ensure_config_exists", calls["ensure_config_exists"] + 1)
-    sys.modules[f"{package_name}.utils.config"] = config_module
+    def _build_node_mappings():
+        calls["build_node_mappings"] += 1
+        return ({"MockNode": object}, {"MockNode": "Mock Node"})
 
-    logging_module = types.ModuleType(f"{package_name}.utils.logging")
-    logging_module.debug_log = lambda *_args, **_kwargs: None
-    sys.modules[f"{package_name}.utils.logging"] = logging_module
+    def _initialize_runtime(prompt_server=None):
+        calls["initialize_runtime"] += 1
+        calls["initialize_runtime_args"].append(prompt_server)
 
-    network_module = types.ModuleType(f"{package_name}.utils.network")
+    entrypoint_module.build_node_mappings = _build_node_mappings
+    entrypoint_module.initialize_runtime = _initialize_runtime
+    sys.modules[f"{package_name}.bootstrap.entrypoint"] = entrypoint_module
 
-    async def _cleanup_client_session():
-        return None
-
-    network_module.cleanup_client_session = _cleanup_client_session
-    sys.modules[f"{package_name}.utils.network"] = network_module
-
-    workers_pkg = types.ModuleType(f"{package_name}.workers")
-    workers_pkg.__path__ = []
-    workers_pkg.get_worker_manager = lambda: None
-    sys.modules[f"{package_name}.workers"] = workers_pkg
-
-    startup_module = types.ModuleType(f"{package_name}.workers.startup")
-    startup_module.delayed_auto_launch = lambda: calls.__setitem__("delayed_auto_launch", calls["delayed_auto_launch"] + 1)
-    startup_module.register_async_signals = lambda: calls.__setitem__("register_async_signals", calls["register_async_signals"] + 1)
-    startup_module.sync_cleanup = lambda: None
-    sys.modules[f"{package_name}.workers.startup"] = startup_module
-
-    upscale_pkg = types.ModuleType(f"{package_name}.upscale")
-    upscale_pkg.__path__ = []
-    sys.modules[f"{package_name}.upscale"] = upscale_pkg
-
-    job_store_module = types.ModuleType(f"{package_name}.upscale.job_store")
-    job_store_module.ensure_tile_jobs_initialized = lambda: calls.__setitem__(
-        "ensure_tile_jobs_initialized", calls["ensure_tile_jobs_initialized"] + 1
-    )
-    sys.modules[f"{package_name}.upscale.job_store"] = job_store_module
-
-    nodes_module = types.ModuleType(f"{package_name}.nodes")
-    nodes_module.NODE_CLASS_MAPPINGS = {}
-    nodes_module.NODE_DISPLAY_NAME_MAPPINGS = {}
-    nodes_module.ImageBatchDivider = object
-    nodes_module.DistributedCollectorNode = object
-    nodes_module.DistributedSeed = object
-    nodes_module.DistributedModelName = object
-    nodes_module.DistributedValue = object
-    nodes_module.AudioBatchDivider = object
-    nodes_module.DistributedEmptyImage = object
-    nodes_module.DistributedListSplitter = object
-    nodes_module.DistributedListCollector = object
-    nodes_module.DistributedBranch = object
-    nodes_module.DistributedBranchCollector = object
-    nodes_module.AnyType = object
-    nodes_module.ByPassTypeTuple = tuple
-    nodes_module.any_type = "*"
-    sys.modules[f"{package_name}.nodes"] = nodes_module
-
-    api_module = types.ModuleType(f"{package_name}.api")
-    sys.modules[f"{package_name}.api"] = api_module
-
-    queue_orch_module = types.ModuleType(f"{package_name}.api.queue_orchestration")
-    queue_orch_module.ensure_distributed_state = lambda _ps: calls.__setitem__(
-        "ensure_distributed_state", calls["ensure_distributed_state"] + 1
-    )
-    sys.modules[f"{package_name}.api.queue_orchestration"] = queue_orch_module
-
-    import atexit
-
-    original_register = atexit.register
-    atexit.register = lambda fn: calls.__setitem__("sync_cleanup_registered", calls["sync_cleanup_registered"] + 1) or fn
-
-    original_env = os.environ.get("COMFYUI_IS_WORKER")
-    try:
-        if is_worker:
-            os.environ["COMFYUI_IS_WORKER"] = "1"
-        else:
-            os.environ.pop("COMFYUI_IS_WORKER", None)
-
-        spec = importlib.util.spec_from_file_location(f"{package_name}.distributed", module_path)
-        module = importlib.util.module_from_spec(spec)
-        assert spec is not None and spec.loader is not None
-        sys.modules[spec.name] = module
-        spec.loader.exec_module(module)
-    finally:
-        atexit.register = original_register
-        if original_env is None:
-            os.environ.pop("COMFYUI_IS_WORKER", None)
-        else:
-            os.environ["COMFYUI_IS_WORKER"] = original_env
-
-    return calls
+    spec = importlib.util.spec_from_file_location(f"{package_name}.distributed", module_path)
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module, calls
 
 
 class DistributedEntryTests(unittest.TestCase):
-    def test_import_initializes_core_state_for_worker(self):
-        calls = _load_distributed_module(is_worker=True)
+    def test_import_builds_mappings_without_runtime_bootstrap(self):
+        module, calls = _load_distributed_module()
 
-        self.assertEqual(calls["ensure_config_exists"], 1)
-        self.assertEqual(calls["ensure_distributed_state"], 1)
-        self.assertEqual(calls["ensure_tile_jobs_initialized"], 1)
-        self.assertEqual(calls["delayed_auto_launch"], 0)
-        self.assertEqual(calls["register_async_signals"], 0)
+        self.assertIn("MockNode", module.NODE_CLASS_MAPPINGS)
+        self.assertEqual(module.NODE_DISPLAY_NAME_MAPPINGS["MockNode"], "Mock Node")
+        self.assertEqual(calls["build_node_mappings"], 1)
+        self.assertEqual(calls["initialize_runtime"], 0)
 
-    def test_import_triggers_master_startup_hooks(self):
-        calls = _load_distributed_module(is_worker=False)
+    def test_initialize_runtime_delegates_to_bootstrap_entrypoint(self):
+        module, calls = _load_distributed_module()
+        sentinel = object()
 
-        self.assertEqual(calls["ensure_config_exists"], 1)
-        self.assertEqual(calls["ensure_distributed_state"], 1)
-        self.assertEqual(calls["ensure_tile_jobs_initialized"], 1)
-        self.assertEqual(calls["delayed_auto_launch"], 1)
-        self.assertEqual(calls["register_async_signals"], 1)
-        self.assertGreaterEqual(calls["sync_cleanup_registered"], 2)
+        module.initialize_runtime(sentinel)
+
+        self.assertEqual(calls["initialize_runtime"], 1)
+        self.assertEqual(calls["initialize_runtime_args"], [sentinel])
 
 
 if __name__ == "__main__":

tests/unit/system/test_logging.py

@@ -1,63 +1,29 @@
 import io
-import json
-import os
-import tempfile
 import unittest
 from contextlib import redirect_stdout
+from unittest.mock import patch
 
 from utils import logging as logging_utils
 
 
 class LoggingUtilsTests(unittest.TestCase):
-    def setUp(self):
-        self.original_config_file = logging_utils.GPU_CONFIG_FILE
-        self.state = logging_utils._debug_state()
-        self.original_cache = self.state.enabled
-        self.original_cache_time = self.state.updated_at
+    def test_is_debug_enabled_delegates_to_config(self):
+        with patch.object(logging_utils, "get_debug_enabled", return_value=True) as mocked:
+            self.assertTrue(logging_utils.is_debug_enabled())
+        mocked.assert_called_once_with(default=False)
 
-    def tearDown(self):
-        logging_utils.GPU_CONFIG_FILE = self.original_config_file
-        self.state.enabled = self.original_cache
-        self.state.updated_at = self.original_cache_time
-
-    def _reset_cache(self):
-        self.state.enabled = None
-        self.state.updated_at = 0.0
-
-    def test_is_debug_enabled_reads_config_flag(self):
-        with tempfile.NamedTemporaryFile("w", delete=False) as temp_file:
-            json.dump({"settings": {"debug": True}}, temp_file)
-            temp_path = temp_file.name
-        self.addCleanup(lambda: os.path.exists(temp_path) and os.unlink(temp_path))
-
-        logging_utils.GPU_CONFIG_FILE = temp_path
-        self._reset_cache()
-
-        self.assertTrue(logging_utils.is_debug_enabled())
-
-    def test_is_debug_enabled_handles_invalid_json(self):
-        with tempfile.NamedTemporaryFile("w", delete=False) as temp_file:
-            temp_file.write("{invalid-json")
-            temp_path = temp_file.name
-        self.addCleanup(lambda: os.path.exists(temp_path) and os.unlink(temp_path))
-
-        logging_utils.GPU_CONFIG_FILE = temp_path
-        self._reset_cache()
-
-        self.assertFalse(logging_utils.is_debug_enabled())
+    def test_is_debug_enabled_handles_config_errors(self):
+        with patch.object(logging_utils, "get_debug_enabled", side_effect=RuntimeError("boom")):
+            self.assertFalse(logging_utils.is_debug_enabled())
 
     def test_debug_log_emits_only_when_enabled(self):
-        self.state.enabled = True
-        self.state.updated_at = 1e18
         output = io.StringIO()
-        with redirect_stdout(output):
+        with patch.object(logging_utils, "is_debug_enabled", return_value=True), redirect_stdout(output):
             logging_utils.debug_log("hello")
         self.assertIn("[Distributed] hello", output.getvalue())
 
-        self.state.enabled = False
-        self.state.updated_at = 1e18
         output = io.StringIO()
-        with redirect_stdout(output):
+        with patch.object(logging_utils, "is_debug_enabled", return_value=False), redirect_stdout(output):
             logging_utils.debug_log("hidden")
         self.assertEqual(output.getvalue(), "")
 

tests/unit/system/test_prompt_transform.py

@@ -81,16 +81,6 @@ def _delegate_prompt():
     }
 
 
-def _list_splitter_prompt():
-    """1(List source) → 2(DistributedListSplitter) → 3(DistributedListCollector) → 4(SaveImage)."""
-    return {
-        "1": {"class_type": "ImageListSource", "inputs": {}},
-        "2": {"class_type": "DistributedListSplitter", "inputs": {"images": ["1", 0]}},
-        "3": {"class_type": "DistributedListCollector", "inputs": {"images": ["2", 0]}},
-        "4": {"class_type": "SaveImage", "inputs": {"images": ["3", 0]}},
-    }
-
-
 def _branch_prompt():
     """1 → 2(DistributedBranch) -> slot0:3->4, slot1:5->6."""
     return {
@@ -285,14 +275,6 @@ class PrunePromptForWorkerTests(unittest.TestCase):
         self.assertIn("2", result)
         self.assertNotIn("3", result)
 
-    def test_list_collector_is_treated_as_distributed_anchor(self):
-        prompt = _list_splitter_prompt()
-        result = pt.prune_prompt_for_worker(prompt)
-        self.assertIn("1", result)
-        self.assertIn("2", result)
-        self.assertIn("3", result)
-        self.assertNotIn("4", result)
-
     def test_branch_anchor_keeps_downstream_for_later_branch_pruning(self):
         prompt = _branch_prompt()
         result = pt.prune_prompt_for_worker(prompt)
@@ -385,22 +367,6 @@ class GenerateJobIdMapTests(unittest.TestCase):
         job_map = pt.generate_job_id_map(idx, "run")
         self.assertEqual(job_map["5"], "run_5")
 
-    def test_maps_list_collector_nodes(self):
-        prompt = {
-            "8": {"class_type": "DistributedListCollector", "inputs": {}},
-        }
-        idx = pt.PromptIndex(prompt)
-        job_map = pt.generate_job_id_map(idx, "run")
-        self.assertEqual(job_map["8"], "run_8")
-
-    def test_maps_list_splitter_nodes(self):
-        prompt = {
-            "7": {"class_type": "DistributedListSplitter", "inputs": {}},
-        }
-        idx = pt.PromptIndex(prompt)
-        job_map = pt.generate_job_id_map(idx, "run")
-        self.assertEqual(job_map["7"], "run_7")
-
     def test_maps_branch_collector_nodes(self):
         prompt = {
             "10": {"class_type": "DistributedBranchCollector", "inputs": {}},
@@ -508,9 +474,9 @@ class ApplyOverridesSeedTests(unittest.TestCase):
         result = _apply(self._seed_prompt(), "worker-a")
         self.assertTrue(result["1"]["inputs"]["is_worker"])
 
-    def test_worker_id_reflects_index_in_enabled_list(self):
+    def test_worker_id_uses_canonical_worker_id(self):
         result = _apply(self._seed_prompt(), "worker-b", enabled_worker_ids=["worker-a", "worker-b"])
-        self.assertEqual(result["1"]["inputs"]["worker_id"], "worker_1")
+        self.assertEqual(result["1"]["inputs"]["worker_id"], "worker-b")
 
     def test_master_sets_is_worker_false(self):
         result = _apply(self._seed_prompt(), "master")
@@ -566,9 +532,9 @@ class ApplyOverridesValueTests(unittest.TestCase):
         result = _apply(self._value_prompt(), "worker-a")
         self.assertTrue(result["1"]["inputs"]["is_worker"])
 
-    def test_worker_id_reflects_index_in_enabled_list(self):
+    def test_worker_id_uses_canonical_worker_id(self):
         result = _apply(self._value_prompt(), "worker-b", enabled_worker_ids=["worker-a", "worker-b"])
-        self.assertEqual(result["1"]["inputs"]["worker_id"], "worker_1")
+        self.assertEqual(result["1"]["inputs"]["worker_id"], "worker-b")
 
     def test_master_sets_is_worker_false(self):
         result = _apply(self._value_prompt(), "master")
@@ -579,62 +545,6 @@ class ApplyOverridesValueTests(unittest.TestCase):
         self.assertEqual(result["1"]["inputs"]["worker_id"], "")
 
 
-class ApplyOverridesListSplitterTests(unittest.TestCase):
-    def _splitter_prompt(self):
-        return {"1": {"class_type": "DistributedListSplitter", "inputs": {}}}
-
-    def test_master_gets_participant_index_zero(self):
-        result = _apply(self._splitter_prompt(), "master", enabled_worker_ids=["worker-a", "worker-b"])
-        self.assertEqual(result["1"]["inputs"]["participant_index"], 0)
-        self.assertEqual(result["1"]["inputs"]["total_participants"], 3)
-        self.assertFalse(result["1"]["inputs"]["is_worker"])
-        self.assertNotIn("master_url", result["1"]["inputs"])
-        self.assertEqual(result["1"]["inputs"]["worker_id"], "master")
-        self.assertEqual(result["1"]["inputs"]["multi_job_id"], "run_1")
-
-    def test_worker_gets_incremented_participant_index(self):
-        result = _apply(self._splitter_prompt(), "worker-a", enabled_worker_ids=["worker-a", "worker-b"])
-        self.assertEqual(result["1"]["inputs"]["participant_index"], 1)
-        self.assertEqual(result["1"]["inputs"]["total_participants"], 3)
-        self.assertTrue(result["1"]["inputs"]["is_worker"])
-        self.assertEqual(result["1"]["inputs"]["worker_id"], "worker-a")
-        self.assertEqual(result["1"]["inputs"]["master_url"], "http://master.example.com")
-
-    def test_delegate_mode_shifts_worker_to_index_zero(self):
-        result = _apply(
-            self._splitter_prompt(),
-            "worker-a",
-            enabled_worker_ids=["worker-a", "worker-b"],
-            delegate_master=True,
-        )
-        self.assertEqual(result["1"]["inputs"]["participant_index"], 0)
-        self.assertEqual(result["1"]["inputs"]["total_participants"], 2)
-
-
-class ApplyOverridesListCollectorTests(unittest.TestCase):
-    def _collector_prompt(self):
-        return {"1": {"class_type": "DistributedListCollector", "inputs": {}}}
-
-    def test_worker_sets_master_url_and_worker_id(self):
-        result = _apply(self._collector_prompt(), "worker-a", enabled_worker_ids=["worker-a"])
-        self.assertTrue(result["1"]["inputs"]["is_worker"])
-        self.assertEqual(result["1"]["inputs"]["master_url"], "http://master.example.com")
-        self.assertEqual(result["1"]["inputs"]["worker_id"], "worker-a")
-
-    def test_master_sets_delegate_only_and_clears_worker_fields(self):
-        prompt = {
-            "1": {
-                "class_type": "DistributedListCollector",
-                "inputs": {"master_url": "stale", "worker_id": "stale"},
-            }
-        }
-        result = _apply(prompt, "master", enabled_worker_ids=["worker-a"], delegate_master=True)
-        self.assertFalse(result["1"]["inputs"]["is_worker"])
-        self.assertTrue(result["1"]["inputs"]["delegate_only"])
-        self.assertNotIn("master_url", result["1"]["inputs"])
-        self.assertNotIn("worker_id", result["1"]["inputs"])
-
-
 class ApplyOverridesBranchTests(unittest.TestCase):
     def test_master_gets_branch_zero_worker_gets_branch_one(self):
         prompt = _branch_prompt()

tests/unit/system/test_queue_request.py

@@ -26,27 +26,26 @@ class QueueRequestPayloadTests(unittest.TestCase):
 
     def test_normalizes_enabled_worker_ids(self):
         payload_data = self._base_payload()
-        payload_data["enabled_worker_ids"] = ["a", 2, 3]
+        payload_data["enabled_worker_ids"] = [" worker-a ", "worker-b", "worker-a"]
         payload_data["delegate_master"] = True
         payload = parse_queue_request_payload(
             payload_data
         )
-        self.assertEqual(payload.enabled_worker_ids, ["a", "2", "3"])
+        self.assertEqual(payload.enabled_worker_ids, ["worker-a", "worker-b"])
         self.assertTrue(payload.delegate_master)
 
     def test_supports_legacy_workers_field(self):
         payload_data = self._base_payload()
         payload_data.pop("enabled_worker_ids", None)
-        payload_data["workers"] = [{"id": "w1"}, "w2", {"id": 3}, {"name": "no-id"}]
+        payload_data["workers"] = [{"id": "w1"}, "w2", {"id": "w3"}, {"name": "no-id"}]
         payload = parse_queue_request_payload(
             payload_data
         )
-        self.assertEqual(payload.enabled_worker_ids, ["w1", "w2", "3"])
+        self.assertEqual(payload.enabled_worker_ids, ["w1", "w2", "w3"])
 
-    def test_supports_auto_prepare_prompt_fallback(self):
+    def test_supports_workflow_prompt_fallback(self):
         payload_data = self._base_payload()
         payload_data.pop("prompt", None)
-        payload_data["auto_prepare"] = True
         payload_data["workflow"] = {
             "prompt": {
                 "10": {"class_type": "DistributedCollector"},
@@ -56,7 +55,6 @@ class QueueRequestPayloadTests(unittest.TestCase):
             payload_data
         )
         self.assertIn("10", payload.prompt)
-        self.assertTrue(payload.auto_prepare)
 
     def test_normalizes_trace_execution_id(self):
         payload_data = self._base_payload()
@@ -74,10 +72,6 @@ class QueueRequestPayloadTests(unittest.TestCase):
         )
         self.assertIsNone(payload.trace_execution_id)
 
-    def test_auto_prepare_defaults_true(self):
-        payload = parse_queue_request_payload(self._base_payload())
-        self.assertTrue(payload.auto_prepare)
-
     def test_workers_field_must_be_list(self):
         payload_data = self._base_payload()
         payload_data.pop("enabled_worker_ids", None)
@@ -91,7 +85,7 @@ class QueueRequestPayloadTests(unittest.TestCase):
         with self.assertRaisesRegex(ValueError, "trace_execution_id must be a string"):
             parse_queue_request_payload(payload_data)
 
-    def test_auto_prepare_false_still_falls_back_to_workflow_prompt(self):
+    def test_auto_prepare_is_ignored_for_backward_compat(self):
         payload_data = self._base_payload()
         payload_data.pop("prompt", None)
         payload_data["auto_prepare"] = False
@@ -100,13 +94,6 @@ class QueueRequestPayloadTests(unittest.TestCase):
         }
         payload = parse_queue_request_payload(payload_data)
         self.assertIn("10", payload.prompt)
-        self.assertFalse(payload.auto_prepare)
-
-    def test_auto_prepare_must_be_boolean(self):
-        payload_data = self._base_payload()
-        payload_data["auto_prepare"] = "true"
-        with self.assertRaisesRegex(ValueError, "auto_prepare must be a boolean"):
-            parse_queue_request_payload(payload_data)
 
     def test_invalid_delegate_master_type_raises(self):
         payload_data = self._base_payload()
@@ -120,6 +107,12 @@ class QueueRequestPayloadTests(unittest.TestCase):
         with self.assertRaisesRegex(ValueError, "enabled_worker_ids must be a list"):
             parse_queue_request_payload(payload_data)
 
+    def test_legacy_index_worker_tokens_are_rejected(self):
+        payload_data = self._base_payload()
+        payload_data["enabled_worker_ids"] = ["worker-a", "0", "worker_1"]
+        with self.assertRaisesRegex(ValueError, "legacy index token"):
+            parse_queue_request_payload(payload_data)
+
     def test_invalid_top_level_payload_raises(self):
         with self.assertRaisesRegex(ValueError, "Expected a JSON object body"):
             parse_queue_request_payload(["not", "an", "object"])

tests/unit/upscale/test_dynamic.py

@@ -4,6 +4,7 @@ import sys
 import types
 import unittest
 from pathlib import Path
+from types import SimpleNamespace
 
 import numpy as np
 from PIL import Image
@@ -67,6 +68,8 @@ def _bootstrap_package(package_name):
     sys.modules[f"{package_name}.utils.config"] = config_module
 
     constants_module = types.ModuleType(f"{package_name}.utils.constants")
+    constants_module.JOB_POLL_INTERVAL = 0.0
+    constants_module.JOB_POLL_MAX_ATTEMPTS = 3
     constants_module.TILE_WAIT_TIMEOUT = 5.0
     constants_module.TILE_SEND_TIMEOUT = 5.0
     sys.modules[f"{package_name}.utils.constants"] = constants_module
@@ -110,6 +113,7 @@ def _load_dynamic_mode_module():
     package_name = "dist_dynamic_mode_testpkg"
     _bootstrap_package(package_name)
     _load_module(package_name, "upscale/processing_args.py", "upscale.processing_args")
+    _load_module(package_name, "upscale/mode_contexts.py", "upscale.mode_contexts")
     return _load_module(package_name, "upscale/modes/dynamic.py", "upscale.modes.dynamic")
 
 
@@ -143,24 +147,33 @@ class _DummyDynamicNode(dynamic_module.DynamicModeMixin):
     def _poll_job_ready(self, *_args, **_kwargs):
         return self._poll_ready
 
-    async def _request_image_from_master(self, *_args, **_kwargs):
-        if self._assignments:
-            return self._assignments.pop(0)
-        return (None, 0)
+    async def check_job_status(self, *_args, **_kwargs):
+        return self._poll_ready
 
-    def _slice_conditioning(self, positive, negative, _image_idx):
+    async def request_assignment(self, *_args, **_kwargs):
+        if self._assignments:
+            image_idx, estimated_remaining = self._assignments.pop(0)
+            return SimpleNamespace(
+                kind="image" if image_idx is not None else "none",
+                task_idx=image_idx,
+                estimated_remaining=estimated_remaining,
+                batched_static=False,
+            )
+        return SimpleNamespace(kind="none", task_idx=None, estimated_remaining=0, batched_static=False)
+
+    def slice_conditioning(self, positive, negative, _image_idx):
         return positive, negative
 
-    def _process_and_blend_tile(self, _tile_idx, _pos, _source_tensor, local_image, *_args, **_kwargs):
+    def process_and_blend_tile(self, _tile_idx, _pos, _source_tensor, local_image, *_args, **_kwargs):
         return local_image
 
-    async def _send_heartbeat_to_master(self, *_args, **_kwargs):
+    async def send_heartbeat(self, *_args, **_kwargs):
         self.heartbeat_count += 1
 
-    async def _send_full_image_to_master(self, local_image, image_idx, _multi_job_id, _master_url, _worker_id, is_last):
+    async def send_full_image(self, local_image, image_idx, _multi_job_id, _master_url, _worker_id, is_last):
         self.sent_images.append((image_idx, is_last, local_image.size))
 
-    async def _send_worker_complete_signal(self, *_args, **_kwargs):
+    async def send_worker_complete_signal(self, *_args, **_kwargs):
         self.completion_sent = True
 
 

tests/unit/upscale/test_job_store.py

@@ -27,9 +27,33 @@ def _bootstrap_package(package_name):
     logging_module.debug_log = lambda *_args, **_kwargs: None
     sys.modules[f"{package_name}.utils.logging"] = logging_module
 
-    server_module = types.ModuleType("server")
-    server_module.PromptServer = types.SimpleNamespace(instance=types.SimpleNamespace())
-    sys.modules["server"] = server_module
+    prompt_server = types.SimpleNamespace()
+
+    runtime_state_module = types.ModuleType(f"{package_name}.utils.runtime_state")
+
+    def _ensure_distributed_runtime_state(server_instance=None):
+        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()
+        if not hasattr(ps, "distributed_job_allowed_workers"):
+            ps.distributed_job_allowed_workers = {}
+        if not hasattr(ps, "distributed_pending_tile_jobs"):
+            ps.distributed_pending_tile_jobs = {}
+        if not hasattr(ps, "distributed_tile_jobs_lock"):
+            ps.distributed_tile_jobs_lock = asyncio.Lock()
+        return types.SimpleNamespace(
+            distributed_pending_jobs=ps.distributed_pending_jobs,
+            distributed_jobs_lock=ps.distributed_jobs_lock,
+            distributed_job_allowed_workers=ps.distributed_job_allowed_workers,
+            distributed_pending_tile_jobs=ps.distributed_pending_tile_jobs,
+            distributed_tile_jobs_lock=ps.distributed_tile_jobs_lock,
+        )
+
+    runtime_state_module.ensure_distributed_runtime_state = _ensure_distributed_runtime_state
+    runtime_state_module.get_prompt_server_instance = lambda: prompt_server
+    sys.modules[f"{package_name}.utils.runtime_state"] = runtime_state_module
 
 
 def _load_module(package_name, module_rel_path, module_name):
@@ -58,7 +82,7 @@ job_models_module, job_store_module = _load_job_store_modules()
 
 class JobStoreTests(unittest.IsolatedAsyncioTestCase):
     async def asyncSetUp(self):
-        self.prompt_server = job_store_module.server.PromptServer.instance
+        self.prompt_server = job_store_module.get_prompt_server_instance()
         self.prompt_server.distributed_pending_tile_jobs = {}
         self.prompt_server.distributed_tile_jobs_lock = asyncio.Lock()
 

tests/unit/upscale/test_single_gpu.py

@@ -105,6 +105,7 @@ def _load_single_gpu_mode_module():
     _bootstrap_package(package_name)
     _load_module(package_name, "upscale/processing_args.py", "upscale.processing_args")
     _load_module(package_name, "upscale/tile_processing.py", "upscale.tile_processing")
+    _load_module(package_name, "upscale/mode_contexts.py", "upscale.mode_contexts")
     return _load_module(package_name, "upscale/modes/single_gpu.py", "upscale.modes.single_gpu")
 
 

tests/unit/upscale/test_static.py

@@ -4,6 +4,7 @@ import sys
 import types
 import unittest
 from pathlib import Path
+from types import SimpleNamespace
 
 import torch
 
@@ -166,6 +167,16 @@ def _load_static_mode_module():
     sys.modules[tile_processing_spec.name] = tile_processing_module
     tile_processing_spec.loader.exec_module(tile_processing_module)
 
+    mode_contexts_path = Path(__file__).resolve().parents[3] / "upscale" / "mode_contexts.py"
+    mode_contexts_spec = importlib.util.spec_from_file_location(
+        f"{package_name}.upscale.mode_contexts",
+        mode_contexts_path,
+    )
+    mode_contexts_module = importlib.util.module_from_spec(mode_contexts_spec)
+    assert mode_contexts_spec is not None and mode_contexts_spec.loader is not None
+    sys.modules[mode_contexts_spec.name] = mode_contexts_module
+    mode_contexts_spec.loader.exec_module(mode_contexts_module)
+
     spec = importlib.util.spec_from_file_location(f"{package_name}.upscale.modes.static", module_path)
     module = importlib.util.module_from_spec(spec)
     assert spec is not None and spec.loader is not None
@@ -198,14 +209,20 @@ class _FakeStaticWorker(static_mode.StaticModeMixin):
     def _poll_job_ready(self, *_args, **_kwargs):
         return self.job_ready
 
-    async def _request_tile_from_master(self, *_args, **_kwargs):
+    async def request_assignment(self, *_args, **_kwargs):
         self.request_calls += 1
-        return self.tile_sequence.pop(0)
+        tile_idx, estimated_remaining, batched_static = self.tile_sequence.pop(0)
+        return SimpleNamespace(
+            kind="tile" if tile_idx is not None else "none",
+            task_idx=tile_idx,
+            estimated_remaining=estimated_remaining,
+            batched_static=batched_static,
+        )
 
-    async def _send_heartbeat_to_master(self, *_args, **_kwargs):
+    async def send_heartbeat(self, *_args, **_kwargs):
         self.heartbeat_calls += 1
 
-    async def send_tiles_batch_to_master(
+    async def send_tiles_batch(
         self,
         processed_tiles,
         _multi_job_id,

upscale/job_state.py

@@ -22,6 +22,10 @@ class JobStateMixin:
             return dict(job_data.completed_images)
         return {}
 
+    async def all_completed_tasks(self, multi_job_id):
+        """Public completed-task accessor."""
+        return await self._get_all_completed_tasks(multi_job_id)
+
     async def _get_next_image_index(self, multi_job_id):
         """Get next image index from pending queue for master."""
         prompt_server = ensure_tile_jobs_initialized()
@@ -39,6 +43,10 @@ class JobStateMixin:
         except asyncio.TimeoutError:
             return None
 
+    async def next_image_index(self, multi_job_id):
+        """Public next-image accessor."""
+        return await self._get_next_image_index(multi_job_id)
+
     async def _get_next_tile_index(self, multi_job_id):
         """Get next tile index from pending queue for master in static mode."""
         prompt_server = ensure_tile_jobs_initialized()
@@ -56,6 +64,10 @@ class JobStateMixin:
         except asyncio.TimeoutError:
             return None
 
+    async def next_tile_index(self, multi_job_id):
+        """Public next-tile accessor."""
+        return await self._get_next_tile_index(multi_job_id)
+
     async def _get_total_completed_count(self, multi_job_id):
         """Get total count of all completed images (master + workers)."""
         prompt_server = ensure_tile_jobs_initialized()
@@ -67,6 +79,10 @@ class JobStateMixin:
                 return len(job_data.completed_tasks)
             return 0
 
+    async def total_completed_count(self, multi_job_id):
+        """Public completed-count accessor."""
+        return await self._get_total_completed_count(multi_job_id)
+
     async def _get_all_completed_images(self, multi_job_id):
         """Get all completed images."""
         prompt_server = ensure_tile_jobs_initialized()
@@ -76,6 +92,10 @@ class JobStateMixin:
                 return job_data.completed_images.copy()
             return {}
 
+    async def all_completed_images(self, multi_job_id):
+        """Public completed-image accessor."""
+        return await self._get_all_completed_images(multi_job_id)
+
     async def _get_pending_count(self, multi_job_id):
         """Get count of pending images in the queue."""
         prompt_server = ensure_tile_jobs_initialized()
@@ -87,6 +107,10 @@ class JobStateMixin:
                 return job_data.pending_tasks.qsize()
             return 0
 
+    async def pending_count(self, multi_job_id):
+        """Public pending-count accessor."""
+        return await self._get_pending_count(multi_job_id)
+
     async def _drain_worker_results_queue(self, multi_job_id):
         """Drain pending worker results from queue and update completed images."""
         prompt_server = ensure_tile_jobs_initialized()
@@ -130,6 +154,14 @@ class JobStateMixin:
 
         return collected
 
+    async def drain_worker_results_queue(self, multi_job_id):
+        """Public worker-result drain API."""
+        return await self._drain_worker_results_queue(multi_job_id)
+
     async def _check_and_requeue_timed_out_workers(self, multi_job_id, batch_size):
         """Check for timed out workers and requeue their assigned images."""
         return await check_and_requeue_timed_out_workers(multi_job_id, batch_size)
+
+    async def check_and_requeue_timed_out_workers(self, multi_job_id, batch_size):
+        """Public timeout-requeue API."""
+        return await self._check_and_requeue_timed_out_workers(multi_job_id, batch_size)

upscale/job_store.py

@@ -4,9 +4,8 @@ import time
 from dataclasses import dataclass
 from typing import Any
 
-import server
-
 from ..utils.logging import debug_log
+from ..utils.runtime_state import ensure_distributed_runtime_state, get_prompt_server_instance
 from .job_models import BaseJobState, ImageJobState, TileJobState
 
 # Configure maximum payload size (50MB default, configurable via environment variable)
@@ -25,20 +24,21 @@ class JobQueueInitConfig:
 
 def ensure_tile_jobs_initialized() -> Any:
     """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]
+    prompt_server = get_prompt_server_instance()
+    state = ensure_distributed_runtime_state(prompt_server)
+    if not isinstance(state.distributed_pending_tile_jobs, dict):
+        debug_log("Resetting invalid distributed_pending_tile_jobs state to empty dict.")
+        state.distributed_pending_tile_jobs = {}
+        ensure_distributed_runtime_state(prompt_server)
+
+    invalid_job_ids = [
+        job_id
+        for job_id, job_data in state.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 state.distributed_pending_tile_jobs[job_id]
     return prompt_server
 
 

upscale/mode_contexts.py

@@ -0,0 +1,238 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+
+class TileOpsCollaborator:
+    """Typed facade over tile-related operations used by USDU modes."""
+
+    def __init__(self, delegate: Any) -> None:
+        self._delegate = delegate
+
+    def round_to_multiple(self, value: int) -> int:
+        return self._delegate.round_to_multiple(value)
+
+    def calculate_tiles(
+        self,
+        width: int,
+        height: int,
+        tile_width: int,
+        tile_height: int,
+        force_uniform_tiles: bool,
+    ) -> list[tuple[int, int]]:
+        return self._delegate.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles)
+
+    def create_tile_mask(
+        self,
+        width: int,
+        height: int,
+        tx: int,
+        ty: int,
+        tile_width: int,
+        tile_height: int,
+        mask_blur: int,
+    ) -> Any:
+        return self._delegate.create_tile_mask(width, height, tx, ty, tile_width, tile_height, mask_blur)
+
+    def blend_tile(
+        self,
+        base_image: Any,
+        tile_pil: Any,
+        x1: int,
+        y1: int,
+        size: tuple[int, int],
+        tile_mask: Any,
+        padding: int,
+    ) -> Any:
+        return self._delegate.blend_tile(base_image, tile_pil, x1, y1, size, tile_mask, padding)
+
+    def slice_conditioning(self, positive: Any, negative: Any, image_idx: int) -> tuple[Any, Any]:
+        return self._delegate.slice_conditioning(positive, negative, image_idx)
+
+    def process_and_blend_tile(
+        self,
+        tile_idx: int,
+        tile_pos: tuple[int, int],
+        upscaled_image: Any,
+        result_image: Any,
+        model: Any,
+        positive: Any,
+        negative: Any,
+        vae: Any,
+        seed: int,
+        steps: int,
+        cfg: float,
+        sampler_name: str,
+        scheduler: str,
+        denoise: float,
+        tile_width: int,
+        tile_height: int,
+        padding: int,
+        mask_blur: int,
+        image_width: int,
+        image_height: int,
+        force_uniform_tiles: bool,
+        tiled_decode: bool,
+        batch_idx: int,
+    ) -> Any:
+        return self._delegate.process_and_blend_tile(
+            tile_idx,
+            tile_pos,
+            upscaled_image,
+            result_image,
+            model,
+            positive,
+            negative,
+            vae,
+            seed,
+            steps,
+            cfg,
+            sampler_name,
+            scheduler,
+            denoise,
+            tile_width,
+            tile_height,
+            padding,
+            mask_blur,
+            image_width,
+            image_height,
+            force_uniform_tiles,
+            tiled_decode,
+            batch_idx=batch_idx,
+        )
+
+
+class JobStateCollaborator:
+    """Typed facade over job-state operations used by USDU modes."""
+
+    def __init__(self, delegate: Any) -> None:
+        self._delegate = delegate
+
+    async def next_tile_index(self, multi_job_id: str) -> int | None:
+        return await self._delegate.next_tile_index(multi_job_id)
+
+    async def all_completed_tasks(self, multi_job_id: str) -> dict[int, dict[str, Any]]:
+        return await self._delegate.all_completed_tasks(multi_job_id)
+
+    async def check_and_requeue_timed_out_workers(self, multi_job_id: str, expected_total: int) -> int:
+        return await self._delegate.check_and_requeue_timed_out_workers(multi_job_id, expected_total)
+
+    async def next_image_index(self, multi_job_id: str) -> int | None:
+        return await self._delegate.next_image_index(multi_job_id)
+
+    async def drain_worker_results_queue(self, multi_job_id: str) -> int:
+        return await self._delegate.drain_worker_results_queue(multi_job_id)
+
+    async def total_completed_count(self, multi_job_id: str) -> int:
+        return await self._delegate.total_completed_count(multi_job_id)
+
+    async def pending_count(self, multi_job_id: str) -> int:
+        return await self._delegate.pending_count(multi_job_id)
+
+    async def all_completed_images(self, multi_job_id: str) -> dict[int, Any]:
+        return await self._delegate.all_completed_images(multi_job_id)
+
+
+class WorkerCommsCollaborator:
+    """Typed facade over worker/master communication operations used by USDU modes."""
+
+    def __init__(self, delegate: Any) -> None:
+        self._delegate = delegate
+
+    async def request_assignment(self, multi_job_id: str, master_url: str, worker_id: str) -> Any:
+        return await self._delegate.request_assignment(multi_job_id, master_url, worker_id)
+
+    async def send_tiles_batch(
+        self,
+        processed_tiles: list[dict[str, Any]],
+        multi_job_id: str,
+        master_url: str,
+        padding: int,
+        worker_id: str,
+        is_final_flush: bool = False,
+    ) -> None:
+        await self._delegate.send_tiles_batch(
+            processed_tiles,
+            multi_job_id,
+            master_url,
+            padding,
+            worker_id,
+            is_final_flush=is_final_flush,
+        )
+
+    async def send_heartbeat(self, multi_job_id: str, master_url: str, worker_id: str) -> None:
+        await self._delegate.send_heartbeat(multi_job_id, master_url, worker_id)
+
+    async def check_job_status(self, multi_job_id: str, master_url: str) -> bool:
+        return await self._delegate.check_job_status(multi_job_id, master_url)
+
+    async def async_yield(self) -> None:
+        await self._delegate.async_yield()
+
+    async def send_full_image(
+        self,
+        image_pil: Any,
+        image_idx: int,
+        multi_job_id: str,
+        master_url: str,
+        worker_id: str,
+        is_last: bool,
+    ) -> None:
+        await self._delegate.send_full_image(
+            image_pil,
+            image_idx,
+            multi_job_id,
+            master_url,
+            worker_id,
+            is_last,
+        )
+
+    async def send_worker_complete_signal(self, multi_job_id: str, master_url: str, worker_id: str) -> None:
+        await self._delegate.send_worker_complete_signal(multi_job_id, master_url, worker_id)
+
+
+class ResultCollectorCollaborator:
+    """Typed facade over result-collection operations used by dynamic mode."""
+
+    def __init__(self, delegate: Any) -> None:
+        self._delegate = delegate
+
+    async def mark_image_completed(self, multi_job_id: str, image_idx: int, image_pil: Any) -> None:
+        await self._delegate.mark_image_completed(multi_job_id, image_idx, image_pil)
+
+    async def collect_dynamic_images(
+        self,
+        multi_job_id: str,
+        remaining_to_collect: int,
+        num_workers: int,
+        batch_size: int,
+        master_processed_count: int,
+    ) -> dict[int, Any]:
+        return await self._delegate.collect_dynamic_images(
+            multi_job_id,
+            remaining_to_collect,
+            num_workers,
+            batch_size,
+            master_processed_count,
+        )
+
+
+@dataclass(frozen=True)
+class SingleGpuModeContext:
+    tile_ops: TileOpsCollaborator
+
+
+@dataclass(frozen=True)
+class StaticModeContext:
+    tile_ops: TileOpsCollaborator
+    job_state: JobStateCollaborator
+    worker_comms: WorkerCommsCollaborator
+
+
+@dataclass(frozen=True)
+class DynamicModeContext:
+    tile_ops: TileOpsCollaborator
+    job_state: JobStateCollaborator
+    worker_comms: WorkerCommsCollaborator
+    result_collector: ResultCollectorCollaborator

upscale/modes/dynamic.py

@@ -1,14 +1,26 @@
 from __future__ import annotations
 
-import asyncio, torch
+import asyncio, time, torch
 from typing import Any
 from PIL import Image
 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 TILE_WAIT_TIMEOUT, TILE_SEND_TIMEOUT
+from ...utils.constants import (
+    JOB_POLL_INTERVAL,
+    JOB_POLL_MAX_ATTEMPTS,
+    TILE_WAIT_TIMEOUT,
+    TILE_SEND_TIMEOUT,
+)
 from ..job_store import ensure_tile_jobs_initialized, init_dynamic_job
+from ..mode_contexts import (
+    DynamicModeContext,
+    JobStateCollaborator,
+    ResultCollectorCollaborator,
+    TileOpsCollaborator,
+    WorkerCommsCollaborator,
+)
 from ..processing_args import UpscaleCoreArgs
 
 
@@ -17,47 +29,82 @@ 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`).
+    - 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`).
+    - WorkerCommsMixin (`request_assignment`, `send_full_image`, `send_heartbeat`).
     """
 
+    def _build_dynamic_mode_context(self) -> DynamicModeContext:
+        """Build explicit collaborators for dynamic mode execution."""
+        return DynamicModeContext(
+            tile_ops=TileOpsCollaborator(self),
+            job_state=JobStateCollaborator(self),
+            worker_comms=WorkerCommsCollaborator(self),
+            result_collector=ResultCollectorCollaborator(self),
+        )
+
+    def _poll_job_ready_dynamic(
+        self,
+        multi_job_id: str,
+        master_url: str,
+        worker_id: str | None = None,
+        max_attempts: int = JOB_POLL_MAX_ATTEMPTS,
+        mode_context: DynamicModeContext | None = None,
+    ) -> bool:
+        """Poll master for job readiness using explicit worker-comms collaborator."""
+        context = mode_context or self._build_dynamic_mode_context()
+        for attempt in range(max_attempts):
+            ready = run_async_in_server_loop(
+                context.worker_comms.check_job_status(multi_job_id, master_url),
+                timeout=5.0,
+            )
+            if ready:
+                if worker_id:
+                    debug_log(f"Worker[{worker_id[:8]}] job {multi_job_id} ready after {attempt} attempts")
+                else:
+                    debug_log(f"Job {multi_job_id} ready after {attempt} attempts")
+                return True
+            time.sleep(JOB_POLL_INTERVAL)
+        return False
+
     def _handle_master_dynamic_idle_state(
         self,
         multi_job_id,
         batch_size,
         consecutive_retries,
         max_consecutive_retries,
+        mode_context: DynamicModeContext | None = None,
     ):
         """Handle queue-empty state while master waits for worker progress."""
+        context = mode_context or self._build_dynamic_mode_context()
         drained_count = run_async_in_server_loop(
-            self._drain_worker_results_queue(multi_job_id),
+            context.job_state.drain_worker_results_queue(multi_job_id),
             timeout=5.0,
         )
-        run_async_in_server_loop(self._async_yield(), timeout=0.1)
+        run_async_in_server_loop(context.worker_comms.async_yield(), timeout=0.1)
 
         requeued_count = run_async_in_server_loop(
-            self._check_and_requeue_timed_out_workers(multi_job_id, batch_size),
+            context.job_state.check_and_requeue_timed_out_workers(multi_job_id, batch_size),
             timeout=5.0,
         )
-        run_async_in_server_loop(self._async_yield(), timeout=0.1)
+        run_async_in_server_loop(context.worker_comms.async_yield(), timeout=0.1)
 
         if requeued_count > 0:
             log(f"Requeued {requeued_count} images from timed out workers")
             return False, True, 0
 
         completed_now = run_async_in_server_loop(
-            self._get_total_completed_count(multi_job_id),
+            context.job_state.total_completed_count(multi_job_id),
             timeout=1.0,
         )
         log(f"USDU Dist: Images progress {completed_now}/{batch_size}")
         if completed_now >= batch_size:
             return True, False, consecutive_retries
 
-        run_async_in_server_loop(self._async_yield(), timeout=0.1)
+        run_async_in_server_loop(context.worker_comms.async_yield(), timeout=0.1)
 
         pending_count = run_async_in_server_loop(
-            self._get_pending_count(multi_job_id),
+            context.job_state.pending_count(multi_job_id),
             timeout=1.0,
         )
         if pending_count > 0:
@@ -82,10 +129,12 @@ class DynamicModeMixin:
         processed_count,
         result_images,
         upscaled_image,
+        mode_context: DynamicModeContext | None = None,
     ):
         """Collect remaining worker outputs and convert final images back to tensor."""
+        context = mode_context or self._build_dynamic_mode_context()
         all_completed = run_async_in_server_loop(
-            self._get_all_completed_images(multi_job_id),
+            context.job_state.all_completed_images(multi_job_id),
             timeout=5.0,
         )
         remaining_to_collect = batch_size - len(all_completed)
@@ -93,7 +142,7 @@ class DynamicModeMixin:
             debug_log(f"Waiting for {remaining_to_collect} more images from workers")
             collection_timeout = float(get_worker_timeout_seconds())
             collected_images = run_async_in_server_loop(
-                self._async_collect_dynamic_images(
+                context.result_collector.collect_dynamic_images(
                     multi_job_id,
                     remaining_to_collect,
                     num_workers,
@@ -128,8 +177,10 @@ class DynamicModeMixin:
         force_uniform_tiles: bool,
         multi_job_id: str,
         enabled_workers: list[str],
+        mode_context: DynamicModeContext | None = None,
     ) -> tuple[torch.Tensor]:
         """Dynamic mode for large batches - assigns whole images to workers dynamically, including master."""
+        context = mode_context or self._build_dynamic_mode_context()
         # Get batch size and dimensions
         batch_size, height, width, _ = upscaled_image.shape
         num_workers = len(enabled_workers)
@@ -142,7 +193,7 @@ class DynamicModeMixin:
         debug_log(f"Processing {batch_size} images dynamically across master + {num_workers} workers.")
         
         # Calculate tiles for processing
-        all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles)
+        all_tiles = context.tile_ops.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles)
         
         # Initialize job queue for communication
         try:
@@ -167,7 +218,7 @@ class DynamicModeMixin:
         while processed_count < batch_size:
             # Try to get an image to process
             image_idx = run_async_in_server_loop(
-                self._get_next_image_index(multi_job_id),
+                context.job_state.next_image_index(multi_job_id),
                 timeout=5.0  # Short timeout to allow frequent checks
             )
 
@@ -183,7 +234,7 @@ class DynamicModeMixin:
                 image_seed = core_args.seed
                 
                 # Pre-slice conditioning once per image (not per tile)
-                positive_sliced, negative_sliced = self._slice_conditioning(
+                positive_sliced, negative_sliced = context.tile_ops.slice_conditioning(
                     core_args.positive,
                     core_args.negative,
                     image_idx,
@@ -193,7 +244,7 @@ class DynamicModeMixin:
                     source_tensor = pil_to_tensor(local_image)
                     if single_tensor.is_cuda:
                         source_tensor = source_tensor.cuda()
-                    local_image = self._process_and_blend_tile(
+                    local_image = context.tile_ops.process_and_blend_tile(
                         tile_idx, pos, source_tensor, local_image,
                         core_args.model,
                         positive_sliced,
@@ -213,20 +264,20 @@ class DynamicModeMixin:
                     )
                     
                     # Yield after each tile to minimize worker downtime
-                    run_async_in_server_loop(self._async_yield(), timeout=0.1)
+                    run_async_in_server_loop(context.worker_comms.async_yield(), timeout=0.1)
                     # Note: No per-tile drain here – that's what makes this "per-image"
                 
                 result_images[image_idx] = local_image
                 
                 # Mark as completed
                 run_async_in_server_loop(
-                    self._mark_image_completed(multi_job_id, image_idx, local_image),
+                    context.result_collector.mark_image_completed(multi_job_id, image_idx, local_image),
                     timeout=5.0
                 )
                 
                 # NEW: Drain after the full image is marked complete (catches workers who finished during master's processing)
                 drained_count = run_async_in_server_loop(
-                    self._drain_worker_results_queue(multi_job_id),
+                    context.job_state.drain_worker_results_queue(multi_job_id),
                     timeout=5.0
                 )
                 
@@ -235,19 +286,20 @@ class DynamicModeMixin:
                 
                 # NEW: Log overall progress (includes master's image + any drained workers)
                 completed_now = run_async_in_server_loop(
-                    self._get_total_completed_count(multi_job_id),
+                    context.job_state.total_completed_count(multi_job_id),
                     timeout=1.0
                 )
                 log(f"USDU Dist: Images progress {completed_now}/{batch_size}")
                 
                 # Yield to allow workers to get new images after completing one
-                run_async_in_server_loop(self._async_yield(), timeout=0.1)
+                run_async_in_server_loop(context.worker_comms.async_yield(), timeout=0.1)
             else:
                 should_break, should_continue, consecutive_retries = self._handle_master_dynamic_idle_state(
                     multi_job_id=multi_job_id,
                     batch_size=batch_size,
                     consecutive_retries=consecutive_retries,
                     max_consecutive_retries=max_consecutive_retries,
+                    mode_context=context,
                 )
                 if should_break:
                     break
@@ -262,6 +314,7 @@ class DynamicModeMixin:
             processed_count=processed_count,
             result_images=result_images,
             upscaled_image=upscaled_image,
+            mode_context=context,
         )
         
         debug_log(f"UltimateSDUpscale Master - Job {multi_job_id} complete")
@@ -282,33 +335,43 @@ class DynamicModeMixin:
         worker_id: str,
         enabled_worker_ids: list[str] | str,
         dynamic_threshold: int,
+        mode_context: DynamicModeContext | None = None,
     ) -> tuple[torch.Tensor]:
         """Worker processing in dynamic mode - processes whole images."""
+        context = mode_context or self._build_dynamic_mode_context()
         # Round tile dimensions
-        tile_width = self.round_to_multiple(tile_width)
-        tile_height = self.round_to_multiple(tile_height)
+        tile_width = context.tile_ops.round_to_multiple(tile_width)
+        tile_height = context.tile_ops.round_to_multiple(tile_height)
 
         # Get dimensions and tile grid
         batch_size, height, width, _ = upscaled_image.shape
-        all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles)
+        all_tiles = context.tile_ops.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles)
         log(f"USDU Dist Worker[{worker_id[:8]}]: Processing image queue | Batch {batch_size}")
 
         # Keep track of processed images for is_last detection
         processed_count = 0
 
         # Poll for job readiness to avoid races during master init
-        max_poll_attempts = 20  # ~20s at 1s sleep
-        if not self._poll_job_ready(multi_job_id, master_url, worker_id=worker_id, max_attempts=max_poll_attempts):
+        max_poll_attempts = JOB_POLL_MAX_ATTEMPTS
+        if not self._poll_job_ready_dynamic(
+            multi_job_id,
+            master_url,
+            worker_id=worker_id,
+            max_attempts=max_poll_attempts,
+            mode_context=context,
+        ):
             log(f"Job {multi_job_id} not ready after {max_poll_attempts} attempts, aborting")
             return (upscaled_image,)
 
         # Loop to request and process images
         while True:
             # Request an image to process
-            image_idx, estimated_remaining = run_async_in_server_loop(
-                self._request_image_from_master(multi_job_id, master_url, worker_id),
-                timeout=TILE_WAIT_TIMEOUT
+            assignment = run_async_in_server_loop(
+                context.worker_comms.request_assignment(multi_job_id, master_url, worker_id),
+                timeout=TILE_WAIT_TIMEOUT,
             )
+            image_idx = assignment.task_idx if assignment.kind == "image" else None
+            estimated_remaining = assignment.estimated_remaining
 
             if image_idx is None:
                 debug_log(f"USDU Dist Worker - No more images to process")
@@ -330,7 +393,7 @@ class DynamicModeMixin:
             image_seed = core_args.seed
 
             # Pre-slice conditioning once per image (not per tile)
-            positive_sliced, negative_sliced = self._slice_conditioning(
+            positive_sliced, negative_sliced = context.tile_ops.slice_conditioning(
                 core_args.positive,
                 core_args.negative,
                 image_idx,
@@ -340,7 +403,7 @@ class DynamicModeMixin:
                 source_tensor = pil_to_tensor(local_image)
                 if single_tensor.is_cuda:
                     source_tensor = source_tensor.cuda()
-                local_image = self._process_and_blend_tile(
+                local_image = context.tile_ops.process_and_blend_tile(
                     tile_idx, pos, source_tensor, local_image,
                     core_args.model,
                     positive_sliced,
@@ -359,7 +422,7 @@ class DynamicModeMixin:
                     batch_idx=image_idx,
                 )
                 run_async_in_server_loop(
-                    self._send_heartbeat_to_master(multi_job_id, master_url, worker_id),
+                    context.worker_comms.send_heartbeat(multi_job_id, master_url, worker_id),
                     timeout=5.0
                 )
 
@@ -368,13 +431,19 @@ class DynamicModeMixin:
                 # 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),
+                    context.worker_comms.send_full_image(
+                        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),
+                    context.worker_comms.send_heartbeat(multi_job_id, master_url, worker_id),
                     timeout=5.0
                 )
                 if is_last:
@@ -387,7 +456,7 @@ class DynamicModeMixin:
         debug_log(f"Worker[{worker_id[:8]}] processed {processed_count} images, sending completion signal")
         try:
             run_async_in_server_loop(
-                self._send_worker_complete_signal(multi_job_id, master_url, worker_id),
+                context.worker_comms.send_worker_complete_signal(multi_job_id, master_url, worker_id),
                 timeout=TILE_SEND_TIMEOUT
             )
         except Exception as e:

upscale/modes/single_gpu.py

@@ -1,23 +1,19 @@
 from __future__ import annotations
 
 import math, torch
-from typing import Any, Protocol, cast
+from typing import Any
 from ...utils.logging import log
 from ...utils.image import blend_processed_batch_item, pil_to_tensor, tensor_to_pil
+from ..mode_contexts import SingleGpuModeContext, TileOpsCollaborator
 from ..processing_args import UpscaleCoreArgs
 from ..tile_processing import TileBatchArgs, extract_and_process_tile_batch
 
 
-class _SingleGpuOps(Protocol):
-    def round_to_multiple(self, value: int) -> int: ...
-    def calculate_tiles(self, width: int, height: int, tile_width: int, tile_height: int, force_uniform_tiles: bool): ...
-    def create_tile_mask(self, width: int, height: int, tx: int, ty: int, tile_width: int, tile_height: int, mask_blur: int): ...
-    def blend_tile(self, base_image: Any, tile_pil: Any, x1: int, y1: int, size: tuple[int, int], tile_mask: Any, padding: int): ...
-    def extract_batch_tile_with_padding(self, source_batch: Any, tx: int, ty: int, tile_width: int, tile_height: int, padding: int, force_uniform_tiles: bool): ...
-    def process_tiles_batch(self, *args: Any, **kwargs: Any): ...
-
-
 class SingleGpuModeMixin:
+    def _build_single_gpu_mode_context(self) -> SingleGpuModeContext:
+        """Build explicit collaborators for single-GPU mode execution."""
+        return SingleGpuModeContext(tile_ops=TileOpsCollaborator(self))
+
     def process_single_gpu(
         self,
         upscaled_image: torch.Tensor,
@@ -27,9 +23,11 @@ class SingleGpuModeMixin:
         padding: int,
         mask_blur: int,
         force_uniform_tiles: bool,
+        mode_context: SingleGpuModeContext | None = None,
     ) -> tuple[torch.Tensor]:
         """Process all tiles on a single GPU (no distribution), batching per tile like USDU."""
-        ops = cast(_SingleGpuOps, self)
+        context = mode_context or self._build_single_gpu_mode_context()
+        ops = context.tile_ops
         # Round tile dimensions
         tile_width = ops.round_to_multiple(tile_width)
         tile_height = ops.round_to_multiple(tile_height)

upscale/modes/static.py

@@ -18,6 +18,12 @@ from ..job_store import (
     cleanup_job, drain_results_queue, get_completed_count, mark_task_completed,
 )
 from ..job_models import TileJobState
+from ..mode_contexts import (
+    JobStateCollaborator,
+    StaticModeContext,
+    TileOpsCollaborator,
+    WorkerCommsCollaborator,
+)
 from ..processing_args import UpscaleCoreArgs
 from ..tile_processing import TileBatchArgs, extract_and_process_tile_batch
 
@@ -29,14 +35,30 @@ class StaticModeMixin:
     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`).
+    - WorkerCommsMixin (`send_tiles_batch`, `request_assignment`, `send_heartbeat`).
     """
 
-    def _poll_job_ready(self, multi_job_id, master_url, worker_id=None, max_attempts=JOB_POLL_MAX_ATTEMPTS):
+    def _build_static_mode_context(self) -> StaticModeContext:
+        """Build explicit collaborators for static mode execution."""
+        return StaticModeContext(
+            tile_ops=TileOpsCollaborator(self),
+            job_state=JobStateCollaborator(self),
+            worker_comms=WorkerCommsCollaborator(self),
+        )
+
+    def _poll_job_ready(
+        self,
+        multi_job_id,
+        master_url,
+        worker_id=None,
+        max_attempts=JOB_POLL_MAX_ATTEMPTS,
+        mode_context: StaticModeContext | None = None,
+    ):
         """Poll master for job readiness to avoid worker/master initialization race."""
+        context = mode_context or self._build_static_mode_context()
         for attempt in range(max_attempts):
             ready = run_async_in_server_loop(
-                self._check_job_status(multi_job_id, master_url),
+                context.worker_comms.check_job_status(multi_job_id, master_url),
                 timeout=5.0
             )
             if ready:
@@ -111,12 +133,14 @@ class StaticModeMixin:
         padding,
         worker_id,
         is_final_flush=False,
+        mode_context: StaticModeContext | None = None,
     ):
         """Send accumulated tile payloads to master and return a fresh accumulator."""
+        context = mode_context or self._build_static_mode_context()
         if not processed_tiles:
             if is_final_flush:
                 run_async_in_server_loop(
-                    self.send_tiles_batch_to_master(
+                    context.worker_comms.send_tiles_batch(
                         [],
                         multi_job_id,
                         master_url,
@@ -128,7 +152,7 @@ class StaticModeMixin:
                 )
             return processed_tiles
         run_async_in_server_loop(
-            self.send_tiles_batch_to_master(
+            context.worker_comms.send_tiles_batch(
                 processed_tiles,
                 multi_job_id,
                 master_url,
@@ -167,8 +191,10 @@ class StaticModeMixin:
         tiled_decode,
         width,
         height,
+        mode_context: StaticModeContext | None = None,
     ):
         """Process one tile_id across the batch and blend into result_images."""
+        context = mode_context or self._build_static_mode_context()
         source_batch = torch.cat([pil_to_tensor(img) for img in result_images], dim=0)
         if upscaled_image.is_cuda:
             source_batch = source_batch.cuda()
@@ -202,7 +228,7 @@ class StaticModeMixin:
                 result_images,
                 processed_batch,
                 b,
-                self.blend_tile,
+                context.tile_ops.blend_tile,
                 x1,
                 y1,
                 ew,
@@ -221,15 +247,17 @@ class StaticModeMixin:
                                     seed, steps, cfg, sampler_name, scheduler, denoise,
                                     tile_width, tile_height, padding, mask_blur,
                                     force_uniform_tiles, tiled_decode, multi_job_id, master_url,
-                                    worker_id, enabled_workers):
+                                    worker_id, enabled_workers,
+                                    mode_context: StaticModeContext | None = None):
         """Worker static mode processing with optional dynamic queue pulling."""
+        context = mode_context or self._build_static_mode_context()
         # Round tile dimensions
-        tile_width = self.round_to_multiple(tile_width)
-        tile_height = self.round_to_multiple(tile_height)
+        tile_width = context.tile_ops.round_to_multiple(tile_width)
+        tile_height = context.tile_ops.round_to_multiple(tile_height)
         
         # Get dimensions and calculate tiles
         _, height, width, _ = upscaled_image.shape
-        all_tiles = self.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles)
+        all_tiles = context.tile_ops.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles)
         num_tiles_per_image = len(all_tiles)
         batch_size = upscaled_image.shape[0]
         total_tiles = batch_size * num_tiles_per_image
@@ -241,24 +269,31 @@ class StaticModeMixin:
             working_images.append(image_pil.copy())
         tile_masks = []
         for tx, ty in all_tiles:
-            tile_masks.append(self.create_tile_mask(width, height, tx, ty, tile_width, tile_height, mask_blur))
+            tile_masks.append(context.tile_ops.create_tile_mask(width, height, tx, ty, tile_width, tile_height, mask_blur))
         
         # Dynamic queue mode (static processing): process batched-per-tile
         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 = JOB_POLL_MAX_ATTEMPTS
-        if not self._poll_job_ready(multi_job_id, master_url, worker_id=worker_id, max_attempts=max_poll_attempts):
+        if not self._poll_job_ready(
+            multi_job_id,
+            master_url,
+            worker_id=worker_id,
+            max_attempts=max_poll_attempts,
+            mode_context=context,
+        ):
             log(f"Job {multi_job_id} not ready after {max_poll_attempts} attempts, aborting")
             return (upscaled_image,)
 
         # Main processing loop - pull tile ids from queue
         while True:
             # Request a tile to process
-            tile_idx, estimated_remaining, batched_static = run_async_in_server_loop(
-                self._request_tile_from_master(multi_job_id, master_url, worker_id),
-                timeout=TILE_WAIT_TIMEOUT
+            assignment = run_async_in_server_loop(
+                context.worker_comms.request_assignment(multi_job_id, master_url, worker_id),
+                timeout=TILE_WAIT_TIMEOUT,
             )
+            tile_idx = assignment.task_idx if assignment.kind == "tile" else None
 
             if tile_idx is None:
                 debug_log(f"Worker[{worker_id[:8]}] - No more tiles to process")
@@ -298,7 +333,7 @@ class StaticModeMixin:
                 tile_pil = tensor_to_pil(processed_batch, b)
                 if tile_pil.size != (ew, eh):
                     tile_pil = tile_pil.resize((ew, eh), Image.LANCZOS)
-                working_images[b] = self.blend_tile(
+                working_images[b] = context.tile_ops.blend_tile(
                     working_images[b],
                     tile_pil,
                     x1,
@@ -322,7 +357,7 @@ class StaticModeMixin:
             # Send heartbeat
             try:
                 run_async_in_server_loop(
-                    self._send_heartbeat_to_master(multi_job_id, master_url, worker_id),
+                    context.worker_comms.send_heartbeat(multi_job_id, master_url, worker_id),
                     timeout=5.0
                 )
             except Exception as e:
@@ -331,20 +366,39 @@ class StaticModeMixin:
             # Send tiles in batches within loop
             if len(processed_tiles) >= MAX_BATCH:
                 processed_tiles = self._flush_tiles_to_master(
-                    processed_tiles, multi_job_id, master_url, padding, worker_id, is_final_flush=False
+                    processed_tiles,
+                    multi_job_id,
+                    master_url,
+                    padding,
+                    worker_id,
+                    is_final_flush=False,
+                    mode_context=context,
                 )
 
         # Send any remaining tiles
         processed_tiles = self._flush_tiles_to_master(
-            processed_tiles, multi_job_id, master_url, padding, worker_id, is_final_flush=True
+            processed_tiles,
+            multi_job_id,
+            master_url,
+            padding,
+            worker_id,
+            is_final_flush=True,
+            mode_context=context,
         )
         
         debug_log(f"Worker {worker_id} completed all assigned and requeued tiles")
         return (upscaled_image,)
 
-    async def _async_collect_and_monitor_static(self, multi_job_id, total_tiles, expected_total):
+    async def _async_collect_and_monitor_static(
+        self,
+        multi_job_id,
+        total_tiles,
+        expected_total,
+        mode_context: StaticModeContext | None = None,
+    ):
         """Async helper for collection and monitoring in static mode.
         Returns collected tasks dict. Caller should check if all tasks are complete."""
+        context = mode_context or self._build_static_mode_context()
         last_progress_log = time.time()
         progress_interval = 5.0
         last_heartbeat_check = time.time()
@@ -362,7 +416,7 @@ class StaticModeMixin:
             # Check and requeue timed-out workers periodically
             current_time = time.time()
             if current_time - last_heartbeat_check >= HEARTBEAT_INTERVAL:
-                requeued_count = await self._check_and_requeue_timed_out_workers(multi_job_id, expected_total)
+                requeued_count = await context.job_state.check_and_requeue_timed_out_workers(multi_job_id, expected_total)
                 if requeued_count > 0:
                     log(f"Requeued {requeued_count} tasks from timed-out workers")
                 last_heartbeat_check = current_time
@@ -395,7 +449,7 @@ class StaticModeMixin:
             await asyncio.sleep(0.1)
         
         # Get all completed tasks for return
-        return await self._get_all_completed_tasks(multi_job_id)
+        return await context.job_state.all_completed_tasks(multi_job_id)
 
     def _init_master_static_state(
         self,
@@ -408,7 +462,9 @@ class StaticModeMixin:
         tile_width,
         tile_height,
         mask_blur,
+        mode_context: StaticModeContext | None = None,
     ):
+        context = mode_context or self._build_static_mode_context()
         _, height, width, _ = upscaled_image.shape
         result_images = [tensor_to_pil(upscaled_image[b:b + 1], 0).copy() for b in range(batch_size)]
 
@@ -420,7 +476,7 @@ class StaticModeMixin:
         debug_log(f"Initialized tile-id queue with {num_tiles_per_image} ids for batch {batch_size}")
 
         tile_masks = [
-            self.create_tile_mask(width, height, tx, ty, tile_width, tile_height, mask_blur)
+            context.tile_ops.create_tile_mask(width, height, tx, ty, tile_width, tile_height, mask_blur)
             for tx, ty in all_tiles
         ]
         return result_images, tile_masks, width, height
@@ -452,14 +508,16 @@ class StaticModeMixin:
         tiled_decode,
         width,
         height,
+        mode_context: StaticModeContext | None = None,
     ):
+        context = mode_context or self._build_static_mode_context()
         processed_count = 0
         consecutive_no_tile = 0
         max_consecutive_no_tile = 2
 
         while processed_count < total_tiles:
             comfy.model_management.throw_exception_if_processing_interrupted()
-            tile_id = run_async_in_server_loop(self._get_next_tile_index(multi_job_id), timeout=5.0)
+            tile_id = run_async_in_server_loop(context.job_state.next_tile_index(multi_job_id), timeout=5.0)
             if tile_id is None:
                 consecutive_no_tile += 1
                 if consecutive_no_tile >= max_consecutive_no_tile:
@@ -495,6 +553,7 @@ class StaticModeMixin:
                 tiled_decode,
                 width,
                 height,
+                mode_context=context,
             )
             log(f"USDU Dist: Tiles progress {processed_count}/{total_tiles} (tile {tile_id})")
         return processed_count
@@ -527,14 +586,21 @@ class StaticModeMixin:
         tiled_decode,
         width,
         height,
+        mode_context: StaticModeContext | None = None,
     ):
+        context = mode_context or self._build_static_mode_context()
         remaining_tiles = total_tiles - master_processed_count
         if remaining_tiles <= 0:
-            return run_async_in_server_loop(self._get_all_completed_tasks(multi_job_id), timeout=5.0)
+            return run_async_in_server_loop(context.job_state.all_completed_tasks(multi_job_id), timeout=5.0)
 
         debug_log(f"Master waiting for {remaining_tiles} tiles from workers")
         collected_tasks = run_async_in_server_loop(
-            self._async_collect_and_monitor_static(multi_job_id, total_tiles, expected_total=total_tiles),
+            self._async_collect_and_monitor_static(
+                multi_job_id,
+                total_tiles,
+                expected_total=total_tiles,
+                mode_context=context,
+            ),
             timeout=None,
         )
 
@@ -545,7 +611,7 @@ class StaticModeMixin:
         log(f"Processing remaining {total_tiles - completed_count} tasks locally after worker failures")
         while True:
             comfy.model_management.throw_exception_if_processing_interrupted()
-            tile_id = run_async_in_server_loop(self._get_next_tile_index(multi_job_id), timeout=5.0)
+            tile_id = run_async_in_server_loop(context.job_state.next_tile_index(multi_job_id), timeout=5.0)
             if tile_id is None:
                 break
             self._master_process_one_tile(
@@ -574,6 +640,7 @@ class StaticModeMixin:
                 tiled_decode,
                 width,
                 height,
+                mode_context=context,
             )
         return collected_tasks
 
@@ -588,7 +655,9 @@ class StaticModeMixin:
         tile_width,
         tile_height,
         padding,
+        mode_context: StaticModeContext | None = None,
     ):
+        context = mode_context or self._build_static_mode_context()
         def _sort_key(item):
             global_idx, tile_data = item
             batch_idx = tile_data.get('batch_idx', global_idx // num_tiles_per_image)
@@ -609,7 +678,7 @@ class StaticModeMixin:
             tile_mask = tile_masks[tile_idx]
             extracted_width = tile_data.get('extracted_width', tile_width + 2 * padding)
             extracted_height = tile_data.get('extracted_height', tile_height + 2 * padding)
-            result_images[batch_idx] = self.blend_tile(
+            result_images[batch_idx] = context.tile_ops.blend_tile(
                 result_images[batch_idx],
                 tile_pil,
                 x,
@@ -632,8 +701,10 @@ class StaticModeMixin:
                                     seed, steps, cfg, sampler_name, scheduler, denoise,
                                     tile_width, tile_height, padding, mask_blur,
                                     force_uniform_tiles, tiled_decode, multi_job_id, enabled_workers,
-                                    all_tiles, num_tiles_per_image):
+                                    all_tiles, num_tiles_per_image,
+                                    mode_context: StaticModeContext | None = None):
         """Static mode master processing with optional dynamic queue pulling."""
+        context = mode_context or self._build_static_mode_context()
         batch_size = upscaled_image.shape[0]
         total_tiles = batch_size * num_tiles_per_image
         result_images, tile_masks, width, height = self._init_master_static_state(
@@ -646,6 +717,7 @@ class StaticModeMixin:
             tile_width=tile_width,
             tile_height=tile_height,
             mask_blur=mask_blur,
+            mode_context=context,
         )
 
         try:
@@ -675,6 +747,7 @@ class StaticModeMixin:
                 tiled_decode=tiled_decode,
                 width=width,
                 height=height,
+                mode_context=context,
             )
 
             collected_tasks = self._collect_remaining_static_tiles(
@@ -704,6 +777,7 @@ class StaticModeMixin:
                 tiled_decode=tiled_decode,
                 width=width,
                 height=height,
+                mode_context=context,
             )
 
             self._blend_static_collected_tiles(
@@ -716,6 +790,7 @@ class StaticModeMixin:
                 tile_width=tile_width,
                 tile_height=tile_height,
                 padding=padding,
+                mode_context=context,
             )
 
             result_tensor = self._result_images_to_tensor(result_images, batch_size, upscaled_image)

upscale/result_collector.py

@@ -16,7 +16,7 @@ class ResultCollectorMixin:
     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.
+    - `self.async_yield(...)` optional helper from WorkerCommsMixin.
     """
 
     def _log_worker_timeout_status(self, job_data, current_time: float, multi_job_id: str) -> list[str]:
@@ -196,6 +196,10 @@ class ResultCollectorMixin:
             if isinstance(job_data, ImageJobState):
                 job_data.completed_images[image_idx] = image_pil
 
+    async def mark_image_completed(self, multi_job_id, image_idx, image_pil):
+        """Public image-completion API."""
+        await self._mark_image_completed(multi_job_id, image_idx, image_pil)
+
     async def _async_collect_dynamic_images(self, multi_job_id, remaining_to_collect, num_workers, batch_size, master_processed_count):
         """Collect remaining processed images from workers."""
         _ = master_processed_count
@@ -207,3 +211,20 @@ class ResultCollectorMixin:
             remaining_to_collect=remaining_to_collect,
             batch_size=batch_size,
         )
+
+    async def collect_dynamic_images(
+        self,
+        multi_job_id,
+        remaining_to_collect,
+        num_workers,
+        batch_size,
+        master_processed_count,
+    ):
+        """Public dynamic image-collection API."""
+        return await self._async_collect_dynamic_images(
+            multi_job_id,
+            remaining_to_collect,
+            num_workers,
+            batch_size,
+            master_processed_count,
+        )

upscale/tile_ops.py

@@ -371,6 +371,10 @@ class TileOpsMixin:
         
         return positive_sliced, negative_sliced
 
+    def slice_conditioning(self, positive, negative, batch_idx):
+        """Public conditioning-slice API."""
+        return self._slice_conditioning(positive, negative, batch_idx)
+
     def _process_and_blend_tile(self, tile_idx, tile_pos, upscaled_image, result_image,
                                model, positive, negative, vae, seed, steps, cfg,
                                sampler_name, scheduler, denoise, tile_width, tile_height,
@@ -399,6 +403,59 @@ class TileOpsMixin:
         
         return result_image
 
+    def process_and_blend_tile(
+        self,
+        tile_idx,
+        tile_pos,
+        upscaled_image,
+        result_image,
+        model,
+        positive,
+        negative,
+        vae,
+        seed,
+        steps,
+        cfg,
+        sampler_name,
+        scheduler,
+        denoise,
+        tile_width,
+        tile_height,
+        padding,
+        mask_blur,
+        image_width,
+        image_height,
+        force_uniform_tiles,
+        tiled_decode,
+        batch_idx: int = 0,
+    ):
+        """Public tile-processing API."""
+        return self._process_and_blend_tile(
+            tile_idx,
+            tile_pos,
+            upscaled_image,
+            result_image,
+            model,
+            positive,
+            negative,
+            vae,
+            seed,
+            steps,
+            cfg,
+            sampler_name,
+            scheduler,
+            denoise,
+            tile_width,
+            tile_height,
+            padding,
+            mask_blur,
+            image_width,
+            image_height,
+            force_uniform_tiles,
+            tiled_decode,
+            batch_idx=batch_idx,
+        )
+
     def _process_single_tile(self, global_idx, num_tiles_per_image, upscaled_image, all_tiles,
                                   model, positive, negative, vae, seed, steps, cfg, sampler_name,
                                   scheduler, denoise, tiled_decode, tile_width, tile_height, padding,

upscale/worker_comms.py

@@ -1,25 +1,45 @@
 import asyncio, io, json, time
-from typing import Any
+from dataclasses import dataclass
+from typing import Any, Literal
 import aiohttp
 from PIL import Image
 from ..utils.logging import debug_log, log
+from ..utils.auth import distributed_auth_headers
+from ..utils.config import load_config
 from ..utils.network import get_client_session
 from ..utils.constants import TILE_SEND_TIMEOUT
-from ..utils.usdu_managment import MAX_PAYLOAD_SIZE, send_heartbeat_to_master
+from ..utils.usdu_management import MAX_PAYLOAD_SIZE, send_heartbeat_to_master
 from ..utils.image import tensor_to_pil
 
 
+WorkAssignmentKind = Literal["image", "tile", "none"]
+
+
+@dataclass(frozen=True, slots=True)
+class WorkAssignment:
+    """Canonical typed representation of a single work-item assignment."""
+
+    kind: WorkAssignmentKind
+    task_idx: int | None
+    estimated_remaining: int = 0
+    batched_static: bool = False
+
+
 class WorkerCommsMixin:
-    async def _send_heartbeat_to_master(
+    @staticmethod
+    def _master_auth_headers() -> dict[str, str]:
+        return distributed_auth_headers(load_config())
+
+    async def send_heartbeat(
         self,
         multi_job_id: str,
         master_url: str,
         worker_id: str,
     ) -> None:
-        """Proxy heartbeat helper used by worker processing mixins."""
+        """Send worker heartbeat to the master."""
         await send_heartbeat_to_master(multi_job_id, master_url, worker_id)
 
-    async def send_tiles_batch_to_master(
+    async def send_tiles_batch(
         self,
         processed_tiles: list[dict[str, Any]],
         multi_job_id: str,
@@ -63,12 +83,8 @@ class WorkerCommsMixin:
         i = 0
         chunk_index = 0
         while i < total_tiles:
-            data = aiohttp.FormData()
-            data.add_field('multi_job_id', multi_job_id)
-            data.add_field('worker_id', str(worker_id))
-            data.add_field('padding', str(padding))
-
             metadata = []
+            chunk_images: list[tuple[int, int, bytes]] = []
             used = 0
             j = i
             while j < total_tiles:
@@ -80,7 +96,7 @@ class WorkerCommsMixin:
                     break
                 # Accept this tile in this chunk
                 metadata.append(meta)
-                data.add_field(f'tile_{j - i}', io.BytesIO(img_bytes), filename=f'tile_{j}.png', content_type='image/png')
+                chunk_images.append((j - i, j, img_bytes))
                 used += len(img_bytes) + overhead
                 j += 1
 
@@ -89,14 +105,28 @@ class WorkerCommsMixin:
                 # Single oversized tile, send anyway
                 meta = encoded[j]['meta']
                 metadata.append(meta)
-                data.add_field('tile_0', io.BytesIO(encoded[j]['bytes']), filename=f'tile_{j}.png', content_type='image/png')
+                chunk_images.append((0, j, encoded[j]['bytes']))
                 j += 1
 
             chunk_size = j - i
             is_chunk_last = (j >= total_tiles)
-            data.add_field('is_last', str(bool(is_final_flush and is_chunk_last)))
-            data.add_field('batch_size', str(chunk_size))
-            data.add_field('tiles_metadata', json.dumps(metadata), content_type='application/json')
+
+            def _build_chunk_form() -> aiohttp.FormData:
+                data = aiohttp.FormData()
+                data.add_field('multi_job_id', multi_job_id)
+                data.add_field('worker_id', str(worker_id))
+                data.add_field('padding', str(padding))
+                data.add_field('is_last', str(bool(is_final_flush and is_chunk_last)))
+                data.add_field('batch_size', str(chunk_size))
+                data.add_field('tiles_metadata', json.dumps(metadata), content_type='application/json')
+                for relative_idx, source_idx, img_bytes in chunk_images:
+                    data.add_field(
+                        f'tile_{relative_idx}',
+                        io.BytesIO(img_bytes),
+                        filename=f'tile_{source_idx}.png',
+                        content_type='image/png',
+                    )
+                return data
 
             # Retry logic with exponential backoff
             max_retries = 5
@@ -105,7 +135,11 @@ class WorkerCommsMixin:
                 try:
                     session = await get_client_session()
                     url = f"{master_url}/distributed/submit_tiles"
-                    async with session.post(url, data=data) as response:
+                    async with session.post(
+                        url,
+                        data=_build_chunk_form(),
+                        headers=self._master_auth_headers(),
+                    ) as response:
                         response.raise_for_status()
                         break
                 except Exception as e:
@@ -130,7 +164,7 @@ class WorkerCommsMixin:
 
         session = await get_client_session()
         url = f"{master_url}/distributed/submit_tiles"
-        async with session.post(url, data=data) as response:
+        async with session.post(url, data=data, headers=self._master_auth_headers()) as response:
             response.raise_for_status()
             debug_log(f"Worker {worker_id} sent static completion signal")
 
@@ -157,7 +191,7 @@ class WorkerCommsMixin:
                 async with session.post(url, json={
                     'worker_id': str(worker_id),
                     'multi_job_id': multi_job_id
-                }) as response:
+                }, headers=self._master_auth_headers()) as response:
                     if response.status == 200:
                         return await response.json()
                     if response.status == 404:
@@ -181,41 +215,83 @@ class WorkerCommsMixin:
 
         return None
 
-    async def _request_image_from_master(self, multi_job_id, master_url, worker_id):
-        """Request an image index to process from master in dynamic mode."""
-        data = await self._request_work_item_from_master(multi_job_id, master_url, worker_id)
+    @staticmethod
+    def _parse_work_assignment(data: dict[str, Any] | None) -> WorkAssignment:
+        """Normalize assignment payloads into a single discriminated contract."""
         if not data:
-            return None, 0
-        image_idx = data.get('image_idx')
-        estimated_remaining = data.get('estimated_remaining', 0)
-        return image_idx, estimated_remaining
+            return WorkAssignment(kind="none", task_idx=None)
 
-    async def _request_tile_from_master(self, multi_job_id, master_url, worker_id):
-        """Request a tile index to process from master in static mode (reusing dynamic infrastructure)."""
+        kind_value = data.get("kind")
+        if kind_value is None:
+            # Backward-compatible parsing for older masters.
+            if "image_idx" in data:
+                kind_value = "image"
+                data = {
+                    "kind": "image",
+                    "task_idx": data.get("image_idx"),
+                    "estimated_remaining": data.get("estimated_remaining", 0),
+                }
+            elif "tile_idx" in data:
+                kind_value = "tile"
+                data = {
+                    "kind": "tile",
+                    "task_idx": data.get("tile_idx"),
+                    "estimated_remaining": data.get("estimated_remaining", 0),
+                    "batched_static": data.get("batched_static", False),
+                }
+            else:
+                kind_value = "none"
+
+        if kind_value not in {"image", "tile", "none"}:
+            return WorkAssignment(kind="none", task_idx=None)
+
+        task_idx_raw = data.get("task_idx")
+        if task_idx_raw is None:
+            task_idx = None
+        else:
+            try:
+                task_idx = int(task_idx_raw)
+            except (TypeError, ValueError):
+                task_idx = None
+
+        try:
+            estimated_remaining = int(data.get("estimated_remaining", 0) or 0)
+        except (TypeError, ValueError):
+            estimated_remaining = 0
+
+        return WorkAssignment(
+            kind=kind_value,
+            task_idx=task_idx,
+            estimated_remaining=estimated_remaining,
+            batched_static=bool(data.get("batched_static", False)),
+        )
+
+    async def request_assignment(self, multi_job_id, master_url, worker_id) -> WorkAssignment:
+        """Request one assignment and parse into the canonical discriminated contract."""
         data = await self._request_work_item_from_master(multi_job_id, master_url, worker_id)
-        if not data:
-            return None, 0, False
-        tile_idx = data.get('tile_idx')
-        estimated_remaining = data.get('estimated_remaining', 0)
-        batched_static = data.get('batched_static', False)
-        return tile_idx, estimated_remaining, batched_static
+        return self._parse_work_assignment(data)
 
-    async def _send_full_image_to_master(self, image_pil, image_idx, multi_job_id, 
-                                        master_url, worker_id, is_last):
+    async def send_full_image(self, image_pil, image_idx, multi_job_id, 
+                              master_url, worker_id, is_last):
         """Send a processed full image back to master in dynamic mode."""
         # Serialize image to PNG
         byte_io = io.BytesIO()
         image_pil.save(byte_io, format='PNG', compress_level=0)
-        byte_io.seek(0)
-        
-        # Prepare form data
-        data = aiohttp.FormData()
-        data.add_field('multi_job_id', multi_job_id)
-        data.add_field('worker_id', str(worker_id))
-        data.add_field('image_idx', str(image_idx))
-        data.add_field('is_last', str(is_last))
-        data.add_field('full_image', byte_io, filename=f'image_{image_idx}.png', 
-                      content_type='image/png')
+        image_bytes = byte_io.getvalue()
+
+        def _build_image_form() -> aiohttp.FormData:
+            data = aiohttp.FormData()
+            data.add_field('multi_job_id', multi_job_id)
+            data.add_field('worker_id', str(worker_id))
+            data.add_field('image_idx', str(image_idx))
+            data.add_field('is_last', str(is_last))
+            data.add_field(
+                'full_image',
+                io.BytesIO(image_bytes),
+                filename=f'image_{image_idx}.png',
+                content_type='image/png',
+            )
+            return data
         
         # Retry logic
         max_retries = 5
@@ -226,7 +302,11 @@ class WorkerCommsMixin:
                 session = await get_client_session()
                 url = f"{master_url}/distributed/submit_image"
                 
-                async with session.post(url, data=data) as response:
+                async with session.post(
+                    url,
+                    data=_build_image_form(),
+                    headers=self._master_auth_headers(),
+                ) as response:
                     response.raise_for_status()
                     debug_log(f"Successfully sent image {image_idx} to master")
                     return
@@ -240,7 +320,7 @@ class WorkerCommsMixin:
                     log(f"Failed to send image {image_idx} after {max_retries} attempts: {e}")
                     raise
 
-    async def _send_worker_complete_signal(self, multi_job_id, master_url, worker_id):
+    async def send_worker_complete_signal(self, multi_job_id, master_url, worker_id):
         """Send completion signal to master in dynamic mode."""
         # Send a dummy request with is_last=True
         data = aiohttp.FormData()
@@ -252,16 +332,16 @@ class WorkerCommsMixin:
         session = await get_client_session()
         url = f"{master_url}/distributed/submit_image"
         
-        async with session.post(url, data=data) as response:
+        async with session.post(url, data=data, headers=self._master_auth_headers()) as response:
             response.raise_for_status()
             debug_log(f"Worker {worker_id} sent completion signal")
 
-    async def _check_job_status(self, multi_job_id, master_url):
+    async def check_job_status(self, multi_job_id, master_url):
         """Check if job is ready on the master."""
         try:
             session = await get_client_session()
             url = f"{master_url}/distributed/job_status?multi_job_id={multi_job_id}"
-            async with session.get(url) as response:
+            async with session.get(url, headers=self._master_auth_headers()) as response:
                 if response.status == 200:
                     data = await response.json()
                     return data.get('ready', False)
@@ -270,6 +350,6 @@ class WorkerCommsMixin:
             debug_log(f"Job status check failed: {e}")
             return False
 
-    async def _async_yield(self):
+    async def async_yield(self):
         """Simple async yield to allow event loop processing."""
         await asyncio.sleep(0)

utils/async_helpers.py

@@ -6,7 +6,6 @@ import threading
 import time
 import uuid
 import execution
-import server
 from typing import Optional, Any, Coroutine
 from .network import get_server_loop
 
@@ -43,10 +42,11 @@ def run_async_in_server_loop(coro: Coroutine, timeout: Optional[float] = None) -
     
     # Schedule on server's event loop
     loop = get_server_loop()
-    asyncio.run_coroutine_threadsafe(wrapper(), loop)
+    task_future = asyncio.run_coroutine_threadsafe(wrapper(), loop)
     
     # Wait for completion
     if not event.wait(timeout):
+        task_future.cancel()
         raise TimeoutError(f"Async operation timed out after {timeout} seconds")
     
     if error:
@@ -54,7 +54,10 @@ def run_async_in_server_loop(coro: Coroutine, timeout: Optional[float] = None) -
     return result
 
 
-prompt_server = server.PromptServer.instance
+def _prompt_server_instance():
+    import server
+
+    return server.PromptServer.instance
 
 
 def _summarize_node_errors(node_errors: dict) -> str:
@@ -111,6 +114,7 @@ async def queue_prompt_payload(
     client_id: str | None = None,
 ) -> str:
     """Validate and queue a prompt via ComfyUI's prompt queue."""
+    prompt_server = _prompt_server_instance()
     payload = {"prompt": prompt_obj}
     payload = prompt_server.trigger_on_prompt(payload)
     prompt = payload["prompt"]

utils/auth.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+import hmac
+from typing import Any
+
+
+AUTH_HEADER_NAME = "X-Distributed-Token"
+
+
+def get_distributed_api_token(config: dict[str, Any] | None) -> str:
+    """Return configured distributed API token or empty string when unset."""
+    settings = (config or {}).get("settings", {}) if isinstance(config, dict) else {}
+    token = settings.get("distributed_api_token") or settings.get("api_token")
+    return str(token).strip() if token is not None else ""
+
+
+def distributed_auth_headers(config: dict[str, Any] | None) -> dict[str, str]:
+    """Build outbound auth headers for distributed internal API calls."""
+    token = get_distributed_api_token(config)
+    if not token:
+        return {}
+    return {AUTH_HEADER_NAME: token}
+
+
+def is_authorized_request(request: Any, config: dict[str, Any] | None) -> bool:
+    """Validate a request against the configured shared distributed token."""
+    expected_token = get_distributed_api_token(config)
+    if not expected_token:
+        return True
+
+    provided_token = ""
+    headers = getattr(request, "headers", None)
+    if headers is not None:
+        provided_token = str(headers.get(AUTH_HEADER_NAME, "")).strip()
+
+    return hmac.compare_digest(provided_token, expected_token)

utils/cloudflare/process_reader.py

@@ -62,7 +62,10 @@ class ProcessReader:
             return
 
         loop = self._loop
-        for raw_line in iter(process.stdout.readline, ""):
+        while True:
+            raw_line = process.stdout.readline()
+            if not raw_line:
+                break
             line = self._normalize_line(raw_line)
             if not line:
                 continue

utils/config.py

@@ -9,10 +9,16 @@ from collections.abc import AsyncIterator
 from contextlib import asynccontextmanager
 from functools import lru_cache
 from typing import Any
-from .logging import log
 
 # Import defaults for timeout fallbacks
-from .constants import GPU_CONFIG_FILE, HEARTBEAT_TIMEOUT
+from .constants import (
+    GPU_CONFIG_FILE,
+    HEARTBEAT_TIMEOUT,
+    ORCHESTRATION_MEDIA_SYNC_CONCURRENCY,
+    ORCHESTRATION_MEDIA_SYNC_TIMEOUT,
+    ORCHESTRATION_WORKER_PREP_CONCURRENCY,
+    ORCHESTRATION_WORKER_PROBE_CONCURRENCY,
+)
 
 CONFIG_FILE = GPU_CONFIG_FILE
 
@@ -36,6 +42,11 @@ def _config_lock() -> asyncio.Lock:
 def _config_path() -> str:
     return CONFIG_FILE
 
+
+def _log(message: str) -> None:
+    print(f"[Distributed] {message}")
+
+
 def get_default_config() -> dict[str, Any]:
     """Returns the default configuration dictionary. Single source of truth."""
     return {
@@ -47,10 +58,10 @@ def get_default_config() -> dict[str, Any]:
             "stop_workers_on_master_exit": True,
             "master_delegate_only": False,
             "websocket_orchestration": True,
-            "worker_probe_concurrency": 8,
-            "worker_prep_concurrency": 4,
-            "media_sync_concurrency": 2,
-            "media_sync_timeout_seconds": 120
+            "worker_probe_concurrency": ORCHESTRATION_WORKER_PROBE_CONCURRENCY,
+            "worker_prep_concurrency": ORCHESTRATION_WORKER_PREP_CONCURRENCY,
+            "media_sync_concurrency": ORCHESTRATION_MEDIA_SYNC_CONCURRENCY,
+            "media_sync_timeout_seconds": ORCHESTRATION_MEDIA_SYNC_TIMEOUT,
         },
         "tunnel": {
             "status": "stopped",
@@ -107,7 +118,7 @@ def load_config() -> dict[str, Any]:
                 loaded = json.load(f)
             state.cache = _merge_with_defaults(loaded, get_default_config())
         except Exception as e:
-            log(f"Error loading config, using defaults: {e}")
+            _log(f"Error loading config, using defaults: {e}")
             state.cache = get_default_config()
         state.mtime = mtime
 
@@ -128,8 +139,8 @@ def save_config(config: dict[str, Any]) -> bool:
         try:
             os.unlink(tmp_path)
         except OSError as cleanup_error:
-            log(f"Error removing temporary config file '{tmp_path}': {cleanup_error}")
-        log(f"Error saving config: {e}")
+            _log(f"Error removing temporary config file '{tmp_path}': {cleanup_error}")
+        _log(f"Error saving config: {e}")
         return False
 
 
@@ -149,11 +160,8 @@ def ensure_config_exists() -> None:
     """Creates default config file if it doesn't exist. Used by __init__.py"""
     if not os.path.exists(_config_path()):
         default_config = get_default_config()
-        if save_config(default_config):
-            from .logging import debug_log
-            debug_log("Created default config file")
-        else:
-            log("Could not create default config file")
+        if not save_config(default_config):
+            _log("Could not create default config file")
 
 def get_worker_timeout_seconds(default: int = HEARTBEAT_TIMEOUT) -> int:
     """Return the unified worker timeout (seconds).
@@ -181,3 +189,12 @@ def is_master_delegate_only() -> bool:
         return bool(cfg.get('settings', {}).get('master_delegate_only', False))
     except Exception:
         return False
+
+
+def get_debug_enabled(default: bool = False) -> bool:
+    """Return whether distributed debug logging is enabled in config settings."""
+    try:
+        cfg = load_config()
+        return bool(cfg.get("settings", {}).get("debug", default))
+    except Exception:
+        return bool(default)

utils/logging.py

@@ -1,50 +1,15 @@
 """
 Shared logging utilities for ComfyUI-Distributed.
 """
-import os
-import json
-import time
-from dataclasses import dataclass
-from functools import lru_cache
-from .constants import GPU_CONFIG_FILE
+from .config import get_debug_enabled
 
-_DEBUG_TTL: float = 5.0
-
-
-@dataclass
-class _DebugState:
-    enabled: bool | None = None
-    updated_at: float = 0.0
-
-
-@lru_cache(maxsize=1)
-def _debug_state() -> _DebugState:
-    return _DebugState()
 
 def is_debug_enabled() -> bool:
-    """Check if debug is enabled."""
-    state = _debug_state()
-
-    now = time.monotonic()
-    if state.enabled is not None and (now - state.updated_at) < _DEBUG_TTL:
-        return state.enabled
-
-    enabled = state.enabled if state.enabled is not None else False
-    if os.path.exists(GPU_CONFIG_FILE):
-        try:
-            with open(GPU_CONFIG_FILE, "r", encoding="utf-8") as f:
-                config = json.load(f)
-                enabled = bool(config.get("settings", {}).get("debug", False))
-        except (OSError, json.JSONDecodeError, ValueError) as exc:
-            # Fall back to the previous cached value to avoid flapping behavior.
-            print(
-                "[Distributed] Failed to read debug flag from config; "
-                f"using cached value {enabled}: {exc}"
-                )
-
-    state.enabled = enabled
-    state.updated_at = now
-    return enabled
+    """Check whether debug logging is currently enabled."""
+    try:
+        return get_debug_enabled(default=False)
+    except Exception:
+        return False
 
 def debug_log(message: str) -> None:
     """Log debug messages only if debug is enabled in config."""

utils/network.py

@@ -5,7 +5,6 @@ Network and API utilities for ComfyUI-Distributed.
 """
 import asyncio
 import aiohttp
-import server
 from dataclasses import dataclass
 from functools import lru_cache
 from typing import Any
@@ -23,6 +22,12 @@ class _NetworkState:
 def _network_state() -> _NetworkState:
     return _NetworkState()
 
+
+def _prompt_server_instance() -> Any:
+    import server
+
+    return server.PromptServer.instance
+
 async def get_client_session() -> aiohttp.ClientSession:
     """Get or create a shared aiohttp client session."""
     state = _network_state()
@@ -50,28 +55,45 @@ async def handle_api_error(request: web.Request, error: Any, status: int = 500)
     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)
+        return web.json_response(
+            {
+                "status": "error",
+                "error": messages,
+                "message": "; ".join(messages),
+            },
+            status=status,
+        )
 
     if isinstance(error, DistributedError):
         debug_log(f"API Error [{status}] ({error.__class__.__name__}): {error}")
         return web.json_response(
-            {"error": str(error), "error_type": error.__class__.__name__},
+            {
+                "status": "error",
+                "error": str(error),
+                "message": str(error),
+                "error_type": error.__class__.__name__,
+            },
             status=status,
         )
 
     message = str(error)
     debug_log(f"API Error [{status}]: {message}")
-    return web.json_response({"error": message}, status=status)
+    return web.json_response(
+        {
+            "status": "error",
+            "error": message,
+            "message": message,
+        },
+        status=status,
+    )
 
 def get_server_port() -> int:
     """Get the ComfyUI server port."""
-    import server
-    return server.PromptServer.instance.port
+    return _prompt_server_instance().port
 
 def get_server_loop() -> asyncio.AbstractEventLoop:
     """Get the ComfyUI server event loop."""
-    import server
-    return server.PromptServer.instance.loop
+    return _prompt_server_instance().loop
 
 
 def normalize_host(value: Any) -> Any:
@@ -94,7 +116,7 @@ def build_worker_url(worker: dict[str, Any], endpoint: str = "") -> str:
     port = int(worker.get("port", worker.get("listen_port", 8188)) or 8188)
 
     if not host:
-        host = getattr(server.PromptServer.instance, "address", "127.0.0.1") or "127.0.0.1"
+        host = getattr(_prompt_server_instance(), "address", "127.0.0.1") or "127.0.0.1"
 
     if host.startswith(("http://", "https://")):
         base = host.rstrip("/")
@@ -148,7 +170,7 @@ def build_master_url(
         from .config import load_config
         config = load_config()
 
-    prompt_server_instance = prompt_server_instance or server.PromptServer.instance
+    prompt_server_instance = prompt_server_instance or _prompt_server_instance()
     master_cfg = (config or {}).get("master", {}) or {}
     configured_host = (master_cfg.get("host") or "").strip()
     configured_port = master_cfg.get("port")

utils/parsing.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+from typing import Any
+
+
+def require_bool_literal(value: Any, *, field_name: str = "value") -> bool:
+    """Strict boolean parser for API boundaries."""
+    if isinstance(value, bool):
+        return value
+    if isinstance(value, str):
+        normalized = value.strip().lower()
+        if normalized == "true":
+            return True
+        if normalized == "false":
+            return False
+    raise ValueError(f"Field '{field_name}' must be a boolean literal ('true' or 'false').")
+
+
+def coerce_bool(value: Any, default: bool = False) -> bool:
+    """Permissive boolean parser for internal/runtime values."""
+    if isinstance(value, bool):
+        return value
+    if isinstance(value, str):
+        normalized = value.strip().lower()
+        if normalized in {"1", "true", "yes", "on"}:
+            return True
+        if normalized in {"0", "false", "no", "off", ""}:
+            return False
+        return default
+    if value is None:
+        return default
+    if isinstance(value, (int, float)):
+        return bool(value)
+    return default

utils/process.py

@@ -126,13 +126,23 @@ def is_process_alive(pid: int) -> bool:
 
 def terminate_process(process, timeout=5):
     """Gracefully terminate a process with timeout."""
-    if process.poll() is None:  # Still running
-        process.terminate()
+    wait_timeout = max(float(timeout), 0.1)
+    if process.poll() is not None:
+        return True
+
+    process.terminate()
+    try:
+        process.wait(timeout=wait_timeout)
+        return True
+    except (subprocess.TimeoutExpired, TimeoutError):
+        process.kill()
         try:
-            process.wait(timeout=timeout)
-        except (subprocess.TimeoutExpired, TimeoutError):
-            process.kill()
-            process.wait()
+            process.wait(timeout=wait_timeout)
+            return True
+        except (subprocess.TimeoutExpired, TimeoutError) as exc:
+            raise TimeoutError(
+                f"Process did not terminate within {wait_timeout} seconds after kill."
+            ) from exc
 
 def get_python_executable():
     """Get the Python executable path."""

utils/runtime_state.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass, field
+from typing import Any
+
+
+@dataclass
+class DistributedRuntimeState:
+    """Typed distributed runtime state bound to PromptServer.instance."""
+
+    distributed_pending_jobs: dict[str, asyncio.Queue] = field(default_factory=dict)
+    distributed_jobs_lock: asyncio.Lock = field(default_factory=asyncio.Lock)
+    distributed_job_allowed_workers: dict[str, set[str]] = field(default_factory=dict)
+    distributed_pending_tile_jobs: dict[str, Any] = field(default_factory=dict)
+    distributed_tile_jobs_lock: asyncio.Lock = field(default_factory=asyncio.Lock)
+
+
+def get_prompt_server_instance() -> Any:
+    import server
+
+    return server.PromptServer.instance
+
+
+def ensure_distributed_runtime_state(server_instance: Any | None = None) -> DistributedRuntimeState:
+    """Return canonical distributed runtime state for the server instance."""
+    prompt_server = server_instance or get_prompt_server_instance()
+    state = getattr(prompt_server, "distributed_runtime_state", None)
+    if not isinstance(state, DistributedRuntimeState):
+        state = DistributedRuntimeState()
+        setattr(prompt_server, "distributed_runtime_state", state)
+
+    # Backward-compatible aliases for legacy call sites.
+    prompt_server.distributed_pending_jobs = state.distributed_pending_jobs
+    prompt_server.distributed_jobs_lock = state.distributed_jobs_lock
+    prompt_server.distributed_job_allowed_workers = state.distributed_job_allowed_workers
+    prompt_server.distributed_pending_tile_jobs = state.distributed_pending_tile_jobs
+    prompt_server.distributed_tile_jobs_lock = state.distributed_tile_jobs_lock
+    return state

utils/usdu_management.py

@@ -1,10 +1,6 @@
-"""Backward-compatibility shim for USDU helpers.
+"""Canonical USDU helper exports.
 
-Route handlers and job logic now live in:
-- upscale.job_store
-- upscale.job_timeout
-- upscale.payload_parsers
-- api.usdu_routes
+This module is the stable import target for USDU queue/heartbeat helpers.
 """
 
 from ..upscale.conditioning import clone_conditioning, clone_control_chain
@@ -12,15 +8,17 @@ 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,
+    get_completed_count,
     init_dynamic_job,
+    init_job_queue,
     init_static_job_batched,
+    mark_task_completed,
 )
 from ..upscale.job_timeout import check_and_requeue_timed_out_workers
 from ..upscale.payload_parsers import parse_tiles_from_form
+from .auth import distributed_auth_headers
+from .config import load_config
 from .logging import debug_log
 from .network import get_client_session
 
@@ -28,26 +26,18 @@ from .network import get_client_session
 async def send_heartbeat_to_master(multi_job_id: str, master_url: str, worker_id: str) -> None:
     """Send heartbeat to master from worker-side processing loops."""
     try:
-        data = {'multi_job_id': multi_job_id, 'worker_id': str(worker_id)}
+        data = {"multi_job_id": multi_job_id, "worker_id": str(worker_id)}
         session = await get_client_session()
         url = f"{master_url}/distributed/heartbeat"
-        async with session.post(url, json=data) as response:
+        async with session.post(
+            url,
+            json=data,
+            headers=distributed_auth_headers(load_config()),
+        ) as response:
             response.raise_for_status()
     except Exception as e:
         debug_log(f"Heartbeat failed: {e}")
 
-
-# Backward-compatible aliases for legacy imports.
-_check_and_requeue_timed_out_workers = check_and_requeue_timed_out_workers
-_cleanup_job = cleanup_job
-_drain_results_queue = drain_results_queue
-_get_completed_count = get_completed_count
-_init_job_queue = init_job_queue
-_mark_task_completed = mark_task_completed
-_parse_tiles_from_form = parse_tiles_from_form
-_send_heartbeat_to_master = send_heartbeat_to_master
-
-
 __all__ = [
     "MAX_PAYLOAD_SIZE",
     "check_and_requeue_timed_out_workers",
@@ -58,14 +48,6 @@ __all__ = [
     "mark_task_completed",
     "parse_tiles_from_form",
     "send_heartbeat_to_master",
-    "_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",

utils/usdu_utils.py

@@ -2,31 +2,17 @@ import numpy as np
 from PIL import Image, ImageFilter
 import torch
 import torch.nn.functional as F
-from torchvision.transforms import GaussianBlur
 import math
 from typing import Any
 
+from .image import pil_to_tensor, tensor_to_pil
+
 if (not hasattr(Image, 'Resampling')):  # For older versions of Pillow
     Image.Resampling = Image
 
 BLUR_KERNEL_SIZE = 15
 
 
-def tensor_to_pil(img_tensor: torch.Tensor, batch_index: int = 0) -> Image.Image:
-    # Takes a batch of images in the form of a tensor of shape [batch_size, height, width, channels]
-    # and returns an RGB PIL Image. Assumes channels=3
-    return Image.fromarray((255 * img_tensor[batch_index].cpu().numpy()).astype(np.uint8))
-
-
-def pil_to_tensor(image: Image.Image) -> torch.Tensor:
-    # Takes a PIL image and returns a tensor of shape [1, height, width, channels]
-    image = np.array(image).astype(np.float32) / 255.0
-    image = torch.from_numpy(image).unsqueeze(0)
-    if len(image.shape) == 3:  # If the image is grayscale, add a channel dimension
-        image = image.unsqueeze(-1)
-    return image
-
-
 def controlnet_hint_to_pil(tensor: torch.Tensor, batch_index: int = 0) -> Image.Image:
     return tensor_to_pil(tensor.movedim(1, -1), batch_index)
 

utils/worker_ids.py

@@ -1,20 +1,98 @@
 import json
-from typing import Any
+import re
+
+_LEGACY_INDEX_TOKEN_RE = re.compile(r"^(?:worker_)?\d+$")
 
 
-def parse_enabled_worker_ids(enabled_worker_ids: str | list[Any] | None) -> list[str]:
-    """Parse and de-duplicate enabled worker IDs from a JSON list payload."""
-    try:
-        raw = json.loads(enabled_worker_ids) if isinstance(enabled_worker_ids, str) else enabled_worker_ids
-    except Exception:
-        raw = []
+def _require_canonical_worker_id(value: object) -> str:
+    worker_id = str(value).strip()
+    if not worker_id:
+        raise ValueError("enabled_worker_ids cannot contain blank IDs")
+    if _LEGACY_INDEX_TOKEN_RE.match(worker_id):
+        raise ValueError(
+            f"enabled_worker_ids entry '{worker_id}' uses a legacy index token; "
+            "provide stable explicit worker IDs"
+        )
+    return worker_id
+
+
+def require_enabled_worker_ids(enabled_worker_ids: list[str] | None) -> list[str]:
+    """Strictly validate canonical enabled-worker list."""
+    if enabled_worker_ids is None:
+        return []
+    if not isinstance(enabled_worker_ids, list):
+        raise ValueError("enabled_worker_ids must decode to a list")
 
     workers = []
     seen = set()
-    for worker_id in raw if isinstance(raw, list) else []:
-        worker_id_str = str(worker_id)
+    for worker_id in enabled_worker_ids:
+        worker_id_str = _require_canonical_worker_id(worker_id)
         if worker_id_str in seen:
             continue
         seen.add(worker_id_str)
         workers.append(worker_id_str)
     return workers
+
+
+def decode_enabled_worker_ids(value: str | list[str] | None) -> list[str]:
+    """Decode boundary payload into canonical enabled-worker list."""
+    raw = json.loads(value) if isinstance(value, str) else value
+    return require_enabled_worker_ids(raw)
+
+
+def coerce_enabled_worker_ids(enabled_worker_ids: str | list[str] | None) -> list[str]:
+    """Best-effort boundary decoding for permissive internal call sites."""
+    try:
+        return decode_enabled_worker_ids(enabled_worker_ids)
+    except (TypeError, ValueError, json.JSONDecodeError):
+        return []
+
+
+def parse_worker_index(
+    worker_id: str,
+    enabled_worker_ids: list[str],
+) -> int | None:
+    """Resolve worker index from canonical ID or legacy index token."""
+    workers = list(enabled_worker_ids or [])
+    worker_token = str(worker_id).strip()
+    if not worker_token:
+        return None
+
+    if workers and worker_token in workers:
+        return workers.index(worker_token)
+
+    normalized = worker_token[len("worker_") :] if worker_token.startswith("worker_") else worker_token
+
+    try:
+        parsed_index = int(normalized)
+    except (TypeError, ValueError):
+        return None
+    if parsed_index < 0:
+        return None
+    return parsed_index
+
+
+def resolve_worker_identity(
+    worker_id: str,
+    enabled_worker_ids: list[str],
+) -> tuple[str, int | None]:
+    """Return canonical worker id and resolved index when possible."""
+    worker_token = str(worker_id).strip()
+    worker_index = parse_worker_index(worker_token, enabled_worker_ids)
+    workers = list(enabled_worker_ids or [])
+    if worker_index is None:
+        return worker_token, None
+    if 0 <= worker_index < len(workers):
+        return workers[worker_index], worker_index
+    return worker_token, worker_index
+
+
+def worker_value_key(
+    worker_id: str,
+    enabled_worker_ids: list[str],
+) -> str:
+    """Return a stable key used for worker-value maps."""
+    worker_token, worker_index = resolve_worker_identity(worker_id, enabled_worker_ids)
+    if worker_index is not None:
+        return str(worker_index + 1)
+    return worker_token

web/constants.js

@@ -225,8 +225,6 @@ export const ENDPOINTS = {
 
 export const NODE_CLASSES = {
     DISTRIBUTED_COLLECTOR: 'DistributedCollector',
-    DISTRIBUTED_LIST_SPLITTER: 'DistributedListSplitter',
-    DISTRIBUTED_LIST_COLLECTOR: 'DistributedListCollector',
     DISTRIBUTED_BRANCH: 'DistributedBranch',
     DISTRIBUTED_BRANCH_COLLECTOR: 'DistributedBranchCollector',
     DISTRIBUTED_SEED: 'DistributedSeed',

web/executionUtils.js

@@ -7,13 +7,11 @@ export function setupInterceptor(extension) {
     api.queuePrompt = async (number, prompt) => {
         if (extension.isEnabled) {
             const hasCollector = findNodesByClass(prompt.output, NODE_CLASSES.DISTRIBUTED_COLLECTOR).length > 0;
-            const hasListSplitter = findNodesByClass(prompt.output, NODE_CLASSES.DISTRIBUTED_LIST_SPLITTER).length > 0;
-            const hasListCollector = findNodesByClass(prompt.output, NODE_CLASSES.DISTRIBUTED_LIST_COLLECTOR).length > 0;
             const hasBranch = findNodesByClass(prompt.output, NODE_CLASSES.DISTRIBUTED_BRANCH).length > 0;
             const hasBranchCollector = findNodesByClass(prompt.output, NODE_CLASSES.DISTRIBUTED_BRANCH_COLLECTOR).length > 0;
             const hasDistUpscale = findNodesByClass(prompt.output, NODE_CLASSES.UPSCALE_DISTRIBUTED).length > 0;
 
-            if (hasCollector || hasListSplitter || hasListCollector || hasBranch || hasBranchCollector || hasDistUpscale) {
+            if (hasCollector || hasBranch || hasBranchCollector || hasDistUpscale) {
                 const result = await executeParallelDistributed(extension, prompt);
                 // Immediate status check for instant feedback
                 checkAllWorkerStatuses(extension);
@@ -87,7 +85,6 @@ export async function executeParallelDistributed(extension, promptWrapper) {
             workers: activeWorkers.map((worker) => ({ id: worker.id })),
             client_id: api.clientId,
             delegate_master: Boolean(extension.config?.settings?.master_delegate_only),
-            auto_prepare: true,
             trace_execution_id: traceExecutionId,
         });
         if (queueResponse?.prompt_id) {

workers/__init__.py

@@ -1,15 +1,9 @@
 from .process_manager import WorkerProcessManager
 from functools import lru_cache
 
-# Keep monitor module importable via workers package for subprocess integration.
-from . import worker_monitor as worker_monitor  # noqa: F401
-
-
 @lru_cache(maxsize=1)
 def _worker_manager() -> WorkerProcessManager:
-    manager = WorkerProcessManager()
-    manager.queues = {}
-    return manager
+    return WorkerProcessManager()
 
 def get_worker_manager() -> WorkerProcessManager:
     return _worker_manager()

workers/process/lifecycle.py

@@ -74,6 +74,13 @@ class ProcessLifecycle:
                     os.path.dirname(os.path.dirname(__file__)),
                     "worker_monitor.py",
                 )
+                project_root = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
+                existing_pythonpath = env.get("PYTHONPATH", "")
+                env["PYTHONPATH"] = (
+                    f"{project_root}{os.pathsep}{existing_pythonpath}"
+                    if existing_pythonpath
+                    else project_root
+                )
                 monitored_cmd = [get_python_executable(), monitor_script] + cmd
                 log_handle.write(f"[Worker Monitor] Monitoring master PID: {env['COMFYUI_MASTER_PID']}\n")
                 log_handle.flush()

workers/process/persistence.py

@@ -15,7 +15,7 @@ class ProcessPersistence:
 
         for worker_id, proc_info in managed_processes.items():
             pid = proc_info.get("pid")
-            if pid and self._manager._is_process_running(pid):
+            if pid and self._manager.is_process_running(pid):
                 self._manager.processes[worker_id] = {
                     "pid": pid,
                     "process": None,
@@ -33,7 +33,7 @@ class ProcessPersistence:
         managed_processes = {}
 
         for worker_id, proc_info in self._manager.processes.items():
-            is_running, _ = self._manager._check_worker_process(worker_id, proc_info)
+            is_running, _ = self._manager.check_worker_process(worker_id, proc_info)
             if not is_running:
                 continue
             managed_processes[worker_id] = {

workers/process_manager.py

@@ -5,10 +5,11 @@ from .process.root_discovery import ComfyRootDiscovery
 
 
 class WorkerProcessManager:
-    """Thin composition wrapper around worker process subsystems."""
+    """Worker process service composed from launch, lifecycle, and persistence subsystems."""
 
     def __init__(self):
         self.processes = {}
+        self.queues = {}
         self._root_discovery = ComfyRootDiscovery()
         self._launch_builder = LaunchCommandBuilder()
         self._lifecycle = ProcessLifecycle(self)
@@ -18,9 +19,6 @@ class WorkerProcessManager:
     def find_comfy_root(self):
         return self._root_discovery.find_comfy_root()
 
-    def _find_windows_terminal(self):
-        return self._launch_builder._find_windows_terminal()
-
     def build_launch_command(self, worker_config, comfy_root):
         return self._launch_builder.build_launch_command(worker_config, comfy_root)
 
@@ -42,11 +40,11 @@ class WorkerProcessManager:
     def save_processes(self):
         return self._persistence.save_processes()
 
-    def _is_process_running(self, pid):
+    def is_process_running(self, pid):
         return self._lifecycle._is_process_running(pid)
 
-    def _check_worker_process(self, worker_id, proc_info):
+    def check_worker_process(self, worker_id, proc_info):
         return self._lifecycle._check_worker_process(worker_id, proc_info)
 
-    def _kill_process_tree(self, pid):
+    def kill_process_tree(self, pid):
         return self._lifecycle._kill_process_tree(pid)

workers/startup.py

@@ -47,7 +47,7 @@ def auto_launch_workers() -> None:
                     # Check if already running
                     if str(worker_id) in wm.processes:
                         proc_info = wm.processes[str(worker_id)]
-                        if wm._is_process_running(proc_info['pid']):
+                        if wm.is_process_running(proc_info['pid']):
                             debug_log(f"Worker {worker_name} already running, skipping")
                             continue
                     

workers/worker_monitor.py

@@ -6,53 +6,17 @@ and terminates the worker if the master dies.
 import os
 import sys
 import time
-import subprocess
 import platform
 import signal
 from typing import Any
 
 try:
+    from ..utils.process import is_process_alive, launch_process_with_timeout, terminate_process
+    from ..utils.constants import WORKER_CHECK_INTERVAL, PROCESS_TERMINATION_TIMEOUT
+except ImportError:
+    # Fallback when executed as a standalone script.
     from utils.process import is_process_alive, launch_process_with_timeout, terminate_process
     from utils.constants import WORKER_CHECK_INTERVAL, PROCESS_TERMINATION_TIMEOUT
-except ImportError:
-    # Fallback if running from different context
-    def is_process_alive(pid: int | str) -> bool:
-        """Best-effort fallback process liveness probe."""
-        try:
-            pid = int(pid)
-        except (TypeError, ValueError):
-            return False
-
-        if os.name == "nt":
-            try:
-                output = subprocess.check_output(
-                    ["tasklist", "/FI", f"PID eq {pid}"],
-                    stderr=subprocess.DEVNULL,
-                    text=True,
-                    timeout=5.0,
-                )
-            except Exception:
-                return False
-            return str(pid) in output
-
-        try:
-            os.kill(pid, 0)
-            return True
-        except OSError:
-            return False
-    
-    def launch_process_with_timeout(
-        command: list[str],
-        timeout_seconds: float = 10.0,
-        **kwargs: Any,
-    ):
-        _ = kwargs
-        raise RuntimeError(
-            "launch_process_with_timeout unavailable: utils.process import failed"
-        )
-
-    WORKER_CHECK_INTERVAL = 2.0
-    PROCESS_TERMINATION_TIMEOUT = 5.0
 
 def monitor_and_run(master_pid: int, command: list[str]) -> int:
     """Run command and monitor master process; return process exit code."""
@@ -98,15 +62,8 @@ def monitor_and_run(master_pid: int, command: list[str]) -> int:
         if worker_process.poll() is None:  # Still running
             try:
                 terminate_process(worker_process, timeout=PROCESS_TERMINATION_TIMEOUT)
-            except NameError:
-                # Fallback if terminate_process wasn't imported
-                worker_process.terminate()
-                try:
-                    worker_process.wait(timeout=PROCESS_TERMINATION_TIMEOUT)
-                except subprocess.TimeoutExpired:
-                    print("[Distributed] Worker didn't terminate gracefully, forcing kill...")
-                    worker_process.kill()
-                    worker_process.wait()
+            except TimeoutError:
+                print("[Distributed] Worker did not terminate within timeout.")
         
         print("[Distributed] Worker terminated.")